Skip to content

Commit

Permalink
Make convolutions member of Grid
Browse files Browse the repository at this point in the history
  • Loading branch information
t7phy committed Aug 6, 2024
1 parent c4f871e commit b4d665d
Show file tree
Hide file tree
Showing 3 changed files with 95 additions and 103 deletions.
3 changes: 2 additions & 1 deletion pineappl/src/convolutions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use super::grid::Grid;
use super::pids;
use super::subgrid::{Mu2, Subgrid};
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};

enum Pdfs<'a> {
Two {
Expand Down Expand Up @@ -343,7 +344,7 @@ impl<'a> LumiCache<'a> {
}

/// Data type that indentifies different types of convolutions.
#[derive(Debug, Eq, PartialEq)]
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum Convolution {
// TODO: eventually get rid of this value
/// No convolution.
Expand Down
193 changes: 92 additions & 101 deletions pineappl/src/grid.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,16 @@ use super::fk_table::FkTable;
use super::import_only_subgrid::ImportOnlySubgridV2;
use super::lagrange_subgrid::{LagrangeSparseSubgridV1, LagrangeSubgridV1, LagrangeSubgridV2};
use super::ntuple_subgrid::NtupleSubgridV1;
use super::packed_array::PackedArray;
use super::packed_subgrid::PackedQ1X2SubgridV1;
use super::pids::{self, PidBasis};
use super::subgrid::{ExtraSubgridParams, Mu2, Subgrid, SubgridEnum, SubgridParams};
use bitflags::bitflags;
use float_cmp::approx_eq;
use git_version::git_version;
use lz4_flex::frame::{FrameDecoder, FrameEncoder};
use ndarray::{s, Array3, ArrayView3, ArrayView5, ArrayViewMut3, Axis, CowArray, Dimension, Ix4};
use pineappl_v0::grid::Grid as GridV0;
use serde::{Deserialize, Serialize};
use std::borrow::Cow;
use std::collections::BTreeMap;
Expand Down Expand Up @@ -100,21 +103,15 @@ impl Mmv3 {
}

fn default_metadata() -> BTreeMap<String, String> {
[
(
"pineappl_gitversion".to_owned(),
git_version!(
args = ["--always", "--dirty", "--long", "--tags"],
cargo_prefix = "cargo:",
fallback = "unknown"
)
.to_owned(),
),
// by default we assume there are unpolarized protons in the initial state
// do not change these to the new metadata to not break backwards compatibility
("initial_state_1".to_owned(), "2212".to_owned()),
("initial_state_2".to_owned(), "2212".to_owned()),
]
[(
"pineappl_gitversion".to_owned(),
git_version!(
args = ["--always", "--dirty", "--long", "--tags"],
cargo_prefix = "cargo:",
fallback = "unknown"
)
.to_owned(),
)]
.iter()
.cloned()
.collect()
Expand Down Expand Up @@ -169,6 +166,7 @@ pub struct Grid {
orders: Vec<Order>,
subgrid_params: SubgridParams,
metadata: BTreeMap<String, String>,
convolutions: Vec<Convolution>,
more_members: MoreMembers,
}

Expand All @@ -181,19 +179,22 @@ impl Grid {
bin_limits: Vec<f64>,
subgrid_params: SubgridParams,
) -> Self {
// TODO: check that channels has same number of PIDs everywhere
Self {
subgrids: Array3::from_shape_simple_fn(
(orders.len(), bin_limits.len() - 1, channels.len()),
|| EmptySubgridV1.into(),
),
orders,
channels,
bin_limits: BinLimits::new(bin_limits),
metadata: default_metadata(),
more_members: MoreMembers::V3(Mmv3::new(
LagrangeSubgridV2::new(&subgrid_params, &ExtraSubgridParams::from(&subgrid_params))
.into(),
)),
// TODO: add this as new parameter
convolutions: vec![Convolution::UnpolPDF(2212); channels[0].entry()[0].0.len()],
channels,
subgrid_params,
}
}
Expand Down Expand Up @@ -231,10 +232,11 @@ impl Grid {
|| EmptySubgridV1.into(),
),
orders,
channels,
bin_limits: BinLimits::new(bin_limits),
subgrid_params,
metadata: default_metadata(),
convolutions: vec![Convolution::UnpolPDF(2212); channels[0].entry()[0].0.len()],
channels,
more_members: MoreMembers::V3(Mmv3::new(subgrid_template)),
})
}
Expand Down Expand Up @@ -491,9 +493,6 @@ impl Grid {
}

