r2rs-stats 0.1.1

// "Whatever you do, work at it with all your heart, as working for the Lord,
// not for human masters, since you know that you will receive an inheritance
// from the Lord as a reward. It is the Lord Christ you are serving."
// (Col 3:23-24)

use log::{error, warn};
use nalgebra::DMatrix;
use num_traits::Float;

use crate::{
    funcs::linear_regression_weighted_no_intercept,
    regression::glm::{family::core::Family, link::core::Link},
};

fn glm_control(epsilon: f64, max_iter: u64, trace: bool) -> (f64, u64, bool) {
    if epsilon <= 0.0 {
        error!("value of 'epsilon' must be > 0");
    }
    if max_iter == 0 {
        error!("maximum number of iterations must be > 0");
    }
    (epsilon, max_iter, trace)
}

pub fn generalized_linear_model<L: Link, F: Family<L>>(
    y: &DMatrix<f64>,
    x: &DMatrix<f64>,
    weights: &DMatrix<f64>,
    family: &F,
) -> (
    DMatrix<f64>,
    f64,
    f64,
    f64,
    DMatrix<f64>,
    DMatrix<f64>,
    DMatrix<f64>,
) {
    let mut y = y.to_owned();
    let x = x.to_owned().insert_column(0, 1.0);
    let mut weights = weights.to_owned();

    let mut final_weights = weights.clone();

    let mut ret = None;
    let mut dev = f64::infinity();

    let (epsilon, max_iter, _trace) = glm_control(1e-8, 25, false);
    let _conv = false;
    let nobs = y.nrows();
    let nvars = x.ncols();

    let (n, mu_start) = family.initialize(&mut y, &mut weights);

    let mut eta = family.link().link_func(&mu_start);
    let mut mu = family.link().link_inverse(&eta);

    if nvars == 0 {
        eta = DMatrix::from_iterator(y.shape().0, y.shape().1, (0..nobs).map(|_| 0.0));
        if !family.link().valid_eta(&eta) {
            error!("invalid linear predictor values in empty model");
        }
        if !family.valid_mu(&mu) {
            error!("invalid fitted means in empty model");
        }
        dev = family
            .residual_deviance(&y, &mu, &weights)
            .into_iter()
            .sum::<f64>();
        let mu_eta = family.link().mu_eta(&eta);
        let mu_var = family.variance(&mu);
        final_weights = DMatrix::from_iterator(
            weights.nrows(),
            weights.ncols(),
            weights.iter().zip(mu_eta.iter()).zip(mu_var.iter()).map(
                |((weights_i, mu_eta_i), mu_var_i)| {
                    ((weights_i * mu_eta_i.powi(2)) / mu_var_i).sqrt()
                },
            ),
        );
    } else {
        let mut coef_old = None;
        if !(family.valid_mu(&mu) && family.link().valid_eta(&eta)) {
            error!("cannot find valid starting values: please specify some");
        }
        let mut deviance_old = family
            .residual_deviance(&y, &mu, &weights)
            .into_iter()
            .sum::<f64>();
        let mut boundary = false;
        let mut conv = false;

        // Iteratively Reweighted Least Squares (IRLS)
        for iter in 0..max_iter {
            let mut bad = weights
                .iter()
                .enumerate()
                .filter_map(|(i, &weight)| if weight > 0.0 { None } else { Some(i) })
                .collect::<Vec<_>>();
            let var_mu = family.variance(&mu);
            let mut var_mu_good = var_mu.clone().remove_rows_at(&bad);
            if var_mu_good.iter().any(|&vm| vm.is_nan()) {
                error!("NAs in V(mu)");
            }
            if var_mu_good.iter().any(|&vm| vm == 0.0) {
                error!("0s in V(mu)");
            }
            let mu_eta_val = family.link().mu_eta(&eta);
            let mut mu_eta_val_good = mu_eta_val.clone().remove_rows_at(&bad);
            if mu_eta_val_good.iter().any(|mu_eta| mu_eta.is_nan()) {
                error!("NAs in d(mu)/d(eta)");
            }
            bad = weights
                .iter()
                .zip(mu_eta_val.iter())
                .enumerate()
                .filter_map(|(i, (&weight, &mu_eta))| {
                    if weight > 0.0 && mu_eta != 0.0 {
                        None
                    } else {
                        Some(i)
                    }
                })
                .collect::<Vec<_>>();
            if bad.len() == y.nrows() {
                conv = false;
                warn!("no observations informative at iteration {}", iter);
                break;
            }

            let y_good = y.clone().remove_rows_at(&bad);
            let x_good = x.clone().remove_rows_at(&bad);
            let weights_good = weights.clone().remove_rows_at(&bad);
            let eta_good = eta.clone().remove_rows_at(&bad);
            let mu_good = mu.clone().remove_rows_at(&bad);
            var_mu_good = var_mu.remove_rows_at(&bad);
            mu_eta_val_good = mu_eta_val.remove_rows_at(&bad);

            let z = DMatrix::from_iterator(
                y_good.nrows(),
                y_good.ncols(),
                eta_good
                    .iter()
                    .zip(y_good.iter())
                    .zip(mu_good.iter())
                    .zip(mu_eta_val_good.iter())
                    .map(|(((&eta, &y), &mu), &mu_eta)| eta + (y - mu) / mu_eta),
            );
            final_weights = DMatrix::from_iterator(
                weights_good.nrows(),
                weights_good.ncols(),
                weights_good
                    .iter()
                    .zip(mu_eta_val_good.iter())
                    .zip(var_mu_good.iter())
                    .map(|((weight, mu_eta), v_mu)| ((weight * mu_eta.powi(2)) / v_mu).sqrt()),
            );

            let mut coefs = linear_regression_weighted_no_intercept(&x_good, &z, &final_weights);
            if coefs.iter().any(|&c| !c.is_finite()) {
                error!("non-finite coefficients at iteration {}", iter);
            }
            if nobs < coefs.nrows() - 1 {
                error!(
                    "X matrix has rank {}, but only {} observations",
                    coefs.nrows() - 1,
                    nobs
                )
            }

            let mut start = coefs.clone();
            eta = x.clone() * start.clone();
            mu = family.link().link_inverse(&eta);
            dev = family
                .residual_deviance(&y, &mu, &weights)
                .into_iter()
                .sum::<f64>();

            boundary = false;
            if !dev.is_finite() {
                if coef_old.is_none() {
                    error!("no valid set of coefficients has been found: please supply starting values");
                }
                warn!("step size truncated due to divergence");
                let mut ii = 1;
                while !dev.is_finite() {
                    if ii > max_iter {
                        error!("inner loop 1; cannot correct step size");
                    }
                    ii += 1;
                    start = (start + coef_old.clone().unwrap()) / 2.0;
                    eta = x.clone() * start.clone();
                    mu = family.link().link_inverse(&eta);
                    dev = family.residual_deviance(&y, &mu, &weights).sum();
                }
                boundary = true;
            }

            if !family.link().valid_eta(&eta) && family.valid_mu(&mu) {
                if coef_old.is_none() {
                    error!("no valid set of coefficients has been found: please supply starting values");
                }
                warn!("step size truncated: out of bounds");
                let mut ii = 1;
                while !family.link().valid_eta(&eta) && family.valid_mu(&mu) {
                    if ii > max_iter {
                        error!("inner loop 2; cannot correct step size")
                    }
                    ii += 1;
                    start = (start.clone() + coef_old.clone().unwrap()) / 2.0;
                    eta = x.clone() * start.clone();
                    mu = family.link().link_inverse(&eta);
                }
                boundary = true;
                dev = family.residual_deviance(&y, &mu, &weights).iter().sum();
            }

            if (dev - deviance_old).abs() / (0.1 + dev.abs()) < epsilon {
                conv = true;
                coefs = start;
                ret = Some(coefs);
                break;
            } else {
                deviance_old = dev;
                coefs = start.clone();
                coef_old = Some(start);
                ret = Some(coefs)
            }
        }

        if !conv {
            warn!("glm.fit: algorithm did not converge");
        }
        if boundary {
            warn!("glm.fit: algorithm stopped at boundary value");
        }
    }

    let wtd_mu_single = (weights.iter().zip(y.iter()).map(|(w_i, y_i)| w_i * y_i)).sum::<f64>()
        / weights.iter().sum::<f64>();
    let mut wtd_mu = mu.clone();
    wtd_mu.iter_mut().for_each(|wtd_mu| *wtd_mu = wtd_mu_single);
    let null_dev = family
        .residual_deviance(&y, &wtd_mu, &weights)
        .iter()
        .sum::<f64>();
    let aic = family.aic(
        y.as_slice(),
        n.as_slice(),
        mu.as_slice(),
        weights.as_slice(),
        dev,
    ) + (2.0 * x.ncols() as f64);

    final_weights.iter_mut().for_each(|w_i| *w_i = w_i.powi(2));

    (ret.unwrap(), aic, null_dev, dev, final_weights, eta, mu)
}