pineappl 0.5.8

//! Supporting classes and functions for [`Grid::evolve`].

use super::grid::{Grid, GridError, Order};
use super::import_only_subgrid::ImportOnlySubgridV2;
use super::lumi::LumiEntry;
use super::lumi_entry;
use super::sparse_array3::SparseArray3;
use super::subgrid::{Mu2, Subgrid, SubgridEnum};
use float_cmp::approx_eq;
use itertools::Itertools;
use ndarray::{s, Array1, Array2, Array3, Array5, ArrayView1, Axis};
use std::iter;

/// Information about the evolution kernel operator (EKO) passed to [`Grid::evolve`] as `operator`,
/// which is used to convert a [`Grid`] into an [`FkTable`]. The dimensions of the EKO must
/// correspond to the values given in [`fac1`], [`pids0`], [`x0`], [`pids1`] and [`x1`], exactly in
/// this order. Members with a `1` are defined at the squared factorization scales given in
/// [`fac1`] (often called process scales) and are found in the [`Grid`] that [`Grid::evolve`] is
/// called with. Members with a `0` are defined at the squared factorization scale [`fac0`] (often
/// called fitting scale or starting scale) and are found in the [`FkTable`] resulting from
/// [`Grid::evolve`].
///
/// The EKO may convert a `Grid` from a basis given by the particle identifiers [`pids1`] to a
/// possibly different basis given by [`pids0`]. This basis must also be identified using
/// [`lumi_id_types`], which tells [`FkTable::convolute`] how to perform a convolution. The members
/// [`ren1`] and [`alphas`] must be the strong couplings given at the respective renormalization
/// scales. Finally, [`xir`] and [`xif`] can be used to vary the renormalization and factorization
/// scales, respectively, around their central values.
pub struct OperatorInfo {
    /// Squared factorization scale of the `FkTable`.
    pub fac0: f64,
    /// Particle identifiers of the `FkTable`.
    pub pids0: Vec<i32>,
    /// `x`-grid coordinates of the `FkTable`
    pub x0: Vec<f64>,
    /// Squared factorization scales of the `Grid`.
    pub fac1: Vec<f64>,
    /// Particle identifiers of the `Grid`. If the `Grid` contains more particle identifiers than
    /// given here, the contributions of them are silently ignored.
    pub pids1: Vec<i32>,
    /// `x`-grid coordinates of the `Grid`.
    pub x1: Vec<f64>,

    /// Renormalization scales of the `Grid`.
    pub ren1: Vec<f64>,
    /// Strong couplings corresponding to the order given in [`ren1`].
    pub alphas: Vec<f64>,
    /// Multiplicative factor for the central renormalization scale.
    pub xir: f64,
    /// Multiplicative factor for the central factorization scale.
    pub xif: f64,
    /// Identifier of the particle basis for the `FkTable`.
    pub lumi_id_types: String,
}

fn gluon_has_pid_zero(grid: &Grid) -> bool {
    grid.key_values()
        .and_then(|key_values| key_values.get("lumi_id_types"))
        .and_then(|value| {
            (value == "pdg_mc_ids").then(|| {
                grid.lumi()
                    .iter()
                    .any(|entry| entry.entry().iter().any(|&(a, b, _)| (a == 0) || (b == 0)))
            })
        })
        .unwrap_or(false)
}

pub(crate) fn pids(
    operator: &Array5<f64>,
    info: &OperatorInfo,
    gluon_has_pid_zero: bool,
    pid1_nonzero: &dyn Fn(i32) -> bool,
) -> (Vec<(usize, usize)>, Vec<(i32, i32)>) {
    // list of all non-zero PID indices
    let pid_indices: Vec<_> = (0..operator.dim().3)
        .cartesian_product(0..operator.dim().1)
        .filter(|&(pid0_idx, pid1_idx)| {
            // 1) at least one element of the operator must be non-zero, and 2) the pid must be
            // contained in the lumi somewhere
            operator
                .slice(s![.., pid1_idx, .., pid0_idx, ..])
                .iter()
                .any(|&value| value != 0.0)
                && pid1_nonzero(if gluon_has_pid_zero && info.pids1[pid1_idx] == 21 {
                    0
                } else {
                    info.pids1[pid1_idx]
                })
        })
        .collect();

    // list of all non-zero (pid0, pid1) combinations
    let pids = pid_indices
        .iter()
        .map(|&(pid0_idx, pid1_idx)| {
            (
                info.pids0[pid0_idx],
                if gluon_has_pid_zero && info.pids1[pid1_idx] == 21 {
                    0
                } else {
                    info.pids1[pid1_idx]
                },
            )
        })
        .collect();

    (pid_indices, pids)
}

