laddu-extensions 0.19.0

//! Experimental extensions to the `laddu` ecosystem.
//!
//! <div class="warning">
//!
//! This module contains experimental code which may be untested or unreliable. Use at your own
//! risk! The features contained here may eventually be moved into the standard crate modules.
//!
//! </div>

use laddu_core::{
    math::histogram, traits::Variable, LadduError, LadduResult, Parameter, ParameterMap,
};
use nalgebra::DVector;

use crate::{
    likelihood::{LikelihoodExpression, LikelihoodTerm},
    NLL,
};

/// A [`LikelihoodTerm`] whose size is proportional to the χ²-distance from a binned projection of
/// the fit to a provided set of datapoints representing the true values in each bin.
///
/// This is intended to be used as follows. Suppose we perform a binned fit to a simple amplitude
/// which is not parameterized over the binning variable. We then form a new
/// [`Expression`](`laddu_core::Expression`) which *is*
/// parameterized over said variable, and we wish to perform an unbinned fit. If we can isolate
/// terms which are not interfering, we could imagine fitting the unbinned data with a cost
/// function that minimizes the distance to the result from the binned fit. From there, it is up to
/// the user to decide what to do with this minimum. Caution should be used, as this will not be
/// the minimum of the [`NLL`], but of the guide term only. However, this minimum could be used as
/// an intermediate for getting close to a global minimum if the likelihood landscape has many
/// local minima. Then a true fit could be performed, starting at this intermediate point.
#[derive(Clone)]
pub struct BinnedGuideTerm {
    nll: Box<NLL>,
    values: Vec<f64>,
    amplitude_sets: Vec<Vec<String>>,
    bins: usize,
    range: (f64, f64),
    count_sets: Vec<Vec<f64>>,
    error_sets: Vec<Vec<f64>>,
}

impl BinnedGuideTerm {
    /// Construct a new [`BinnedGuideTerm`]
    ///
    /// This term takes a list of subsets of amplitudes, activates each set, and compares the projected
    /// histogram to the known one provided at construction. Both `count_sets` and `error_sets` should
    /// have the same shape, and their first dimension should be the same as that of `amplitude_sets`.
    ///
    /// The intended usage is to provide some sets of amplitudes to isolate, like `[["amp1", "amp2"], ["amp3"]]`,
    /// along with some known counts for a binned fit (`count_sets ~ [[histogram counts involving "amp1" and "amp2"], [histogram counts involving "amp3"]]` and simlar for `error_sets`).
    #[allow(clippy::new_ret_no_self)]
    pub fn new<
        V: Variable + 'static,
        L: AsRef<str>,
        T: AsRef<[L]>,
        U: AsRef<[f64]>,
        E: AsRef<[f64]>,
    >(
        nll: Box<NLL>,
        variable: &V,
        amplitude_sets: &[T],
        bins: usize,
        range: (f64, f64),
        count_sets: &[U],
        error_sets: Option<&[E]>,
    ) -> LadduResult<LikelihoodExpression> {
        if bins == 0 {
            return Err(LadduError::Custom(
                "BinnedGuideTerm requires bins > 0".to_string(),
            ));
        }
        let values = variable.value_on(&nll.accmc_evaluator.dataset)?;
        let amplitude_sets: Vec<Vec<String>> = amplitude_sets
            .iter()
            .map(|t| t.as_ref().iter().map(|s| s.as_ref().to_string()).collect())
            .collect();
        let count_sets: Vec<Vec<f64>> = count_sets.iter().map(|f| f.as_ref().to_vec()).collect();
        let error_sets: Vec<Vec<f64>> = if let Some(error_sets) = error_sets {
            error_sets.iter().map(|f| f.as_ref().to_vec()).collect()
        } else {
            count_sets
                .iter()
                .map(|v| v.iter().map(|f| f.sqrt()).collect())
                .collect()
        };
        if amplitude_sets.len() != count_sets.len() {
            return Err(LadduError::LengthMismatch {
                context: "BinnedGuideTerm amplitude_sets/count_sets".to_string(),
                expected: amplitude_sets.len(),
                actual: count_sets.len(),
            });
        }
        if count_sets.len() != error_sets.len() {
            return Err(LadduError::LengthMismatch {
                context: "BinnedGuideTerm count_sets/error_sets".to_string(),
                expected: count_sets.len(),
                actual: error_sets.len(),
            });
        }
        for (set_idx, counts) in count_sets.iter().enumerate() {
            if counts.len() != bins {
                return Err(LadduError::LengthMismatch {
                    context: format!("BinnedGuideTerm count_sets[{set_idx}]"),
                    expected: bins,
                    actual: counts.len(),
                });
            }
        }
        for (set_idx, errors) in error_sets.iter().enumerate() {
            if errors.len() != bins {
                return Err(LadduError::LengthMismatch {
                    context: format!("BinnedGuideTerm error_sets[{set_idx}]"),
                    expected: bins,
                    actual: errors.len(),
                });
            }
        }
        Self {
            nll,
            amplitude_sets,
            values,
            bins,
            range,
            count_sets,
            error_sets,
        }
        .into_expression()
    }
}

