Skip to main content

laddu_cpu/
lib.rs

1use pyo3::prelude::*;
2
3#[cfg_attr(feature = "mpi", pymodule(name = "laddu_mpi"))]
4#[cfg_attr(not(feature = "mpi"), pymodule(name = "laddu_cpu"))]
5mod laddu {
6    use super::*;
7    #[pyfunction]
8    fn version() -> String {
9        env!("CARGO_PKG_VERSION").to_string()
10    }
11
12    #[pymodule_export]
13    use laddu_python::{
14        amplitudes::{
15            py_constant, py_expr_one, py_expr_product, py_expr_sum, py_expr_zero, py_parameter,
16            py_test_amplitude, PyEvaluator, PyExpression, PyParameterLike,
17        },
18        available_parallelism,
19        data::{PyBinnedDataset, PyDataset, PyEvent},
20        mpi::{finalize_mpi, get_rank, get_size, is_mpi_available, is_root, use_mpi, using_mpi},
21        utils::{
22            variables::{
23                PyAngles, PyCosTheta, PyMandelstam, PyMass, PyPhi, PyPolAngle, PyPolMagnitude,
24                PyPolarization, PyTopology, PyVariableExpression,
25            },
26            vectors::{PyVec3, PyVec4},
27        },
28    };
29
30    #[pymodule_export]
31    use laddu_amplitudes::{
32        breit_wigner::py_breit_wigner,
33        common::{py_complex_scalar, py_polar_complex_scalar, py_scalar},
34        kmatrix::{
35            py_kopf_kmatrix_a0, py_kopf_kmatrix_a2, py_kopf_kmatrix_f0, py_kopf_kmatrix_f2,
36            py_kopf_kmatrix_pi1, py_kopf_kmatrix_rho,
37        },
38        phase_space::py_phase_space_factor,
39        piecewise::{
40            py_piecewise_complex_scalar, py_piecewise_polar_complex_scalar, py_piecewise_scalar,
41        },
42        ylm::py_ylm,
43        zlm::{py_polphase, py_zlm},
44    };
45
46    #[pymodule_export]
47    use laddu_extensions::{
48        ganesh_ext::py_ganesh::{
49            py_integrated_autocorrelation_times, PyAutocorrelationTerminator, PyControlFlow,
50            PyEnsembleStatus, PyMCMCSummary, PyMinimizationStatus, PyMinimizationSummary, PySwarm,
51            PySwarmParticle, PyWalker,
52        },
53        likelihoods::{
54            py_likelihood_one, py_likelihood_product, py_likelihood_scalar, py_likelihood_sum,
55            py_likelihood_zero, PyLikelihoodEvaluator, PyLikelihoodExpression, PyNLL,
56            PyStochasticNLL,
57        },
58    };
59
60    #[pymodule_export]
61    use laddu_extensions::experimental::{py_binned_guide_term, py_regularizer};
62
63    #[pymodule_export]
64    use laddu_python::data::{read_parquet, read_root, write_parquet, write_root};
65}