Implement the backend logic for PDF interpolation with LHAPDF and NeoPDF

Author: Radonirinaunimi

Cargo.lock

@@ -28,6 +28,7 @@ git-version = "0.3.5"
 itertools = "0.10.1"
 lhapdf = { package = "managed-lhapdf", version = "0.3.4" }
 lz4_flex = { optional = true, version = "0.9.2" }
+neopdf = "=0.3.0-alpha2"
 ndarray = "0.15.4"
 ndarray-npy = { default-features = false, features = ["npz"], optional = true, version = "0.8.1" }
 pineappl = { path = "../pineappl", version = "=1.3.0" }

pineappl_cli/Cargo.toml

@@ -58,10 +58,11 @@ pub struct Opts {
 impl Subcommand for Opts {
     fn run(&self, cfg: &GlobalConfiguration) -> Result<ExitCode> {
         let grid = helpers::read_grid(&self.input)?;
-        let mut conv_funs_0 = helpers::create_conv_funs(&self.conv_funs[0])?;
+        let mut conv_funs_0 =
+            helpers::create_conv_funs_with_backend(&self.conv_funs[0], cfg.pdf_backend)?;
         let bins: Vec<_> = self.bins.iter().cloned().flatten().collect();
 
-        let results = helpers::convolve_scales(
+        let results = helpers::convolve_scales_with_backend(
             &grid,
             &mut conv_funs_0,
             &self.conv_funs[0].conv_types,
@@ -91,8 +92,9 @@ impl Subcommand for Opts {
             .iter()
             .flat_map(|conv_funs| {
                 let conv_types = &conv_funs.conv_types;
-                let mut conv_funs = helpers::create_conv_funs(conv_funs).unwrap();
-                helpers::convolve(
+                let mut conv_funs =
+                    helpers::create_conv_funs_with_backend(conv_funs, cfg.pdf_backend).unwrap();
+                helpers::convolve_with_backend(
                     &grid,
                     &mut conv_funs,
                     conv_types,

pineappl_cli/src/convolve.rs

@@ -1,3 +1,4 @@
+use super::pdf_backend::{self, Backend, ForcePositive, PdfBackend, PdfSetBackend};
 use super::GlobalConfiguration;
 use anyhow::{anyhow, bail, Context, Error, Result};
 use itertools::Itertools;
@@ -64,6 +65,7 @@ impl FromStr for ConvFuns {
     }
 }
 
+/// Creates convolution functions using LHAPDF backend (legacy).
 pub fn create_conv_funs(funs: &ConvFuns) -> Result<Vec<Pdf>> {
     Ok(funs
         .lhapdf_names
@@ -82,6 +84,22 @@ pub fn create_conv_funs(funs: &ConvFuns) -> Result<Vec<Pdf>> {
         .collect::<Result<_, _>>()?)
 }
 
+/// Creates convolution functions using the specified backend.
+pub fn create_conv_funs_with_backend(
+    funs: &ConvFuns,
+    backend: Backend,
+) -> Result<Vec<Box<dyn PdfBackend>>> {
+    funs.lhapdf_names
+        .iter()
+        .zip(&funs.members)
+        .map(|(name, member)| {
+            let member = member.unwrap_or(0);
+            pdf_backend::create_pdf(name, member, backend)
+        })
+        .collect()
+}
+
+/// Creates convolution functions for a PDF set using LHAPDF backend (legacy).
 pub fn create_conv_funs_for_set(
     funs: &ConvFuns,
     index_of_set: usize,
@@ -115,6 +133,30 @@ pub fn create_conv_funs_for_set(
     Ok((set, conv_funs))
 }
 
+/// Creates convolution functions for a PDF set using the specified backend.
+///
+/// Returns a tuple of (PdfSetBackend, Vec<Vec<Box<dyn PdfBackend>>>).
+pub fn create_conv_funs_for_set_with_backend(
+    funs: &ConvFuns,
+    index_of_set: usize,
+    backend: Backend,
+) -> Result<(Box<dyn PdfSetBackend>, Vec<Vec<Box<dyn PdfBackend>>>)> {
+    let setname = &funs.lhapdf_names[index_of_set];
+    let set = pdf_backend::create_pdf_set(setname, backend)?;
+
+    let set_members = set.mk_pdfs()?;
+    let conv_funs = set_members
+        .into_iter()
+        .map(|member_pdf| {
+            let mut conv_funs = create_conv_funs_with_backend(funs, backend)?;
+            conv_funs[index_of_set] = member_pdf;
+            Ok::<_, Error>(conv_funs)
+        })
+        .collect::<Result<_, _>>()?;
+
+    Ok((set, conv_funs))
+}
+
 pub fn read_grid(input: &Path) -> Result<Grid> {
     Grid::read(File::open(input).context(format!("unable to open '{}'", input.display()))?)
         .context(format!("unable to read '{}'", input.display()))
@@ -238,6 +280,7 @@ pub enum ConvoluteMode {
     Normal,
 }
 
+/// Performs convolution with scale variations using LHAPDF backend (legacy).
 pub fn convolve_scales(
     grid: &Grid,
     conv_funs: &mut [Pdf],
@@ -357,6 +400,131 @@ pub fn convolve_scales(
     }
 }
 
+/// Performs convolution with scale variations using the backend abstraction.
+#[allow(clippy::too_many_arguments)]
+pub fn convolve_scales_with_backend(
+    grid: &Grid,
+    conv_funs: &mut [Box<dyn PdfBackend>],
+    conv_types: &[ConvType],
+    orders: &[(u8, u8)],
+    bins: &[usize],
+    channels: &[bool],
+    scales: &[(f64, f64, f64)],
+    mode: ConvoluteMode,
+    cfg: &GlobalConfiguration,
+) -> Vec<f64> {
+    let orders: Vec<_> = grid
+        .orders()
+        .iter()
+        .map(|order| {
+            orders.is_empty()
+                || orders
+                    .iter()
+                    .any(|other| (order.alphas == other.0) && (order.alpha == other.1))
+        })
+        .collect();
+
+    if cfg.force_positive {
+        for fun in conv_funs.iter_mut() {
+            fun.set_force_positive(ForcePositive::ClipNegative);
+        }
+    }
+
+    // TODO: promote this to an error
+    assert!(
+        cfg.use_alphas_from < conv_funs.len(),
+        "expected `use_alphas_from` to be an integer within `[0, {})`, but got `{}`",
+        conv_funs.len(),
+        cfg.use_alphas_from
+    );
+
+    // Get x_min/x_max before creating closures (requires mut)
+    let x_min_max: Vec<_> = conv_funs
+        .iter_mut()
+        .map(|fun| (fun.x_min(), fun.x_max()))
+        .collect();
+
+    // Create closures for xfx evaluation
+    let mut funs: Vec<_> = conv_funs
+        .iter()
+        .zip(&x_min_max)
+        .map(|(fun, &(x_min, x_max))| {
+            move |id: i32, x: f64, q2: f64| {
+                if !cfg.allow_extrapolation && (x < x_min || x > x_max) {
+                    0.0
+                } else {
+                    fun.xfx_q2(id, x, q2)
+                }
+            }
+        })
+        .collect();
+
+    let xfx: Vec<_> = funs
+        .iter_mut()
+        .map(|fun| fun as &mut dyn FnMut(i32, f64, f64) -> f64)
+        .collect();
+
+    // Create alphas closures
+    let mut alphas_funs: Vec<_> = conv_funs
+        .iter()
+        .map(|fun| {
+            let fun_ref = fun.as_ref();
+            move |q2: f64| fun_ref.alphas_q2(q2)
+        })
+        .collect();
+
+    // Get particle IDs from the PDFs using the backend interface
+    let convolutions: Vec<_> = conv_funs
+        .iter()
+        .zip(conv_types)
+        .map(|(fun, &conv_type)| {
+            let pid = fun.particle_id();
+            Conv::new(conv_type, pid)
+        })
+        .collect();
+
+    let mut cache = ConvolutionCache::new(convolutions, xfx, &mut alphas_funs[cfg.use_alphas_from]);
+    let mut results = grid.convolve(&mut cache, &orders, bins, channels, scales);
+
+    match mode {
+        ConvoluteMode::Asymmetry => {
+            let bin_count = grid.bwfl().len();
+
+            // calculating the asymmetry for a subset of bins doesn't work
+            assert!((bins.is_empty() || (bins.len() == bin_count)) && (bin_count % 2 == 0));
+
+            results
+                .iter()
+                .skip((bin_count / 2) * scales.len())
+                .zip(
+                    results
+                        .chunks_exact(scales.len())
+                        .take(bin_count / 2)
+                        .rev()
+                        .flatten(),
+                )
+                .map(|(pos, neg)| (pos - neg) / (pos + neg))
+                .collect()
+        }
+        ConvoluteMode::Integrated => {
+            results
+                .iter_mut()
+                .zip(
+                    grid.bwfl()
+                        .normalizations()
+                        .into_iter()
+                        .enumerate()
+                        .filter(|(index, _)| bins.is_empty() || bins.contains(index))
+                        .flat_map(|(_, norm)| iter::repeat(norm).take(scales.len())),
+                )
+                .for_each(|(value, norm)| *value *= norm);
+
+            results
+        }
+        ConvoluteMode::Normal => results,
+    }
+}
+
 pub fn scales_vector(grid: &Grid, scales: usize) -> &[(f64, f64, f64)] {
     let Scales { fac, frg, .. } = grid.scales();
 
@@ -373,6 +541,7 @@ pub fn scales_vector(grid: &Grid, scales: usize) -> &[(f64, f64, f64)] {
     }
 }
 
+/// Performs convolution using LHAPDF backend (legacy).
 pub fn convolve(
     grid: &Grid,
     conv_funs: &mut [Pdf],
@@ -397,6 +566,32 @@ pub fn convolve(
     )
 }
 
+/// Performs convolution using the backend abstraction.
+#[allow(clippy::too_many_arguments)]
+pub fn convolve_with_backend(
+    grid: &Grid,
+    conv_funs: &mut [Box<dyn PdfBackend>],
+    conv_types: &[ConvType],
+    orders: &[(u8, u8)],
+    bins: &[usize],
+    lumis: &[bool],
+    scales: usize,
+    mode: ConvoluteMode,
+    cfg: &GlobalConfiguration,
+) -> Vec<f64> {
+    convolve_scales_with_backend(
+        grid,
+        conv_funs,
+        conv_types,
+        orders,
+        bins,
+        lumis,
+        scales_vector(grid, scales),
+        mode,
+        cfg,
+    )
+}
+
 pub fn convolve_limits(grid: &Grid, bins: &[usize], mode: ConvoluteMode) -> Vec<Vec<(f64, f64)>> {
     let limits: Vec<_> = grid
         .bwfl()

pineappl_cli/src/helpers.rs

@@ -14,6 +14,7 @@ mod helpers;
 mod import;
 mod merge;
 mod orders;
+pub mod pdf_backend;
 mod plot;
 mod pull;
 mod read;
@@ -41,6 +42,9 @@ pub struct GlobalConfiguration {
     /// Choose the PDF/FF set for the strong coupling.
     #[arg(default_value = "0", long, value_name = "IDX")]
     pub use_alphas_from: usize,
+    /// Select the PDF interpolation backend: 'lhapdf' or 'neopdf'.
+    #[arg(default_value = "lhapdf", long, value_name = "BACKEND")]
+    pub pdf_backend: pdf_backend::Backend,
 }
 
 #[enum_dispatch]