Add merge to Python API

pineappl_py/pineappl/grid.py

@@ -6,8 +6,7 @@
 
 
 class Order(PyWrapper):
-    r"""
-    Python wrapper object to interface :class:`~pineappl.pineappl.PyOrder`.
+    r"""Python wrapper object to interface :class:`~pineappl.pineappl.PyOrder`.
 
     Parameters
     ----------
@@ -51,8 +50,7 @@ def create_mask(orders, max_as, max_al, logs):
 
 
 class Grid(PyWrapper):
-    r"""
-    Python wrapper object to interface :class:`~pineappl.pineappl.PyGrid`.
+    r"""Python wrapper object to interface :class:`~pineappl.pineappl.PyGrid`.
 
     To create an object, you should call either :meth:`create`
     or :meth:`read`.
@@ -68,8 +66,7 @@ def __init__(self, pygrid):
 
     @classmethod
     def create(cls, lumi, orders, bin_limits, subgrid_params):
-        """
-        Create a grid object from its ingredients
+        """Create a grid object from its ingredients.
 
         Parameters
         ---------
@@ -87,8 +84,7 @@ def create(cls, lumi, orders, bin_limits, subgrid_params):
         return cls(PyGrid(lumi, orders, np.array(bin_limits), subgrid_params.raw))
 
     def subgrid(self, order, bin_, lumi):
-        """
-        Retrieve the subgrid at the given position.
+        """Retrieve the subgrid at the given position.
 
         Convenience wrapper for :meth:`pineappl.pineappl.PyGrid.set_subgrid()`.
 
@@ -109,8 +105,7 @@ def subgrid(self, order, bin_, lumi):
         return self.raw.subgrid(order, bin_, lumi)
 
     def __getitem__(self, key):
-        """
-        Retrieve the subgrid at the given position.
+        """Retrieve the subgrid at the given position.
 
         Syntactic sugar for :meth:`subgrid`
 
@@ -130,8 +125,7 @@ def __getitem__(self, key):
         return self.subgrid(*key)
 
     def set_subgrid(self, order, bin_, lumi, subgrid):
-        """
-        Set the subgrid at the given position.
+        """Set the subgrid at the given position.
 
         Convenience wrapper for :meth:`pineappl.pineappl.PyGrid.set_subgrid()`.
 
@@ -149,8 +143,7 @@ def set_subgrid(self, order, bin_, lumi, subgrid):
         self.raw.set_subgrid(order, bin_, lumi, subgrid.into())
 
     def __setitem__(self, key, subgrid):
-        """
-        Set the subgrid at the given position.
+        """Set the subgrid at the given position.
 
         Syntactic sugar for :meth:`set_subgrid`
 
@@ -167,8 +160,7 @@ def __setitem__(self, key, subgrid):
         self.set_subgrid(*key, subgrid)
 
     def set_remapper(self, remapper):
-        """
-        Set the normalizations.
+        """Set the normalizations.
 
         Convenience wrapper for :meth:`pineappl.pineappl.PyGrid.set_remapper()`.
 
@@ -180,8 +172,7 @@ def set_remapper(self, remapper):
         self.raw.set_remapper(remapper.raw)
 
     def orders(self):
-        """
-        Extract the available perturbative orders and scale variations.
+        """Extract the available perturbative orders and scale variations.
 
         Convenience wrapper for :meth:`pineappl.pineappl.PyGrid.orders()`.
 
@@ -202,8 +193,7 @@ def convolute_with_one(
         lumi_mask=np.array([], dtype=bool),
         xi=((1.0, 1.0),),
     ):
-        r"""
-        Convolute grid with pdf.
+        r"""Convolute grid with pdf.
 
         Parameters
         ----------