pub(crate) fn lumi0_with_one(pids: &[(i32, i32)]) -> Vec<i32> {
    let mut pids0: Vec<_> = pids.iter().map(|&(pid0, _)| pid0).collect();
    pids0.sort_unstable();
    pids0.dedup();

    pids0
}

pub(crate) fn lumi0_with_two(pids_a: &[(i32, i32)], pids_b: &[(i32, i32)]) -> Vec<(i32, i32)> {
    let mut pids0_a: Vec<_> = pids_a.iter().map(|&(pid0, _)| pid0).collect();
    pids0_a.sort_unstable();
    pids0_a.dedup();
    let mut pids0_b: Vec<_> = pids_b.iter().map(|&(pid0, _)| pid0).collect();
    pids0_b.sort_unstable();
    pids0_b.dedup();

    pids0_a
        .iter()
        .copied()
        .cartesian_product(pids0_b.iter().copied())
        .collect()
}

pub(crate) fn operators(
    operator: &Array5<f64>,
    info: &OperatorInfo,
    fac1: &[f64],
    pid_indices: &[(usize, usize)],
    x1: &[f64],
) -> Result<Vec<Array3<f64>>, GridError> {
    // permutation between the grid fac1 values and the operator fac1 values
    let fac1_indices: Vec<_> = if let Some(fac1_indices) = fac1
        .iter()
        .map(|&fac1p| {
            info.fac1
                .iter()
                .position(|&fac1| approx_eq!(f64, fac1p, fac1, ulps = 64))
        })
        .collect()
    {
        fac1_indices
    } else {
        return Err(GridError::EvolutionFailure(
            "operator information does not match grid's factorization scale values".to_string(),
        ));
    };

    // permutation between the grid x values and the operator x1 values
    let x1_indices: Vec<_> = if let Some(x1_indices) = x1
        .iter()
        .map(|&x1p| {
            info.x1
                .iter()
                .position(|&x1| approx_eq!(f64, x1p, x1, ulps = 64))
        })
        .collect()
    {
        x1_indices
    } else {
        return Err(GridError::EvolutionFailure(
            "operator information does not match grid's x-grid values".to_string(),
        ));
    };

    // create the corresponding operators accessible in the form [muf2, x0, x1]
    let operators: Vec<_> = pid_indices
        .iter()
        .map(|&(pid0_idx, pid1_idx)| {
            operator
                .slice(s![.., pid1_idx, .., pid0_idx, ..])
                .select(Axis(0), &fac1_indices)
                .select(Axis(1), &x1_indices)
                .permuted_axes([0, 2, 1])
                .as_standard_layout()
                .into_owned()
        })
        .collect();

    Ok(operators)
}

