Start refactoring data structure

partons/examples/interp.rs

@@ -0,0 +1,23 @@
+// Load configs and download index file
+use anyhow::Result;
+use partons::configs::Configs;
+
+fn main() -> Result<()> {
+    let cfg = Configs::load()?;
+
+    let mut source = cfg.sources[0].clone();
+    source.register_cache(cfg.data_path()?);
+    let index = source.index()?;
+
+    // display the first element, if non-empty
+    for set in ["NNPDF40_nnlo_as_01180"] {
+        //, "MSHT20nnlo_as118", "CT18NNLO"] {
+        let header = index.get(set)?;
+        let mut set = source.set(&header)?;
+        println!("{set:#?}");
+        let grid0 = set.member(0)?;
+        println!("{grid0}");
+    }
+
+    Ok(())
+}

partons/src/block.rs

@@ -1,68 +1,29 @@
 //! Interpolation block
-use std::collections::HashMap;
-
-use anyhow::{anyhow, Result};
-use ndarray::{Array1, Array3};
+use anyhow::Result;
+use ndarray::{Array1, Array2};
 use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize, Debug)]
-pub struct Block {
-    pub pids: Array1<i32>,
-    pub xgrid: Array1<f64>,
-    pub mu2grid: Array1<f64>,
-    pub(crate) values: Array3<f64>,
-    pid_map: HashMap<i32, usize>,
+pub struct Block1 {
+    pub coords: Array1<f64>,
+    pub(crate) values: Array1<f64>,
 }
 
-impl Block {
-    pub(crate) fn new(
-        pids: Array1<i32>,
-        xgrid: Array1<f64>,
-        mu2grid: Array1<f64>,
-        values: Array3<f64>,
-    ) -> Self {
-        assert_eq!(values.shape(), &[pids.len(), xgrid.len(), mu2grid.len()]);
+#[derive(Serialize, Deserialize, Debug)]
+pub struct Block2 {
+    pub coords: (Array1<f64>, Array1<f64>),
+    pub(crate) values: Array2<f64>,
+}
 
-        let pid_map = pids.iter().enumerate().map(|(i, v)| (*v, i)).collect();
+impl Block2 {
+    pub(crate) fn new(coords: (Array1<f64>, Array1<f64>), values: Array2<f64>) -> Self {
+        assert_eq!(values.shape(), &[coords.0.len(), coords.1.len()]);
 
-        Self {
-            pids,
-            xgrid,
-            mu2grid,
-            values,
-            pid_map,
-        }
-    }
-
-    pub(crate) fn pid_index(&self, pid: i32) -> Result<usize> {
-        self.pid_map
-            .get(&pid)
-            .ok_or(anyhow!("PID '{}' not found", pid))
-            .map(|idx| *idx)
+        Self { coords, values }
     }
 
     // TODO: delegate interpolation to ndinterp
-    pub(crate) fn interp(&self, pid: i32, x: f64, mu2: f64) -> Result<f64> {
-        // determine interpolation slice
-        let slice = self.pid_index(pid)?;
-
-        let idx = self
-            .xgrid
-            .iter()
-            .enumerate()
-            .filter(|(_, y)| y > &&x)
-            .next()
-            .unwrap()
-            .0;
-        let idmu = self
-            .mu2grid
-            .iter()
-            .enumerate()
-            .filter(|(_, nu2)| nu2 > &&mu2)
-            .next()
-            .unwrap()
-            .0;
-
-        return Ok(self.values[[slice, idx, idmu]]);
+    pub(crate) fn interp(&self, x1: f64, x2: f64) -> Result<f64> {
+        return Ok(self.values[[0, 0]]);
     }
 }

partons/src/data/lhapdf/grid.rs

@@ -5,15 +5,15 @@ use std::str::FromStr;
 
 use anyhow::{bail, Result};
 use bytes::Bytes;
-use ndarray::{Array1, Array2, Array3};
+use ndarray::{Array1, Array2};
 
-use crate::block::Block;
+use crate::block::Block2;
 use crate::data::format::ConversionError;
 use crate::member::{Member, Metadata};
 
 pub(crate) struct Grid {
     metadata: Metadata,
-    blocks: Vec<Block>,
+    blocks: Vec<Block2>,
 }
 
 impl Grid {
@@ -25,6 +25,7 @@ impl Grid {
         }
     }
 
+    /// Parse 1 dimensional sequence of values from a space separated string
     fn sequence<T: FromStr>(line: Option<&str>) -> Result<Array1<T>>
     where
         <T as FromStr>::Err: Debug,
@@ -42,6 +43,7 @@ impl Grid {
         }
     }
 
+    /// Parse 2 dimensional table of values from a space and line separated string
     fn table(lines: Option<&str>) -> Result<Array2<f64>> {
         if let Some(text) = lines {
             let mut rows = Vec::new();
@@ -63,19 +65,26 @@ impl Grid {
         }
     }
 
-    fn values(values: Array2<f64>, xs: usize) -> Result<Array3<f64>> {
-        let &[points, pids] = values.shape() else { bail!("") };
+    /// Split the table of values in a sequence of 2 dimensional array, and lift the dimension
+    /// related to the discrete quantity (PID) as the outermost.
+    ///
+    /// This is required in partons, because the innermost indices are dedicated to the
+    /// interpolated dimensions.
+    fn values(values: Array2<f64>, xs: usize) -> Result<Vec<Array2<f64>>> {
+        let &[points, pids] = values.shape() else {
+            bail!("")
+        };
         let mu2s = points / xs;
         let lha_shaped = values.into_shape((xs, mu2s, pids))?;
         // TODO: solve the following horrible memory duplication, and horrible loop
-        let mut native_shaped = Array3::zeros((pids, xs, mu2s));
+        let mut native_shaped = vec![Array2::zeros((xs, mu2s)); pids];
         for ((x, mu, p), val) in lha_shaped.indexed_iter() {
-            native_shaped[(p, x, mu)] = *val;
+            native_shaped[p][(x, mu)] = *val;
         }
         Ok(native_shaped)
     }
 
-    fn block(section: &str) -> Result<Block> {
+    fn block(section: &str) -> Result<(Vec<Block2>, Array1<i32>)> {
         let mut split = section.trim().splitn(4, '\n');
 
         let xgrid = Self::sequence(split.next())?;
@@ -84,7 +93,13 @@ impl Grid {
 
         let values = Self::values(Self::table(split.next())?, xgrid.len())?;
 
-        Ok(Block::new(pids, xgrid, mu2grid, values))
+        Ok((
+            values
+                .into_iter()
+                .map(|ar| Block2::new((xgrid.clone(), mu2grid.clone()), ar))
+                .collect(),
+            pids,
+        ))
     }
 
     pub(crate) fn load(bytes: Bytes) -> Result<Self> {
@@ -94,8 +109,11 @@ impl Grid {
         let metadata = Self::metadata(sections.next())?;
 
         let mut blocks = Vec::new();
+        let mut pids = Vec::new();
         for section in sections {
-            blocks.push(Self::block(section)?);
+            let (block, pids_) = Self::block(section)?;
+            blocks.extend(block);
+            pids.extend(pids_);
         }
 
         Ok(Self { metadata, blocks })

partons/src/member.rs

@@ -9,18 +9,18 @@ use bytes::Bytes;
 use ndarray::Array1;
 use serde::{Deserialize, Serialize};
 
-use crate::block::Block;
+use crate::block::Block2;
 
 pub(crate) type Metadata = HashMap<String, String>;
 
 /// Member of a set
 ///
 /// This contains the whole member data, including the interpolation
-/// [`Block`](crate::block::Block)s and further optional metadata.
+/// [blocks](crate::block::Block2) and further optional metadata.
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Member {
     pub(crate) metadata: Metadata,
-    pub(crate) blocks: Vec<Block>,
+    pub(crate) blocks: Vec<Block2>,
 }
 
 #[derive(Decode, Encode)]
@@ -49,7 +49,7 @@ impl Member {
         }
 
         let mut values: Array1<f64> = Array1::zeros(x.raw_dim());
-        values[0] = self.blocks[(nf[0] - 3) as usize].interp(pid[0], x[0], mu2[0])?;
+        values[0] = self.blocks[(nf[0] - 3) as usize].interp(x[0], mu2[0])?;
 
         Ok(values)
     }