gammalooprs/utils/
fitting.rs1use crate::utils::{F, FloatLike};
2use color_eyre::eyre::{Result, eyre};
3use std::cmp::Ordering;
4
5pub fn constant_dropped_fit_points<T: FloatLike>(
6 start: &F<T>,
7 stop: &F<T>,
8 num: usize,
9) -> Result<Vec<F<T>>> {
10 if num < 3 {
11 return Err(eyre!(
12 "need at least 3 points to generate a grid usable by the constant-dropped fit"
13 ));
14 }
15
16 if !start.0.is_finite() || !stop.0.is_finite() || start <= &start.zero() || stop <= &stop.zero()
17 {
18 return Err(eyre!(
19 "fit points must be generated from finite positive bounds"
20 ));
21 }
22
23 if stop == start {
24 return Err(eyre!(
25 "fit point range must have distinct bounds for multiplicative spacing"
26 ));
27 }
28
29 let is_ascending = stop > start;
30 let log_start = start.log10();
31 let log_stop = stop.log10();
32 let step_count = start.from_usize(num - 1);
33 let log_step = (&log_stop - &log_start) / &step_count;
34 if log_step.is_nan() || log_step.is_infinite() || log_step.abs().less_than_epsilon() {
35 return Err(eyre!(
36 "fit point range must yield a finite non-zero log step"
37 ));
38 }
39
40 let ten = start.from_i64(10);
41 let mut points = Vec::with_capacity(num);
42 points.push(start.clone());
43
44 for index in 1..(num - 1) {
45 let step_index = log_step.from_usize(index);
46 let exponent = &log_start + &(&log_step * &step_index);
47 let point = ten.powf(&exponent);
48 if point.is_nan() || point.is_infinite() {
49 return Err(eyre!("generated a non-finite fit point"));
50 }
51 let previous = points.last().expect("point list is non-empty");
52 let advances = if is_ascending {
53 &point > previous
54 } else {
55 &point < previous
56 };
57 if !advances {
58 return Err(eyre!(
59 "range and point count collapse at current precision; generated fit points are not strictly monotonic"
60 ));
61 }
62
63 points.push(point);
64 }
65
66 let previous = points.last().expect("point list is non-empty");
67 let advances = if is_ascending {
68 stop > previous
69 } else {
70 stop < previous
71 };
72 if !advances {
73 return Err(eyre!(
74 "range and point count collapse at current precision; final bound does not preserve the requested monotonic direction"
75 ));
76 }
77
78 points.push(stop.clone());
79 Ok(points)
80}
81
82pub fn log_log_slope_constant_dropped<T: FloatLike>(x: &[F<T>], y: &[F<T>]) -> Result<SlopeFit<T>> {
83 let samples = collect_valid_samples(x, y)?;
84 if samples.len() < 3 {
85 return Err(eyre!(
86 "need at least 3 valid samples for constant-dropped slope fit"
87 ));
88 }
89
90 let transformed_samples = collect_log_difference_samples(&samples)?;
91 linear_regression(&transformed_samples)
92}
93
94pub struct SlopeFit<T: FloatLike> {
95 pub slope: F<T>,
96 pub r_squared: F<T>,
97}
98
99impl<T: FloatLike> SlopeFit<T> {
100 pub fn r_squared(&self) -> &F<T> {
101 &self.r_squared
102 }
103}
104
105fn collect_valid_samples<T: FloatLike>(x: &[F<T>], y: &[F<T>]) -> Result<Vec<(F<T>, F<T>)>> {
106 if x.len() != y.len() {
107 return Err(eyre!(
108 "x and y must have the same length, got {} and {}",
109 x.len(),
110 y.len()
111 ));
112 }
113
114 let mut samples = x
115 .iter()
116 .zip(y)
117 .filter_map(|(x_value, y_value)| {
118 if !x_value.0.is_finite() || !y_value.0.is_finite() || x_value <= &x_value.zero() {
119 return None;
120 }
121
122 Some((x_value.clone(), y_value.clone()))
123 })
124 .collect::<Vec<_>>();
125
126 samples.sort_by(|(lhs, _), (rhs, _)| lhs.partial_cmp(rhs).unwrap_or(Ordering::Equal));
127
128 let mut deduped = Vec::with_capacity(samples.len());
129 for sample in samples {
130 if deduped
131 .last()
132 .is_some_and(|(previous_x, _)| previous_x == &sample.0)
133 {
134 continue;
135 }
136 deduped.push(sample);
137 }
138
139 Ok(deduped)
140}
141
142fn collect_log_difference_samples<T: FloatLike>(
143 samples: &[(F<T>, F<T>)],
144) -> Result<Vec<(F<T>, F<T>)>> {
145 let reference_log_x = samples[0].0.log10();
146 let next_log_x = samples[1].0.log10();
147 let reference_step = &next_log_x - &reference_log_x;
148 if reference_step.is_nan()
149 || reference_step.is_infinite()
150 || reference_step.abs().less_than_epsilon()
151 {
152 return Err(eyre!("x grid must have a non-zero multiplicative spacing"));
153 }
154
155 let tolerance = {
156 let scale = &reference_step.abs() + &reference_step.one();
157 F::from_f64(1.0e-9) * scale
158 };
159
160 let mut transformed = Vec::with_capacity(samples.len().saturating_sub(1));
161 for pair in samples.windows(2) {
162 let left_log_x = pair[0].0.log10();
163 let right_log_x = pair[1].0.log10();
164 let step = &right_log_x - &left_log_x;
165 if (&step - &reference_step).abs() > tolerance {
166 return Err(eyre!(
167 "x grid must have constant multiplicative spacing for constant-dropped log-log fitting"
168 ));
169 }
170
171 let difference_magnitude = (&pair[1].1 - &pair[0].1).abs();
172 if !difference_magnitude.positive()
173 || difference_magnitude.is_nan()
174 || difference_magnitude.is_infinite()
175 {
176 return Err(eyre!(
177 "adjacent y differences must stay finite and non-zero after dropping the constant term"
178 ));
179 }
180
181 let log_difference = difference_magnitude.log10();
182 if log_difference.is_nan() || log_difference.is_infinite() {
183 return Err(eyre!(
184 "adjacent y differences must be positive in magnitude for the log transform"
185 ));
186 }
187
188 transformed.push((left_log_x, log_difference));
189 }
190
191 if transformed.len() < 2 {
192 return Err(eyre!(
193 "need at least 2 adjacent log-difference samples for linear regression"
194 ));
195 }
196
197 Ok(transformed)
198}
199
200fn linear_regression<T: FloatLike>(samples: &[(F<T>, F<T>)]) -> Result<SlopeFit<T>> {
201 let zero = samples[0].0.zero();
202 let one = zero.one();
203 let n = zero.from_usize(samples.len());
204
205 let mut sum_x = zero.clone();
206 let mut sum_y = zero.clone();
207 let mut sum_xy = zero.clone();
208 let mut sum_x2 = zero.clone();
209 for (x_value, y_value) in samples {
210 sum_x += x_value.clone();
211 sum_y += y_value.clone();
212 sum_xy += x_value * y_value;
213 sum_x2 += x_value * x_value;
214 }
215
216 let denominator = &(&n * &sum_x2) - &(&sum_x * &sum_x);
217 if denominator.abs().less_than_epsilon() {
218 return Err(eyre!(
219 "log-log regression is singular for the transformed difference samples"
220 ));
221 }
222
223 let slope = (&(&n * &sum_xy) - &(&sum_x * &sum_y)) / &denominator;
224 if slope.is_nan() || slope.is_infinite() {
225 return Err(eyre!("log-log regression produced a non-finite slope"));
226 }
227
228 let mean_y = &sum_y / &n;
229 let mut ss_tot = zero.clone();
230 let mut ss_res = zero;
231 let intercept = (&sum_y - &(&slope * &sum_x)) / &n;
232 for (x_value, y_value) in samples {
233 let prediction = &(&slope * x_value) + &intercept;
234 let centered = y_value - &mean_y;
235 let residual = y_value - &prediction;
236
237 ss_tot += ¢ered * ¢ered;
238 ss_res += &residual * &residual;
239 }
240
241 let r_squared = if ss_tot.less_than_epsilon() {
242 one
243 } else {
244 &one - &(&ss_res / &ss_tot)
245 };
246
247 if r_squared.is_nan() || r_squared.is_infinite() {
248 Err(eyre!("log-log regression produced a non-finite r-squared"))
249 } else {
250 Ok(SlopeFit { slope, r_squared })
251 }
252}
253
254#[cfg(test)]
255mod tests {
256 use super::*;
257
258 fn ff64_values(values: &[f64]) -> Vec<F<f64>> {
259 values.iter().copied().map(F::from_f64).collect()
260 }
261
262 #[test]
263 fn constant_dropped_fit_recovers_power_law_slope() {
264 let coefficient = 3.2_f64;
265 let slope = -1.75_f64;
266 let offset = 0.6_f64;
267 let ratio: f64 = 1.6;
268
269 let x = constant_dropped_fit_points(
270 &F::from_f64(0.2_f64),
271 &F::from_f64(0.2_f64 * ratio.powi(7)),
272 8,
273 )
274 .expect("fit point generation should succeed");
275 let x_values = x
276 .iter()
277 .map(|x_value| x_value.into_f64())
278 .collect::<Vec<_>>();
279 let y = x
280 .iter()
281 .map(|x_value| coefficient * x_value.into_f64().powf(slope) + offset)
282 .collect::<Vec<_>>();
283
284 let fit = log_log_slope_constant_dropped(&x, &ff64_values(&y))
285 .expect("power-law fit should succeed");
286
287 assert!((x_values[0] - 0.2_f64).abs() < 1.0e-12);
288 assert!((x_values[7] - 0.2_f64 * ratio.powi(7)).abs() < 1.0e-12);
289 assert!((fit.slope.into_f64() - slope).abs() < 1.0e-10);
290 assert!(fit.r_squared().into_f64() > 0.999_999_999);
291 }
292
293 #[test]
294 fn constant_dropped_fit_supports_descending_geometric_grid() {
295 let coefficient = 3.2_f64;
296 let slope = -1.75_f64;
297 let offset = 0.6_f64;
298 let ratio: f64 = 1.6;
299
300 let x = constant_dropped_fit_points(
301 &F::from_f64(0.2_f64 * ratio.powi(7)),
302 &F::from_f64(0.2_f64),
303 8,
304 )
305 .expect("descending fit point generation should succeed");
306 let x_values = x
307 .iter()
308 .map(|x_value| x_value.into_f64())
309 .collect::<Vec<_>>();
310 let y = x
311 .iter()
312 .map(|x_value| coefficient * x_value.into_f64().powf(slope) + offset)
313 .collect::<Vec<_>>();
314
315 let fit = log_log_slope_constant_dropped(&x, &ff64_values(&y))
316 .expect("power-law fit should succeed on a descending geometric grid");
317
318 assert!((x_values[0] - 0.2_f64 * ratio.powi(7)).abs() < 1.0e-12);
319 assert!((x_values[7] - 0.2_f64).abs() < 1.0e-12);
320 assert!(x_values.windows(2).all(|window| window[0] > window[1]));
321 assert!((fit.slope.into_f64() - slope).abs() < 1.0e-10);
322 assert!(fit.r_squared().into_f64() > 0.999_999_999);
323 }
324
325 #[test]
326 fn constant_dropped_fit_points_require_strict_positive_range() {
327 let points = constant_dropped_fit_points(
328 &F::<f64>::from_f64(1.0_f64),
329 &F::<f64>::from_f64(1.0_f64),
330 8,
331 );
332
333 assert!(points.is_err());
334 }
335
336 #[test]
337 fn constant_dropped_fit_rejects_non_geometric_grid() {
338 let coefficient = 3.2_f64;
339 let slope = -1.75_f64;
340 let offset = 0.6_f64;
341
342 let x = [
343 0.2_f64, 0.37_f64, 0.55_f64, 0.92_f64, 1.3_f64, 1.85_f64, 2.75_f64, 3.6_f64,
344 ];
345 let y = x
346 .iter()
347 .map(|x_value| coefficient * (*x_value).powf(slope) + offset)
348 .collect::<Vec<_>>();
349
350 let fit = log_log_slope_constant_dropped(&ff64_values(&x), &ff64_values(&y));
351
352 assert!(fit.is_err());
353 }
354}