diff --git a/pineappl_py/pineappl/grid.py b/pineappl_py/pineappl/grid.py index b0b6c4cc..2f962d94 100644 --- 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 index 78e74e81..8d6a70cc 100644 --- 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 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, + x0: Vec, + fac1: f64, + pids1: Vec, + x1: Vec, + 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, - _xi: (f64, f64), - _ren1: Vec, - _alphas: Vec, + slices: &PyIterator, + order_mask: PyReadonlyArray1, + xi: (f64, f64), + ren1: Vec, + alphas: Vec, ) -> PyResult { - 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::().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::().unwrap(); - // let operator = item1.extract::>().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)>() + .unwrap(); + Ok::<_, std::io::Error>(( + info.info, + // TODO: avoid copying + CowArray::from(op.as_array().to_owned()), + )) + }), + // TODO: make `order_mask` a `Vec` + &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 index 885c603d..edb669e7 100644 --- 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::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/pineappl_py/tests/test_fk_table.py b/pineappl_py/tests/test_fk_table.py index a5fb537e..e1fa93aa 100644 --- 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