Skip to content

Commit

Permalink
Complete final missing PyInterface and make tests more Robust
Browse files Browse the repository at this point in the history
  • Loading branch information
@Radonirinaunimi
Radonirinaunimi committed Nov 3, 2024
1 parent 7d1e3f5 commit 84a4e97
Show file tree
Hide file tree
Showing 5 changed files with 310 additions and 135 deletions.
103 changes: 96 additions & 7 deletions pineappl_py/src/grid.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,12 @@ use super::pids::PyPidBasis;
use super::subgrid::PySubgridEnum;
use itertools::izip;
use ndarray::CowArray;
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray4};
use numpy::{IntoPyArray, PyArray1, PyArrayDyn, PyReadonlyArray4};
use pineappl::boc::Kinematics;
use pineappl::convolutions::ConvolutionCache;
use pineappl::evolution::AlphasTable;
use pineappl::grid::Grid;
use pineappl::pids::PidBasis;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use std::fs::File;
Expand Down Expand Up @@ -288,7 +290,7 @@ impl PyGrid {
})
.collect();

let mut lumi_cache = ConvolutionCache::new(
let mut convolution_cache = ConvolutionCache::new(
pdg_convs.into_iter().map(|pdg| pdg.conv.clone()).collect(),
xfx_funcs
.iter_mut()
Expand All @@ -299,7 +301,7 @@ impl PyGrid {

self.grid
.convolve(
&mut lumi_cache,
&mut convolution_cache,
&order_mask.unwrap_or_default(),
&bin_indices.unwrap_or_default(),
&channel_mask.unwrap_or_default(),
Expand All @@ -308,6 +310,58 @@ impl PyGrid {
.into_pyarray_bound(py)
}

/// Convolve a single subgrid `(order, bin, channel)` with the distributions.
///
/// # Panics
///
/// TODO
#[must_use]
#[pyo3(signature = (pdg_convs, xfxs, alphas, ord, bin, channel, xi = None))]
pub fn convolve_subgrid<'py>(
&self,
pdg_convs: Vec<PyRef<PyConv>>,
xfxs: Vec<PyObject>,
alphas: PyObject,
ord: usize,
bin: usize,
channel: usize,
xi: Option<(f64, f64, f64)>,
py: Python<'py>,
) -> Bound<'py, PyArrayDyn<f64>> {
let mut alphas = |q2: f64| {
let result: f64 = alphas.call1(py, (q2,)).unwrap().extract(py).unwrap();
result
};

let mut xfx_funcs: Vec<_> = xfxs
.iter()
.map(|xfx| {
move |id: i32, x: f64, q2: f64| {
xfx.call1(py, (id, x, q2)).unwrap().extract(py).unwrap()
}
})
.collect();

let mut convolution_cache = ConvolutionCache::new(
pdg_convs.into_iter().map(|pdg| pdg.conv.clone()).collect(),
xfx_funcs
.iter_mut()
.map(|fx| fx as &mut dyn FnMut(i32, f64, f64) -> f64)
.collect(),
&mut alphas,
);

self.grid
.convolve_subgrid(
&mut convolution_cache,
ord,
bin,
channel,
xi.unwrap_or((1.0, 1.0, 0.0)),
)
.into_pyarray_bound(py)
}

/// Collect information for convolution with an evolution operator.
///
/// # Panics
Expand Down Expand Up @@ -441,7 +495,42 @@ impl PyGrid {
self.grid.write_lz4(File::create(path).unwrap()).unwrap();
}

/// Optimize content.
/// Return the convention by which the channels' PIDS are encoded.
#[getter]
#[must_use]
pub const fn pid_basis(&self) -> PyPidBasis {
match self.grid.pid_basis() {
PidBasis::Pdg => PyPidBasis::Pdg,
PidBasis::Evol => PyPidBasis::Evol,
}
}

/// Return the convention by which the Kinematics are encoded.
/// TODO
#[getter]
#[must_use]
pub fn kinematics(&self) -> Vec<PyKinematics> {
self.grid
.kinematics()
.iter()
.map(|&kin| match kin {
Kinematics::X(v) => PyKinematics::X(v),
Kinematics::Scale(v) => PyKinematics::Scale(v),
})
.collect()
}

/// Return the convention by which the Scales are encoded.
/// TODO
#[getter]
#[must_use]
pub fn scales(&self) -> PyScales {
PyScales {
scales: self.grid.scales().clone(),
}
}

/// Optimize the contents of the Grid.
pub fn optimize(&mut self) {
self.grid.optimize();
}
Expand Down Expand Up @@ -564,9 +653,9 @@ impl PyGrid {
///
/// Returns
/// -------
/// list(list(tuple(float,float,int))) :
/// channels as tuples (pid, pid, factor) (multiple tuples can be associated to the same
/// contribution)
/// list(list(tuple(list[float],int))) :
/// channels as tuples (List of PIDs, factor) (multiple tuples can be associated
/// to the same contribution)
#[must_use]
pub fn channels(&self) -> Vec<Vec<(Vec<i32>, f64)>> {
self.grid
Expand Down
89 changes: 22 additions & 67 deletions pineappl_py/tests/conftest.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,12 +42,11 @@ def unpolarized_pdf(self, pid, x, q2):
class FakeGrid:
"""Class that mocks a PineAPPL grid. This should contain functions
that return all the possible number of convolutions.
TODO: combine the different convolutions
"""

def grid_with_one_convolution(
def grid_with_generic_convolution(
self,
nb_convolutions: int,
channels: List[Channel],
orders: List[Order],
convolutions: List[Conv],
Expand All @@ -59,6 +58,11 @@ def grid_with_one_convolution(
xmax: float = 1,
xnodes: int = 50,
) -> Grid:
"""A function to generate fake GRIDs that can take any number of convolutions.
Note that the `nb_convolutions` can be different from the number of convolution
types passed to `convolutions`. Indeed, if all the types of convolutions are
the same, then only one single element can be passed to `convolutions`.
"""
kinematics = [
Kinematics.Scale(0), # Scale
Kinematics.X(0), # momentum fraction x
Expand All @@ -84,72 +88,23 @@ def grid_with_one_convolution(
interpolation_meth=InterpolationMethod.Lagrange,
), # Interpolation on momentum fraction x
]
# Construct the `Scales` object
scale_funcs = Scales(
ren=ScaleFuncForm.Scale(0),
fac=ScaleFuncForm.Scale(0),
frg=ScaleFuncForm.NoScale(0),
)

return Grid(
pid_basis=PidBasis.Evol,
channels=channels,
orders=orders,
bin_limits=np.array(bins),
convolutions=convolutions,
interpolations=interpolations,
kinematics=kinematics,
scale_funcs=scale_funcs,
)
# Extend the Kinematics and Interpolations
if nb_convolutions > 1:
for i in range(1, nb_convolutions):
kinematics.append(Kinematics.X(i))
interpolations.append(
Interp(
min=xmin,
max=xmax,
nodes=xnodes,
order=3,
reweight_meth=ReweightingMethod.ApplGridX,
map=MappingMethod.ApplGridF2,
interpolation_meth=InterpolationMethod.Lagrange,
)
)

def grid_with_two_convolutions(
self,
channels: List[Channel],
orders: List[Order],
convolutions: List[Conv],
bins: List[float] = [1e-7, 1e-3, 1],
q2min: float = 1e2,
q2max: float = 1e8,
q2nodes: int = 40,
xmin: float = 2e-7,
xmax: float = 1,
xnodes: int = 50,
) -> Grid:
kinematics = [
Kinematics.Scale(0), # Scale
Kinematics.X(0), # x1 momentum fraction
Kinematics.X(1), # x2 momentum fraction
]
# Define the interpolation specs for each item of the Kinematics
interpolations = [
Interp(
min=q2min,
max=q2max,
nodes=q2nodes,
order=3,
reweight_meth=ReweightingMethod.NoReweight,
map=MappingMethod.ApplGridH0,
interpolation_meth=InterpolationMethod.Lagrange,
), # Interpolation on the Scale
Interp(
min=xmin,
max=xmax,
nodes=xnodes,
order=3,
reweight_meth=ReweightingMethod.ApplGridX,
map=MappingMethod.ApplGridF2,
interpolation_meth=InterpolationMethod.Lagrange,
), # Interpolation on x1 momentum fraction
Interp(
min=xmin,
max=xmax,
nodes=xnodes,
order=3,
reweight_meth=ReweightingMethod.ApplGridX,
map=MappingMethod.ApplGridF2,
interpolation_meth=InterpolationMethod.Lagrange,
), # Interpolation on x2 momentum fraction
]
# Construct the `Scales` object
scale_funcs = Scales(
ren=ScaleFuncForm.Scale(0),
Expand Down
7 changes: 5 additions & 2 deletions pineappl_py/tests/test_fk_table.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,11 @@ def test_convolve(self, fake_grids):
channels = [Channel(down_channel), Channel(up_channel)]
# Now we define the perturbative orders
orders = [Order(0, 0, 0, 0, 0)]
g = fake_grids.grid_with_one_convolution(
channels=channels, orders=orders, convolutions=convolutions
g = fake_grids.grid_with_generic_convolution(
nb_convolutions=1,
channels=channels,
orders=orders,
convolutions=convolutions,
)

# DIS grid
Expand Down
Loading

0 comments on commit 84a4e97

Please sign in to comment.