Merge pull request #64 from Radonirinaunimi/chebyshev

Optimize `Chebyshev` for Multi-dimensional data using batch vectorization

Author: Radonirinaunimi

.github/actions/cache-data/action.yml

@@ -10,7 +10,7 @@ runs:
       uses: actions/cache@v4
       with:
         path: neopdf-data
-        key: data-v8
+        key: data-v10
     - name: Download data if cache miss
       if: steps.cache-data.outputs.cache-hit != 'true'
       run: |

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added a logic to compute Chebyshev interpolations in batches (https://github.com/Radonirinaunimi/neopdf/pull/64).
 - Added proper LHAPDF drop-in compatibility layer for no-code migration.
 - Added an interface to the Wolfram Language to allow Rust APIs to be called in
   Mathematica.

maintainer/download-data.sh

@@ -22,6 +22,10 @@ NEOPDF_SETS=(
   nNNPDF30_nlo_as_0118
 )
 
+TMDLIB_SETS=(
+  MAP22_grids_FF_Km_N3LL
+)
+
 # Store the data in the root of the repository
 cd ..
 test -d neopdf-data || mkdir neopdf-data
@@ -37,3 +41,8 @@ done
 for neo in "${NEOPDF_SETS[@]}"; do
   wget --no-verbose --no-clobber -P neopdf-data "https://data.nnpdf.science/neopdf/data/${neo}.neopdf.lz4"
 done
+
+# Dowload TMDlib sets
+for tmd in "${TMDLIB_SETS[@]}"; do
+  wget --no-verbose --no-clobber -P neopdf-data "https://data.nnpdf.science/neopdf/data/${tmd}.neopdf.lz4"
+done

neopdf/benches/bench_pdf.rs

@@ -10,6 +10,32 @@ fn xfxq2(c: &mut Criterion) {
     });
 }
 
+fn xfxq2_cheby(c: &mut Criterion) {
+    let pdf = PDF::load("MAP22_grids_FF_Km_N3LL.neopdf.lz4", 0);
+
+    c.bench_function("xfxq2_cheby", |b| {
+        b.iter(|| {
+            pdf.xfxq2(
+                std::hint::black_box(2),
+                std::hint::black_box(&[1e-2, 5e-1, 10.0]),
+            )
+        })
+    });
+}
+
+fn xfxq2_cheby_batch(c: &mut Criterion) {
+    let pdf = PDF::load("MAP22_grids_FF_Km_N3LL.neopdf.lz4", 0);
+
+    c.bench_function("xfxq2_cheby_batch", |b| {
+        b.iter(|| {
+            pdf.xfxq2_cheby_batch(
+                std::hint::black_box(2),
+                std::hint::black_box(&[&[1e-2, 5e-1, 10.0]]),
+            )
+        })
+    });
+}
+
 fn xfxq2s(c: &mut Criterion) {
     let pdf = PDF::load("NNPDF40_nnlo_as_01180", 0);
 
@@ -45,5 +71,12 @@ fn xfxq2_members(c: &mut Criterion) {
     });
 }
 
-criterion_group!(benches, xfxq2, xfxq2s, xfxq2_members);
+criterion_group!(
+    benches,
+    xfxq2,
+    xfxq2s,
+    xfxq2_members,
+    xfxq2_cheby,
+    xfxq2_cheby_batch
+);
 criterion_main!(benches);

neopdf/src/gridpdf.rs

@@ -8,6 +8,7 @@
 use core::panic;
 use ndarray::{Array1, Array2};
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 use thiserror::Error;
 
 use super::alphas::AlphaS;
@@ -244,7 +245,7 @@ impl GridPDF {
             Some(ForcePositive::ClipNegative) => value.max(0.0),
             Some(ForcePositive::ClipSmall) => value.max(1e-10),
             Some(ForcePositive::NoClipping) => value,
-            None => value,
+            _ => value,
         }
     }
 
@@ -286,9 +287,10 @@ impl GridPDF {
             Error::SubgridNotFound { x, q2 }
         })?;
 
