Multi-grids loader with GridArrayCollection

neopdf/src/converter.rs

@@ -29,7 +29,7 @@ pub fn convert_lhapdf<P: AsRef<std::path::Path>>(
     output_path: P,
 ) -> Result<(), Box<dyn std::error::Error>> {
     let lhapdf_set = LhapdfSet::new(pdf_name);
-    let members = lhapdf_set.members();
+    let members: Vec<_> = lhapdf_set.members().collect();
     if members.is_empty() {
         return Err("No members found in the LHAPDF set".into());
     }
@@ -83,7 +83,7 @@ pub fn combine_lhapdf_npdfs<P: AsRef<std::path::Path>>(
         };
         a_values.push(a);
         let set = LhapdfSet::new(pdf_name);
-        let members = set.members();
+        let members: Vec<_> = set.members().collect();
         if members.is_empty() {
             return Err(format!("No members found in set: {}", pdf_name).into());
         }
@@ -200,7 +200,7 @@ pub fn combine_lhapdf_alphas<P: AsRef<std::path::Path>>(
         };
         alphas_values.push(alphas);
         let set = LhapdfSet::new(pdf_name);
-        let members = set.members();
+        let members: Vec<_> = set.members().collect();
         if members.is_empty() {
             return Err(format!("No members found in set: {}", pdf_name).into());
         }

neopdf/src/parser.rs

@@ -4,12 +4,12 @@
 //! set formats, including subgrid data extraction and metadata reading.
 use serde::{Deserialize, Serialize};
 use std::fs;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 use super::gridpdf::GridArray;
 use super::manage::{ManageData, PdfSetFormat};
 use super::metadata::MetaData;
-use super::writer::GridArrayReader;
+use super::writer::{GridArrayCollection, GridArrayReader, LazyGridArrayIterator};
 
 /// Represents the data for a single subgrid within a PDF data file.
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -85,18 +85,20 @@ impl LhapdfSet {
     ///
     /// # Returns
     ///
-    /// A vector of tuples, where each tuple contains the `MetaData` and `PdfData`
+    /// An iterator over tuples, where each tuple contains the `MetaData` and `GridArray`
     /// for a member of the set.
-    pub fn members(&self) -> Vec<(MetaData, GridArray)> {
-        (0..self.info.num_members as usize)
-            .map(|i| self.member(i))
-            .collect()
+    pub fn members(&self) -> impl Iterator<Item = (MetaData, GridArray)> + Send + '_ {
+        (0..self.info.num_members as usize).map(move |i| self.member(i))
+    }
+
+    /// Consumes the set and returns a lazy iterator over its members.
+    pub fn into_lazy_members(self) -> impl Iterator<Item = (MetaData, GridArray)> + Send {
+        (0..self.info.num_members as usize).map(move |i| self.member(i))
     }
 
     /// Reads the `.info` file for a PDF set and deserializes it into an `Info` struct.
     ///
     /// # Arguments
-    ///
     /// * `path` - The path to the `.info` file.
     ///
     /// # Returns
@@ -200,6 +202,7 @@ impl LhapdfSet {
 pub struct NeopdfSet {
     pub info: MetaData,
     grid_reader: GridArrayReader,
+    neopdf_setpath: PathBuf,
 }
 
 impl NeopdfSet {
@@ -210,14 +213,35 @@ impl NeopdfSet {
         let grid_reader = GridArrayReader::from_file(neopdf_setpath).unwrap();
         let info = grid_reader.metadata().as_ref().clone();
 
-        Self { info, grid_reader }
+        Self {
+            info,
+            grid_reader,
+            neopdf_setpath: neopdf_setpath.to_path_buf(),
+        }
     }
 
     /// TODO
     pub fn member(&self, member: usize) -> (MetaData, GridArray) {
         let load_grid = self.grid_reader.load_grid(member).unwrap();
         (self.info.clone(), load_grid.grid)
     }
+
+    /// TODO
+    pub fn members(&self) -> impl Iterator<Item = (MetaData, GridArray)> {
+        GridArrayCollection::decompress(&self.neopdf_setpath)
+            .unwrap()
+            .into_iter()
+            .map(|grid| (grid.metadata.as_ref().clone(), grid.grid))
+    }
+
+    /// TODO
+    pub fn into_lazy_members(
+        self,
+    ) -> impl Iterator<Item = Result<(MetaData, GridArray), Box<dyn std::error::Error>>> {
+        LazyGridArrayIterator::from_file(&self.neopdf_setpath)
+            .unwrap()
+            .map(|grid_result| grid_result.map(|grid| (grid.metadata.as_ref().clone(), grid.grid)))
+    }
 }
 
 #[cfg(test)]

neopdf/src/pdf.rs

@@ -20,40 +20,72 @@
 //!
 //! See the documentation for [`PDF`] for more details on available methods and usage patterns.
 use ndarray::{Array1, Array2};
-use rayon::prelude::*;
 
 use super::gridpdf::{ForcePositive, GridArray, GridPDF};
 use super::metadata::MetaData;
 use super::parser::{LhapdfSet, NeopdfSet};
 use super::subgrid::{RangeParameters, SubGrid};
 
+#[derive(Debug)]
+pub struct SendableError(Box<dyn std::error::Error>);
+
+impl std::fmt::Display for SendableError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+impl std::error::Error for SendableError {}
+
+unsafe impl Send for SendableError {}
+unsafe impl Sync for SendableError {}
+
+type Type = Result<(MetaData, GridArray), SendableError>;
+
 /// Trait for abstracting over different PDF set backends (e.g., LHAPDF, NeoPDF).
 ///
 /// Provides a unified interface for accessing the number of members and retrieving individual
 /// members as metadata and grid arrays.
 trait PdfSet: Send + Sync {
-    /// Returns the number of members in the PDF set.
-    fn num_members(&self) -> usize;
     /// Retrieves the metadata and grid array for the specified member index.
     fn member(&self, idx: usize) -> (MetaData, GridArray);
+
+    /// Retrieves the metadata and grid arrays for all members.
+    fn members(&self) -> Box<dyn Iterator<Item = (MetaData, GridArray)> + Send + '_>;
+
+    /// Consumes the set and returns a lazy iterator over its members.
+    fn into_lazy_members(self) -> Box<dyn Iterator<Item = Type> + Send>;
 }
 
 impl PdfSet for LhapdfSet {
-    fn num_members(&self) -> usize {
-        self.info.num_members as usize
-    }
     fn member(&self, idx: usize) -> (MetaData, GridArray) {
         self.member(idx)
     }
+
+    fn members(&self) -> Box<dyn Iterator<Item = (MetaData, GridArray)> + Send + '_> {
+        Box::new(self.members())
+    }
+
+    fn into_lazy_members(self) -> Box<dyn Iterator<Item = Type> + Send> {
+        Box::new(self.into_lazy_members().map(Ok))
+    }
 }
 
 impl PdfSet for NeopdfSet {
-    fn num_members(&self) -> usize {
-        self.info.num_members as usize
-    }
     fn member(&self, idx: usize) -> (MetaData, GridArray) {
         self.member(idx)
     }
+
+    fn members(&self) -> Box<dyn Iterator<Item = (MetaData, GridArray)> + Send + '_> {
+        Box::new(self.members())
+    }
+
+    fn into_lazy_members(self) -> Box<dyn Iterator<Item = Type> + Send> {
+        Box::new(
+            self.into_lazy_members()
+                .map(|res| res.map_err(SendableError)),
+        )
+    }
 }
 
 /// Loads a single PDF member from a generic PDF set backend.
@@ -83,17 +115,32 @@ fn pdfset_loader<T: PdfSet>(set: T, member: usize) -> PDF {
 ///
 /// A vector of [`PDF`] instances, one for each member in the set.
 fn pdfsets_loader<T: PdfSet + Send + Sync>(set: T) -> Vec<PDF> {
-    (0..set.num_members())
-        .into_par_iter()
-        .map(|idx| {
-            let (info, knot_array) = set.member(idx);
-            PDF {
-                grid_pdf: GridPDF::new(info, knot_array),
-            }
+    set.members()
+        .map(|(info, knot_array)| PDF {
+            grid_pdf: GridPDF::new(info, knot_array),
         })
         .collect()
 }
 
+/// Lazily loads all PDF members from a generic PDF set backend.
+///
+/// # Arguments
+///
+/// * `set` - The PDF set backend implementing [`PdfSet`].
+///
+/// # Returns
+///
+/// An iterator over [`PDF`] instances, one for each member in the set.
+fn pdfsets_loader_lazy<T: PdfSet + Send + Sync + 'static>(
+    set: T,
+) -> Box<dyn Iterator<Item = Result<PDF, SendableError>> + Send> {
+    Box::new(set.into_lazy_members().map(|member_result| {
+        member_result.map(|(info, knot_array)| PDF {
+            grid_pdf: GridPDF::new(info, knot_array),
+        })
+    }))
+}
+
 /// Represents a Parton Distribution Function (PDF) set.
 ///
 /// This struct provides a high-level interface for accessing PDF data,
@@ -146,6 +193,29 @@ impl PDF {
         }
     }
 
+    /// Lazily loads all members of a PDF set.
+    ///
+    /// This function reads the `.info` file and all `.dat` member files
+    /// to construct a `Vec<PDF>`, with each `PDF` instance representing a member
+    /// of the set. The loading is performed in parallel.
+    ///
+    /// # Arguments
+    ///
+    /// * `pdf_name` - The name of the PDF set.
+    ///
+    /// # Returns
+    ///
+    /// A `Vec<PDF>` where each element is a `PDF` instance for a member of the set.
+    pub fn load_pdfs_lazy(
+        pdf_name: &str,
+    ) -> Box<dyn Iterator<Item = Result<PDF, SendableError>> + Send> {
+        if pdf_name.ends_with(".neopdf.lz4") {
+            pdfsets_loader_lazy(NeopdfSet::new(pdf_name))
+        } else {
+            pdfsets_loader_lazy(LhapdfSet::new(pdf_name))
+        }
+    }
+
     /// Clip the negative values for the `PDF` object.
     ///
     /// # Arguments

neopdf_cli/src/converter.rs

@@ -1,5 +1,5 @@
 //! CLI logic for `NeoPDF` conversion utilities.
-//_!
+//!
 //! This module defines the command-line interface for converting LHAPDF sets to `NeoPDF` format
 //! and for combining multiple nuclear PDFs into a single `NeoPDF` file with explicit A dependence.
 
@@ -120,6 +120,10 @@ fn load_pdf_names(
 ///
 /// * `cli` - The parsed command line interface structure containing the command and its arguments
 ///
+/// # Errors
+///
+/// TODO
+///
 /// # Returns
 ///
 /// * `Ok(())` - If the command executed successfully

neopdf_pyapi/src/parser.rs

@@ -43,7 +43,7 @@ impl PyLhapdfSet {
     #[must_use]
     pub fn members(&self) -> Vec<(PyMetaData, PyGridArray)> {
         // TODO: Use the parallelized `members` in the crate
-        let num_members = self.inner.members().len();
+        let num_members = self.inner.members().collect::<Vec<_>>().len();
         (0..num_members).map(|idx| self.member(idx)).collect()
     }
 }

neopdf_pyapi/src/pdf.rs

@@ -1,8 +1,9 @@
 use numpy::{IntoPyArray, PyArray2};
+use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 
 use neopdf::gridpdf::ForcePositive;
-use neopdf::pdf::PDF;
+use neopdf::pdf::{SendableError, PDF};
 
 use super::gridpdf::PySubGrid;
 use super::metadata::PyMetaData;
@@ -70,6 +71,30 @@ pub enum PyGridParams {
     Q2,
 }
 
+/// An iterator over `PDF` objects.
+#[pyclass(name = "PDFIterator")]
+pub struct PyPDFIterator {
+    iter: Box<dyn Iterator<Item = Result<PDF, SendableError>> + Send>,
+}
+
+unsafe impl Send for PyPDFIterator {}
+unsafe impl Sync for PyPDFIterator {}
+
+#[pymethods]
+impl PyPDFIterator {
+    const fn __iter__(slf: PyRef<Self>) -> PyRef<Self> {
+        slf
+    }
+
+    fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<PyPDF>> {
+        match slf.iter.next() {
+            Some(Ok(pdf)) => Ok(Some(PyPDF { pdf })),
+            Some(Err(e)) => Err(PyRuntimeError::new_err(format!("PDF loading error: {e}"))),
+            None => Ok(None),
+        }
+    }
+}
+
 /// Python wrapper for the `neopdf::pdf::PDF` struct.
 ///
 /// This class provides a Python-friendly interface to the core PDF
@@ -154,6 +179,30 @@ impl PyPDF {
             .collect()
     }
 
+    /// Loads all members of the PDF set lazily.
+    ///
+    /// This function returns an iterator that loads each member of a given
+    /// PDF set lazily.
+    ///
+    /// Parameters
+    /// ----------
+    /// pdf_name : str
+    ///     The name of the PDF set.
+    ///
+    /// Returns
+    /// -------
+    /// PDFIterator
+    ///     An iterator over `PDF` instances, one for each member.
+    #[must_use]
+    #[staticmethod]
+    #[pyo3(name = "mkPDFs_lazy")]
+    pub fn mkpdfs_lazy(pdf_name: &str) -> PyPDFIterator {
+        let iter = PDF::load_pdfs_lazy(pdf_name);
+        PyPDFIterator {
+            iter: Box::new(iter),
+        }
+    }
+
     /// Returns the list of `PID` values.
     ///
     /// Returns
@@ -435,6 +484,7 @@ pub fn register(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
         "import sys; sys.modules['neopdf.pdf'] = m"
     );
     m.add_class::<PyPDF>()?;
+    m.add_class::<PyPDFIterator>()?;
     m.add_class::<PyForcePositive>()?;
     m.add_class::<PyGridParams>()?;
     parent_module.add_submodule(&m)