diff --git a/pineappl_py/src/grid.rs b/pineappl_py/src/grid.rs index e066f772..42910571 100644 --- a/pineappl_py/src/grid.rs +++ b/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) { 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) { + 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) { 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) { + 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. diff --git a/pineappl_py/tests/test_grid.py b/pineappl_py/tests/test_grid.py index 6cb0203c..baf64023 100644 --- a/pineappl_py/tests/test_grid.py +++ b/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,