1use serde::{Deserialize, Serialize};
48use thiserror::Error;
49
50use super::hist::Histogram;
51
52#[derive(Debug, Clone, PartialEq, Error, Serialize, Deserialize)]
54pub enum PearsonError {
55 #[error("invalid moments: mean = {mean}, std_dev = {std_dev}, skewness = {skewness}, kurtosis = {kurtosis}")]
57 InvalidMoments {
58 mean: f64,
60 std_dev: f64,
62 skewness: f64,
64 kurtosis: f64,
66 },
67 #[error("moments outside the Pearson IV region: skewness = {skewness}, kurtosis = {kurtosis} (needs kurtosis > {min_kurtosis:?})")]
69 OutsideTypeIv {
70 skewness: f64,
72 kurtosis: f64,
74 min_kurtosis: Option<f64>,
77 },
78 #[error("histogram unsuitable for a fit: {0}")]
80 BadHistogram(String),
81}
82
83pub fn type_iv_min_kurtosis(skewness: f64) -> Option<f64> {
85 let g2 = skewness * skewness;
86 (g2 < 32.0).then(|| (48.0 + 39.0 * g2 + 6.0 * (g2 + 4.0).powf(1.5)) / (32.0 - g2))
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
96pub struct PearsonIv {
97 pub mean: f64,
99 pub std_dev: f64,
101 pub skewness: f64,
103 pub kurtosis: f64,
105 pub m: f64,
107 pub nu: f64,
109 pub a: f64,
111 pub lambda: f64,
113 pub log_norm: f64,
115}
116
117fn re_ln_gamma(x: f64, y: f64) -> f64 {
120 let shift = if x < 12.0 {
121 (12.0 - x).ceil() as usize
122 } else {
123 0
124 };
125 let mut correction = 0.0;
126 for k in 0..shift {
127 let xr = x + k as f64;
128 correction += 0.5 * (xr * xr + y * y).ln();
129 }
130 let (wr, wi) = (x + shift as f64, y);
131 let modulus2 = wr * wr + wi * wi;
132 let ln_mod = 0.5 * modulus2.ln();
133 let arg = wi.atan2(wr);
134 let mut s = (wr - 0.5) * ln_mod - wi * arg - wr + 0.5 * (2.0 * std::f64::consts::PI).ln();
136 let (ir, ii) = (wr / modulus2, -wi / modulus2); let (i2r, i2i) = (ir * ir - ii * ii, 2.0 * ir * ii); let mut pr = ir;
140 let mut pi = ii;
141 for c in [1.0 / 12.0, -1.0 / 360.0, 1.0 / 1260.0, -1.0 / 1680.0] {
142 s += c * pr;
143 let t = pr * i2r - pi * i2i;
144 pi = pr * i2i + pi * i2r;
145 pr = t;
146 }
147 s - correction
148}
149
150impl PearsonIv {
151 pub fn from_moments(
155 mean: f64,
156 std_dev: f64,
157 skewness: f64,
158 kurtosis: f64,
159 ) -> Result<Self, PearsonError> {
160 if !(mean.is_finite()
161 && std_dev.is_finite()
162 && std_dev > 0.0
163 && skewness.is_finite()
164 && kurtosis.is_finite())
165 {
166 return Err(PearsonError::InvalidMoments {
167 mean,
168 std_dev,
169 skewness,
170 kurtosis,
171 });
172 }
173 let (g, b) = (skewness, kurtosis);
174 let outside = || PearsonError::OutsideTypeIv {
175 skewness: g,
176 kurtosis: b,
177 min_kurtosis: type_iv_min_kurtosis(g),
178 };
179 let den = 2.0 * b - 3.0 * g * g - 6.0;
180 if den.is_nan() || den <= 0.0 {
181 return Err(outside());
182 }
183 let r = 6.0 * (b - g * g - 1.0) / den;
184 let d = 16.0 * (r - 1.0) - g * g * (r - 2.0) * (r - 2.0);
185 if !(d > 0.0 && r > 3.0) {
186 return Err(outside());
187 }
188 let m = 1.0 + 0.5 * r;
189 let sd = d.sqrt();
190 let nu = -r * (r - 2.0) * g / sd;
191 let a = std_dev * sd / 4.0;
192 let lambda = mean - (r - 2.0) * g * std_dev / 4.0;
193 let ln_beta = re_ln_gamma(m - 0.5, 0.0) + re_ln_gamma(0.5, 0.0) - re_ln_gamma(m, 0.0);
194 let log_norm = 2.0 * (re_ln_gamma(m, 0.5 * nu) - re_ln_gamma(m, 0.0)) - a.ln() - ln_beta;
195 Ok(Self {
196 mean,
197 std_dev,
198 skewness,
199 kurtosis,
200 m,
201 nu,
202 a,
203 lambda,
204 log_norm,
205 })
206 }
207
208 pub fn pdf(&self, x: f64) -> f64 {
210 let t = (x - self.lambda) / self.a;
211 (self.log_norm - self.m * t.mul_add(t, 1.0).ln() - self.nu * t.atan()).exp()
212 }
213}
214
215fn bin_masses(hist: &Histogram, pdf: impl Fn(f64) -> f64) -> Vec<f64> {
218 let b = &hist.binning;
219 let h = b.width();
220 let off = 0.5 * h * (0.6f64).sqrt();
221 (0..b.bins)
222 .map(|i| {
223 let c = b.center(i);
224 0.5 * h * ((5.0 / 9.0) * (pdf(c - off) + pdf(c + off)) + (8.0 / 9.0) * pdf(c))
225 })
226 .collect()
227}
228
229fn chi_square(hist: &Histogram, p: &[f64]) -> f64 {
232 let n = hist.total() as f64;
233 let s: f64 = hist
234 .counts
235 .iter()
236 .zip(p)
237 .map(|(&c, &pi)| {
238 let d = c as f64 - n * pi;
239 d * d / (c as f64).max(1.0)
240 })
241 .sum();
242 if s.is_finite() {
243 s
244 } else {
245 f64::INFINITY
246 }
247}
248
249#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
251pub struct DualPearson {
252 pub head_fraction: f64,
254 pub head: PearsonIv,
256 pub tail: PearsonIv,
258}
259
260#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
262pub struct DualPearsonFit {
263 pub profile: DualPearson,
265 pub chi_square: f64,
267 pub single_chi_square: Option<f64>,
270 pub bins: usize,
272 pub evaluations: u64,
274}
275
276fn component(p: &[f64]) -> Option<PearsonIv> {
279 let g = 4.0 * p[2].tanh();
280 let b = type_iv_min_kurtosis(g)? + p[3].exp().max(1e-6);
281 PearsonIv::from_moments(p[0], p[1].exp(), g, b).ok()
282}
283
284fn unconstrained(mean: f64, sd: f64, g: f64, b: f64) -> [f64; 4] {
285 let g = g.clamp(-3.6, 3.6);
286 let excess = (b - type_iv_min_kurtosis(g).expect("|g| < 4")).max(0.2);
287 [mean, sd.ln(), (g / 4.0).atanh(), excess.ln()]
288}
289
290fn logistic(x: f64) -> f64 {
291 1.0 / (1.0 + (-x).exp())
292}
293
294fn nelder_mead(
301 f: &dyn Fn(&[f64]) -> f64,
302 x0: &[f64],
303 steps: &[f64],
304 max_evals: u64,
305 evals: &mut u64,
306) -> (Vec<f64>, f64) {
307 let n = x0.len();
308 let used = std::cell::Cell::new(0u64);
309 let eval = |x: &[f64]| {
310 used.set(used.get() + 1);
311 let v = f(x);
312 if v.is_nan() {
313 f64::INFINITY
314 } else {
315 v
316 }
317 };
318 let mut pts: Vec<(Vec<f64>, f64)> = Vec::with_capacity(n + 1);
319 let v0 = eval(x0);
320 pts.push((x0.to_vec(), v0));
321 for i in 0..n {
322 let mut x = x0.to_vec();
323 x[i] += steps[i];
324 let v = eval(&x);
325 pts.push((x, v));
326 }
327 while used.get() < max_evals {
328 pts.sort_by(|a, b| a.1.total_cmp(&b.1));
330 let (best, worst) = (pts[0].1, pts[n].1);
331 if (worst - best).abs() <= 1e-12 * (best.abs() + 1e-12) {
332 break;
333 }
334 let mut c = vec![0.0; n];
335 for (x, _) in &pts[..n] {
336 for (ci, xi) in c.iter_mut().zip(x) {
337 *ci += xi / n as f64;
338 }
339 }
340 let along = |t: f64| -> Vec<f64> {
341 c.iter()
342 .zip(&pts[n].0)
343 .map(|(ci, wi)| ci + t * (ci - wi))
344 .collect()
345 };
346 let xr = along(1.0);
347 let fr = eval(&xr);
348 if fr < pts[0].1 {
349 let xe = along(2.0);
350 let fe = eval(&xe);
351 pts[n] = if fe < fr { (xe, fe) } else { (xr, fr) };
352 } else if fr < pts[n - 1].1 {
353 pts[n] = (xr, fr);
354 } else {
355 let (xc, fc) = if fr < pts[n].1 {
356 let x = along(0.5);
357 let v = eval(&x);
358 (x, v)
359 } else {
360 let x = along(-0.5);
361 let v = eval(&x);
362 (x, v)
363 };
364 if fc < fr.min(pts[n].1) {
365 pts[n] = (xc, fc);
366 } else {
367 let x_best = pts[0].0.clone();
368 for (x, v) in pts.iter_mut().skip(1) {
369 for (xi, bi) in x.iter_mut().zip(&x_best) {
370 *xi = bi + 0.5 * (*xi - bi);
371 }
372 *v = eval(x);
373 }
374 }
375 }
376 }
377 pts.sort_by(|a, b| a.1.total_cmp(&b.1));
378 *evals += used.get();
379 pts.swap_remove(0)
380}
381
382impl DualPearson {
383 pub fn pdf(&self, x: f64) -> f64 {
385 self.head_fraction * self.head.pdf(x) + (1.0 - self.head_fraction) * self.tail.pdf(x)
386 }
387
388 pub fn moments(&self) -> (f64, f64, f64, f64) {
391 let w = [self.head_fraction, 1.0 - self.head_fraction];
392 let c = [&self.head, &self.tail];
393 let mean = w[0] * c[0].mean + w[1] * c[1].mean;
394 const BINOM: [[f64; 5]; 5] = [
397 [1.0, 0.0, 0.0, 0.0, 0.0],
398 [1.0, 1.0, 0.0, 0.0, 0.0],
399 [1.0, 2.0, 1.0, 0.0, 0.0],
400 [1.0, 3.0, 3.0, 1.0, 0.0],
401 [1.0, 4.0, 6.0, 4.0, 1.0],
402 ];
403 let mut mu = [0.0f64; 5];
404 for (wk, ck) in w.iter().zip(c) {
405 let s = ck.std_dev;
406 let own = [
407 1.0,
408 0.0,
409 s * s,
410 ck.skewness * s.powi(3),
411 ck.kurtosis * s.powi(4),
412 ];
413 let d = ck.mean - mean;
414 for (p, slot) in mu.iter_mut().enumerate() {
415 let acc: f64 = (0..=p)
416 .map(|j| BINOM[p][j] * own[j] * d.powi((p - j) as i32))
417 .sum();
418 *slot += wk * acc;
419 }
420 }
421 let sd = mu[2].sqrt();
422 (mean, sd, mu[3] / (sd * sd * sd), mu[4] / (mu[2] * mu[2]))
423 }
424
425 pub fn fit(hist: &Histogram) -> Result<DualPearsonFit, PearsonError> {
450 let in_range: u64 = hist.counts.iter().sum();
451 let nonzero = hist.counts.iter().filter(|&&c| c > 0).count();
452 if in_range < 100 || nonzero < 10 {
453 return Err(PearsonError::BadHistogram(format!(
454 "{in_range} in-range counts in {nonzero} non-empty bins (need >= 100 and >= 10)"
455 )));
456 }
457 let b = &hist.binning;
459 let mut mom = super::moments::Moments::new();
460 for (i, &c) in hist.counts.iter().enumerate() {
461 let mut one = super::moments::Moments::new();
462 if c > 0 {
463 one.n = c;
465 one.mean = b.center(i);
466 mom.merge(&one);
467 }
468 }
469 let s = mom
470 .summary()
471 .ok_or_else(|| PearsonError::BadHistogram("zero variance".into()))?;
472 let (mu, sigma) = (s.mean, s.std_dev);
473
474 let objective = |p: &[f64]| -> f64 {
476 let (Some(h), Some(t)) = (component(&p[1..5]), component(&p[5..9])) else {
477 return f64::INFINITY;
478 };
479 let r = logistic(p[0]);
480 let masses = bin_masses(hist, |x| {
481 let z = (x - mu) / sigma;
482 (r * h.pdf(z) + (1.0 - r) * t.pdf(z)) / sigma
483 });
484 chi_square(hist, &masses)
485 };
486
487 let single = PearsonIv::from_moments(0.0, 1.0, s.skewness, s.kurtosis).ok();
488 let single_chi = single
489 .map(|p| chi_square(hist, &bin_masses(hist, |x| p.pdf((x - mu) / sigma) / sigma)));
490
491 let head0 = unconstrained(0.0, 1.0, s.skewness, s.kurtosis);
492 let steps = [0.5, 0.25, 0.2, 0.2, 0.5, 0.25, 0.2, 0.2, 0.5];
493 let mut evals = 0;
494 let mut best: Option<(Vec<f64>, f64)> = None;
495 for tail_shift in [1.0, -1.0] {
496 let tail0 = unconstrained(tail_shift, 1.5, 0.0, 4.0);
497 let mut x0 = vec![(0.9f64 / 0.1).ln()];
498 x0.extend_from_slice(&head0);
499 x0.extend_from_slice(&tail0);
500 let (x1, _) = nelder_mead(&objective, &x0, &steps, 6000, &mut evals);
501 let (x2, v2) = nelder_mead(&objective, &x1, &steps, 6000, &mut evals);
502 if best.as_ref().is_none_or(|(_, v)| v2 < *v) {
503 best = Some((x2, v2));
504 }
505 }
506 let (x, chi) = best.expect("two starts");
507 let to_x = |c: PearsonIv| {
508 PearsonIv::from_moments(
509 mu + sigma * c.mean,
510 sigma * c.std_dev,
511 c.skewness,
512 c.kurtosis,
513 )
514 };
515 let fit = match (single, single_chi) {
516 (Some(p), Some(sc)) if sc <= chi => {
517 let p = to_x(p)?;
518 DualPearsonFit {
519 profile: DualPearson {
520 head_fraction: 1.0,
521 head: p,
522 tail: p,
523 },
524 chi_square: sc,
525 single_chi_square: Some(sc),
526 bins: b.bins,
527 evaluations: evals,
528 }
529 }
530 _ => {
531 let h = component(&x[1..5]).expect("finite objective");
532 let t = component(&x[5..9]).expect("finite objective");
533 DualPearsonFit {
534 profile: DualPearson {
535 head_fraction: logistic(x[0]),
536 head: to_x(h)?,
537 tail: to_x(t)?,
538 },
539 chi_square: chi,
540 single_chi_square: single_chi,
541 bins: b.bins,
542 evaluations: evals,
543 }
544 }
545 };
546 Ok(fit)
547 }
548}
549
550#[cfg(test)]
551mod tests {
552 use super::*;
553
554 #[test]
555 fn log_gamma_matches_known_values() {
556 let half = 0.5 * std::f64::consts::PI.ln();
557 assert!((re_ln_gamma(0.5, 0.0) - half).abs() < 1e-12);
558 assert!((re_ln_gamma(10.0, 0.0) - 362_880f64.ln()).abs() < 1e-12);
559 assert!((re_ln_gamma(1.0, 0.0)).abs() < 1e-12);
560 for y in [0.3, 2.0, 15.0] {
562 let pi_y = std::f64::consts::PI * y;
563 let want = 0.5 * (pi_y / pi_y.sinh()).ln();
564 assert!((re_ln_gamma(1.0, y) - want).abs() < 1e-12, "y = {y}");
565 }
566 }
567
568 #[test]
569 fn boundary_formula_is_d_equals_zero() {
570 assert_eq!(type_iv_min_kurtosis(0.0), Some(3.0));
572 assert!(PearsonIv::from_moments(0.0, 1.0, 0.0, 3.0).is_err());
573 assert!(PearsonIv::from_moments(0.0, 1.0, 0.0, 3.001).is_ok());
574 for g in [0.3, -1.0, 2.0, -4.5] {
575 let b = type_iv_min_kurtosis(g).unwrap();
576 let r = 6.0 * (b - g * g - 1.0) / (2.0 * b - 3.0 * g * g - 6.0);
577 let d = 16.0 * (r - 1.0) - g * g * (r - 2.0) * (r - 2.0);
578 assert!(d.abs() < 1e-9 * r * r, "g = {g}: D = {d}");
579 assert!(PearsonIv::from_moments(0.0, 1.0, g, b * 0.999).is_err());
580 assert!(PearsonIv::from_moments(0.0, 1.0, g, b * 1.001 + 1e-6).is_ok());
581 }
582 assert_eq!(type_iv_min_kurtosis(6.0), None);
583 }
584
585 #[test]
586 fn symmetric_case_is_student_like() {
587 let p = PearsonIv::from_moments(2.0, 1.0, 0.0, 4.0).unwrap();
589 assert_eq!(p.nu, 0.0);
590 assert_eq!(p.lambda, 2.0);
591 assert!((p.pdf(2.5) - p.pdf(1.5)).abs() < 1e-15);
592 }
593}