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/// Intergration 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
101    if args.is_ome {
102        if is_singlet {
103            raw = unravel(
104                ekore::operator_matrix_elements::unpolarized::spacelike::A_singlet(
105                    args.order_qcd,
106                    &mut c,
107                    args.nf,
108                    args.L,
109                ),
110                args.order_qcd,
111            );
112        } else {
113            raw = unravel(
114                ekore::operator_matrix_elements::unpolarized::spacelike::A_non_singlet(
115                    args.order_qcd,
116                    &mut c,
117                    args.nf,
118                    args.L,
119                ),
120                args.order_qcd,
121            );
122        }
123    } else if is_singlet {
124        if args.order_qed > 0 {
125            let gamma_singlet_qed =
126                ekore::anomalous_dimensions::unpolarized::spacelike::gamma_singlet_qed;
127            raw = unravel_qed(
128                gamma_singlet_qed(args.order_qcd, args.order_qed, &mut c, args.nf),
129                args.order_qcd,
130                args.order_qed,
131            );
132        } else {
133            let gamma_singlet_qcd = match args.is_polarized {
134                true => ekore::anomalous_dimensions::polarized::spacelike::gamma_singlet_qcd,
135                false => ekore::anomalous_dimensions::unpolarized::spacelike::gamma_singlet_qcd,
136            };
137            raw = unravel(
138                gamma_singlet_qcd(args.order_qcd, &mut c, args.nf),
139                args.order_qcd,
140            );
141        }
142    } else if args.order_qed > 0 {
143        if is_qed_valence {
144            let gamma_valence_qed =
145                ekore::anomalous_dimensions::unpolarized::spacelike::gamma_valence_qed;
146            raw = unravel_qed(
147                gamma_valence_qed(args.order_qcd, args.order_qed, &mut c, args.nf),
148                args.order_qcd,
149                args.order_qed,
150            );
151        } else {
152            let gamma_ns_qed = ekore::anomalous_dimensions::unpolarized::spacelike::gamma_ns_qed;
153            raw = unravel_qed_ns(
154                gamma_ns_qed(args.order_qcd, args.order_qed, args.mode0, &mut c, args.nf),
155                args.order_qcd,
156                args.order_qed,
157            );
158        }
159    } else {
160        // we can not do 1D
161        let gamma_ns_qcd = match args.is_polarized {
162            true => ekore::anomalous_dimensions::polarized::spacelike::gamma_ns_qcd,
163            false => ekore::anomalous_dimensions::unpolarized::spacelike::gamma_ns_qcd,
164        };
165        let res = gamma_ns_qcd(args.order_qcd, args.mode0, &mut c, args.nf);
166        for el in res.iter().take(args.order_qcd) {
167            raw.re.push(el.re);
168            raw.im.push(el.im);
169        }
170    }
171
172    // pass on
173    (args.py)(
174        raw.re.as_ptr(),
175        raw.im.as_ptr(),
176        c.n().re,
177        c.n().im,
178        jac.re,
179        jac.im,
180        args.order_qcd,
181        args.order_qed,
182        is_singlet,
183        args.mode0,
184        args.mode1,
185        args.nf,
186        args.is_log,
187        args.logx,
188        args.areas,
189        args.areas_x,
190        args.areas_y,
191        args.method_num,
192        args.as1,
193        args.as0,
194        args.ev_op_iterations,
195        args.ev_op_max_order_qcd,
196        args.sv_mode_num,
197        args.is_threshold,
198        args.Lsv,
199        // additional QED params
200        args.as_list,
201        args.as_list_len,
202        args.mu2_from,
203        args.mu2_to,
204        args.a_half,
205        args.a_half_x,
206        args.a_half_y,
207        args.alphaem_running,
208    )
209}
210
211/// Python callback signature
212type PyQuadKerT = unsafe extern "C" fn(
213    *const f64, // re_gamma
214    *const f64, // im_gamma
215    f64,        // re_n
216    f64,        // im_n
217    f64,        // re_jac
218    f64,        // im_jac
219    usize,      // order_qcd
220    usize,      // order_qed
221    bool,       // is_singlet
222    u16,        // mode0
223    u16,        // mode1
224    u8,         // nf
225    bool,       // is_log
226    f64,        // logx
227    *const f64, // areas
228    u8,         // areas_x
229    u8,         // areas_y
230    u8,         // method_num
231    f64,        // as1
232    f64,        // as0
233    u8,         // ev_op_iterations
234    u8,         // ev_op_max_order_qcd
235    u8,         // sv_mode_num
236    bool,       // is_threshold
237    f64,        // lsv
238    *const f64, // as_list
239    u8,         // as_list_len
240    f64,        // mu2_from
241    f64,        // mu2_to
242    *const f64, // a_half
243    u8,         // a_half_x
244    u8,         // a_half_y
245    bool,       // alphaem_running
246) -> f64;
247
248/// Additional integration parameters
249#[allow(non_snake_case)]
250#[repr(C)]
251#[derive(Clone, Copy)]
252pub struct QuadArgs {
253    pub order_qcd: usize,
254    pub order_qed: usize,
255    pub mode0: u16,
256    pub mode1: u16,
257    pub is_polarized: bool,
258    pub is_time_like: bool,
259    pub nf: u8,
260    pub py: PyQuadKerT,
261    pub is_log: bool,
262    pub logx: f64,
263    pub areas: *const f64,
264    pub areas_x: u8,
265    pub areas_y: u8,
266    pub L: f64,
267    pub method_num: u8,
268    pub as1: f64,
269    pub as0: f64,
270    pub ev_op_iterations: u8,
271    pub ev_op_max_order_qcd: u8,
272    pub sv_mode_num: u8,
273    pub is_threshold: bool,
274    pub is_ome: bool,
275    pub Lsv: f64,
276    // additional param required for QED
277    pub as_list: *const f64,
278    pub as_list_len: u8,
279    pub mu2_from: f64,
280    pub mu2_to: f64,
281    pub a_half: *const f64,
282    pub a_half_x: u8,
283    pub a_half_y: u8,
284    pub alphaem_running: bool,
285}
286
287/// Empty placeholder function for python callback.
288///
289/// # Safety
290/// This is the connection back to Python, so we don't know what is on the other side.
291pub unsafe extern "C" fn my_py(
292    _re_gamma: *const f64,
293    _im_gamma: *const f64,
294    _re_n: f64,
295    _im_n: f64,
296    _re_jac: f64,
297    _im_jac: f64,
298    _order_qcd: usize,
299    _order_qed: usize,
300    _is_singlet: bool,
301    _mode0: u16,
302    _mode1: u16,
303    _nf: u8,
304    _is_log: bool,
305    _logx: f64,
306    _areas: *const f64,
307    _areas_x: u8,
308    _areas_y: u8,
309    _method_num: u8,
310    _as1: f64,
311    _as0: f64,
312    _ev_op_iterations: u8,
313    _ev_op_max_order_qcd: u8,
314    _sv_mode_num: u8,
315    _is_threshold: bool,
316    _lsv: f64,
317    _as_list: *const f64,
318    _as_list_len: u8,
319    _mu2_from: f64,
320    _mu2_to: f64,
321    _a_half: *const f64,
322    _a_half_x: u8,
323    _a_half_y: u8,
324    _alphaem_running: bool,
325) -> f64 {
326    0.
327}
328
329/// Return empty additional arguments.
330///
331/// This is required to make the arguments part of the API, otherwise it won't be added to the compiled
332/// package (since it does not appear in the signature of `rust_quad_ker`).
333///
334/// # Safety
335/// This is the connection from and back to Python, so we don't know what is on the other side.
336#[no_mangle]
337pub unsafe extern "C" fn empty_args() -> QuadArgs {
338    QuadArgs {
339        order_qcd: 0,
340        order_qed: 0,
341        mode0: 0,
342        mode1: 0,
343        is_polarized: false,
344        is_time_like: false,
345        nf: 0,
346        py: my_py,
347        is_log: true,
348        logx: 0.,
349        areas: [].as_ptr(),
350        areas_x: 0,
351        areas_y: 0,
352        L: 0.,
353        method_num: 0,
354        as1: 0.,
355        as0: 0.,
356        ev_op_iterations: 0,
357        ev_op_max_order_qcd: 0,
358        sv_mode_num: 0,
359        is_threshold: false,
360        is_ome: false,
361        Lsv: 0.,
362        as_list: [].as_ptr(),
363        as_list_len: 0,
364        mu2_from: 0.,
365        mu2_to: 0.,
366        a_half: [].as_ptr(),
367        a_half_x: 0,
368        a_half_y: 0,
369        alphaem_running: false,
370    }
371}
ekors/lib.rs

ekors/
lib.rs