-        let pid_idx = self.knot_array.pid_index(flavor_id).ok_or_else(|| {
-            Error::InterpolationError(format!("Invalid flavor ID: {}", flavor_id))
-        })?;
+        let pid_idx = self
+            .knot_array
+            .pid_index(flavor_id)
+            .ok_or_else(|| Error::InterpolationError(format!("Invalid flavor ID: {flavor_id}")))?;
 
         let use_log = matches!(
             self.info.interpolator_type,
@@ -335,6 +337,81 @@ impl GridPDF {
         Array2::from_shape_vec(grid_shape, data).unwrap()
     }
 
+    /// Interpolates PDF values for multiple points in parallel using Chebyshev batch interpolation.
+    ///
+    /// # Arguments
+    ///
+    /// * `flavor_id` - The flavor ID.
+    /// * `points` - A slice containing the collection of knots to interpolate on.
+    ///   A knot is a collection of points containing `(nucleon, alphas, x, Q2)`.
+    ///
+    /// # Returns
+    ///
+    /// A `Vec<f64>` of interpolated PDF values.
+    pub fn xfxq2_cheby_batch(&self, flavor_id: i32, points: &[&[f64]]) -> Result<Vec<f64>, Error> {
+        if points.is_empty() {
+            return Ok(Vec::new());
+        }
+
+        let pid_idx = self
+            .knot_array
+            .pid_index(flavor_id)
+            .ok_or_else(|| Error::InterpolationError(format!("Invalid flavor ID: {flavor_id}")))?;
+
+        if !matches!(self.info.interpolator_type, InterpolatorType::LogChebyshev) {
+            return Err(Error::InterpolationError(
+                "xfxq2_cheby_batch only supports LogChebyshev interpolator".to_string(),
+            ));
+        }
+
+        let mut subgrid_groups: HashMap<usize, Vec<(usize, &[f64])>> = HashMap::new();
+        for (i, point) in points.iter().enumerate() {
+            let subgrid_idx = self.knot_array.find_subgrid(point).ok_or_else(|| {
+                let (x, q2) = self.get_x_q2(point);
+                Error::SubgridNotFound { x, q2 }
+            })?;
+
+            subgrid_groups
+                .entry(subgrid_idx)
+                .or_default()
+                .push((i, *point));
+        }
+
+        let mut all_results: Vec<(usize, f64)> = Vec::new();
+
+        for (subgrid_idx, group) in subgrid_groups {
+            let subgrid = &self.knot_array.subgrids[subgrid_idx];
+
+            let (indices, group_points): (Vec<_>, Vec<_>) = group.into_iter().unzip();
+
+            let log_points: Vec<Vec<f64>> = group_points
+                .iter()
+                .map(|p| p.iter().map(|&v| v.ln()).collect::<Vec<f64>>())
+                .collect();
+
+            let batch_interpolator =
+                InterpolatorFactory::create_batch_interpolator(subgrid, pid_idx)
+                    .map_err(Error::InterpolationError)?;
+
+            let results = batch_interpolator
+                .interpolate(log_points)
+                .map_err(|e| Error::InterpolationError(e.to_string()))?;
+
+            for (original_index, result) in indices.into_iter().zip(results) {
+                all_results.push((original_index, result));
+            }
+        }
+
+        // sort the results according to the original index
+        all_results.sort_by_key(|&(i, _)| i);
+        let final_results = all_results
+            .into_iter()
+            .map(|(_, r)| self.apply_force_positive(r))
+            .collect();
+
+        Ok(final_results)
+    }
+
     /// Get the values of the momentum fraction `x` and momentum scale `Q2`.
     ///
     /// # Arguments

neopdf/src/interpolator.rs

@@ -12,6 +12,7 @@
 //! The [`SubGrid`] struct is defined in `subgrid.rs`.
 
 use ndarray::{s, OwnedRepr};
+use ninterp::data::{InterpData2D, InterpData3D};
 use ninterp::error::InterpolateError;
 use ninterp::interpolator::{
     Extrapolate, Interp2D, Interp2DOwned, Interp3D, Interp3DOwned, InterpND, InterpNDOwned,
@@ -23,7 +24,7 @@ use ninterp::strategy::Linear;
 use super::metadata::InterpolatorType;
 use super::strategy::{
     BilinearInterpolation, LogBicubicInterpolation, LogBilinearInterpolation,
-    LogChebyshevInterpolation, LogTricubicInterpolation,
+    LogChebyshevBatchInterpolation, LogChebyshevInterpolation, LogTricubicInterpolation,
 };
 use super::subgrid::SubGrid;
 
@@ -115,6 +116,40 @@ where
     }
 }
 
+/// An enum to dispatch batch interpolation to the correct Chebyshev interpolator.
+pub enum BatchInterpolator {
+    Chebyshev2D(
+        LogChebyshevBatchInterpolation<2>,
+        InterpData2D<OwnedRepr<f64>>,
+    ),
+    Chebyshev3D(
+        LogChebyshevBatchInterpolation<3>,
+        InterpData3D<OwnedRepr<f64>>,
+    ),
+}
+
+impl BatchInterpolator {
+    /// Interpolates a batch of points.
+    pub fn interpolate(&self, points: Vec<Vec<f64>>) -> Result<Vec<f64>, InterpolateError> {
+        match self {
+            BatchInterpolator::Chebyshev2D(strategy, data) => {
+                let points_2d: Vec<[f64; 2]> = points
+                    .into_iter()
+                    .map(|p| p.try_into().expect("Invalid point dimension for 2D"))
+                    .collect();
+                strategy.interpolate(data, &points_2d)
+            }
+            BatchInterpolator::Chebyshev3D(strategy, data) => {
+                let points_3d: Vec<[f64; 3]> = points
+                    .into_iter()
+                    .map(|p| p.try_into().expect("Invalid point dimension for 3D"))
+                    .collect();
+                strategy.interpolate(data, &points_3d)
+            }
+        }
+    }
+}
+
 /// Factory for creating dynamic interpolators based on interpolation type and grid dimensions.
 pub struct InterpolatorFactory;
 