fn read_uncompressed_v0(mut reader: impl BufRead) -> Result<Self, GridError> {
use super::packed_array::PackedArray;
use super::packed_subgrid::PackedQ1X2SubgridV1;
use pineappl_v0::grid::Grid as GridV0;
use pineappl_v0::subgrid::Subgrid as _;

// TODO: convert error from v0 to v1
Expand Down Expand Up @@ -568,6 +567,7 @@ impl Grid {
.unwrap_or_default()
.into_iter()
.collect(),
convolutions: Self::read_convolutions_from_metadata(&grid),
// TODO: make these proper members
more_members: MoreMembers::V3(Mmv3 {
remapper: grid.remapper().map(|r| {
Expand All @@ -585,6 +585,69 @@ impl Grid {
Ok(result)
}

fn read_convolutions_from_metadata(grid: &GridV0) -> Vec<Convolution> {
grid.key_values().map_or_else(
// if there isn't any metadata, we assume two unpolarized proton-PDFs are used
|| vec![Convolution::UnpolPDF(2212), Convolution::UnpolPDF(2212)],
|kv| {
// file format v0 only supports exactly two convolutions
(1..=2)
.map(|index| {
// if there are key-value pairs `convolution_particle_1` and
// `convolution_type_1` and the same with a higher index, we convert this
// metadata into `Convolution`
match (
kv.get(&format!("convolution_particle_{index}"))
.map(|s| s.parse::<i32>()),
kv.get(&format!("convolution_type_{index}"))
.map(String::as_str),
) {
(_, Some("None")) => Convolution::None,
(Some(Ok(pid)), Some("UnpolPDF")) => Convolution::UnpolPDF(pid),
(Some(Ok(pid)), Some("PolPDF")) => Convolution::PolPDF(pid),
(Some(Ok(pid)), Some("UnpolFF")) => Convolution::UnpolFF(pid),
(Some(Ok(pid)), Some("PolFF")) => Convolution::PolFF(pid),
(None, None) => {
// if these key-value pairs are missing use the old metadata
match kv
.get(&format!("initial_state_{index}"))
.map(|s| s.parse::<i32>())
{
Some(Ok(pid)) => {
let condition = !grid.channels().iter().all(|entry| {
entry.entry().iter().all(|&(a, b, _)| match index {1 => a, 2 => b, _ => unreachable!()} == pid)
});

if condition {
Convolution::UnpolPDF(pid)
} else {
Convolution::None
}
}
None => Convolution::UnpolPDF(2212),
Some(Err(err)) => panic!(
"metadata 'initial_state_{index}' could not be parsed: {err}"
),
}
}
(None, Some(_)) => {
panic!("metadata 'convolution_type_{index}' is missing")
}
(Some(_), None) => {
panic!("metadata 'convolution_particle_{index}' is missing")
}
(Some(Ok(_)), Some(type_)) => {
panic!("metadata 'convolution_type_{index} = {type_}' is unknown")
}
(Some(Err(err)), Some(_)) => {
panic!("metadata 'convolution_particle_{index}' could not be parsed: {err}")
}
}
})
.collect()
},
)
}
/// Serializes `self` into `writer`. Writing is buffered.
///
/// # Errors
Expand Down Expand Up @@ -818,84 +881,13 @@ impl Grid {
/// Panics if the metadata key--value pairs `convolution_particle_1` and `convolution_type_1`,
/// or `convolution_particle_2` and `convolution_type_2` are not correctly set.
#[must_use]
pub fn convolutions(&self) -> Vec<Convolution> {
let kv = self.metadata();

// the current file format only supports exactly two convolutions
(1..=2)
.map(|index| {
// if there are key-value pairs `convolution_particle_1` and
// `convolution_type_1` and the same with a higher index, we convert this
// metadata into `Convolution`
match (
kv.get(&format!("convolution_particle_{index}"))
.map(|s| s.parse::<i32>()),
kv.get(&format!("convolution_type_{index}"))
.map(String::as_str),
) {
(_, Some("None")) => Convolution::None,
(Some(Ok(pid)), Some("UnpolPDF")) => Convolution::UnpolPDF(pid),
(Some(Ok(pid)), Some("PolPDF")) => Convolution::PolPDF(pid),
(Some(Ok(pid)), Some("UnpolFF")) => Convolution::UnpolFF(pid),
(Some(Ok(pid)), Some("PolFF")) => Convolution::PolFF(pid),
(None, None) => {
// if these key-value pairs are missing use the old metadata
match kv
.get(&format!("initial_state_{index}"))
.map(|s| s.parse::<i32>())
{
Some(Ok(pid)) => {
let condition = !self.channels().iter().all(|entry| {
entry.entry().iter().all(|(pids, _)| pids[index - 1] == pid)
});

if condition {
Convolution::UnpolPDF(pid)
} else {
Convolution::None
}
}
None => Convolution::UnpolPDF(2212),
Some(Err(err)) => panic!(
"metadata 'initial_state_{index}' could not be parsed: {err}"
),
}
}
(None, Some(_)) => {
panic!("metadata 'convolution_type_{index}' is missing")
}
(Some(_), None) => {
panic!("metadata 'convolution_particle_{index}' is missing")
}
(Some(Ok(_)), Some(type_)) => {
panic!("metadata 'convolution_type_{index} = {type_}' is unknown")
}
(Some(Err(err)), Some(_)) => {
panic!("metadata 'convolution_particle_{index}' could not be parsed: {err}")
}
}
})
.collect()
pub fn convolutions(&self) -> &[Convolution] {
&self.convolutions
}

/// Set the convolution type for this grid for the corresponding `index`.
pub fn set_convolution(&mut self, index: usize, convolution: Convolution) {
// remove outdated metadata
self.metadata_mut()
.remove(&format!("initial_state_{}", index + 1));

let (type_, particle) = match convolution {
Convolution::UnpolPDF(pid) => ("UnpolPDF".to_owned(), pid.to_string()),
Convolution::PolPDF(pid) => ("PolPDF".to_owned(), pid.to_string()),
Convolution::UnpolFF(pid) => ("UnpolFF".to_owned(), pid.to_string()),
Convolution::PolFF(pid) => ("PolFF".to_owned(), pid.to_string()),
Convolution::None => ("None".to_owned(), String::new()),
};

self.metadata_mut()
.insert(format!("convolution_type_{}", index + 1), type_);
self.metadata_mut()
.insert(format!("convolution_particle_{}", index + 1), particle);
/// Return the convolution types.
pub fn convolutions_mut(&mut self) -> &mut [Convolution] {
&mut self.convolutions
}

fn increase_shape(&mut self, new_dim: &(usize, usize, usize)) {
Expand Down Expand Up @@ -1477,6 +1469,7 @@ impl Grid {
orders: vec![Order::new(0, 0, 0, 0, 0)],
subgrid_params: SubgridParams::default(),
metadata: self.metadata.clone(),
convolutions: self.convolutions.clone(),
more_members: self.more_members.clone(),
};

Expand Down Expand Up @@ -1515,9 +1508,6 @@ impl Grid {
// TODO: convert this unwrap into error
let grid = lhs.unwrap();

// check that the convolutions are unchanged
assert_eq!(self.convolutions(), grid.convolutions());

// UNWRAP: merging evolved slices should be a proper FkTable again
Ok(FkTable::try_from(grid).unwrap_or_else(|_| unreachable!()))
}
Expand Down Expand Up @@ -1617,6 +1607,7 @@ impl Grid {
orders: vec![Order::new(0, 0, 0, 0, 0)],
subgrid_params: SubgridParams::default(),
metadata: self.metadata.clone(),
convolutions: self.convolutions.clone(),
more_members: self.more_members.clone(),
};

Expand Down Expand Up @@ -2087,8 +2078,8 @@ mod tests {
[Convolution::UnpolPDF(2212), Convolution::UnpolPDF(2212)]
);

grid.set_convolution(0, Convolution::UnpolPDF(-2212));
grid.set_convolution(1, Convolution::UnpolPDF(-2212));
grid.convolutions_mut()[0] = Convolution::UnpolPDF(-2212);
grid.convolutions_mut()[1] = Convolution::UnpolPDF(-2212);

assert_eq!(
grid.convolutions(),
Expand Down
2 changes: 1 addition & 1 deletion pineappl_cli/src/write.rs
Original file line number Diff line number Diff line change
Expand Up @@ -519,7 +519,7 @@ impl Subcommand for Opts {
);
}

grid.set_convolution(index, grid.convolutions()[index].charge_conjugate());
grid.convolutions_mut()[index] = grid.convolutions()[index].charge_conjugate();
}
OpsArg::DedupChannels(ulps) => {
grid.dedup_channels(*ulps);
Expand Down

0 comments on commit b4d665d

Please sign in to comment.