climate-indices 0.1.0

//! Generalized logistic (log-logistic, GLO) kernels behind SPEI's reference distribution.
//!
//! Reference: Hosking (1990); the parameter estimate is `PELGLO` and the CDF is
//! `CDFGLO` from Hosking's IBM Research Report RC20525 (1996) as distributed in
//! the R package `lmom`. Vicente-Serrano, Beguería, and López-Moreno (2010)
//! define SPEI on it; ADR-0016 records the choice.
//!
//! The Python implementations in `climate_indices.compute` and
//! `climate_indices.lmoments` stay the reference oracle;
//! `tests/test_native_parity_distributions.py` checks these kernels against them at
//! `rtol = atol = 1e-10`. Calibration-period selection, validity masking of the
//! fitted parameters, output scales, warnings, and logging stay in Python.

use std::f64::consts::PI;

use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Zip};

use crate::ClimateError;
use crate::lmoments::sample_lmoments;

/// A shape of at most this magnitude is the ordinary logistic, `SMALL` in `PELGLO`.
const SHAPE_TOLERANCE: f64 = 1e-6;

/// GLO fit of every column of a calibration block.
///
/// A column that cannot be fitted has `valid == false` and zero parameters.
#[derive(Debug, Clone, PartialEq)]
pub struct LogLogisticFit {
    /// Location per calibration column, in the input's units.
    pub locs: Array1<f64>,
    /// Scale per calibration column, in the input's units.
    pub scales: Array1<f64>,
    /// Dimensionless generalized-logistic shape per calibration column.
    pub shapes: Array1<f64>,
    /// Whether each column has a usable fit.
    pub valid: Array1<bool>,
}

/// `(loc, scale, shape)` of one column from its L-moments.
fn fit_column(column: ArrayView1<'_, f64>, sorted: &mut Vec<f64>) -> Option<[f64; 3]> {
    let [first, second, skewness] = sample_lmoments(column, sorted)?;
    let shape = -skewness;
    // negated so that a NaN L-moment is invalid
    if !(second > 0.0 && shape.abs() < 1.0) {
        return None;
    }
    if shape.abs() <= SHAPE_TOLERANCE {
        return Some([first, second, 0.0]);
    }
    let gg = shape * PI / (shape * PI).sin();
    let scale = second / gg;
    Some([first - scale * (1.0 - gg) / shape, scale, shape])
}

/// L-moment GLO location, scale, and shape for each column.
///
/// # Numerical contract
///
/// - Python source: `lmoments.fit_glo_spatial`, which `lmoments.fit_glo` mirrors
///   one time step at a time, via `compute._loglogistic_parameters_spatial`.
/// - Inputs: `calibration`, shaped (years, columns), one column per calendar step
///   (and grid cell), already restricted to the calibration years.
/// - Outputs: a [`LogLogisticFit`] with one value per column.
/// - Zero semantics: every value takes part in the fit, zeros included; unlike
///   gamma and Pearson Type III there is no separate zero mass.
/// - NaN semantics: NaN is missing. A column with fewer than four non-missing
///   values, a non-positive second L-moment, or `|shape| >= 1` (or a NaN
///   there) is invalid.
/// - Numerics: `shape = -tau_3`; a shape of at most 1e-6 in magnitude is
///   exactly zero (the ordinary logistic), otherwise `PELGLO`'s `gg`.
///
/// The fit uses Hosking's (1990, 1996) sample L-moments and `PELGLO` estimate.
/// Values, locations, and scales share units, usually mm; shapes are dimensionless.
/// Calendar columns follow [`crate`]; the caller selects the Calibration Period.
/// Invalid columns, including empty and constant columns, retain zero parameters
/// and `valid == false`. Negative values are ordinary sample values.
/// See the `Measured accuracy` section in [`crate`], entry `spei_loglogistic`
/// in `tests/parity_registry.py`. It measures a complete workflow, not isolated fits.
///
/// # Example
///
/// ```
/// use climate_indices::loglogistic::loglogistic_parameters;
/// use ndarray::array;
///
/// let sample_mm = array![[1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0]];
/// let fit = loglogistic_parameters(sample_mm.view());
/// assert!(fit.valid[0]);
/// assert_eq!(fit.locs[0], 4.0);
/// ```
pub fn loglogistic_parameters(calibration: ArrayView2<'_, f64>) -> LogLogisticFit {
    let columns = calibration.ncols();
    let mut fit = LogLogisticFit {
        locs: Array1::zeros(columns),
        scales: Array1::zeros(columns),
        shapes: Array1::zeros(columns),
        valid: Array1::from_elem(columns, false),
    };
    let mut sorted = Vec::with_capacity(calibration.nrows());
    Zip::from(&mut fit.locs)
        .and(&mut fit.scales)
        .and(&mut fit.shapes)
        .and(&mut fit.valid)
        .and(calibration.columns())
        .for_each(|loc, scale, shape, valid, column| {
            if let Some(parameters) = fit_column(column, &mut sorted) {
                [*loc, *scale, *shape] = parameters;
                *valid = true;
            }
        });
    fit
}

