Accommodate for two different EKOs in Evolution

pineappl/src/evolution.rs

@@ -524,21 +524,22 @@ pub(crate) fn evolve_slice_with_one(
 
 pub(crate) fn evolve_slice_with_two(
     grid: &Grid,
-    operator: &ArrayView4<f64>,
+    operator_a: &ArrayView4<f64>,
+    operator_b: &ArrayView4<f64>,
     info: &OperatorSliceInfo,
     order_mask: &[bool],
     xi: (f64, f64),
     alphas_table: &AlphasTable,
 ) -> Result<(Array3<SubgridEnum>, Vec<LumiEntry>), GridError> {
     let gluon_has_pid_zero = gluon_has_pid_zero(grid);
 
-    let (pid_indices_a, pids_a) = pid_slices(operator, info, gluon_has_pid_zero, &|pid1| {
+    let (pid_indices_a, pids_a) = pid_slices(operator_a, info, gluon_has_pid_zero, &|pid1| {
         grid.lumi()
             .iter()
             .flat_map(LumiEntry::entry)
             .any(|&(a, _, _)| a == pid1)
     })?;
-    let (pid_indices_b, pids_b) = pid_slices(operator, info, gluon_has_pid_zero, &|pid1| {
+    let (pid_indices_b, pids_b) = pid_slices(operator_b, info, gluon_has_pid_zero, &|pid1| {
         grid.lumi()
             .iter()
             .flat_map(LumiEntry::entry)
@@ -572,7 +573,7 @@ pub(crate) fn evolve_slice_with_two(
                     .zip(x1_a.iter())
                     .any(|(&lhs, &rhs)| !approx_eq!(f64, lhs, rhs, ulps = EVOLUTION_TOL_ULPS))
             {
-                operators_a = operator_slices(operator, info, &pid_indices_a, &x1_a)?;
+                operators_a = operator_slices(operator_a, info, &pid_indices_a, &x1_a)?;
                 last_x1a = x1_a;
             }
 
@@ -582,7 +583,7 @@ pub(crate) fn evolve_slice_with_two(
                     .zip(x1_b.iter())
                     .any(|(&lhs, &rhs)| !approx_eq!(f64, lhs, rhs, ulps = EVOLUTION_TOL_ULPS))
             {
-                operators_b = operator_slices(operator, info, &pid_indices_b, &x1_b)?;
+                operators_b = operator_slices(operator_b, info, &pid_indices_b, &x1_b)?;
                 last_x1b = x1_b;
             }
 

pineappl/src/grid.rs

@@ -1384,6 +1384,8 @@ impl Grid {
         }
     }
 
+    /// TODO: Fix to account for the different EKOs
+    /// TODO: Does the info have to be different?
     /// Converts this `Grid` into an [`FkTable`] using an evolution kernel operator (EKO) given as
     /// `operator`. The dimensions and properties of this operator must be described using `info`.
     /// The parameter `order_mask` can be used to include or exclude orders from this operation,
@@ -1397,14 +1399,32 @@ impl Grid {
     #[deprecated(since = "0.7.4", note = "use evolve_with_slice_iter instead")]
     pub fn evolve(
         &self,
-        operator: ArrayView5<f64>,
+        operator_a: ArrayView5<f64>,
+        operator_b: ArrayView5<f64>,
         info: &OperatorInfo,
         order_mask: &[bool],
     ) -> Result<FkTable, GridError> {
         self.evolve_with_slice_iter(
             info.fac1
                 .iter()
-                .zip(operator.axis_iter(Axis(0)))
+                .zip(operator_a.axis_iter(Axis(0)))
+                .map(|(&fac1, op)| {
+                    Ok::<_, GridError>((
+                        OperatorSliceInfo {
+                            fac0: info.fac0,
+                            pids0: info.pids0.clone(),
+                            x0: info.x0.clone(),
+                            fac1,
+                            pids1: info.pids1.clone(),
+                            x1: info.x1.clone(),
+                            lumi_id_types: info.lumi_id_types.clone(),
+                        },
+                        CowArray::from(op),
+                    ))
+                }),
+            info.fac1
+                .iter()
+                .zip(operator_b.axis_iter(Axis(0)))
                 .map(|(&fac1, op)| {
                     Ok::<_, GridError>((
                         OperatorSliceInfo {
@@ -1446,16 +1466,21 @@ impl Grid {
     pub fn evolve_with_slice_iter<'a, E: Into<anyhow::Error>>(
         &self,
         slices: impl IntoIterator<Item = Result<(OperatorSliceInfo, CowArray<'a, f64, Ix4>), E>>,
+        extra_slices: impl IntoIterator<Item = Result<(OperatorSliceInfo, CowArray<'a, f64, Ix4>), E>>,
         order_mask: &[bool],
         xi: (f64, f64),
         alphas_table: &AlphasTable,
     ) -> Result<FkTable, GridError> {
         use super::evolution::EVOLVE_INFO_TOL_ULPS;
+        use itertools::izip;
 
         let mut lhs: Option<Self> = None;
         let mut fac1 = Vec::new();
 
-        for result in slices {
+        // TODO: simplify the ugly repetition below
+        // TODO: what to do about the info? Shoulb they be the same?
+        for (result, extra_result) in izip!(slices, extra_slices) {
+            // Deal with the main EKO
             let (info, operator) = result.map_err(|err| GridError::Other(err.into()))?;
 
             let op_info_dim = (
@@ -1475,10 +1500,39 @@ impl Grid {
 
             let view = operator.view();
 
+            // Deal with the additional EKO
+            let (extra_info, extra_operator) =
+                extra_result.map_err(|err| GridError::Other(err.into()))?;
+
+            let op_info_dim_extra = (
+                extra_info.pids1.len(),
+                extra_info.x1.len(),
+                extra_info.pids0.len(),
+                extra_info.x0.len(),
+            );
+
+            if extra_operator.dim() != op_info_dim_extra {
+                return Err(GridError::EvolutionFailure(format!(
+                    "operator information {:?} does not match the operator's dimensions: {:?}",
+                    op_info_dim_extra,
+                    extra_operator.dim(),
+                )));
+            }
+
+            let extra_view = extra_operator.view();
+
             let (subgrids, lumi) = if self.convolutions()[0] != Convolution::None
                 && self.convolutions()[1] != Convolution::None
             {
-                evolution::evolve_slice_with_two(self, &view, &info, order_mask, xi, alphas_table)
+                evolution::evolve_slice_with_two(
+                    self,
+                    &view,
+                    &extra_view,
+                    &info,
+                    order_mask,
+                    xi,
+                    alphas_table,
+                )
             } else {
                 evolution::evolve_slice_with_one(self, &view, &info, order_mask, xi, alphas_table)
             }?;

pineappl_cli/src/evolve.rs

@@ -425,6 +425,7 @@ mod eko {
 fn evolve_grid(
     grid: &Grid,
     eko: &Path,
+    extra_eko: &Path,
     pdf: &Pdf,
     orders: &[(u32, u32)],
     xir: f64,
@@ -446,8 +447,10 @@ fn evolve_grid(
         .collect();
 
     let mut eko_slices = EkoSlices::new(eko)?;
+    let mut extra_eko_slices = EkoSlices::new(extra_eko)?;
     let alphas_table = AlphasTable::from_grid(grid, xir, &|q2| pdf.alphas_q2(q2));
 
+    // TODO: Take care of the old EKO behavior
     if use_old_evolve {
         if let EkoSlices::V0 {
             fac1,
@@ -475,14 +478,21 @@ fn evolve_grid(
             unimplemented!();
         }
     } else {
-        Ok(grid.evolve_with_slice_iter(&mut eko_slices, &order_mask, (xir, xif), &alphas_table)?)
+        Ok(grid.evolve_with_slice_iter(
+            &mut eko_slices,
+            extra_eko_slices,
+            &order_mask,
+            (xir, xif),
+            &alphas_table,
+        )?)
     }
 }
 
 #[cfg(not(feature = "evolve"))]
 fn evolve_grid(
     _: &Grid,
     _: &Path,
+    _: &Path,
     _: &Pdf,
     _: &[(u32, u32)],
     _: f64,
@@ -554,9 +564,11 @@ impl Subcommand for Opts {
             cfg,
         );
 
+        // TODO: properly pass the additional EKO from input args
         let fk_table = evolve_grid(
             &grid,
             &self.eko,
+            &self.eko,
             &pdf,
             &self.orders,
             self.xir,

pineappl_py/pineappl/grid.py

@@ -344,6 +344,7 @@ def evolve(
             self.raw.evolve(
                 np.array(operator_grid),
                 operators["q2_ref"],
+                operators["q2_ref"], # TODO: Modify the input EKO
                 np.array(operators["inputpids"], dtype=np.int32),
                 np.array(operators["inputgrid"]),
                 np.array(q2grid, dtype=np.float64),

pineappl_py/src/grid.rs

@@ -458,7 +458,8 @@ impl PyGrid {
     ///     produced FK table
     pub fn evolve(
         &self,
-        operator: PyReadonlyArray5<f64>,
+        operator_a: PyReadonlyArray5<f64>,
+        operator_b: PyReadonlyArray5<f64>,
         fac0: f64,
         pids0: PyReadonlyArray1<i32>,
         x0: PyReadonlyArray1<f64>,
@@ -488,7 +489,8 @@ impl PyGrid {
         let evolved_grid = self
             .grid
             .evolve(
-                operator.as_array(),
+                operator_a.as_array(),
+                operator_b.as_array(),
                 &op_info,
                 order_mask.as_slice().unwrap(),
             )