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
use crate::float_trait::Float;
use conv::ConvUtil;
#[allow(clippy::many_single_char_names)]
pub fn fit_straight_line<T: Float>(
x: &[T],
y: &[T],
err2: Option<&[T]>,
) -> StraightLineFitterResult<T> {
let n = x.len().value_as::<T>().unwrap();
let weight = err2.map(|err2| err2.iter().map(|e| e.recip()).collect::<Vec<_>>());
let (s, sx, sy) = match weight.as_ref() {
Some(weight) => {
let (mut s, mut sx, mut sy) = (T::zero(), T::zero(), T::zero());
for ((&x, &y), &w) in x.iter().zip(y.iter()).zip(weight.iter()) {
s += w;
sx += w * x;
sy += w * y;
}
(s, sx, sy)
}
None => {
let (mut sx, mut sy) = (T::zero(), T::zero());
for (&x, &y) in x.iter().zip(y.iter()) {
sx += x;
sy += y;
}
(n, sx, sy)
}
};
let (stt, sty) = match weight.as_ref() {
Some(weight) => {
let (mut stt, mut sty) = (T::zero(), T::zero());
for ((&x, &y), &w) in x.iter().zip(y.iter()).zip(weight.iter()) {
let t = x - sx / s;
stt += w * t.powi(2);
sty += w * t * y;
}
(stt, sty)
}
None => {
let (mut stt, mut sty) = (T::zero(), T::zero());
for (&x, &y) in x.iter().zip(y.iter()) {
let t = x - sx / s;
stt += t.powi(2);
sty += t * y;
}
(stt, sty)
}
};
let slope = sty / stt;
let intercept = (sy - sx * slope) / s;
let mut slope_sigma2 = stt.recip();
let chi2: T = match weight.as_ref() {
Some(weight) => x
.iter()
.zip(y.iter())
.zip(weight.iter())
.map(|((&x, &y), &w)| (y - intercept - slope * x).powi(2) * w)
.sum(),
None => x
.iter()
.zip(y.iter())
.map(|(&x, &y)| (y - intercept - slope * x).powi(2))
.sum(),
};
let reduced_chi2 = chi2 / (n - T::two());
if err2.is_none() {
slope_sigma2 *= reduced_chi2;
}
StraightLineFitterResult {
slope,
slope_sigma2,
reduced_chi2,
}
}
pub struct StraightLineFitterResult<T> {
pub slope: T,
pub slope_sigma2: T,
pub reduced_chi2: T,
}
#[cfg(test)]
#[allow(clippy::unreadable_literal)]
#[allow(clippy::excessive_precision)]
mod tests {
use super::*;
use light_curve_common::all_close;
#[test]
fn straight_line_fitter() {
let x = [0.5, 1.5, 2.5, 5.0, 7.0, 16.0];
let y = [-1.0, 3.0, 2.0, 6.0, 10.0, 25.0];
let desired_slope = 1.63021767;
let desired_slope_sigma2 = 0.0078127;
let desired_reduced_chi2 = 1.271190781049937;
let result = fit_straight_line(&x, &y, None);
all_close(&[result.slope], &[desired_slope], 1e-6);
all_close(&[result.slope_sigma2], &[desired_slope_sigma2], 1e-6);
all_close(&[result.reduced_chi2], &[desired_reduced_chi2], 1e-6);
}
#[test]
fn noisy_straight_line_fitter() {
let x = [0.5, 1.5, 2.5, 5.0, 7.0, 16.0];
let y = [-1.0, 3.0, 2.0, 6.0, 10.0, 25.0];
let err2 = [0.5, 1.0, 0.3, 0.1, 0.9, 2.5];
let desired_slope = 1.60504579;
let desired_slope_sigma2 = 0.00868733;
let desired_reduced_chi2 = 1.8057513419557492;
let result = fit_straight_line(&x, &y, Some(&err2));
all_close(&[result.slope], &[desired_slope], 1e-6);
all_close(&[result.slope_sigma2], &[desired_slope_sigma2], 1e-6);
all_close(&[result.reduced_chi2], &[desired_reduced_chi2], 1e-6);
}
}