Skip to main content

laddu_python/extensions/
experimental.rs

1use laddu_extensions::experimental::{BinnedGuideTerm, Regularizer};
2use pyo3::{exceptions::PyValueError, prelude::*};
3
4use crate::{
5    extensions::likelihood::{PyLikelihoodExpression, PyNLL},
6    variables::PyVariable,
7};
8/// A χ²-like term which uses a known binned result to guide the fit
9///
10///
11/// This term takes a list of subsets of amplitudes, activates each set, and compares the projected
12/// histogram to the known one provided at construction. Both `count_sets` and `error_sets` should
13/// have the same shape, and their first dimension should be the same as that of `amplitude_sets`.
14///
15/// Parameters
16/// ----------
17/// nll: NLL
18/// variable : {laddu.Mass, laddu.CosTheta, laddu.Phi, laddu.PolAngle, laddu.PolMagnitude, laddu.Mandelstam}
19///     The variable to use for binning
20/// amplitude_sets : list of list of str
21///     A list of lists of amplitudes to activate, with each inner list representing a set that
22///     corresponds to the provided binned data
23/// bins : int
24/// range : tuple of (min, max)
25///     The range of the variable to use for binning
26/// count_sets : list of list of float
27///      A list of binned counts for each amplitude set
28/// error_sets : list of list of float, optional
29///      A list of bin errors for each amplitude set (square root of `count_sets` if None is
30///      provided)
31///
32/// Returns
33/// -------
34/// LikelihoodExpression
35///     A term that can be combined with other likelihood expressions.
36#[pyfunction(name = "BinnedGuideTerm", signature = (nll, variable, amplitude_sets, bins, range, count_sets, error_sets = None))]
37pub fn py_binned_guide_term(
38    nll: PyNLL,
39    variable: Bound<'_, PyAny>,
40    amplitude_sets: Vec<Vec<String>>,
41    bins: usize,
42    range: (f64, f64),
43    count_sets: Vec<Vec<f64>>,
44    error_sets: Option<Vec<Vec<f64>>>,
45) -> PyResult<PyLikelihoodExpression> {
46    let variable = variable.extract::<PyVariable>()?;
47    Ok(PyLikelihoodExpression(BinnedGuideTerm::new(
48        nll.0.clone(),
49        &variable,
50        &amplitude_sets,
51        bins,
52        range,
53        &count_sets,
54        error_sets.as_deref(),
55    )?))
56}
57
58/// An weighted :math:`\ell_p` regularization term which acts as a maximum a posteriori (MAP) prior.
59///
60/// This can be interpreted as a prior of the form
61///
62/// .. math:: f(\vec{x}) = \frac{p\lambda^{1/p}}{2\Gamma(1/p)}e^{-\lambda|\vec{x}|^p}
63///
64/// which becomes a Laplace distribution for :math:`p=1` and a Gaussian for :math:`p=2`. These are commonly
65/// interpreted as :math:`\ell_p` regularizers for linear regression models, with :math:`p=1` and :math:`p=2`
66/// corresponding to LASSO and ridge regression, respectively. When used in nonlinear regression,
67/// these should be interpeted as the prior listed above when used in maximum a posteriori (MAP)
68/// estimation. Explicitly, when the logarithm is taken, this term becomes
69///
70/// .. math:: \lambda \left(\sum_{j} w_j |x_j|^p\right)^{1/p}
71///
72/// plus some additional constant terms which do not depend on free parameters.
73///
74/// Weights can be specified to vary the influence of each parameter used in the regularization.
75/// These weights are typically assigned by first fitting without a regularization term to obtain
76/// parameter values :math:`\vec{\beta}`, choosing a value :math:`\gamma>0`, and setting the weights to
77/// :math:`\vec{w} = 1/|\vec{\beta}|^\gamma` according to [Zou]_.
78///
79/// References
80/// ----------
81/// .. [Zou] Zou, H. (2006). The Adaptive Lasso and Its Oracle Properties.
82///    Journal of the American Statistical Association, 101(476), 1418–1429.
83///    doi:10.1198/016214506000000735
84///
85/// Parameters
86/// ----------
87/// parameters : list of str
88///     The names of the parameters to regularize
89/// lda : float
90///     The regularization parameter :math:`\lambda`
91/// p : {1, 2}
92///     The degree of the norm :math:`\ell_p`
93/// weights : list of float, optional
94///     Weights to apply in the regularization to each parameter
95///
96/// Raises
97/// ------
98/// ValueError
99///     If :math:`p` is not 1 or 2
100/// Exception
101///     If the number of parameters and weights is not equal
102///
103/// Returns
104/// -------
105/// LikelihoodExpression
106///     A term that can be combined with other likelihood expressions.
107#[allow(rustdoc::broken_intra_doc_links)]
108#[pyfunction(name = "Regularizer", signature = (parameters, lda, p=1, weights=None))]
109pub fn py_regularizer(
110    parameters: Vec<String>,
111    lda: f64,
112    p: usize,
113    weights: Option<Vec<f64>>,
114) -> PyResult<PyLikelihoodExpression> {
115    if p == 1 {
116        Ok(PyLikelihoodExpression(Regularizer::<1>::new(
117            parameters, lda, weights,
118        )?))
119    } else if p == 2 {
120        Ok(PyLikelihoodExpression(Regularizer::<2>::new(
121            parameters, lda, weights,
122        )?))
123    } else {
124        Err(PyValueError::new_err(
125            "'Regularizer' only supports p = 1 or 2",
126        ))
127    }
128}