Merge pull request #385 from NNPDF/ome_rust

Rustify ome.u.s.{as1, as2}

Author: felixhekhorn

crates/dekoder/src/eko.rs

@@ -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)
     }

crates/eko/src/lib.rs

@@ -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.,
     }
 }

crates/ekore/refs.bib

@@ -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"
+}

crates/ekore/src/anomalous_dimensions.rs

@@ -1,3 +1,3 @@
-//! The anomalous dimensions for DGLAP evolution.
+//! The anomalous dimensions for |DGLAP| evolution.
 
 pub mod unpolarized;

crates/ekore/src/anomalous_dimensions/unpolarized/spacelike/as1.rs

@@ -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
+        );
     }
 }