/// Generalized logistic CDF, `cdfglo`, clipped to [0, 1].
///
/// # Numerical contract
///
/// With `z = (x - loc) / scale`: `F = 1 / (1 + exp(-y))`, where `y = z` for a
/// shape of at most 1e-6 in magnitude, else `y = -ln(max(0, 1 - shape * z)) / shape`.
/// A NaN input gives NaN. This is the bare formula: the caller decides which
/// parameter sets are valid, as `compute._loglogistic_fit` does.
/// This is Hosking's (1996) `CDFGLO` formula with the `PELGLO` small-shape limit.
/// `x`, `loc`, and `scale` share units; `shape` and the probability are dimensionless.
/// This scalar function has no calendar interpretation or discrete zero mass.
/// Zero is an ordinary value. For positive scale and `|shape| > 1e-6`, crossing
/// the finite support gives 1 for positive shape and 0 for negative shape.
/// Non-positive scales are not rejected; zero scale follows IEEE division,
/// including NaN when `x == loc`. NaN values or parameters propagate.
/// See the `Measured accuracy` section in [`crate`], entry `spei_loglogistic`
/// in `tests/parity_registry.py`. It measures a complete workflow, not this isolated CDF.
pub fn loglogistic_cdf(x: f64, loc: f64, scale: f64, shape: f64) -> f64 {
    let z = (x - loc) / scale;
    let y = if shape.abs() <= SHAPE_TOLERANCE {
        z
    } else {
        // f64::max would drop a NaN, which np.maximum propagates
        let argument = 1.0 - shape * z;
        let argument = if argument.is_nan() {
            argument
        } else {
            argument.max(0.0)
        };
        -argument.ln() / shape
    };
    (1.0 / (1.0 + (-y).exp())).clamp(0.0, 1.0)
}

