metalforge 0.1.4

//! GLUE — Generalized Likelihood Uncertainty Estimation (Beven & Binley 1992,
//! Hydrological Processes 6, 279–298).
//!
//! GLUE is **not** a point optimizer: it quantifies parameter and prediction
//! uncertainty. It Monte-Carlo samples the parameter space, scores each sample
//! by a goodness measure, keeps the *behavioral* sets (those better than a
//! threshold), and weights them by a pseudo-likelihood. The weighted behavioral
//! ensemble is the posterior: its per-parameter quantiles bound the parameters,
//! and — propagated through the model by the caller — its weighted output
//! quantiles give prediction uncertainty bands.
//!
//! It fits forge's [`Problem`] trait by reading the objective as a **cost**
//! (lower = better, the engine convention): behavioral sets are those with cost
//! below the threshold, and the pseudo-likelihood is `cutoff − cost` (the
//! linear/trapezoidal form), normalized over behavioral sets. A model that
//! maximizes a goodness measure (NSE, KGE) wraps it so the cost is `−goodness`
//! (see [`Maximize`]). Prediction-band propagation lives in the client (which
//! holds the model); forge supplies [`GlueResult::weighted_quantile`] for it.
//!
//! Deterministic for a given seed (each Monte-Carlo sample draws from an
//! independent sub-stream), and identical with or without the `rayon` feature.
//!
//! [`Problem`]: crate::problem::Problem
//! [`Maximize`]: crate::problem::Maximize

use crate::problem::Problem;
use crate::rng::{mix_seed, Rng};

/// How the behavioral threshold is defined.
#[derive(Debug, Clone, Copy)]
pub enum Behavioral {
    /// Keep samples whose cost is at or below this absolute cutoff.
    BelowCost(f64),
    /// Keep the best `fraction` (0–1) of samples by cost.
    TopFraction(f64),
    /// Keep the best `n` samples by cost.
    TopN(usize),
}

/// GLUE configuration.
#[derive(Debug, Clone, Copy)]
pub struct Glue {
    /// Number of Monte-Carlo parameter samples.
    pub samples: usize,
    /// Behavioral-threshold rule.
    pub behavioral: Behavioral,
    /// RNG seed; same seed + same problem ⇒ same result.
    pub seed: u64,
}

impl Default for Glue {
    fn default() -> Self {
        Glue {
            samples: 10_000,
            behavioral: Behavioral::TopFraction(0.1),
            seed: 42,
        }
    }
}

/// One behavioral parameter set with its cost and GLUE weight.
#[derive(Debug, Clone, PartialEq)]
pub struct GlueSample {
    /// Parameter vector.
    pub x: Vec<f64>,
    /// Objective cost (lower = better).
    pub cost: f64,
    /// Normalized GLUE weight (behavioral weights sum to 1).
    pub weight: f64,
}

/// Result of a GLUE run: the weighted behavioral ensemble (the posterior).
#[derive(Debug, Clone)]
pub struct GlueResult {
    /// Behavioral sets, sorted by cost ascending (so `[0]` is the best fit).
    /// Weights sum to 1.
    pub behavioral: Vec<GlueSample>,
    /// The behavioral cost cutoff that was applied.
    pub cutoff: f64,
    /// Total Monte-Carlo samples drawn.
    pub total_samples: usize,
    /// Objective evaluations performed (equals `total_samples`).
    pub evaluations: usize,
}

