antecedent-stats 0.4.0

//! Weighted least squares and two-stage least squares.
//!
//! SPDX-License-Identifier: MIT OR Apache-2.0

#![allow(clippy::cast_precision_loss, clippy::many_single_char_names)]

use crate::error::StatsError;
use crate::linalg::{DenseLinearAlgebra, LeastSquaresFit, LeastSquaresWorkspace};

/// Fit weighted least squares by row-scaling with `sqrt(weight)`.
///
/// `weights` length = `nrows`. Weights must be finite and non-negative.
/// Zero is allowed and drops that row (scale factor 0); negatives and non-finite
/// values are rejected as upstream errors rather than coerced.
///
/// # Errors
///
/// Shape mismatch, invalid weights, or backend failure.
pub fn fit_wls(
    x_colmajor: &[f64],
    nrows: usize,
    ncols: usize,
    y: &[f64],
    weights: &[f64],
    backend: &impl DenseLinearAlgebra,
    workspace: &mut LeastSquaresWorkspace,
) -> Result<LeastSquaresFit, StatsError> {
    if y.len() != nrows || weights.len() != nrows {
        return Err(StatsError::Shape { message: "y/weights length != nrows" });
    }
    if x_colmajor.len() < nrows.saturating_mul(ncols) {
        return Err(StatsError::Shape { message: "X buffer too short" });
    }
    let mut x_w = vec![0.0; nrows * ncols];
    let mut y_w = vec![0.0; nrows];
    for r in 0..nrows {
        let wr = weights[r];
        if !(wr.is_finite() && wr >= 0.0) {
            return Err(StatsError::Shape {
                message: "WLS weights must be finite and non-negative",
            });
        }
        let w = wr.sqrt();
        y_w[r] = y[r] * w;
        for c in 0..ncols {
            x_w[c * nrows + r] = x_colmajor[c * nrows + r] * w;
        }
    }
    backend.least_squares(&x_w, nrows, ncols, &y_w, workspace)
}

/// Weak-instrument diagnostic: joint significance of the excluded instruments in the
/// first-stage regression of the endogenous variable on `[instruments | exogenous]`.
///
/// This is purely informational — [`fit_2sls`] never hard-fails on a weak instrument
/// (only on an EXACTLY degenerate one, at the caller level). A caller may compare
/// `f_statistic` against a rule-of-thumb threshold itself (e.g. the Staiger-Stock /
/// Stock-Yogo guideline of 10 for a single endogenous regressor); a future release could
/// add an opt-in hard threshold on top of this diagnostic.
#[derive(Clone, Debug, PartialEq)]
pub struct FirstStageDiagnostics {
    /// F-statistic for the joint null that all excluded-instrument coefficients are zero,
    /// controlling for the included exogenous regressors.
    pub f_statistic: f64,
    /// Numerator degrees of freedom (number of excluded instruments).
    pub df1: usize,
    /// Denominator degrees of freedom (first-stage residual df, `n − (df1 + x_ncols)`).
    pub df2: usize,
    /// Partial R² attributable to the excluded instruments:
    /// `(RSS_restricted − RSS_full) / RSS_restricted`.
    pub partial_r2: f64,
}

/// Result of two-stage least squares.
#[derive(Clone, Debug)]
pub struct TwoSlsFit {
    /// First-stage coefficients (full instrument set `[Z | X]` → endogenous).
    pub first_stage: LeastSquaresFit,
    /// Second-stage coefficients (fitted endogenous + covariates → outcome).
    pub second_stage: LeastSquaresFit,
    /// Fitted endogenous values used in stage 2.
    pub fitted_endogenous: Vec<f64>,
    /// Structural residual sum of squares `‖y − Tβ̂ − Xγ̂‖²` using the *actual*
    /// endogenous values (not the fitted ones). This is the σ̂² numerator for the
    /// conventional 2SLS analytic standard error.
    pub structural_rss: f64,
    /// Structural residuals `e_i = y_i − T_i β̂ − X_i γ̂` (same as RSS terms).
    pub structural_residuals: Vec<f64>,
    /// Weak-instrument diagnostic for the excluded instrument set (see
    /// [`FirstStageDiagnostics`]).
    pub first_stage_diagnostics: FirstStageDiagnostics,
}

