Modify ImportSubgridV1 to accept a generic N-dimensional objects

pineappl_py/src/import_subgrid.rs

@@ -1,7 +1,8 @@
 //! PyPackedSubgrid* interface.
 
 use super::subgrid::PySubgridEnum;
-use numpy::{PyReadonlyArray2, PyReadonlyArray3};
+use ndarray::Dimension;
+use numpy::PyReadonlyArrayDyn;
 use pineappl::import_subgrid::ImportSubgridV1;
 use pineappl::packed_array::PackedArray;
 use pyo3::prelude::*;
@@ -23,63 +24,28 @@ impl PyImportSubgridV1 {
     /// Parameters
     /// ----------
     /// array : numpy.ndarray(float)
-    ///     3D array with all weights
+    ///     `N`-dimensional array with all weights
     /// scales : list(float)
     ///     scales grid
-    /// x1_grid : list(float)
-    ///     first momentum fraction grid
-    /// x2_grid : Optional(list(float))
-    ///     second momentum fraction grid
+    /// x_grids : list(list(float))
+    ///     list with length `N` containing the momentum fractions (x1, x2, ...)
+    ///     which are also expressed as lists
     #[new]
     #[must_use]
-    pub fn new<'py>(
-        array: PyObject,
-        scales: Vec<f64>,
-        x1_grid: Vec<f64>,
-        x2_grid: Option<Vec<f64>>,
-        py: Python<'py>,
-    ) -> Self {
-        // let node_values: Vec<Vec<f64>> = vec![scales, x1_grid, x2_grid];
-        // let mut sparse_array: PackedArray<f64> =
-        //     PackedArray::new(node_values.iter().map(Vec::len).collect());
+    pub fn new(array: PyReadonlyArrayDyn<f64>, scales: Vec<f64>, x_grids: Vec<Vec<f64>>) -> Self {
+        // Total number of nodes = (scales + x-grids)
+        let mut node_values: Vec<Vec<f64>> = vec![scales];
+        node_values.extend(x_grids); // Extend nodes with x-grids
 
-        // for ((iscale, ix1, ix2), value) in array
-        //     .as_array()
-        //     .indexed_iter()
-        //     .filter(|((_, _, _), value)| **value != 0.0)
-        // {
-        //     sparse_array[[iscale, ix1, ix2]] = *value;
-        // }
-
-        // Self {
-        //     import_subgrid: ImportSubgridV1::new(sparse_array, node_values),
-        // }
-        let mut node_values: Vec<Vec<f64>> = vec![scales, x1_grid];
-
-        if let Some(x2) = x2_grid {
-            node_values.push(x2);
-        }
         let mut sparse_array: PackedArray<f64> =
             PackedArray::new(node_values.iter().map(Vec::len).collect());
 
-        if sparse_array.shape().to_vec().len() == 3 {
-            let array_3d: PyReadonlyArray3<f64> = array.extract(py).unwrap();
-            for ((iscale, ix1, ix2), value) in array_3d
-                .as_array()
-                .indexed_iter()
-                .filter(|((_, _, _), value)| **value != 0.0)
-            {
-                sparse_array[[iscale, ix1, ix2]] = *value;
-            }
-        } else {
-            let array_2d: PyReadonlyArray2<f64> = array.extract(py).unwrap();
-            for ((iscale, ix1), value) in array_2d
-                .as_array()
-                .indexed_iter()
-                .filter(|((_, _), value)| **value != 0.0)
-            {
-                sparse_array[[iscale, ix1]] = *value;
-            }
+        for (index, value) in array
+            .as_array()
+            .indexed_iter()
+            .filter(|(_, value)| **value != 0.0)
+        {
+            sparse_array[index.as_array_view().to_slice().unwrap()] = *value;
         }
 
         Self {

pineappl_py/tests/conftest.py

@@ -1,6 +1,7 @@
-import pytest
 import subprocess
 
+import pytest
+
 
 class PDF:
     def xfxQ2(self, pid, x, q2):

pineappl_py/tests/test_evolution.py

@@ -5,15 +5,15 @@
 """
 
 import itertools
+from typing import List
 
 import numpy as np
 from pineappl.boc import Channel, Kinematics
 from pineappl.convolutions import Conv, ConvType
-from pineappl.evolution import OperatorSliceInfo, EvolveInfo
+from pineappl.evolution import EvolveInfo, OperatorSliceInfo
 from pineappl.grid import Grid, Order
 from pineappl.interpolation import Interp
 from pineappl.pids import PidBasis
-from typing import List
 
 
 class TestEvolution:

pineappl_py/tests/test_fk_table.py

@@ -4,16 +4,17 @@
 three (general) convolutions.
 """
 
-import numpy as np
 import tempfile
+from typing import List
+
+import numpy as np
 from pineappl.boc import Channel, Kinematics
 from pineappl.convolutions import Conv, ConvType
-from pineappl.fk_table import FkTable, FkAssumptions
+from pineappl.fk_table import FkAssumptions, FkTable
 from pineappl.grid import Grid, Order
 from pineappl.import_subgrid import ImportSubgridV1
 from pineappl.interpolation import Interp
 from pineappl.pids import PidBasis
-from typing import List
 
 
 class TestFkTable:
@@ -86,9 +87,9 @@ def test_convolve(self):
         xs = np.linspace(0.5, 1.0, 5)
         vs = xs.copy()
         subgrid = ImportSubgridV1(
-            vs[np.newaxis, :],  # DIS shape: (len(q2), len(x_grid))
-            np.array([90.0]),
-            xs,
+            array=vs[np.newaxis, :],  # DIS shape: (len(q2), len(x_grid))
+            scales=np.array([90.0]),
+            x_grids=[xs],  # Vector containing one single `x-grid`
         )
         g.set_subgrid(0, 0, 0, subgrid.into())
 

pineappl_py/tests/test_grid.py

@@ -1,15 +1,15 @@
-import numpy as np
-import pytest
 import tempfile
+from typing import List
 
+import numpy as np
+import pytest
 from pineappl.bin import BinRemapper
 from pineappl.boc import Channel, Kinematics
 from pineappl.convolutions import Conv, ConvType
 from pineappl.grid import Grid, Order
 from pineappl.import_subgrid import ImportSubgridV1
 from pineappl.interpolation import Interp
 from pineappl.pids import PidBasis
-from typing import List
 
 # Construct the type of convolutions and the convolution object
 # We assume unpolarized protons in the initial state
@@ -167,10 +167,9 @@ def test_set_subgrid(self):
         xs = np.linspace(0.1, 1.0, 5)
         vs = np.random.rand(len(xs))
         subgrid = ImportSubgridV1(
-            vs[np.newaxis, :, np.newaxis],
-            np.array([90.0]),
-            xs,
-            np.array([1.0]),
+            array=vs[np.newaxis, :, np.newaxis],
+            scales=np.array([90.0]),
+            x_grids=[xs, np.array([1.0])],
         )
         g.set_subgrid(0, 0, 0, subgrid.into())
 
@@ -179,7 +178,9 @@ def test_set_subgrid(self):
         x2s = np.linspace(0.5, 1, 2)
         Q2s = np.linspace(10, 20, 2)
         subgrid = ImportSubgridV1(
-            np.random.rand(len(Q2s), len(x1s), len(x2s)), Q2s, x1s, x2s
+            array=np.random.rand(len(Q2s), len(x1s), len(x2s)),
+            scales=Q2s,
+            x_grids=[x1s, x2s],
         )
         g.set_subgrid(0, 1, 0, subgrid.into())
         g.optimize()
@@ -232,8 +233,7 @@ def test_toy_convolution(self, params, expected):
         subgrid = ImportSubgridV1(
             array=vs[np.newaxis, :, np.newaxis],
             scales=np.array([90.0]),
-            x1_grid=xs,
-            x2_grid=np.array([1.0]),
+            x_grids=[xs, np.array([1.0])],
         )
         g.set_subgrid(0, 0, 0, subgrid.into())
 

pineappl_py/tests/test_subgrid.py

@@ -6,7 +6,7 @@
 from pineappl.import_subgrid import ImportSubgridV1
 from pineappl.interpolation import Interp
 from pineappl.pids import PidBasis
-from pineappl.subgrid import SubgridEnum, Mu2
+from pineappl.subgrid import Mu2, SubgridEnum
 
 # Define some default for the minimum value of `Q2`
 Q2_MIN = 1e2
@@ -95,19 +95,18 @@ def fake_grid(self):
         orders = [Order(0, 0, 0, 0, 0)]
         return fake_dis_grid(orders, channels)
 
-    def fake_importonlysubgrid(self) -> tuple:
-        x1s = np.linspace(0.1, 1, 2)
-        x2s = np.linspace(0.5, 1, 2)
+    def fake_importonlysubgrid(self, nb_dim: int = 2) -> tuple:
+        x_grids = [np.linspace(0.1, 1, 2) for _ in range(nb_dim)]
+        xgrid_size = [x.size for x in x_grids]
         Q2s = np.linspace(10, 20, 2)
         scale = [q2 for q2 in Q2s]
-        array = np.random.rand(len(Q2s), len(x1s), len(x2s))
+        array = np.random.rand(len(Q2s), *xgrid_size)
         subgrid = ImportSubgridV1(
             array=array,
             scales=scale,
-            x1_grid=x1s,
-            x2_grid=x2s,
+            x_grids=x_grids,
         )
-        return subgrid, [x1s, x2s, scale, array]
+        return subgrid, [*x_grids, scale, array]
 
     def test_mu2(self):
         mu2_obj = Mu2(ren=1.0, fac=2.0, frg=0.0)
@@ -132,6 +131,14 @@ def test_subgrid_methods(self):
         extr_subgrid = grid.subgrid(0, 0, 0)
         assert isinstance(extr_subgrid, SubgridEnum)
 
+    @pytest.mark.parametrize("nb_dim", [1, 2, 3, 4])
+    def test_subgrid_arrays(self, nb_dim):
+        """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)
+        assert isinstance(subgrid, ImportSubgridV1)
+
     @pytest.mark.skip(reason="No implementation of Array3 for subgrid.")
     def test_to_array3(self):
         # TODO: extract and check the dense array of the subgrid
@@ -142,6 +149,4 @@ def test_to_array3(self):
         grid.set_subgrid(0, 0, 0, test_subgrid.into())
         extr_subgrid = grid.subgrid(0, 0, 0)
         test_array = extr_subgrid.to_array3()
-        print(test_array)
-        print(array)
         np.testing.assert_allclose(test_array, array)