Skip to main content

datui_lib/
distribution_fit.rs

1//! Fitting a column's values to each candidate distribution, and testing each fit.
2//!
3//! One set of parameters per family, estimated once, is what the test, the Q-Q plot
4//! and the histogram overlay all use. The test is Kolmogorov-Smirnov, calibrated by a
5//! parametric bootstrap: the statistic of the values against their fit is ranked among
6//! the statistics of samples drawn from that fit and refitted the same way. That is
7//! what makes it a p-value when the parameters come from the data being tested, which
8//! the textbook KS table assumes they do not, and it holds for the discrete families,
9//! where the table does not apply at all.
10//!
11//! A p-value says how surprising the values would be if they came from the fitted
12//! distribution. It is not the probability that they did, and the largest one is not
13//! the best model: among the families a test does not reject, the one chosen is the
14//! one with the lowest AIC.
15
16use crate::statistics::DistributionType;
17
18/// Every family tested, in the order a tie is listed.
19pub const FAMILIES: [DistributionType; 14] = [
20    DistributionType::Normal,
21    DistributionType::LogNormal,
22    DistributionType::Uniform,
23    DistributionType::Exponential,
24    DistributionType::Gamma,
25    DistributionType::ChiSquared,
26    DistributionType::Beta,
27    DistributionType::StudentsT,
28    DistributionType::Weibull,
29    DistributionType::PowerLaw,
30    DistributionType::Poisson,
31    DistributionType::Bernoulli,
32    DistributionType::Binomial,
33    DistributionType::Geometric,
34];
35
36/// Simulated samples per test. The smallest p-value it can give is 1/200.
37pub const REPLICATES: usize = 199;
38
39/// The most values a test is run on. The bootstrap costs this times [`REPLICATES`]
40/// fits per family; past a few hundred values, a KS test rejects any real-world data
41/// anyway, which says more about the sample size than about the fit.
42pub const TEST_VALUES: usize = 500;
43
44/// A distribution with its parameters, as fitted to a column.
45#[derive(Debug, Clone, Copy, PartialEq)]
46pub enum Fitted {
47    Normal {
48        mean: f64,
49        sd: f64,
50    },
51    LogNormal {
52        mu: f64,
53        sigma: f64,
54    },
55    Uniform {
56        low: f64,
57        high: f64,
58    },
59    Exponential {
60        rate: f64,
61    },
62    Gamma {
63        shape: f64,
64        scale: f64,
65    },
66    ChiSquared {
67        df: f64,
68    },
69    Beta {
70        alpha: f64,
71        beta: f64,
72    },
73    StudentsT {
74        df: f64,
75        location: f64,
76        scale: f64,
77    },
78    Weibull {
79        shape: f64,
80        scale: f64,
81    },
82    PowerLaw {
83        xmin: f64,
84        alpha: f64,
85    },
86    Poisson {
87        rate: f64,
88    },
89    Bernoulli {
90        p: f64,
91    },
92    Binomial {
93        trials: u64,
94        p: f64,
95    },
96    /// Failures before the first success when `start` is 0; trials up to and including
97    /// it when `start` is 1.
98    Geometric {
99        p: f64,
100        start: u64,
101    },
102}
103
104/// How a family fared against a column.
105#[derive(Debug, Clone, PartialEq)]
106pub enum FitOutcome {
107    Tested(FitTest),
108    /// The family cannot describe these values, and why: `needs positive values`.
109    NotApplicable(&'static str),
110}
111
112impl FitOutcome {
113    pub fn p_value(&self) -> Option<f64> {
114        match self {
115            Self::Tested(test) => Some(test.p_value),
116            Self::NotApplicable(_) => None,
117        }
118    }
119
120    pub fn test(&self) -> Option<&FitTest> {
121        match self {
122            Self::Tested(test) => Some(test),
123            Self::NotApplicable(_) => None,
124        }
125    }
126}
127
128/// A fit and its test.
129#[derive(Debug, Clone, PartialEq)]
130pub struct FitTest {
131    /// Fitted to every value read, and the parameters every view of this family uses.
132    pub fitted: Fitted,
133    /// The bootstrap p-value: `(1 + beyond) / (1 + replicates)`.
134    pub p_value: f64,
135    /// Simulated samples whose statistic was at least the column's.
136    pub beyond: usize,
137    pub replicates: usize,
138    /// Values the test was run on.
139    pub tested_on: usize,
140    /// On every value read; lower is better, and only comparable within one kind
141    /// (densities or probabilities).
142    pub aic: f64,
143}
144
145impl FitTest {
146    /// Whether no simulated sample reached the column's statistic, so the p-value is
147    /// only a bound: `< 0.005`, not `0.005`.
148    pub fn at_bound(&self) -> bool {
149        self.beyond == 0
150    }
151}
152
153impl Fitted {
154    pub fn family(&self) -> DistributionType {
155        match self {
156            Self::Normal { .. } => DistributionType::Normal,
157            Self::LogNormal { .. } => DistributionType::LogNormal,
158            Self::Uniform { .. } => DistributionType::Uniform,
159            Self::Exponential { .. } => DistributionType::Exponential,
160            Self::Gamma { .. } => DistributionType::Gamma,
161            Self::ChiSquared { .. } => DistributionType::ChiSquared,
162            Self::Beta { .. } => DistributionType::Beta,
163            Self::StudentsT { .. } => DistributionType::StudentsT,
164            Self::Weibull { .. } => DistributionType::Weibull,
165            Self::PowerLaw { .. } => DistributionType::PowerLaw,
166            Self::Poisson { .. } => DistributionType::Poisson,
167            Self::Bernoulli { .. } => DistributionType::Bernoulli,
168            Self::Binomial { .. } => DistributionType::Binomial,
169            Self::Geometric { .. } => DistributionType::Geometric,
170        }
171    }
172
173    /// Whether this is a distribution over whole numbers.
174    pub fn discrete(&self) -> bool {
175        matches!(
176            self,
177            Self::Poisson { .. }
178                | Self::Bernoulli { .. }
179                | Self::Binomial { .. }
180                | Self::Geometric { .. }
181        )
182    }
183
184    fn parameters(&self) -> usize {
185        match self {
186            Self::Exponential { .. }
187            | Self::ChiSquared { .. }
188            | Self::Poisson { .. }
189            | Self::Bernoulli { .. } => 1,
190            // The start is read off the values, as the minimum is for a power law.
191            Self::Geometric { .. } => 2,
192            Self::StudentsT { .. } => 3,
193            _ => 2,
194        }
195    }
196
197    /// The parameters, as a reader would write them: `μ 4.91, σ 0.43`.
198    pub fn describe(&self) -> String {
199        let g = |x: f64| format_param(x);
200        match *self {
201            Self::Normal { mean, sd } => format!("mean {}, sd {}", g(mean), g(sd)),
202            Self::LogNormal { mu, sigma } => format!("log mean {}, log sd {}", g(mu), g(sigma)),
203            Self::Uniform { low, high } => format!("{} to {}", g(low), g(high)),
204            Self::Exponential { rate } => format!("rate {}", g(rate)),
205            Self::Gamma { shape, scale } => format!("shape {}, scale {}", g(shape), g(scale)),
206            Self::ChiSquared { df } => format!("df {}", g(df)),
207            Self::Beta { alpha, beta } => format!("a {}, b {}", g(alpha), g(beta)),
208            Self::StudentsT {
209                df,
210                location,
211                scale,
212            } => format!("df {}, at {}, scale {}", g(df), g(location), g(scale)),
213            Self::Weibull { shape, scale } => format!("shape {}, scale {}", g(shape), g(scale)),
214            Self::PowerLaw { xmin, alpha } => format!("from {}, exponent {}", g(xmin), g(alpha)),
215            Self::Poisson { rate } => format!("rate {}", g(rate)),
216            Self::Bernoulli { p } => format!("p {}", g(p)),
217            Self::Binomial { trials, p } => format!("{trials} trials, p {}", g(p)),
218            Self::Geometric { p, start } => format!("p {}, from {start}", g(p)),
219        }
220    }
221
222    /// `P(X <= x)`.
223    pub fn cdf(&self, x: f64) -> f64 {
224        match *self {
225            Self::Normal { mean, sd } => normal_cdf((x - mean) / sd),
226            Self::LogNormal { mu, sigma } => {
227                if x <= 0.0 {
228                    0.0
229                } else {
230                    normal_cdf((x.ln() - mu) / sigma)
231                }
232            }
233            Self::Uniform { low, high } => ((x - low) / (high - low)).clamp(0.0, 1.0),
234            Self::Exponential { rate } => {
235                if x <= 0.0 {
236                    0.0
237                } else {
238                    -(-rate * x).exp_m1()
239                }
240            }
241            Self::Gamma { shape, scale } => {
242                if x <= 0.0 {
243                    0.0
244                } else {
245                    gamma_p(shape, x / scale)
246                }
247            }
248            Self::ChiSquared { df } => {
249                if x <= 0.0 {
250                    0.0
251                } else {
252                    gamma_p(df / 2.0, x / 2.0)
253                }
254            }
255            Self::Beta { alpha, beta } => {
256                if x <= 0.0 {
257                    0.0
258                } else if x >= 1.0 {
259                    1.0
260                } else {
261                    beta_inc(alpha, beta, x)
262                }
263            }
264            Self::StudentsT {
265                df,
266                location,
267                scale,
268            } => t_cdf((x - location) / scale, df),
269            Self::Weibull { shape, scale } => {
270                if x <= 0.0 {
271                    0.0
272                } else {
273                    -(-(x / scale).powf(shape)).exp_m1()
274                }
275            }
276            Self::PowerLaw { xmin, alpha } => {
277                if x < xmin {
278                    0.0
279                } else {
280                    1.0 - (x / xmin).powf(1.0 - alpha)
281                }
282            }
283            Self::Poisson { rate } => {
284                if x < 0.0 {
285                    0.0
286                } else {
287                    // P(X <= k) = Q(k + 1, rate), the upper regularized gamma.
288                    gamma_q(x.floor() + 1.0, rate)
289                }
290            }
291            Self::Bernoulli { p } => {
292                if x < 0.0 {
293                    0.0
294                } else if x < 1.0 {
295                    1.0 - p
296                } else {
297                    1.0
298                }
299            }
300            Self::Binomial { trials, p } => {
301                let k = x.floor();
302                if k < 0.0 {
303                    0.0
304                } else if k >= trials as f64 {
305                    1.0
306                } else {
307                    // P(X <= k) = I_{1-p}(n - k, k + 1).
308                    beta_inc(trials as f64 - k, k + 1.0, 1.0 - p)
309                }
310            }
311            Self::Geometric { p, start } => {
312                let k = x.floor() - start as f64;
313                if k < 0.0 {
314                    0.0
315                } else {
316                    -((k + 1.0) * (-p).ln_1p()).exp_m1()
317                }
318            }
319        }
320    }
321
322    /// `P(X < x)`: the same as [`Self::cdf`] where no single value carries weight.
323    pub fn cdf_below(&self, x: f64) -> f64 {
324        if !self.discrete() {
325            return self.cdf(x);
326        }
327        if x == x.floor() {
328            self.cdf(x - 1.0)
329        } else {
330            self.cdf(x)
331        }
332    }
333
334    /// The smallest `x` with `P(X <= x) >= p`.
335    pub fn quantile(&self, p: f64) -> f64 {
336        if p.is_nan() {
337            return f64::NAN;
338        }
339        match *self {
340            Self::Normal { mean, sd } => mean + sd * normal_quantile(p),
341            Self::LogNormal { mu, sigma } => (mu + sigma * normal_quantile(p)).exp(),
342            Self::Uniform { low, high } => low + (high - low) * p.clamp(0.0, 1.0),
343            Self::Exponential { rate } => -(-p).ln_1p() / rate,
344            Self::Weibull { shape, scale } => scale * (-(-p).ln_1p()).powf(1.0 / shape),
345            Self::PowerLaw { xmin, alpha } => xmin * (1.0 - p).powf(-1.0 / (alpha - 1.0)),
346            Self::StudentsT {
347                df,
348                location,
349                scale,
350            } => location + scale * t_quantile(p, df),
351            Self::Gamma { .. } | Self::ChiSquared { .. } | Self::Beta { .. } => {
352                self.continuous_quantile_by_search(p)
353            }
354            Self::Poisson { .. }
355            | Self::Bernoulli { .. }
356            | Self::Binomial { .. }
357            | Self::Geometric { .. } => self.discrete_quantile(p),
358        }
359    }
360
361    /// Bisection on the CDF, for the families with no closed-form inverse.
362    fn continuous_quantile_by_search(&self, p: f64) -> f64 {
363        if p <= 0.0 {
364            return match self {
365                Self::Beta { .. } => 0.0,
366                _ => 0.0,
367            };
368        }
369        if p >= 1.0 {
370            return match self {
371                Self::Beta { .. } => 1.0,
372                _ => f64::INFINITY,
373            };
374        }
375        let (mut low, mut high) = match self {
376            Self::Beta { .. } => (0.0, 1.0),
377            _ => {
378                let mut high = 1.0;
379                while self.cdf(high) < p && high < 1e300 {
380                    high *= 2.0;
381                }
382                (0.0, high)
383            }
384        };
385        for _ in 0..200 {
386            let mid = 0.5 * (low + high);
387            if self.cdf(mid) < p {
388                low = mid;
389            } else {
390                high = mid;
391            }
392            if high - low <= 1e-12 * high.abs().max(1e-300) {
393                break;
394            }
395        }
396        0.5 * (low + high)
397    }
398
399    /// The smallest whole number whose CDF reaches `p`, found by bisection from a
400    /// guess: a number of CDF evaluations that grows with the log of the spread. A walk
401    /// one number at a time took a million steps per draw on a geometric fitted to
402    /// values in the millions, and the bootstrap draws a hundred thousand.
403    fn discrete_quantile(&self, p: f64) -> f64 {
404        let start = match *self {
405            Self::Geometric { start, .. } => start as f64,
406            _ => 0.0,
407        };
408        if p <= 0.0 {
409            return start;
410        }
411        if p >= 1.0 {
412            return match *self {
413                Self::Bernoulli { .. } => 1.0,
414                Self::Binomial { trials, .. } => trials as f64,
415                _ => f64::INFINITY,
416            };
417        }
418        let (guess, spread) = match *self {
419            // The geometric's own inverse, exact up to rounding.
420            Self::Geometric { p: q, .. } => {
421                let k = ((-p).ln_1p() / (-q).ln_1p()).ceil() - 1.0;
422                (start + k.max(0.0), 1.0)
423            }
424            Self::Poisson { rate } => (rate + rate.sqrt() * normal_quantile(p), rate.sqrt()),
425            Self::Binomial { trials, p: q } => {
426                let n = trials as f64;
427                let sd = (n * q * (1.0 - q)).sqrt();
428                (n * q + sd * normal_quantile(p), sd)
429            }
430            _ => (start, 1.0),
431        };
432        let top = match *self {
433            Self::Bernoulli { .. } => 1.0,
434            Self::Binomial { trials, .. } => trials as f64,
435            _ => 1e15,
436        };
437        let guess = guess.floor().clamp(start, top);
438        let step = spread.max(1.0).ceil();
439        // The answer lies in (low, high]: cdf(low) < p <= cdf(high), with `low` below
440        // the support when nothing in it falls short.
441        let (mut low, mut high);
442        if self.cdf(guess) >= p {
443            high = guess;
444            low = start - 1.0;
445            let mut step = step;
446            while high > start {
447                let next = (high - step).max(start);
448                if self.cdf(next) < p {
449                    low = next;
450                    break;
451                }
452                high = next;
453                step *= 2.0;
454            }
455        } else {
456            low = guess;
457            high = top;
458            let mut step = step;
459            while low < top {
460                let next = (low + step).min(top);
461                if self.cdf(next) >= p {
462                    high = next;
463                    break;
464                }
465                low = next;
466                step *= 2.0;
467            }
468        }
469        while high - low > 1.0 {
470            let mid = (low + (high - low) / 2.0).floor();
471            if self.cdf(mid) >= p {
472                high = mid;
473            } else {
474                low = mid;
475            }
476        }
477        high
478    }
479
480    /// `ln` of the density, or of the probability at a whole number for a discrete
481    /// family; `-inf` outside the support.
482    fn ln_density(&self, x: f64) -> f64 {
483        match *self {
484            Self::Normal { mean, sd } => {
485                let z = (x - mean) / sd;
486                -0.5 * z * z - sd.ln() - 0.5 * LN_2PI
487            }
488            Self::LogNormal { mu, sigma } => {
489                if x <= 0.0 {
490                    return f64::NEG_INFINITY;
491                }
492                let z = (x.ln() - mu) / sigma;
493                -0.5 * z * z - sigma.ln() - x.ln() - 0.5 * LN_2PI
494            }
495            Self::Uniform { low, high } => {
496                if x < low || x > high {
497                    f64::NEG_INFINITY
498                } else {
499                    -(high - low).ln()
500                }
501            }
502            Self::Exponential { rate } => {
503                if x < 0.0 {
504                    f64::NEG_INFINITY
505                } else {
506                    rate.ln() - rate * x
507                }
508            }
509            Self::Gamma { shape, scale } => gamma_ln_pdf(x, shape, scale),
510            Self::ChiSquared { df } => gamma_ln_pdf(x, df / 2.0, 2.0),
511            Self::Beta { alpha, beta } => {
512                if x <= 0.0 || x >= 1.0 {
513                    return f64::NEG_INFINITY;
514                }
515                (alpha - 1.0) * x.ln() + (beta - 1.0) * (-x).ln_1p() - ln_beta(alpha, beta)
516            }
517            Self::StudentsT {
518                df,
519                location,
520                scale,
521            } => {
522                let t = (x - location) / scale;
523                ln_gamma((df + 1.0) / 2.0)
524                    - ln_gamma(df / 2.0)
525                    - 0.5 * (df * std::f64::consts::PI).ln()
526                    - scale.ln()
527                    - (df + 1.0) / 2.0 * (t * t / df).ln_1p()
528            }
529            Self::Weibull { shape, scale } => {
530                if x < 0.0 {
531                    return f64::NEG_INFINITY;
532                }
533                let z = x / scale;
534                shape.ln() - scale.ln() + (shape - 1.0) * z.ln() - z.powf(shape)
535            }
536            Self::PowerLaw { xmin, alpha } => {
537                if x < xmin {
538                    return f64::NEG_INFINITY;
539                }
540                (alpha - 1.0).ln() - xmin.ln() - alpha * (x / xmin).ln()
541            }
542            Self::Poisson { rate } => {
543                if x < 0.0 || x != x.floor() {
544                    return f64::NEG_INFINITY;
545                }
546                x * rate.ln() - rate - ln_gamma(x + 1.0)
547            }
548            Self::Bernoulli { p } => match x {
549                0.0 => (-p).ln_1p(),
550                1.0 => p.ln(),
551                _ => f64::NEG_INFINITY,
552            },
553            Self::Binomial { trials, p } => {
554                let n = trials as f64;
555                if x < 0.0 || x > n || x != x.floor() {
556                    return f64::NEG_INFINITY;
557                }
558                ln_gamma(n + 1.0) - ln_gamma(x + 1.0) - ln_gamma(n - x + 1.0)
559                    + x * p.ln()
560                    + (n - x) * (-p).ln_1p()
561            }
562            Self::Geometric { p, start } => {
563                let k = x - start as f64;
564                if k < 0.0 || x != x.floor() {
565                    return f64::NEG_INFINITY;
566                }
567                p.ln() + k * (-p).ln_1p()
568            }
569        }
570    }
571
572    /// The density, or the probability at a whole number for a discrete family.
573    pub fn density(&self, x: f64) -> f64 {
574        self.ln_density(x).exp()
575    }
576
577    fn sample(&self, rng: &mut Rng) -> f64 {
578        match *self {
579            Self::Gamma { shape, scale } => rng.gamma(shape) * scale,
580            Self::ChiSquared { df } => rng.gamma(df / 2.0) * 2.0,
581            Self::Beta { alpha, beta } => {
582                let a = rng.gamma(alpha);
583                let b = rng.gamma(beta);
584                a / (a + b)
585            }
586            Self::StudentsT {
587                df,
588                location,
589                scale,
590            } => {
591                let z = rng.normal();
592                let v = rng.gamma(df / 2.0) * 2.0;
593                location + scale * z / (v / df).sqrt()
594            }
595            // Inverting the CDF costs a few dozen incomplete gammas or betas per draw,
596            // each longer as the counts grow: seconds to minutes per test at counts in
597            // the millions.
598            Self::Poisson { rate } if rate >= 10.0 => rng.poisson(rate),
599            Self::Binomial { trials, p } if trials as f64 * p.min(1.0 - p) >= 10.0 => {
600                rng.binomial(trials, p)
601            }
602            _ => self.quantile(rng.uniform()),
603        }
604    }
605
606    /// Estimate a family's parameters from `values`, or say why the family cannot
607    /// describe them.
608    pub fn fit(family: DistributionType, values: &[f64]) -> Result<Fitted, &'static str> {
609        let n = values.len();
610        if n < 5 {
611            return Err("needs 5 or more values");
612        }
613        let nf = n as f64;
614        let mean = values.iter().sum::<f64>() / nf;
615        let variance = values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / (nf - 1.0);
616        let min = values.iter().copied().fold(f64::INFINITY, f64::min);
617        let max = values.iter().copied().fold(f64::NEG_INFINITY, f64::max);
618        let counts = values.iter().all(|v| *v >= 0.0 && *v == v.floor());
619        if variance <= 0.0 || !variance.is_finite() {
620            return Err("needs values that vary");
621        }
622        let positive = || {
623            if min > 0.0 {
624                Ok(())
625            } else {
626                Err("needs positive values")
627            }
628        };
629        let fitted = match family {
630            DistributionType::Normal => Fitted::Normal {
631                mean,
632                sd: variance.sqrt(),
633            },
634            DistributionType::LogNormal => {
635                positive()?;
636                let logs: Vec<f64> = values.iter().map(|v| v.ln()).collect();
637                let mu = logs.iter().sum::<f64>() / nf;
638                let var = logs.iter().map(|v| (v - mu).powi(2)).sum::<f64>() / (nf - 1.0);
639                Fitted::LogNormal {
640                    mu,
641                    sigma: var.sqrt(),
642                }
643            }
644            DistributionType::Uniform => Fitted::Uniform {
645                low: min,
646                high: max,
647            },
648            DistributionType::Exponential => {
649                if min < 0.0 {
650                    return Err("needs values of zero or more");
651                }
652                Fitted::Exponential { rate: 1.0 / mean }
653            }
654            DistributionType::Gamma => {
655                positive()?;
656                // Minka's closed-form approximation to the maximum-likelihood shape.
657                let mean_log = values.iter().map(|v| v.ln()).sum::<f64>() / nf;
658                let s = mean.ln() - mean_log;
659                if s <= 0.0 {
660                    return Err("needs values that vary");
661                }
662                let shape = (3.0 - s + ((s - 3.0).powi(2) + 24.0 * s).sqrt()) / (12.0 * s);
663                Fitted::Gamma {
664                    shape,
665                    scale: mean / shape,
666                }
667            }
668            DistributionType::ChiSquared => {
669                positive()?;
670                Fitted::ChiSquared { df: mean }
671            }
672            DistributionType::Beta => {
673                if min <= 0.0 || max >= 1.0 {
674                    return Err("needs values strictly between 0 and 1");
675                }
676                let common = mean * (1.0 - mean) / variance - 1.0;
677                if common <= 0.0 {
678                    return Err("spread too wide for a beta");
679                }
680                // Moments to start, then maximum likelihood, as every family here is
681                // fitted where it can be: AIC compares likelihoods, and a family
682                // fitted by moments competes with a handicap.
683                beta_mle(values, mean * common, (1.0 - mean) * common)
684            }
685            DistributionType::StudentsT => {
686                let fourth = values.iter().map(|v| (v - mean).powi(4)).sum::<f64>() / nf;
687                let excess = fourth / (variance * variance) - 3.0;
688                if excess <= 0.0 {
689                    return Err("tails no heavier than a normal's");
690                }
691                // Kurtosis 6 / (df - 4) above a normal's, for df above 4.
692                let df = (4.0 + 6.0 / excess).clamp(2.5, 1_000.0);
693                Fitted::StudentsT {
694                    df,
695                    location: mean,
696                    scale: (variance * (df - 2.0) / df).sqrt(),
697                }
698            }
699            DistributionType::Weibull => {
700                positive()?;
701                weibull_mle(values).ok_or("no weibull fits these values")?
702            }
703            DistributionType::PowerLaw => {
704                positive()?;
705                let sum_log = values.iter().map(|v| (v / min).ln()).sum::<f64>();
706                if sum_log <= 0.0 {
707                    return Err("needs values that vary");
708                }
709                Fitted::PowerLaw {
710                    xmin: min,
711                    alpha: 1.0 + nf / sum_log,
712                }
713            }
714            DistributionType::Poisson => {
715                if !counts {
716                    return Err("needs counts: whole numbers, zero or more");
717                }
718                Fitted::Poisson { rate: mean }
719            }
720            DistributionType::Bernoulli => {
721                if !values.iter().all(|v| *v == 0.0 || *v == 1.0) {
722                    return Err("needs only 0 and 1");
723                }
724                Fitted::Bernoulli { p: mean }
725            }
726            DistributionType::Binomial => {
727                if !counts {
728                    return Err("needs counts: whole numbers, zero or more");
729                }
730                if variance >= mean {
731                    return Err("spread too wide for a binomial");
732                }
733                // Moments, then at least as many trials as the largest count.
734                let p = 1.0 - variance / mean;
735                let trials = ((mean / p).round() as u64).max(max as u64).max(1);
736                Fitted::Binomial {
737                    trials,
738                    p: (mean / trials as f64).clamp(1e-12, 1.0 - 1e-12),
739                }
740            }
741            DistributionType::Geometric => {
742                if !counts {
743                    return Err("needs counts: whole numbers, zero or more");
744                }
745                // Counted from 1 (trials) when no zero appears, from 0 (failures) when
746                // one does.
747                let start = if min >= 1.0 { 1 } else { 0 };
748                Fitted::Geometric {
749                    p: 1.0 / (1.0 + mean - start as f64),
750                    start,
751                }
752            }
753            DistributionType::Constant | DistributionType::Unknown => {
754                return Err("not a distribution");
755            }
756        };
757        if fitted.valid() {
758            Ok(fitted)
759        } else {
760            Err("no fit for these values")
761        }
762    }
763
764    fn valid(&self) -> bool {
765        let finite_positive = |x: f64| x.is_finite() && x > 0.0;
766        match *self {
767            Self::Normal { mean, sd } => mean.is_finite() && finite_positive(sd),
768            Self::LogNormal { mu, sigma } => mu.is_finite() && finite_positive(sigma),
769            Self::Uniform { low, high } => low.is_finite() && high.is_finite() && high > low,
770            Self::Exponential { rate } => finite_positive(rate),
771            Self::Gamma { shape, scale } => finite_positive(shape) && finite_positive(scale),
772            Self::ChiSquared { df } => finite_positive(df),
773            Self::Beta { alpha, beta } => finite_positive(alpha) && finite_positive(beta),
774            Self::StudentsT {
775                df,
776                location,
777                scale,
778            } => finite_positive(df) && location.is_finite() && finite_positive(scale),
779            Self::Weibull { shape, scale } => finite_positive(shape) && finite_positive(scale),
780            Self::PowerLaw { xmin, alpha } => {
781                finite_positive(xmin) && alpha.is_finite() && alpha > 1.0
782            }
783            Self::Poisson { rate } => finite_positive(rate),
784            Self::Bernoulli { p } | Self::Geometric { p, .. } => p > 0.0 && p < 1.0,
785            Self::Binomial { trials, p } => trials >= 1 && p > 0.0 && p < 1.0,
786        }
787    }
788}
789
790/// Maximum-likelihood Beta by Newton's method on both shapes, from `alpha` and `beta`:
791/// `ψ(a) - ψ(a + b)` and `ψ(b) - ψ(a + b)` equal the mean logs of `x` and `1 - x`.
792fn beta_mle(values: &[f64], alpha: f64, beta: f64) -> Fitted {
793    let n = values.len() as f64;
794    let mean_log = values.iter().map(|v| v.ln()).sum::<f64>() / n;
795    let mean_log_rest = values.iter().map(|v| (-v).ln_1p()).sum::<f64>() / n;
796    let (mut a, mut b) = (alpha, beta);
797    for _ in 0..100 {
798        let both = digamma(a + b);
799        let g1 = digamma(a) - both - mean_log;
800        let g2 = digamma(b) - both - mean_log_rest;
801        let t = trigamma(a + b);
802        let (j11, j22, j12) = (trigamma(a) - t, trigamma(b) - t, -t);
803        let det = j11 * j22 - j12 * j12;
804        if det.abs() < 1e-300 {
805            break;
806        }
807        let da = (g1 * j22 - g2 * j12) / det;
808        let db = (g2 * j11 - g1 * j12) / det;
809        // Halved until both stay positive.
810        let mut step = 1.0;
811        while (a - step * da <= 0.0 || b - step * db <= 0.0) && step > 1e-6 {
812            step /= 2.0;
813        }
814        let (next_a, next_b) = (a - step * da, b - step * db);
815        let done = (next_a - a).abs() < 1e-10 * a && (next_b - b).abs() < 1e-10 * b;
816        a = next_a;
817        b = next_b;
818        if done {
819            break;
820        }
821    }
822    Fitted::Beta { alpha: a, beta: b }
823}
824
825/// `ψ(x)`, the digamma function, for `x > 0`: shifted up past 6, then its asymptotic
826/// series.
827fn digamma(mut x: f64) -> f64 {
828    let mut result = 0.0;
829    while x < 6.0 {
830        result -= 1.0 / x;
831        x += 1.0;
832    }
833    let inv = 1.0 / x;
834    let inv2 = inv * inv;
835    result + x.ln()
836        - 0.5 * inv
837        - inv2 * (1.0 / 12.0 - inv2 * (1.0 / 120.0 - inv2 * (1.0 / 252.0 - inv2 / 240.0)))
838}
839
840/// `ψ'(x)`, the trigamma function, for `x > 0`.
841fn trigamma(mut x: f64) -> f64 {
842    let mut result = 0.0;
843    while x < 6.0 {
844        result += 1.0 / (x * x);
845        x += 1.0;
846    }
847    let inv = 1.0 / x;
848    let inv2 = inv * inv;
849    result
850        + inv
851        + 0.5 * inv2
852        + inv * inv2 * (1.0 / 6.0 - inv2 * (1.0 / 30.0 - inv2 * (1.0 / 42.0 - inv2 / 30.0)))
853}
854
855/// Maximum-likelihood Weibull by Newton's method on the shape.
856fn weibull_mle(values: &[f64]) -> Option<Fitted> {
857    let n = values.len() as f64;
858    let logs: Vec<f64> = values.iter().map(|v| v.ln()).collect();
859    let mean_log = logs.iter().sum::<f64>() / n;
860    // Scaled so the powers do not overflow for large values.
861    let top = values.iter().copied().fold(0.0, f64::max);
862    let scaled: Vec<f64> = values.iter().map(|v| v / top).collect();
863    let scaled_logs: Vec<f64> = scaled.iter().map(|v| v.ln()).collect();
864    let mean_scaled_log = scaled_logs.iter().sum::<f64>() / n;
865    let sd_log = (logs.iter().map(|l| (l - mean_log).powi(2)).sum::<f64>() / n).sqrt();
866    // The shape a Gumbel on the logs suggests, as a start.
867    let mut k = if sd_log > 0.0 { 1.2825 / sd_log } else { 1.0 };
868    for _ in 0..100 {
869        let (mut s0, mut s1, mut s2) = (0.0, 0.0, 0.0);
870        for (x, l) in scaled.iter().zip(&scaled_logs) {
871            let xk = x.powf(k);
872            s0 += xk;
873            s1 += xk * l;
874            s2 += xk * l * l;
875        }
876        if s0 <= 0.0 {
877            return None;
878        }
879        let f = s1 / s0 - 1.0 / k - mean_scaled_log;
880        let df = (s2 * s0 - s1 * s1) / (s0 * s0) + 1.0 / (k * k);
881        let step = f / df;
882        let next = (k - step).clamp(k / 2.0, k * 2.0);
883        let done = (next - k).abs() < 1e-10 * k;
884        k = next;
885        if done {
886            break;
887        }
888    }
889    let s0 = scaled.iter().map(|x| x.powf(k)).sum::<f64>();
890    let scale = top * (s0 / n).powf(1.0 / k);
891    Some(Fitted::Weibull { shape: k, scale })
892}
893
894/// The two-sided Kolmogorov-Smirnov statistic of sorted `values` against `fitted`:
895/// the largest gap between the empirical and fitted CDFs, on either side of each step.
896/// Tied values are one step, compared at the value and just below it, so a discrete
897/// fit is measured where its jumps are.
898pub fn ks_statistic(sorted: &[f64], fitted: &Fitted) -> f64 {
899    let n = sorted.len() as f64;
900    let mut d: f64 = 0.0;
901    let mut i = 0;
902    while i < sorted.len() {
903        let x = sorted[i];
904        let mut j = i;
905        while j < sorted.len() && sorted[j] == x {
906            j += 1;
907        }
908        let below = i as f64 / n;
909        let at = j as f64 / n;
910        d = d
911            .max((at - fitted.cdf(x)).abs())
912            .max((fitted.cdf_below(x) - below).abs());
913        i = j;
914    }
915    d
916}
917
918/// Test one family against `values`, all of them as read.
919pub fn test_family(family: DistributionType, values: &[f64], seed: u64) -> FitOutcome {
920    let fitted = match Fitted::fit(family, values) {
921        Ok(fitted) => fitted,
922        Err(reason) => return FitOutcome::NotApplicable(reason),
923    };
924    let ln_likelihood: f64 = values.iter().map(|v| fitted.ln_density(*v)).sum();
925    let aic = 2.0 * fitted.parameters() as f64 - 2.0 * ln_likelihood;
926
927    // The test runs on a seeded random subset, refitted: the bootstrap below refits
928    // each simulated sample, and the column has to be treated the same way.
929    let mut rng = Rng::new(seed ^ family as u64);
930    let mut subset = values.to_vec();
931    if subset.len() > TEST_VALUES {
932        for i in 0..TEST_VALUES {
933            let j = i + (rng.next_u64() % (subset.len() - i) as u64) as usize;
934            subset.swap(i, j);
935        }
936        subset.truncate(TEST_VALUES);
937    }
938    subset.sort_by(f64::total_cmp);
939    let Ok(sub_fit) = Fitted::fit(family, &subset) else {
940        return FitOutcome::NotApplicable("no fit for these values");
941    };
942    let observed = ks_statistic(&subset, &sub_fit);
943
944    let mut beyond = 0;
945    let mut simulated = vec![0.0; subset.len()];
946    for _ in 0..REPLICATES {
947        for value in simulated.iter_mut() {
948            *value = sub_fit.sample(&mut rng);
949        }
950        simulated.sort_by(f64::total_cmp);
951        // A simulated sample the family cannot refit counts against the fit: the test
952        // leans towards not rejecting rather than towards a p-value it did not earn.
953        let statistic = match Fitted::fit(family, &simulated) {
954            Ok(refit) => ks_statistic(&simulated, &refit),
955            Err(_) => f64::INFINITY,
956        };
957        if statistic >= observed {
958            beyond += 1;
959        }
960    }
961    FitOutcome::Tested(FitTest {
962        fitted,
963        p_value: (1 + beyond) as f64 / (1 + REPLICATES) as f64,
964        beyond,
965        replicates: REPLICATES,
966        tested_on: subset.len(),
967        aic,
968    })
969}
970
971/// Test every family, in parallel: each is independent, and each is a few hundred
972/// thousand CDF evaluations.
973pub fn test_all(values: &[f64], seed: u64) -> Vec<(DistributionType, FitOutcome)> {
974    std::thread::scope(|scope| {
975        let handles: Vec<_> = FAMILIES
976            .iter()
977            .map(|family| scope.spawn(move || (*family, test_family(*family, values, seed))))
978            .collect();
979        handles
980            .into_iter()
981            .filter_map(|handle| handle.join().ok())
982            .collect()
983    })
984}
985
986/// Below this p-value a family is rejected. At 0.05 the true family of a column is
987/// rejected one time in twenty by design, and a verdict that flips between samples of
988/// the same data says more about the sample.
989pub const REJECT_P: f64 = 0.01;
990
991/// A family that contains a simpler one fits whatever the simpler one does: uniform
992/// values are a Beta(1, 1). When the simpler one is not rejected, it is the answer.
993fn nested_in(family: DistributionType) -> &'static [DistributionType] {
994    match family {
995        DistributionType::Beta => &[DistributionType::Uniform],
996        DistributionType::Gamma => &[DistributionType::Exponential, DistributionType::ChiSquared],
997        DistributionType::Weibull => &[DistributionType::Exponential],
998        DistributionType::StudentsT => &[DistributionType::Normal],
999        // A binomial of one trial is a Bernoulli; of many trials and a small p, a
1000        // Poisson.
1001        DistributionType::Binomial => &[DistributionType::Bernoulli, DistributionType::Poisson],
1002        _ => &[],
1003    }
1004}
1005
1006/// The family the values are consistent with, or `Unknown` when every test rejects
1007/// them. Among those not rejected, the lowest AIC; counts are described by a count
1008/// distribution when one is not rejected, since a density and a probability are not
1009/// on one scale.
1010pub fn select(results: &[(DistributionType, FitOutcome)], counts: bool) -> DistributionType {
1011    let held: Vec<&FitTest> = results
1012        .iter()
1013        .filter_map(|(_, outcome)| outcome.test())
1014        .filter(|test| test.p_value >= REJECT_P)
1015        .collect();
1016    let discrete_held = counts && held.iter().any(|test| test.fitted.discrete());
1017    let best = held
1018        .iter()
1019        .filter(|test| test.fitted.discrete() == discrete_held)
1020        .min_by(|a, b| a.aic.total_cmp(&b.aic));
1021    let Some(best) = best else {
1022        return DistributionType::Unknown;
1023    };
1024    let family = best.fitted.family();
1025    // The simpler of two nested families when it holds and the richer one is not
1026    // decisively better: an AIC gap of 10 or more is the usual line for "no support".
1027    // A test on a few hundred values can miss a t's tails that every value shows.
1028    nested_in(family)
1029        .iter()
1030        .find(|simpler| {
1031            held.iter()
1032                .any(|test| test.fitted.family() == **simpler && test.aic - best.aic < DECISIVE_AIC)
1033        })
1034        .copied()
1035        .unwrap_or(family)
1036}
1037
1038/// An AIC gap at which the family behind has essentially no support.
1039const DECISIVE_AIC: f64 = 10.0;
1040
1041/// The families as a list is read: tested ones by p-value, highest first, then the
1042/// ones that do not apply, in [`FAMILIES`] order.
1043pub fn listing_order(results: &[(DistributionType, FitOutcome)]) -> Vec<DistributionType> {
1044    let mut tested: Vec<(DistributionType, f64)> = results
1045        .iter()
1046        .filter_map(|(family, outcome)| outcome.p_value().map(|p| (*family, p)))
1047        .collect();
1048    tested.sort_by(|a, b| b.1.total_cmp(&a.1));
1049    let mut order: Vec<DistributionType> = tested.into_iter().map(|(family, _)| family).collect();
1050    let missing: Vec<DistributionType> = FAMILIES
1051        .iter()
1052        .filter(|family| !order.contains(family))
1053        .copied()
1054        .collect();
1055    order.extend(missing);
1056    order
1057}
1058
1059/// Theoretical quantiles for a Q-Q plot of `n` sorted values: the fitted quantile at
1060/// each plotting position `i / (n + 1)`.
1061pub fn qq_quantiles(fitted: &Fitted, n: usize) -> Vec<f64> {
1062    (1..=n)
1063        .map(|i| fitted.quantile(i as f64 / (n as f64 + 1.0)))
1064        .collect()
1065}
1066
1067fn format_param(x: f64) -> String {
1068    if x != 0.0 && (x.abs() >= 1e5 || x.abs() < 1e-3) {
1069        format!("{x:.3e}")
1070    } else if x.abs() >= 100.0 {
1071        format!("{x:.1}")
1072    } else {
1073        format!("{x:.3}")
1074    }
1075}
1076
1077const LN_2PI: f64 = 1.837_877_066_409_345_5;
1078
1079/// `ln Γ(x)` for `x > 0`, by the Lanczos approximation (g = 7, nine terms), accurate
1080/// to about 1e-15.
1081pub fn ln_gamma(x: f64) -> f64 {
1082    const COEFFICIENTS: [f64; 9] = [
1083        0.999_999_999_999_809_9,
1084        676.520_368_121_885_1,
1085        -1_259.139_216_722_402_8,
1086        771.323_428_777_653_1,
1087        -176.615_029_162_140_6,
1088        12.507_343_278_686_905,
1089        -0.138_571_095_265_720_12,
1090        9.984_369_578_019_572e-6,
1091        1.505_632_735_149_311_6e-7,
1092    ];
1093    if x < 0.5 {
1094        // Reflection, for the half-integers the t density needs near zero.
1095        let pi = std::f64::consts::PI;
1096        return (pi / (pi * x).sin()).ln() - ln_gamma(1.0 - x);
1097    }
1098    let x = x - 1.0;
1099    let mut sum = COEFFICIENTS[0];
1100    for (i, c) in COEFFICIENTS.iter().enumerate().skip(1) {
1101        sum += c / (x + i as f64);
1102    }
1103    let t = x + 7.5;
1104    0.5 * LN_2PI + (x + 0.5) * t.ln() - t + sum.ln()
1105}
1106
1107fn ln_beta(a: f64, b: f64) -> f64 {
1108    ln_gamma(a) + ln_gamma(b) - ln_gamma(a + b)
1109}
1110
1111fn gamma_ln_pdf(x: f64, shape: f64, scale: f64) -> f64 {
1112    if x <= 0.0 {
1113        return f64::NEG_INFINITY;
1114    }
1115    (shape - 1.0) * x.ln() - x / scale - shape * scale.ln() - ln_gamma(shape)
1116}
1117
1118/// The regularized lower incomplete gamma `P(a, x)`.
1119pub fn gamma_p(a: f64, x: f64) -> f64 {
1120    let (p, _) = incomplete_gamma(a, x);
1121    p
1122}
1123
1124/// Its complement `Q(a, x) = 1 - P(a, x)`, computed directly: in the upper tail `P` is
1125/// within 1e-16 of 1, and `1 - P` would be noise.
1126pub fn gamma_q(a: f64, x: f64) -> f64 {
1127    let (_, q) = incomplete_gamma(a, x);
1128    q
1129}
1130
1131/// `(P(a, x), Q(a, x))`: by the series below `a + 1`, by the continued fraction for
1132/// `Q` above (Numerical Recipes, `gser` and `gcf`), each side taken from whichever is
1133/// computed without cancellation.
1134fn incomplete_gamma(a: f64, x: f64) -> (f64, f64) {
1135    if x <= 0.0 {
1136        return (0.0, 1.0);
1137    }
1138    if !x.is_finite() {
1139        return (1.0, 0.0);
1140    }
1141    let ln_prefix = a * x.ln() - x - ln_gamma(a);
1142    if x < a + 1.0 {
1143        let mut term = 1.0 / a;
1144        let mut sum = term;
1145        let mut ap = a;
1146        for _ in 0..10_000 {
1147            ap += 1.0;
1148            term *= x / ap;
1149            sum += term;
1150            if term.abs() < sum.abs() * 1e-16 {
1151                break;
1152            }
1153        }
1154        let p = (sum.ln() + ln_prefix).exp().clamp(0.0, 1.0);
1155        (p, 1.0 - p)
1156    } else {
1157        // Lentz's method for the continued fraction of Q(a, x).
1158        let tiny = 1e-300;
1159        let mut b = x + 1.0 - a;
1160        let mut c = 1.0 / tiny;
1161        let mut d = 1.0 / b;
1162        let mut h = d;
1163        for i in 1..10_000 {
1164            let an = -(i as f64) * (i as f64 - a);
1165            b += 2.0;
1166            d = an * d + b;
1167            if d.abs() < tiny {
1168                d = tiny;
1169            }
1170            c = b + an / c;
1171            if c.abs() < tiny {
1172                c = tiny;
1173            }
1174            d = 1.0 / d;
1175            let delta = d * c;
1176            h *= delta;
1177            if (delta - 1.0).abs() < 1e-16 {
1178                break;
1179            }
1180        }
1181        let q = (ln_prefix + h.ln()).exp().clamp(0.0, 1.0);
1182        (1.0 - q, q)
1183    }
1184}
1185
1186/// The regularized incomplete beta `I_x(a, b)` (Numerical Recipes, `betai`).
1187pub fn beta_inc(a: f64, b: f64, x: f64) -> f64 {
1188    if x <= 0.0 {
1189        return 0.0;
1190    }
1191    if x >= 1.0 {
1192        return 1.0;
1193    }
1194    let ln_front = a * x.ln() + b * (-x).ln_1p() - ln_beta(a, b);
1195    if x < (a + 1.0) / (a + b + 2.0) {
1196        (ln_front.exp() * beta_fraction(a, b, x) / a).clamp(0.0, 1.0)
1197    } else {
1198        (1.0 - ln_front.exp() * beta_fraction(b, a, 1.0 - x) / b).clamp(0.0, 1.0)
1199    }
1200}
1201
1202fn beta_fraction(a: f64, b: f64, x: f64) -> f64 {
1203    let tiny = 1e-300;
1204    let (qab, qap, qam) = (a + b, a + 1.0, a - 1.0);
1205    let mut c = 1.0;
1206    let mut d = 1.0 - qab * x / qap;
1207    if d.abs() < tiny {
1208        d = tiny;
1209    }
1210    d = 1.0 / d;
1211    let mut h = d;
1212    for m in 1..10_000 {
1213        let m = m as f64;
1214        let m2 = 2.0 * m;
1215        let aa = m * (b - m) * x / ((qam + m2) * (a + m2));
1216        d = 1.0 + aa * d;
1217        if d.abs() < tiny {
1218            d = tiny;
1219        }
1220        c = 1.0 + aa / c;
1221        if c.abs() < tiny {
1222            c = tiny;
1223        }
1224        d = 1.0 / d;
1225        h *= d * c;
1226        let aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2));
1227        d = 1.0 + aa * d;
1228        if d.abs() < tiny {
1229            d = tiny;
1230        }
1231        c = 1.0 + aa / c;
1232        if c.abs() < tiny {
1233            c = tiny;
1234        }
1235        d = 1.0 / d;
1236        let delta = d * c;
1237        h *= delta;
1238        if (delta - 1.0).abs() < 1e-16 {
1239            break;
1240        }
1241    }
1242    h
1243}
1244
1245/// The standard normal CDF, through `erfc(x) = Q(1/2, x²)`: accurate in the tails,
1246/// where `1 - erf` loses everything.
1247pub fn normal_cdf(z: f64) -> f64 {
1248    if z.is_nan() {
1249        return f64::NAN;
1250    }
1251    let tail = 0.5 * gamma_q(0.5, 0.5 * z * z);
1252    if z < 0.0 { tail } else { 1.0 - tail }
1253}
1254
1255/// The standard normal quantile. Acklam's rational approximation, then one Halley step
1256/// on the exact CDF: good to about 1e-15. Exactly 0 at the median, `-inf` and `inf` at
1257/// 0 and 1, NaN outside them.
1258pub fn normal_quantile(p: f64) -> f64 {
1259    if !(0.0..=1.0).contains(&p) {
1260        return f64::NAN;
1261    }
1262    if p == 0.0 {
1263        return f64::NEG_INFINITY;
1264    }
1265    if p == 1.0 {
1266        return f64::INFINITY;
1267    }
1268    if p == 0.5 {
1269        return 0.0;
1270    }
1271    const A: [f64; 6] = [
1272        -3.969_683_028_665_376e1,
1273        2.209_460_984_245_205e2,
1274        -2.759_285_104_469_687e2,
1275        1.383_577_518_672_69e2,
1276        -3.066_479_806_614_716e1,
1277        2.506_628_277_459_239,
1278    ];
1279    const B: [f64; 5] = [
1280        -5.447_609_879_822_406e1,
1281        1.615_858_368_580_409e2,
1282        -1.556_989_798_598_866e2,
1283        6.680_131_188_771_972e1,
1284        -1.328_068_155_288_572e1,
1285    ];
1286    const C: [f64; 6] = [
1287        -7.784_894_002_430_293e-3,
1288        -3.223_964_580_411_365e-1,
1289        -2.400_758_277_161_838,
1290        -2.549_732_539_343_734,
1291        4.374_664_141_464_968,
1292        2.938_163_982_698_783,
1293    ];
1294    const D: [f64; 4] = [
1295        7.784_695_709_041_462e-3,
1296        3.224_671_290_700_398e-1,
1297        2.445_134_137_142_996,
1298        3.754_408_661_907_416,
1299    ];
1300    let low = 0.02425;
1301    let x = if p < low {
1302        let q = (-2.0 * p.ln()).sqrt();
1303        (((((C[0] * q + C[1]) * q + C[2]) * q + C[3]) * q + C[4]) * q + C[5])
1304            / ((((D[0] * q + D[1]) * q + D[2]) * q + D[3]) * q + 1.0)
1305    } else if p <= 1.0 - low {
1306        let q = p - 0.5;
1307        let r = q * q;
1308        (((((A[0] * r + A[1]) * r + A[2]) * r + A[3]) * r + A[4]) * r + A[5]) * q
1309            / (((((B[0] * r + B[1]) * r + B[2]) * r + B[3]) * r + B[4]) * r + 1.0)
1310    } else {
1311        let q = (-2.0 * (-p).ln_1p()).sqrt();
1312        -(((((C[0] * q + C[1]) * q + C[2]) * q + C[3]) * q + C[4]) * q + C[5])
1313            / ((((D[0] * q + D[1]) * q + D[2]) * q + D[3]) * q + 1.0)
1314    };
1315    // One Halley step against the exact CDF.
1316    let e = normal_cdf(x) - p;
1317    let u = e * (2.0 * std::f64::consts::PI).sqrt() * (0.5 * x * x).exp();
1318    x - u / (1.0 + 0.5 * x * u)
1319}
1320
1321/// Student's t CDF with `df` degrees of freedom.
1322pub fn t_cdf(t: f64, df: f64) -> f64 {
1323    if t.is_nan() {
1324        return f64::NAN;
1325    }
1326    let tail = 0.5 * beta_inc(0.5 * df, 0.5, df / (df + t * t));
1327    if t > 0.0 { 1.0 - tail } else { tail }
1328}
1329
1330fn t_quantile(p: f64, df: f64) -> f64 {
1331    if p <= 0.0 {
1332        return f64::NEG_INFINITY;
1333    }
1334    if p >= 1.0 {
1335        return f64::INFINITY;
1336    }
1337    if p == 0.5 {
1338        return 0.0;
1339    }
1340    let mut high = 1.0;
1341    let target = p.max(1.0 - p);
1342    while t_cdf(high, df) < target && high < 1e300 {
1343        high *= 2.0;
1344    }
1345    let (mut low, mut hi) = (0.0, high);
1346    for _ in 0..200 {
1347        let mid = 0.5 * (low + hi);
1348        if t_cdf(mid, df) < target {
1349            low = mid;
1350        } else {
1351            hi = mid;
1352        }
1353        if hi - low <= 1e-13 * hi {
1354            break;
1355        }
1356    }
1357    let t = 0.5 * (low + hi);
1358    if p < 0.5 { -t } else { t }
1359}
1360
1361/// A small seeded generator (SplitMix64): the same seed draws the same samples on
1362/// every machine, which a test of a test needs.
1363pub struct Rng(u64);
1364
1365impl Rng {
1366    pub fn new(seed: u64) -> Self {
1367        Self(seed)
1368    }
1369
1370    pub fn next_u64(&mut self) -> u64 {
1371        self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
1372        let mut z = self.0;
1373        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
1374        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
1375        z ^ (z >> 31)
1376    }
1377
1378    /// Uniform on (0, 1), never either end.
1379    pub fn uniform(&mut self) -> f64 {
1380        ((self.next_u64() >> 11) as f64 + 0.5) / (1u64 << 53) as f64
1381    }
1382
1383    pub fn normal(&mut self) -> f64 {
1384        normal_quantile(self.uniform())
1385    }
1386
1387    /// A Gamma(shape, 1) draw, by Marsaglia and Tsang's method.
1388    pub fn gamma(&mut self, shape: f64) -> f64 {
1389        if shape < 1.0 {
1390            let boost = self.uniform().powf(1.0 / shape);
1391            return self.gamma(shape + 1.0) * boost;
1392        }
1393        let d = shape - 1.0 / 3.0;
1394        let c = 1.0 / (9.0 * d).sqrt();
1395        loop {
1396            let x = self.normal();
1397            let v = 1.0 + c * x;
1398            if v <= 0.0 {
1399                continue;
1400            }
1401            let v = v * v * v;
1402            let u = self.uniform();
1403            if u.ln() < 0.5 * x * x + d - d * v + d * v.ln() {
1404                return d * v;
1405            }
1406        }
1407    }
1408
1409    /// A Poisson(rate) draw for a rate of 10 or more, by Hörmann's transformed
1410    /// rejection (PTRS): exact, and a constant cost whatever the rate.
1411    pub fn poisson(&mut self, rate: f64) -> f64 {
1412        let ln_rate = rate.ln();
1413        let b = 0.931 + 2.53 * rate.sqrt();
1414        let a = -0.059 + 0.02483 * b;
1415        let ln_inv_alpha = (1.1239 + 1.1328 / (b - 3.4)).ln();
1416        let v_r = 0.9277 - 3.6224 / (b - 2.0);
1417        loop {
1418            let u = self.uniform() - 0.5;
1419            let v = self.uniform();
1420            let us = 0.5 - u.abs();
1421            let k = ((2.0 * a / us + b) * u + rate + 0.43).floor();
1422            if us >= 0.07 && v <= v_r {
1423                return k;
1424            }
1425            if k < 0.0 || (us < 0.013 && v > us) {
1426                continue;
1427            }
1428            if v.ln() + ln_inv_alpha - (a / (us * us) + b).ln()
1429                <= -rate + k * ln_rate - ln_gamma(k + 1.0)
1430            {
1431                return k;
1432            }
1433        }
1434    }
1435
1436    /// A Binomial(trials, p) draw for `trials * min(p, 1 - p)` of 10 or more, by
1437    /// Hörmann's transformed rejection (BTRS): exact, and a constant cost whatever
1438    /// the trials.
1439    pub fn binomial(&mut self, trials: u64, p: f64) -> f64 {
1440        if p > 0.5 {
1441            return trials as f64 - self.binomial(trials, 1.0 - p);
1442        }
1443        let n = trials as f64;
1444        let q = 1.0 - p;
1445        let spq = (n * p * q).sqrt();
1446        let b = 1.15 + 2.53 * spq;
1447        let a = -0.0873 + 0.0248 * b + 0.01 * p;
1448        let c = n * p + 0.5;
1449        let v_r = 0.92 - 4.2 / b;
1450        let alpha = (2.83 + 5.1 / b) * spq;
1451        let ln_pq = (p / q).ln();
1452        let m = ((n + 1.0) * p).floor();
1453        let h = ln_gamma(m + 1.0) + ln_gamma(n - m + 1.0);
1454        loop {
1455            let u = self.uniform() - 0.5;
1456            let v = self.uniform();
1457            let us = 0.5 - u.abs();
1458            let k = ((2.0 * a / us + b) * u + c).floor();
1459            if k < 0.0 || k > n {
1460                continue;
1461            }
1462            if us >= 0.07 && v <= v_r {
1463                return k;
1464            }
1465            let v = (v * alpha / (a / (us * us) + b)).ln();
1466            if v <= h - ln_gamma(k + 1.0) - ln_gamma(n - k + 1.0) + (k - m) * ln_pq {
1467                return k;
1468            }
1469        }
1470    }
1471}
1472
1473#[cfg(test)]
1474mod tests {
1475    use super::*;
1476
1477    fn close(a: f64, b: f64, tolerance: f64) -> bool {
1478        (a - b).abs() <= tolerance * b.abs().max(1.0)
1479    }
1480
1481    /// Reference values from Wichura's AS241 (Python's `NormalDist().inv_cdf`).
1482    #[test]
1483    fn normal_quantiles_match_reference_values() {
1484        for (p, z) in [
1485            (0.975, 1.959_963_984_540_054),
1486            (0.841_344_746_068_542_9, 1.0),
1487            (0.001, -3.090_232_306_167_813),
1488            (1e-10, -6.361_340_902_404_056),
1489            (0.3, -0.524_400_512_708_041),
1490        ] {
1491            assert!(
1492                close(normal_quantile(p), z, 1e-12),
1493                "{p}: {}",
1494                normal_quantile(p)
1495            );
1496        }
1497        assert_eq!(normal_quantile(0.5), 0.0);
1498        assert_eq!(normal_quantile(0.0), f64::NEG_INFINITY);
1499        assert_eq!(normal_quantile(1.0), f64::INFINITY);
1500        assert!(normal_quantile(1.5).is_nan());
1501        let grid: Vec<f64> = (1..1000)
1502            .map(|i| normal_quantile(i as f64 / 1000.0))
1503            .collect();
1504        assert!(grid.windows(2).all(|pair| pair[0] < pair[1]));
1505        for i in 1..500 {
1506            let p = i as f64 / 1000.0;
1507            assert!(close(normal_quantile(p), -normal_quantile(1.0 - p), 1e-12));
1508        }
1509    }
1510
1511    #[test]
1512    fn cdfs_match_reference_values() {
1513        assert!(close(normal_cdf(1.959_963_984_540_054), 0.975, 1e-14));
1514        assert!(close(normal_cdf(-8.0), 6.220_960_574_271_784e-16, 1e-9));
1515        // Student's t with 5 df at 2.015048373: 0.95.
1516        assert!(close(t_cdf(2.015_048_372_669_157, 5.0), 0.95, 1e-10));
1517        // Chi-squared, 3 df, at 7.814727903: 0.95.
1518        let chi = Fitted::ChiSquared { df: 3.0 };
1519        assert!(close(chi.cdf(7.814_727_903_251_178), 0.95, 1e-10));
1520        // Beta(2, 5) at 0.5: 57/64.
1521        let beta = Fitted::Beta {
1522            alpha: 2.0,
1523            beta: 5.0,
1524        };
1525        assert!(close(beta.cdf(0.5), 57.0 / 64.0, 1e-12));
1526        // Poisson(3): P(X <= 2) = 8.5 e^-3.
1527        let poisson = Fitted::Poisson { rate: 3.0 };
1528        assert!(close(poisson.cdf(2.0), 8.5 * (-3.0f64).exp(), 1e-12));
1529        // Binomial(10, 0.3): P(X <= 3) = 0.6496107184.
1530        let binomial = Fitted::Binomial { trials: 10, p: 0.3 };
1531        assert!(close(binomial.cdf(3.0), 0.649_610_718_4, 1e-9));
1532        assert!(close(ln_gamma(0.5), 0.572_364_942_924_700_1, 1e-14));
1533        assert!(close(ln_gamma(10.0), 12.801_827_480_081_469, 1e-14));
1534    }
1535
1536    /// Every family's quantile inverts its CDF and rises with the probability, and
1537    /// the log-normal's stay positive.
1538    #[test]
1539    fn quantiles_invert_their_cdfs() {
1540        let fits = [
1541            Fitted::Normal { mean: 3.0, sd: 2.0 },
1542            Fitted::LogNormal {
1543                mu: 0.5,
1544                sigma: 0.8,
1545            },
1546            Fitted::Uniform {
1547                low: -1.0,
1548                high: 4.0,
1549            },
1550            Fitted::Exponential { rate: 2.0 },
1551            Fitted::Gamma {
1552                shape: 2.5,
1553                scale: 0.7,
1554            },
1555            Fitted::ChiSquared { df: 5.0 },
1556            Fitted::Beta {
1557                alpha: 2.0,
1558                beta: 5.0,
1559            },
1560            Fitted::StudentsT {
1561                df: 5.0,
1562                location: 1.0,
1563                scale: 2.0,
1564            },
1565            Fitted::Weibull {
1566                shape: 1.7,
1567                scale: 3.0,
1568            },
1569            Fitted::PowerLaw {
1570                xmin: 1.0,
1571                alpha: 2.5,
1572            },
1573        ];
1574        for fit in fits {
1575            let qs: Vec<f64> = (1..100).map(|i| fit.quantile(i as f64 / 100.0)).collect();
1576            assert!(qs.windows(2).all(|pair| pair[0] < pair[1]), "{fit:?}");
1577            for (i, q) in qs.iter().enumerate() {
1578                let p = (i + 1) as f64 / 100.0;
1579                assert!(
1580                    close(fit.cdf(*q), p, 1e-8),
1581                    "{fit:?} at {p}: {}",
1582                    fit.cdf(*q)
1583                );
1584            }
1585        }
1586        let lognormal = Fitted::LogNormal {
1587            mu: 0.0,
1588            sigma: 1.0,
1589        };
1590        assert!(lognormal.quantile(1e-9) > 0.0);
1591        let discrete = [
1592            Fitted::Poisson { rate: 4.0 },
1593            Fitted::Binomial { trials: 20, p: 0.4 },
1594            Fitted::Geometric { p: 0.3, start: 1 },
1595            Fitted::Bernoulli { p: 0.3 },
1596            // Fitted to counts in the millions, where walking to the answer is a
1597            // million steps.
1598            Fitted::Poisson { rate: 1e6 },
1599            Fitted::Binomial {
1600                trials: 2_000_000,
1601                p: 0.4,
1602            },
1603            Fitted::Geometric { p: 1e-6, start: 0 },
1604        ];
1605        for fit in discrete {
1606            for i in 1..100 {
1607                let p = i as f64 / 100.0;
1608                let k = fit.quantile(p);
1609                assert!(fit.cdf(k) >= p && fit.cdf(k - 1.0) < p, "{fit:?} at {p}");
1610            }
1611        }
1612    }
1613
1614    /// Both sides of each step count: values bunched high against a uniform are far
1615    /// from it below the first one, which a check at the top of each step misses.
1616    #[test]
1617    fn the_ks_statistic_is_two_sided() {
1618        let uniform = Fitted::Uniform {
1619            low: 0.0,
1620            high: 1.0,
1621        };
1622        assert!(close(ks_statistic(&[0.5, 0.6, 0.9], &uniform), 0.5, 1e-12));
1623        assert!(close(ks_statistic(&[0.1, 0.4, 0.7], &uniform), 0.3, 1e-12));
1624        // Ties are one step: three 1s against a fair coin are 0.5 away at 1.
1625        let coin = Fitted::Bernoulli { p: 0.5 };
1626        assert!(close(ks_statistic(&[1.0, 1.0, 1.0], &coin), 0.5, 1e-12));
1627    }
1628
1629    fn draw(fit: Fitted, n: usize, seed: u64) -> Vec<f64> {
1630        let mut rng = Rng::new(seed);
1631        (0..n).map(|_| fit.sample(&mut rng)).collect()
1632    }
1633
1634    /// The Poisson and binomial samplers for large counts draw the distribution, not
1635    /// an approximation: whole numbers whose CDF matches the family's.
1636    #[test]
1637    fn large_count_draws_follow_their_family() {
1638        let fits = [
1639            Fitted::Poisson { rate: 10.0 },
1640            Fitted::Poisson { rate: 37.5 },
1641            Fitted::Poisson { rate: 1e6 },
1642            Fitted::Binomial { trials: 40, p: 0.7 },
1643            Fitted::Binomial {
1644                trials: 2_000_000,
1645                p: 0.4,
1646            },
1647        ];
1648        for fit in fits {
1649            let mut values = draw(fit, 5_000, 9);
1650            assert!(values.iter().all(|v| *v >= 0.0 && *v == v.floor()));
1651            values.sort_by(f64::total_cmp);
1652            // 1.63 / sqrt(n) is the 1% critical value of the KS statistic.
1653            let d = ks_statistic(&values, &fit);
1654            assert!(d < 1.63 / (5_000f64).sqrt(), "{fit:?}: D = {d}");
1655        }
1656    }
1657
1658    /// The bootstrap p-value is calibrated: values drawn from a family are rejected at
1659    /// 5% about 5% of the time, and values from another are rejected.
1660    #[test]
1661    fn the_bootstrap_p_value_is_calibrated() {
1662        let normal = Fitted::Normal {
1663            mean: 10.0,
1664            sd: 3.0,
1665        };
1666        let rejected = (0..40)
1667            .filter(|seed| {
1668                let values = draw(normal, 120, 1_000 + seed);
1669                test_family(DistributionType::Normal, &values, *seed)
1670                    .p_value()
1671                    .unwrap()
1672                    < 0.05
1673            })
1674            .count();
1675        assert!(
1676            rejected <= 6,
1677            "{rejected} of 40 normal samples rejected as normal"
1678        );
1679
1680        let skewed = draw(
1681            Fitted::LogNormal {
1682                mu: 0.0,
1683                sigma: 1.0,
1684            },
1685            400,
1686            7,
1687        );
1688        let outcome = test_family(DistributionType::Normal, &skewed, 7);
1689        let test = outcome.test().unwrap();
1690        assert!(test.at_bound(), "{test:?}");
1691        assert!(test.p_value < 0.01);
1692    }
1693
1694    /// A family that does not apply says why and has no p-value to rank.
1695    #[test]
1696    fn a_family_that_does_not_apply_has_no_p_value() {
1697        let values = draw(Fitted::Normal { mean: 0.0, sd: 1.0 }, 50, 3);
1698        let outcome = test_family(DistributionType::LogNormal, &values, 3);
1699        assert_eq!(outcome, FitOutcome::NotApplicable("needs positive values"));
1700        let results = vec![
1701            (DistributionType::LogNormal, outcome),
1702            (
1703                DistributionType::Normal,
1704                test_family(DistributionType::Normal, &values, 3),
1705            ),
1706        ];
1707        assert_eq!(
1708            listing_order(&results)[..2],
1709            [DistributionType::Normal, DistributionType::LogNormal]
1710        );
1711    }
1712
1713    /// Selection is among the families not rejected, by AIC, the simpler of two
1714    /// nested ones when both hold: exponential values are a Gamma of shape 1 too.
1715    #[test]
1716    fn selection_prefers_the_simpler_family_that_holds() {
1717        let values = draw(Fitted::Exponential { rate: 2.0 }, 2_000, 11);
1718        let results = test_all(&values, 11);
1719        assert_eq!(select(&results, false), DistributionType::Exponential);
1720
1721        let counts = draw(Fitted::Poisson { rate: 5.0 }, 2_000, 12);
1722        let results = test_all(&counts, 12);
1723        assert_eq!(select(&results, true), DistributionType::Poisson);
1724    }
1725}