r2rs-stats 0.1.1

Statistics programming for Rust based on R's stats package
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

// "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 nalgebra::DMatrix;
use num_traits::ToPrimitive;
use r2rs_nmath::{distribution::TBuilder, traits::Distribution};
use strafe_type::{Alpha64, FloatConstraint, Rational64, Real64};

use crate::funcs::residual_sum_of_squares;

pub fn predict(x: &DMatrix<f64>, b: &DMatrix<f64>) -> DMatrix<f64> {
    x * b
}

fn interval_inner<A: Into<Alpha64>>(
    x: &DMatrix<f64>,
    y: &DMatrix<f64>,
    b: &DMatrix<f64>,
    w: &DMatrix<f64>,
    alpha: A,
    scale: Option<f64>,
    df: Option<Rational64>,
) -> (DMatrix<f64>, Real64, f64) {
    let mut sqrt_w = w.to_owned();
    sqrt_w
        .column_iter_mut()
        .for_each(|mut r| r.iter_mut().for_each(|w| *w = w.sqrt()));

    let mut weighted_x = x.clone();
    weighted_x
        .column_iter_mut()
        .for_each(|mut row| row.component_mul_assign(&sqrt_w));

    let degrees_of_freedom = x.shape().0 - b.shape().0;

    let residual_variance = if let Some(scale) = scale {
        scale.powi(2)
    } else {
        let residual_sum_square = residual_sum_of_squares(x, y, b);
        residual_sum_square / degrees_of_freedom as f64
    };

    let mut xr_inverse = x.clone().qr().q();

    let ip =
        xr_inverse.component_mul(&xr_inverse) * DMatrix::repeat(x.ncols(), 1, residual_variance);

    let mut t_builder = TBuilder::new();
    if let Some(df) = df {
        // Used by MASS's RLM
        t_builder.with_df(df.to_f64().unwrap()).unwrap();
    } else {
        t_builder.with_df(degrees_of_freedom).unwrap();
    }
    let t = t_builder.build();
    let t_fraction = t.quantile(alpha.into().unwrap() / 2.0, true);

    (ip, t_fraction, residual_variance)
}

pub fn confidence_interval<A: Into<Alpha64>>(
    x: &DMatrix<f64>,
    y: &DMatrix<f64>,
    b: &DMatrix<f64>,
    w: &DMatrix<f64>,
    alpha: A,
    scale: Option<f64>,
    df: Option<Rational64>,
) -> (DMatrix<f64>, DMatrix<f64>) {
    let (mut ip, t_fraction, _) = interval_inner(x, y, b, w, alpha, scale, df);

    let predict = predict(x, b);

    (
        DMatrix::from_iterator(
            predict.len(),
            1,
            predict
                .iter()
                .zip(ip.iter())
                .map(|(base, ip)| base + (t_fraction.unwrap() * ip.sqrt())),
        ),
        DMatrix::from_iterator(
            predict.len(),
            1,
            predict
                .iter()
                .zip(ip.iter())
                .map(|(base, ip)| base - (t_fraction.unwrap() * ip.sqrt())),
        ),
    )
}

pub fn prediction_interval<A: Into<Alpha64>>(
    x: &DMatrix<f64>,
    y: &DMatrix<f64>,
    b: &DMatrix<f64>,
    w: &DMatrix<f64>,
    alpha: A,
    scale: Option<f64>,
    df: Option<Rational64>,
) -> (DMatrix<f64>, DMatrix<f64>) {
    let (mut ip, t_fraction, residual_variance) = interval_inner(x, y, b, w, alpha, scale, df);

    let predict = predict(x, b);

    (
        DMatrix::from_iterator(
            predict.len(),
            1,
            predict
                .iter()
                .zip(ip.iter())
                .map(|(base, ip)| base + (t_fraction.unwrap() * (*ip + residual_variance).sqrt())),
        ),
        DMatrix::from_iterator(
            predict.len(),
            1,
            predict
                .iter()
                .zip(ip.iter())
                .map(|(base, ip)| base - (t_fraction.unwrap() * (*ip + residual_variance).sqrt())),
        ),
    )
}