@@ -455,6 +490,94 @@ impl InterpolatorFactory {
             _ => panic!("Unsupported 5D interpolator: {:?}", interp_type),
         }
     }
+
+    pub fn create_batch_interpolator(
+        subgrid: &SubGrid,
+        pid_idx: usize,
+    ) -> Result<BatchInterpolator, String> {
+        match subgrid.interpolation_config() {
+            InterpolationConfig::TwoD => {
+                let mut strategy = LogChebyshevBatchInterpolation::<2>::default();
+                let grid_slice = subgrid.grid_slice(pid_idx).to_owned();
+
+                let data = InterpData2D::new(
+                    subgrid.xs.mapv(f64::ln),
+                    subgrid.q2s.mapv(f64::ln),
+                    grid_slice,
+                )
+                .map_err(|e| e.to_string())?;
+                strategy.init(&data).map_err(|e| e.to_string())?;
+
+                Ok(BatchInterpolator::Chebyshev2D(strategy, data))
+            }
+            InterpolationConfig::ThreeDNucleons => {
+                let mut strategy = LogChebyshevBatchInterpolation::<3>::default();
+                let grid_data = subgrid.grid.slice(s![.., 0, pid_idx, 0, .., ..]).to_owned();
+
+                let reshaped_data = grid_data
+                    .into_shape_with_order((
+                        subgrid.nucleons.len(),
+                        subgrid.xs.len(),
+                        subgrid.q2s.len(),
+                    ))
+                    .expect("Failed to reshape 3D data");
+
+                let data = InterpData3D::new(
+                    subgrid.nucleons.mapv(f64::ln),
+                    subgrid.xs.mapv(f64::ln),
+                    subgrid.q2s.mapv(f64::ln),
+                    reshaped_data,
+                )
+                .map_err(|e| e.to_string())?;
+                strategy.init(&data).map_err(|e| e.to_string())?;
+
+                Ok(BatchInterpolator::Chebyshev3D(strategy, data))
+            }
+            InterpolationConfig::ThreeDAlphas => {
+                let mut strategy = LogChebyshevBatchInterpolation::<3>::default();
+                let grid_data = subgrid.grid.slice(s![0, .., pid_idx, 0, .., ..]).to_owned();
+
+                let reshaped_data = grid_data
+                    .into_shape_with_order((
+                        subgrid.alphas.len(),
+                        subgrid.xs.len(),
+                        subgrid.q2s.len(),
+                    ))
+                    .expect("Failed to reshape 3D data");
+
+                let data = InterpData3D::new(
+                    subgrid.alphas.mapv(f64::ln),
+                    subgrid.xs.mapv(f64::ln),
+                    subgrid.q2s.mapv(f64::ln),
+                    reshaped_data,
+                )
+                .map_err(|e| e.to_string())?;
+                strategy.init(&data).map_err(|e| e.to_string())?;
+
+                Ok(BatchInterpolator::Chebyshev3D(strategy, data))
+            }
+            InterpolationConfig::ThreeDKt => {
+                let mut strategy = LogChebyshevBatchInterpolation::<3>::default();
+                let grid_data = subgrid.grid.slice(s![0, 0, pid_idx, .., .., ..]).to_owned();
+
+                let reshaped_data = grid_data
+                    .into_shape_with_order((subgrid.kts.len(), subgrid.xs.len(), subgrid.q2s.len()))
+                    .expect("Failed to reshape 3D data");
+
+                let data = InterpData3D::new(
+                    subgrid.kts.mapv(f64::ln),
+                    subgrid.xs.mapv(f64::ln),
+                    subgrid.q2s.mapv(f64::ln),
+                    reshaped_data,
+                )
+                .map_err(|e| e.to_string())?;
+                strategy.init(&data).map_err(|e| e.to_string())?;
+
+                Ok(BatchInterpolator::Chebyshev3D(strategy, data))
+            }
+            _ => Err("Unsupported dimension for batch interpolation".to_string()),
+        }
+    }
 }
 
 #[cfg(test)]

neopdf/src/pdf.rs

@@ -287,6 +287,23 @@ impl PDF {
         self.grid_pdf.xfxq2s(pids, slice_points)
     }
 
+    /// Interpolates the PDF value (xf) for multiple points using Chebyshev batch interpolation.
+    ///
+    /// Abstraction to the `GridPDF::xfxq2_cheby_batch` method.
+    ///
+    /// # Arguments
+    ///
+    /// * `pid` - The flavor ID.
+    /// * `points` - A slice containing the collection of knots to interpolate on.
+    ///   A knot is a collection of points containing `(nucleon, alphas, x, Q2)`.
+    ///
+    /// # Returns
+    ///
+    /// A `Vec<f64>` of interpolated PDF values.
+    pub fn xfxq2_cheby_batch(&self, pid: i32, points: &[&[f64]]) -> Vec<f64> {
+        self.grid_pdf.xfxq2_cheby_batch(pid, points).unwrap()
+    }
+
     /// Interpolates the strong coupling constant `alpha_s` for a given Q2.
     ///
     /// Abstraction to the `GridPDF::alphas_q2` method.