impl LikelihoodTerm for BinnedGuideTerm {
    fn evaluate(&self, parameters: &[f64]) -> LadduResult<f64> {
        let mut result = 0.0;
        for ((counts, errors), amplitudes) in self
            .count_sets
            .iter()
            .zip(self.error_sets.iter())
            .zip(self.amplitude_sets.iter())
        {
            let weights = self
                .nll
                .project_weights_subset(parameters, amplitudes, None)?;
            let eval_hist = histogram(&self.values, self.bins, self.range, Some(&weights));
            // TODO: handle entries where e == 0
            let chisqr: f64 = eval_hist
                .counts
                .iter()
                .zip(counts.iter())
                .zip(errors.iter())
                .map(|((o, c), e)| (o - c).powi(2) / e.powi(2))
                .sum();
            result += chisqr;
        }
        Ok(result)
    }

    fn evaluate_gradient(&self, parameters: &[f64]) -> LadduResult<DVector<f64>> {
        let mut gradient = DVector::zeros(parameters.len());
        let bin_width = (self.range.1 - self.range.0) / self.bins as f64;
        for ((counts, errors), amplitudes) in self
            .count_sets
            .iter()
            .zip(self.error_sets.iter())
            .zip(self.amplitude_sets.iter())
        {
            let (weights, weights_gradient) = self
                .nll
                .project_weights_and_gradients_subset(parameters, amplitudes, None)?;
            let mut eval_counts = vec![0.0; self.bins];
            let mut eval_count_gradient: Vec<DVector<f64>> =
                vec![DVector::zeros(parameters.len()); self.bins];

            for (j, &value) in self.values.iter().enumerate() {
                if value >= self.range.0 && value < self.range.1 {
                    let bin_idx =
                        (((value - self.range.0) / bin_width).floor() as usize).min(self.bins - 1);
                    eval_counts[bin_idx] += weights[j];
                    for k in 0..parameters.len() {
                        eval_count_gradient[bin_idx][k] += weights_gradient[j][k];
                    }
                }
            }
            for i in 0..self.bins {
                let o_i = eval_counts[i];
                let c_i = counts[i];
                let e_i = errors[i];
                let residual = o_i - c_i;
                let residual_gradient = &eval_count_gradient[i];
                for k in 0..parameters.len() {
                    gradient[k] += 2.0 * residual * residual_gradient[k] / e_i.powi(2);
                }
            }
        }
        Ok(gradient)
    }

    fn fix_parameter(&self, name: &str, value: f64) -> LadduResult<()> {
        self.nll.fix_parameter(name, value)
    }

    fn free_parameter(&self, name: &str) -> LadduResult<()> {
        self.nll.free_parameter(name)
    }

