neopdf/
gridpdf.rs

1use ndarray::Array1;
2use numpy::{PyArrayMethods, PyReadonlyArray6};
3use pyo3::prelude::*;
4
5use neopdf::gridpdf::GridArray;
6use neopdf::subgrid::{ParamRange, SubGrid};
7
8/// Python wrapper for the `SubGrid` struct.
9#[pyclass(name = "SubGrid")]
10pub struct PySubGrid {
11    pub(crate) subgrid: SubGrid,
12}
13
14#[pymethods]
15impl PySubGrid {
16    /// Constructs a new `SubGrid` instance from the provided axes and grid data.
17    ///
18    /// # Parameters
19    ///
20    /// - `xs`: The x-axis values.
21    /// - `q2s`: The Q^2-axis values.
22    /// - `kts`: The kT-axis values.
23    /// - `nucleons`: The nucleon number axis values.
24    /// - `alphas`: The alpha_s axis values.
25    /// - `grid`: The 6D grid data as a NumPy array.
26    ///
27    /// # Returns
28    ///
29    /// Returns a new `PySubGrid` instance.
30    ///
31    /// # Panics
32    ///
33    /// Panics if any of the input vectors are empty.
34    ///
35    /// # Errors
36    ///
37    /// Returns a `PyErr` if the grid cannot be constructed from the input data.
38    #[new]
39    #[allow(clippy::needless_pass_by_value)]
40    pub fn new(
41        xs: Vec<f64>,
42        q2s: Vec<f64>,
43        kts: Vec<f64>,
44        nucleons: Vec<f64>,
45        alphas: Vec<f64>,
46        grid: PyReadonlyArray6<f64>,
47    ) -> PyResult<Self> {
48        let alphas_range = ParamRange::new(*alphas.first().unwrap(), *alphas.last().unwrap());
49        let x_range = ParamRange::new(*xs.first().unwrap(), *xs.last().unwrap());
50        let q2_range = ParamRange::new(*q2s.first().unwrap(), *q2s.last().unwrap());
51        let kt_range = ParamRange::new(*kts.first().unwrap(), *kts.last().unwrap());
52        let nucleons_range = ParamRange::new(*nucleons.first().unwrap(), *nucleons.last().unwrap());
53
54        let subgrid = SubGrid {
55            xs: Array1::from(xs),
56            q2s: Array1::from(q2s),
57            kts: Array1::from(kts),
58            grid: grid.to_owned_array(),
59            nucleons: Array1::from(nucleons),
60            alphas: Array1::from(alphas),
61            nucleons_range,
62            alphas_range,
63            kt_range,
64            x_range,
65            q2_range,
66        };
67
68        Ok(Self { subgrid })
69    }
70
71    /// Returns the minimum and maximum values of the alpha_s axis.
72    #[must_use]
73    pub const fn alphas_range(&self) -> (f64, f64) {
74        (self.subgrid.alphas_range.min, self.subgrid.alphas_range.max)
75    }
76
77    /// Returns the minimum and maximum values of the x axis.
78    #[must_use]
79    pub const fn x_range(&self) -> (f64, f64) {
80        (self.subgrid.x_range.min, self.subgrid.x_range.max)
81    }
82
83    /// Returns the minimum and maximum values of the Q^2 axis.
84    #[must_use]
85    pub const fn q2_range(&self) -> (f64, f64) {
86        (self.subgrid.q2_range.min, self.subgrid.q2_range.max)
87    }
88
89    /// Returns the shape of the subgrid
90    #[must_use]
91    pub fn grid_shape(&self) -> (usize, usize, usize, usize, usize, usize) {
92        self.subgrid.grid.dim()
93    }
94}
95
96/// Python wrapper for the `GridArray` struct.
97#[pyclass(name = "GridArray")]
98#[repr(transparent)]
99pub struct PyGridArray {
100    pub(crate) gridarray: GridArray,
101}
102
103#[pymethods]
104impl PyGridArray {
105    /// Constructs a new `GridArray` from a list of particle IDs and subgrids.
106    ///
107    /// # Parameters
108    ///
109    /// - `pids`: The list of particle IDs.
110    /// - `subgrids`: The list of subgrid objects.
111    ///
112    /// # Returns
113    ///
114    /// Returns a new `PyGridArray` instance.
115    #[new]
116    #[must_use]
117    pub fn new(pids: Vec<i32>, subgrids: Vec<PyRef<PySubGrid>>) -> Self {
118        let subgrids = subgrids
119            .into_iter()
120            .map(|py_ref| py_ref.subgrid.clone())
121            .collect();
122
123        let gridarray = GridArray {
124            pids: Array1::from(pids),
125            subgrids,
126        };
127        Self { gridarray }
128    }
129
130    /// Returns the particle IDs associated with this grid array.
131    #[must_use]
132    pub fn pids(&self) -> Vec<i32> {
133        self.gridarray.pids.to_vec()
134    }
135
136    /// Returns the subgrids contained in this grid array.
137    #[must_use]
138    pub fn subgrids(&self) -> Vec<PySubGrid> {
139        self.gridarray
140            .subgrids
141            .iter()
142            .cloned()
143            .map(|sg| PySubGrid { subgrid: sg })
144            .collect()
145    }
146}
147
148/// Registers the gridpdf module with the parent Python module.
149///
150/// Adds the `gridpdf` submodule to the parent Python module, exposing grid
151/// interpolation utilities to Python.
152///
153/// # Errors
154///
155/// Returns a `PyErr` if the submodule cannot be created or added, or if any
156/// class registration fails.
157pub fn register(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
158    let m = PyModule::new(parent_module.py(), "gridpdf")?;
159    m.setattr(
160        pyo3::intern!(m.py(), "__doc__"),
161        "GridPDF interpolation interface.",
162    )?;
163    pyo3::py_run!(
164        parent_module.py(),
165        m,
166        "import sys; sys.modules['neopdf.gridpdf'] = m"
167    );
168    m.add_class::<PySubGrid>()?;
169    m.add_class::<PyGridArray>()?;
170    parent_module.add_submodule(&m)
171}
neopdf/gridpdf.rs

neopdf/
gridpdf.rs
