diff --git a/pineappl/src/grid.rs b/pineappl/src/grid.rs index 6e49e442..121b6f77 100644 --- a/pineappl/src/grid.rs +++ b/pineappl/src/grid.rs @@ -206,15 +206,39 @@ impl MoreMembers { } } -/// Information required to compute a compatible EKO. -/// Members spell out specific characteristic of a suitable EKO. -pub struct EkoInfo { - /// is the interpolation grid in x used at the process scale +/// Information required to calculate the evolution kernel operators (EKO) to perform a conversion +/// of a [`Grid`] using [`Grid::convolute_eko`] to an [`FkTable`]. +pub struct GridAxes { + /// Interpolation grid in x of the `Grid`. pub x_grid: Vec, - /// is the intepolation grid in q2, spanning the q2 range covered by the process data + /// Parton IDs used in the grid. + pub pids: Vec, + /// Interpolation grid for the factorization scale of the `Grid`. pub muf2_grid: Vec, } +/// Extra information required to perform the conversion of a [`Grid`] to an [`FkTable`] using +/// [`Grid::convolute_eko`]. +pub struct EkoInfo { + /// Scale of the FkTable. + pub muf2_0: f64, + /// Strong coupling constants for the factorization scales in the same ordering as given in + /// [`GridAxes`]. + pub alphas: Vec, + /// Renormalization scale variation. + pub xir: f64, + /// Factorization scale variation. + pub xif: f64, + /// Interpolation grid in x of the `FkTable`. + pub target_x_grid: Vec, + /// Parton IDs for the `FkTable`. + pub target_pids: Vec, + /// axes shared with the process grid + pub grid_axes: GridAxes, + /// TODO: replace this member with the actual data + pub lumi_id_types: String, +} + /// Main data structure of `PineAPPL`. This structure contains a `Subgrid` for each `LumiEntry`, /// bin, and coupling order it was created with. #[derive(Deserialize, Serialize)] @@ -1066,9 +1090,10 @@ impl Grid { /// Provide information used to compute a suitable EKO for the current grid. /// More specific, the `x_grid` and `muf2_grid` are extracted and checked. #[must_use] - pub fn eko_info(&self) -> Option { + pub fn axes(&self) -> Option { let mut muf2_grid = Vec::::new(); let mut x_grid = Vec::::new(); + let pids = Vec::::new(); let mut has_pdf1 = true; let mut has_pdf2 = true; @@ -1123,7 +1148,11 @@ impl Grid { } } - Some(EkoInfo { x_grid, muf2_grid }) + Some(GridAxes { + x_grid, + pids, // TODO: for the time being they are just empty, but we might use them for slicing the eko + muf2_grid, + }) } /// Applies an evolution kernel operator (EKO) to the grids to evolve them from different @@ -1135,24 +1164,15 @@ impl Grid { /// /// Panics if the parameters do not match with the given grid. #[must_use] - pub fn convolute_eko( - &self, - q2: f64, - alphas: &[f64], - (xir, xif): (f64, f64), - pids: &[i32], - x_grid: Vec, - q2_grid: Vec, - operator: Array5, - ) -> Option { + pub fn convolute_eko(&self, operator: Array5, eko_info: EkoInfo) -> Option { // Check operator layout let dim = operator.shape(); - assert_eq!(dim[0], q2_grid.len()); - assert_eq!(dim[1], pids.len()); - assert_eq!(dim[2], x_grid.len()); - assert_eq!(dim[3], pids.len()); - assert_eq!(dim[4], x_grid.len()); + assert_eq!(dim[0], eko_info.grid_axes.muf2_grid.len()); + assert_eq!(dim[1], eko_info.target_pids.len()); + assert_eq!(dim[2], eko_info.target_x_grid.len()); + assert_eq!(dim[3], eko_info.grid_axes.pids.len()); + assert_eq!(dim[4], eko_info.grid_axes.x_grid.len()); // swap axes around to optimize convolution let operator = operator.permuted_axes([3, 1, 4, 0, 2]); @@ -1178,27 +1198,45 @@ impl Grid { _ => unimplemented!(), }; - // create target luminosities let pids1 = if has_pdf1 { - pids.to_vec() + eko_info.grid_axes.pids.to_vec() } else { vec![initial_state_1] }; let pids2 = if has_pdf2 { - pids.to_vec() + eko_info.grid_axes.pids.to_vec() + } else { + vec![initial_state_2] + }; + // create target luminosities + let tgt_pids1 = if has_pdf1 { + eko_info.target_pids.to_vec() + } else { + vec![initial_state_1] + }; + let tgt_pids2 = if has_pdf2 { + eko_info.target_pids.to_vec() } else { vec![initial_state_2] }; - let lumi: Vec<_> = pids1 + let lumi: Vec<_> = tgt_pids1 .iter() - .cartesian_product(pids2.iter()) + .cartesian_product(tgt_pids2.iter()) .map(|(a, b)| lumi_entry![*a, *b, 1.0]) .collect(); // create target subgrid dimensions - let tgt_q2_grid = vec![q2]; - let tgt_x1_grid = if has_pdf1 { x_grid.clone() } else { vec![1.0] }; - let tgt_x2_grid = if has_pdf2 { x_grid.clone() } else { vec![1.0] }; + let tgt_q2_grid = vec![eko_info.muf2_0]; + let tgt_x1_grid = if has_pdf1 { + eko_info.target_x_grid.clone() + } else { + vec![1.0] + }; + let tgt_x2_grid = if has_pdf2 { + eko_info.target_x_grid.clone() + } else { + vec![1.0] + }; // create target grid let mut result = Self { @@ -1216,9 +1254,11 @@ impl Grid { subgrid_params: SubgridParams::default(), more_members: self.more_members.clone(), }; + // write additional metadata + result.set_key_value("lumi_id_types", &eko_info.lumi_id_types); // collect source grid informations - let eko_info = self.eko_info().unwrap(); + let grid_axes = self.axes().unwrap(); // Setup progress bar let bar = ProgressBar::new( @@ -1248,22 +1288,26 @@ impl Grid { let mut src_array = SparseArray3::::new( src_array_q2_grid.len(), - if has_pdf1 { eko_info.x_grid.len() } else { 1 }, - if has_pdf2 { eko_info.x_grid.len() } else { 1 }, + if has_pdf1 { grid_axes.x_grid.len() } else { 1 }, + if has_pdf2 { grid_axes.x_grid.len() } else { 1 }, ); // iterate over the source grid orders and add all of them together into // `src_array`, using the right powers of alphas for (order, powers) in self.orders.iter().enumerate() { - let logs = if (xir, xif) == (1.0, 1.0) { + let logs = if (eko_info.xir, eko_info.xif) == (1.0, 1.0) { if (powers.logxir > 0) || (powers.logxif > 0) { continue; } 1.0 } else { - (xir * xir).ln().powi(powers.logxir.try_into().unwrap()) - * (xif * xif).ln().powi(powers.logxif.try_into().unwrap()) + (eko_info.xir * eko_info.xir) + .ln() + .powi(powers.logxir.try_into().unwrap()) + * (eko_info.xif * eko_info.xif) + .ln() + .powi(powers.logxif.try_into().unwrap()) }; let src_subgrid = &self.subgrids[[order, bin, src_lumi]]; @@ -1273,13 +1317,13 @@ impl Grid { && !src_subgrid .x1_grid() .iter() - .zip(x_grid.iter()) + .zip(eko_info.grid_axes.x_grid.iter()) .all(|(a, b)| approx_eq!(f64, *a, *b, ulps = 128))) || (has_pdf2 && !src_subgrid .x2_grid() .iter() - .zip(x_grid.iter()) + .zip(eko_info.grid_axes.x_grid.iter()) .all(|(a, b)| approx_eq!(f64, *a, *b, ulps = 128))); for ((iq2, ix1, ix2), &value) in src_subgrid.iter() { @@ -1288,24 +1332,31 @@ impl Grid { .iter() .position(|&q2| q2 == scale) .unwrap(); - let als_iq2 = q2_grid + let als_iq2 = eko_info + .grid_axes + .muf2_grid .iter() - .position(|&q2| q2 == xir * xir * scale) - .unwrap(); + .position(|&q2| q2 == eko_info.xir * eko_info.xir * scale) + .expect(&format!( + "Couldn't find q2: {:?} with xir: {:?} and muf2_grid: {:?}", + scale, eko_info.xir, eko_info.grid_axes.muf2_grid + )); let ix1 = if invert_x && has_pdf1 { - eko_info.x_grid.len() - ix1 - 1 + eko_info.grid_axes.x_grid.len() - ix1 - 1 } else { ix1 }; let ix2 = if invert_x && has_pdf2 { - eko_info.x_grid.len() - ix2 - 1 + eko_info.grid_axes.x_grid.len() - ix2 - 1 } else { ix2 }; - src_array[[src_iq2, ix1, ix2]] += - alphas[als_iq2].powi(powers.alphas.try_into().unwrap()) * logs * value; + src_array[[src_iq2, ix1, ix2]] += eko_info.alphas[als_iq2] + .powi(powers.alphas.try_into().unwrap()) + * logs + * value; } } @@ -1320,7 +1371,14 @@ impl Grid { // Next we need to apply the tensor let eko_src_q2_indices: Vec<_> = src_array_q2_grid .iter() - .map(|&src_q2| q2_grid.iter().position(|&q2| q2 == src_q2).unwrap()) + .map(|&src_q2| { + eko_info + .grid_axes + .muf2_grid + .iter() + .position(|&q2| q2 == src_q2) + .unwrap() + }) .collect(); // Iterate target lumis for (tgt_lumi, (tgt_pid1_idx, tgt_pid2_idx)) in (0..pids1.len()) @@ -1330,7 +1388,10 @@ impl Grid { for (src_pid1, src_pid2, factor) in src_entries.entry().iter() { // find source lumi position let src_pid1_idx = if has_pdf1 { - pids.iter() + eko_info + .grid_axes + .pids + .iter() .position(|x| { // if `pid == 0` the gluon is meant if *src_pid1 == 0 { @@ -1344,7 +1405,10 @@ impl Grid { 0 }; let src_pid2_idx = if has_pdf2 { - pids.iter() + eko_info + .grid_axes + .pids + .iter() .position(|x| { // `pid == 0` is the gluon exception, which might be 0 or 21 if *src_pid2 == 0 { @@ -1443,6 +1507,7 @@ impl Grid { bar.finish(); + result.optimize(); Some(FkTable::try_from(result).unwrap()) } } @@ -1715,4 +1780,98 @@ mod tests { "-2212" ); } + + // TODO: properly test axes returned + + fn simple_grid() -> (Grid, GridAxes) { + let muf2_grid = vec![20.]; + let x_grid = vec![0.1, 0.5, 1.]; + + let mut subgrid_params = SubgridParams::default(); + subgrid_params.set_x_order(1); + subgrid_params.set_x_bins(1); + subgrid_params.set_q2_order(1); + subgrid_params.set_q2_bins(1); + + let mut array = Array3::zeros((1, 3, 3)); + array[[0, 0, 0]] = 1.; + let sparse_array = SparseArray3::from_ndarray(&array, 0, 3); + let subgrid = ImportOnlySubgridV1::new( + sparse_array, + muf2_grid.clone(), + x_grid.clone(), + x_grid.clone(), + ) + .into(); + + let pids = vec![21, 1, 2]; + let mut grid = Grid::new( + vec![lumi_entry![21, 21, 1.0], lumi_entry![1, 2, 1.0]], + vec![Order { + alphas: 0, + alpha: 0, + logxir: 0, + logxif: 0, + }], + vec![0.0, 1.0], + subgrid_params, + ); + + grid.set_subgrid(0, 0, 0, subgrid); + + ( + grid, + GridAxes { + x_grid, + pids, + muf2_grid, + }, + ) + } + + #[test] + fn grid_axes() { + let (grid, axes) = simple_grid(); + + let ret_axes = grid.axes().unwrap(); + assert_eq!(ret_axes.x_grid, axes.x_grid); + assert_eq!(ret_axes.muf2_grid, axes.muf2_grid); + assert_eq!(ret_axes.pids, vec![]); + } + + #[test] + fn grid_convolute_eko() { + let (grid, axes) = simple_grid(); + let target_x_grid = vec![1e-7, 1e-2, 1.]; + let target_pids = vec![21, 1, 2]; + + let lumi_id_types = "pdg_mc_ids".to_owned(); + + let eko_info = EkoInfo { + muf2_0: 1., + alphas: vec![1.], + xir: 1., + xif: 1., + target_x_grid: target_x_grid.clone(), + target_pids: target_pids.clone(), + grid_axes: GridAxes { + x_grid: axes.x_grid.clone(), + pids: axes.pids.clone(), + muf2_grid: axes.muf2_grid.clone(), + }, + lumi_id_types, + }; + let operator = ndarray::Array::from_shape_vec( + (1, 3, 3, 3, 3), + (0..4) + .map(|_| (0..3)) + .multi_cartesian_product() + .map(|v| if v[0] == v[2] && v[1] == v[3] { 1. } else { 0. }) + .collect(), + ) + .unwrap(); + let fk = grid.convolute_eko(operator, eko_info).unwrap(); + + assert_eq!(fk.bins(), 1); + } } diff --git a/pineappl_py/pineappl/grid.py b/pineappl_py/pineappl/grid.py index b9756952..0a6158c0 100644 --- a/pineappl_py/pineappl/grid.py +++ b/pineappl_py/pineappl/grid.py @@ -155,7 +155,7 @@ def convolute( xfx1, xfx2, alphas, order_mask, bin_indices, lumi_mask, xi ) - def convolute_eko(self, operators): + def convolute_eko(self, operators, lumi_id_types="pdg_mc_ids"): """ Create an FKTable with the EKO. @@ -165,6 +165,9 @@ def convolute_eko(self, operators): ---------- operators : dict EKO Output + lumi_id_types : str + kind of lumi types (e.g. "pdg_mc_ids" for flavor basis, "evol" + for evolution basis) Returns ------ @@ -181,10 +184,13 @@ def convolute_eko(self, operators): operators["q2_ref"], alphas_values, operators["targetpids"], - operators["interpolation_xgrid"], + operators["targetgrid"], + operators["inputpids"], + operators["inputgrid"], q2grid, operator_grid.flatten(), operator_grid.shape, + lumi_id_types, ) ) diff --git a/pineappl_py/src/grid.rs b/pineappl_py/src/grid.rs index 9839c544..417d63c7 100644 --- a/pineappl_py/src/grid.rs +++ b/pineappl_py/src/grid.rs @@ -1,4 +1,4 @@ -use pineappl::grid::{EkoInfo, Grid, Ntuple, Order}; +use pineappl::grid::{EkoInfo, Grid, GridAxes, Ntuple, Order}; use super::bin::PyBinRemapper; use super::fk_table::PyFkTable; @@ -166,11 +166,17 @@ impl PyGrid { /// ------- /// x_grid: list(float) /// interpolation grid + /// pids: list(int) + /// particle ids /// muf2_grid : list(float) /// factorization scale list - pub fn eko_info(&self) -> (Vec, Vec) { - let EkoInfo { x_grid, muf2_grid } = self.grid.eko_info().unwrap(); - (x_grid, muf2_grid) + pub fn axes(&self) -> (Vec, Vec, Vec) { + let GridAxes { + x_grid, + pids, + muf2_grid, + } = self.grid.axes().unwrap(); + (x_grid, pids, muf2_grid) } /// Convolute grid with pdf. @@ -234,7 +240,7 @@ impl PyGrid { /// /// Parameters /// ---------- - /// q2 : float + /// muf2_0 : float /// reference scale /// alphas : list(float) /// list with :math:`\alpha_s(Q2)` for the process scales @@ -242,12 +248,18 @@ impl PyGrid { /// sorting of the particles in the tensor /// x_grid : list(float) /// interpolation grid - /// q2_grid : list(float) + /// target_pids : list(int) + /// sorting of the particles in the tensor for final FkTable + /// target_x_grid : list(float) + /// final FKTable interpolation grid + /// muf2_grid : list(float) /// list of process scales /// operator_flattened : list(float) /// evolution tensor as a flat list /// operator_shape : list(int) /// shape of the evolution tensor + /// additional_metadata : dict(str: str) + /// further metadata /// /// Returns /// ------- @@ -255,26 +267,36 @@ impl PyGrid { /// produced FK table pub fn convolute_eko( &self, - q2: f64, + muf2_0: f64, alphas: Vec, pids: Vec, x_grid: Vec, - q2_grid: Vec, + target_pids: Vec, + target_x_grid: Vec, + muf2_grid: Vec, operator_flattened: Vec, operator_shape: Vec, + lumi_id_types: String, ) -> PyFkTable { let operator = Array::from_shape_vec(operator_shape, operator_flattened).unwrap(); + let eko_info = EkoInfo { + muf2_0, + alphas, + xir: 1., + xif: 1., + target_pids, + target_x_grid, + grid_axes: GridAxes { + x_grid, + pids, + muf2_grid, + }, + lumi_id_types, + }; + let evolved_grid = self .grid - .convolute_eko( - q2, - &alphas, - (1., 1.), - &pids, - x_grid, - q2_grid, - operator.into_dimensionality::().unwrap(), - ) + .convolute_eko(operator.into_dimensionality::().unwrap(), eko_info) .unwrap(); PyFkTable { fk_table: evolved_grid, @@ -318,6 +340,11 @@ impl PyGrid { self.grid.optimize(); } + /// Optimize grid content. + // pub fn merge(&mut self, mut other: Self) { + // self.grid.merge(other.grid); + // } + /// Extract the number of dimensions for bins. /// /// **Usage:** `pineko` diff --git a/pineappl_py/tests/test_grid.py b/pineappl_py/tests/test_grid.py index 2c79e5ad..4c7cccb5 100644 --- a/pineappl_py/tests/test_grid.py +++ b/pineappl_py/tests/test_grid.py @@ -97,7 +97,7 @@ def test_convolute(self): np.testing.assert_allclose(g.convolute(lambda pid,x,q2: 1, lambda pid,x,q2: 1., lambda q2: 1.), [5e6/9999,0.]) np.testing.assert_allclose(g.convolute(lambda pid,x,q2: 1, lambda pid,x,q2: 1., lambda q2: 2.), [2**3 * 5e6/9999,0.]) - def test_eko_info(self): + def test_axes(self): g = self.fake_grid() # add 2 DIS grids @@ -119,10 +119,11 @@ def test_eko_info(self): ) g.set_subgrid(0, 1, 0, subgrid) # now get the thing - ei = g.eko_info() + ei = g.axes() np.testing.assert_allclose(ei[0], xs) - np.testing.assert_allclose(ei[1], [90., 100.]) + np.testing.assert_allclose(ei[1], []) + np.testing.assert_allclose(ei[2], [90., 100.]) def test_io(self, tmp_path): g = self.fake_grid() @@ -137,6 +138,9 @@ def test_convolute_eko(self): fake_eko = { "q2_ref": 1., "targetpids": [1], + "targetgrid": [.1,1.], + "inputpids": [1], + "inputgrid": [.1,1.], "interpolation_xgrid": [.1,1.], "Q2grid": { 90: { @@ -145,6 +149,6 @@ def test_convolute_eko(self): } } } - g.set_key_value("lumi_id_types", "PDG") + g.set_key_value("lumi_id_types", "pdg_mc_ids") fk = g.convolute_eko(fake_eko) assert isinstance(fk.raw, pineappl.pineappl.PyFkTable)