Exposes methods that would be required by pineko#183 and more

Author: Radonirinaunimi

pineappl_py/src/grid.rs

@@ -698,6 +698,16 @@ impl PyGrid {
             .collect()
     }
 
+    /// Rotate the Grid into the specified basis
+    ///
+    /// Parameters
+    /// ----------
+    /// pid_basis: PyPidBasis
+    ///     PID basis of the resulting Grid
+    pub fn rotate_pid_basis(&mut self, pid_basis: PyPidBasis) {
+        self.grid.rotate_pid_basis(pid_basis.into());
+    }
+
     /// Scale all subgrids.
     ///
     /// Parameters
@@ -716,12 +726,23 @@ impl PyGrid {
     ///
     /// Parameters
     /// ----------
-    /// factors : numpy.ndarray[float]
+    /// factors : list[float]
     ///     bin-dependent factors by which to scale
     pub fn scale_by_bin(&mut self, factors: Vec<f64>) {
         self.grid.scale_by_bin(&factors);
     }
 
+    /// Delete orders with the corresponding `order_indices`. Repeated indices and indices larger
+    /// or equal than the number of orders are ignored.
+    ///
+    /// Parameters
+    /// ----------
+    /// order_indices : list[int]
+    ///     list of indices of orders to be removed
+    pub fn delete_orders(&mut self, order_indices: Vec<usize>) {
+        self.grid.delete_orders(&order_indices);
+    }
+
     /// Delete bins.
     ///
     /// # Panics
@@ -732,11 +753,27 @@ impl PyGrid {
     ///
     /// Parameters
     /// ----------
-    /// bin_indices : numpy.ndarray[int]
-    ///     list of indices of bins to removed
+    /// bin_indices : list[int]
+    ///     list of indices of bins to be removed
     pub fn delete_bins(&mut self, bin_indices: Vec<usize>) {
         self.grid.delete_bins(&bin_indices);
     }
+
+    /// Deletes channels with the corresponding `channel_indices`. Repeated indices and indices
+    /// larger or equal than the number of channels are ignored.
+    ///
+    /// Parameters
+    /// ----------
+    /// bin_indices : list[int]
+    ///     list of indices of bins to be removed
+    pub fn delete_channels(&mut self, channel_indices: Vec<usize>) {
+        self.grid.delete_channels(&channel_indices);
+    }
+
+    /// Splits the grid such that each channel contains only a single tuple of PIDs.
+    pub fn split_channels(&mut self) {
+        self.grid.split_channels();
+    }
 }
 
 /// Register submodule in parent.

pineappl_py/tests/test_grid.py

@@ -212,6 +212,55 @@ def test_bins(self, fake_grids):
         np.testing.assert_allclose(g.bin_left(1), [2, 3])
         np.testing.assert_allclose(g.bin_right(1), [3, 5])
 
+    def test_rotate_pidbasis(self, fake_grids):
+        g = fake_grids.grid_with_generic_convolution(
+            nb_convolutions=2,
+            channels=CHANNELS,
+            orders=ORDERS,
+            convolutions=[CONVOBJECT, CONVOBJECT],
+        )
+        # Rotate the Grid into the PDG basis
+        g.rotate_pid_basis(PidBasis.Pdg)
+        assert g.pid_basis == PidBasis.Pdg
+
+    def test_delete_orders(
+        self,
+        download_objects,
+        gridname: str = "GRID_STAR_WMWP_510GEV_WP-AL-POL.pineappl.lz4",
+        order_indices: list[int] = [1],
+    ):
+        grid = download_objects(f"{gridname}")
+        g = Grid.read(grid)
+        orders = [o.as_tuple() for o in g.orders()]
+        g.delete_orders(order_indices)
+        for idx in order_indices:
+            assert orders[idx] not in g.orders()
+
+    def test_delete_channels(
+        self,
+        download_objects,
+        gridname: str = "GRID_STAR_WMWP_510GEV_WP-AL-POL.pineappl.lz4",
+        channel_indices: list[int] = [1, 4, 5],
+    ):
+        grid = download_objects(f"{gridname}")
+        g = Grid.read(grid)
+        channels = g.channels()
+        g.delete_channels(channel_indices)
+        for idx in channel_indices:
+            assert channels[idx] not in g.channels()
+
+    def test_split_channels(
+        self,
+        pdf,
+        download_objects,
+        gridname: str = "GRID_DYE906R_D_bin_1.pineappl.lz4",
+    ):
+        grid = download_objects(f"{gridname}")
+        g = Grid.read(grid)
+        assert len(g.channels()) == 15
+        g.split_channels()
+        assert len(g.channels()) == 170
+
     def test_grid(
         self,
         download_objects,