Adjust PyGrid::convolve_*

NNPDF · Sep 2, 2024 · 8a1e929 · 8a1e929
1 parent 712b22e
commit 8a1e929
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 14 deletions.
diff --git a/pineappl_py/src/grid.rs b/pineappl_py/src/grid.rs
@@ -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,9 +280,9 @@ 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()),
+                &order_mask.map_or(vec![], |b| b.to_vec()),
+                &bin_indices.map_or(vec![], |c| c.to_vec()),
+                &channel_mask.map_or(vec![], |d| d.to_vec()),
                 &xi.map_or(vec![(1.0, 1.0)], |m| m),
             )
             .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,9 +347,9 @@ 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()),
+                &order_mask.map_or(vec![], |b| b.to_vec()),
+                &bin_indices.map_or(vec![], |c| c.to_vec()),
+                &channel_mask.map_or(vec![], |d| d.to_vec()),
                 &xi.map_or(vec![(1.0, 1.0)], |m| m),
             )
             .into_pyarray_bound(py)

diff --git a/pineappl_py/tests/test_grid.py b/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):