diff --git a/.github/workflows/lha_bot_rust.yml b/.github/workflows/lha_bot_rust.yml index 04d1293de..fc749c9a5 100644 --- a/.github/workflows/lha_bot_rust.yml +++ b/.github/workflows/lha_bot_rust.yml @@ -41,6 +41,8 @@ jobs: - name: Install task runner run: pip install poethepoet - name: Run benchmark + env: + VIRTUAL_ENV: ${{ env.pythonLocation }} run: | ./rustify.sh poe compile diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b96f3206..135e3f84e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,9 +15,19 @@ Items without prefix refer to a global change. ## [Unreleased](https://github.com/NNPDF/eko/compare/v0.15.3...HEAD) +### Added +- rust: Expose `qcd_gamma_singlet`, `qcd_gamma_ns`, `ome_singlet`, and `ome_non_singlet` as a native Python extension module (`ekors`) via PyO3 bindings ([#526](https://github.com/NNPDF/eko/pull/526)) + ### Fixed - rust: Fix and improve publication of Rust crates ([#522](https://github.com/NNPDF/eko/pull/522)) +### Changed +- py: Moved all quad kernels to quad_ker.py ([#526](https://github.com/NNPDF/eko/pull/526)) +- Reorganize Python / Rust interface: Rust is now called inside numba (and not the inverse). This also simplified the Rust patch ([#526](https://github.com/NNPDF/eko/pull/526)) + +### Removed +- rust: Drop CFFI and Python callback interface in favour of PyO3 bindings ([#526](https://github.com/NNPDF/eko/pull/526)) + ## [0.15.3](https://github.com/NNPDF/eko/compare/v0.15.2...v0.15.3) - 2026-02-05 ### Added diff --git a/Cargo.lock b/Cargo.lock index f5dffa896..4143735fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -220,6 +220,8 @@ version = "0.0.1" dependencies = [ "ekore", "num", + "numpy", + "pyo3", ] [[package]] @@ -361,6 +363,12 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "heck" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" + [[package]] name = "ignore" version = "0.4.23" @@ -377,6 +385,15 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + [[package]] name = "libc" version = "0.2.170" @@ -400,6 +417,15 @@ version = "0.4.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + [[package]] name = "log" version = "0.4.26" @@ -431,6 +457,15 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + [[package]] name = "miniz_oxide" version = "0.8.5" @@ -546,12 +581,50 @@ dependencies = [ "autocfg", ] +[[package]] +name = "numpy" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec170733ca37175f5d75a5bea5911d6ff45d2cd52849ce98b685394e4f2f37f4" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "rustc-hash", +] + [[package]] name = "once_cell" version = "1.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "945462a4b81e43c4e3ba96bd7b49d834c6f61198356aa858733bc4acf3cbe62e" +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + [[package]] name = "paste" version = "1.0.15" @@ -609,6 +682,12 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2" +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + [[package]] name = "predicates" version = "3.1.3" @@ -661,6 +740,69 @@ dependencies = [ "pest_derive", ] +[[package]] +name = "pyo3" +version = "0.21.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e00b96a521718e08e03b1a622f01c8a8deb50719335de3f60b3b3950f069d8" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "parking_lot", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.21.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7883df5835fafdad87c0d888b266c8ec0f4c9ca48a5bed6bbb592e8dedee1b50" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.21.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01be5843dc60b916ab4dad1dca6d20b9b4e6ddc8e15f50c47fe6d85f1fb97403" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.21.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77b34069fc0682e11b31dbd10321cbf94808394c56fd996796ce45217dfac53c" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.21.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08260721f32db5e1a5beae69a55553f56b99bd0e1c3e6e0a5e8851a9d0f5a85c" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + [[package]] name = "quote" version = "1.0.38" @@ -714,6 +856,12 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustix" version = "0.38.44" @@ -727,6 +875,12 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + [[package]] name = "same-file" version = "1.0.6" @@ -736,6 +890,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "serde" version = "1.0.218" @@ -767,6 +927,12 @@ dependencies = [ "digest", ] +[[package]] +name = "smallvec" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" + [[package]] name = "syn" version = "2.0.98" @@ -789,6 +955,12 @@ dependencies = [ "xattr", ] +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + [[package]] name = "tempfile" version = "3.17.1" @@ -873,6 +1045,12 @@ version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "00e2473a93778eb0bad35909dff6a10d28e63f792f16ed15e404fca9d5eeedbe" +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + [[package]] name = "version_check" version = "0.9.5" @@ -907,6 +1085,12 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + [[package]] name = "windows-sys" version = "0.59.0" diff --git a/crates/eko/Cargo.toml b/crates/eko/Cargo.toml index 95fd429d6..36fa99bc2 100644 --- a/crates/eko/Cargo.toml +++ b/crates/eko/Cargo.toml @@ -22,3 +22,5 @@ crate-type = ["cdylib"] [dependencies] num = "0.4.1" ekore = { path = "../ekore", version = "0.0.1" } +pyo3 = { version = "0.21", features = ["extension-module"] } +numpy = "0.21" diff --git a/crates/eko/pyproject.toml b/crates/eko/pyproject.toml index e1c7ff6da..128265089 100644 --- a/crates/eko/pyproject.toml +++ b/crates/eko/pyproject.toml @@ -4,13 +4,13 @@ build-backend = "maturin" [project] name = "eko-rs" +version = "0.0.1" requires-python = ">=3.10" classifiers = [ "Programming Language :: Rust", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] -dependencies = ["cffi"] [tool.maturin] -bindings = "cffi" +bindings = "pyo3" diff --git a/crates/eko/src/lib.rs b/crates/eko/src/lib.rs index 242f404b5..541654b98 100644 --- a/crates/eko/src/lib.rs +++ b/crates/eko/src/lib.rs @@ -1,8 +1,15 @@ //! Interface to the eko Python package. +use ekore::anomalous_dimensions::polarized::spacelike as ad_ps; +use ekore::anomalous_dimensions::unpolarized::spacelike as ad_us; +// use ekore::anomalous_dimensions::unpolarized::timelike as ad_ut; TODO use ekore::harmonics::cache::Cache; +// use ekore::operator_matrix_elements::polarized::spacelike as ome_ps; TODO +use ekore::operator_matrix_elements::unpolarized::spacelike as ome_us; +// use ekore::operator_matrix_elements::unpolarized::timelike as ome_ut; TODO use num::Complex; -use std::ffi::c_void; +use numpy::{IntoPyArray, PyArray1}; +use pyo3::prelude::*; pub mod bib; pub mod mellin; @@ -74,334 +81,207 @@ fn unravel_qed_ns(res: Vec>>, order_qcd: usize, order_qed: usiz target } -/// Integration kernel inside quad. -/// -/// # Safety -/// This is the connection from Python, so we don't know what is on the other side. -#[no_mangle] -pub unsafe extern "C" fn rust_quad_ker(u: f64, rargs: *mut c_void) -> f64 { - let args = *(rargs as *mut QuadArgs); +/// Type for list of complex numbers (real + imaginary part) +type ComplexResult = PyResult<(Py>, Py>)>; - let is_singlet = (100 == args.mode0) - || (21 == args.mode0) - || (90 == args.mode0) - || (22 == args.mode0) - || (101 == args.mode0); +// TODO: once the Rust kernels grow we take the computation of the integration path back to here +// in the old implementation it was ~ +// // prepare Mellin stuff +// let path = mellin::TalbotPath::new(u, args.logx, is_singlet); +// let jac = path.jac() * path.prefactor(); - let is_qed_valence = (10200 == args.mode0) || (10204 == args.mode0); - // 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 raw = RawCmplx { - re: Vec::::new(), - im: Vec::::new(), +/// Compute the QCD singlet anomalous dimension matrix in Mellin space. +/// +/// # Errors +/// Returns `PyNotImplementedError` if `is_time_like` is set (both polarized +/// and unpolarized time-like are not yet implemented). +#[pyfunction] +#[allow(clippy::too_many_arguments)] +pub fn qcd_gamma_singlet( + py: Python<'_>, + is_polarized: bool, + is_time_like: bool, + order_qcd: usize, + re_n: f64, + im_n: f64, + nf: u8, + variations: [u8; 4], +) -> ComplexResult { + if is_polarized && is_time_like { + return Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Polarized time-like is not implemented", + )); + } + if is_time_like { + return Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Unpolarized time-like is not yet implemented", + )); + } + let mut c = Cache::new(Complex::new(re_n, im_n)); + let raw = if is_polarized { + unravel( + ad_ps::gamma_singlet_qcd(order_qcd, &mut c, nf, variations), + order_qcd, + ) + } else { + unravel( + ad_us::gamma_singlet_qcd(order_qcd, &mut c, nf, variations), + order_qcd, + ) }; - let n3lo_ad_variation = std::slice::from_raw_parts(args.n3lo_ad_variation, 7) - .try_into() - .unwrap(); + Ok(( + raw.re.into_pyarray_bound(py).unbind(), + raw.im.into_pyarray_bound(py).unbind(), + )) +} - 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 { - if args.order_qed > 0 { - let gamma_singlet_qed = - ekore::anomalous_dimensions::unpolarized::spacelike::gamma_singlet_qed; - raw = unravel_qed( - gamma_singlet_qed( - args.order_qcd, - args.order_qed, - &mut c, - args.nf, - n3lo_ad_variation, - ), - args.order_qcd, - args.order_qed, - ); - } else { - let gamma_singlet_qcd = match args.is_polarized { - true => ekore::anomalous_dimensions::polarized::spacelike::gamma_singlet_qcd, - false => ekore::anomalous_dimensions::unpolarized::spacelike::gamma_singlet_qcd, - }; - raw = unravel( - gamma_singlet_qcd( - args.order_qcd, - &mut c, - args.nf, - n3lo_ad_variation[0..4].try_into().unwrap(), - ), - args.order_qcd, - ); - } - } else if args.order_qed > 0 { - if is_qed_valence { - let gamma_valence_qed = - ekore::anomalous_dimensions::unpolarized::spacelike::gamma_valence_qed; - raw = unravel_qed( - gamma_valence_qed( - args.order_qcd, - args.order_qed, - &mut c, - args.nf, - n3lo_ad_variation[4..7].try_into().unwrap(), - ), - args.order_qcd, - args.order_qed, - ); - } else { - let gamma_ns_qed = ekore::anomalous_dimensions::unpolarized::spacelike::gamma_ns_qed; - raw = unravel_qed_ns( - gamma_ns_qed( - args.order_qcd, - args.order_qed, - args.mode0, - &mut c, - args.nf, - n3lo_ad_variation[4..7].try_into().unwrap(), - ), - args.order_qcd, - args.order_qed, - ); - } - } else { - // we can not do 1D - let gamma_ns_qcd = match args.is_polarized { - true => ekore::anomalous_dimensions::polarized::spacelike::gamma_ns_qcd, - false => ekore::anomalous_dimensions::unpolarized::spacelike::gamma_ns_qcd, - }; - let res = gamma_ns_qcd( - args.order_qcd, - args.mode0, - &mut c, - args.nf, - n3lo_ad_variation[4..7].try_into().unwrap(), - ); - for el in res.iter().take(args.order_qcd) { - raw.re.push(el.re); - raw.im.push(el.im); - } +/// Compute the QCD non-singlet anomalous dimension in Mellin space. +/// +/// # Errors +/// Returns `PyNotImplementedError` for polarized time-like (not implemented) +/// or plain time-like (not yet implemented). +#[pyfunction] +#[allow(clippy::too_many_arguments)] +pub fn qcd_gamma_ns( + py: Python<'_>, + is_polarized: bool, + is_time_like: bool, + order_qcd: usize, + mode0: u16, + re_n: f64, + im_n: f64, + nf: u8, + variations: [u8; 3], + _use_fhmruvv: bool, // TODO +) -> ComplexResult { + if is_polarized && is_time_like { + return Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Polarized time-like is not implemented", + )); + } + if is_time_like { + return Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Unpolarized time-like is not yet implemented", + )); } - // pass on - (args.py)( - raw.re.as_ptr(), - raw.im.as_ptr(), - c.n().re, - c.n().im, - jac.re, - jac.im, - args.order_qcd, - args.order_qed, - is_singlet, - args.mode0, - args.mode1, - args.nf, - args.is_log, - args.logx, - args.areas, - args.areas_x, - args.areas_y, - args.method_num, - args.as1, - args.as0, - args.ev_op_iterations, - args.ev_op_max_order_qcd, - args.sv_mode_num, - args.is_threshold, - args.Lsv, - // additional QED params - args.as_list, - args.as_list_len, - args.mu2_from, - args.mu2_to, - args.a_half, - args.a_half_x, - args.a_half_y, - args.alphaem_running, - ) -} + let mut c = Cache::new(Complex::new(re_n, im_n)); + + let raw: Vec> = if is_polarized { + ad_ps::gamma_ns_qcd(order_qcd, mode0, &mut c, nf, variations) + } else { + ad_us::gamma_ns_qcd(order_qcd, mode0, &mut c, nf, variations) + }; -/// Python callback signature -type PyQuadKerT = unsafe extern "C" fn( - *const f64, // re_gamma - *const f64, // im_gamma - f64, // re_n - f64, // im_n - f64, // re_jac - f64, // im_jac - usize, // order_qcd - usize, // order_qed - bool, // is_singlet - u16, // mode0 - u16, // mode1 - u8, // nf - bool, // is_log - f64, // logx - *const f64, // areas - u8, // areas_x - u8, // areas_y - u8, // method_num - f64, // as1 - f64, // as0 - u8, // ev_op_iterations - u8, // ev_op_max_order_qcd - u8, // sv_mode_num - bool, // is_threshold - f64, // lsv - *const f64, // as_list - u8, // as_list_len - f64, // mu2_from - f64, // mu2_to - *const f64, // a_half - u8, // a_half_x - u8, // a_half_y - bool, // alphaem_running -) -> f64; + let re: Vec = raw.iter().map(|z| z.re).collect(); + let im: Vec = raw.iter().map(|z| z.im).collect(); -/// Additional integration parameters -#[allow(non_snake_case)] -#[repr(C)] -#[derive(Clone, Copy)] -pub struct QuadArgs { - pub order_qcd: usize, - pub order_qed: usize, - pub mode0: u16, - pub mode1: u16, - pub is_polarized: bool, - pub is_time_like: bool, - pub nf: u8, - pub py: PyQuadKerT, - pub is_log: bool, - pub logx: f64, - pub areas: *const f64, - pub areas_x: u8, - pub areas_y: u8, - pub L: f64, - pub method_num: u8, - pub as1: f64, - pub as0: f64, - pub ev_op_iterations: u8, - pub ev_op_max_order_qcd: u8, - pub sv_mode_num: u8, - pub is_threshold: bool, - pub is_ome: bool, - pub Lsv: f64, - // additional param required for QED - pub as_list: *const f64, - pub as_list_len: u8, - pub mu2_from: f64, - pub mu2_to: f64, - pub a_half: *const f64, - pub a_half_x: u8, - pub a_half_y: u8, - pub alphaem_running: bool, - // additional param required for N3LO - pub n3lo_ad_variation: *const u8, + Ok(( + re.into_pyarray_bound(py).unbind(), + im.into_pyarray_bound(py).unbind(), + )) } -/// Empty placeholder function for python callback. +/// Compute the singlet operator matrix element (OME) array in Mellin space. /// -/// # Safety -/// This is the connection back to Python, so we don't know what is on the other side. -pub unsafe extern "C" fn my_py( - _re_gamma: *const f64, - _im_gamma: *const f64, - _re_n: f64, - _im_n: f64, - _re_jac: f64, - _im_jac: f64, - _order_qcd: usize, - _order_qed: usize, - _is_singlet: bool, - _mode0: u16, - _mode1: u16, - _nf: u8, - _is_log: bool, - _logx: f64, - _areas: *const f64, - _areas_x: u8, - _areas_y: u8, - _method_num: u8, - _as1: f64, - _as0: f64, - _ev_op_iterations: u8, - _ev_op_max_order_qcd: u8, - _sv_mode_num: u8, - _is_threshold: bool, - _lsv: f64, - _as_list: *const f64, - _as_list_len: u8, - _mu2_from: f64, - _mu2_to: f64, - _a_half: *const f64, - _a_half_x: u8, - _a_half_y: u8, - _alphaem_running: bool, -) -> f64 { - 0. +/// # Errors +/// Returns a `PyNotImplementedError` in the following cases: +/// * Polarized time-like (not implemented). +/// * Polarized spacelike (not yet implemented). +/// * Unpolarized time-like (not yet implemented). +#[pyfunction] +#[allow(clippy::too_many_arguments)] +pub fn ome_singlet( + py: Python<'_>, + is_polarized: bool, + is_time_like: bool, + order_qcd: usize, + re_n: f64, + im_n: f64, + nf: u8, + l: f64, + _is_msbar: bool, +) -> ComplexResult { + if is_polarized && is_time_like { + return Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Polarized time-like OME singlet is not implemented", + )); + } + if is_polarized { + return Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Polarized spacelike OME singlet is not yet implemented in Rust", + )); + } + if is_time_like { + return Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Unpolarized time-like OME singlet is not yet implemented in Rust", + )); + } + let mut c = Cache::new(Complex::new(re_n, im_n)); + let raw = unravel::<3>(ome_us::A_singlet(order_qcd, &mut c, nf, l), order_qcd); + Ok(( + raw.re.into_pyarray_bound(py).unbind(), + raw.im.into_pyarray_bound(py).unbind(), + )) } -/// 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 `rust_quad_ker`). +/// Compute the non-singlet operator matrix element (OME) array **A** in Mellin space. /// -/// # Safety -/// This is the connection from and back to Python, so we don't know what is on the other side. -#[no_mangle] -pub unsafe extern "C" fn empty_args() -> QuadArgs { - QuadArgs { - order_qcd: 0, - order_qed: 0, - mode0: 0, - mode1: 0, - is_polarized: false, - is_time_like: false, - nf: 0, - py: my_py, - is_log: true, - logx: 0., - areas: [].as_ptr(), - areas_x: 0, - areas_y: 0, - L: 0., - method_num: 0, - as1: 0., - as0: 0., - ev_op_iterations: 0, - ev_op_max_order_qcd: 0, - sv_mode_num: 0, - is_threshold: false, - is_ome: false, - Lsv: 0., - as_list: [].as_ptr(), - as_list_len: 0, - mu2_from: 0., - mu2_to: 0., - a_half: [].as_ptr(), - a_half_x: 0, - a_half_y: 0, - alphaem_running: false, - n3lo_ad_variation: [0, 0, 0, 0, 0, 0, 0].as_ptr(), +/// # Errors +/// Returns a `PyNotImplementedError` in the following cases: +/// * Polarized time-like (not implemented). +/// * Polarized spacelike (not yet implemented). +/// * Unpolarized time-like (not yet implemented). +#[pyfunction] +#[allow(clippy::too_many_arguments)] +pub fn ome_non_singlet( + py: Python<'_>, + is_polarized: bool, + is_time_like: bool, + matching_order_qcd: usize, + re_n: f64, + im_n: f64, + nf: u8, + l: f64, +) -> ComplexResult { + if is_polarized && is_time_like { + return Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Polarized time-like OME non-singlet is not implemented", + )); + } + if is_polarized { + return Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Polarized spacelike OME non-singlet is not yet implemented in Rust", + )); + } + if is_time_like { + return Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Unpolarized time-like OME non-singlet is not yet implemented in Rust", + )); } + let mut c = Cache::new(Complex::new(re_n, im_n)); + let raw = unravel::<2>( + ome_us::A_non_singlet(matching_order_qcd, &mut c, nf, l), + matching_order_qcd, + ); + Ok(( + raw.re.into_pyarray_bound(py).unbind(), + raw.im.into_pyarray_bound(py).unbind(), + )) +} + +/// Python extension module exposing EKO's Rust kernels. +/// +/// Currently exposes: +/// - [`qcd_gamma_singlet`]: QCD singlet anomalous dimension matrix in Mellin space +/// - [`qcd_gamma_ns`]: QCD non-singlet anomalous dimension array in Mellin space +/// - [`ome_singlet`]: Singlet OME in Mellin space +/// - [`ome_non_singlet`]: Non Singlet OME in Mellin space +#[pymodule] +fn ekors(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(qcd_gamma_singlet, m)?)?; + m.add_function(wrap_pyfunction!(qcd_gamma_ns, m)?)?; + m.add_function(wrap_pyfunction!(ome_singlet, m)?)?; + m.add_function(wrap_pyfunction!(ome_non_singlet, m)?)?; + Ok(()) } diff --git a/extras/gsoc/architecture.md b/extras/gsoc/architecture.md index f33855541..91ae70f1b 100644 --- a/extras/gsoc/architecture.md +++ b/extras/gsoc/architecture.md @@ -96,27 +96,29 @@ np.real(ker * integrand) → returned float `quad_ker` is decorated `@nb.njit`; Numba compiles it to machine code. The anomalous dimensions called from within `quad_ker` are also Numba-compiled (via ekore). `scipy` calls the resulting function through Python's normal calling convention on every quadrature node, carrying Python overhead at the entry point. -### 5.2 rust_quad_ker +### 5.2 rust quad_ker ```text -scipy.integrate.quad(LowLevelCallable(rust_quad_ker, &cfg), 0.5, 1-ε) - ↓ scipy's Fortran/C backend calls the C function pointer directly -rust_quad_ker(u, *args) [crates/eko/src/lib.rs] +scipy.integrate.quad(quad_ker_ad, 0.5, 1-ε) + ↓ calls func at each quadrature node +quad_ker_ad(u, order, mode0, ...) [quad_ker.py, @nb.njit] ↓ -ekore Rust → anomalous dimensions +QuadKerBase → integrand ↓ -Python callback → cb_quad_ker_qcd / cb_quad_ker_qed [Numba] +quad_ker_qcd / quad_ker_qed → anomalous dimensions via ekore [Numba] + ↓ nb.objmode → exits Numba's compiled context +ekors.qcd_gamma_singlet / ekors.qcd_gamma_ns [crates/eko/src/lib.rs] ↓ +ekore Rust → anomalous dimensions + ↓ back to Numba via nb.objmode kernels/singlet.py | non_singlet.py | ... → evolution operator matrix ↓ f64 returned to scipy ``` -**Flow:** `scipy → Rust → Numba → Rust → scipy` - -`scipy` skips Python overhead entirely by calling `rust_quad_ker` via a `LowLevelCallable` C function pointer. Rust handles the anomalous dimensions but then delegates back out to Numba callbacks for the evolution operator matrix, before returning to Rust and finally to `scipy`. +**Flow:** `scipy → Numba → Rust → Numba → scipy` -The short-term direction is to invert the middle portion, i.e. moving to `scipy → Numba → Rust → Numba → scipy` to avoid code duplication, keeping the Numba as the primary entry point while Rust handles anomalous dimensions. +`scipy.integrate.quad` receives `quad_ker_ad` — a Numba `@njit`-compiled function wrapped via `functools.partial` — as a plain Python callable. Numba is the primary entry point into the integration kernel. Rust is reached from within Numba via `nb.objmode` blocks, which temporarily exit the compiled context to call the `ekors` PyO3 native extension module for the anomalous dimensions. Control then returns to Numba for the evolution operator matrix computation before the final scalar is handed back to `scipy`. The long-term goal is to grow the Rust portion until it replaces Numba entirely: `scipy → Rust → scipy`. @@ -126,7 +128,7 @@ The long-term goal is to grow the Rust portion until it replaces Numba entirely: `ekore` provides user-facing entry points into the anomalous dimensions and operator matrix elements, for example [`ad.u.s.gamma_ns_qcd`](https://github.com/NNPDF/eko/blob/ba40d2be721b647b82259bb998c69bc7feedd1dd/crates/ekore/src/anomalous_dimensions/unpolarized/spacelike.rs#L22), which assembles the underlying mathematical ingredients (the "bottom" layer of actual math implementations). -The `eko` Python library accesses the `ekore` Rust library through the `eko` Rust library. In the Rust workflow, `lib.rs` acts as the bridge between Python and Rust, with the integrand selection delegated to Rust via `cfg`. +The `eko` Python library accesses the `ekore` Rust library through the `ekors` PyO3 native extension module. `crates/eko/src/lib.rs` exposes four functions (`qcd_gamma_singlet`, `qcd_gamma_ns`, `ome_singlet`, and `ome_non_singlet`) called from within Numba `@njit` functions via `nb.objmode` blocks. The user-facing entry points of `ekore` are the primary focus of [#519](https://github.com/NNPDF/eko/issues/519), which tracks adding and testing C/C++ and Fortran interfaces to the Rust library. diff --git a/poetry.lock b/poetry.lock index f8dfe7456..9737be51e 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.4.1 and should not be changed by hand. [[package]] name = "alabaster" @@ -64,6 +64,7 @@ files = [ {file = "asttokens-2.4.1-py2.py3-none-any.whl", hash = "sha256:051ed49c3dcae8913ea7cd08e46a606dba30b79993209636c4875bc1d637bc24"}, {file = "asttokens-2.4.1.tar.gz", hash = "sha256:b03869718ba9a6eb027e134bfdf69f38a236d681c83c160d510768af11254ba0"}, ] +markers = {main = "extra == \"mark\""} [package.dependencies] six = ">=1.12.0" @@ -445,7 +446,7 @@ files = [ {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] -markers = {main = "(extra == \"mark\" or extra == \"box\") and platform_system == \"Windows\" or sys_platform == \"win32\"", dev = "sys_platform == \"win32\"", docs = "sys_platform == \"win32\"", test = "sys_platform == \"win32\""} +markers = {main = "platform_system == \"Windows\" and (extra == \"mark\" or extra == \"box\") or extra == \"mark\" and sys_platform == \"win32\"", dev = "sys_platform == \"win32\"", docs = "sys_platform == \"win32\"", test = "sys_platform == \"win32\""} [[package]] name = "comm" @@ -825,6 +826,7 @@ files = [ {file = "decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a"}, {file = "decorator-5.2.1.tar.gz", hash = "sha256:65f266143752f734b0a7cc83c46f4618af75b8c5911b00ccb61d0ac9b6da0360"}, ] +markers = {main = "extra == \"mark\""} [[package]] name = "defusedxml" @@ -904,11 +906,11 @@ description = "Backport of PEP 654 (exception groups)" optional = false python-versions = ">=3.7" groups = ["main", "dev", "docs", "test"] -markers = "python_version == \"3.10\"" files = [ {file = "exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598"}, {file = "exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219"}, ] +markers = {main = "extra == \"mark\" and python_version == \"3.10\"", dev = "python_version == \"3.10\"", docs = "python_version == \"3.10\"", test = "python_version == \"3.10\""} [package.dependencies] typing-extensions = {version = ">=4.6.0", markers = "python_version < \"3.13\""} @@ -927,6 +929,7 @@ files = [ {file = "executing-2.2.1-py2.py3-none-any.whl", hash = "sha256:760643d3452b4d777d295bb167ccc74c64a81df23fb5e08eff250c425a4b2017"}, {file = "executing-2.2.1.tar.gz", hash = "sha256:3632cc370565f6648cc328b32435bd120a1e4ebb20c77e3fdde9a13cd1e533c4"}, ] +markers = {main = "extra == \"mark\""} [package.extras] tests = ["asttokens (>=2.1.0)", "coverage", "coverage-enable-subprocess", "ipython", "littleutils", "pytest", "rich ; python_version >= \"3.11\""] @@ -1121,7 +1124,7 @@ version = "3.15" description = "Internationalized Domain Names in Applications (IDNA)" optional = false python-versions = ">=3.8" -groups = ["dev", "docs"] +groups = ["docs"] files = [ {file = "idna-3.15-py3-none-any.whl", hash = "sha256:048adeaf8c2d788c40fee287673ccaa74c24ffd8dcf09ffa555a2fbb59f10ac8"}, {file = "idna-3.15.tar.gz", hash = "sha256:ca962446ea538f7092a95e057da437618e886f4d349216d2b1e294abfdb65fdc"}, @@ -1199,6 +1202,7 @@ files = [ {file = "ipython-8.38.0-py3-none-any.whl", hash = "sha256:750162629d800ac65bb3b543a14e7a74b0e88063eac9b92124d4b2aa3f6d8e86"}, {file = "ipython-8.38.0.tar.gz", hash = "sha256:9cfea8c903ce0867cc2f23199ed8545eb741f3a69420bfcf3743ad1cec856d39"}, ] +markers = {main = "extra == \"mark\""} [package.dependencies] colorama = {version = "*", markers = "sys_platform == \"win32\""} @@ -1254,6 +1258,7 @@ files = [ {file = "jedi-0.19.2-py2.py3-none-any.whl", hash = "sha256:a8ef22bde8490f57fe5c7681a3c83cb58874daf72b4784de3cce5b6ef6edb5b9"}, {file = "jedi-0.19.2.tar.gz", hash = "sha256:4770dc3de41bde3966b02eb84fbcf557fb33cce26ad23da12c742fb50ecb11f0"}, ] +markers = {main = "extra == \"mark\""} [package.dependencies] parso = ">=0.8.4,<0.9.0" @@ -1295,7 +1300,7 @@ files = [ [package.dependencies] attrs = ">=22.2.0" -jsonschema-specifications = ">=2023.03.6" +jsonschema-specifications = ">=2023.3.6" referencing = ">=0.28.4" rpds-py = ">=0.25.0" @@ -1791,6 +1796,7 @@ files = [ {file = "matplotlib_inline-0.2.1-py3-none-any.whl", hash = "sha256:d56ce5156ba6085e00a9d54fead6ed29a9c47e215cd1bba2e976ef39f5710a76"}, {file = "matplotlib_inline-0.2.1.tar.gz", hash = "sha256:e1ee949c340d771fc39e241ea75683deb94762c8fa5f2927ec57c83c4dffa9fe"}, ] +markers = {main = "extra == \"mark\""} [package.dependencies] traitlets = "*" @@ -1798,6 +1804,37 @@ traitlets = "*" [package.extras] test = ["flake8", "nbdime", "nbval", "notebook", "pytest"] +[[package]] +name = "maturin" +version = "1.13.3" +description = "Build and publish crates with pyo3, cffi and uniffi bindings as well as rust binaries as python packages" +optional = false +python-versions = ">=3.7" +groups = ["dev"] +files = [ + {file = "maturin-1.13.3-py3-none-linux_armv6l.whl", hash = "sha256:3cc13929ca82aefa4adbf0f2c35419369796213c6fb0eb24e914945f50ef5d8c"}, + {file = "maturin-1.13.3-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:53b08bd075649ce96513ad9abf241a43cb685ed6e9e7790f8dbc2d66e95d8323"}, + {file = "maturin-1.13.3-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:4cd478e6e4c56251e48ed079b8efd55b30bc5c09cf695a1bdafaeb582ee735a0"}, + {file = "maturin-1.13.3-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:a2675e25f313034ae6f57388cf14818f87d8961c4a96795287f3e155f59beb11"}, + {file = "maturin-1.13.3-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:4667ef609ab446c1b5e0bfe4f9fb99699ab6d8548433f8d1a684256e0b67217f"}, + {file = "maturin-1.13.3-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:3db93337ed97e60ffc878aa8b493cd7ae44d3a5e1a37256db3a4491f57565018"}, + {file = "maturin-1.13.3-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:1cc0a110b224ca90406b668a3e3c1f5a515062e59e26292f6dbaf5fd4909c6f3"}, + {file = "maturin-1.13.3-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:c00ea6428dea17bf616fe93770837634454b28c2de1a876e42ef8036c616079a"}, + {file = "maturin-1.13.3-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:49fd6ab08da28098ccf37afca24cdba72376ba9c1eedf9dd25ff82ed771961ff"}, + {file = "maturin-1.13.3-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:b6741d7bf4af97da937528fd1e523c6ab54f53d9a21870fa735d6e67fd88e273"}, + {file = "maturin-1.13.3-py3-none-win32.whl", hash = "sha256:0ef257e692cc756c87af5bea95ddfe7d3ac49d3376a7a87f728d63f06e7b6f8b"}, + {file = "maturin-1.13.3-py3-none-win_amd64.whl", hash = "sha256:def4a435ea9d2ee93b18ba579dc8c9cf898889a66f312cd379b5e374ec3e3ad6"}, + {file = "maturin-1.13.3-py3-none-win_arm64.whl", hash = "sha256:2389fe92d017cea9d94e521fa0175314a4c52f79a1057b901fbc9f8686ef7d0b"}, + {file = "maturin-1.13.3.tar.gz", hash = "sha256:771e1e9e71a278e56db01552e0d1acfd1464259f9575b6e72842f893cd299079"}, +] + +[package.dependencies] +tomli = {version = ">=1.1.0", markers = "python_full_version < \"3.11.0\""} + +[package.extras] +patchelf = ["patchelf"] +zig = ["ziglang (>=0.10.0)"] + [[package]] name = "mccabe" version = "0.7.0" @@ -2120,8 +2157,8 @@ files = [ [package.dependencies] numpy = [ {version = ">=1.22.4", markers = "python_version < \"3.11\""}, - {version = ">=1.26.0", markers = "python_version >= \"3.12\""}, {version = ">=1.23.2", markers = "python_version == \"3.11\""}, + {version = ">=1.26.0", markers = "python_version >= \"3.12\""}, ] python-dateutil = ">=2.8.2" pytz = ">=2020.1" @@ -2175,6 +2212,7 @@ files = [ {file = "parso-0.8.5-py2.py3-none-any.whl", hash = "sha256:646204b5ee239c396d040b90f9e272e9a8017c630092bf59980beb62fd033887"}, {file = "parso-0.8.5.tar.gz", hash = "sha256:034d7354a9a018bdce352f48b2a8a450f05e9d6ee85db84764e9b6bd96dafe5a"}, ] +markers = {main = "extra == \"mark\""} [package.extras] qa = ["flake8 (==5.0.4)", "mypy (==0.971)", "types-setuptools (==67.2.0.1)"] @@ -2288,11 +2326,11 @@ description = "Pexpect allows easy control of interactive console applications." optional = false python-versions = "*" groups = ["main", "dev", "docs"] -markers = "sys_platform != \"win32\" and sys_platform != \"emscripten\"" files = [ {file = "pexpect-4.9.0-py2.py3-none-any.whl", hash = "sha256:7236d1e080e4936be2dc3e326cec0af72acf9212a7e1d060210e70a47e253523"}, {file = "pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f"}, ] +markers = {main = "extra == \"mark\" and sys_platform != \"win32\" and sys_platform != \"emscripten\"", dev = "sys_platform != \"win32\" and sys_platform != \"emscripten\"", docs = "sys_platform != \"win32\" and sys_platform != \"emscripten\""} [package.dependencies] ptyprocess = ">=0.5" @@ -2451,6 +2489,7 @@ files = [ {file = "prompt_toolkit-3.0.52-py3-none-any.whl", hash = "sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955"}, {file = "prompt_toolkit-3.0.52.tar.gz", hash = "sha256:28cde192929c8e7321de85de1ddbe736f1375148b02f2e17edd840042b1be855"}, ] +markers = {main = "extra == \"mark\""} [package.dependencies] wcwidth = "*" @@ -2497,11 +2536,11 @@ description = "Run a subprocess in a pseudo terminal" optional = false python-versions = "*" groups = ["main", "dev", "docs"] -markers = "sys_platform != \"win32\" and sys_platform != \"emscripten\"" files = [ {file = "ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35"}, {file = "ptyprocess-0.7.0.tar.gz", hash = "sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220"}, ] +markers = {main = "extra == \"mark\" and sys_platform != \"win32\" and sys_platform != \"emscripten\"", dev = "sys_platform != \"win32\" and sys_platform != \"emscripten\"", docs = "sys_platform != \"win32\" and sys_platform != \"emscripten\""} [[package]] name = "pure-eval" @@ -2514,6 +2553,7 @@ files = [ {file = "pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0"}, {file = "pure_eval-0.2.3.tar.gz", hash = "sha256:5f4e983f40564c576c7c8635ae88db5956bb2229d7e9237d03b3c0b0190eaf42"}, ] +markers = {main = "extra == \"mark\""} [package.extras] tests = ["pytest"] @@ -2532,7 +2572,7 @@ files = [ [package.dependencies] latexcodec = ">=1.0.4" -PyYAML = ">=3.01" +PyYAML = ">=3.1" [package.extras] doc = ["sphinx"] @@ -2578,6 +2618,7 @@ files = [ {file = "pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176"}, {file = "pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f"}, ] +markers = {main = "extra == \"mark\" or extra == \"box\""} [package.extras] windows-terminal = ["colorama (>=0.4.6)"] @@ -2610,8 +2651,8 @@ astroid = ">=3.3.8,<=3.4.0.dev0" colorama = {version = ">=0.4.5", markers = "sys_platform == \"win32\""} dill = [ {version = ">=0.2", markers = "python_version < \"3.11\""}, - {version = ">=0.3.7", markers = "python_version >= \"3.12\""}, {version = ">=0.3.6", markers = "python_version == \"3.11\""}, + {version = ">=0.3.7", markers = "python_version >= \"3.12\""}, ] isort = ">=4.2.5,<5.13 || >5.13,<7" mccabe = ">=0.6,<0.8" @@ -3206,6 +3247,7 @@ files = [ {file = "six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274"}, {file = "six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81"}, ] +markers = {main = "extra == \"mark\""} [[package]] name = "snowballstemmer" @@ -3510,6 +3552,7 @@ files = [ {file = "stack_data-0.6.3-py3-none-any.whl", hash = "sha256:d5558e0c25a4cb0853cddad3d77da9891a08cb85dd9f9f91b9f8cd66e511e695"}, {file = "stack_data-0.6.3.tar.gz", hash = "sha256:836a778de4fec4dcd1dcd89ed8abff8a221f58308462e1c4aa2a3cf30148f0b9"}, ] +markers = {main = "extra == \"mark\""} [package.dependencies] asttokens = ">=2.1.0" @@ -3544,7 +3587,7 @@ version = "2.4.0" description = "A lil' TOML parser" optional = false python-versions = ">=3.8" -groups = ["docs", "test"] +groups = ["dev", "docs", "test"] markers = "python_version == \"3.10\"" files = [ {file = "tomli-2.4.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:b5ef256a3fd497d4973c11bf142e9ed78b150d36f5773f1ca6088c230ffc5867"}, @@ -3639,6 +3682,7 @@ files = [ {file = "traitlets-5.14.3-py3-none-any.whl", hash = "sha256:b74e89e397b1ed28cc831db7aea759ba6640cb3de13090ca145426688ff1ac4f"}, {file = "traitlets-5.14.3.tar.gz", hash = "sha256:9ed0579d3502c94b4b3732ac120375cda96f923114522847de4b3bb98b96b6b7"}, ] +markers = {main = "extra == \"mark\""} [package.extras] docs = ["myst-parser", "pydata-sphinx-theme", "sphinx"] @@ -3655,7 +3699,7 @@ files = [ {file = "typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548"}, {file = "typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466"}, ] -markers = {main = "python_version < \"3.12\"", dev = "python_version < \"3.12\"", test = "python_version == \"3.10\""} +markers = {main = "extra == \"mark\" and python_version < \"3.12\"", dev = "python_version < \"3.12\"", test = "python_version == \"3.10\""} [[package]] name = "tzdata" @@ -3721,6 +3765,7 @@ files = [ {file = "wcwidth-0.2.14-py2.py3-none-any.whl", hash = "sha256:a7bb560c8aee30f9957e5f9895805edd20602f2d7f720186dfd906e82b4982e1"}, {file = "wcwidth-0.2.14.tar.gz", hash = "sha256:4d478375d31bc5395a3c55c40ccdf3354688364cd61c4f6adacaa9215d0b3605"}, ] +markers = {main = "extra == \"mark\""} [[package]] name = "webencodings" @@ -3741,4 +3786,4 @@ mark = ["banana-hep", "matplotlib", "pandas", "sqlalchemy"] [metadata] lock-version = "2.1" python-versions = "^3.10,<3.14" -content-hash = "81882f196367ffda15424824878ab1dae1f754a52a24ceb352d5eca54cd6156b" +content-hash = "1f3de3fd1a1f47197340ba386c21990cdf7e77aa28be0593194c57d36729ce50" diff --git a/pyproject.toml b/pyproject.toml index c2e71815b..67dfb3629 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,6 +81,7 @@ ipython = "^8.0" # benchmark virtualenv = "^20.13.2" devtools = "^0.10.0" +maturin = "^1.1" [tool.poetry.group.version.dependencies] tomlkit = "^0.12.5" @@ -147,7 +148,7 @@ docs-clean = { "shell" = "cd doc; make clean" } docs-cleanall = { "shell" = "cd doc; make cleanall" } docs-install-nb = { "shell" = "python -m ipykernel install --user --name=$(basename $(poetry env info -p))" } bump-version = { "shell" = "python crates/bump-versions.py $(git describe --tags)" } -compile = "pip install -e crates/eko/" +compile = "maturin develop --manifest-path crates/eko/Cargo.toml" rdocs.cmd = "cargo doc --workspace --no-deps" rdocs.env = { RUSTDOCFLAGS = "--html-in-header crates/doc-header.html" } rdocs-view = "xdg-open target/doc/ekors/index.html" diff --git a/rustify.sh b/rustify.sh index bb4ae2565..5bcb8518f 100755 --- a/rustify.sh +++ b/rustify.sh @@ -3,11 +3,8 @@ # git diff --merge-base master pyproject.toml > pyproject.toml.patch patch -p1 src/eko/evolution_operator/__init__.py.patch -patch -p1 src/eko/evolution_operator/operator_matrix_element.py.patch -patch -p1 src/eko/evolution_operator/quad_ker.py.patch +patch -p1 MatchingMethods: """Return the matching method. @@ -50,158 +37,6 @@ def matching_method(s: Optional[InversionMethod]) -> MatchingMethods: return MatchingMethods.FORWARD -@nb.njit(cache=True) -def build_ome(A, matching_order, a_s, backward_method): - r"""Construct the matching expansion in :math:`a_s` with the appropriate - method. - - Parameters - ---------- - A : numpy.ndarray - list of |OME| - matching_order : tuple(int,int) - perturbation matching order - a_s : float - strong coupling, needed only for the exact inverse - backward_method : MatchingMethods - empty or method for inverting the matching condition (exact or expanded) - - Returns - ------- - ome : numpy.ndarray - matching operator matrix - """ - # to get the inverse one can use this FORM snippet - # Symbol a; - # NTensor c,d,e; - # Local x=-(a*c+a**2* d + a**3 * e); - # Local bi = 1+x+x**2+x**3; - # Print; - # .end - ome = np.eye(len(A[0]), dtype=np.complex128) - A = A[:, :, :] - A = np.ascontiguousarray(A) - if backward_method is MatchingMethods.BACKWARD_EXPANDED: - # expended inverse - if matching_order[0] >= 1: - ome -= a_s * A[0] - if matching_order[0] >= 2: - ome += a_s**2 * (-A[1] + A[0] @ A[0]) - if matching_order[0] >= 3: - ome += a_s**3 * (-A[2] + A[0] @ A[1] + A[1] @ A[0] - A[0] @ A[0] @ A[0]) - else: - # forward or exact inverse - if matching_order[0] >= 1: - ome += a_s * A[0] - if matching_order[0] >= 2: - ome += a_s**2 * A[1] - if matching_order[0] >= 3: - ome += a_s**3 * A[2] - # need inverse exact ? so add the missing pieces - if backward_method is MatchingMethods.BACKWARD_EXACT: - ome = np.linalg.inv(ome) - return ome - - -@nb.njit(cache=True) -def quad_ker( - u, - order, - mode0, - mode1, - is_log, - logx, - areas, - a_s, - nf, - L, - sv_mode, - Lsv, - backward_method, - is_msbar, - is_polarized, - is_time_like, -): - r"""Raw kernel inside quad. - - Parameters - ---------- - u : float - quad argument - order : tuple(int,int) - perturbation matching order - mode0 : int - pid for first element in the singlet sector - mode1 : int - pid for second element in the singlet sector - is_log : boolean - logarithmic interpolation - logx : float - Mellin inversion point - areas : tuple - basis function configuration - a_s : float - strong coupling, needed only for the exact inverse - nf: int - number of active flavor below threshold - L : float - :math:``\ln(\mu_F^2 / m_h^2)`` - backward_method : InversionMethod or None - empty or method for inverting the matching condition (exact or expanded) - is_msbar: bool - add the |MSbar| contribution - is_polarized : boolean - is polarized evolution ? - is_time_like : boolean - is time-like evolution ? - - Returns - ------- - ker : float - evaluated integration kernel - """ - ker_base = QuadKerBase(u, is_log, logx, mode0) - integrand = ker_base.integrand(areas) - if integrand == 0.0: - return 0.0 - # compute the ome - if ker_base.is_singlet or ker_base.is_QEDsinglet: - indices = {21: 0, 100: 1, 90: 2} - if is_polarized: - if is_time_like: - raise NotImplementedError("Polarized, time-like is not implemented") - A = ome_ps.A_singlet(order, ker_base.n, nf, L) - else: - if is_time_like: - A = ome_ut.A_singlet(order, ker_base.n, L) - else: - A = ome_us.A_singlet(order, ker_base.n, nf, L, is_msbar) - else: - indices = {200: 0, 91: 1} - if is_polarized: - if is_time_like: - raise NotImplementedError("Polarized, time-like is not implemented") - A = ome_ps.A_non_singlet(order, ker_base.n, L) - else: - if is_time_like: - A = ome_ut.A_non_singlet(order, ker_base.n, L) - else: - A = ome_us.A_non_singlet(order, ker_base.n, nf, L) - - # correct for scale variations - if sv_mode == sv.Modes.exponentiated: - A = gamma_variation(A, order, nf, Lsv) - - # build the expansion in alpha_s depending on the strategy - ker = build_ome(A, order, a_s, backward_method) - - # select the needed matrix element - ker = ker[indices[mode0], indices[mode1]] - - # recombine everything - return np.real(ker * integrand) - - class OperatorMatrixElement(Operator): r"""Internal representation of a single |OME|. diff --git a/src/eko/evolution_operator/operator_matrix_element.py.patch b/src/eko/evolution_operator/operator_matrix_element.py.patch deleted file mode 100644 index e8c7b7ba4..000000000 --- a/src/eko/evolution_operator/operator_matrix_element.py.patch +++ /dev/null @@ -1,197 +0,0 @@ -diff --git a/src/eko/evolution_operator/operator_matrix_element.py b/src/eko/evolution_operator/operator_matrix_element.py -index 839e05a1..13ad879e 100644 ---- a/src/eko/evolution_operator/operator_matrix_element.py -+++ b/src/eko/evolution_operator/operator_matrix_element.py -@@ -3,23 +3,19 @@ evolution.""" - - import copy - import enum --import functools - import logging - from typing import Optional - -+import ekors - import numba as nb - import numpy as np - --import ekore.operator_matrix_elements.polarized.space_like as ome_ps --import ekore.operator_matrix_elements.unpolarized.space_like as ome_us --import ekore.operator_matrix_elements.unpolarized.time_like as ome_ut -- - from .. import basis_rotation as br - from .. import scale_variations as sv - from ..io.types import InversionMethod - from ..matchings import Segment --from ..scale_variations.exponentiated import gamma_variation --from . import Managers, Operator, QuadKerBase -+from . import Managers, Operator -+from .quad_ker import cb_quad_ker_ome - - logger = logging.getLogger(__name__) - -@@ -79,8 +75,6 @@ def build_ome(A, matching_order, a_s, backward_method): - # Print; - # .end - ome = np.eye(len(A[0]), dtype=np.complex128) -- A = A[:, :, :] -- A = np.ascontiguousarray(A) - if backward_method is MatchingMethods.BACKWARD_EXPANDED: - # expended inverse - if matching_order[0] >= 1: -@@ -103,105 +97,6 @@ def build_ome(A, matching_order, a_s, backward_method): - return ome - - --@nb.njit(cache=True) --def quad_ker( -- u, -- order, -- mode0, -- mode1, -- is_log, -- logx, -- areas, -- a_s, -- nf, -- L, -- sv_mode, -- Lsv, -- backward_method, -- is_msbar, -- is_polarized, -- is_time_like, --): -- r"""Raw kernel inside quad. -- -- Parameters -- ---------- -- u : float -- quad argument -- order : tuple(int,int) -- perturbation matching order -- mode0 : int -- pid for first element in the singlet sector -- mode1 : int -- pid for second element in the singlet sector -- is_log : boolean -- logarithmic interpolation -- logx : float -- Mellin inversion point -- areas : tuple -- basis function configuration -- a_s : float -- strong coupling, needed only for the exact inverse -- nf: int -- number of active flavor below threshold -- L : float -- :math:``\ln(\mu_F^2 / m_h^2)`` -- backward_method : InversionMethod or None -- empty or method for inverting the matching condition (exact or expanded) -- is_msbar: bool -- add the |MSbar| contribution -- is_polarized : boolean -- is polarized evolution ? -- is_time_like : boolean -- is time-like evolution ? -- -- Returns -- ------- -- ker : float -- evaluated integration kernel -- """ -- ker_base = QuadKerBase(u, is_log, logx, mode0) -- integrand = ker_base.integrand(areas) -- if integrand == 0.0: -- return 0.0 -- # compute the ome -- if ker_base.is_singlet or ker_base.is_QEDsinglet: -- indices = {21: 0, 100: 1, 90: 2} -- if is_polarized: -- if is_time_like: -- raise NotImplementedError("Polarized, time-like is not implemented") -- A = ome_ps.A_singlet(order, ker_base.n, nf, L) -- else: -- if is_time_like: -- A = ome_ut.A_singlet(order, ker_base.n, L) -- else: -- A = ome_us.A_singlet(order, ker_base.n, nf, L, is_msbar) -- else: -- indices = {200: 0, 91: 1} -- if is_polarized: -- if is_time_like: -- raise NotImplementedError("Polarized, time-like is not implemented") -- A = ome_ps.A_non_singlet(order, ker_base.n, L) -- else: -- if is_time_like: -- A = ome_ut.A_non_singlet(order, ker_base.n, L) -- else: -- A = ome_us.A_non_singlet(order, ker_base.n, nf, L) -- -- # correct for scale variations -- if sv_mode == sv.Modes.exponentiated: -- A = gamma_variation(A, order, nf, Lsv) -- -- # build the expansion in alpha_s depending on the strategy -- ker = build_ome(A, order, a_s, backward_method) -- -- # select the needed matrix element -- ker = ker[indices[mode0], indices[mode1]] -- -- # recombine everything -- return np.real(ker * integrand) -- -- - class OperatorMatrixElement(Operator): - r"""Internal representation of a single |OME|. - -@@ -300,41 +195,15 @@ class OperatorMatrixElement(Operator): - ) - return labels - -- def quad_ker(self, label, logx, areas): -- """Return partially initialized integrand function. -- -- Parameters -- ---------- -- label: tuple -- operator element pids -- logx: float -- Mellin inversion point -- areas : tuple -- basis function configuration -- -- Returns -- ------- -- functools.partial -- partially initialized integration kernel -- """ -- return functools.partial( -- quad_ker, -- order=self.order, -- mode0=label[0], -- mode1=label[1], -- is_log=self.int_disp.log, -- logx=logx, -- areas=areas, -- a_s=self.a_s, -- nf=self.nf, -- L=self.L, -- sv_mode=self.sv_mode, -- Lsv=np.log(self.xif2), -- backward_method=self.backward_method, -- is_msbar=self.is_msbar, -- is_polarized=self.config["polarized"], -- is_time_like=self.config["time_like"], -- ) -+ def update_cfg(self, cfg): -+ """Adjust integration config.""" -+ cfg.is_ome = True -+ cfg.py = ekors.ffi.cast("void *", cb_quad_ker_ome.address) -+ cfg.L = self.L -+ cfg.as1 = self.a_s -+ cfg.as0 = 0.0 -+ cfg.Lsv = np.log(self.xif2) -+ cfg.method_num = self.backward_method - - @property - def a_s(self): diff --git a/src/eko/evolution_operator/quad_ker.py b/src/eko/evolution_operator/quad_ker.py index bd40d57b5..bc23643cd 100644 --- a/src/eko/evolution_operator/quad_ker.py +++ b/src/eko/evolution_operator/quad_ker.py @@ -1,13 +1,20 @@ r"""Integration kernels.""" +import enum import logging import numba as nb import numpy as np -from .. import interpolation +import ekore.anomalous_dimensions.polarized.space_like as ad_ps +import ekore.anomalous_dimensions.unpolarized.space_like as ad_us +import ekore.anomalous_dimensions.unpolarized.time_like as ad_ut +import ekore.operator_matrix_elements.polarized.space_like as ome_ps +import ekore.operator_matrix_elements.unpolarized.space_like as ome_us +import ekore.operator_matrix_elements.unpolarized.time_like as ome_ut + +from .. import interpolation, mellin from .. import scale_variations as sv -from ..io.types import InversionMethod from ..kernels import non_singlet as ns from ..kernels import non_singlet_qed as qed_ns from ..kernels import singlet as s @@ -16,6 +23,7 @@ from ..matchings import lepton_number from ..scale_variations import expanded as sv_expanded from ..scale_variations import exponentiated as sv_exponentiated +from ..scale_variations.exponentiated import gamma_variation logger = logging.getLogger(__name__) @@ -43,139 +51,70 @@ def select_singlet_element(ker, mode0, mode1): return ker[j, k] -CB_SIGNATURE = nb.types.double( - nb.types.CPointer(nb.types.double), # re_*_raw - nb.types.CPointer(nb.types.double), # im_*_raw - nb.types.double, # re_n - nb.types.double, # im_n - nb.types.double, # re_jac - nb.types.double, # im_jac - nb.types.uintc, # order_qcd - nb.types.uintc, # order_qed - nb.types.bool_, # is_singlet - nb.types.uintc, # mode0 - nb.types.uintc, # mode1 - nb.types.uintc, # nf - nb.types.bool_, # is_log - nb.types.double, # logx - nb.types.CPointer(nb.types.double), # areas_raw - nb.types.uintc, # areas_x - nb.types.uintc, # areas_y - nb.types.uintc, # method_num - nb.types.double, # as1 - nb.types.double, # as0 - nb.types.uintc, # ev_op_iterations - nb.types.uintc, # ev_op_max_order_qcd - nb.types.uintc, # sv_mode_num - nb.types.bool_, # is_threshold - nb.types.double, # Lsv - nb.types.CPointer(nb.types.double), # as_list - nb.types.uintc, # as_list_len - nb.types.double, # mu2_from - nb.types.double, # mu2_to - nb.types.CPointer(nb.types.double), # a_half - nb.types.uintc, # a_half_x - nb.types.uintc, # a_half_y - nb.types.bool_, # alphaem_running -) - - -@nb.cfunc( - CB_SIGNATURE, - cache=True, - nopython=True, -) -def cb_quad_ker_qcd( - re_gamma_raw, - im_gamma_raw, - re_n, - im_n, - re_jac, - im_jac, - order_qcd, - _order_qed, - is_singlet, - mode0, - mode1, - nf, - is_log, - logx, - areas_raw, - areas_x, - areas_y, - ev_method, - as1, - as0, - ev_op_iterations, - ev_op_max_order_qcd, - sv_mode, - is_threshold, - Lsv, - _as_list, - _as_list_len, - _mu2_from, - _mu2_to, - _a_half, - _a_half_x, - _a_half_y, - _alphaem_running, -): - """C Callback inside integration kernel.""" - # recover complex variables - n = re_n + im_n * 1j - jac = re_jac + im_jac * 1j - # combute basis functions - areas = nb.carray(areas_raw, (areas_x, areas_y)) - pj = interpolation.evaluate_grid(n, is_log, logx, areas) - order = (order_qcd, 0) - ev_op_max_order = (ev_op_max_order_qcd, 0) - if is_singlet: - # reconstruct singlet matrices - re_gamma_singlet = nb.carray(re_gamma_raw, (order_qcd, 2, 2)) - im_gamma_singlet = nb.carray(im_gamma_raw, (order_qcd, 2, 2)) - gamma_singlet = re_gamma_singlet + im_gamma_singlet * 1j - if sv_mode == sv.Modes.exponentiated: - gamma_singlet = sv_exponentiated.gamma_variation( - gamma_singlet, order, nf, Lsv - ) - # construct eko - ker = s.dispatcher( - order, - ev_method, - gamma_singlet, - as1, - as0, - nf, - ev_op_iterations, - ev_op_max_order, - ) - # scale var expanded is applied on the kernel - if sv_mode == sv.Modes.expanded and not is_threshold: - ker = np.ascontiguousarray( - sv_expanded.singlet_variation(gamma_singlet, as1, order, nf, Lsv, dim=2) - ) @ np.ascontiguousarray(ker) - ker = select_singlet_element(ker, mode0, mode1) +@nb.njit(cache=True) +def select_QEDsinglet_element(ker, mode0, mode1): + """Select element of the QEDsinglet matrix. + + Parameters + ---------- + ker : numpy.ndarray + QEDsinglet integration kernel + mode0 : int + id for first sector element + mode1 : int + id for second sector element + Returns + ------- + ker : complex + QEDsinglet integration kernel element + """ + if mode0 == 21: + index1 = 0 + elif mode0 == 22: + index1 = 1 + elif mode0 == 100: + index1 = 2 else: - # construct non-singlet matrices - re_gamma_ns = nb.carray(re_gamma_raw, order_qcd) - im_gamma_ns = nb.carray(im_gamma_raw, order_qcd) - gamma_ns = re_gamma_ns + im_gamma_ns * 1j - if sv_mode == sv.Modes.exponentiated: - gamma_ns = sv_exponentiated.gamma_variation(gamma_ns, order, nf, Lsv) - # construct eko - ker = ns.dispatcher( - order, - ev_method, - gamma_ns, - as1, - as0, - nf, - ) - if sv_mode == sv.Modes.expanded and not is_threshold: - ker = sv_expanded.non_singlet_variation(gamma_ns, as1, order, nf, Lsv) * ker - # recombine everything - res = ker * pj * jac - return np.real(res) + index1 = 3 + if mode1 == 21: + index2 = 0 + elif mode1 == 22: + index2 = 1 + elif mode1 == 100: + index2 = 2 + else: + index2 = 3 + return ker[index1, index2] + + +@nb.njit(cache=True) +def select_QEDvalence_element(ker, mode0, mode1): + """Select element of the QEDvalence matrix. + + Parameters + ---------- + ker : numpy.ndarray + QEDvalence integration kernel + mode0 : int + id for first sector element + mode1 : int + id for second sector element + Returns + ------- + ker : complex + QEDvalence integration kernel element + """ + index1 = 0 if mode0 == 10200 else 1 + index2 = 0 if mode1 == 10200 else 1 + return ker[index1, index2] + + +class MatchingMethods(enum.IntEnum): + """Enumerate matching methods.""" + + FORWARD = enum.auto() + BACKWARD_EXACT = enum.auto() + BACKWARD_EXPANDED = enum.auto() @nb.njit(cache=True) @@ -191,7 +130,7 @@ def build_ome(A, matching_order, a_s, backward_method): perturbation matching order a_s : float strong coupling, needed only for the exact inverse - backward_method : InversionMethod or None + backward_method : MatchingMethods empty or method for inverting the matching condition (exact or expanded) Returns @@ -209,7 +148,7 @@ def build_ome(A, matching_order, a_s, backward_method): ome = np.eye(len(A[0]), dtype=np.complex128) A = A[:, :, :] A = np.ascontiguousarray(A) - if backward_method is InversionMethod.EXPANDED: + if backward_method is MatchingMethods.BACKWARD_EXPANDED: # expended inverse if matching_order[0] >= 1: ome -= a_s * A[0] @@ -226,249 +165,452 @@ def build_ome(A, matching_order, a_s, backward_method): if matching_order[0] >= 3: ome += a_s**3 * A[2] # need inverse exact ? so add the missing pieces - if backward_method is InversionMethod.EXACT: + if backward_method is MatchingMethods.BACKWARD_EXACT: ome = np.linalg.inv(ome) return ome -@nb.cfunc( - CB_SIGNATURE, - cache=True, - nopython=True, -) -def cb_quad_ker_ome( - re_ome_raw, - im_ome_raw, - re_n, - im_n, - re_jac, - im_jac, - order_qcd, - _order_qed, - is_singlet, - mode0, - mode1, - nf, - is_log, - logx, - areas_raw, - areas_x, - areas_y, - backward_method, - as1, - _as0, - _ev_op_iterations, - _ev_op_max_order_qcd, - sv_mode, - _is_threshold, - Lsv, - _as_list, - _as_list_len, - _mu2_from, - _mu2_to, - _a_half, - _a_half_x, - _a_half_y, - _alphaem_running, -): - """C Callback inside integration kernel.""" - # recover complex variables - n = re_n + im_n * 1j - jac = re_jac + im_jac * 1j - # compute basis functions - areas = nb.carray(areas_raw, (areas_x, areas_y)) - pj = interpolation.evaluate_grid(n, is_log, logx, areas) - order = (order_qcd, 0) - if is_singlet: - indices = {21: 0, 100: 1, 90: 2} - # reconstruct singlet matrices - re_ome_singlet = nb.carray(re_ome_raw, (order_qcd, 3, 3)) - im_ome_singlet = nb.carray(im_ome_raw, (order_qcd, 3, 3)) - A = re_ome_singlet + im_ome_singlet * 1j - else: - indices = {200: 0, 91: 1} - # construct non-singlet matrices - re_ome_ns = nb.carray(re_ome_raw, (order_qcd, 2, 2)) - im_ome_ns = nb.carray(im_ome_raw, (order_qcd, 2, 2)) - A = re_ome_ns + im_ome_ns * 1j +spec = [ + ("is_singlet", nb.boolean), + ("is_QEDsinglet", nb.boolean), + ("is_QEDvalence", nb.boolean), + ("is_log", nb.boolean), + ("logx", nb.float64), + ("u", nb.float64), +] - # correct for scale variations - if sv_mode == sv.Modes.exponentiated: - A = sv.exponentiated.gamma_variation(A, order, nf, Lsv) - # build the expansion in alpha_s depending on the strategy - ker = build_ome(A, order, as1, backward_method) +@nb.experimental.jitclass(spec) +class QuadKerBase: + """Manage the common part of Mellin inversion integral. - # select the needed matrix element - ker = ker[indices[mode0], indices[mode1]] + Parameters + ---------- + u : float + quad argument + is_log : boolean + is a logarithmic interpolation + logx : float + Mellin inversion point + mode0 : str + first sector element + """ - # recombine everything - res = ker * pj * jac - return np.real(res) + def __init__(self, u, is_log, logx, mode0): + self.is_singlet = mode0 in [100, 21, 90] + self.is_QEDsinglet = mode0 in [21, 22, 100, 101, 90] + self.is_QEDvalence = mode0 in [10200, 10204] + self.is_log = is_log + self.u = u + self.logx = logx + + @property + def path(self): + """Return the associated instance of :class:`eko.mellin.Path`.""" + if self.is_singlet or self.is_QEDsinglet: + return mellin.Path(self.u, self.logx, True) + else: + return mellin.Path(self.u, self.logx, False) + + @property + def n(self): + """Returns the Mellin moment :math:`N`.""" + return self.path.n + + def integrand( + self, + areas, + ): + """Get transformation to Mellin space integral. + + Parameters + ---------- + areas : tuple + basis function configuration + + Returns + ------- + complex + common mellin inversion integrand + """ + if self.logx == 0.0: + return 0.0 + pj = interpolation.evaluate_grid(self.path.n, self.is_log, self.logx, areas) + if pj == 0.0: + return 0.0 + return self.path.prefactor * pj * self.path.jac @nb.njit(cache=True) -def select_QEDsinglet_element(ker, mode0, mode1): - """Select element of the QEDsinglet matrix. +def quad_ker_ad( + u, + order, + mode0, + mode1, + method, + is_log, + logx, + areas, + as_list, + mu2_from, + mu2_to, + a_half, + alphaem_running, + nf, + L, + ev_op_iterations, + ev_op_max_order, + sv_mode, + is_threshold, + n3lo_ad_variation, + is_polarized, + is_time_like, + use_fhmruvv, +): + """Raw evolution kernel inside quad. Parameters ---------- - ker : numpy.ndarray - QEDsinglet integration kernel - mode0 : int - id for first sector element + u : float + quad argument + order : int + perturbation order + mode0: int + pid for first sector element mode1 : int - id for second sector element + pid for second sector element + method : str + method + is_log : boolean + is a logarithmic interpolation + logx : float + Mellin inversion point + areas : tuple + basis function configuration + as1 : float + target coupling value + as0 : float + initial coupling value + mu2_from : float + initial value of mu2 + mu2_from : float + final value of mu2 + aem_list : list + list of electromagnetic coupling values + alphaem_running : bool + whether alphaem is running or not + nf : int + number of active flavors + L : float + logarithm of the squared ratio of factorization and renormalization scale + ev_op_iterations : int + number of evolution steps + ev_op_max_order : int + perturbative expansion order of U + sv_mode: int, `enum.IntEnum` + scale variation mode, see `eko.scale_variations.Modes` + is_threshold : boolean + is this an intermediate threshold operator? + n3lo_ad_variation : tuple + |N3LO| anomalous dimension variation ``(gg, gq, qg, qq, nsp, nsm, nsv)`` + is_polarized : boolean + is polarized evolution ? + is_time_like : boolean + is time-like evolution ? + use_fhmruvv : bool + if True use the |FHMRUVV| |N3LO| anomalous dimension + Returns ------- - ker : complex - QEDsinglet integration kernel element + float + evaluated integration kernel """ - if mode0 == 21: - index1 = 0 - elif mode0 == 22: - index1 = 1 - elif mode0 == 100: - index1 = 2 - else: - index1 = 3 - if mode1 == 21: - index2 = 0 - elif mode1 == 22: - index2 = 1 - elif mode1 == 100: - index2 = 2 + ker_base = QuadKerBase(u, is_log, logx, mode0) + integrand = ker_base.integrand(areas) + if integrand == 0.0: + return 0.0 + if order[1] == 0: + ker = quad_ker_qcd( + ker_base, + order, + mode0, + mode1, + method, + as_list[-1], + as_list[0], + nf, + L, + ev_op_iterations, + ev_op_max_order, + sv_mode, + is_threshold, + is_polarized, + is_time_like, + n3lo_ad_variation, + use_fhmruvv, + ) else: - index2 = 3 - return ker[index1, index2] + ker = quad_ker_qed( + ker_base, + order, + mode0, + mode1, + method, + as_list, + mu2_from, + mu2_to, + a_half, + alphaem_running, + nf, + L, + ev_op_iterations, + ev_op_max_order, + sv_mode, + is_threshold, + n3lo_ad_variation, + use_fhmruvv, + ) + + # recombine everything + return np.real(ker * integrand) @nb.njit(cache=True) -def select_QEDvalence_element(ker, mode0, mode1): - """Select element of the QEDvalence matrix. +def quad_ker_qcd( + ker_base, + order, + mode0, + mode1, + method, + as1, + as0, + nf, + L, + ev_op_iterations, + ev_op_max_order, + sv_mode, + is_threshold, + is_polarized, + is_time_like, + n3lo_ad_variation, + use_fhmruvv, +): + """Raw evolution kernel inside quad. Parameters ---------- - ker : numpy.ndarray - QEDvalence integration kernel - mode0 : int - id for first sector element + quad_ker : float + quad argument + order : int + perturbation order + mode0: int + pid for first sector element mode1 : int - id for second sector element + pid for second sector element + method : str + method + as1 : float + target coupling value + as0 : float + initial coupling value + nf : int + number of active flavors + L : float + logarithm of the squared ratio of factorization and renormalization scale + ev_op_iterations : int + number of evolution steps + ev_op_max_order : int + perturbative expansion order of U + sv_mode: int, `enum.IntEnum` + scale variation mode, see `eko.scale_variations.Modes` + is_threshold : boolean + is this an intermediate threshold operator? + n3lo_ad_variation : tuple + |N3LO| anomalous dimension variation ``(gg, gq, qg, qq, nsp, nsm, nsv)`` + use_fhmruvv : bool + if True use the |FHMRUVV| |N3LO| anomalous dimensions + Returns ------- - ker : complex - QEDvalence integration kernel element + float + evaluated integration kernel """ - index1 = 0 if mode0 == 10200 else 1 - index2 = 0 if mode1 == 10200 else 1 - return ker[index1, index2] + # compute the actual evolution kernel for pure QCD + if ker_base.is_singlet: + if is_polarized: + if is_time_like: + raise NotImplementedError("Polarized, time-like is not implemented") + else: + gamma_singlet = ad_ps.gamma_singlet(order, ker_base.n, nf) + else: + if is_time_like: + gamma_singlet = ad_ut.gamma_singlet(order, ker_base.n, nf) + else: + gamma_singlet = ad_us.gamma_singlet( + order, ker_base.n, nf, n3lo_ad_variation, use_fhmruvv + ) + # scale var exponentiated is directly applied on gamma + if sv_mode == sv.Modes.exponentiated: + gamma_singlet = sv_exponentiated.gamma_variation( + gamma_singlet, order, nf, L + ) + ker = s.dispatcher( + order, + method, + gamma_singlet, + as1, + as0, + nf, + ev_op_iterations, + ev_op_max_order, + ) + # scale var expanded is applied on the kernel + if sv_mode == sv.Modes.expanded and not is_threshold: + ker = np.ascontiguousarray( + sv_expanded.singlet_variation(gamma_singlet, as1, order, nf, L, dim=2) + ) @ np.ascontiguousarray(ker) + ker = select_singlet_element(ker, mode0, mode1) + else: + if is_polarized: + if is_time_like: + raise NotImplementedError("Polarized, time-like is not implemented") + else: + gamma_ns = ad_ps.gamma_ns(order, mode0, ker_base.n, nf) + else: + if is_time_like: + gamma_ns = ad_ut.gamma_ns(order, mode0, ker_base.n, nf) + else: + gamma_ns = ad_us.gamma_ns( + order, mode0, ker_base.n, nf, n3lo_ad_variation, use_fhmruvv + ) + if sv_mode == sv.Modes.exponentiated: + gamma_ns = sv_exponentiated.gamma_variation(gamma_ns, order, nf, L) + ker = ns.dispatcher( + order, + method, + gamma_ns, + as1, + as0, + nf, + ) + if sv_mode == sv.Modes.expanded and not is_threshold: + ker = sv_expanded.non_singlet_variation(gamma_ns, as1, order, nf, L) * ker + return ker -@nb.cfunc( - CB_SIGNATURE, - cache=True, - nopython=True, -) -def cb_quad_ker_qed( - re_gamma_raw, - im_gamma_raw, - re_n, - im_n, - re_jac, - im_jac, - order_qcd, - order_qed, - is_singlet, +@nb.njit(cache=True) +def quad_ker_qed( + ker_base, + order, mode0, mode1, + method, + as_list, + mu2_from, + mu2_to, + a_half, + alphaem_running, nf, - is_log, - logx, - areas_raw, - areas_x, - areas_y, - ev_method, - _as1, - _as0, + L, ev_op_iterations, - ev_op_max_order_qcd, + ev_op_max_order, sv_mode, is_threshold, - Lsv, - as_list_raw, - as_list_len, - mu2_from, - mu2_to, - a_half_raw, - a_half_x, - a_half_y, - alphaem_running, + n3lo_ad_variation, + use_fhmruvv, ): - """C Callback inside integration kernel.""" - # recover complex variables - n = re_n + im_n * 1j - jac = re_jac + im_jac * 1j - # compute basis functions - areas = nb.carray(areas_raw, (areas_x, areas_y)) - pj = interpolation.evaluate_grid(n, is_log, logx, areas) - order = (order_qcd, order_qed) - ev_op_max_order = (ev_op_max_order_qcd, order_qed) - is_valence = (mode0 == 10200) or (mode0 == 10204) - - as_list = nb.carray(as_list_raw, as_list_len) - a_half = nb.carray(a_half_raw, (a_half_x, a_half_y)) - - if is_singlet: - # reconstruct singlet matrices - re_gamma_singlet = nb.carray(re_gamma_raw, (order_qcd + 1, order_qed + 1, 4, 4)) - im_gamma_singlet = nb.carray(im_gamma_raw, (order_qcd + 1, order_qed + 1, 4, 4)) - gamma_singlet = re_gamma_singlet + im_gamma_singlet * 1j + """Raw evolution kernel inside quad. + Parameters + ---------- + ker_base : QuadKerBase + quad argument + order : int + perturbation order + mode0: int + pid for first sector element + mode1 : int + pid for second sector element + method : str + method + as1 : float + target coupling value + as0 : float + initial coupling value + mu2_from : float + initial value of mu2 + mu2_from : float + final value of mu2 + aem_list : list + list of electromagnetic coupling values + alphaem_running : bool + whether alphaem is running or not + nf : int + number of active flavors + L : float + logarithm of the squared ratio of factorization and renormalization scale + ev_op_iterations : int + number of evolution steps + ev_op_max_order : int + perturbative expansion order of U + sv_mode: int, `enum.IntEnum` + scale variation mode, see `eko.scale_variations.Modes` + is_threshold : boolean + is this an intermediate threshold operator? + n3lo_ad_variation : tuple + |N3LO| anomalous dimension variation ``(gg, gq, qg, qq, nsp, nsm, nsv)`` + use_fhmruvv : bool + if True use the |FHMRUVV| |N3LO| anomalous dimensions + + Returns + ------- + float + evaluated integration kernel + """ + # compute the actual evolution kernel for QEDxQCD + if ker_base.is_QEDsinglet: + gamma_s = ad_us.gamma_singlet_qed( + order, ker_base.n, nf, n3lo_ad_variation, use_fhmruvv + ) # scale var exponentiated is directly applied on gamma if sv_mode == sv.Modes.exponentiated: - gamma_singlet = sv.exponentiated.gamma_variation_qed( - gamma_singlet, order, nf, lepton_number(mu2_to), Lsv, alphaem_running + gamma_s = sv_exponentiated.gamma_variation_qed( + gamma_s, order, nf, lepton_number(mu2_to), L, alphaem_running ) - ker = qed_s.dispatcher( order, - ev_method, - gamma_singlet, + method, + gamma_s, as_list, a_half, nf, ev_op_iterations, ev_op_max_order, ) + # scale var expanded is applied on the kernel + # TODO : in this way a_half[-1][1] is the aem value computed in + # the middle point of the last step. Instead we want aem computed in mu2_final. + # However the distance between the two is very small and affects only the running aem if sv_mode == sv.Modes.expanded and not is_threshold: ker = np.ascontiguousarray( - sv.expanded.singlet_variation_qed( - gamma_singlet, - as_list[-1], - a_half[-1][1], - alphaem_running, - order, - nf, - Lsv, + sv_expanded.singlet_variation_qed( + gamma_s, as_list[-1], a_half[-1][1], alphaem_running, order, nf, L ) ) @ np.ascontiguousarray(ker) ker = select_QEDsinglet_element(ker, mode0, mode1) - - elif is_valence: - # reconstruct valence matrices - re_gamma_valence = nb.carray(re_gamma_raw, (order_qcd + 1, order_qed + 1, 2, 2)) - im_gamma_valence = nb.carray(im_gamma_raw, (order_qcd + 1, order_qed + 1, 2, 2)) - gamma_valence = re_gamma_valence + im_gamma_valence * 1j - + elif ker_base.is_QEDvalence: + gamma_v = ad_us.gamma_valence_qed( + order, ker_base.n, nf, n3lo_ad_variation, use_fhmruvv + ) + # scale var exponentiated is directly applied on gamma if sv_mode == sv.Modes.exponentiated: - gamma_valence = sv.exponentiated.gamma_variation_qed( - gamma_valence, order, nf, lepton_number(mu2_to), Lsv, alphaem_running + gamma_v = sv_exponentiated.gamma_variation_qed( + gamma_v, order, nf, lepton_number(mu2_to), L, alphaem_running ) ker = qed_v.dispatcher( order, - ev_method, - gamma_valence, + method, + gamma_v, as_list, a_half, nf, @@ -478,31 +620,23 @@ def cb_quad_ker_qed( # scale var expanded is applied on the kernel if sv_mode == sv.Modes.expanded and not is_threshold: ker = np.ascontiguousarray( - sv.expanded.valence_variation_qed( - gamma_valence, - as_list[-1], - a_half[-1][1], - alphaem_running, - order, - nf, - Lsv, + sv_expanded.valence_variation_qed( + gamma_v, as_list[-1], a_half[-1][1], alphaem_running, order, nf, L ) ) @ np.ascontiguousarray(ker) ker = select_QEDvalence_element(ker, mode0, mode1) - else: - # construct non-singlet matrices - re_gamma_ns = nb.carray(re_gamma_raw, (order_qcd + 1, order_qed + 1)) - im_gamma_ns = nb.carray(im_gamma_raw, (order_qcd + 1, order_qed + 1)) - gamma_ns = re_gamma_ns + im_gamma_ns * 1j + gamma_ns = ad_us.gamma_ns_qed( + order, mode0, ker_base.n, nf, n3lo_ad_variation, use_fhmruvv + ) + # scale var exponentiated is directly applied on gamma if sv_mode == sv.Modes.exponentiated: gamma_ns = sv_exponentiated.gamma_variation_qed( - gamma_ns, order, nf, lepton_number(mu2_to), Lsv, alphaem_running + gamma_ns, order, nf, lepton_number(mu2_to), L, alphaem_running ) - # construct eko ker = qed_ns.dispatcher( order, - ev_method, + method, gamma_ns, as_list, a_half[:, 1], @@ -515,16 +649,107 @@ def cb_quad_ker_qed( if sv_mode == sv.Modes.expanded and not is_threshold: ker = ( sv_expanded.non_singlet_variation_qed( - gamma_ns, - as_list[-1], - a_half[-1][1], - alphaem_running, - order, - nf, - Lsv, + gamma_ns, as_list[-1], a_half[-1][1], alphaem_running, order, nf, L ) * ker ) + return ker + + +@nb.njit(cache=True) +def quad_ker_ome( + u, + order, + mode0, + mode1, + is_log, + logx, + areas, + a_s, + nf, + L, + sv_mode, + Lsv, + backward_method, + is_msbar, + is_polarized, + is_time_like, +): + r"""Raw kernel inside quad. + + Parameters + ---------- + u : float + quad argument + order : tuple(int,int) + perturbation matching order + mode0 : int + pid for first element in the singlet sector + mode1 : int + pid for second element in the singlet sector + is_log : boolean + logarithmic interpolation + logx : float + Mellin inversion point + areas : tuple + basis function configuration + a_s : float + strong coupling, needed only for the exact inverse + nf: int + number of active flavor below threshold + L : float + :math:``\ln(\mu_F^2 / m_h^2)`` + backward_method : InversionMethod or None + empty or method for inverting the matching condition (exact or expanded) + is_msbar: bool + add the |MSbar| contribution + is_polarized : boolean + is polarized evolution ? + is_time_like : boolean + is time-like evolution ? + + Returns + ------- + ker : float + evaluated integration kernel + """ + ker_base = QuadKerBase(u, is_log, logx, mode0) + integrand = ker_base.integrand(areas) + if integrand == 0.0: + return 0.0 + # compute the ome + if ker_base.is_singlet or ker_base.is_QEDsinglet: + indices = {21: 0, 100: 1, 90: 2} + if is_polarized: + if is_time_like: + raise NotImplementedError("Polarized, time-like is not implemented") + A = ome_ps.A_singlet(order, ker_base.n, nf, L) + else: + if is_time_like: + A = ome_ut.A_singlet(order, ker_base.n, L) + else: + A = ome_us.A_singlet(order, ker_base.n, nf, L, is_msbar) + else: + indices = {200: 0, 91: 1} + if is_polarized: + if is_time_like: + raise NotImplementedError("Polarized, time-like is not implemented") + A = ome_ps.A_non_singlet(order, ker_base.n, L) + else: + if is_time_like: + A = ome_ut.A_non_singlet(order, ker_base.n, L) + else: + A = ome_us.A_non_singlet(order, ker_base.n, nf, L) + + # correct for scale variations + if sv_mode == sv.Modes.exponentiated: + A = gamma_variation(A, order, nf, Lsv) + + # build the expansion in alpha_s depending on the strategy + ker = build_ome(A, order, a_s, backward_method) + + # select the needed matrix element + ker = ker[indices[mode0], indices[mode1]] + # recombine everything - res = ker * pj * jac - return np.real(res) + return np.real(ker * integrand) diff --git a/src/eko/evolution_operator/quad_ker.py.patch b/src/eko/evolution_operator/quad_ker.py.patch new file mode 100644 index 000000000..9f790cc9a --- /dev/null +++ b/src/eko/evolution_operator/quad_ker.py.patch @@ -0,0 +1,145 @@ +diff --git a/src/eko/evolution_operator/quad_ker.py b/src/eko/evolution_operator/quad_ker.py +index bc23643c..79cc0bdc 100644 +--- a/src/eko/evolution_operator/quad_ker.py ++++ b/src/eko/evolution_operator/quad_ker.py +@@ -3,15 +3,11 @@ r"""Integration kernels.""" + import enum + import logging + ++import ekors + import numba as nb + import numpy as np + +-import ekore.anomalous_dimensions.polarized.space_like as ad_ps + import ekore.anomalous_dimensions.unpolarized.space_like as ad_us +-import ekore.anomalous_dimensions.unpolarized.time_like as ad_ut +-import ekore.operator_matrix_elements.polarized.space_like as ome_ps +-import ekore.operator_matrix_elements.unpolarized.space_like as ome_us +-import ekore.operator_matrix_elements.unpolarized.time_like as ome_ut + + from .. import interpolation, mellin + from .. import scale_variations as sv +@@ -437,18 +433,20 @@ def quad_ker_qcd( + """ + # compute the actual evolution kernel for pure QCD + if ker_base.is_singlet: +- if is_polarized: +- if is_time_like: +- raise NotImplementedError("Polarized, time-like is not implemented") +- else: +- gamma_singlet = ad_ps.gamma_singlet(order, ker_base.n, nf) +- else: +- if is_time_like: +- gamma_singlet = ad_ut.gamma_singlet(order, ker_base.n, nf) +- else: +- gamma_singlet = ad_us.gamma_singlet( +- order, ker_base.n, nf, n3lo_ad_variation, use_fhmruvv +- ) ++ order_qcd = order[0] ++ n_re = ker_base.n.real ++ n_im = ker_base.n.imag ++ with nb.objmode(gamma_singlet="complex128[:,:,:]"): ++ re_out, im_out = ekors.qcd_gamma_singlet( ++ is_polarized, ++ is_time_like, ++ order_qcd, ++ n_re, ++ n_im, ++ nf, ++ n3lo_ad_variation[:4], ++ ) ++ gamma_singlet = (re_out + 1j * im_out).reshape(order_qcd, 2, 2) + # scale var exponentiated is directly applied on gamma + if sv_mode == sv.Modes.exponentiated: + gamma_singlet = sv_exponentiated.gamma_variation( +@@ -471,18 +469,22 @@ def quad_ker_qcd( + ) @ np.ascontiguousarray(ker) + ker = select_singlet_element(ker, mode0, mode1) + else: +- if is_polarized: +- if is_time_like: +- raise NotImplementedError("Polarized, time-like is not implemented") +- else: +- gamma_ns = ad_ps.gamma_ns(order, mode0, ker_base.n, nf) +- else: +- if is_time_like: +- gamma_ns = ad_ut.gamma_ns(order, mode0, ker_base.n, nf) +- else: +- gamma_ns = ad_us.gamma_ns( +- order, mode0, ker_base.n, nf, n3lo_ad_variation, use_fhmruvv +- ) ++ order_qcd = order[0] ++ n_re = ker_base.n.real ++ n_im = ker_base.n.imag ++ with nb.objmode(gamma_ns="complex128[:]"): ++ re_out, im_out = ekors.qcd_gamma_ns( ++ is_polarized, ++ is_time_like, ++ order_qcd, ++ mode0, ++ n_re, ++ n_im, ++ nf, ++ n3lo_ad_variation[4:], ++ use_fhmruvv, ++ ) ++ gamma_ns = re_out + 1j * im_out + if sv_mode == sv.Modes.exponentiated: + gamma_ns = sv_exponentiated.gamma_variation(gamma_ns, order, nf, L) + ker = ns.dispatcher( +@@ -720,26 +722,37 @@ def quad_ker_ome( + # compute the ome + if ker_base.is_singlet or ker_base.is_QEDsinglet: + indices = {21: 0, 100: 1, 90: 2} +- if is_polarized: +- if is_time_like: +- raise NotImplementedError("Polarized, time-like is not implemented") +- A = ome_ps.A_singlet(order, ker_base.n, nf, L) +- else: +- if is_time_like: +- A = ome_ut.A_singlet(order, ker_base.n, L) +- else: +- A = ome_us.A_singlet(order, ker_base.n, nf, L, is_msbar) ++ matching_order_qcd = order[0] ++ n_re = ker_base.n.real ++ n_im = ker_base.n.imag ++ with nb.objmode(A="complex128[:,:,:]"): ++ re_out, im_out = ekors.ome_singlet( ++ is_polarized, ++ is_time_like, ++ matching_order_qcd, ++ n_re, ++ n_im, ++ nf, ++ L, ++ is_msbar, ++ ) ++ A = (re_out + 1j * im_out).reshape(matching_order_qcd, 3, 3) + else: + indices = {200: 0, 91: 1} +- if is_polarized: +- if is_time_like: +- raise NotImplementedError("Polarized, time-like is not implemented") +- A = ome_ps.A_non_singlet(order, ker_base.n, L) +- else: +- if is_time_like: +- A = ome_ut.A_non_singlet(order, ker_base.n, L) +- else: +- A = ome_us.A_non_singlet(order, ker_base.n, nf, L) ++ matching_order_qcd = order[0] ++ n_re = ker_base.n.real ++ n_im = ker_base.n.imag ++ with nb.objmode(A="complex128[:,:,:]"): ++ re_out, im_out = ekors.ome_non_singlet( ++ is_polarized, ++ is_time_like, ++ matching_order_qcd, ++ n_re, ++ n_im, ++ nf, ++ L, ++ ) ++ A = (re_out + 1j * im_out).reshape(matching_order_qcd, 2, 2) + + # correct for scale variations + if sv_mode == sv.Modes.exponentiated: diff --git a/tests/eko/evolution_operator/test_ome.py b/tests/eko/evolution_operator/test_ome.py index c0e7aa674..c6b504aab 100644 --- a/tests/eko/evolution_operator/test_ome.py +++ b/tests/eko/evolution_operator/test_ome.py @@ -7,9 +7,9 @@ from eko.evolution_operator.operator_matrix_element import ( MatchingMethods, OperatorMatrixElement, - build_ome, quad_ker, ) +from eko.evolution_operator.quad_ker import build_ome from eko.io.runcards import OperatorCard, TheoryCard from eko.io.types import InversionMethod from eko.runner.parts import _managers, _matching_configs