Skip to content

Commit

Permalink
Merge pull request #385 from NNPDF/ome_rust
Browse files Browse the repository at this point in the history
Rustify ome.u.s.{as1, as2}
felixhekhorn authored Oct 14, 2024
2 parents d8963ef + 8896dea commit 972edc4
Showing 26 changed files with 1,488 additions and 236 deletions.
4 changes: 2 additions & 2 deletions crates/dekoder/src/eko.rs
Original file line number Diff line number Diff line change
@@ -92,7 +92,7 @@ impl EKO {
pub fn write(&self, dst: PathBuf) -> Result<()> {
self.assert_working_dir()?;
// create writer
let dst_file = File::create(&dst)?;
let dst_file = File::create(dst)?;
let dst_file = BufWriter::with_capacity(TAR_WRITER_CAPACITY, dst_file);
let mut ar = tar::Builder::new(dst_file);
// do it!
@@ -101,7 +101,7 @@ impl EKO {

/// Extract tar file from `src` to `dst`.
pub fn extract(src: PathBuf, dst: PathBuf) -> Result<Self> {
let mut ar = tar::Archive::new(File::open(&src)?);
let mut ar = tar::Archive::new(File::open(src)?);
ar.unpack(&dst)?;
Self::load_opened(dst)
}
97 changes: 75 additions & 22 deletions crates/eko/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,35 @@
//! Interface to the eko Python package.
use ekore::harmonics::cache::Cache;
use num::Complex;
use std::ffi::c_void;

pub mod bib;
pub mod mellin;

/// Wrapper to pass arguments back to Python
struct RawCmplx {
re: Vec<f64>,
im: Vec<f64>,
}

/// Map tensors to c-ordered list
fn unravel<const DIM: usize>(res: Vec<[[Complex<f64>; DIM]; DIM]>, order_qcd: usize) -> RawCmplx {
let mut target = RawCmplx {
re: Vec::<f64>::new(),
im: Vec::<f64>::new(),
};
for obj in res.iter().take(order_qcd) {
for col in obj.iter().take(DIM) {
for el in col.iter().take(DIM) {
target.re.push(el.re);
target.im.push(el.im);
}
}
}
target
}

/// QCD intergration kernel inside quad.
///
/// # Safety
@@ -14,44 +38,66 @@ pub mod mellin;
pub unsafe extern "C" fn rust_quad_ker_qcd(u: f64, rargs: *mut c_void) -> f64 {
let args = *(rargs as *mut QuadQCDargs);
let is_singlet = (100 == args.mode0) || (21 == args.mode0) || (90 == args.mode0);
// prepare gamma
// prepare Mellin stuff
let path = mellin::TalbotPath::new(u, args.logx, is_singlet);
let jac = path.jac() * path.prefactor();
let mut c = Cache::new(path.n());
let mut re = Vec::<f64>::new();
let mut im = Vec::<f64>::new();
if is_singlet {
let res = ekore::anomalous_dimensions::unpolarized::spacelike::gamma_singlet_qcd(
let mut raw = RawCmplx {
re: Vec::<f64>::new(),
im: Vec::<f64>::new(),
};

if args.is_ome {
if is_singlet {
raw = unravel(
ekore::operator_matrix_elements::unpolarized::spacelike::A_singlet(
args.order_qcd,
&mut c,
args.nf,
args.L,
),
args.order_qcd,
);
} else {
raw = unravel(
ekore::operator_matrix_elements::unpolarized::spacelike::A_non_singlet(
args.order_qcd,
&mut c,
args.nf,
args.L,
),
args.order_qcd,
);
}
} else if is_singlet {
raw = unravel(
ekore::anomalous_dimensions::unpolarized::spacelike::gamma_singlet_qcd(
args.order_qcd,
&mut c,
args.nf,
),
args.order_qcd,
&mut c,
args.nf,
);
for gamma_s in res.iter().take(args.order_qcd) {
for col in gamma_s.iter().take(2) {
for el in col.iter().take(2) {
re.push(el.re);
im.push(el.im);
}
}
}
} else {
// we can not do 1D
let res = ekore::anomalous_dimensions::unpolarized::spacelike::gamma_ns_qcd(
args.order_qcd,
args.mode0,
&mut c,
args.nf,
);
for el in res.iter().take(args.order_qcd) {
re.push(el.re);
im.push(el.im);
raw.re.push(el.re);
raw.im.push(el.im);
}
}

// pass on
(args.py)(
re.as_ptr(),
im.as_ptr(),
c.n.re,
c.n.im,
raw.re.as_ptr(),
raw.im.as_ptr(),
c.n().re,
c.n().im,
jac.re,
jac.im,
args.order_qcd,
@@ -72,6 +118,7 @@ pub unsafe extern "C" fn rust_quad_ker_qcd(u: f64, rargs: *mut c_void) -> f64 {
args.ev_op_max_order_qcd,
args.sv_mode_num,
args.is_threshold,
args.Lsv,
)
}

@@ -101,6 +148,7 @@ type PyQuadKerQCDT = unsafe extern "C" fn(
u8,
u8,
bool,
f64,
) -> f64;

/// Additional integration parameters
@@ -128,6 +176,8 @@ pub struct QuadQCDargs {
pub ev_op_max_order_qcd: u8,
pub sv_mode_num: u8,
pub is_threshold: bool,
pub is_ome: bool,
pub Lsv: f64,
}

/// Empty placeholder function for python callback.
@@ -159,14 +209,15 @@ pub unsafe extern "C" fn my_py(
_ev_op_max_order_qcd: u8,
_sv_mode_num: u8,
_is_threshold: bool,
_lsv: f64,
) -> f64 {
0.
}

/// Return empty additional arguments.
///
/// This is required to make the arguments part of the API, otherwise it won't be added to the compiled
/// package (since it does not appear in the signature of `quad_ker_qcd`).
/// package (since it does not appear in the signature of `rust_quad_ker_qcd`).
///
/// # Safety
/// This is the connection from and back to Python, so we don't know what is on the other side.
@@ -193,5 +244,7 @@ pub unsafe extern "C" fn empty_qcd_args() -> QuadQCDargs {
ev_op_max_order_qcd: 0,
sv_mode_num: 0,
is_threshold: false,
is_ome: false,
Lsv: 0.,
}
}
13 changes: 13 additions & 0 deletions crates/ekore/refs.bib
Original file line number Diff line number Diff line change
@@ -94,3 +94,16 @@ @article{Buza1996wv
pages = "301--320",
year = "1998"
}
@article{Bierenbaum2009zt,
author = "Bierenbaum, Isabella and Blumlein, Johannes and Klein, Sebastian",
title = "{The Gluonic Operator Matrix Elements at O(alpha(s)**2) for DIS Heavy Flavor Production}",
eprint = "0901.0669",
archivePrefix = "arXiv",
primaryClass = "hep-ph",
reportNumber = "DESY-08-187, SFB-CPP-08-107, IFIC-08-68",
doi = "10.1016/j.physletb.2009.01.057",
journal = "Phys. Lett. B",
volume = "672",
pages = "401--406",
year = "2009"
}
2 changes: 1 addition & 1 deletion crates/ekore/src/anomalous_dimensions.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
//! The anomalous dimensions for DGLAP evolution.
//! The anomalous dimensions for |DGLAP| evolution.
pub mod unpolarized;
53 changes: 27 additions & 26 deletions crates/ekore/src/anomalous_dimensions/unpolarized/spacelike/as1.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
//! LO QCD
//! |LO| |QCD|.
use num::complex::Complex;

@@ -9,7 +9,7 @@ use crate::harmonics::cache::{Cache, K};
///
/// Implements Eq. (3.4) of [\[Moch:2004pa\]][crate::bib::Moch2004pa].
pub fn gamma_ns(c: &mut Cache, _nf: u8) -> Complex<f64> {
let N = c.n;
let N = c.n();
let S1 = c.get(K::S1);
let gamma = -(3.0 - 4.0 * S1 + 2.0 / N / (N + 1.0));
CF * gamma
@@ -19,7 +19,7 @@ pub fn gamma_ns(c: &mut Cache, _nf: u8) -> Complex<f64> {
///
/// Implements Eq. (3.5) of [\[Vogt:2004mw\]](crate::bib::Vogt2004mw).
pub fn gamma_qg(c: &mut Cache, nf: u8) -> Complex<f64> {
let N = c.n;
let N = c.n();
let gamma = -(N.powu(2) + N + 2.0) / (N * (N + 1.0) * (N + 2.0));
2.0 * TR * 2.0 * (nf as f64) * gamma
}
@@ -28,7 +28,7 @@ pub fn gamma_qg(c: &mut Cache, nf: u8) -> Complex<f64> {
///
/// Implements Eq. (3.5) of [\[Vogt:2004mw\]](crate::bib::Vogt2004mw).
pub fn gamma_gq(c: &mut Cache, _nf: u8) -> Complex<f64> {
let N = c.n;
let N = c.n();
let gamma = -(N.powu(2) + N + 2.0) / (N * (N + 1.0) * (N - 1.0));
2.0 * CF * gamma
}
@@ -37,7 +37,7 @@ pub fn gamma_gq(c: &mut Cache, _nf: u8) -> Complex<f64> {
///
/// Implements Eq. (3.5) of [\[Vogt:2004mw\]](crate::bib::Vogt2004mw).
pub fn gamma_gg(c: &mut Cache, nf: u8) -> Complex<f64> {
let N = c.n;
let N = c.n();
let S1 = c.get(K::S1);
let gamma = S1 - 1.0 / N / (N - 1.0) - 1.0 / (N + 1.0) / (N + 2.0);
CA * (4.0 * gamma - 11.0 / 3.0) + 4.0 / 3.0 * TR * (nf as f64)
@@ -54,63 +54,64 @@ pub fn gamma_singlet(c: &mut Cache, nf: u8) -> [[Complex<f64>; 2]; 2] {

#[cfg(test)]
mod tests {
use crate::cmplx;
use crate::{anomalous_dimensions::unpolarized::spacelike::as1::*, harmonics::cache::Cache};
use float_cmp::assert_approx_eq;
use super::*;
use crate::harmonics::cache::Cache;
use crate::{assert_approx_eq_cmplx, cmplx};
use num::complex::Complex;
use num::Zero;
const NF: u8 = 5;

#[test]
fn number_conservation() {
const N: Complex<f64> = cmplx![1., 0.];
const N: Complex<f64> = cmplx!(1., 0.);
let mut c = Cache::new(N);
let me = gamma_ns(&mut c, NF);
assert_approx_eq!(f64, me.re, 0., epsilon = 1e-12);
assert_approx_eq!(f64, me.im, 0., epsilon = 1e-12);
assert_approx_eq_cmplx!(f64, me, Complex::zero(), epsilon = 1e-12);
}

#[test]
fn quark_momentum_conservation() {
const N: Complex<f64> = cmplx![2., 0.];
const N: Complex<f64> = cmplx!(2., 0.);
let mut c = Cache::new(N);
let me = gamma_ns(&mut c, NF) + gamma_gq(&mut c, NF);
assert_approx_eq!(f64, me.re, 0., epsilon = 1e-12);
assert_approx_eq!(f64, me.im, 0., epsilon = 1e-12);
assert_approx_eq_cmplx!(f64, me, Complex::zero(), epsilon = 1e-12);
}

#[test]
fn gluon_momentum_conservation() {
const N: Complex<f64> = cmplx![2., 0.];
const N: Complex<f64> = cmplx!(2., 0.);
let mut c = Cache::new(N);
let me = gamma_qg(&mut c, NF) + gamma_gg(&mut c, NF);
assert_approx_eq!(f64, me.re, 0., epsilon = 1e-12);
assert_approx_eq!(f64, me.im, 0., epsilon = 1e-12);
assert_approx_eq_cmplx!(f64, me, Complex::zero(), epsilon = 1e-12);
}

#[test]
fn gamma_qg_() {
const N: Complex<f64> = cmplx![1., 0.];
const N: Complex<f64> = cmplx!(1., 0.);
let mut c = Cache::new(N);
let me = gamma_qg(&mut c, NF);
assert_approx_eq!(f64, me.re, -20. / 3.0, ulps = 32);
assert_approx_eq!(f64, me.im, 0., epsilon = 1e-12);
assert_approx_eq_cmplx!(f64, me, cmplx!(-20. / 3., 0.), ulps = 32, epsilon = 1e-12);
}

#[test]
fn gamma_gq_() {
const N: Complex<f64> = cmplx![0., 1.];
const N: Complex<f64> = cmplx!(0., 1.);
let mut c = Cache::new(N);
let me = gamma_gq(&mut c, NF);
assert_approx_eq!(f64, me.re, 4. / 3.0, ulps = 32);
assert_approx_eq!(f64, me.im, -4. / 3.0, ulps = 32);
assert_approx_eq_cmplx!(f64, me, cmplx!(4. / 3.0, -4. / 3.0), ulps = 32);
}

#[test]
fn gamma_gg_() {
const N: Complex<f64> = cmplx![0., 1.];
const N: Complex<f64> = cmplx!(0., 1.);
let mut c = Cache::new(N);
let me = gamma_gg(&mut c, NF);
assert_approx_eq!(f64, me.re, 5.195725159621, ulps = 32, epsilon = 1e-11);
assert_approx_eq!(f64, me.im, 10.52008856962, ulps = 32, epsilon = 1e-11);
assert_approx_eq_cmplx!(
f64,
me,
cmplx!(5.195725159621, 10.52008856962),
ulps = 32,
epsilon = 1e-11
);
}
}
Loading

0 comments on commit 972edc4

Please sign in to comment.