@@ -235,7 +225,6 @@ def convolute_with_one(
             list(float) :
                 cross sections for all bins, for each scale-variation tuple (first all bins, then
                 the scale variation)
-
         """
         return self.raw.convolute_with_one(
             pdg_id,
@@ -247,9 +236,16 @@ def convolute_with_one(
             xi,
         )
 
-    def convolute_eko(self, operators, mur2_grid, alphas_values, lumi_id_types="pdg_mc_ids", order_mask=(), xi=(1.0, 1.0)):
-        """
-        Create an FKTable with the EKO.
+    def convolute_eko(
+        self,
+        operators,
+        mur2_grid,
+        alphas_values,
+        lumi_id_types="pdg_mc_ids",
+        order_mask=(),
+        xi=(1.0, 1.0),
+    ):
+        """Create an FKTable with the EKO.
 
         Convenience wrapper for :meth:`pineappl.pineappl.PyGrid.convolute_eko()`.
 
@@ -296,13 +292,20 @@ def convolute_eko(self, operators, mur2_grid, alphas_values, lumi_id_types="pdg_
                 np.array(operator_grid),
                 lumi_id_types,
                 np.array(order_mask, dtype=bool),
-                xi
+                xi,
             )
         )
 
-    def evolve(self, operators, mur2_grid, alphas_values, lumi_id_types="pdg_mc_ids", order_mask=(), xi=(1.0, 1.0)):
-        """
-        Create an FKTable with the EKO.
+    def evolve(
+        self,
+        operators,
+        mur2_grid,
+        alphas_values,
+        lumi_id_types="pdg_mc_ids",
+        order_mask=(),
+        xi=(1.0, 1.0),
+    ):
+        """Create an FKTable with the EKO.
 
         Convenience wrapper for :meth:`pineappl.pineappl.PyGrid.evolve()`.
 
@@ -355,8 +358,7 @@ def evolve(self, operators, mur2_grid, alphas_values, lumi_id_types="pdg_mc_ids"
 
     @classmethod
     def read(cls, path):
-        """
-        Load an existing grid from file.
+        """Load an existing grid from file.
 
         Convenience wrapper for :meth:`pineappl.pineappl.PyGrid.read()`.
 
@@ -372,6 +374,10 @@ def read(cls, path):
         """
         return cls(PyGrid.read(path))
 
+    def merge(self, other: "Grid"):
+        """Merge a second grid in the current one."""
+        self.raw.merge(other.raw)
+
     def delete_bins(self, bin_indices):
         """Delete bins.
 
@@ -381,6 +387,5 @@ def delete_bins(self, bin_indices):
         ----------
         bin_indices : sequence(int)
             list of indices of bins to removed
-
         """
         self.raw.delete_bins(np.array(bin_indices, dtype=np.uint))

pineappl_py/src/grid.rs

@@ -100,6 +100,7 @@ impl PyOrder {
 /// **Usage**: `yadism`, `pineko`, FKTable interface
 #[pyclass]
 #[repr(transparent)]
+#[derive(Clone)]
 pub struct PyGrid {
     pub(crate) grid: Grid,
 }
@@ -609,12 +610,21 @@ impl PyGrid {
         self.grid.optimize();
     }
 
+    /// Merge grid with another one
+    pub fn merge(&mut self, other: Self) -> PyResult<()> {
+        match self.grid.merge(other.grid) {
+            Ok(()) => Ok(()),
+            Err(x) => Err(PyValueError::new_err(format!("{:?}", x))),
+        }
+    }
+
     /// Merge grid with another one, loaded from file
     ///
     /// Note
     /// ----
     /// For a current limitation with the implementation of the bound object `Grid` is not possible
     /// to operate with two `Grid`s in memory, since is not possible to pass a `Grid` by argument
+    #[deprecated = "Deprecated in favor of PyGrid::merge"]
     pub fn merge_from_file(&mut self, path: PathBuf) -> PyResult<()> {
         match self.grid.merge(Self::read(path).grid) {
             Ok(()) => Ok(()),

pineappl_py/tests/test_grid.py

@@ -1,8 +1,8 @@
-import pineappl
-
 import numpy as np
 import pytest
 
+import pineappl
+
 
 class TestOrder:
     def test_init(self):
@@ -15,10 +15,10 @@ def test_init(self):
 
 
 class TestGrid:
-    def fake_grid(self):
+    def fake_grid(self, bins=None):
         lumis = [pineappl.lumi.LumiEntry([(1, 21, 0.1)])]
         orders = [pineappl.grid.Order(3, 0, 0, 0)]
-        bin_limits = np.array([1e-7, 1e-3, 1])
+        bin_limits = np.array([1e-7, 1e-3, 1] if bins is None else bins, dtype=float)
         subgrid_params = pineappl.subgrid.SubgridParams()
         g = pineappl.grid.Grid.create(lumis, orders, bin_limits, subgrid_params)
         return g
@@ -184,3 +184,31 @@ def test_fill_all(self):
         g.fill_all(1.0, 1.0, 1.0, 0, 1e-2, np.array([10.0]))
         res = g.convolute_with_one(2212, lambda pid, x, q2: x, lambda q2: 1.0)
         pytest.approx(res) == 0.0
+
+    def test_merge(self):
+        g = self.fake_grid([1, 2, 3])
+        g1 = self.fake_grid([3, 4, 5])
+        assert g.bins() == 2
+        assert g1.bins() == 2
+
+        g.merge(g1)
+        assert g.bins() == 4
+
+        g2 = self.fake_grid([1, 2, 3])
+        g3 = self.fake_grid([1, 2, 3])
+        assert g2.bins() == 2
+        assert g3.bins() == 2
+
+        g2.merge(g3)
+        assert g2.bins() == 2
+
+        g4 = self.fake_grid([2, 3, 4])
+        g5 = self.fake_grid([4, 5, 6])
+        assert g4.bins() == 2
+        assert g5.bins() == 2
+
+        with pytest.raises(ValueError, match="NonConsecutiveBins"):
+            g2.merge(g4)
+
+        with pytest.raises(ValueError, match="NonConsecutiveBins"):
+            g2.merge(g5)