1use std::{convert::identity, io::Write, ops::SubAssign};
2
3use faer::{
4 diag::Diag,
5 get_global_parallelism,
6 linalg::solvers::{DenseSolveCore, Solve},
7 mat::AsMatRef,
8 Col, ColMut, ColRef, Mat, MatMut, MatRef, RowMut, Side,
9};
10use pulp::{Arch, Simd, WithSimd};
11use rand_distr::{Distribution, StandardNormal};
12use rayon::{iter::IntoParallelIterator, prelude::*};
13use statrs::distribution::{ContinuousCDF, StudentsT};
14use tracing::{debug, error, trace, warn};
15
16use crate::{mean, variance};
17
18pub static mut DISABLE_PREDICTED: bool = false;
19
20pub fn should_disable_predicted() -> bool {
21 if cfg!(test) {
22 return false;
23 }
24 if unsafe { DISABLE_PREDICTED } {
26 return true;
27 }
28 let enabled = std::env::var("LMUTILS_ENABLE_PREDICTED").is_ok();
29 let disabled = std::env::var("LMUTILS_DISABLE_PREDICTED").is_ok();
30 if disabled {
32 return true;
33 }
34 if enabled {
36 return false;
37 }
38 true
39}
40
41#[derive(Debug, Clone)]
42pub struct R2 {
43 pub(crate) r2: f64,
44 pub(crate) adj_r2: f64,
45 pub(crate) predicted: Vec<f64>,
46 pub(crate) betas: Vec<f64>,
47 pub(crate) data: Option<String>,
48 pub(crate) outcome: Option<String>,
49 pub(crate) n: u32,
50 pub(crate) m: u32,
51}
52
53impl R2 {
54 #[inline]
55 #[cfg_attr(coverage_nightly, coverage(off))]
56 pub fn r2(&self) -> f64 {
57 self.r2
58 }
59
60 #[inline]
61 #[cfg_attr(coverage_nightly, coverage(off))]
62 pub fn adj_r2(&self) -> f64 {
63 self.adj_r2
64 }
65
66 #[inline]
67 #[cfg_attr(coverage_nightly, coverage(off))]
68 pub fn predicted(&self) -> &[f64] {
69 &self.predicted
70 }
71
72 #[inline]
73 #[cfg_attr(coverage_nightly, coverage(off))]
74 pub fn betas(&self) -> &[f64] {
75 &self.betas
76 }
77
78 #[inline]
79 #[cfg_attr(coverage_nightly, coverage(off))]
80 pub fn data(&self) -> Option<&str> {
81 self.data.as_deref()
82 }
83
84 #[inline]
85 #[cfg_attr(coverage_nightly, coverage(off))]
86 pub fn set_data(&mut self, data: String) {
87 self.data = Some(data);
88 }
89
90 #[inline]
91 #[cfg_attr(coverage_nightly, coverage(off))]
92 pub fn outcome(&self) -> Option<&str> {
93 self.outcome.as_deref()
94 }
95
96 #[inline]
97 #[cfg_attr(coverage_nightly, coverage(off))]
98 pub fn set_outcome(&mut self, outcome: String) {
99 self.outcome = Some(outcome);
100 }
101
102 #[inline]
103 #[cfg_attr(coverage_nightly, coverage(off))]
104 pub fn n(&self) -> u32 {
105 self.n
106 }
107
108 #[inline]
109 #[cfg_attr(coverage_nightly, coverage(off))]
110 pub fn m(&self) -> u32 {
111 self.m
112 }
113}
114
115#[tracing::instrument(skip(data, outcomes))]
116pub fn get_r2s(data: MatRef<f64>, outcomes: MatRef<f64>) -> Vec<R2> {
117 debug!("Calculating R2s");
118 let n = data.nrows();
119 let m = data.ncols();
120
121 let c_all = data.transpose() * outcomes;
122 trace!("Computed c_all");
123 let c_matrix = data.transpose() * data;
124 let betas = match c_matrix.llt(Side::Lower) {
125 Ok(chol) => chol.solve(c_all),
126 Err(_) => {
127 warn!("Using pseudo inverse");
128 c_matrix
129 .as_mat_ref()
130 .thin_svd()
131 .expect("could not compute thin SVD for pseudoinverse")
132 .pseudoinverse()
133 * c_all
134 },
135 };
136
137 debug!("Calculated betas");
138
139 let r2s = (0..outcomes.ncols())
140 .into_par_iter()
141 .map(|i| {
142 let betas = betas.col(i);
143 let mut predicted = (data * betas)
144 .try_as_col_major()
145 .unwrap()
146 .as_slice()
147 .to_vec();
148 let actual = outcomes.col(i).try_as_col_major().unwrap().as_slice();
149 let r2 = R2Simd::new(actual, &predicted).calculate();
150 let adj_r2 = calculate_adj_r2(r2, n, m);
151 let mut betas = betas.try_as_col_major().unwrap().as_slice().to_vec();
152 if should_disable_predicted() {
153 predicted = Vec::new();
154 betas = Vec::new();
155 }
156 R2 {
157 r2,
158 adj_r2,
159 predicted,
160 betas,
161 outcome: None,
162 data: None,
163 n: n as u32,
164 m: m as u32,
165 }
166 })
167 .collect();
168 debug!("Calculated R2s");
169 r2s
170}
171
172#[derive(Debug, Clone)]
173pub struct PValue {
174 p_value: f64,
175 beta: f64,
176 intercept: f64,
177 pub(crate) data: Option<String>,
178 pub(crate) data_column: Option<u32>,
179 pub(crate) outcome: Option<String>,
180}
181
182impl PValue {
183 #[inline]
184 #[cfg_attr(coverage_nightly, coverage(off))]
185 pub fn p_value(&self) -> f64 {
186 self.p_value
187 }
188
189 #[inline]
190 #[cfg_attr(coverage_nightly, coverage(off))]
191 pub fn beta(&self) -> f64 {
192 self.beta
193 }
194
195 #[inline]
196 #[cfg_attr(coverage_nightly, coverage(off))]
197 pub fn intercept(&self) -> f64 {
198 self.intercept
199 }
200
201 #[inline]
202 #[cfg_attr(coverage_nightly, coverage(off))]
203 pub fn data(&self) -> Option<&str> {
204 self.data.as_deref()
205 }
206
207 #[inline]
208 #[cfg_attr(coverage_nightly, coverage(off))]
209 pub fn set_data(&mut self, data: String) {
210 self.data = Some(data);
211 }
212
213 #[inline]
214 #[cfg_attr(coverage_nightly, coverage(off))]
215 pub fn data_column(&self) -> Option<u32> {
216 self.data_column
217 }
218
219 #[inline]
220 #[cfg_attr(coverage_nightly, coverage(off))]
221 pub fn set_data_column(&mut self, data_column: u32) {
222 self.data_column = Some(data_column);
223 }
224
225 #[inline]
226 #[cfg_attr(coverage_nightly, coverage(off))]
227 pub fn outcome(&self) -> Option<&str> {
228 self.outcome.as_deref()
229 }
230
231 #[inline]
232 #[cfg_attr(coverage_nightly, coverage(off))]
233 pub fn set_outcome(&mut self, outcome: String) {
234 self.outcome = Some(outcome);
235 }
236}
237
238#[tracing::instrument(skip(xs, ys))]
239pub fn p_value(xs: &[f64], ys: &[f64]) -> PValue {
240 debug!("Calculating p-values");
241 let mut x = Mat::new();
242 x.resize_with(
243 xs.len(),
244 2,
245 #[inline(always)]
246 |i, j| if j == 0 { xs[i] } else { 1.0 },
247 );
248 let y: MatRef<'_, f64> = MatRef::from_column_major_slice(ys, ys.len(), 1);
249 let c_all = x.transpose() * y;
250 let mut c_matrix = faer::Mat::zeros(2, 2);
251 faer::linalg::matmul::triangular::matmul(
252 c_matrix.as_mut(),
253 faer::linalg::matmul::triangular::BlockStructure::TriangularLower,
254 faer::Accum::Replace,
255 x.transpose(),
256 faer::linalg::matmul::triangular::BlockStructure::Rectangular,
257 &x,
258 faer::linalg::matmul::triangular::BlockStructure::Rectangular,
259 1.0,
260 get_global_parallelism(),
261 );
262 let chol = c_matrix.llt(Side::Lower).unwrap();
263 let inv_matrix = chol.solve(Mat::<f64>::identity(2, 2));
264 let betas = chol.solve(c_all);
265 let m = betas.get(0, 0);
266 let intercept = betas.get(1, 0);
267 let df = xs.len() as f64 - 2.0;
268 let residuals = ys
269 .iter()
270 .zip(xs.iter())
271 .map(|(y, x)| (y - (intercept + m * x)))
272 .collect::<Vec<_>>();
273
274 let se =
275 (inv_matrix[(0, 0)] * ((residuals.iter().map(|x| x.powi(2)).sum::<f64>()) / df)).sqrt();
276 let t = m / se;
277 let t_distr = StudentsT::new(0.0, 1.0, (xs.len() - 2) as f64).unwrap();
278 PValue {
279 p_value: 2.0 * (1.0 - t_distr.cdf(t.abs())),
280 beta: *m,
281 intercept: *intercept,
282 data: None,
283 data_column: None,
284 outcome: None,
285 }
286}
287
288pub fn standardize_column(mut x: ColMut<f64>) {
289 let mut mean = mean(x.as_ref().try_as_col_major().unwrap().as_slice());
290 let mut std = variance(x.as_ref().try_as_col_major().unwrap().as_slice(), 1);
291 let std = std.sqrt();
292 if std == 0.0 {
293 x.fill(0.0);
294 return;
295 }
296 let xx = x.as_mut();
297 let std_recip = 1.0 / std;
298 if let Some(x) = xx.try_as_col_major_mut() {
299 let x = x.as_slice_mut();
300 Arch::new().dispatch(|| {
301 for x in x.iter_mut() {
302 *x = (*x - mean) * std_recip;
303 }
304 });
305 } else {
306 for x in x.iter_mut() {
307 *x = (*x - mean) * std_recip;
308 }
309 }
310}
311
312pub fn standardize_row(mut x: RowMut<f64>) {
313 let mut mean = [0.0];
314 let mut std = [0.0];
315 faer::stats::col_mean(
316 ColMut::from_slice_mut(&mut mean),
317 x.as_ref().as_mat(),
318 faer::stats::NanHandling::Ignore,
319 );
320 faer::stats::col_varm(
321 ColMut::from_slice_mut(&mut std),
322 x.as_ref().as_mat(),
323 ColRef::from_slice(&mean),
324 faer::stats::NanHandling::Ignore,
325 );
326 let std = std[0].sqrt();
327 let mean = mean[0];
328 if std == 0.0 {
329 x.fill(0.0);
330 return;
331 }
332 let xx = x.as_mut();
333 let std_recip = 1.0 / std;
334 if let Some(x) = xx.try_as_row_major_mut() {
335 let x = x.as_slice_mut();
336 Arch::new().dispatch(|| {
337 for x in x.iter_mut() {
338 *x = (*x - mean) * std_recip;
339 }
340 });
341 } else {
342 for x in x.iter_mut() {
343 *x = (*x - mean) * std_recip;
344 }
345 }
346}
347
348pub fn calculate_adj_r2(r2: f64, nrows: usize, ncols: usize) -> f64 {
349 1.0 - (1.0 - r2) * (nrows as f64 - 1.0) / (nrows as f64 - ncols as f64 - 1.0)
350}
351
352#[derive(Debug, Clone)]
353pub struct LinearModel {
354 slopes: Vec<f64>,
355 intercept: f64,
356 predicted: Vec<f64>,
357 r2: f64,
358 adj_r2: f64,
359}
360
361impl LinearModel {
362 #[inline]
363 #[cfg_attr(coverage_nightly, coverage(off))]
364 pub fn slopes(&self) -> &[f64] {
365 &self.slopes
366 }
367
368 #[inline]
369 #[cfg_attr(coverage_nightly, coverage(off))]
370 pub fn intercept(&self) -> f64 {
371 self.intercept
372 }
373
374 #[inline]
375 #[cfg_attr(coverage_nightly, coverage(off))]
376 pub fn predicted(&self) -> &[f64] {
377 &self.predicted
378 }
379
380 #[inline]
381 #[cfg_attr(coverage_nightly, coverage(off))]
382 pub fn r2(&self) -> f64 {
383 self.r2
384 }
385
386 #[inline]
387 #[cfg_attr(coverage_nightly, coverage(off))]
388 pub fn adj_r2(&self) -> f64 {
389 self.adj_r2
390 }
391
392 pub fn predict(&self, x: &[f64]) -> f64 {
393 self.intercept
394 + self
395 .slopes
396 .iter()
397 .zip(x.iter())
398 .map(|(a, b)| a * b)
399 .sum::<f64>()
400 }
401}
402
403#[tracing::instrument(skip(xs, ys))]
404pub fn linear_regression(xs: MatRef<'_, f64>, ys: &[f64]) -> LinearModel {
405 let ncols = xs.ncols();
406 let mut x = xs.to_owned();
407 x.resize_with(
408 xs.nrows(),
409 xs.ncols() + 1,
410 #[inline(always)]
411 |_, _| 1.0,
412 );
413 let y: MatRef<'_, f64> = MatRef::from_column_major_slice(ys, ys.len(), 1);
414 let c_all = x.transpose() * y;
415 let c_matrix = x.transpose() * &x;
416 let betas = match c_matrix.llt(Side::Lower) {
417 Ok(chol) => chol.solve(c_all),
418 Err(_) => {
419 warn!("Using pseudo inverse");
420 c_matrix
421 .as_mat_ref()
422 .thin_svd()
423 .expect("could not compute thin SVD for pseudoinverse")
424 .pseudoinverse()
425 * &c_all
426 },
427 };
428 let betas = betas.col(0).try_as_col_major().unwrap().as_slice();
429 let intercept = betas[ncols];
430 let mut predicted = (0..ys.len())
431 .map(|i| intercept + (0..ncols).map(|j| betas[j] * x[(i, j)]).sum::<f64>())
432 .collect::<Vec<_>>();
433 let r2 = R2Simd::new(ys, &predicted).calculate();
434 let adj_r2 = calculate_adj_r2(r2, ys.len(), ncols);
435 if should_disable_predicted() {
436 predicted = Vec::new();
437 }
438 LinearModel {
439 slopes: betas[..ncols].to_vec(),
440 intercept,
441 predicted,
442 r2,
443 adj_r2,
444 }
445}
446
447#[derive(Debug, Clone)]
448pub struct R2Simd<'a> {
449 actual: &'a [f64],
450 predicted: &'a [f64],
451}
452
453impl<'a> R2Simd<'a> {
454 pub fn new(actual: &'a [f64], predicted: &'a [f64]) -> Self {
455 assert!(actual.len() == predicted.len());
456 Self { actual, predicted }
457 }
458
459 pub fn calculate(self) -> f64 {
460 let arch = Arch::new();
461 arch.dispatch(self)
462 }
463
464 pub fn calculate_no_simd(self) -> f64 {
465 let mean = self.actual.iter().sum::<f64>() / self.actual.len() as f64;
466 let mut rss = 0.0;
467 for (actual, predicted) in self.actual.iter().zip(self.predicted.iter()) {
468 rss += (actual - predicted).powi(2);
469 }
470 let mut tss = 0.0;
471 for actual in self.actual.iter() {
472 tss += (actual - mean).powi(2);
473 }
474 1.0 - rss / tss
475 }
476}
477
478impl WithSimd for R2Simd<'_> {
479 type Output = f64;
480
481 #[inline(always)]
482 fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
483 let (actual_head, actual_tail) = S::f64s_as_simd(self.actual);
484 let (predicted_head, predicted_tail) = S::f64s_as_simd(self.predicted);
485
486 let mean = mean(self.actual);
487 let simd_mean = simd.f64s_splat(mean);
488
489 let mut rss0 = simd.f64s_splat(0.0);
490 let mut rss1 = simd.f64s_splat(0.0);
491 let mut rss2 = simd.f64s_splat(0.0);
492 let mut rss3 = simd.f64s_splat(0.0);
493 let mut tss0 = simd.f64s_splat(0.0);
494 let mut tss1 = simd.f64s_splat(0.0);
495 let mut tss2 = simd.f64s_splat(0.0);
496 let mut tss3 = simd.f64s_splat(0.0);
497
498 let (actual_head4, actual_head1) = pulp::as_arrays::<4, _>(actual_head);
499 let (predicted_head4, predicted_head1) = pulp::as_arrays::<4, _>(predicted_head);
500
501 for (
502 [actual0, actual1, actual2, actual3],
503 [predicted0, predicted1, predicted2, predicted3],
504 ) in actual_head4.iter().zip(predicted_head4.iter())
505 {
506 let rs0 = simd.f64s_sub(*actual0, *predicted0);
507 let rs1 = simd.f64s_sub(*actual1, *predicted1);
508 let rs2 = simd.f64s_sub(*actual2, *predicted2);
509 let rs3 = simd.f64s_sub(*actual3, *predicted3);
510 rss0 = simd.f64s_mul_add(rs0, rs0, rss0);
511 rss1 = simd.f64s_mul_add(rs1, rs1, rss1);
512 rss2 = simd.f64s_mul_add(rs2, rs2, rss2);
513 rss3 = simd.f64s_mul_add(rs3, rs3, rss3);
514 let ts0 = simd.f64s_sub(*actual0, simd_mean);
515 let ts1 = simd.f64s_sub(*actual1, simd_mean);
516 let ts2 = simd.f64s_sub(*actual2, simd_mean);
517 let ts3 = simd.f64s_sub(*actual3, simd_mean);
518 tss0 = simd.f64s_mul_add(ts0, ts0, tss0);
519 tss1 = simd.f64s_mul_add(ts1, ts1, tss1);
520 tss2 = simd.f64s_mul_add(ts2, ts2, tss2);
521 tss3 = simd.f64s_mul_add(ts3, ts3, tss3);
522 }
523
524 tss0 = simd.f64s_add(tss0, tss2);
525 tss1 = simd.f64s_add(tss1, tss3);
526 rss0 = simd.f64s_add(rss0, rss2);
527 rss1 = simd.f64s_add(rss1, rss3);
528
529 let mut tss = simd.f64s_add(tss0, tss1);
530 let mut rss = simd.f64s_add(rss0, rss1);
531
532 for (actual, predicted) in actual_head1.iter().zip(predicted_head1.iter()) {
533 let rs = simd.f64s_sub(*actual, *predicted);
534 rss = simd.f64s_mul_add(rs, rs, rss);
535 let ts = simd.f64s_sub(*actual, simd_mean);
536 tss = simd.f64s_mul_add(ts, ts, tss);
537 }
538
539 let mut tss = simd.f64s_reduce_sum(tss);
540 let mut rss = simd.f64s_reduce_sum(rss);
541
542 for (actual, predicted) in actual_tail.iter().zip(predicted_tail.iter()) {
543 rss += (*actual - *predicted).powi(2);
544 tss += (*actual - mean).powi(2);
545 }
546
547 1.0 - rss / tss
548 }
549}
550
551#[derive(Debug, Clone)]
552pub struct LogisticModel {
553 slopes: Vec<f64>,
554 intercept: f64,
555 predicted: Vec<f64>,
556 r2: f64,
557 adj_r2: f64,
558}
559
560impl LogisticModel {
561 #[inline]
562 #[cfg_attr(coverage_nightly, coverage(off))]
563 pub fn slopes(&self) -> &[f64] {
564 &self.slopes
565 }
566
567 #[inline]
568 #[cfg_attr(coverage_nightly, coverage(off))]
569 pub fn intercept(&self) -> f64 {
570 self.intercept
571 }
572
573 #[inline]
574 #[cfg_attr(coverage_nightly, coverage(off))]
575 pub fn predicted(&self) -> &[f64] {
576 &self.predicted
577 }
578
579 #[inline]
580 #[cfg_attr(coverage_nightly, coverage(off))]
581 pub fn r2(&self) -> f64 {
582 self.r2
583 }
584
585 #[inline]
586 #[cfg_attr(coverage_nightly, coverage(off))]
587 pub fn adj_r2(&self) -> f64 {
588 self.adj_r2
589 }
590
591 pub fn predict(&self, x: &[f64]) -> f64 {
592 self.intercept
593 + self
594 .slopes
595 .iter()
596 .zip(x.iter())
597 .map(|(a, b)| a * b)
598 .sum::<f64>()
599 }
600}
601
602#[inline(always)]
603fn logit(x: f64) -> f64 {
604 (x / (1.0 - x)).ln()
605}
606
607#[inline(always)]
608fn logistic(x: f64) -> f64 {
609 1.0 / (1.0 + (-x).exp())
610}
611
612#[inline(always)]
613fn logit_prime(x: f64) -> f64 {
614 1.0 / (x * (1.0 - x))
615}
616
617#[inline(always)]
618fn v(x: f64) -> f64 {
619 x * (1.0 - x)
620}
621
622#[inline(always)]
623fn ll(p: &[f64], y: &[f64]) -> f64 {
624 p.iter()
625 .zip(y)
626 .map(|(p, y)| y * p.ln() + (1.0 - y) * (1.0 - p).ln())
627 .sum()
628}
629
630#[tracing::instrument(skip(xs, ys))]
631pub fn logistic_regression_irls(xs: MatRef<'_, f64>, ys: &[f64]) -> LogisticModel {
632 let mut mu = vec![0.5; ys.len()];
633 let mut delta = 1.0;
634 let mut l = 0.0;
635 let mut x = xs.to_owned();
636 x.resize_with(
637 xs.nrows(),
638 xs.ncols() + 1,
639 #[inline(always)]
640 |_, _| 1.0,
641 );
642 let mut slopes = vec![0.0; xs.ncols()];
643 let mut intercept = 0.0;
644 let mut z = vec![0.0; ys.len()];
645 let mut w = vec![0.0; ys.len()];
647 let xt = x.transpose();
648 let mut xtw = Mat::<f64>::zeros(x.ncols(), ys.len());
649 let mut xtwx = Mat::zeros(x.ncols(), x.ncols());
650 let mut xtwz = Col::zeros(x.ncols());
651 while delta > 1e-5 {
652 for ((z, mu), y) in z.iter_mut().zip(mu.iter()).zip(ys) {
653 *z = logit(*mu) + (y - mu) * logit_prime(*mu);
654 }
655 for (i, mu) in mu.iter().enumerate() {
656 w[i] = 1.0 / (logit_prime(*mu).powi(2) * v(*mu));
657 }
658
659 xtw.par_col_chunks_mut(1)
671 .zip(w.par_iter())
672 .enumerate()
673 .for_each(|(j, (mut xtw, w))| {
674 xtw.col_iter_mut().for_each(|x| {
675 x.iter_mut()
676 .enumerate()
677 .for_each(|(i, x)| *x = xt[(i, j)] * w)
678 })
679 });
680 faer::linalg::matmul::matmul(
681 xtwx.as_mut(),
682 faer::Accum::Replace,
683 &xtw,
684 &x,
685 1.0,
686 get_global_parallelism(),
687 );
688 faer::linalg::matmul::matmul(
689 xtwz.as_mut(),
690 faer::Accum::Replace,
691 &xtw,
692 ColRef::from_slice(z.as_slice()),
693 1.0,
694 get_global_parallelism(),
695 );
696
697 let beta = match xtwx.llt(Side::Lower) {
698 Ok(chol) => chol.solve(&xtwz),
699 Err(_) => {
700 warn!("Using pseudo inverse");
701 xtwx.as_mat_ref()
702 .thin_svd()
703 .expect("could not compute thin SVD for pseudoinverse")
704 .pseudoinverse()
705 * &xtwz
706 },
707 };
708 let b = beta.try_as_col_major().unwrap().as_slice();
709 slopes.as_mut_slice().copy_from_slice(&b[..xs.ncols()]);
710 intercept = b[xs.ncols()];
711 let eta = &x * beta;
712 let eta = eta.try_as_col_major().unwrap().as_slice();
713 for (mu, eta) in mu.iter_mut().zip(eta) {
714 *mu = logistic(*eta);
715 }
716 let old_ll = l;
717 l = ll(mu.as_slice(), ys);
718 delta = (l - old_ll).abs();
719 }
720
721 let r2 = R2Simd::new(ys, &mu).calculate();
722 let adj_r2 = calculate_adj_r2(r2, ys.len(), xs.ncols());
723
724 if should_disable_predicted() {
725 mu = Vec::new();
726 }
727
728 LogisticModel {
729 slopes,
730 intercept,
731 predicted: mu,
732 r2,
733 adj_r2,
734 }
735}
736
737#[tracing::instrument(skip(xs, ys))]
738pub fn logistic_regression_newton_raphson(xs: MatRef<'_, f64>, ys: &[f64]) -> LogisticModel {
739 let mut x = xs.to_owned();
740 x.resize_with(
741 xs.nrows(),
742 xs.ncols() + 1,
743 #[inline(always)]
744 |_, _| 1.0,
745 );
746 let mut beta = vec![0.0; x.ncols()];
747 let mut mu = vec![0.0; ys.len()];
748 let mut w = vec![0.0; ys.len()];
750 let mut linear_predictor = Col::zeros(ys.len());
751 let mut ys_sub_mu = vec![0.0; ys.len()];
752 let xt = x.transpose();
753 let mut jacobian = Col::zeros(x.ncols());
754 let mut xtw = Mat::<f64>::zeros(x.ncols(), ys.len());
755 let mut hessian = Mat::zeros(x.ncols(), x.ncols());
756 for _ in 0..100 {
757 faer::linalg::matmul::matmul(
758 linear_predictor.as_mut(),
759 faer::Accum::Replace,
760 &x,
761 ColRef::from_slice(beta.as_slice()),
762 1.0,
763 get_global_parallelism(),
764 );
765 for (mu, l) in mu
766 .iter_mut()
767 .zip(linear_predictor.try_as_col_major().unwrap().as_slice())
768 {
769 *mu = logistic(*l);
770 }
771 for (i, mu) in mu.iter().enumerate() {
772 w[i] = mu * (1.0 - mu);
773 }
774 for (i, (mu, y)) in mu.iter().zip(ys).enumerate() {
775 ys_sub_mu[i] = *y - mu;
776 }
777
778 faer::linalg::matmul::matmul(
779 jacobian.as_mut(),
780 faer::Accum::Replace,
781 xt,
782 ColRef::from_slice(ys_sub_mu.as_slice()),
783 1.0,
784 get_global_parallelism(),
785 );
786 xtw.par_col_chunks_mut(1)
798 .zip(w.par_iter())
799 .enumerate()
800 .for_each(|(j, (mut xtw, w))| {
801 xtw.col_iter_mut().for_each(|x| {
802 x.iter_mut()
803 .enumerate()
804 .for_each(|(i, x)| *x = xt[(i, j)] * w)
805 })
806 });
807 faer::linalg::matmul::matmul(
808 hessian.as_mut(),
809 faer::Accum::Replace,
810 &xtw,
811 &x,
812 1.0,
813 get_global_parallelism(),
814 );
815
816 let beta_new = ColRef::from_slice(beta.as_slice())
817 + match hessian.llt(Side::Lower) {
818 Ok(chol) => chol.solve(&jacobian),
819 Err(_) => {
820 warn!("Using pseudo inverse");
821 hessian
822 .as_mat_ref()
823 .thin_svd()
824 .expect("could not compute thin SVD for pseudoinverse")
825 .pseudoinverse()
826 * &jacobian
827 },
828 };
829
830 if (&beta_new - ColRef::from_slice(beta.as_slice())).norm_l1() < 1e-5 {
831 break;
832 }
833 beta.copy_from_slice(beta_new.try_as_col_major().unwrap().as_slice());
834 }
835 let r2 = R2Simd::new(ys, &mu).calculate();
836 let adj_r2 = calculate_adj_r2(r2, ys.len(), xs.ncols());
837
838 let predicted = if should_disable_predicted() {
839 Vec::new()
840 } else {
841 (&x * ColRef::from_slice(beta.as_slice()))
842 .try_as_col_major()
843 .unwrap()
844 .as_slice()
845 .iter()
846 .map(|x| logistic(*x))
847 .collect()
848 };
849
850 LogisticModel {
851 predicted,
852 intercept: beta[x.ncols() - 1],
853 slopes: {
854 beta.truncate(x.ncols() - 1);
855 beta
856 },
857 r2,
858 adj_r2,
859 }
860}
861
862fn logistic_regression_glm(xs: MatRef<'_, f64>, ys: &[f64]) -> LogisticModel {
864 let mut x = xs.to_owned();
865 x.resize_with(
866 xs.nrows(),
867 xs.ncols() + 1,
868 #[inline(always)]
869 |_, _| 1.0,
870 );
871 let weights = vec![1.0; ys.len()];
872 let mustart = weights
873 .iter()
874 .zip(ys)
875 .map(|(w, y)| (w * y + 0.5) / (w + 1.0))
876 .collect::<Vec<_>>();
877 let mut eta = mustart.iter().map(|x| logit(*x)).collect::<Vec<_>>();
878 let mut mu = mustart;
879 let mut delta = 1.0;
880 let mut l = 0.0;
881 let mut slopes = vec![0.0; xs.ncols()];
882 let mut intercept = 0.0;
883 let mut i = 0;
884 while delta > 1e-5 {
885 fn mu_eta_val(eta: f64) -> f64 {
886 let exp = eta.exp();
887 exp / (1.0 + exp).powi(2)
888 }
889 let z = eta
890 .iter()
891 .zip(ys)
892 .zip(mu.iter())
893 .map(|((e, y), m)| e + (y - m) / mu_eta_val(*e))
894 .collect::<Vec<_>>();
895 let w = mu
896 .iter()
897 .zip(weights.iter())
898 .zip(eta.iter())
899 .map(|((m, w), e)| (w * mu_eta_val(*e).powi(2)).sqrt() / v(*m))
900 .collect::<Vec<_>>();
901 let mut xw = x.clone();
902 xw.row_iter_mut()
903 .zip(w.iter())
904 .for_each(|(x, w)| x.iter_mut().for_each(|x| *x *= *w));
905 let zw = z.iter().zip(w).map(|(z, w)| z * w).collect::<Vec<_>>();
906 std::io::stdout().flush().unwrap();
907 let qr = xw.qr();
908 let beta = (qr.thin_R().thin_svd().unwrap().inverse()
916 * qr.compute_thin_Q().transpose()
917 * ColRef::from_slice(zw.as_slice()));
918 let beta_slice = beta.try_as_col_major().unwrap().as_slice();
919 slopes.copy_from_slice(&beta_slice[..xs.ncols()]);
920 intercept = beta[xs.ncols()];
921 let eta = (&x * beta).try_as_col_major().unwrap().as_slice().to_vec();
922 for (i, eta) in eta.iter().enumerate() {
923 mu[i] = logistic(*eta);
924 }
925 let old_ll = l;
926 l = ll(mu.as_slice(), ys);
927 delta = (l - old_ll).abs();
928 i += 1;
929 if i > 100 {
930 error!("Did not converge");
931 break;
932 }
933 }
934 let r2 = R2Simd::new(ys, &mu).calculate();
935 let adj_r2 = calculate_adj_r2(r2, ys.len(), xs.ncols());
936
937 LogisticModel {
938 slopes,
939 intercept,
940 predicted: mu,
941 r2,
942 adj_r2,
943 }
944}
945
946pub fn compute_r2_tjur(actual: &[f64], predicted: &[f64]) -> f64 {
948 let mut ones_sum = 0.0;
949 let mut zeros_sum = 0.0;
950 let mut ones_count = 0.0;
951 let mut zeros_count = 0.0;
952 for (a, p) in actual.iter().zip(predicted.iter()) {
953 if *a == 1.0 {
954 ones_sum += p;
955 ones_count += 1.0;
956 } else {
957 zeros_sum += p;
958 zeros_count += 1.0;
959 }
960 }
961 ((ones_sum / ones_count) - (zeros_sum / zeros_count))
962}
963
964#[cfg(test)]
965mod tests {
966 use faer::RowRef;
967 use rand::SeedableRng;
968 use test_log::test;
969
970 use super::*;
971 use crate::{IntoMatrix, Matrix, OwnedMatrix};
972
973 macro_rules! assert_float_eq {
974 ($a:expr, $b:expr, $tol:expr) => {
975 assert!(($a - $b).abs() < $tol, "{:.22} != {:.22}", $a, $b);
976 };
977 }
978
979 macro_rules! float_eq {
980 ($a:expr, $b:expr) => {
981 assert_float_eq!($a, $b, 1e-12);
982 };
983 }
984
985 macro_rules! rough_eq {
986 ($a:expr, $b:expr) => {
987 assert_float_eq!($a, $b, 1e-3);
988 };
989 }
990
991 #[test]
992 fn test_standardize_column() {
993 let mut data = [1.0, 2.0, 3.0];
994 standardize_column(ColMut::from_slice_mut(&mut data));
995 assert_eq!(data, [-1.0, 0.0, 1.0]);
996 }
997
998 #[test]
999 fn test_standardize_row() {
1000 let mut data = [1.0, 2.0, 3.0];
1001 standardize_row(RowMut::from_slice_mut(&mut data));
1002 assert_eq!(data, [-1.0, 0.0, 1.0]);
1003 }
1004
1005 #[test]
1006 fn test_get_r2s_normalized() {
1007 std::env::set_var("LMUTILS_ENABLE_PREDICTED", "1");
1008 let mut data = OwnedMatrix::new(4, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 9.0], None)
1009 .into_matrix();
1010 data.standardize_columns();
1011 let mut outcomes =
1012 OwnedMatrix::new(4, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 8.0], None)
1013 .into_matrix();
1014 outcomes.standardize_columns();
1015 let r2s = get_r2s(data.as_mat_ref_loaded(), outcomes.as_mat_ref_loaded());
1016 assert_eq!(r2s.len(), 2);
1017 float_eq!(r2s[0].r2(), 1.0);
1018 float_eq!(r2s[0].adj_r2(), 1.0);
1019 let pred: [f64; 4] = [
1020 -1.1618950038622262,
1021 -0.3872983346207394,
1022 0.3872983346207394,
1023 1.1618950038622262,
1024 ];
1025 assert_eq!(r2s[0].predicted().len(), 4);
1026 for (a, b) in r2s[0].predicted().iter().copied().zip(pred.iter().copied()) {
1027 float_eq!(a, b);
1028 }
1029 assert_eq!(r2s[0].n(), 4);
1030 assert_eq!(r2s[0].m(), 2);
1031 assert!(r2s[0].data().is_none());
1032 assert!(r2s[0].outcome().is_none());
1033 float_eq!(r2s[1].r2(), 0.9629629629629629);
1034 float_eq!(r2s[1].adj_r2(), 0.8888888888888891);
1035 let pred: [f64; 4] = [
1036 -0.9999999999999998,
1037 -0.6666666666666672,
1038 0.6666666666666672,
1039 0.9999999999999998,
1040 ];
1041 assert_eq!(r2s[1].predicted().len(), 4);
1042 for (a, b) in r2s[1].predicted().iter().copied().zip(pred.iter().copied()) {
1043 float_eq!(a, b);
1044 }
1045 assert_eq!(r2s[1].n(), 4);
1046 assert_eq!(r2s[1].m(), 2);
1047 assert!(r2s[1].data().is_none());
1048 assert!(r2s[1].outcome().is_none());
1049 }
1050
1051 #[test]
1052 fn test_get_r2s_not_normalized() {
1053 std::env::set_var("LMUTILS_ENABLE_PREDICTED", "1");
1054 let data = MatRef::from_column_major_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 9.0], 4, 2);
1055 let outcomes =
1056 MatRef::from_column_major_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 8.0], 4, 2);
1057 let r2s = get_r2s(data.as_mat_ref(), outcomes.as_mat_ref());
1058 assert_eq!(r2s.len(), 2);
1059 float_eq!(r2s[0].r2(), 1.0);
1060 float_eq!(r2s[0].adj_r2(), 1.0);
1061 let pred: [f64; 4] = [1.0, 2.0, 3.0, 4.0];
1062 assert_eq!(r2s[0].predicted().len(), 4);
1063 for (a, b) in r2s[0].predicted().iter().copied().zip(pred.iter().copied()) {
1064 float_eq!(a, b);
1065 }
1066 assert_eq!(r2s[0].n(), 4);
1067 assert_eq!(r2s[0].m(), 2);
1068 assert!(r2s[0].data().is_none());
1069 assert!(r2s[0].outcome().is_none());
1070 float_eq!(r2s[1].r2(), 0.9592740150509075);
1071 float_eq!(r2s[1].adj_r2(), 0.8778220451527224);
1072 let pred: [f64; 4] = [
1073 5.235059760956178,
1074 5.864541832669323,
1075 7.6454183266932265,
1076 8.274900398406372,
1077 ];
1078 assert_eq!(r2s[1].predicted().len(), 4);
1079 for (a, b) in r2s[1].predicted().iter().copied().zip(pred.iter().copied()) {
1080 float_eq!(a, b);
1081 }
1082 assert_eq!(r2s[1].n(), 4);
1083 assert_eq!(r2s[1].m(), 2);
1084 assert!(r2s[1].data().is_none());
1085 assert!(r2s[1].outcome().is_none());
1086 }
1087
1088 #[test]
1089 fn test_p_value() {
1090 let xs = [1.0, 2.0, 3.0, 4.0, 5.0];
1091 let ys = [1.0, 2.0, 3.0, 4.0, 5.0];
1092 let p_value = p_value(&xs, &ys);
1093 float_eq!(p_value.p_value(), 0.0);
1094 assert!(p_value.data().is_none());
1095 assert!(p_value.outcome().is_none());
1096 }
1097
1098 #[test]
1099 fn test_linear_regression() {
1100 let xs = MatRef::from_column_major_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], 5, 1);
1101 let ys = [1.0, 2.0, 3.0, 4.0, 5.0];
1102 let model = linear_regression(xs.as_mat_ref(), &ys);
1103 assert_eq!(model.slopes().len(), 1);
1104 float_eq!(model.slopes()[0], 1.0);
1105 float_eq!(model.intercept(), 0.0);
1106 assert_eq!(model.predicted().len(), 5);
1107 let predicted = [1.0, 2.0, 3.0, 4.0, 5.0];
1108 for (a, b) in model
1109 .predicted()
1110 .iter()
1111 .copied()
1112 .zip(predicted.iter().copied())
1113 {
1114 float_eq!(a, b);
1115 }
1116 float_eq!(model.r2(), 1.0);
1117 float_eq!(model.adj_r2(), 1.0);
1118 }
1119
1120 #[test]
1121 fn test_linear_regression_predict() {
1122 let xs = MatRef::from_column_major_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], 5, 1);
1123 let ys = [1.0, 2.0, 3.0, 4.0, 5.0];
1124 let model = linear_regression(xs.as_mat_ref(), &ys);
1125 float_eq!(model.predict(&[6.0]), 6.0);
1126 let xs = MatRef::from_column_major_slice(&[1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0], 4, 2);
1127 let ys = [2.0, 4.0, 6.0, 8.0];
1128 let model = linear_regression(xs.as_mat_ref(), &ys);
1129 float_eq!(model.predict(&[1.0, 1.0]), 2.0);
1130 }
1131
1132 #[test]
1133 fn test_r2_simd() {
1134 for _ in 0..100 {
1135 let random_actual: Vec<f64> = (0..1000).map(|_| rand::random()).collect();
1136 let random_predicted: Vec<f64> = (0..1000).map(|_| rand::random()).collect();
1137
1138 let r2 = R2Simd::new(&random_actual, &random_predicted).calculate();
1139 let r2_no_simd = R2Simd::new(&random_actual, &random_predicted).calculate_no_simd();
1140 float_eq!(r2, r2_no_simd);
1141 }
1142 }
1143
1144 #[test]
1145 fn test_r2() {
1146 let actual = [1.0, 2.0, 3.0, 4.0, 5.0];
1147 let predicted = [1.0, 2.0, 3.0, 4.0, 5.0];
1148 let r2 = R2Simd::new(&actual, &predicted).calculate();
1149 float_eq!(r2, 1.0);
1150 let actual = [1.0, 2.0, 3.0, 5.0, 6.0];
1151 let predicted = [1.0, 2.0, 3.0, 4.0, 5.0];
1152 let r2 = R2Simd::new(&actual, &predicted).calculate();
1153 float_eq!(r2, 0.8837209302325582);
1154 }
1155
1156 #[test]
1157 fn test_logistic_regression() {
1158 let nrows = 50;
1159 let xs = statrs::distribution::Normal::new(0.0, 1.0).unwrap();
1160 let ys = statrs::distribution::Bernoulli::new(0.5).unwrap();
1161 let xs = xs
1162 .sample_iter(rand::thread_rng())
1163 .take(nrows)
1164 .collect::<Vec<_>>();
1165 let ys = ys
1166 .sample_iter(rand::thread_rng())
1167 .take(nrows)
1168 .collect::<Vec<_>>();
1169 let xs = MatRef::from_column_major_slice(xs.as_slice(), nrows, 1);
1170 let m1 = logistic_regression_irls(xs, ys.as_slice());
1171 let m2 = logistic_regression_newton_raphson(xs, ys.as_slice());
1172 for (a, b) in m1.slopes.iter().zip(m2.slopes.iter()) {
1173 rough_eq!(a, b);
1174 }
1175 rough_eq!(m1.intercept, m2.intercept);
1176 for (a, b) in m1.predicted.iter().zip(m2.predicted.iter()) {
1177 rough_eq!(a, b);
1178 }
1179 }
1180}