Add generalized constructor for ConvolutionCache

NNPDF · Oct 4, 2024 · c196936 · c196936
1 parent bf02084
commit c196936
Show file tree

Hide file tree

Showing 7 changed files with 30 additions and 128 deletions.
diff --git a/pineappl/src/convolutions.rs b/pineappl/src/convolutions.rs
@@ -25,58 +25,24 @@ pub struct ConvolutionCache<'a> {
 }
 
 impl<'a> ConvolutionCache<'a> {
-    /// Construct a luminosity cache with two PDFs, `xfx1` and `xfx2`. The types of hadrons the
-    /// PDFs correspond to must be given as `pdg1` and `pdg2`. The function to evaluate the
-    /// strong coupling must be given as `alphas`. The grid that the cache will be used with must
-    /// be given as `grid`; this parameter determines which of the initial states are hadronic, and
-    /// if an initial states is not hadronic the corresponding 'PDF' is set to `xfx = x`. If some
-    /// of the PDFs must be charge-conjugated, this is automatically done in this function.
-    pub fn with_two(
-        pdg1: i32,
-        xfx1: &'a mut dyn FnMut(i32, f64, f64) -> f64,
-        pdg2: i32,
-        xfx2: &'a mut dyn FnMut(i32, f64, f64) -> f64,
-        alphas: &'a mut dyn FnMut(f64) -> f64,
-    ) -> Self {
-        Self {
-            xfx: vec![xfx1, xfx2],
-            xfx_cache: vec![FxHashMap::default(); 2],
-            alphas,
-            alphas_cache: vec![],
-            mur2_grid: vec![],
-            muf2_grid: vec![],
-            x_grid: vec![],
-            imur2: Vec::new(),
-            imuf2: Vec::new(),
-            ix: Vec::new(),
-            pdg: vec![pdg1, pdg2],
-            perm: Vec::new(),
-        }
-    }
-
-    /// Construct a luminosity cache with a single PDF `xfx`. The type of hadron the PDF
-    /// corresponds to must be given as `pdg`. The function to evaluate the strong coupling must be
-    /// given as `alphas`. The grid that the cache should be used with must be given as `grid`;
-    /// this parameter determines which of the initial states are hadronic, and if an initial
-    /// states is not hadronic the corresponding 'PDF' is set to `xfx = x`. If some of the PDFs
-    /// must be charge-conjugated, this is automatically done in this function.
-    pub fn with_one(
-        pdg: i32,
-        xfx: &'a mut dyn FnMut(i32, f64, f64) -> f64,
+    /// TODO
+    pub fn new(
+        pdg: Vec<i32>,
+        xfx: Vec<&'a mut dyn FnMut(i32, f64, f64) -> f64>,
         alphas: &'a mut dyn FnMut(f64) -> f64,
     ) -> Self {
         Self {
-            xfx: vec![xfx],
-            xfx_cache: vec![FxHashMap::default()],
+            xfx_cache: vec![FxHashMap::default(); xfx.len()],
+            xfx,
             alphas,
-            alphas_cache: vec![],
-            mur2_grid: vec![],
-            muf2_grid: vec![],
-            x_grid: vec![],
+            alphas_cache: Vec::new(),
+            mur2_grid: Vec::new(),
+            muf2_grid: Vec::new(),
+            x_grid: Vec::new(),
             imur2: Vec::new(),
             imuf2: Vec::new(),
             ix: Vec::new(),
-            pdg: vec![pdg],
+            pdg,
             perm: Vec::new(),
         }
     }

diff --git a/pineappl/tests/drell_yan_lo.rs b/pineappl/tests/drell_yan_lo.rs
@@ -351,7 +351,7 @@ fn perform_grid_tests(
     grid.scale_by_order(10.0, 1.0, 10.0, 10.0, 4.0);
 
     // TEST 5: `convolve`
-    let mut convolution_cache = ConvolutionCache::with_one(2212, &mut xfx, &mut alphas);
+    let mut convolution_cache = ConvolutionCache::new(vec![2212], vec![&mut xfx], &mut alphas);
     let bins = grid.convolve(&mut convolution_cache, &[], &[], &[], &[(1.0, 1.0, 1.0)]);
 
     for (result, reference) in bins.iter().zip(reference.iter()) {
@@ -363,7 +363,7 @@ fn perform_grid_tests(
     let mut xfx2 = |id, x, q2| pdf.xfx_q2(id, x, q2);
     let mut alphas2 = |_| 0.0;
     let mut convolution_cache2 =
-        ConvolutionCache::with_two(2212, &mut xfx1, 2212, &mut xfx2, &mut alphas2);
+        ConvolutionCache::new(vec![2212, 2212], vec![&mut xfx1, &mut xfx2], &mut alphas2);
     let bins2 = grid.convolve(&mut convolution_cache2, &[], &[], &[], &[(1.0, 1.0, 1.0)]);
 
     for (result, reference) in bins2.iter().zip(reference.iter()) {

diff --git a/pineappl_capi/src/lib.rs b/pineappl_capi/src/lib.rs
@@ -442,7 +442,7 @@ pub unsafe extern "C" fn pineappl_grid_convolve_with_one(
     results: *mut f64,
 ) {
     let grid = unsafe { &*grid };
-    let mut pdf = |id, x, q2| xfx(id, x, q2, state);
+    let mut xfx = |id, x, q2| xfx(id, x, q2, state);
     let mut als = |q2| alphas(q2, state);
     let order_mask = if order_mask.is_null() {
         vec![]
@@ -455,7 +455,7 @@ pub unsafe extern "C" fn pineappl_grid_convolve_with_one(
         unsafe { slice::from_raw_parts(channel_mask, grid.channels().len()) }.to_vec()
     };
     let results = unsafe { slice::from_raw_parts_mut(results, grid.bin_info().bins()) };
-    let mut convolution_cache = ConvolutionCache::with_one(pdg_id, &mut pdf, &mut als);
+    let mut convolution_cache = ConvolutionCache::new(vec![pdg_id], vec![&mut xfx], &mut als);
 
     results.copy_from_slice(&grid.convolve(
         &mut convolution_cache,
@@ -503,8 +503,8 @@ pub unsafe extern "C" fn pineappl_grid_convolve_with_two(
     results: *mut f64,
 ) {
     let grid = unsafe { &*grid };
-    let mut pdf1 = |id, x, q2| xfx1(id, x, q2, state);
-    let mut pdf2 = |id, x, q2| xfx2(id, x, q2, state);
+    let mut xfx1 = |id, x, q2| xfx1(id, x, q2, state);
+    let mut xfx2 = |id, x, q2| xfx2(id, x, q2, state);
     let mut als = |q2| alphas(q2, state);
     let order_mask = if order_mask.is_null() {
         vec![]
@@ -518,7 +518,7 @@ pub unsafe extern "C" fn pineappl_grid_convolve_with_two(
     };
     let results = unsafe { slice::from_raw_parts_mut(results, grid.bin_info().bins()) };
     let mut convolution_cache =
-        ConvolutionCache::with_two(pdg_id1, &mut pdf1, pdg_id2, &mut pdf2, &mut als);
+        ConvolutionCache::new(vec![pdg_id1, pdg_id2], vec![&mut xfx1, &mut xfx2], &mut als);
 
     results.copy_from_slice(&grid.convolve(
         &mut convolution_cache,

diff --git a/pineappl_cli/src/helpers.rs b/pineappl_cli/src/helpers.rs
@@ -271,6 +271,10 @@ pub fn convolve_scales(
             }
         })
         .collect();
+    let xfx: Vec<_> = funs
+        .iter_mut()
+        .map(|fun| fun as &mut dyn FnMut(i32, f64, f64) -> f64)
+        .collect();
     let mut alphas_funs: Vec<_> = conv_funs
         .iter()
         .map(|fun| move |q2| fun.alphas_q2(q2))
@@ -287,25 +291,7 @@ pub fn convolve_scales(
         })
         .collect();
 
-    // TODO: write a new constructor of `LumiCache` that accepts a vector of all the arguments
-    let mut cache = match funs.as_mut_slice() {
-        [funs0] => {
-            ConvolutionCache::with_one(pdg_ids[0], funs0, &mut alphas_funs[cfg.use_alphas_from])
-        }
-        [funs0, funs1] => ConvolutionCache::with_two(
-            pdg_ids[0],
-            funs0,
-            pdg_ids[1],
-            funs1,
-            &mut alphas_funs[cfg.use_alphas_from],
-        ),
-        // TODO: convert this into an error
-        _ => panic!(
-            "convolutions with {} convolution functions is not supported",
-            conv_funs.len()
-        ),
-    };
-
+    let mut cache = ConvolutionCache::new(pdg_ids, xfx, &mut alphas_funs[cfg.use_alphas_from]);
     let mut results = grid.convolve(&mut cache, &orders, bins, channels, scales);
 
     match mode {
@@ -436,6 +422,10 @@ pub fn convolve_subgrid(
             }
         })
         .collect();
+    let xfx: Vec<_> = funs
+        .iter_mut()
+        .map(|fun| fun as &mut dyn FnMut(i32, f64, f64) -> f64)
+        .collect();
     let mut alphas_funs: Vec<_> = conv_funs
         .iter()
         .map(|fun| move |q2| fun.alphas_q2(q2))
@@ -452,26 +442,9 @@ pub fn convolve_subgrid(
         })
         .collect();
 
-    // TODO: write a new constructor of `LumiCache` that accepts a vector of all the arguments
-    let mut cache = match funs.as_mut_slice() {
-        [funs0] => {
-            ConvolutionCache::with_one(pdg_ids[0], funs0, &mut alphas_funs[cfg.use_alphas_from])
-        }
-        [funs0, funs1] => ConvolutionCache::with_two(
-            pdg_ids[0],
-            funs0,
-            pdg_ids[1],
-            funs1,
-            &mut alphas_funs[cfg.use_alphas_from],
-        ),
-        // TODO: convert this into an error
-        _ => panic!(
-            "convolutions with {} convolution functions is not supported",
-            conv_funs.len()
-        ),
-    };
-
+    let mut cache = ConvolutionCache::new(pdg_ids, xfx, &mut alphas_funs[cfg.use_alphas_from]);
     let subgrid = grid.convolve_subgrid(&mut cache, order, bin, lumi, (1.0, 1.0, 1.0));
+
     subgrid
         .into_dimensionality::<Ix3>()
         .map_err(|_| anyhow!("Only 3-dimensional subgrids are supported",))

diff --git a/pineappl_cli/tests/convolve.rs b/pineappl_cli/tests/convolve.rs
@@ -34,9 +34,6 @@ const DEFAULT_STR: &str = "b   etal    dsig/detal
 const USE_ALPHAS_FROM_ERROR_STR: &str = "expected `use_alphas_from` to be `0` or `1`, is `2`
 ";
 
-const THREE_PDF_ERROR_STR: &str = "convolutions with 3 convolution functions is not supported
-";
-
 const FORCE_POSITIVE_STR: &str = "b   etal    dsig/detal 
      []        [pb]    
 -+----+----+-----------
@@ -211,20 +208,6 @@ fn use_alphas_from_error() {
         .stderr(str::contains(USE_ALPHAS_FROM_ERROR_STR));
 }
 
-#[test]
-fn three_pdf_error() {
-    Command::cargo_bin("pineappl")
-        .unwrap()
-        .args([
-            "convolve",
-            "../test-data/LHCB_WP_7TEV_opt.pineappl.lz4",
-            "NNPDF31_nlo_as_0118_luxqed,NNPDF31_nlo_as_0118_luxqed,NNPDF31_nlo_as_0118_luxqed",
-        ])
-        .assert()
-        .failure()
-        .stderr(str::contains(THREE_PDF_ERROR_STR));
-}
-
 #[test]
 fn force_positive() {
     Command::cargo_bin("pineappl")

diff --git a/pineappl_cli/tests/import.rs b/pineappl_cli/tests/import.rs
@@ -893,7 +893,7 @@ fn import_hadronic_fktable() {
     let pdf = Pdf::with_setname_and_member("NNPDF31_nlo_as_0118_luxqed", 0).unwrap();
     let mut xfx = |id, x, q2| pdf.xfx_q2(id, x, q2);
     let mut alphas = |_| 0.0;
-    let mut convolution_cache = ConvolutionCache::with_one(2212, &mut xfx, &mut alphas);
+    let mut convolution_cache = ConvolutionCache::new(vec![2212], vec![&mut xfx], &mut alphas);
     let results = grid.convolve(&mut convolution_cache, &[], &[], &[], &[(1.0, 1.0, 1.0)]);
 
     let mut fk_table = FkTable::try_from(grid).unwrap();

diff --git a/pineappl_cli/tests/plot.rs b/pineappl_cli/tests/plot.rs
@@ -1416,9 +1416,6 @@ if __name__ == "__main__":
     main()
 "#;
 
-const THREE_PDF_ERROR_STR: &str = "convolutions with 3 convolution functions is not supported
-";
-
 #[test]
 fn help() {
     Command::cargo_bin("pineappl")
@@ -1469,23 +1466,6 @@ fn subgrid_pull() {
         .stdout(SUBGRID_PULL_STR);
 }
 
-#[test]
-fn three_pdf_error() {
-    Command::cargo_bin("pineappl")
-        .unwrap()
-        .args([
-            "plot",
-            "--subgrid-pull=0,0,0",
-            "--threads=1",
-            "../test-data/LHCB_WP_7TEV_opt.pineappl.lz4",
-            "NNPDF31_nlo_as_0118_luxqed",
-            "NNPDF40_nnlo_as_01180,NNPDF40_nnlo_as_01180,NNPDF40_nnlo_as_01180",
-        ])
-        .assert()
-        .failure()
-        .stderr(str::contains(THREE_PDF_ERROR_STR));
-}
-
 #[test]
 fn drell_yan_afb() {
     Command::cargo_bin("pineappl")