From 8a1e9295396b4b3f88b8d880f3eb863712a1e1af Mon Sep 17 00:00:00 2001
From: Felix Hekhorn <felix.hekhorn@uni-tuebingen.de>
Date: Mon, 2 Sep 2024 16:06:49 +0300
Subject: [PATCH] Adjust PyGrid::convolve_*

---
 pineappl_py/src/grid.rs        | 24 ++++++++++-----------
 pineappl_py/tests/test_grid.py | 39 ++++++++++++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/pineappl_py/src/grid.rs b/pineappl_py/src/grid.rs
index 6eda73c1..f6bdc48a 100644
--- 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
index f0678db9..fbc4d17d 100644
--- 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):