    fn rename_parameter(&self, old: &str, new: &str) -> LadduResult<()> {
        self.nll.rename_parameter(old, new)
    }

    fn rename_parameters(
        &self,
        mapping: &std::collections::HashMap<String, String>,
    ) -> LadduResult<()> {
        self.nll.rename_parameters(mapping)
    }

    fn parameter_map(&self) -> ParameterMap {
        self.nll.parameter_map()
    }
}

/// A weighted regularization term.
///
/// This can be interpreted as a prior of the form
///
/// ```math
/// f(\vec{x}) = \frac{p\lambda^{1/p}}{2\Gamma(1/p)}e^{-\frac{\lambda|\vec{x}|^p}}
/// ```
/// which becomes a Laplace distribution for $`p=1`$ and a Gaussian for $`p=2`$. These are commonly
/// interpreted as $`\ell_p`$ regularizers for linear regression models, with $`p=1`$ and $`p=2`$
/// corresponding to LASSO and ridge regression, respectively. When used in nonlinear regression,
/// these should be interpeted as the prior listed above when used in maximum a posteriori (MAP)
/// estimation. Explicitly, when the logarithm is taken, this term becomes
///
/// ```math
/// \lambda \left(\sum_{j} w_j |x_j|^p\right)^{1/p}
/// ```
/// plus some additional constant terms which do not depend on free parameters.
///
/// Weights can be specified to vary the influence of each parameter used in the regularization.
/// These weights are typically assigned by first fitting without a regularization term to obtain
/// parameter values $`\vec{\beta}`$, choosing a value $`\gamma>0`, and setting the weights to
/// $`\vec{w} = 1/|\vec{\beta}|^\gamma`$ according to a paper by Zou[^1].
///
/// [^1]: [Zou, H. (2006). The Adaptive Lasso and Its Oracle Properties. In Journal of the American Statistical Association (Vol. 101, Issue 476, pp. 1418–1429). Informa UK Limited.](https://doi.org/10.1198/016214506000000735)
#[derive(Clone)]
pub struct Regularizer<const P: usize> {
    parameter_map: ParameterMap,
    lambda: f64,
    weights: Vec<f64>,
}

impl<const P: usize> Regularizer<P> {
    fn construct<T, U, F>(parameters: T, lambda: f64, weights: Option<F>) -> LadduResult<Box<Self>>
    where
        T: IntoIterator<Item = U>,
        U: AsRef<str>,
        F: AsRef<[f64]>,
    {
        let parameters: Vec<String> = parameters
            .into_iter()
            .map(|s| s.as_ref().to_string())
            .collect();
        let weights: Vec<f64> = weights
            .as_ref()
            .map_or(vec![1.0; parameters.len()].as_ref(), AsRef::as_ref)
            .to_vec();
        if parameters.len() != weights.len() {
            return Err(LadduError::LengthMismatch {
                context: "Regularizer parameter/weight vector".to_string(),
                expected: parameters.len(),
                actual: weights.len(),
            });
        }
        let mut parameter_map = ParameterMap::default();
        for parameter in &parameters {
            if parameter_map.insert(Parameter::new(parameter)).is_some() {
                return Err(LadduError::ParameterConflict {
                    name: parameter.clone(),
                    reason: "duplicate regularizer parameter name".to_string(),
                });
            }
        }
        Ok(Self {
            parameter_map,
            lambda,
            weights,
        }
        .into())
    }
}

impl Regularizer<1> {
    /// Create a new $`\ell_1`$ [`Regularizer`] expressed as a [`LikelihoodExpression`].
    #[allow(clippy::new_ret_no_self)]
    pub fn new<T, U, F>(
        parameters: T,
        lambda: f64,
        weights: Option<F>,
    ) -> LadduResult<LikelihoodExpression>
    where
        T: IntoIterator<Item = U>,
        U: AsRef<str>,
        F: AsRef<[f64]>,
    {
        Self::construct(parameters, lambda, weights)?.into_expression()
    }
}