/// Two-stage least squares.
///
/// Stage 1: `endogenous ~ [instruments | exogenous]` — the full instrument set. Included
/// exogenous regressors instrument themselves, so `exogenous` (column-major, may be
/// empty / intercept-only) is appended to the excluded instruments in the first-stage
/// design. Pass the intercept in exactly one of the two blocks.
/// Stage 2: `y ~ [fitted_endogenous | exogenous]`.
///
/// Convention: stage-2 design is `[fitted_T | X]` with `1 + x_ncols` columns; the treatment
/// coefficient is `second_stage.coefficients[0]`.
///
/// # Errors
///
/// Shape mismatch or backend failure.
#[allow(clippy::too_many_arguments)]
pub fn fit_2sls(
    instruments_colmajor: &[f64],
    z_nrows: usize,
    z_ncols: usize,
    endogenous: &[f64],
    exogenous_colmajor: &[f64],
    x_ncols: usize,
    y: &[f64],
    backend: &impl DenseLinearAlgebra,
    workspace: &mut LeastSquaresWorkspace,
) -> Result<TwoSlsFit, StatsError> {
    if endogenous.len() != z_nrows || y.len() != z_nrows {
        return Err(StatsError::Shape { message: "endogenous/y length != nrows" });
    }
    if exogenous_colmajor.len() < z_nrows.saturating_mul(x_ncols) {
        return Err(StatsError::Shape { message: "exogenous buffer too short" });
    }
    // Full instrument set: excluded instruments plus included exogenous regressors.
    let stage1_ncols = z_ncols + x_ncols;
    let mut x1 = vec![0.0; z_nrows * stage1_ncols];
    x1[..z_nrows * z_ncols].copy_from_slice(&instruments_colmajor[..z_nrows * z_ncols]);
    x1[z_nrows * z_ncols..].copy_from_slice(&exogenous_colmajor[..z_nrows * x_ncols]);
    let first_stage = backend.least_squares(&x1, z_nrows, stage1_ncols, endogenous, workspace)?;
    let first_stage_diagnostics = first_stage_f_test(
        z_nrows,
        z_ncols,
        x_ncols,
        exogenous_colmajor,
        endogenous,
        first_stage.rss,
        backend,
        workspace,
    )?;
    let mut fitted = vec![0.0; z_nrows];
    for r in 0..z_nrows {
        let mut pred = 0.0;
        for c in 0..stage1_ncols {
            pred += x1[c * z_nrows + r] * first_stage.coefficients[c];
        }
        fitted[r] = pred;
    }
    let stage2_ncols = 1 + x_ncols;
    let mut x2 = vec![0.0; z_nrows * stage2_ncols];
    for r in 0..z_nrows {
        x2[r] = fitted[r];
        for c in 0..x_ncols {
            x2[(1 + c) * z_nrows + r] = exogenous_colmajor[c * z_nrows + r];
        }
    }
    let second_stage = backend.least_squares(&x2, z_nrows, stage2_ncols, y, workspace)?;
    // Structural residuals evaluate the second-stage coefficients at the ACTUAL
    // endogenous values; `second_stage.rss` uses fitted T and is not σ̂².
    let mut structural_residuals = vec![0.0; z_nrows];
    let mut structural_rss = 0.0;
    for r in 0..z_nrows {
        let mut pred = second_stage.coefficients[0] * endogenous[r];
        for c in 0..x_ncols {
            pred += exogenous_colmajor[c * z_nrows + r] * second_stage.coefficients[1 + c];
        }
        let e = y[r] - pred;
        structural_residuals[r] = e;
        structural_rss += e * e;
    }
    Ok(TwoSlsFit {
        first_stage,
        second_stage,
        fitted_endogenous: fitted,
        structural_rss,
        structural_residuals,
        first_stage_diagnostics,
    })
}

