Add tentative Python interface for Grid::evolve_with_slice_iter

NNPDF · Mar 4, 2024 · 90d9efd · 90d9efd
1 parent 58972f2
commit 90d9efd
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 3 deletions.
diff --git a/pineappl_py/src/evolution.rs b/pineappl_py/src/evolution.rs
@@ -1,5 +1,5 @@
 use numpy::{IntoPyArray, PyArray1};
-use pineappl::evolution::EvolveInfo;
+use pineappl::evolution::{EvolveInfo, OperatorSliceInfo};
 
 use pyo3::prelude::*;
 
@@ -36,3 +36,11 @@ impl PyEvolveInfo {
         self.evolve_info.ren1.clone().into_pyarray(py)
     }
 }
+
+/// TODO
+#[pyclass]
+#[repr(transparent)]
+#[derive(Clone)]
+pub struct PyOperatorSliceInfo {
+    pub(crate) slice_info: OperatorSliceInfo,
+}
diff --git a/pineappl_py/src/grid.rs b/pineappl_py/src/grid.rs
@@ -7,21 +7,25 @@ use pineappl::grid::{EkoInfo, GridAxes};
 use pineappl::lumi::LumiCache;
 
 use super::bin::PyBinRemapper;
-use super::evolution::PyEvolveInfo;
+use super::evolution::{PyEvolveInfo, PyOperatorSliceInfo};
 use super::fk_table::PyFkTable;
 use super::lumi::PyLumiEntry;
 use super::subgrid::{PySubgridEnum, PySubgridParams};
 
 use itertools::izip;
-use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1, PyReadonlyArray5};
+use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1, PyReadonlyArray4, PyReadonlyArray5};
 
 use std::collections::HashMap;
+use std::error::Error;
 use std::fs::File;
 use std::io::BufReader;
 use std::path::PathBuf;
 
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
+use pyo3::types::{PyIterator, PyTuple};
+
+use ndarray::CowArray;
 
 /// PyO3 wrapper to :rustdoc:`pineappl::grid::Order <grid/struct.Order.html>`
 ///
@@ -562,6 +566,44 @@ impl PyGrid {
         }
     }
 
+    /// TODO
+    ///
+    /// Parameters
+    /// ----------
+    /// slices : TODO
+    /// order_mask : TODO
+    ///
+    /// Returns
+    /// -------
+    /// TODO
+    pub fn evolve_with_slice_iter(
+        &self,
+        slices: &PyIterator,
+        order_mask: PyReadonlyArray1<bool>,
+    ) -> PyResult<PyFkTable> {
+        Ok(self
+            .grid
+            .evolve_with_slice_iter(
+                slices.map(|result| {
+                    // TODO: check whether we can avoid the `.unwrap` calls
+                    let any = result.unwrap();
+                    let tuple = any.downcast::<PyTuple>().unwrap();
+                    let item0 = tuple.get_item(0).unwrap();
+                    let item1 = tuple.get_item(1).unwrap();
+                    let slice_info = item0.extract::<PyOperatorSliceInfo>().unwrap();
+                    let operator = item1.extract::<PyReadonlyArray4<f64>>().unwrap();
+
+                    // TODO: change `PyErr` into something appropriate
+                    Ok::<_, PyErr>((slice_info.slice_info, CowArray::from(operator.as_array())))
+                }),
+                // TODO: what if it's non-contiguous?
+                order_mask.as_slice().unwrap(),
+            )
+            .map(|fk_table| PyFkTable { fk_table })
+            // TODO: get rid of this `.unwrap` call
+            .unwrap())
+    }
+
     /// Load grid from file.
     ///
     /// **Usage:** `pineko`, FKTable generation