Add preliminary Python interface to packed_subgrid

pineappl_py/src/grid.rs

@@ -7,6 +7,7 @@ use super::evolution::{PyEvolveInfo, PyOperatorSliceInfo};
 use super::fk_table::PyFkTable;
 use super::interpolation::PyInterp;
 use super::pids::PyPidBasis;
+use super::subgrid::PySubgridEnum;
 use ndarray::CowArray;
 use numpy::{IntoPyArray, PyArray1, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray4};
 use pineappl::boc::{ScaleFuncForm, Scales};
@@ -107,6 +108,25 @@ impl PyGrid {
         self.grid.fill(order, observable, channel, &ntuple, weight);
     }
 
+    /// Retrieve a subgrid.
+    #[must_use]
+    pub fn subgrid(&self, order: usize, bin: usize, channel: usize) -> PySubgridEnum {
+        PySubgridEnum {
+            subgrid_enum: self.grid.subgrids()[[order, bin, channel]].clone(),
+        }
+    }
+
+    /// Set a subgrid.
+    pub fn set_subgrid(
+        &mut self,
+        order: usize,
+        bin: usize,
+        channel: usize,
+        subgrid: PySubgridEnum,
+    ) {
+        self.grid.subgrids_mut()[[order, bin, channel]] = subgrid.subgrid_enum;
+    }
+
     /// Set the bin normalizations.
     ///
     /// # Panics

pineappl_py/src/lib.rs

@@ -11,6 +11,7 @@ pub mod evolution;
 pub mod fk_table;
 pub mod grid;
 pub mod interpolation;
+pub mod packed_subgrid;
 pub mod pids;
 pub mod subgrid;
 

pineappl_py/src/packed_subgrid.rs

@@ -0,0 +1,87 @@
+//! PyPackedSubgrid* interface.
+
+use super::subgrid::PySubgridEnum;
+use numpy::PyReadonlyArray3;
+use pineappl::packed_array::PackedArray;
+use pineappl::packed_subgrid::PackedQ1X2SubgridV1;
+use pineappl::subgrid::NodeValues;
+use pyo3::prelude::*;
+
+/// PyO3 wrapper to :rustdoc:`pineappl`.
+#[pyclass(name = "PyPackedSubgrid")]
+#[derive(Clone)]
+#[repr(transparent)]
+pub struct PyPackedSubgrid {
+    pub(crate) packed_subgrid: PackedQ1X2SubgridV1,
+}
+
+#[pymethods]
+impl PyPackedSubgrid {
+    /// Constructor.
+    /// Constructor.
+    ///
+    /// Parameters
+    /// ----------
+    /// array : numpy.ndarray(float)
+    ///     3D array with all weights
+    /// scales : list(float)
+    ///     scales grid
+    /// x1_grid : list(float)
+    ///     first momentum fraction grid
+    /// x2_grid : list(float)
+    ///     second momentum fraction grid
+    #[new]
+    pub fn new(
+        array: PyReadonlyArray3<f64>,
+        scales: Vec<f64>,
+        x1_grid: Vec<f64>,
+        x2_grid: Vec<f64>,
+    ) -> Self {
+        let node_values: Vec<NodeValues> = vec![
+            NodeValues::UseThese(scales),
+            NodeValues::UseThese(x1_grid),
+            NodeValues::UseThese(x2_grid),
+        ];
+        let mut sparse_array: PackedArray<f64> =
+            PackedArray::new(node_values.iter().map(NodeValues::len).collect());
+
+        for ((iscale, ix1, ix2), value) in array
+            .as_array()
+            .indexed_iter()
+            .filter(|((_, _, _), value)| **value != 0.0)
+        {
+            sparse_array[[iscale, ix1, ix2]] = *value;
+        }
+
+        Self {
+            packed_subgrid: PackedQ1X2SubgridV1::new(sparse_array, node_values),
+        }
+    }
+
+    /// TODO
+    #[must_use]
+    pub fn into(&self) -> PySubgridEnum {
+        PySubgridEnum {
+            subgrid_enum: self.packed_subgrid.clone().into(),
+        }
+    }
+}
+
+/// Register submodule in parent.
+/// # Errors
+///
+/// Raises an error if (sub)module is not found.
+pub fn register(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
+    let m = PyModule::new_bound(parent_module.py(), "packed_subgrid")?;
+    m.setattr(
+        pyo3::intern!(m.py(), "__doc__"),
+        "Interface for packed subgrid specs.",
+    )?;
+    pyo3::py_run!(
+        parent_module.py(),
+        m,
+        "import sys; sys.modules['pineappl.packed_subgrid'] = m"
+    );
+    m.add_class::<PyPackedSubgrid>()?;
+    parent_module.add_submodule(&m)
+}