/// Partial F-test for the excluded-instrument block of a first-stage regression:
/// compares the unrestricted fit `[instruments | exogenous]` (`full_rss`, already
/// computed by the caller) against a restricted fit on `exogenous` alone (or the
/// zero model when `x_ncols == 0`, i.e. `RSS = Σ endogenous²`).
#[allow(clippy::too_many_arguments)]
fn first_stage_f_test(
    nrows: usize,
    z_ncols: usize,
    x_ncols: usize,
    exogenous_colmajor: &[f64],
    endogenous: &[f64],
    full_rss: f64,
    backend: &impl DenseLinearAlgebra,
    workspace: &mut LeastSquaresWorkspace,
) -> Result<FirstStageDiagnostics, StatsError> {
    let restricted_rss = if x_ncols == 0 {
        endogenous.iter().map(|v| v * v).sum::<f64>()
    } else {
        let restricted =
            backend.least_squares(exogenous_colmajor, nrows, x_ncols, endogenous, workspace)?;
        restricted.rss
    };
    let df1 = z_ncols;
    let df2 = nrows.saturating_sub(z_ncols + x_ncols);
    let f_statistic = if df1 == 0 || df2 == 0 {
        f64::NAN
    } else if full_rss > 0.0 {
        ((restricted_rss - full_rss) / df1 as f64) / (full_rss / df2 as f64)
    } else {
        f64::INFINITY
    };
    let partial_r2 = if restricted_rss > 0.0 {
        ((restricted_rss - full_rss) / restricted_rss).max(0.0)
    } else {
        0.0
    };
    Ok(FirstStageDiagnostics { f_statistic, df1, df2, partial_r2 })
}

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

    #[test]
    fn wls_matches_ols_with_unit_weights() {
        let n = 20usize;
        let mut x = vec![0.0; n * 2];
        let mut y = vec![0.0; n];
        for i in 0..n {
            x[i] = 1.0;
            x[n + i] = i as f64;
            y[i] = 1.0 + 2.0 * (i as f64);
        }
        let w = vec![1.0; n];
        let mut ws = LeastSquaresWorkspace::default();
        let ols = FaerBackend.least_squares(&x, n, 2, &y, &mut ws).unwrap();
        let wls = fit_wls(&x, n, 2, &y, &w, &FaerBackend, &mut ws).unwrap();
        assert!((ols.coefficients[0] - wls.coefficients[0]).abs() < 1e-10);
        assert!((ols.coefficients[1] - wls.coefficients[1]).abs() < 1e-10);
    }

    #[test]
    fn wls_rejects_negative_and_nonfinite_weights() {
        let n = 4usize;
        let x = vec![1.0; n * 2];
        let y = vec![1.0, 2.0, 3.0, 4.0];
        let mut ws = LeastSquaresWorkspace::default();
        for bad in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY, -1.0] {
            let mut w = vec![1.0; n];
            w[1] = bad;
            let err = fit_wls(&x, n, 2, &y, &w, &FaerBackend, &mut ws).unwrap_err();
            assert_eq!(
                err,
                StatsError::Shape { message: "WLS weights must be finite and non-negative" }
            );
        }
    }

    #[test]
    fn wls_zero_weight_drops_row() {
        // Zero weight is an explicit drop-row policy: omitting the zero-weight
        // observation must match fitting on the complementary subset with unit weights.
        let n = 4usize;
        let mut x_full = vec![0.0; n * 2];
        let y_full = [1.0, 10.0, 3.0, 4.0];
        for i in 0..n {
            x_full[i] = 1.0;
            x_full[n + i] = i as f64;
        }
        let w = [1.0, 0.0, 1.0, 1.0];
        let mut ws = LeastSquaresWorkspace::default();
        let wls = fit_wls(&x_full, n, 2, &y_full, &w, &FaerBackend, &mut ws).unwrap();

        let n_sub = 3usize;
        let mut x_sub = vec![0.0; n_sub * 2];
        let y_sub = [1.0, 3.0, 4.0];
        for (j, i) in [0usize, 2, 3].into_iter().enumerate() {
            x_sub[j] = 1.0;
            x_sub[n_sub + j] = i as f64;
        }
        let ols = FaerBackend.least_squares(&x_sub, n_sub, 2, &y_sub, &mut ws).unwrap();
        assert!((ols.coefficients[0] - wls.coefficients[0]).abs() < 1e-10);
        assert!((ols.coefficients[1] - wls.coefficients[1]).abs() < 1e-10);
    }

    #[test]
    fn twosls_recovers_just_identified() {
        // Z → T → Y with no confounding on Z→Y; T = Z + e, Y = 2T + u.
        // Instruments carry only the excluded Z column; the intercept lives in the
        // exogenous block (stage 1 uses [Z | 1], stage 2 uses [fitted_T | 1]).
        let n = 200usize;
        let mut z = vec![0.0; n];
        let mut t = vec![0.0; n];
        let mut y = vec![0.0; n];
        let mut x = vec![0.0; n]; // intercept only exogenous
        for i in 0..n {
            let zi = (i as f64) / n as f64 - 0.5;
            z[i] = zi;
            t[i] = zi + 0.01 * ((i % 7) as f64 - 3.0);
            y[i] = 2.0 * t[i] + 0.01 * ((i % 5) as f64 - 2.0);
            x[i] = 1.0;
        }
        let mut ws = LeastSquaresWorkspace::default();
        let fit = fit_2sls(&z, n, 1, &t, &x, 1, &y, &FaerBackend, &mut ws).unwrap();
        assert!((fit.second_stage.coefficients[0] - 2.0).abs() < 0.05);
        assert!(fit.structural_rss <= fit.second_stage.rss);
    }

    #[test]
    fn first_stage_diagnostics_strong_vs_weak_instrument() {
        // Same DGP shape (Z -> T -> Y, intercept-only exogenous block), varied only by how
        // strongly Z drives T. Strong: T = Z + small noise (first stage explains most of
        // T's variance). Weak: T = 0.001*Z + large noise (first stage explains ~none of
        // it). The F-statistic must separate these cases sharply.
        let n = 200usize;
        let mut z = vec![0.0; n];
        let mut x = vec![0.0; n];
        for i in 0..n {
            z[i] = (i as f64) / n as f64 - 0.5;
            x[i] = 1.0;
        }

        let mut t_strong = vec![0.0; n];
        let mut y_strong = vec![0.0; n];
        let mut t_weak = vec![0.0; n];
        let mut y_weak = vec![0.0; n];
        for i in 0..n {
            let small_noise = 0.01 * ((i % 7) as f64 - 3.0);
            let large_noise = 1.0 * ((i % 7) as f64 - 3.0);
            t_strong[i] = z[i] + small_noise;
            y_strong[i] = 2.0 * t_strong[i] + 0.01 * ((i % 5) as f64 - 2.0);
            t_weak[i] = 0.001 * z[i] + large_noise;
            y_weak[i] = 2.0 * t_weak[i] + 0.01 * ((i % 5) as f64 - 2.0);
        }

        let mut ws = LeastSquaresWorkspace::default();
        let strong =
            fit_2sls(&z, n, 1, &t_strong, &x, 1, &y_strong, &FaerBackend, &mut ws).unwrap();
        let weak = fit_2sls(&z, n, 1, &t_weak, &x, 1, &y_weak, &FaerBackend, &mut ws).unwrap();

        assert_eq!(strong.first_stage_diagnostics.df1, 1);
        assert_eq!(strong.first_stage_diagnostics.df2, n - 2);
        assert_eq!(weak.first_stage_diagnostics.df1, 1);
        assert_eq!(weak.first_stage_diagnostics.df2, n - 2);

        assert!(
            strong.first_stage_diagnostics.f_statistic > 1000.0,
            "strong F={}",
            strong.first_stage_diagnostics.f_statistic
        );
        assert!(
            weak.first_stage_diagnostics.f_statistic < 5.0,
            "weak F={}",
            weak.first_stage_diagnostics.f_statistic
        );
        assert!(strong.first_stage_diagnostics.partial_r2 > 0.9);
        assert!(weak.first_stage_diagnostics.partial_r2 < 0.1);
    }

    #[test]
    fn twosls_first_stage_includes_exogenous_regressors() {
        // Y = 2T + 1.5X + u with X correlated with T beyond Z; projecting T on [1, Z]
        // only (the old first stage) is inconsistent here.
        let n = 400usize;
        let mut z = vec![0.0; n];
        let mut t = vec![0.0; n];
        let mut y = vec![0.0; n];
        let mut x = vec![0.0; n * 2];
        for i in 0..n {
            let zi = (i as f64) / n as f64 - 0.5;
            let xi = ((i % 13) as f64 - 6.0) / 6.0;
            let e = 0.01 * ((i % 7) as f64 - 3.0);
            z[i] = zi;
            t[i] = zi + 0.8 * xi + e;
            y[i] = 2.0 * t[i] + 1.5 * xi + 0.01 * ((i % 5) as f64 - 2.0);
            x[i] = 1.0;
            x[n + i] = xi;
        }
        let mut ws = LeastSquaresWorkspace::default();
        let fit = fit_2sls(&z, n, 1, &t, &x, 2, &y, &FaerBackend, &mut ws).unwrap();
        assert!(
            (fit.second_stage.coefficients[0] - 2.0).abs() < 0.05,
            "beta_T={}",
            fit.second_stage.coefficients[0]
        );
        assert!((fit.second_stage.coefficients[2] - 1.5).abs() < 0.05);
    }
}