Increase Python coverage and add cosmetic changes

Author: Radonirinaunimi

pineappl_py/tests/test_subgrid.py

@@ -1,5 +1,9 @@
 import numpy as np
 import pytest
+
+from dataclasses import dataclass
+from typing import List, Tuple
+
 from pineappl.boc import Channel, Order
 from pineappl.convolutions import Conv, ConvType
 from pineappl.grid import Grid
@@ -14,13 +18,20 @@
 CONVOBJECT = Conv(conv_type=TYPECONV, pid=2212)
 
 
+@dataclass
+class OperatorInfo:
+    x_grids: List[np.ndarray]
+    scale: List[float]
+    array: np.ndarray
+
+
 def test_issue_164(pdf, fake_grids):
     # https://github.com/NNPDF/pineappl/issues/164
     # DIS-like convolution now ONLY requires one entry of `PID`
     channels = [Channel([([2], 1.0)])]  # DIS-case
     orders = [Order(0, 0, 0, 0, 0)]
 
-    def convolve_grid(q2_min: float = Q2_MIN) -> Grid:
+    def convolve_grid(q2_min: float = Q2_MIN) -> np.ndarray:
         grid = fake_grids.grid_with_generic_convolution(
             nb_convolutions=1,
             orders=orders,
@@ -52,7 +63,7 @@ def convolve_grid(q2_min: float = Q2_MIN) -> Grid:
 
 
 class TestSubgrid:
-    def fake_grid(self, fake_grids):
+    def fake_grid(self, fake_grids) -> Grid:
         channels = [Channel([([2], 1.0)]), Channel([([3], 0.5)])]
         orders = [Order(0, 0, 0, 0, 0)]
         return fake_grids.grid_with_generic_convolution(
@@ -62,19 +73,19 @@ def fake_grid(self, fake_grids):
             convolutions=[CONVOBJECT],
         )
 
-    def fake_importonlysubgrid(self, nb_dim: int = 2) -> tuple:
-        x_grids = [np.linspace(0.1, 1, 2) for _ in range(nb_dim)]
+    def fake_importonlysubgrid(self, nb_xdim: int = 1) -> Tuple[ImportSubgridV1, OperatorInfo]:
+        x_grids = [np.linspace(0.1, 1, 2) for _ in range(nb_xdim)]
         xgrid_size = [x.size for x in x_grids]
         Q2s = np.linspace(10, 20, 2)
         scale = [q2 for q2 in Q2s]  # One single scale Q2
         array = np.random.rand(len(Q2s), *xgrid_size)
+        infos = OperatorInfo(x_grids, scale, array)
         subgrid = ImportSubgridV1(array=array, node_values=[scale, *x_grids])
-        return subgrid, [*x_grids, scale, array]
+        return subgrid, infos
 
     def test_subgrid_methods(self, fake_grids):
         grid = self.fake_grid(fake_grids)
         test_subgrid, infos = self.fake_importonlysubgrid()
-        x1s, x2s, mu2s, _ = (obj for obj in infos)
         grid.set_subgrid(0, 0, 0, test_subgrid.into())
         extr_subgrid = grid.subgrid(0, 0, 0)
         assert isinstance(extr_subgrid, SubgridEnum)
@@ -83,22 +94,27 @@ def test_subgrid_methods(self, fake_grids):
         extr_subgrid.scale(factor=100)
         assert isinstance(extr_subgrid.into(), SubgridEnum)
 
-    @pytest.mark.parametrize("nb_dim", [1, 2, 3, 4])
-    def test_subgrid_arrays(self, nb_dim):
+    @pytest.mark.parametrize("nb_xdim", [1, 2, 3, 4])
+    def test_subgrid_arrays(self, nb_xdim: int):
         """This simply checks that the commands run without raising any
         errors and that the objects have been succesfully instantiated.
         """
-        subgrid, info = self.fake_importonlysubgrid(nb_dim=nb_dim)
+        subgrid, info = self.fake_importonlysubgrid(nb_xdim=nb_xdim)
         assert isinstance(subgrid, ImportSubgridV1)
 
-    @pytest.mark.skip(reason="No implementation of Array3 for subgrid.")
-    def test_to_array3(self, fake_grids):
-        # TODO: extract and check the dense array of the subgrid
-        # requires `impl From<&SubgridEnum> for Array3<f64>`
+    def test_to_array(self, fake_grids):
         grid = self.fake_grid(fake_grids)
         test_subgrid, infos = self.fake_importonlysubgrid()
-        _, _, _, array = (obj for obj in infos)
         grid.set_subgrid(0, 0, 0, test_subgrid.into())
         extr_subgrid = grid.subgrid(0, 0, 0)
-        test_array = extr_subgrid.to_array3()
-        np.testing.assert_allclose(test_array, array)
+
+        # Check that the shape of the subgrid matches specs
+        extr_subgrid_shape = extr_subgrid.shape
+        assert tuple(extr_subgrid_shape) == infos.array.shape
+
+        # Check that the `node_values` correspond with the Kinematics
+        node_values = extr_subgrid.node_values
+        np.testing.assert_allclose(node_values, [infos.scale, *infos.x_grids])
+
+        test_array = extr_subgrid.to_array(shape=extr_subgrid_shape)
+        np.testing.assert_allclose(test_array, infos.array)