Merge pull request #308 from NNPDF/python-list-args

Fix Python list arguments

pineappl_py/src/bin.rs

@@ -1,6 +1,5 @@
 //! Binnning interface.
 
-use numpy::{PyArrayMethods, PyReadonlyArray1};
 use pineappl::bin::BinRemapper;
 use pyo3::prelude::*;
 
@@ -23,9 +22,9 @@ impl PyBinRemapper {
     /// limits : list(tuple(float, float))
     ///     bin limits
     #[new]
-    pub fn new(normalizations: PyReadonlyArray1<f64>, limits: Vec<(f64, f64)>) -> Self {
+    pub fn new(normalizations: Vec<f64>, limits: Vec<(f64, f64)>) -> Self {
         Self {
-            bin_remapper: BinRemapper::new(normalizations.to_vec().unwrap(), limits).unwrap(),
+            bin_remapper: BinRemapper::new(normalizations, limits).unwrap(),
         }
     }
 }

pineappl_py/src/fk_table.rs

@@ -1,7 +1,7 @@
 //! FK table interface.
 
 use super::grid::PyGrid;
-use numpy::{IntoPyArray, PyArray1, PyArray4, PyArrayMethods, PyReadonlyArray1};
+use numpy::{IntoPyArray, PyArray1, PyArray4};
 use pineappl::convolutions::LumiCache;
 use pineappl::fk_table::{FkAssumptions, FkTable};
 use pineappl::grid::Grid;
@@ -238,8 +238,8 @@ impl PyFkTable {
         &self,
         pdg_id: i32,
         xfx: &Bound<'py, PyAny>,
-        bin_indices: Option<PyReadonlyArray1<usize>>,
-        channel_mask: Option<PyReadonlyArray1<bool>>,
+        bin_indices: Option<Vec<usize>>,
+        channel_mask: Option<Vec<bool>>,
         py: Python<'py>,
     ) -> Bound<'py, PyArray1<f64>> {
         let mut xfx = |id, x, q2| xfx.call1((id, x, q2)).unwrap().extract().unwrap();
@@ -248,8 +248,8 @@ impl PyFkTable {
         self.fk_table
             .convolve(
                 &mut lumi_cache,
-                &bin_indices.map_or(vec![], |b| b.to_vec().unwrap()),
-                &channel_mask.map_or(vec![], |l| l.to_vec().unwrap()),
+                &bin_indices.unwrap_or_default(),
+                &channel_mask.unwrap_or_default(),
             )
             .into_pyarray_bound(py)
     }
@@ -278,8 +278,8 @@ impl PyFkTable {
         xfx1: &Bound<'py, PyAny>,
         pdg_id2: i32,
         xfx2: &Bound<'py, PyAny>,
-        bin_indices: Option<PyReadonlyArray1<usize>>,
-        channel_mask: Option<PyReadonlyArray1<bool>>,
+        bin_indices: Option<Vec<usize>>,
+        channel_mask: Option<Vec<bool>>,
         py: Python<'py>,
     ) -> Bound<'py, PyArray1<f64>> {
         let mut xfx1 = |id, x, q2| xfx1.call1((id, x, q2)).unwrap().extract().unwrap();
@@ -290,8 +290,8 @@ impl PyFkTable {
         self.fk_table
             .convolve(
                 &mut lumi_cache,
-                &bin_indices.map_or(vec![], |b| b.to_vec().unwrap()),
-                &channel_mask.map_or(vec![], |l| l.to_vec().unwrap()),
+                &bin_indices.unwrap_or_default(),
+                &channel_mask.unwrap_or_default(),
             )
             .into_pyarray_bound(py)
     }

pineappl_py/src/grid.rs

@@ -7,7 +7,7 @@ use super::fk_table::PyFkTable;
 use super::subgrid::{PySubgridEnum, PySubgridParams};
 use itertools::izip;
 use ndarray::CowArray;
-use numpy::{IntoPyArray, PyArray1, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray4};
+use numpy::{IntoPyArray, PyArray1, PyReadonlyArray4};
 use pineappl::convolutions::LumiCache;
 use pineappl::evolution::AlphasTable;
 use pineappl::grid::{Grid, Ntuple};
@@ -45,14 +45,14 @@ impl PyGrid {
     pub fn new_grid(
         channels: Vec<PyRef<PyChannel>>,
         orders: Vec<PyRef<PyOrder>>,
-        bin_limits: PyReadonlyArray1<f64>,
+        bin_limits: Vec<f64>,
         subgrid_params: PySubgridParams,
     ) -> Self {
         Self {
             grid: Grid::new(
                 channels.iter().map(|pyc| pyc.entry.clone()).collect(),
                 orders.iter().map(|pyo| pyo.order.clone()).collect(),
-                bin_limits.to_vec().unwrap(),
+                bin_limits,
                 subgrid_params.subgrid_params,
             ),
         }
@@ -116,20 +116,20 @@ impl PyGrid {
     ///     cross section weight for all events
     pub fn fill_array(
         &mut self,
-        x1s: PyReadonlyArray1<f64>,
-        x2s: PyReadonlyArray1<f64>,
-        q2s: PyReadonlyArray1<f64>,
+        x1s: Vec<f64>,
+        x2s: Vec<f64>,
+        q2s: Vec<f64>,
         order: usize,
-        observables: PyReadonlyArray1<f64>,
+        observables: Vec<f64>,
         channel: usize,
-        weights: PyReadonlyArray1<f64>,
+        weights: Vec<f64>,
     ) {
         for (&x1, &x2, &q2, &observable, &weight) in izip!(
-            x1s.as_array().iter(),
-            x2s.as_array().iter(),
-            q2s.as_array().iter(),
-            observables.as_array().iter(),
-            weights.as_array().iter(),
+            x1s.iter(),
+            x2s.iter(),
+            q2s.iter(),
+            observables.iter(),
+            weights.iter(),
         ) {
             self.grid.fill(
                 order,
@@ -163,7 +163,7 @@ impl PyGrid {
         q2: f64,
         order: usize,
         observable: f64,
-        weights: PyReadonlyArray1<f64>,
+        weights: Vec<f64>,
     ) {
         self.grid.fill_all(
             order,
@@ -174,7 +174,7 @@ impl PyGrid {
                 q2,
                 weight: (),
             },
-            &weights.to_vec().unwrap(),
+            &weights,
         );
     }
 
@@ -267,9 +267,9 @@ impl PyGrid {
         pdg_id: i32,
         xfx: &Bound<'py, PyAny>,
         alphas: &Bound<'py, PyAny>,
-        order_mask: Option<PyReadonlyArray1<bool>>,
-        bin_indices: Option<PyReadonlyArray1<usize>>,
-        channel_mask: Option<PyReadonlyArray1<bool>>,
+        order_mask: Option<Vec<bool>>,
+        bin_indices: Option<Vec<usize>>,
+        channel_mask: Option<Vec<bool>>,
         xi: Option<Vec<(f64, f64)>>,
         py: Python<'py>,
     ) -> Bound<'py, PyArray1<f64>> {
@@ -280,10 +280,10 @@ impl PyGrid {
         self.grid
             .convolve(
                 &mut lumi_cache,
-                &order_mask.map_or(vec![], |b| b.to_vec().unwrap()),
-                &bin_indices.map_or(vec![], |c| c.to_vec().unwrap()),
-                &channel_mask.map_or(vec![], |d| d.to_vec().unwrap()),
-                &xi.map_or(vec![(1.0, 1.0)], |m| m),
+                &order_mask.unwrap_or_default(),
+                &bin_indices.unwrap_or_default(),
+                &channel_mask.unwrap_or_default(),
+                &xi.unwrap_or(vec![(1.0, 1.0)]),
             )
             .into_pyarray_bound(py)
     }
@@ -332,9 +332,9 @@ impl PyGrid {
         pdg_id2: i32,
         xfx2: &Bound<'py, PyAny>,
         alphas: &Bound<'py, PyAny>,
-        order_mask: Option<PyReadonlyArray1<bool>>,
-        bin_indices: Option<PyReadonlyArray1<usize>>,
-        channel_mask: Option<PyReadonlyArray1<bool>>,
+        order_mask: Option<Vec<bool>>,
+        bin_indices: Option<Vec<usize>>,
+        channel_mask: Option<Vec<bool>>,
         xi: Option<Vec<(f64, f64)>>,
         py: Python<'py>,
     ) -> Bound<'py, PyArray1<f64>> {
@@ -347,10 +347,10 @@ impl PyGrid {
         self.grid
             .convolve(
                 &mut lumi_cache,
-                &order_mask.map_or(vec![], |b| b.to_vec().unwrap()),
-                &bin_indices.map_or(vec![], |c| c.to_vec().unwrap()),
-                &channel_mask.map_or(vec![], |d| d.to_vec().unwrap()),
-                &xi.map_or(vec![(1.0, 1.0)], |m| m),
+                &order_mask.unwrap_or_default(),
+                &bin_indices.unwrap_or_default(),
+                &channel_mask.unwrap_or_default(),
+                &xi.unwrap_or(vec![(1.0, 1.0)]),
             )
             .into_pyarray_bound(py)
     }
@@ -365,10 +365,10 @@ impl PyGrid {
     /// Returns
     /// -------
     /// PyEvolveInfo :
-    ///     evolution informations
-    pub fn evolve_info(&self, order_mask: PyReadonlyArray1<bool>) -> PyEvolveInfo {
+    ///     evolution information
+    pub fn evolve_info(&self, order_mask: Vec<bool>) -> PyEvolveInfo {
         PyEvolveInfo {
-            evolve_info: self.grid.evolve_info(order_mask.as_slice().unwrap()),
+            evolve_info: self.grid.evolve_info(order_mask.as_slice()),
         }
     }
 
@@ -394,7 +394,7 @@ impl PyGrid {
     pub fn evolve_with_slice_iter<'py>(
         &self,
         slices: &Bound<'py, PyIterator>,
-        order_mask: PyReadonlyArray1<bool>,
+        order_mask: Vec<bool>,
         xi: (f64, f64),
         ren1: Vec<f64>,
         alphas: Vec<f64>,
@@ -413,8 +413,7 @@ impl PyGrid {
                         CowArray::from(op.as_array().to_owned()),
                     ))
                 }),
-                // TODO: make `order_mask` a `Vec<f64>`
-                &order_mask.to_vec().unwrap(),
+                &order_mask,
                 xi,
                 &AlphasTable { ren1, alphas },
             )
@@ -448,7 +447,7 @@ impl PyGrid {
         &self,
         slices_a: &Bound<'py, PyIterator>,
         slices_b: &Bound<'py, PyIterator>,
-        order_mask: PyReadonlyArray1<bool>,
+        order_mask: Vec<bool>,
         xi: (f64, f64),
         ren1: Vec<f64>,
         alphas: Vec<f64>,
@@ -478,8 +477,7 @@ impl PyGrid {
                         CowArray::from(op.as_array().to_owned()),
                     ))
                 }),
-                // TODO: make `order_mask` a `Vec<f64>`
-                &order_mask.to_vec().unwrap(),
+                &order_mask,
                 xi,
                 &AlphasTable { ren1, alphas },
             )
@@ -648,8 +646,8 @@ impl PyGrid {
     /// ----------
     /// factors : numpy.ndarray[float]
     ///     bin-dependent factors by which to scale
-    pub fn scale_by_bin(&mut self, factors: PyReadonlyArray1<f64>) {
-        self.grid.scale_by_bin(&factors.to_vec().unwrap());
+    pub fn scale_by_bin(&mut self, factors: Vec<f64>) {
+        self.grid.scale_by_bin(&factors);
     }
 
     /// Delete bins.
@@ -660,8 +658,8 @@ impl PyGrid {
     /// ----------
     /// bin_indices : numpy.ndarray[int]
     ///     list of indices of bins to removed
-    pub fn delete_bins(&mut self, bin_indices: PyReadonlyArray1<usize>) {
-        self.grid.delete_bins(&bin_indices.to_vec().unwrap())
+    pub fn delete_bins(&mut self, bin_indices: Vec<usize>) {
+        self.grid.delete_bins(&bin_indices)
     }
 }
 

pineappl_py/src/import_only_subgrid.rs

@@ -1,7 +1,7 @@
 //! PyImportOnlySubgrid* interface.
 
 use super::subgrid::PySubgridEnum;
-use numpy::{PyArrayMethods, PyReadonlyArray1, PyReadonlyArray3};
+use numpy::PyReadonlyArray3;
 use pineappl::import_only_subgrid::ImportOnlySubgridV1;
 use pineappl::import_only_subgrid::ImportOnlySubgridV2;
 use pineappl::sparse_array3::SparseArray3;
@@ -34,14 +34,10 @@ impl PyImportOnlySubgridV2 {
     pub fn new(
         array: PyReadonlyArray3<f64>,
         mu2_grid: Vec<(f64, f64)>,
-        x1_grid: PyReadonlyArray1<f64>,
-        x2_grid: PyReadonlyArray1<f64>,
+        x1_grid: Vec<f64>,
+        x2_grid: Vec<f64>,
     ) -> Self {
-        let mut sparse_array = SparseArray3::new(
-            mu2_grid.len(),
-            x1_grid.as_slice().unwrap().len(),
-            x2_grid.as_slice().unwrap().len(),
-        );
+        let mut sparse_array = SparseArray3::new(mu2_grid.len(), x1_grid.len(), x2_grid.len());
 
         for ((imu2, ix1, ix2), value) in array
             .as_array()
@@ -60,8 +56,8 @@ impl PyImportOnlySubgridV2 {
                         fac: *fac,
                     })
                     .collect(),
-                x1_grid.to_vec().unwrap(),
-                x2_grid.to_vec().unwrap(),
+                x1_grid,
+                x2_grid,
             ),
         }
     }
@@ -112,15 +108,11 @@ impl PyImportOnlySubgridV1 {
     #[new]
     pub fn new_import_only_subgrid(
         array: PyReadonlyArray3<f64>,
-        q2_grid: PyReadonlyArray1<f64>,
-        x1_grid: PyReadonlyArray1<f64>,
-        x2_grid: PyReadonlyArray1<f64>,
+        q2_grid: Vec<f64>,
+        x1_grid: Vec<f64>,
+        x2_grid: Vec<f64>,
     ) -> Self {
-        let mut sparse_array = SparseArray3::new(
-            q2_grid.as_slice().unwrap().len(),
-            x1_grid.as_slice().unwrap().len(),
-            x2_grid.as_slice().unwrap().len(),
-        );
+        let mut sparse_array = SparseArray3::new(q2_grid.len(), x1_grid.len(), x2_grid.len());
 
         for ((iq2, ix1, ix2), value) in array
             .as_array()
@@ -132,9 +124,9 @@ impl PyImportOnlySubgridV1 {
 
         Self::new(ImportOnlySubgridV1::new(
             sparse_array,
-            q2_grid.to_vec().unwrap(),
-            x1_grid.to_vec().unwrap(),
-            x2_grid.to_vec().unwrap(),
+            q2_grid,
+            x1_grid,
+            x2_grid,
         ))
     }
 

pineappl_py/tests/test_grid.py

@@ -96,13 +96,48 @@ def test_convolve_with_one(self):
             g.convolve_with_one(2212, lambda pid, x, q2: 0.0, lambda q2: 0.0),
             [0.0] * 2,
         )
+        v = 5e6 / 9999
         np.testing.assert_allclose(
             g.convolve_with_one(2212, lambda pid, x, q2: 1, lambda q2: 1.0),
-            [5e6 / 9999, 0.0],
+            [v, 0.0],
+        )
+        np.testing.assert_allclose(
+            g.convolve_with_one(
+                2212, lambda pid, x, q2: 1, lambda q2: 1.0, bin_indices=[0]
+            ),
+            [v],
+        )
+        # block first bins with additional args
+        np.testing.assert_allclose(
+            g.convolve_with_one(
+                2212,
+                lambda pid, x, q2: 1,
+                lambda q2: 1.0,
+                bin_indices=[0],
+                order_mask=[False],
+            ),
+            [0.0],
+        )
+        np.testing.assert_allclose(
+            g.convolve_with_one(
+                2212,
+                lambda pid, x, q2: 1,
+                lambda q2: 1.0,
+                bin_indices=[0],
+                channel_mask=[False],
+            ),
+            [0.0],
+        )
+        # second bin is empty
+        np.testing.assert_allclose(
+            g.convolve_with_one(
+                2212, lambda pid, x, q2: 1, lambda q2: 1.0, bin_indices=[1]
+            ),
+            [0.0],
         )
         np.testing.assert_allclose(
             g.convolve_with_one(2212, lambda pid, x, q2: 1, lambda q2: 2.0),
-            [2**3 * 5e6 / 9999, 0.0],
+            [2**3 * v, 0.0],
         )
 
     def test_io(self, tmp_path):