pub(crate) fn ndarray_from_subgrid_orders(
    info: &OperatorInfo,
    subgrids: &ArrayView1<SubgridEnum>,
    orders: &[Order],
    order_mask: &[bool],
) -> Result<(Vec<f64>, Vec<f64>, Vec<f64>, Array3<f64>), GridError> {
    // TODO: skip empty subgrids

    let mut fac1: Vec<_> = subgrids
        .iter()
        .flat_map(|subgrid| {
            subgrid
                .mu2_grid()
                .iter()
                .cloned()
                .map(|mu2| info.xif * info.xif * mu2.fac)
                .collect::<Vec<_>>()
        })
        .collect();
    let mut x1_a: Vec<_> = subgrids
        .iter()
        .flat_map(|subgrid| subgrid.x1_grid().into_owned())
        .collect();
    let mut x1_b: Vec<_> = subgrids
        .iter()
        .flat_map(|subgrid| subgrid.x2_grid().into_owned())
        .collect();

    fac1.sort_by(|a, b| a.partial_cmp(b).unwrap());
    fac1.dedup_by(|a, b| approx_eq!(f64, *a, *b, ulps = 64));
    x1_a.sort_by(|a, b| a.partial_cmp(b).unwrap());
    x1_a.dedup_by(|a, b| approx_eq!(f64, *a, *b, ulps = 64));
    x1_b.sort_by(|a, b| a.partial_cmp(b).unwrap());
    x1_b.dedup_by(|a, b| approx_eq!(f64, *a, *b, ulps = 64));

    let mut array = Array3::<f64>::zeros((fac1.len(), x1_a.len(), x1_b.len()));

    // add subgrids for different orders, but the same bin and lumi, using the right
    // couplings
    for (subgrid, order) in subgrids
        .iter()
        .zip(orders.iter())
        .zip(order_mask.iter().chain(iter::repeat(&true)))
        .filter_map(|((subgrid, order), &enabled)| {
            (enabled && !subgrid.is_empty()).then(|| (subgrid, order))
        })
    {
        let mut logs = 1.0;

        if order.logxir > 0 {
            if approx_eq!(f64, info.xir, 1.0, ulps = 4) {
                continue;
            }

            logs *= (info.xir * info.xir).ln();
        }

        if order.logxif > 0 {
            if approx_eq!(f64, info.xif, 1.0, ulps = 4) {
                continue;
            }

            logs *= (info.xif * info.xif).ln();
        }

        // TODO: use `try_collect` once stabilized
        let fac1_indices: Vec<_> = subgrid
            .mu2_grid()
            .iter()
            .map(|&Mu2 { fac, .. }| {
                fac1.iter()
                    .position(|&scale| approx_eq!(f64, info.xif * info.xif * fac, scale, ulps = 64))
                    .ok_or_else(|| {
                        GridError::EvolutionFailure(format!("no operator for muf2 = {} found", fac))
                    })
            })
            .collect::<Result<_, _>>()?;
        let xa_indices: Vec<_> = subgrid
            .x1_grid()
            .iter()
            .map(|&xa| {
                x1_a.iter()
                    .position(|&x1a| approx_eq!(f64, x1a, xa, ulps = 64))
                    .ok_or_else(|| {
                        GridError::EvolutionFailure(format!("no operator for x1 = {} found", xa))
                    })
            })
            .collect::<Result<_, _>>()?;
        let xb_indices: Vec<_> = subgrid
            .x2_grid()
            .iter()
            .map(|&xb| {
                x1_b.iter()
                    .position(|&x1b| approx_eq!(f64, x1b, xb, ulps = 64))
                    .ok_or_else(|| {
                        GridError::EvolutionFailure(format!("no operator for x1 = {} found", xb))
                    })
            })
            .collect::<Result<_, _>>()?;

        for ((ifac1, ix1, ix2), value) in subgrid.iter() {
            let mur2 = info.xir * info.xir * subgrid.mu2_grid()[ifac1].ren;

            let als = if order.alphas == 0 {
                1.0
            } else if let Some(alphas) = info
                .ren1
                .iter()
                .zip(info.alphas.iter())
                .find_map(|(&ren1, &alphas)| approx_eq!(f64, ren1, mur2, ulps = 64).then(|| alphas))
            {
                alphas.powi(order.alphas.try_into().unwrap())
            } else {
                return Err(GridError::EvolutionFailure(format!(
                    "could not find alphas for mur2 = {}",
                    mur2
                )));
            };

            array[[fac1_indices[ifac1], xa_indices[ix1], xb_indices[ix2]]] += als * logs * value;
        }
    }

    Ok((fac1, x1_a, x1_b, array))
}

