Implement Python method Grid.evolve_with_slice_iter

NNPDF · Jul 11, 2024 · 94e4ab2 · 94e4ab2
1 parent ed8601d
commit 94e4ab2
Show file tree

Hide file tree

Showing 4 changed files with 106 additions and 35 deletions.
diff --git a/pineappl_py/pineappl/grid.py b/pineappl_py/pineappl/grid.py
@@ -1,7 +1,7 @@
 import numpy as np
 
 from .fk_table import FkTable
-from .pineappl import PyGrid, PyOrder
+from .pineappl import PyGrid, PyOrder, PyOperatorSliceInfo, PyPidBasis
 from .utils import PyWrapper
 
 

diff --git a/pineappl_py/src/grid.rs b/pineappl_py/src/grid.rs
@@ -1,7 +1,9 @@
+use ndarray::CowArray;
 use pineappl::boc::Order;
 use pineappl::convolutions::LumiCache;
-use pineappl::evolution::OperatorInfo;
+use pineappl::evolution::{AlphasTable, OperatorInfo, OperatorSliceInfo};
 use pineappl::grid::{Grid, Ntuple};
+use pineappl::pids::PidBasis;
 
 use super::bin::PyBinRemapper;
 use super::evolution::PyEvolveInfo;
@@ -10,7 +12,7 @@ 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::fs::File;
@@ -30,6 +32,58 @@ pub struct PyOrder {
     pub(crate) order: Order,
 }
 
+// TODO: should probably be in a different module
+// TODO: rename to `PidBasis`
+#[pyclass]
+#[derive(Clone)]
+pub enum PyPidBasis {
+    Pdg,
+    Evol,
+}
+
+impl From<PyPidBasis> for PidBasis {
+    fn from(basis: PyPidBasis) -> Self {
+        match basis {
+            PyPidBasis::Pdg => Self::Pdg,
+            PyPidBasis::Evol => Self::Evol,
+        }
+    }
+}
+
+// TODO: should probably be in a different module
+// TODO: rename to `OperatorSliceInfo`
+#[pyclass]
+#[derive(Clone)]
+pub struct PyOperatorSliceInfo {
+    info: OperatorSliceInfo,
+}
+
+#[pymethods]
+impl PyOperatorSliceInfo {
+    #[new]
+    pub fn new(
+        fac0: f64,
+        pids0: Vec<i32>,
+        x0: Vec<f64>,
+        fac1: f64,
+        pids1: Vec<i32>,
+        x1: Vec<f64>,
+        pid_basis: PyPidBasis,
+    ) -> Self {
+        Self {
+            info: OperatorSliceInfo {
+                fac0,
+                pids0,
+                x0,
+                fac1,
+                pids1,
+                x1,
+                pid_basis: pid_basis.into(),
+            },
+        }
+    }
+}
+
 impl PyOrder {
     pub(crate) fn new(order: Order) -> Self {
         Self { order }
@@ -523,39 +577,34 @@ impl PyGrid {
     /// TODO
     pub fn evolve_with_slice_iter(
         &self,
-        _slices: &PyIterator,
-        _order_mask: PyReadonlyArray1<bool>,
-        _xi: (f64, f64),
-        _ren1: Vec<f64>,
-        _alphas: Vec<f64>,
+        slices: &PyIterator,
+        order_mask: PyReadonlyArray1<bool>,
+        xi: (f64, f64),
+        ren1: Vec<f64>,
+        alphas: Vec<f64>,
     ) -> PyResult<PyFkTable> {
-        todo!()
-        //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();
-        //            // TODO: `get_item` changes return type from pyo3-0.14 to 0.15
-        //            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: can we get rid of the `into_owned` call?
-        //            let array = CowArray::from(operator.as_array().into_owned());
-
-        //            // TODO: change `PyErr` into something appropriate
-        //            Ok::<_, PyErr>((slice_info.slice_info, array))
-        //        }),
-        //        // TODO: what if it's non-contiguous?
-        //        order_mask.as_slice().unwrap(),
-        //        xi,
-        //        &AlphasTable { ren1, alphas },
-        //    )
-        //    .map(|fk_table| PyFkTable { fk_table })
-        //    // TODO: get rid of this `.unwrap` call
-        //    .unwrap())
+        Ok(self
+            .grid
+            .evolve_with_slice_iter(
+                slices.into_iter().map(|slice| {
+                    let (info, op) = slice
+                        .unwrap()
+                        .extract::<(PyOperatorSliceInfo, PyReadonlyArray4<f64>)>()
+                        .unwrap();
+                    Ok::<_, std::io::Error>((
+                        info.info,
+                        // TODO: avoid copying
+                        CowArray::from(op.as_array().to_owned()),
+                    ))
+                }),
+                // TODO: make `order_mask` a `Vec<f64>`
+                &order_mask.to_vec().unwrap(),
+                xi,
+                &AlphasTable { ren1, alphas },
+            )
+            .map(|fk_table| PyFkTable { fk_table })
+            // TODO: avoid unwrap and convert `Result` into `PyResult`
+            .unwrap())
     }
 
     /// Load grid from file.

diff --git a/pineappl_py/src/lib.rs b/pineappl_py/src/lib.rs
@@ -20,6 +20,8 @@ fn pineappl(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<evolution::PyEvolveInfo>()?;
     m.add_class::<grid::PyGrid>()?;
     m.add_class::<grid::PyOrder>()?;
+    m.add_class::<grid::PyOperatorSliceInfo>()?;
+    m.add_class::<grid::PyPidBasis>()?;
     m.add_class::<lumi::PyLumiEntry>()?;
     m.add_class::<import_only_subgrid::PyImportOnlySubgridV1>()?;
     m.add_class::<import_only_subgrid::PyImportOnlySubgridV2>()?;

diff --git a/pineappl_py/tests/test_fk_table.py b/pineappl_py/tests/test_fk_table.py
@@ -34,3 +34,23 @@ def test_convolve_with_one(self):
             fk.convolve_with_one(2212, lambda pid, x, q2: 1),
             [5e7 / 9999, 0.0],
         )
+
+        info = pineappl.grid.PyOperatorSliceInfo(
+            1.0, [], [], 1.0, [], [], pineappl.grid.PyPidBasis.Pdg
+        )
+
+        # TODO: write a better test
+        try:
+            g.evolve_with_slice_iter(
+                iter(
+                    [(info, np.ndarray([0, 0, 0, 0])), (info, np.ndarray([0, 0, 0, 0]))]
+                ),
+                np.array([], dtype=bool),
+                (1.0, 1.0),
+                [],
+                [],
+            )
+
+            assert False
+        except:
+            assert True