Skip to main content

ekors/
lib.rs

1//! Interface to the eko Python package.
2
3use ekore::harmonics::cache::Cache;
4use num::Complex;
5use std::ffi::c_void;
6
7pub mod bib;
8pub mod mellin;
9
10/// Wrapper to pass arguments back to Python.
11struct RawCmplx {
12    re: Vec<f64>,
13    im: Vec<f64>,
14}
15
16/// Map tensor with shape (o,d,d) to c-ordered list.
17///
18/// This is needed for the QCD singlet.
19fn unravel<const DIM: usize>(res: Vec<[[Complex<f64>; DIM]; DIM]>, order_qcd: usize) -> RawCmplx {
20    let mut target = RawCmplx {
21        re: Vec::<f64>::new(),
22        im: Vec::<f64>::new(),
23    };
24    for obj in res.iter().take(order_qcd) {
25        for col in obj.iter().take(DIM) {
26            for el in col.iter().take(DIM) {
27                target.re.push(el.re);
28                target.im.push(el.im);
29            }
30        }
31    }
32    target
33}
34
35/// Map tensor with shape (o,o',d,d) to c-ordered list.
36///
37/// This is needed for the QED singlet and valence.
38fn unravel_qed<const DIM: usize>(
39    res: Vec<Vec<[[Complex<f64>; DIM]; DIM]>>,
40    order_qcd: usize,
41    order_qed: usize,
42) -> RawCmplx {
43    let mut target = RawCmplx {
44        re: Vec::<f64>::new(),
45        im: Vec::<f64>::new(),
46    };
47    for obj_ in res.iter().take(order_qcd + 1) {
48        for obj in obj_.iter().take(order_qed + 1) {
49            for col in obj.iter().take(DIM) {
50                for el in col.iter().take(DIM) {
51                    target.re.push(el.re);
52                    target.im.push(el.im);
53                }
54            }
55        }
56    }
57    target
58}
59
60/// Map tensor with shape (o,o',d) to c-ordered list.
61///
62/// This is needed for the QED non-singlet.
63fn unravel_qed_ns(res: Vec<Vec<Complex<f64>>>, order_qcd: usize, order_qed: usize) -> RawCmplx {
64    let mut target = RawCmplx {
65        re: Vec::<f64>::new(),
66        im: Vec::<f64>::new(),
67    };
68    for col in res.iter().take(order_qcd + 1) {
69        for el in col.iter().take(order_qed + 1) {
70            target.re.push(el.re);
71            target.im.push(el.im);
72        }
73    }
74    target
75}
76
77/// Integration kernel inside quad.
78///
79/// # Safety
80/// This is the connection from Python, so we don't know what is on the other side.
81#[no_mangle]
82pub unsafe extern "C" fn rust_quad_ker(u: f64, rargs: *mut c_void) -> f64 {
83    let args = *(rargs as *mut QuadArgs);
84
85    let is_singlet = (100 == args.mode0)
86        || (21 == args.mode0)
87        || (90 == args.mode0)
88        || (22 == args.mode0)
89        || (101 == args.mode0);
90
91    let is_qed_valence = (10200 == args.mode0) || (10204 == args.mode0);
92    // prepare Mellin stuff
93    let path = mellin::TalbotPath::new(u, args.logx, is_singlet);
94    let jac = path.jac() * path.prefactor();
95    let mut c = Cache::new(path.n());
96    let mut raw = RawCmplx {
97        re: Vec::<f64>::new(),
98        im: Vec::<f64>::new(),
99    };
100    let n3lo_ad_variation = std::slice::from_raw_parts(args.n3lo_ad_variation, 7)
101        .try_into()
102        .unwrap();
103
104    if args.is_ome {
105        if is_singlet {
106            raw = unravel(
107                ekore::operator_matrix_elements::unpolarized::spacelike::A_singlet(
108                    args.order_qcd,
109                    &mut c,
110                    args.nf,
111                    args.L,
112                ),
113                args.order_qcd,
114            );
115        } else {
116            raw = unravel(
117                ekore::operator_matrix_elements::unpolarized::spacelike::A_non_singlet(
118                    args.order_qcd,
119                    &mut c,
120                    args.nf,
121                    args.L,
122                ),
123                args.order_qcd,
124            );
125        }
126    } else if is_singlet {
127        if args.order_qed > 0 {
128            let gamma_singlet_qed =
129                ekore::anomalous_dimensions::unpolarized::spacelike::gamma_singlet_qed;
130            raw = unravel_qed(
131                gamma_singlet_qed(
132                    args.order_qcd,
133                    args.order_qed,
134                    &mut c,
135                    args.nf,
136                    n3lo_ad_variation,
137                ),
138                args.order_qcd,
139                args.order_qed,
140            );
141        } else {
142            let gamma_singlet_qcd = match args.is_polarized {
143                true => ekore::anomalous_dimensions::polarized::spacelike::gamma_singlet_qcd,
144                false => ekore::anomalous_dimensions::unpolarized::spacelike::gamma_singlet_qcd,
145            };
146            raw = unravel(
147                gamma_singlet_qcd(
148                    args.order_qcd,
149                    &mut c,
150                    args.nf,
151                    n3lo_ad_variation[0..4].try_into().unwrap(),
152                ),
153                args.order_qcd,
154            );
155        }
156    } else if args.order_qed > 0 {
157        if is_qed_valence {
158            let gamma_valence_qed =
159                ekore::anomalous_dimensions::unpolarized::spacelike::gamma_valence_qed;
160            raw = unravel_qed(
161                gamma_valence_qed(
162                    args.order_qcd,
163                    args.order_qed,
164                    &mut c,
165                    args.nf,
166                    n3lo_ad_variation[4..7].try_into().unwrap(),
167                ),
168                args.order_qcd,
169                args.order_qed,
170            );
171        } else {
172            let gamma_ns_qed = ekore::anomalous_dimensions::unpolarized::spacelike::gamma_ns_qed;
173            raw = unravel_qed_ns(
174                gamma_ns_qed(
175                    args.order_qcd,
176                    args.order_qed,
177                    args.mode0,
178                    &mut c,
179                    args.nf,
180                    n3lo_ad_variation[4..7].try_into().unwrap(),
181                ),
182                args.order_qcd,
183                args.order_qed,
184            );
185        }
186    } else {
187        // we can not do 1D
188        let gamma_ns_qcd = match args.is_polarized {
189            true => ekore::anomalous_dimensions::polarized::spacelike::gamma_ns_qcd,
190            false => ekore::anomalous_dimensions::unpolarized::spacelike::gamma_ns_qcd,
191        };
192        let res = gamma_ns_qcd(
193            args.order_qcd,
194            args.mode0,
195            &mut c,
196            args.nf,
197            n3lo_ad_variation[4..7].try_into().unwrap(),
198        );
199        for el in res.iter().take(args.order_qcd) {
200            raw.re.push(el.re);
201            raw.im.push(el.im);
202        }
203    }
204
205    // pass on
206    (args.py)(
207        raw.re.as_ptr(),
208        raw.im.as_ptr(),
209        c.n().re,
210        c.n().im,
211        jac.re,
212        jac.im,
213        args.order_qcd,
214        args.order_qed,
215        is_singlet,
216        args.mode0,
217        args.mode1,
218        args.nf,
219        args.is_log,
220        args.logx,
221        args.areas,
222        args.areas_x,
223        args.areas_y,
224        args.method_num,
225        args.as1,
226        args.as0,
227        args.ev_op_iterations,
228        args.ev_op_max_order_qcd,
229        args.sv_mode_num,
230        args.is_threshold,
231        args.Lsv,
232        // additional QED params
233        args.as_list,
234        args.as_list_len,
235        args.mu2_from,
236        args.mu2_to,
237        args.a_half,
238        args.a_half_x,
239        args.a_half_y,
240        args.alphaem_running,
241    )
242}
243
244/// Python callback signature
245type PyQuadKerT = unsafe extern "C" fn(
246    *const f64, // re_gamma
247    *const f64, // im_gamma
248    f64,        // re_n
249    f64,        // im_n
250    f64,        // re_jac
251    f64,        // im_jac
252    usize,      // order_qcd
253    usize,      // order_qed
254    bool,       // is_singlet
255    u16,        // mode0
256    u16,        // mode1
257    u8,         // nf
258    bool,       // is_log
259    f64,        // logx
260    *const f64, // areas
261    u8,         // areas_x
262    u8,         // areas_y
263    u8,         // method_num
264    f64,        // as1
265    f64,        // as0
266    u8,         // ev_op_iterations
267    u8,         // ev_op_max_order_qcd
268    u8,         // sv_mode_num
269    bool,       // is_threshold
270    f64,        // lsv
271    *const f64, // as_list
272    u8,         // as_list_len
273    f64,        // mu2_from
274    f64,        // mu2_to
275    *const f64, // a_half
276    u8,         // a_half_x
277    u8,         // a_half_y
278    bool,       // alphaem_running
279) -> f64;
280
281/// Additional integration parameters
282#[allow(non_snake_case)]
283#[repr(C)]
284#[derive(Clone, Copy)]
285pub struct QuadArgs {
286    pub order_qcd: usize,
287    pub order_qed: usize,
288    pub mode0: u16,
289    pub mode1: u16,
290    pub is_polarized: bool,
291    pub is_time_like: bool,
292    pub nf: u8,
293    pub py: PyQuadKerT,
294    pub is_log: bool,
295    pub logx: f64,
296    pub areas: *const f64,
297    pub areas_x: u8,
298    pub areas_y: u8,
299    pub L: f64,
300    pub method_num: u8,
301    pub as1: f64,
302    pub as0: f64,
303    pub ev_op_iterations: u8,
304    pub ev_op_max_order_qcd: u8,
305    pub sv_mode_num: u8,
306    pub is_threshold: bool,
307    pub is_ome: bool,
308    pub Lsv: f64,
309    // additional param required for QED
310    pub as_list: *const f64,
311    pub as_list_len: u8,
312    pub mu2_from: f64,
313    pub mu2_to: f64,
314    pub a_half: *const f64,
315    pub a_half_x: u8,
316    pub a_half_y: u8,
317    pub alphaem_running: bool,
318    // additional param required for N3LO
319    pub n3lo_ad_variation: *const u8,
320}
321
322/// Empty placeholder function for python callback.
323///
324/// # Safety
325/// This is the connection back to Python, so we don't know what is on the other side.
326pub unsafe extern "C" fn my_py(
327    _re_gamma: *const f64,
328    _im_gamma: *const f64,
329    _re_n: f64,
330    _im_n: f64,
331    _re_jac: f64,
332    _im_jac: f64,
333    _order_qcd: usize,
334    _order_qed: usize,
335    _is_singlet: bool,
336    _mode0: u16,
337    _mode1: u16,
338    _nf: u8,
339    _is_log: bool,
340    _logx: f64,
341    _areas: *const f64,
342    _areas_x: u8,
343    _areas_y: u8,
344    _method_num: u8,
345    _as1: f64,
346    _as0: f64,
347    _ev_op_iterations: u8,
348    _ev_op_max_order_qcd: u8,
349    _sv_mode_num: u8,
350    _is_threshold: bool,
351    _lsv: f64,
352    _as_list: *const f64,
353    _as_list_len: u8,
354    _mu2_from: f64,
355    _mu2_to: f64,
356    _a_half: *const f64,
357    _a_half_x: u8,
358    _a_half_y: u8,
359    _alphaem_running: bool,
360) -> f64 {
361    0.
362}
363
364/// Return empty additional arguments.
365///
366/// This is required to make the arguments part of the API, otherwise it won't be added to the compiled
367/// package (since it does not appear in the signature of `rust_quad_ker`).
368///
369/// # Safety
370/// This is the connection from and back to Python, so we don't know what is on the other side.
371#[no_mangle]
372pub unsafe extern "C" fn empty_args() -> QuadArgs {
373    QuadArgs {
374        order_qcd: 0,
375        order_qed: 0,
376        mode0: 0,
377        mode1: 0,
378        is_polarized: false,
379        is_time_like: false,
380        nf: 0,
381        py: my_py,
382        is_log: true,
383        logx: 0.,
384        areas: [].as_ptr(),
385        areas_x: 0,
386        areas_y: 0,
387        L: 0.,
388        method_num: 0,
389        as1: 0.,
390        as0: 0.,
391        ev_op_iterations: 0,
392        ev_op_max_order_qcd: 0,
393        sv_mode_num: 0,
394        is_threshold: false,
395        is_ome: false,
396        Lsv: 0.,
397        as_list: [].as_ptr(),
398        as_list_len: 0,
399        mu2_from: 0.,
400        mu2_to: 0.,
401        a_half: [].as_ptr(),
402        a_half_x: 0,
403        a_half_y: 0,
404        alphaem_running: false,
405        n3lo_ad_variation: [0, 0, 0, 0, 0, 0, 0].as_ptr(),
406    }
407}