diff --git a/.github/actions/cache-data/action.yml b/.github/actions/cache-data/action.yml index 94fced37..bac12935 100644 --- a/.github/actions/cache-data/action.yml +++ b/.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: | diff --git a/CHANGELOG.md b/CHANGELOG.md index 7f4a56f6..f04385dc 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/maintainer/download-data.sh b/maintainer/download-data.sh index 7a4761a8..7d7d99f6 100755 --- a/maintainer/download-data.sh +++ b/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 diff --git a/neopdf/benches/bench_pdf.rs b/neopdf/benches/bench_pdf.rs index 9041b803..a73243fc 100644 --- a/neopdf/benches/bench_pdf.rs +++ b/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); diff --git a/neopdf/src/gridpdf.rs b/neopdf/src/gridpdf.rs index 0c2bc5d9..baad6e48 100644 --- a/neopdf/src/gridpdf.rs +++ b/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` of interpolated PDF values. + pub fn xfxq2_cheby_batch(&self, flavor_id: i32, points: &[&[f64]]) -> Result, 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> = 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> = group_points + .iter() + .map(|p| p.iter().map(|&v| v.ln()).collect::>()) + .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 diff --git a/neopdf/src/interpolator.rs b/neopdf/src/interpolator.rs index 411ac764..43e886bc 100644 --- a/neopdf/src/interpolator.rs +++ b/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>, + ), + Chebyshev3D( + LogChebyshevBatchInterpolation<3>, + InterpData3D>, + ), +} + +impl BatchInterpolator { + /// Interpolates a batch of points. + pub fn interpolate(&self, points: Vec>) -> Result, 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 { + 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)] diff --git a/neopdf/src/pdf.rs b/neopdf/src/pdf.rs index bd6893ef..c7f9a33f 100644 --- a/neopdf/src/pdf.rs +++ b/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` of interpolated PDF values. + pub fn xfxq2_cheby_batch(&self, pid: i32, points: &[&[f64]]) -> Vec { + 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. diff --git a/neopdf/src/strategy.rs b/neopdf/src/strategy.rs index f7d1a42c..67beb733 100644 --- a/neopdf/src/strategy.rs +++ b/neopdf/src/strategy.rs @@ -14,7 +14,7 @@ //! All interpolation strategies are designed to work with `ninterp`'s data structures and traits, //! ensuring compatibility and extensibility. -use ndarray::{Data, RawDataClone}; +use ndarray::{Array2, Axis, Data, RawDataClone}; use ninterp::data::{InterpData1D, InterpData2D, InterpData3D}; use ninterp::error::{InterpolateError, ValidateError}; use ninterp::strategy::traits::{Strategy1D, Strategy2D, Strategy3D}; @@ -1031,13 +1031,6 @@ where let x_min = *x_coords.first().unwrap(); let x_max = *x_coords.last().unwrap(); - if x < x_min || x > x_max { - return Err(InterpolateError::Other(format!( - "Input point {} is outside the grid range [{}, {}]", - x, x_min, x_max - ))); - } - if (x_max - x_min).abs() < 1e-15 { return Ok(f_values[0]); } @@ -1089,21 +1082,8 @@ where let x_min = *x_coords.first().unwrap(); let x_max = *x_coords.last().unwrap(); - if x < x_min || x > x_max { - return Err(InterpolateError::Other(format!( - "Input point x={} is outside the grid range [{}, {}]", - x, x_min, x_max - ))); - } - let y_min = *y_coords.first().unwrap(); let y_max = *y_coords.last().unwrap(); - if y < y_min || y > y_max { - return Err(InterpolateError::Other(format!( - "Input point y={} is outside the grid range [{}, {}]", - y, y_min, y_max - ))); - } let t_x = 2.0 * (x - x_min) / (x_max - x_min) - 1.0; let t_y = 2.0 * (y - y_min) / (y_max - y_min) - 1.0; @@ -1160,30 +1140,10 @@ where let x_min = *x_coords.first().unwrap(); let x_max = *x_coords.last().unwrap(); - if x < x_min || x > x_max { - return Err(InterpolateError::Other(format!( - "Input point x={} is outside the grid range [{}, {}]", - x, x_min, x_max - ))); - } - let y_min = *y_coords.first().unwrap(); let y_max = *y_coords.last().unwrap(); - if y < y_min || y > y_max { - return Err(InterpolateError::Other(format!( - "Input point y={} is outside the grid range [{}, {}]", - y, y_min, y_max - ))); - } - let z_min = *z_coords.first().unwrap(); let z_max = *z_coords.last().unwrap(); - if z < z_min || z > z_max { - return Err(InterpolateError::Other(format!( - "Input point z={} is outside the grid range [{}, {}]", - z, z_min, z_max - ))); - } let t_x = 2.0 * (x - x_min) / (x_max - x_min) - 1.0; let t_y = 2.0 * (y - y_min) / (y_max - y_min) - 1.0; @@ -1210,6 +1170,306 @@ where } } +/// Implements a global N-dimensional batch interpolation using Chebyshev polynomials +/// with logarithmic coordinate scaling. +/// +/// This strategy is optimized for interpolating multiple points at once by leveraging +/// matrix operations with `ndarray`. +/// +/// TODO: Potentially merge this with `LogChebyshevInterpolation`. +#[derive(Debug, Clone)] +pub struct LogChebyshevBatchInterpolation { + // Pre-computed weights for the barycentric formula for each dimension. + weights: [Vec; DIM], + // Grid points in the t-domain [-1, 1] for each dimension. + t_coords: [Vec; DIM], +} + +impl Default for LogChebyshevBatchInterpolation { + fn default() -> Self { + Self { + weights: std::array::from_fn(|_| Vec::new()), + t_coords: std::array::from_fn(|_| Vec::new()), + } + } +} + +impl Serialize for LogChebyshevBatchInterpolation { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut state = serializer.serialize_struct("LogChebyshevBatchInterpolation", 2)?; + state.serialize_field("weights", &self.weights.as_slice())?; + state.serialize_field("t_coords", &self.t_coords.as_slice())?; + state.end() + } +} + +impl<'de, const DIM: usize> Deserialize<'de> for LogChebyshevBatchInterpolation { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + #[derive(Deserialize)] + struct Helper { + weights: Vec>, + t_coords: Vec>, + } + + let helper = Helper::deserialize(deserializer)?; + let weights = helper.weights.try_into().map_err(|v: Vec>| { + serde::de::Error::invalid_length(v.len(), &"an array of the correct length") + })?; + let t_coords = helper.t_coords.try_into().map_err(|v: Vec>| { + serde::de::Error::invalid_length(v.len(), &"an array of the correct length") + })?; + + Ok(Self { weights, t_coords }) + } +} + +impl LogChebyshevBatchInterpolation { + /// Computes the barycentric weights for a given set of Chebyshev points. + /// The formula for the weights is `w_j = (-1)^j * delta_j`, where `delta_j` + /// is 1/2 for the first and last points, and 1 otherwise. + fn compute_barycentric_weights(n: usize) -> Vec { + let mut weights = vec![1.0; n]; + (0..n).for_each(|j| { + if j % 2 == 1 { + weights[j] = -1.0; + } + }); + weights[0] *= 0.5; + + if n > 1 { + weights[n - 1] *= 0.5; + } + + weights + } + + /// Compute barycentric coefficients for multiple points in batch + fn barycentric_coefficients( + t_values: &[f64], + t_coords: &[f64], + weights: &[f64], + ) -> Array2 { + let num_points = t_values.len(); + let num_coords = t_coords.len(); + let mut coeffs = Array2::::zeros((num_points, num_coords)); + + for (p, &t) in t_values.iter().enumerate() { + let mut found_exact = false; + for (j, &t_j) in t_coords.iter().enumerate() { + if (t - t_j).abs() < 1e-15 { + coeffs[[p, j]] = 1.0; + found_exact = true; + break; + } + } + + if !found_exact { + let mut terms = Vec::with_capacity(num_coords); + for (j, &t_j) in t_coords.iter().enumerate() { + terms.push(weights[j] / (t - t_j)); + } + + let sum: f64 = terms.iter().sum(); + + for (j, &term) in terms.iter().enumerate() { + coeffs[[p, j]] = term / sum; + } + } + } + + coeffs + } +} + +impl LogChebyshevBatchInterpolation<1> { + pub fn init(&mut self, data: &InterpData1D) -> Result<(), ValidateError> + where + D: Data + RawDataClone + Clone, + { + let x_coords = data.grid[0].as_slice().unwrap(); + let n = x_coords.len(); + if n < 2 { + return Err(ValidateError::Other( + "LogChebyshevBatchInterpolation requires at least 2 grid points.".to_string(), + )); + } + + self.t_coords[0] = (0..n) + .map(|j| (PI * (n - 1 - j) as f64 / (n - 1) as f64).cos()) + .collect(); + + self.weights[0] = Self::compute_barycentric_weights(n); + + Ok(()) + } + + pub fn interpolate( + &self, + data: &InterpData1D, + points: &[[f64; 1]], + ) -> Result, InterpolateError> + where + D: Data + RawDataClone + Clone, + { + let x_coords = data.grid[0].as_slice().unwrap(); + let f_values = data.values.to_owned(); + + let x_min = *x_coords.first().unwrap(); + let x_max = *x_coords.last().unwrap(); + + let mut t_x_vals = Vec::with_capacity(points.len()); + for &[x] in points { + t_x_vals.push(2.0 * (x - x_min) / (x_max - x_min) - 1.0); + } + + let c_x = Self::barycentric_coefficients(&t_x_vals, &self.t_coords[0], &self.weights[0]); + let results = c_x.dot(&f_values); + + Ok(results.to_vec()) + } +} + +impl LogChebyshevBatchInterpolation<2> { + pub fn init(&mut self, data: &InterpData2D) -> Result<(), ValidateError> + where + D: Data + RawDataClone + Clone, + { + for dim in 0..2 { + let x_coords = data.grid[dim].as_slice().unwrap(); + let n = x_coords.len(); + if n < 2 { + return Err(ValidateError::Other( + "LogChebyshevBatchInterpolation requires at least 2 grid points per dimension." + .to_string(), + )); + } + self.t_coords[dim] = (0..n) + .map(|j| (PI * (n - 1 - j) as f64 / (n - 1) as f64).cos()) + .collect(); + self.weights[dim] = Self::compute_barycentric_weights(n); + } + + Ok(()) + } + + pub fn interpolate( + &self, + data: &InterpData2D, + points: &[[f64; 2]], + ) -> Result, InterpolateError> + where + D: Data + RawDataClone + Clone, + { + let x_coords = data.grid[0].as_slice().unwrap(); + let y_coords = data.grid[1].as_slice().unwrap(); + + let x_min = *x_coords.first().unwrap(); + let x_max = *x_coords.last().unwrap(); + let y_min = *y_coords.first().unwrap(); + let y_max = *y_coords.last().unwrap(); + + let mut t_x_vals = Vec::with_capacity(points.len()); + let mut t_y_vals = Vec::with_capacity(points.len()); + for &[x, y] in points { + t_x_vals.push(2.0 * (x - x_min) / (x_max - x_min) - 1.0); + t_y_vals.push(2.0 * (y - y_min) / (y_max - y_min) - 1.0); + } + + let c_x = Self::barycentric_coefficients(&t_x_vals, &self.t_coords[0], &self.weights[0]); + let c_y = Self::barycentric_coefficients(&t_y_vals, &self.t_coords[1], &self.weights[1]); + let v = data.values.to_owned(); + let results = (&c_x.dot(&v) * &c_y).sum_axis(Axis(1)); + + Ok(results.to_vec()) + } +} + +impl LogChebyshevBatchInterpolation<3> { + pub fn init(&mut self, data: &InterpData3D) -> Result<(), ValidateError> + where + D: Data + RawDataClone + Clone, + { + for dim in 0..3 { + let x_coords = data.grid[dim].as_slice().unwrap(); + let n = x_coords.len(); + if n < 2 { + return Err(ValidateError::Other( + "LogChebyshevBatchInterpolation requires at least 2 grid points per dimension." + .to_string(), + )); + } + self.t_coords[dim] = (0..n) + .map(|j| (PI * (n - 1 - j) as f64 / (n - 1) as f64).cos()) + .collect(); + self.weights[dim] = Self::compute_barycentric_weights(n); + } + Ok(()) + } + + pub fn interpolate( + &self, + data: &InterpData3D, + points: &[[f64; 3]], + ) -> Result, InterpolateError> + where + D: Data + RawDataClone + Clone, + { + let x_coords = data.grid[0].as_slice().unwrap(); + let y_coords = data.grid[1].as_slice().unwrap(); + let z_coords = data.grid[2].as_slice().unwrap(); + + let x_min = *x_coords.first().unwrap(); + let x_max = *x_coords.last().unwrap(); + let y_min = *y_coords.first().unwrap(); + let y_max = *y_coords.last().unwrap(); + let z_min = *z_coords.first().unwrap(); + let z_max = *z_coords.last().unwrap(); + + let mut t_x_vals = Vec::with_capacity(points.len()); + let mut t_y_vals = Vec::with_capacity(points.len()); + let mut t_z_vals = Vec::with_capacity(points.len()); + + for &[x, y, z] in points { + t_x_vals.push(2.0 * (x - x_min) / (x_max - x_min) - 1.0); + t_y_vals.push(2.0 * (y - y_min) / (y_max - y_min) - 1.0); + t_z_vals.push(2.0 * (z - z_min) / (z_max - z_min) - 1.0); + } + + let c_x = Self::barycentric_coefficients(&t_x_vals, &self.t_coords[0], &self.weights[0]); + let c_y = Self::barycentric_coefficients(&t_y_vals, &self.t_coords[1], &self.weights[1]); + let c_z = Self::barycentric_coefficients(&t_z_vals, &self.t_coords[2], &self.weights[2]); + + let v = &data.values; + + let num_points = points.len(); + let (nx, ny, nz) = (x_coords.len(), y_coords.len(), z_coords.len()); + + let v_flat = v.to_owned().into_shape_with_order((nx, ny * nz)).unwrap(); + let temp1 = c_x.dot(&v_flat); + let temp1_3d = temp1.into_shape_with_order((num_points, ny, nz)).unwrap(); + + let mut results = Vec::with_capacity(num_points); + for p in 0..num_points { + let temp_slice = temp1_3d.index_axis(Axis(0), p); + let cy_slice = c_y.row(p); + let cz_slice = c_z.row(p); + + let temp2 = cy_slice.dot(&temp_slice); + let result = cz_slice.dot(&temp2); + results.push(result); + } + + Ok(results) + } +} + #[cfg(test)] mod tests { use super::*; @@ -1417,19 +1677,19 @@ mod tests { let alphas_cubic = AlphaSCubicInterpolation; // Test within interpolation range - let result = alphas_cubic.interpolate(&data, &[2.25f64.ln()]).unwrap(); // Q=1.5 + let result = alphas_cubic.interpolate(&data, &[2.25f64.ln()]).unwrap(); assert!(result > 0.1 && result < 0.14); // Test at grid point - let result = alphas_cubic.interpolate(&data, &[4.0f64.ln()]).unwrap(); // Q=2.0 + let result = alphas_cubic.interpolate(&data, &[4.0f64.ln()]).unwrap(); assert_close(result, 0.11, EPSILON); // Test extrapolation below range - let result = alphas_cubic.interpolate(&data, &[0.5f64.ln()]).unwrap(); // Q=sqrt(0.5) + let result = alphas_cubic.interpolate(&data, &[0.5f64.ln()]).unwrap(); assert!(result < 0.1); // Test extrapolation above range - let result = alphas_cubic.interpolate(&data, &[30.0f64.ln()]).unwrap(); // Q=sqrt(30) + let result = alphas_cubic.interpolate(&data, &[30.0f64.ln()]).unwrap(); assert_close(result, 0.14, EPSILON); } @@ -1679,4 +1939,93 @@ mod tests { assert_close(result, expected, EPSILON); } + + #[test] + fn test_log_chebyshev_batch_interpolation_1d() { + let n = 21; + let x_min: f64 = 0.1; + let x_max: f64 = 10.0; + let x_coords = create_cheby_grid(n, x_min, x_max); + + let f_values: Vec = x_coords.iter().map(|&x| x.ln()).collect(); + let data = create_test_data_1d(x_coords.iter().map(|v| v.ln()).collect(), f_values); + let mut cheby = LogChebyshevBatchInterpolation::<1>::default(); + cheby.init(&data).unwrap(); + + let test_points = [[2.5f64.ln()], [5.0f64.ln()], [7.5f64.ln()]]; + let expected: Vec = test_points.iter().map(|p| p[0]).collect(); + let results = cheby.interpolate(&data, &test_points).unwrap(); + + for (res, exp) in results.iter().zip(expected.iter()) { + assert_close(*res, *exp, EPSILON); + } + } + + #[test] + fn test_log_chebyshev_batch_interpolation_2d() { + let n = 11; + let x_coords = create_cheby_grid(n, 0.1, 10.0); + let y_coords = create_cheby_grid(n, 0.1, 10.0); + + let f_values: Vec = x_coords + .iter() + .flat_map(|&x| y_coords.iter().map(move |&y| x.ln() + y.ln())) + .collect(); + + let data = create_test_data_2d( + x_coords.iter().map(|v| v.ln()).collect(), + y_coords.iter().map(|v| v.ln()).collect(), + f_values, + ); + let mut cheby = LogChebyshevBatchInterpolation::<2>::default(); + cheby.init(&data).unwrap(); + + let test_points = [ + [2.5f64.ln(), 3.5f64.ln()], + [5.0f64.ln(), 6.0f64.ln()], + [7.5f64.ln(), 8.5f64.ln()], + ]; + let expected: Vec = test_points.iter().map(|p| p[0] + p[1]).collect(); + let results = cheby.interpolate(&data, &test_points).unwrap(); + + for (res, exp) in results.iter().zip(expected.iter()) { + assert_close(*res, *exp, EPSILON); + } + } + + #[test] + fn test_log_chebyshev_batch_interpolation_3d() { + let n = 7; + let x_coords = create_cheby_grid(n, 0.1, 10.0); + let y_coords = create_cheby_grid(n, 0.1, 10.0); + let z_coords = create_cheby_grid(n, 0.1, 10.0); + + let f_values: Vec = x_coords + .iter() + .cartesian_product(y_coords.iter()) + .cartesian_product(z_coords.iter()) + .map(|((&x, &y), &z)| x.ln() + y.ln() + z.ln()) + .collect(); + + let data = create_test_data_3d( + x_coords.iter().map(|v| v.ln()).collect(), + y_coords.iter().map(|v| v.ln()).collect(), + z_coords.iter().map(|v| v.ln()).collect(), + f_values, + ); + let mut cheby = LogChebyshevBatchInterpolation::<3>::default(); + cheby.init(&data).unwrap(); + + let test_points = [ + [2.5f64.ln(), 3.5f64.ln(), 4.5f64.ln()], + [5.0f64.ln(), 6.0f64.ln(), 7.0f64.ln()], + [7.5f64.ln(), 8.5f64.ln(), 9.5f64.ln()], + ]; + let expected: Vec = test_points.iter().map(|p| p[0] + p[1] + p[2]).collect(); + let results = cheby.interpolate(&data, &test_points).unwrap(); + + for (res, exp) in results.iter().zip(expected.iter()) { + assert_close(*res, *exp, EPSILON); + } + } } diff --git a/neopdf/tests/pdf.rs b/neopdf/tests/pdf.rs index dfc8ef9e..6f71ea2b 100644 --- a/neopdf/tests/pdf.rs +++ b/neopdf/tests/pdf.rs @@ -301,3 +301,40 @@ fn test_pdf_download() { assert!(gluon_as.is_finite()); assert!(gluon_xf.is_finite()); } + +#[test] +pub fn test_xfxq2_cheby_batch() { + let pdf = PDF::load("MAP22_grids_FF_Km_N3LL.neopdf.lz4", 0); + + let nb_points: usize = 20; + let generate_points = |min: f64, max: f64| -> Vec { + (0..nb_points) + .map(|i| min + i as f64 * (max - min) / (nb_points as f64 - 1.0)) + .collect() + }; + let kts: Vec = generate_points(1e-4, 5.0); + let xs: Vec = generate_points(1e-1, 1.0); + let q2s: Vec = generate_points(1.0, 1e4); + + let mut flatten_points: Vec> = Vec::new(); + for &kt in &kts { + for &x in &xs { + for &q2 in &q2s { + flatten_points.push(vec![kt, x, q2]); + } + } + } + + let points_interp: Vec<&[f64]> = flatten_points.iter().map(Vec::as_slice).collect(); + let slice_points: &[&[f64]] = &points_interp; + + let ids: Vec = (-3..=3).filter(|&x| x != 0).collect(); + for &pid in &ids { + let results_batch = pdf.xfxq2_cheby_batch(pid, slice_points); + let results_seq: Vec = slice_points.iter().map(|p| pdf.xfxq2(pid, p)).collect(); + + for (res_b, res_s) in results_batch.iter().zip(results_seq.iter()) { + assert!((res_b - res_s).abs() < LOW_PRECISION); + } + } +} diff --git a/neopdf_capi/src/include/NeoPDF.hpp b/neopdf_capi/src/include/NeoPDF.hpp index a3de2bc9..c3ab2699 100644 --- a/neopdf_capi/src/include/NeoPDF.hpp +++ b/neopdf_capi/src/include/NeoPDF.hpp @@ -157,6 +157,23 @@ class NeoPDF { return neopdf_pdf_xfxq2_nd(this->raw, pid, params.data(), params.size()); } + /** @brief Compute the `xf` value for a generic set of parameters using batch Chebyshev interpolation. */ + std::vector + xfxQ2_cheby_batch(int pid, const std::vector> &points) const { + std::vector c_points(points.size()); + std::vector lengths(points.size()); + for (size_t i = 0; i < points.size(); ++i) { + c_points[i] = points[i].data(); + lengths[i] = points[i].size(); + } + + std::vector results(points.size()); + neopdf_pdf_xfxq2_cheby_batch(this->raw, pid, c_points.data(), lengths.data(), + points.size(), results.data()); + + return results; + } + /** @brief Compute the value of `alphas` at the Q2 value. */ double alphasQ2(double q2) const { return neopdf_pdf_alphas_q2(this->raw, q2); diff --git a/neopdf_capi/src/lib.rs b/neopdf_capi/src/lib.rs index a50d73ea..8768ed7d 100644 --- a/neopdf_capi/src/lib.rs +++ b/neopdf_capi/src/lib.rs @@ -317,6 +317,46 @@ pub unsafe extern "C" fn neopdf_pdf_xfxq2_nd( pdf_obj.xfxq2(id, params) } +/// Interpolates PDF values for multiple points in parallel using Chebyshev batch interpolation. +/// +/// # Safety +/// +/// The `pdf` pointer must be a valid pointer to a `NeoPDF` object. +/// The `points` pointer must be a valid pointer to an array of pointers to `c_double`. +/// The `lengths` pointer must be a valid pointer to an array of `usize`. +/// `num_points` must be the correct number of points. +/// The `results` pointer must be valid for writing `num_points` elements. +/// +/// # Panics +/// +/// TODO +#[no_mangle] +pub unsafe extern "C" fn neopdf_pdf_xfxq2_cheby_batch( + pdf: *mut NeoPDFWrapper, + pid: i32, + points: *const *const c_double, + lengths: *const usize, + num_points: usize, + results: *mut c_double, +) { + assert!(!pdf.is_null()); + let pdf_obj = unsafe { &(*pdf).0 }; + + let points_slices: &[*const c_double] = unsafe { slice::from_raw_parts(points, num_points) }; + let lengths_slice: &[usize] = unsafe { slice::from_raw_parts(lengths, num_points) }; + + let rust_points: Vec<&[f64]> = points_slices + .iter() + .zip(lengths_slice) + .map(|(&p, &l)| unsafe { slice::from_raw_parts(p, l) }) + .collect(); + + let res_vec = pdf_obj.xfxq2_cheby_batch(pid, &rust_points); + + let results_slice = unsafe { slice::from_raw_parts_mut(results, res_vec.len()) }; + results_slice.copy_from_slice(&res_vec); +} + /// Clip the interpolated values if they turned out negatives. /// /// # Panics diff --git a/neopdf_cli/src/converter.rs b/neopdf_cli/src/converter.rs index 760bd3ea..8eb8c5b2 100644 --- a/neopdf_cli/src/converter.rs +++ b/neopdf_cli/src/converter.rs @@ -73,7 +73,37 @@ pub enum Commands { #[arg(short, long)] output: String, }, - /// Update the metadata of the `NeoPDF` grid. + /// Update the metadata of the `NeoPDF` grid + /// + /// Here are the metadata keys that can be updated: + /// + /// - `SetDesc`: description of the set + /// - `SetIndex`: set index + /// - `NumMembers`: number of members + /// - `XMin`: minimum value of x + /// - `XMax`: maximum value of x + /// - `QMin`: minimum value of Q2 + /// - `QMax`: maximum value of Q2 + /// - `Format`: format of the set (can be any string) + /// - `Polarized`: polarization, takes a boolean + /// - `InterpolatorType`: interpolation method (bilinear, logbilinear, logbicubic, logtricubic, chebyshev, interpndlinear) + /// - `ErrorType`: set error (replica, hessian, etc.) + /// - `Particle`: Hadron PID + /// - `FlavorScheme`: flavour scheme (fixed, variable, etc.) + /// - `OrderQCD`: perturbative order + /// - `AlphaS_OrderQCD`: perturbative order of the strong coupling + /// - `MW`: W boson mass + /// - `MZ`: Z boson mass + /// - `MUp`: mass of the up quark + /// - `MDown`: mass of the down quark + /// - `MStrange`: mass of the strange quark + /// - `MCharm`: mass of the charm quark + /// - `MBottom`: mass of the bottom quark + /// - `MTop`: mass of the top quark + /// - `AlphaS_Type`: type of the `AlphaS` interpolation (ipol, etc.) + /// - `NumFlavors`: number of active quark flavours + /// - `SetType`: type of the set (spacelike, timelike) + #[command(verbatim_doc_comment)] Metadata { /// Path to the `NeoPDF` file. #[arg(short, long)] diff --git a/neopdf_cli/tests/pdf.rs b/neopdf_cli/tests/pdf.rs index ebdde642..0cf9e807 100644 --- a/neopdf_cli/tests/pdf.rs +++ b/neopdf_cli/tests/pdf.rs @@ -171,5 +171,28 @@ fn xfxq2_kt_tmdlib() { ]) .assert() .success() - .stdout("0.0655544\n"); + .stdout("0.07892252798564643\n"); +} + +#[test] +#[cfg(feature = "tmdlib")] +fn xfxq2_kt_neopdf_tmdlib() { + Command::cargo_bin("neopdf") + .unwrap() + .args([ + "compute", + "xfx_q2", + "--pdf-name", + "MAP22_grids_FF_Km_N3LL.neopdf.lz4", + "--member", + "0", + "--pid", + "2", + "1.0", + "0.1", + "10.0", + ]) + .assert() + .success() + .stdout("0.07899136744063368\n"); } diff --git a/neopdf_fapi/neopdf.f90 b/neopdf_fapi/neopdf.f90 index 0e2c4853..97937537 100644 --- a/neopdf_fapi/neopdf.f90 +++ b/neopdf_fapi/neopdf.f90 @@ -222,6 +222,16 @@ function c_neopdf_pdf_xfxq2_nd(pdf, id, params, num_params) bind(c, name="neopdf real (c_double) :: c_neopdf_pdf_xfxq2_nd end function + subroutine c_neopdf_pdf_xfxq2_cheby_batch(pdf, flavor_id, points, num_points, num_dims, results) bind(c, name="neopdf_pdf_xfxq2_cheby_batch") + use iso_c_binding + type(c_ptr), value :: pdf + integer(c_int), value :: flavor_id + real(c_double), intent(in) :: points(*) + integer(c_size_t), value :: num_points + integer(c_size_t), value :: num_dims + real(c_double), intent(out) :: results(*) + end subroutine + subroutine c_neopdf_pdf_set_force_positive(pdf, option) bind(c, name="neopdf_pdf_set_force_positive") use iso_c_binding type (c_ptr), value :: pdf @@ -533,6 +543,20 @@ function neopdf_pdf_xfxq2_nd(pdf, id, params) result(res) res = c_neopdf_pdf_xfxq2_nd(pdf%ptr, id, params, int(size(params), c_size_t)) end function + subroutine neopdf_pdf_xfxq2_cheby_batch(pdf, flavor_id, points, results) + implicit none + type(neopdf_pdf), intent(in) :: pdf + integer, intent(in) :: flavor_id + real(dp), intent(in) :: points(:,:) + real(dp), intent(out) :: results(:) + integer(c_size_t) :: num_dims, num_points + + num_dims = size(points, 1) + num_points = size(points, 2) + + call c_neopdf_pdf_xfxq2_cheby_batch(pdf%ptr, flavor_id, points, num_points, num_dims, results) + end subroutine + subroutine neopdf_pdf_set_force_positive(pdf, option) implicit none type (neopdf_pdf), intent(in) :: pdf diff --git a/neopdf_pyapi/src/pdf.rs b/neopdf_pyapi/src/pdf.rs index 59802332..17ad0f30 100644 --- a/neopdf_pyapi/src/pdf.rs +++ b/neopdf_pyapi/src/pdf.rs @@ -403,6 +403,33 @@ impl PyPDF { self.pdf.xfxq2(id, ¶ms) } + /// Interpolates the PDF value (xf) for a list containg a set of parameters. + /// + /// Parameters + /// ---------- + /// id : int + /// The flavor ID (e.g., 21 for gluon, 1 for d-quark). + /// params: list[list[float]] + /// A list containing the list of points. Each element in the list + /// is in turn a list containing the parameters that the grids depends + /// on. If the PDF grid only contains `x` and `Q2` dependence then its + /// value is `[x, q2]`; if it contains either the `A` and `alpha_s` + /// dependence, then its value is `[A, x, q2]` or `[alpha_s, x, q2]` + /// respectively; if it contains both, then `[A, alpha_s, x, q2]`. + /// + /// Returns + /// ------- + /// float + /// The interpolated PDF value. Returns 0.0 if extrapolation is + /// attempted and not allowed. + #[must_use] + #[pyo3(name = "xfxQ2_Chebyshev_batch")] + #[allow(clippy::needless_pass_by_value)] + pub fn xfxq2_cheby_batch(&self, id: i32, params: Vec>) -> Vec { + let slices: Vec<&[f64]> = params.iter().map(Vec::as_slice).collect(); + self.pdf.xfxq2_cheby_batch(id, &slices) + } + /// Interpolates the PDF value (xf) for lists of flavors, x-values, /// and Q2-values. ///