diff --git a/pineappl_py/src/import_subgrid.rs b/pineappl_py/src/import_subgrid.rs index d3721a81..fd14dc43 100644 --- a/pineappl_py/src/import_subgrid.rs +++ b/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, - x1_grid: Vec, - x2_grid: Option>, - py: Python<'py>, - ) -> Self { - // let node_values: Vec> = vec![scales, x1_grid, x2_grid]; - // let mut sparse_array: PackedArray = - // PackedArray::new(node_values.iter().map(Vec::len).collect()); + pub fn new(array: PyReadonlyArrayDyn, scales: Vec, x_grids: Vec>) -> Self { + // Total number of nodes = (scales + x-grids) + let mut node_values: Vec> = 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![scales, x1_grid]; - - if let Some(x2) = x2_grid { - node_values.push(x2); - } let mut sparse_array: PackedArray = PackedArray::new(node_values.iter().map(Vec::len).collect()); - if sparse_array.shape().to_vec().len() == 3 { - let array_3d: PyReadonlyArray3 = 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 = 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 { diff --git a/pineappl_py/tests/conftest.py b/pineappl_py/tests/conftest.py index 3561254d..2819edaa 100644 --- a/pineappl_py/tests/conftest.py +++ b/pineappl_py/tests/conftest.py @@ -1,6 +1,7 @@ -import pytest import subprocess +import pytest + class PDF: def xfxQ2(self, pid, x, q2): diff --git a/pineappl_py/tests/test_evolution.py b/pineappl_py/tests/test_evolution.py index bc25abec..5f9eb9ac 100644 --- a/pineappl_py/tests/test_evolution.py +++ b/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: diff --git a/pineappl_py/tests/test_fk_table.py b/pineappl_py/tests/test_fk_table.py index ac61e166..45259318 100644 --- a/pineappl_py/tests/test_fk_table.py +++ b/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()) diff --git a/pineappl_py/tests/test_grid.py b/pineappl_py/tests/test_grid.py index 10373eab..6f7bfbbb 100644 --- a/pineappl_py/tests/test_grid.py +++ b/pineappl_py/tests/test_grid.py @@ -1,7 +1,8 @@ -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 @@ -9,7 +10,6 @@ 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()) diff --git a/pineappl_py/tests/test_subgrid.py b/pineappl_py/tests/test_subgrid.py index 1b1700fe..bf2b37e3 100644 --- a/pineappl_py/tests/test_subgrid.py +++ b/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)