pub(crate) fn evolve_with_one(
    grid: &Grid,
    operator: &Array5<f64>,
    info: &OperatorInfo,
    order_mask: &[bool],
) -> Result<(Array3<SubgridEnum>, Vec<LumiEntry>), GridError> {
    let gluon_has_pid_zero = gluon_has_pid_zero(grid);
    let has_pdf1 = grid.has_pdf1();

    let (pid_indices, pids) = pids(operator, info, gluon_has_pid_zero, &|pid| {
        grid.lumi()
            .iter()
            .flat_map(LumiEntry::entry)
            .any(|&(a, b, _)| if has_pdf1 { a } else { b } == pid)
    });

    let lumi0 = lumi0_with_one(&pids);
    let mut sub_fk_tables = Vec::with_capacity(grid.bin_info().bins() * lumi0.len());
    let new_axis = if has_pdf1 { 2 } else { 1 };

    let mut last_x1 = Vec::new();
    let mut last_fac1 = Vec::new();
    let mut ops = Vec::new();

    for subgrids_ol in grid.subgrids().axis_iter(Axis(1)) {
        let mut tables = vec![Array1::zeros(info.x0.len()); lumi0.len()];

        for (lumi1, subgrids_o) in subgrids_ol.axis_iter(Axis(1)).enumerate() {
            let (fac1, x1_a, x1_b, array) =
                ndarray_from_subgrid_orders(info, &subgrids_o, grid.orders(), order_mask)?;

            let x1 = if has_pdf1 { x1_a } else { x1_b };

            if x1.is_empty() {
                continue;
            }

            if (last_fac1.len() != fac1.len())
                || last_fac1
                    .iter()
                    .zip(fac1.iter())
                    .any(|(&lhs, &rhs)| !approx_eq!(f64, lhs, rhs, ulps = 64))
                || (last_x1.len() != x1.len())
                || last_x1
                    .iter()
                    .zip(x1.iter())
                    .any(|(&lhs, &rhs)| !approx_eq!(f64, lhs, rhs, ulps = 64))
            {
                ops = operators(operator, info, &fac1, &pid_indices, &x1)?;
                last_fac1 = fac1;
                last_x1 = x1;
            }

            // TODO: get rid of array-index access
            for (&pid1, &factor) in
                grid.lumi()[lumi1].entry().iter().map(
                    |(a, b, f)| {
                        if has_pdf1 {
                            (a, f)
                        } else {
                            (b, f)
                        }
                    },
                )
            {
                for (fk_table, op) in
                    lumi0
                        .iter()
                        .zip(tables.iter_mut())
                        .filter_map(|(&pid0, fk_table)| {
                            pids.iter()
                                .zip(ops.iter())
                                .find_map(|(&(p0, p1), op)| (p0 == pid0 && p1 == pid1).then(|| op))
                                .map(|op| (fk_table, op))
                        })
                {
                    let mut result = Array1::zeros(info.x0.len());

                    for imu2 in 0..array.dim().0 {
                        let op = op.index_axis(Axis(0), imu2);

                        result += &op.dot(
                            &array
                                .index_axis(Axis(0), imu2)
                                .index_axis(Axis(new_axis - 1), 0),
                        );
                    }

                    fk_table.scaled_add(factor, &result);
                }
            }
        }

        sub_fk_tables.extend(tables.into_iter().map(|table| {
            ImportOnlySubgridV2::new(
                SparseArray3::from_ndarray(
                    &table.insert_axis(Axis(0)).insert_axis(Axis(new_axis)),
                    0,
                    1,
                ),
                vec![Mu2 {
                    // TODO: FK tables don't depend on the renormalization scale
                    //ren: -1.0,
                    ren: info.fac0,
                    fac: info.fac0,
                }],
                if has_pdf1 { info.x0.clone() } else { vec![1.0] },
                if has_pdf1 { vec![1.0] } else { info.x0.clone() },
            )
            .into()
        }));
    }

    let pid = if has_pdf1 {
        grid.initial_state_2()
    } else {
        grid.initial_state_1()
    };

    Ok((
        Array1::from_iter(sub_fk_tables.into_iter())
            .into_shape((1, grid.bin_info().bins(), lumi0.len()))
            .unwrap(),
        lumi0
            .iter()
            .map(|&a| {
                lumi_entry![
                    if has_pdf1 { a } else { pid },
                    if has_pdf1 { pid } else { a },
                    1.0
                ]
            })
            .collect(),
    ))
}