impl Glue {
    /// Runs GLUE on `problem`. The objective is read as a cost (lower better).
    ///
    /// `problem` must be `Sync` so samples can be evaluated in parallel under
    /// the `rayon` feature; the result is identical either way.
    pub fn run(&self, problem: &(dyn Problem + Sync)) -> GlueResult {
        let bounds = problem.bounds();
        let n_samples = self.samples.max(1);

        // Each sample draws from its own deterministic sub-stream, so parallel
        // and sequential runs agree.
        let sample = |i: usize| -> GlueSample {
            let mut rng = Rng::new(mix_seed(self.seed, i as u64));
            let x: Vec<f64> = bounds
                .iter()
                .map(|&(lo, hi)| rng.uniform_in(lo, hi))
                .collect();
            let c = problem.objective(&x);
            GlueSample {
                x,
                cost: if c.is_finite() { c } else { f64::INFINITY },
                weight: 0.0,
            }
        };

        #[cfg(feature = "rayon")]
        let mut all: Vec<GlueSample> = {
            use rayon::prelude::*;
            (0..n_samples).into_par_iter().map(sample).collect()
        };
        #[cfg(not(feature = "rayon"))]
        let mut all: Vec<GlueSample> = (0..n_samples).map(sample).collect();

        // Sort by cost; the behavioral set is a prefix of this order.
        all.sort_by(|a, b| {
            a.cost
                .partial_cmp(&b.cost)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Resolve the cost cutoff.
        let cutoff = match self.behavioral {
            Behavioral::BelowCost(c) => c,
            Behavioral::TopFraction(frac) => {
                let k = ((frac.clamp(0.0, 1.0) * n_samples as f64).round() as usize)
                    .clamp(1, n_samples);
                all[k - 1].cost
            }
            Behavioral::TopN(nn) => {
                let k = nn.clamp(1, n_samples);
                all[k - 1].cost
            }
        };

        let mut behavioral: Vec<GlueSample> = all
            .into_iter()
            .filter(|s| s.cost.is_finite() && s.cost <= cutoff)
            .collect();

        // Pseudo-likelihood L_i = cutoff − cost_i (≥ 0), normalized. If that is
        // degenerate (all at the cutoff), fall back to uniform weights.
        let raw: Vec<f64> = behavioral
            .iter()
            .map(|s| (cutoff - s.cost).max(0.0))
            .collect();
        let total: f64 = raw.iter().sum();
        let k = behavioral.len();
        for (s, &l) in behavioral.iter_mut().zip(&raw) {
            s.weight = if total > 0.0 {
                l / total
            } else {
                1.0 / k.max(1) as f64
            };
        }

        GlueResult {
            behavioral,
            cutoff,
            total_samples: n_samples,
            evaluations: n_samples,
        }
    }
}

impl GlueResult {
    /// Fraction of samples that were behavioral.
    pub fn acceptance_rate(&self) -> f64 {
        self.behavioral.len() as f64 / self.total_samples.max(1) as f64
    }

    /// Whether any behavioral set was found.
    pub fn is_empty(&self) -> bool {
        self.behavioral.is_empty()
    }

    /// Weight-weighted mean of decision variable `j` over the posterior.
    pub fn parameter_mean(&self, j: usize) -> f64 {
        self.behavioral.iter().map(|s| s.weight * s.x[j]).sum()
    }

    /// Weighted `q`-quantile (0–1) of decision variable `j` over the posterior
    /// — the GLUE parameter uncertainty bound. `NaN` if there are no behavioral
    /// sets.
    pub fn parameter_quantile(&self, j: usize, q: f64) -> f64 {
        let pairs: Vec<(f64, f64)> = self.behavioral.iter().map(|s| (s.x[j], s.weight)).collect();
        weighted_quantile(pairs, q)
    }

    /// Weighted `q`-quantile of a per-behavioral-sample quantity. `values` must
    /// be aligned to [`GlueResult::behavioral`] (one value per behavioral set,
    /// same order). This is the bridge to **prediction** uncertainty bands: run
    /// the model on each behavioral set, collect a model output (e.g. simulated
    /// flow at time `t`) into `values`, and read its 5th/95th quantiles.
    pub fn weighted_quantile(&self, values: &[f64], q: f64) -> f64 {
        let pairs: Vec<(f64, f64)> = values
            .iter()
            .zip(&self.behavioral)
            .map(|(&v, s)| (v, s.weight))
            .collect();
        weighted_quantile(pairs, q)
    }
}

/// Weighted quantile of `(value, weight)` pairs (weights need not be
/// normalized). Returns `NaN` for an empty input.
fn weighted_quantile(mut pairs: Vec<(f64, f64)>, q: f64) -> f64 {
    if pairs.is_empty() {
        return f64::NAN;
    }
    pairs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
    let total: f64 = pairs.iter().map(|p| p.1).sum();
    let target = q.clamp(0.0, 1.0) * total;
    let mut cum = 0.0;
    for (v, w) in &pairs {
        cum += w;
        if cum >= target {
            return *v;
        }
    }
    pairs.last().unwrap().0
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::problem::func;

    #[test]
    fn weighted_quantile_basic() {
        // Equal weights on 0,1,2,3,4: median ~2.
        let pairs: Vec<(f64, f64)> = (0..5).map(|v| (v as f64, 0.2)).collect();
        assert_eq!(weighted_quantile(pairs.clone(), 0.5), 2.0);
        assert_eq!(
            weighted_quantile(pairs, 0.0_f64.max(f64::MIN_POSITIVE)),
            0.0
        );
    }

    #[test]
    fn recovers_behavioral_region() {
        // Cost (x-3)² on [-10,10]; behavioral cost ≤ 1 ⇒ x ∈ [2,4].
        let p = func(vec![(-10.0, 10.0)], |x| (x[0] - 3.0).powi(2));
        let res = Glue {
            samples: 20_000,
            behavioral: Behavioral::BelowCost(1.0),
            seed: 1,
        }
        .run(&p);

        assert!(!res.is_empty());
        // Every behavioral parameter sits in the analytical band [2, 4].
        for s in &res.behavioral {
            assert!((2.0..=4.0).contains(&s.x[0]), "x out of band: {}", s.x[0]);
        }
        // Posterior is centered on the optimum and brackets it.
        assert!((res.parameter_mean(0) - 3.0).abs() < 0.1);
        let lo = res.parameter_quantile(0, 0.05);
        let hi = res.parameter_quantile(0, 0.95);
        assert!(
            lo < 3.0 && hi > 3.0 && lo >= 2.0 && hi <= 4.0,
            "[{lo}, {hi}]"
        );
        // Acceptance ≈ band width / domain = 2/20 = 10%.
        assert!(
            (res.acceptance_rate() - 0.1).abs() < 0.03,
            "acc {}",
            res.acceptance_rate()
        );
    }

    #[test]
    fn top_fraction_keeps_expected_count() {
        let p = func(vec![(0.0, 1.0)], |x| x[0]);
        let res = Glue {
            samples: 1000,
            behavioral: Behavioral::TopFraction(0.2),
            seed: 5,
        }
        .run(&p);
        assert_eq!(res.behavioral.len(), 200);
        // Best fit (lowest cost) is first.
        assert!(res.behavioral[0].cost <= res.behavioral[1].cost);
    }

    #[test]
    fn is_deterministic() {
        let p = func(vec![(-5.0, 5.0), (-5.0, 5.0)], |x| {
            x.iter().map(|v| v * v).sum()
        });
        let cfg = Glue {
            samples: 5000,
            behavioral: Behavioral::TopN(100),
            seed: 9,
        };
        let a = cfg.run(&p);
        let b = cfg.run(&p);
        assert_eq!(a.behavioral, b.behavioral);
        assert_eq!(a.cutoff, b.cutoff);
    }
}