/// GLO CDF of each value, one `(loc, scale, shape)` per column.
///
/// # Numerical contract
///
/// - Python source: the probability step of `compute._loglogistic_fit`.
/// - Inputs: `values` shaped (years, columns); `locs`, `scales`, and `shapes`
///   with one value per column.
/// - Outputs: probabilities shaped like `values`, before Python masks the
///   positions whose parameters are invalid and maps to the output scale.
/// - Zero semantics: a zero is an ordinary value; the GLO has no zero mass.
/// - NaN semantics: a NaN value or parameter gives NaN. A finite value beyond
///   the fitted support gives exactly 0 or 1.
/// - Invalid input: parameters are not validated here (a non-positive scale or
///   `|shape| >= 1` still evaluates the formula); Python masks those positions.
///
/// Returns [`ClimateError::ShapeMismatch`] when a parameter's length is not the
/// number of columns.
///
/// Each element uses Hosking's [`loglogistic_cdf`], including its small-shape
/// limit and unvalidated scale arithmetic. No fit-validity mask is applied.
/// Values, locations, and scales share units; shapes and outputs are dimensionless.
/// Calendar columns follow [`crate`]; no calendar selection occurs here.
/// See the `Measured accuracy` section in [`crate`], entry `spei_loglogistic`
/// in `tests/parity_registry.py`. It measures a complete workflow, not this isolated block CDF.
///
/// # Example
///
/// ```
/// use climate_indices::loglogistic::{loglogistic_cdf_block, loglogistic_parameters};
/// use ndarray::array;
///
/// let sample_mm = array![[1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0]];
/// let fit = loglogistic_parameters(sample_mm.view());
/// let probabilities = loglogistic_cdf_block(
///     sample_mm.view(),
///     fit.locs.view(),
///     fit.scales.view(),
///     fit.shapes.view(),
/// )?;
/// assert!(probabilities[[3, 0]] > 0.0 && probabilities[[3, 0]] < 1.0);
/// # Ok::<(), climate_indices::ClimateError>(())
/// ```
pub fn loglogistic_cdf_block(
    values: ArrayView2<'_, f64>,
    locs: ArrayView1<'_, f64>,
    scales: ArrayView1<'_, f64>,
    shapes: ArrayView1<'_, f64>,
) -> Result<Array2<f64>, ClimateError> {
    let columns = values.ncols();
    for (argument, parameter) in [("locs", &locs), ("scales", &scales), ("shapes", &shapes)] {
        if parameter.len() != columns {
            return Err(ClimateError::ShapeMismatch {
                argument,
                expected: columns,
                actual: parameter.len(),
            });
        }
    }

    let mut probabilities = Array2::<f64>::zeros(values.raw_dim());
    Zip::from(probabilities.columns_mut())
        .and(values.columns())
        .and(&locs)
        .and(&scales)
        .and(&shapes)
        .for_each(|mut out, column, &loc, &scale, &shape| {
            Zip::from(&mut out)
                .and(&column)
                .for_each(|p, &x| *p = loglogistic_cdf(x, loc, scale, shape));
        });
    Ok(probabilities)
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::array;

    #[test]
    fn fit_of_a_symmetric_sample_is_the_ordinary_logistic() {
        // tau_3 = 0, so shape = 0 and (loc, scale) = (lambda_1, lambda_2)
        let sample = array![[1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0]];
        let fit = loglogistic_parameters(sample.view());
        assert!(fit.valid[0]);
        assert_eq!(fit.shapes[0], 0.0);
        assert_eq!(fit.locs[0], 4.0);
        assert!((fit.scales[0] - 4.0 / 3.0).abs() < 1e-14);
    }

    #[test]
    fn fit_keeps_zeros_and_rejects_a_short_column() {
        let sample = array![
            [0.0, 1.0],
            [0.0, 2.0],
            [1.0, f64::NAN],
            [4.0, 3.0],
            [9.0, f64::NAN]
        ];
        let fit = loglogistic_parameters(sample.view());
        assert!(fit.valid[0], "zeros are ordinary values in a GLO fit");
        assert!(
            !fit.valid[1],
            "three non-missing values is below the minimum"
        );
        assert_eq!((fit.locs[1], fit.scales[1], fit.shapes[1]), (0.0, 0.0, 0.0));
    }

    // Reference values from climate_indices.lmoments.fit_glo (the Python oracle) on
    // the samples pearson.rs uses: tau_3 < 1/3, tau_3 >= 1/3, and a negative skew.
    #[test]
    fn fit_matches_the_python_oracle_on_each_branch() {
        let sample = array![
            [12.0, 1.0, 99.0],
            [15.5, 1.0, 99.0],
            [9.1, 1.0, 99.0],
            [30.2, 2.0, 98.0],
            [22.4, 2.0, 98.0],
            [18.8, 3.0, 97.0],
            [11.3, 4.0, 96.0],
            [40.7, 8.0, 92.0],
            [25.0, 20.0, 80.0],
            [14.2, 60.0, 40.0],
        ];
        let expected = [
            // (loc, scale, shape)
            (
                17.346_295_894_276_206,
                4.921_322_708_270_736,
                -0.287_872_841_444_269,
            ),
            (
                2.809_361_769_423_021_3,
                2.175_845_858_244_121_4,
                -0.766_997_167_138_809_2,
            ),
            (
                97.190_638_230_577,
                2.175_845_858_244_097_4,
                0.766_997_167_138_811_5,
            ),
        ];
        let fit = loglogistic_parameters(sample.view());
        for (column, (loc, scale, shape)) in expected.into_iter().enumerate() {
            assert!(fit.valid[column]);
            for (name, actual, wanted) in [
                ("loc", fit.locs[column], loc),
                ("scale", fit.scales[column], scale),
                ("shape", fit.shapes[column], shape),
            ] {
                assert!(
                    (actual - wanted).abs() <= 1e-13 * wanted.abs(),
                    "column {column} {name} = {actual:e}, expected {wanted:e}"
                );
            }
        }
    }

    #[test]
    fn cdf_matches_the_numpy_formula() {
        // (x, loc, scale, shape, compute._loglogistic_fit's probability)
        let cases = [
            (
                1.0,
                2.809_361_769_423_021_3,
                2.175_845_858_244_121_4,
                -0.766_997_167_138_809_2,
                0.210_136_047_689_925_94,
            ),
            (
                10.0,
                2.809_361_769_423_021_3,
                2.175_845_858_244_121_4,
                -0.766_997_167_138_809_2,
                0.838_378_299_114_486_9,
            ),
            (
                95.0,
                97.190_638_230_577,
                2.175_845_858_244_097_4,
                0.766_997_167_138_811_5,
                0.321_679_750_448_014_45,
            ),
            (
                17.0,
                17.346_295_894_276_206,
                4.921_322_708_270_736,
                -0.287_872_841_444_269,
                0.482_235_257_835_268_7,
            ),
            (5.0, 0.0, 2.0, 0.0, 0.924_141_819_978_756_6),
        ];
        for (x, loc, scale, shape, expected) in cases {
            let actual = loglogistic_cdf(x, loc, scale, shape);
            assert!(
                (actual - expected).abs() <= 1e-14 * expected,
                "loglogistic_cdf({x}, {loc}, {scale}, {shape}) = {actual:e}, expected {expected:e}"
            );
        }
    }

    #[test]
    fn cdf_is_the_logistic_at_zero_shape_and_clips_beyond_the_support() {
        assert_eq!(loglogistic_cdf(2.0, 2.0, 3.0, 0.0), 0.5);
        // shape 0.5, scale 1: the support ends at z = 2
        assert_eq!(loglogistic_cdf(3.0, 0.0, 1.0, 0.5), 1.0);
        assert_eq!(loglogistic_cdf(-3.0, 0.0, 1.0, -0.5), 0.0);
        assert!(loglogistic_cdf(f64::NAN, 0.0, 1.0, 0.5).is_nan());
    }

    #[test]
    fn cdf_block_rejects_a_parameter_of_the_wrong_length() {
        let error = loglogistic_cdf_block(
            array![[1.0, 2.0]].view(),
            array![0.0, 0.0].view(),
            array![1.0, 1.0].view(),
            array![0.1].view(),
        )
        .unwrap_err();
        assert_eq!(
            error,
            ClimateError::ShapeMismatch {
                argument: "shapes",
                expected: 2,
                actual: 1
            }
        );
    }
}