pub(crate) fn evolve_with_two(
    grid: &Grid,
    operator: &Array5<f64>,
    info: &OperatorInfo,
    order_mask: &[bool],
) -> Result<(Array3<SubgridEnum>, Vec<LumiEntry>), GridError> {
    let gluon_has_pid_zero = gluon_has_pid_zero(grid);

    let (pid_indices_a, pids_a) = pids(operator, info, gluon_has_pid_zero, &|pid1| {
        grid.lumi()
            .iter()
            .flat_map(LumiEntry::entry)
            .any(|&(a, _, _)| a == pid1)
    });
    let (pid_indices_b, pids_b) = pids(operator, info, gluon_has_pid_zero, &|pid1| {
        grid.lumi()
            .iter()
            .flat_map(LumiEntry::entry)
            .any(|&(_, b, _)| b == pid1)
    });

    let lumi0 = lumi0_with_two(&pids_a, &pids_b);
    let mut sub_fk_tables = Vec::with_capacity(grid.bin_info().bins() * lumi0.len());

    let mut last_fac1 = Vec::new();
    let mut last_x1a = Vec::new();
    let mut last_x1b = Vec::new();
    let mut operators_a = Vec::new();
    let mut operators_b = Vec::new();

    for subgrids_ol in grid.subgrids().axis_iter(Axis(1)) {
        let mut tables = vec![Array2::zeros((info.x0.len(), info.x0.len())); lumi0.len()];

        for (lumi1, subgrids_o) in subgrids_ol.axis_iter(Axis(1)).enumerate() {
            let (fac1, x1_a, x1_b, array) =
                ndarray_from_subgrid_orders(info, &subgrids_o, grid.orders(), order_mask)?;

            let fac1_diff = (last_fac1.len() != fac1.len())
                || last_fac1
                    .iter()
                    .zip(fac1.iter())
                    .any(|(&lhs, &rhs)| !approx_eq!(f64, lhs, rhs, ulps = 64));

            if fac1_diff
                || (last_x1a.len() != x1_a.len())
                || last_x1a
                    .iter()
                    .zip(x1_a.iter())
                    .any(|(&lhs, &rhs)| !approx_eq!(f64, lhs, rhs, ulps = 64))
            {
                operators_a = operators(operator, info, &fac1, &pid_indices_a, &x1_a)?;
                last_x1a = x1_a;
            }

            if fac1_diff
                || (last_x1b.len() != x1_b.len())
                || last_x1b
                    .iter()
                    .zip(x1_b.iter())
                    .any(|(&lhs, &rhs)| !approx_eq!(f64, lhs, rhs, ulps = 64))
            {
                operators_b = operators(operator, info, &fac1, &pid_indices_b, &x1_b)?;
                last_x1b = x1_b;
            }

            if fac1_diff {
                last_fac1 = fac1;
            };

            // TODO: get rid of array-index access
            for &(pida1, pidb1, factor) in grid.lumi()[lumi1].entry() {
                for (fk_table, opa, opb) in
                    lumi0
                        .iter()
                        .zip(tables.iter_mut())
                        .filter_map(|(&(pida0, pidb0), fk_table)| {
                            pids_a
                                .iter()
                                .zip(operators_a.iter())
                                .cartesian_product(pids_b.iter().zip(operators_b.iter()))
                                .find_map(|((&(pa0, pa1), opa), (&(pb0, pb1), opb))| {
                                    (pa0 == pida0 && pa1 == pida1 && pb0 == pidb0 && pb1 == pidb1)
                                        .then(|| (opa, opb))
                                })
                                .map(|(opa, opb)| (fk_table, opa, opb))
                        })
                {
                    let mut result = Array2::zeros((info.x0.len(), info.x0.len()));

                    for imu2 in 0..array.dim().0 {
                        let opa = opa.index_axis(Axis(0), imu2);
                        let opb = opb.index_axis(Axis(0), imu2);
                        let arr = array.index_axis(Axis(0), imu2);

                        result += &opa.dot(&arr.dot(&opb.t()));
                    }

                    fk_table.scaled_add(factor, &result);
                }
            }
        }

        sub_fk_tables.extend(tables.into_iter().map(|table| {
            ImportOnlySubgridV2::new(
                SparseArray3::from_ndarray(&table.insert_axis(Axis(0)), 0, 1),
                vec![Mu2 {
                    // TODO: FK tables don't depend on the renormalization scale
                    //ren: -1.0,
                    ren: info.fac0,
                    fac: info.fac0,
                }],
                info.x0.clone(),
                info.x0.clone(),
            )
            .into()
        }));
    }

    Ok((
        Array1::from_iter(sub_fk_tables.into_iter())
            .into_shape((1, grid.bin_info().bins(), lumi0.len()))
            .unwrap(),
        lumi0.iter().map(|&(a, b)| lumi_entry![a, b, 1.0]).collect(),
    ))
}