impl Regularizer<2> {
    /// Create a new $`\ell_2`$ [`Regularizer`] expressed as a [`LikelihoodExpression`].
    #[allow(clippy::new_ret_no_self)]
    pub fn new<T, U, F>(
        parameters: T,
        lambda: f64,
        weights: Option<F>,
    ) -> LadduResult<LikelihoodExpression>
    where
        T: IntoIterator<Item = U>,
        U: AsRef<str>,
        F: AsRef<[f64]>,
    {
        Self::construct(parameters, lambda, weights)?.into_expression()
    }
}

impl LikelihoodTerm for Regularizer<1> {
    fn evaluate(&self, parameters: &[f64]) -> LadduResult<f64> {
        Ok(self.lambda * parameters.iter().map(|p| p.abs()).sum::<f64>())
    }

    fn evaluate_gradient(&self, parameters: &[f64]) -> LadduResult<DVector<f64>> {
        Ok(DVector::from_vec(
            parameters
                .iter()
                .zip(self.weights.iter())
                .map(|(p, w)| w * p.signum())
                .collect(),
        )
        .scale(self.lambda))
    }

    fn parameter_map(&self) -> ParameterMap {
        self.parameter_map.clone()
    }
}

impl LikelihoodTerm for Regularizer<2> {
    fn evaluate(&self, parameters: &[f64]) -> LadduResult<f64> {
        Ok(self.lambda * parameters.iter().map(|p| p.powi(2)).sum::<f64>().sqrt())
    }

    fn evaluate_gradient(&self, parameters: &[f64]) -> LadduResult<DVector<f64>> {
        let denom = parameters
            .iter()
            .zip(self.weights.iter())
            .map(|(p, w)| w * p.powi(2))
            .sum::<f64>()
            .sqrt();
        Ok(DVector::from_vec(parameters.to_vec()).scale(self.lambda / denom))
    }

    fn parameter_map(&self) -> ParameterMap {
        self.parameter_map.clone()
    }
}

#[cfg(test)]
mod tests {
    use approx::assert_relative_eq;

    use super::Regularizer;

    #[test]
    fn l1_regularizer_respects_weights() {
        let expr = Regularizer::<1>::new(["alpha", "beta"], 2.0, Some([1.0, 0.5])).unwrap();
        let values = vec![1.5, -2.0];
        assert_relative_eq!(expr.evaluate(&values).unwrap(), 7.0);
        let grad = expr.evaluate_gradient(&values).unwrap();
        assert_relative_eq!(grad[0], 2.0);
        assert_relative_eq!(grad[1], -1.0);
    }

    #[test]
    fn l2_regularizer_gradient_scales_parameters() {
        let expr = Regularizer::<2>::new(["x", "y"], 3.0, Some([1.0, 2.0])).unwrap();
        let values = vec![3.0_f64, 4.0_f64];
        assert_relative_eq!(expr.evaluate(&values).unwrap(), 15.0);
        let grad = expr.evaluate_gradient(&values).unwrap();
        let denom = (1.0 * values[0].powi(2) + 2.0 * values[1].powi(2)).sqrt();
        assert_relative_eq!(grad[0], 3.0 * values[0] / denom);
        assert_relative_eq!(grad[1], 3.0 * values[1] / denom);
    }

    #[test]
    fn regularizer_rejects_weight_mismatch() {
        let err = Regularizer::<1>::new(["alpha", "beta"], 1.0, Some([1.0]));
        assert!(err.is_err());
    }

    #[test]
    fn regularizer_defaults_to_unit_weights() {
        let expr = Regularizer::<1>::new(["alpha", "beta"], 1.5, None::<Vec<f64>>).unwrap();
        let values = vec![1.0, -2.0];
        assert_relative_eq!(expr.evaluate(&values).unwrap(), 4.5);
        let grad = expr.evaluate_gradient(&values).unwrap();
        assert_relative_eq!(grad[0], 1.5);
        assert_relative_eq!(grad[1], -1.5);
    }
}