Skip to main content

probl_engine/
analytic.rs

1//! Identity-preserving continuous outcomes during enumeration.
2//!
3//! An outcome is a function of one axis: a latent variable, or a linear
4//! form of independent normal latents, as `x + y` is. It's affine on each
5//! of a few pieces of the axis's CDF coordinates, which keeps `abs`, `min`,
6//! `max` and `clamp` exact. Evidence is a set of intervals in a latent's
7//! coordinates, and an event a union of boxes of independent axes' ones.
8//! Worlds own their restrictions, and the latent's distribution there: a
9//! posterior after an exact update. An observation of several normal
10//! latents together makes each the linear form of fresh independent ones,
11//! which carries their correlation. A value carries the coordinates of the
12//! distribution it was made under, and moves to the world's when it's read.
13//! Values and closures remain immutable, and calls return restrictions
14//! alongside their results.
15
16use crate::continuous::Family;
17use crate::dist::Budget;
18use crate::error::{Fault, OpError, OpResult};
19use crate::likelihood::Likelihood;
20use crate::value::{Closure, Value, family_key, float_key};
21use probl_syntax::ast::BinOp;
22use std::collections::BTreeMap;
23use std::sync::Arc;
24
25pub type Constraints = Arc<BTreeMap<u64, Latent>>;
26
27/// What a world knows about a latent: its distribution, and where it can
28/// be, in that distribution's CDF coordinates. After an observation of
29/// several normal latents together, also the linear form of fresh ones that
30/// it is now: its distribution is then that form's, and it's unrestricted.
31#[derive(Clone, Debug, PartialEq)]
32pub struct Latent {
33    pub family: Family,
34    pub domain: Domain,
35    pub form: Option<Arc<Linear>>,
36    /// With `@integrate`, what observations that aren't conjugate say about
37    /// it: a weight on `family` (see `crate::likelihood`).
38    pub weight: Option<Arc<Likelihood>>,
39}
40impl Eq for Latent {}
41impl std::hash::Hash for Latent {
42    fn hash<H: std::hash::Hasher>(&self, h: &mut H) {
43        family_key(&self.family).hash(h);
44        self.domain.key().hash(h);
45        self.form.as_ref().map(|f| f.key()).hash(h);
46        self.weight.as_ref().map(|w| w.key()).hash(h);
47    }
48}
49
50/// One of the independent normal latents of a linear form: its id, its
51/// coefficient, and its distribution where the form was made.
52#[derive(Clone, Copy, Debug, PartialEq)]
53pub struct Term {
54    pub id: u64,
55    pub coef: f64,
56    pub family: Family,
57}
58
59/// `constant + Σ coef × latent` over independent normal latents, by
60/// increasing id, none with a zero coefficient.
61#[derive(Clone, Debug, PartialEq)]
62pub struct Linear {
63    pub constant: f64,
64    pub terms: Vec<Term>,
65}
66type FormKey = (u64, Vec<(u64, u64, (&'static str, Vec<u64>))>);
67impl Linear {
68    fn key(&self) -> FormKey {
69        (
70            float_key(self.constant),
71            self.terms
72                .iter()
73                .map(|t| (t.id, float_key(t.coef), family_key(&t.family)))
74                .collect(),
75        )
76    }
77    /// The sum of a constant and terms in any order: those of one latent
78    /// add up, and zero coefficients go.
79    fn from_terms(constant: f64, mut terms: Vec<Term>) -> OpResult<Linear> {
80        terms.sort_by_key(|t| t.id);
81        let mut out: Vec<Term> = Vec::with_capacity(terms.len());
82        for t in terms {
83            match out.last_mut() {
84                Some(last) if last.id == t.id => {
85                    if last.family != t.family {
86                        return Err(OpError::internal(
87                            "internal error: a latent with two distributions in one sum",
88                        ));
89                    }
90                    last.coef += t.coef;
91                }
92                _ => out.push(t),
93            }
94        }
95        out.retain(|t| t.coef != 0.0);
96        Ok(Linear { constant, terms: out })
97    }
98    /// `a` of it: `a.scale × self + a.offset`.
99    fn affine(&self, a: Affine) -> Linear {
100        if a.scale == 0.0 {
101            return Linear {
102                constant: a.offset,
103                terms: Vec::new(),
104            };
105        }
106        Linear {
107            constant: a.scale * self.constant + a.offset,
108            terms: self
109                .terms
110                .iter()
111                .map(|t| Term {
112                    coef: a.scale * t.coef,
113                    ..*t
114                })
115                .collect(),
116        }
117    }
118    /// `self + k × other`.
119    fn plus(&self, other: &Linear, k: f64) -> OpResult<Linear> {
120        let terms = self
121            .terms
122            .iter()
123            .copied()
124            .chain(other.terms.iter().map(|t| Term { coef: k * t.coef, ..*t }))
125            .collect();
126        Linear::from_terms(self.constant + k * other.constant, terms)
127    }
128    /// Its mean and variance.
129    fn moments(&self) -> (f64, f64) {
130        let mean = self.terms.iter().map(|t| t.coef * t.family.mean()).sum::<f64>();
131        let variance = self.terms.iter().map(|t| t.coef * t.coef * t.family.variance()).sum();
132        (self.constant + mean, variance)
133    }
134    /// Its distribution, a normal one: unsupported where it isn't finite.
135    fn normal(&self) -> OpResult<Family> {
136        let (mean, variance) = self.moments();
137        Family::normal(mean, libm::sqrt(variance))
138            .map_err(|_| unsupported("a sum of normal draws this large or this small"))
139    }
140    /// The axis of a form of several latents, and the form as an affine
141    /// function of it: the axis is the form without its constant, scaled so
142    /// that its first coefficient is 1 or −1.
143    fn axis(&self) -> (Arc<Linear>, Affine) {
144        let k = self.terms[0].coef.abs();
145        let axis = Linear {
146            constant: 0.0,
147            terms: self.terms.iter().map(|t| Term { coef: t.coef / k, ..*t }).collect(),
148        };
149        (
150            Arc::new(axis),
151            Affine {
152                scale: k,
153                offset: self.constant,
154            },
155        )
156    }
157    /// The outcome it is: a number without latents, a function of the
158    /// latent with one, and otherwise a function of its axis.
159    pub fn value(self) -> OpResult<Value> {
160        match self.terms[..] {
161            [] => Ok(Value::Float(self.constant)),
162            [t] => Analytic {
163                id: t.id,
164                family: t.family,
165                domain: Domain::full(),
166                pieces: vec![(
167                    1.0,
168                    Fun::affine(Affine {
169                        scale: t.coef,
170                        offset: self.constant,
171                    }),
172                )],
173                int: false,
174                form: None,
175                weight: None,
176            }
177            .value(),
178            _ => {
179                let (axis, map) = self.axis();
180                Analytic {
181                    id: 0,
182                    family: axis.normal()?,
183                    domain: Domain::full(),
184                    pieces: vec![(1.0, Fun::affine(map))],
185                    int: false,
186                    form: Some(axis),
187                    weight: None,
188                }
189                .value()
190            }
191        }
192    }
193}
194
195/// The CDF coordinate under `to` of the value at `c` under `from`.
196fn moved(c: f64, from: &Family, to: &Family) -> f64 {
197    to.cdf(from.quantile(c)).clamp(0.0, 1.0)
198}
199
200/// Sorted, disjoint CDF intervals. Endpoints have zero probability.
201#[derive(Clone, Debug, PartialEq)]
202pub struct Domain(pub Vec<(f64, f64)>);
203impl Domain {
204    pub fn full() -> Self {
205        Self(vec![(0.0, 1.0)])
206    }
207    pub fn mass(&self) -> f64 {
208        self.0.iter().map(|(a, b)| b - a).sum()
209    }
210    pub fn is_full(&self) -> bool {
211        self.0[..] == [(0.0, 1.0)]
212    }
213    pub fn union(&self, other: &Self) -> Self {
214        self.complement().intersect(&other.complement()).complement()
215    }
216    pub fn intersect(&self, other: &Self) -> Self {
217        let (mut i, mut j) = (0, 0);
218        let mut out = Vec::new();
219        while i < self.0.len() && j < other.0.len() {
220            let (a, b) = self.0[i];
221            let (c, d) = other.0[j];
222            if a.max(c) < b.min(d) {
223                out.push((a.max(c), b.min(d)));
224            }
225            if b < d {
226                i += 1;
227            } else {
228                j += 1;
229            }
230        }
231        Self(out)
232    }
233    pub fn complement(&self) -> Self {
234        let mut out = Vec::new();
235        let mut lo = 0.0;
236        for &(a, b) in &self.0 {
237            if lo < a {
238                out.push((lo, a));
239            }
240            lo = b;
241        }
242        if lo < 1.0 {
243            out.push((lo, 1.0));
244        }
245        Self(out)
246    }
247    fn below(p: f64) -> Self {
248        if p > 0.0 { Self(vec![(0.0, p)]) } else { Self(vec![]) }
249    }
250    fn key(&self) -> Vec<(u64, u64)> {
251        self.0.iter().map(|(a, b)| (float_key(*a), float_key(*b))).collect()
252    }
253    /// The same values of a latent, in `to`'s coordinates instead of
254    /// `from`'s.
255    pub fn moved(&self, from: &Family, to: &Family) -> Self {
256        let mut out = Vec::with_capacity(self.0.len());
257        for &(a, b) in &self.0 {
258            push(&mut out, moved(a, from, to), moved(b, from, to));
259        }
260        Self(out)
261    }
262    /// The domain as a tiling of the CDF coordinates: whether each part is
263    /// in it.
264    fn mask(&self) -> Vec<(f64, bool)> {
265        let mut out = Vec::with_capacity(2 * self.0.len() + 1);
266        let mut at = 0.0;
267        for &(a, b) in &self.0 {
268            if a > at {
269                out.push((a, false));
270            }
271            out.push((b, true));
272            at = b;
273        }
274        if at < 1.0 {
275            out.push((1.0, false));
276        }
277        out
278    }
279}
280impl Eq for Domain {}
281impl std::hash::Hash for Domain {
282    fn hash<H: std::hash::Hasher>(&self, h: &mut H) {
283        self.key().hash(h);
284    }
285}
286impl Ord for Domain {
287    fn cmp(&self, b: &Self) -> std::cmp::Ordering {
288        self.key().cmp(&b.key())
289    }
290}
291impl PartialOrd for Domain {
292    fn partial_cmp(&self, b: &Self) -> Option<std::cmp::Ordering> {
293        Some(self.cmp(b))
294    }
295}
296
297/// Adds an interval to sorted ones, joining it to the last if they touch.
298fn push(out: &mut Vec<(f64, f64)>, a: f64, b: f64) {
299    if a < b {
300        match out.last_mut() {
301            Some(last) if last.1 >= a => last.1 = last.1.max(b),
302            _ => out.push((a, b)),
303        }
304    }
305}
306
307/// `scale * x + offset`, for the latent `x`.
308#[derive(Clone, Copy, Debug, PartialEq)]
309pub struct Affine {
310    pub scale: f64,
311    pub offset: f64,
312}
313impl Affine {
314    pub const IDENTITY: Affine = Affine {
315        scale: 1.0,
316        offset: 0.0,
317    };
318    fn constant(c: f64) -> Self {
319        Affine { scale: 0.0, offset: c }
320    }
321    /// Its value where the latent is `x`. A constant is that constant even
322    /// at an infinite end of the latent's range.
323    pub fn at(self, x: f64) -> f64 {
324        if self.scale == 0.0 {
325            self.offset
326        } else {
327            self.scale * x + self.offset
328        }
329    }
330    fn negated(self) -> Self {
331        Affine {
332            scale: -self.scale,
333            offset: -self.offset,
334        }
335    }
336    fn is_finite(self) -> bool {
337        self.scale.is_finite() && self.offset.is_finite()
338    }
339    /// The `x` where it's `y`, for a nonconstant one.
340    fn inverse(self, y: f64) -> f64 {
341        (y - self.offset) / self.scale
342    }
343    /// It of `inner`.
344    fn of(self, inner: Affine) -> Affine {
345        Affine {
346            scale: self.scale * inner.scale,
347            offset: self.scale * inner.offset + self.offset,
348        }
349    }
350}
351
352/// How a function of one axis becomes a function of another, when the old
353/// axis is `map` of the new one, and their distributions are `from` and
354/// `to`: the same values of the old axis, in the new axis's coordinates.
355#[derive(Clone, Copy, Debug)]
356struct Transport {
357    from: Family,
358    to: Family,
359    map: Affine,
360}
361impl Transport {
362    /// The same axis under another distribution.
363    fn moved(from: Family, to: Family) -> Transport {
364        Transport {
365            from,
366            to,
367            map: Affine::IDENTITY,
368        }
369    }
370    fn coordinate(&self, c: f64) -> f64 {
371        self.to.cdf(self.map.inverse(self.from.quantile(c))).clamp(0.0, 1.0)
372    }
373    fn increasing(&self) -> bool {
374        self.map.scale > 0.0
375    }
376    fn domain(&self, d: &Domain) -> Domain {
377        let mut out = Vec::with_capacity(d.0.len());
378        if self.increasing() {
379            for &(a, b) in &d.0 {
380                push(&mut out, self.coordinate(a), self.coordinate(b));
381            }
382        } else {
383            for &(a, b) in d.0.iter().rev() {
384                push(&mut out, self.coordinate(b), self.coordinate(a));
385            }
386        }
387        Domain(out)
388    }
389    /// A piece's function of the old axis, as one of the new.
390    fn fun(&self, f: Fun) -> Fun {
391        match f.kernel {
392            Kernel::Id => Fun::affine(f.outer.of(self.map)),
393            _ => Fun {
394                inner: f.inner.of(self.map),
395                ..f
396            },
397        }
398    }
399    fn pieces(&self, pieces: &[(f64, Fun)]) -> Pieces {
400        let n = pieces.len();
401        let mut out: Pieces = Vec::with_capacity(n);
402        let mut add = |end: f64, f: Fun| {
403            // A piece that rounds to nothing has no probability either way.
404            if end > out.last().map_or(0.0, |p| p.0) {
405                out.push((end, self.fun(f)));
406            }
407        };
408        if self.increasing() {
409            for (i, &(end, f)) in pieces.iter().enumerate() {
410                add(if i + 1 == n { 1.0 } else { self.coordinate(end) }, f);
411            }
412        } else {
413            // In reverse: a piece ends where it started.
414            for i in (0..n).rev() {
415                add(if i == 0 { 1.0 } else { self.coordinate(pieces[i - 1].0) }, pieces[i].1);
416            }
417        }
418        match out.last_mut() {
419            Some(p) => p.0 = 1.0,
420            None => out.push((1.0, self.fun(pieces[n - 1].1))),
421        }
422        out
423    }
424}
425
426/// A monotone function of a number, between two affine ones in a piece's
427/// function.
428#[derive(Clone, Copy, Debug, PartialEq, Eq)]
429pub enum Kernel {
430    Id,
431    Square,
432    Sqrt,
433    Exp,
434    Ln,
435    Recip,
436}
437impl Kernel {
438    /// Its value at `z`, with `z` clamped into where it's defined: on a part
439    /// of the draw without probability, `sqrt` may meet a negative number.
440    fn apply(self, z: f64) -> f64 {
441        match self {
442            Kernel::Id => z,
443            Kernel::Square => z * z,
444            Kernel::Sqrt => z.max(0.0).sqrt(),
445            Kernel::Exp => crate::math::exp(z),
446            Kernel::Ln => libm::log(z.max(0.0)),
447            Kernel::Recip => 1.0 / z,
448        }
449    }
450    /// The `z` on the side of 0 that `side` is on where it's `y`, for a `y`
451    /// it takes there.
452    fn inverse(self, y: f64, side: f64) -> f64 {
453        match self {
454            Kernel::Id => y,
455            Kernel::Square if side < 0.0 => -y.sqrt(),
456            Kernel::Square => y.sqrt(),
457            Kernel::Sqrt => y * y,
458            Kernel::Exp => libm::log(y),
459            Kernel::Ln => crate::math::exp(y),
460            Kernel::Recip => 1.0 / y,
461        }
462    }
463    /// Whether it increases on the side of 0 that `side` is on.
464    fn increasing(self, side: f64) -> bool {
465        match self {
466            Kernel::Square => side >= 0.0,
467            Kernel::Recip => false,
468            _ => true,
469        }
470    }
471    /// Whether it turns, or has a pole, at 0: a piece mustn't span it.
472    fn turns(self) -> bool {
473        matches!(self, Kernel::Square | Kernel::Recip)
474    }
475}
476
477/// The function an outcome is of the draw on a piece of it: `outer`, of
478/// `kernel`, of `inner`, of the draw's value. Monotone on its piece. An
479/// affine one has the identity kernel and inner function, so that `outer`
480/// is all of it.
481#[derive(Clone, Copy, Debug, PartialEq)]
482pub struct Fun {
483    pub outer: Affine,
484    pub kernel: Kernel,
485    pub inner: Affine,
486}
487impl Fun {
488    pub const IDENTITY: Fun = Fun {
489        outer: Affine::IDENTITY,
490        kernel: Kernel::Id,
491        inner: Affine::IDENTITY,
492    };
493    pub fn affine(outer: Affine) -> Fun {
494        Fun { outer, ..Fun::IDENTITY }
495    }
496    fn constant(c: f64) -> Fun {
497        Fun::affine(Affine::constant(c))
498    }
499    /// `kernel` of an affine function of the draw.
500    fn of(kernel: Kernel, inner: Affine) -> Fun {
501        Fun {
502            outer: Affine::IDENTITY,
503            kernel,
504            inner,
505        }
506    }
507    /// The affine function it is, if it is one.
508    pub fn as_affine(self) -> Option<Affine> {
509        (self.kernel == Kernel::Id).then_some(self.outer)
510    }
511    /// The number it is, if it's constant.
512    pub fn as_constant(self) -> Option<f64> {
513        (self.outer.scale == 0.0).then_some(self.outer.offset)
514    }
515    /// Its value where the draw is `x`.
516    pub fn at(self, x: f64) -> f64 {
517        match self.kernel {
518            Kernel::Id => self.outer.at(x),
519            k => self.outer.at(k.apply(self.inner.at(x))),
520        }
521    }
522    /// `scale * self + offset`.
523    fn then(self, scale: f64, offset: f64) -> Fun {
524        Fun {
525            outer: Affine {
526                scale: self.outer.scale * scale,
527                offset: self.outer.offset * scale + offset,
528            },
529            ..self
530        }
531    }
532    fn negated(self) -> Fun {
533        Fun {
534            outer: self.outer.negated(),
535            ..self
536        }
537    }
538    fn is_finite(self) -> bool {
539        self.outer.is_finite() && self.inner.is_finite()
540    }
541    pub(crate) fn key(self) -> (u8, u64, u64, u64, u64) {
542        (
543            self.kernel as u8,
544            float_key(self.outer.scale),
545            float_key(self.outer.offset),
546            float_key(self.inner.scale),
547            float_key(self.inner.offset),
548        )
549    }
550    /// The sign of the kernel's input on a piece, which doesn't change there.
551    fn side(self, family: &Family, start: f64, end: f64) -> f64 {
552        self.inner.at(family.quantile((start + end) / 2.0))
553    }
554    /// Its values at the ends of a piece. At a pole, the kernel's input is
555    /// zero of the piece's sign: `1 / x` falls to −∞ from below 0.
556    pub fn ends(self, family: &Family, start: f64, end: f64) -> (f64, f64) {
557        let side = match self.kernel {
558            Kernel::Recip => self.side(family, start, end),
559            _ => 1.0,
560        };
561        let at = |u: f64| {
562            let x = family.quantile(u);
563            match self.kernel {
564                Kernel::Id => self.outer.at(x),
565                k => {
566                    let z = self.inner.at(x);
567                    let z = if z == 0.0 { 0.0f64.copysign(side) } else { z };
568                    self.outer.at(k.apply(z))
569                }
570            }
571        };
572        (at(start), at(end))
573    }
574    /// Whether it increases with the draw on a piece.
575    fn increasing(self, family: &Family, start: f64, end: f64) -> bool {
576        match self.kernel {
577            Kernel::Id => self.outer.scale > 0.0,
578            k => {
579                let side = self.side(family, start, end);
580                (self.outer.scale > 0.0) == (k.increasing(side) == (self.inner.scale > 0.0))
581            }
582        }
583    }
584    /// The draw's coordinate on a piece where it's `t`, for a value it takes
585    /// there.
586    fn coordinate(self, family: &Family, start: f64, end: f64, t: f64) -> f64 {
587        match self.kernel {
588            Kernel::Id => family.cdf(self.outer.inverse(t)),
589            k => {
590                let z = k.inverse(self.outer.inverse(t), self.side(family, start, end));
591                family.cdf(self.inner.inverse(z))
592            }
593        }
594    }
595    /// The part of a piece of the draw where it's below `t`, or at most `t`
596    /// unless `strict`.
597    fn below_on(self, family: &Family, start: f64, end: f64, t: f64, strict: bool) -> OpResult<(f64, f64)> {
598        if let Some(c) = self.as_constant() {
599            return Ok(if c < t || (!strict && c == t) {
600                (start, end)
601            } else {
602                (end, end)
603            });
604        }
605        if let Some(f) = self.as_affine() {
606            let p = family.cdf(f.inverse(t));
607            if !p.is_finite() {
608                return Err(unsupported("this numerically unstable comparison"));
609            }
610            return Ok(if f.scale > 0.0 {
611                (start, end.min(p))
612            } else {
613                (start.max(p), end)
614            });
615        }
616        let up = self.increasing(family, start, end);
617        let (a, b) = self.ends(family, start, end);
618        let (low, high) = if up { (a, b) } else { (b, a) };
619        if t <= low {
620            return Ok((end, end));
621        }
622        if t >= high {
623            return Ok((start, end));
624        }
625        let u = self.coordinate(family, start, end, t);
626        if !u.is_finite() {
627            return Err(unsupported("this numerically unstable comparison"));
628        }
629        let u = u.clamp(start, end);
630        Ok(if up { (start, u) } else { (u, end) })
631    }
632    /// The coefficients of the polynomial in the draw it is, if it's one of
633    /// degree 2 or less.
634    fn polynomial(self) -> Option<[f64; 3]> {
635        let Affine { scale: s, offset: o } = self.outer;
636        match self.kernel {
637            Kernel::Id => Some([o, s, 0.0]),
638            Kernel::Square => {
639                let Affine { scale: a, offset: b } = self.inner;
640                Some([s * b * b + o, 2.0 * s * a * b, s * a * a])
641            }
642            _ => None,
643        }
644    }
645    /// The function a polynomial of degree 2 or less is: a square about its
646    /// vertex, unless it's affine.
647    fn from_polynomial([c, b, a]: [f64; 3]) -> Fun {
648        if a == 0.0 {
649            return Fun::affine(Affine { scale: b, offset: c });
650        }
651        let vertex = -b / (2.0 * a);
652        Fun {
653            outer: Affine {
654                scale: a,
655                offset: c - a * vertex * vertex,
656            },
657            kernel: Kernel::Square,
658            inner: Affine {
659                scale: 1.0,
660                offset: -vertex,
661            },
662        }
663    }
664}
665
666/// A function of a latent, by where each of its pieces ends in the latent's
667/// CDF coordinates: the first starts at 0, each other where the one before
668/// it ends, and the last ends at 1.
669pub type Pieces = Vec<(f64, Fun)>;
670
671/// The common refinement of two tilings of the CDF coordinates.
672fn overlay<A: Copy, B: Copy>(a: &[(f64, A)], b: &[(f64, B)]) -> Vec<(f64, (A, B))> {
673    let (mut i, mut j) = (0, 0);
674    let mut out = Vec::with_capacity(a.len() + b.len() - 1);
675    loop {
676        let end = a[i].0.min(b[j].0);
677        out.push((end, (a[i].1, b[j].1)));
678        // Both tilings end at 1, so neither runs out first.
679        if end >= 1.0 {
680            return out;
681        }
682        i += usize::from(a[i].0 == end);
683        j += usize::from(b[j].0 == end);
684    }
685}
686
687/// `then` where `yes` holds and `otherwise` elsewhere: the partition that
688/// piecewise operations share.
689fn select(yes: &Domain, then: &[(f64, Fun)], otherwise: &[(f64, Fun)]) -> Pieces {
690    overlay(&overlay(then, otherwise), &yes.mask())
691        .into_iter()
692        .map(|(end, ((then, otherwise), inside))| (end, if inside { then } else { otherwise }))
693        .collect()
694}
695
696/// Where a function of a latent is below `t` (or at most `t`), in the
697/// latent's CDF coordinates.
698fn below(family: &Family, pieces: &[(f64, Fun)], t: f64, strict: bool) -> OpResult<Domain> {
699    let mut out = Vec::with_capacity(pieces.len());
700    let mut start = 0.0;
701    for &(end, f) in pieces {
702        let (lo, hi) = f.below_on(family, start, end, t, strict)?;
703        push(&mut out, lo, hi);
704        start = end;
705    }
706    Ok(Domain(out))
707}
708
709/// Splits each piece whose kernel turns, or has a pole, where it does, so
710/// that every piece is monotone.
711fn split_turns(family: &Family, pieces: Pieces) -> Pieces {
712    if pieces.iter().all(|(_, f)| !f.kernel.turns()) {
713        return pieces;
714    }
715    let mut out = Vec::with_capacity(pieces.len() + 1);
716    let mut start = 0.0;
717    for (end, f) in pieces {
718        if f.kernel.turns() && f.as_constant().is_none() {
719            let u = family.cdf(f.inner.inverse(0.0));
720            if start < u && u < end {
721                out.push((u, f));
722            }
723        }
724        out.push((end, f));
725        start = end;
726    }
727    out
728}
729
730#[derive(Clone, Debug)]
731pub struct Analytic {
732    pub id: u64,
733    pub family: Family,
734    pub domain: Domain,
735    pub pieces: Pieces,
736    /// Whether it's an int, as `floor(x)` is: then every piece is a whole
737    /// number.
738    pub int: bool,
739    /// For a function of several normal latents, its axis: a linear form
740    /// of them, whose distribution is `family`. `id` is then 0.
741    pub form: Option<Arc<Linear>>,
742    /// The latent's likelihood, with `@integrate`: a weight on `family`.
743    pub weight: Option<Arc<Likelihood>>,
744}
745impl Analytic {
746    pub fn new(id: u64, family: Family) -> Self {
747        Self {
748            id,
749            family,
750            domain: Domain::full(),
751            pieces: vec![(1.0, Fun::IDENTITY)],
752            int: false,
753            form: None,
754            weight: None,
755        }
756    }
757    /// The axis itself: the draw, or the form.
758    pub(crate) fn latent(&self) -> Self {
759        self.with(vec![(1.0, Fun::IDENTITY)])
760    }
761    /// Whether it's a function of the same axis as `other`.
762    pub fn same_axis(&self, other: &Analytic) -> bool {
763        match (&self.form, &other.form) {
764            (None, None) => self.id == other.id,
765            (Some(a), Some(b)) => Arc::ptr_eq(a, b) || a == b,
766            _ => false,
767        }
768    }
769    /// Whether it depends on one of the latents that `other` does: then the
770    /// two aren't independent.
771    fn shares_latents(&self, other: &Analytic) -> bool {
772        let mut mine = std::collections::BTreeSet::new();
773        self.collect_ids(&mut mine);
774        let mut theirs = std::collections::BTreeSet::new();
775        other.collect_ids(&mut theirs);
776        !mine.is_disjoint(&theirs)
777    }
778    fn collect_ids(&self, ids: &mut std::collections::BTreeSet<u64>) {
779        match &self.form {
780            Some(f) => ids.extend(f.terms.iter().map(|t| t.id)),
781            None => {
782                ids.insert(self.id);
783            }
784        }
785    }
786    /// The linear form of independent normal latents it is, if it's one:
787    /// affine in a normal latent or a form, and unrestricted.
788    pub fn linear(&self) -> Option<Linear> {
789        let a = self.affine()?;
790        if !self.domain.is_full() || self.weight.is_some() {
791            return None;
792        }
793        match &self.form {
794            Some(f) => Some(f.affine(a)),
795            None if matches!(self.family, Family::Normal { .. }) => Some(
796                Linear {
797                    constant: 0.0,
798                    terms: vec![Term {
799                        id: self.id,
800                        coef: 1.0,
801                        family: self.family,
802                    }],
803                }
804                .affine(a),
805            ),
806            None => None,
807        }
808    }
809    /// The same outcome of other latents: each as `map` renames it.
810    pub fn renamed(&self, map: &BTreeMap<u64, u64>) -> Self {
811        let rename = |id: u64| *map.get(&id).unwrap_or(&id);
812        let mut x = self.clone();
813        match &self.form {
814            Some(f) => {
815                x.form = Some(Arc::new(Linear {
816                    constant: f.constant,
817                    terms: f.terms.iter().map(|t| Term { id: rename(t.id), ..*t }).collect(),
818                }))
819            }
820            None => x.id = rename(self.id),
821        }
822        x
823    }
824    /// Its value where its axis is `value`, at the CDF coordinate `c`.
825    pub fn at(&self, c: f64, value: f64) -> Value {
826        let f = self
827            .pieces
828            .iter()
829            .find(|(end, _)| c <= *end)
830            .unwrap_or(&self.pieces[self.pieces.len() - 1])
831            .1;
832        number(f.at(value), self.int)
833    }
834    /// The same function of the same values, as a function of another
835    /// axis, `id` or `form`, that `t` moves it to.
836    fn rebased(&self, id: u64, form: Option<Arc<Linear>>, t: &Transport) -> Self {
837        Self {
838            id,
839            family: t.to,
840            domain: t.domain(&self.domain),
841            pieces: t.pieces(&self.pieces),
842            int: self.int,
843            form,
844            weight: None,
845        }
846    }
847    /// The same function of the same values of the draw, in the coordinates
848    /// of `to`, a distribution of the draw with the same support or less.
849    pub fn moved(&self, to: &Family) -> Self {
850        if self.family == *to {
851            return self.clone();
852        }
853        let last = self.pieces.len() - 1;
854        let mut pieces: Pieces = Vec::with_capacity(self.pieces.len());
855        for (i, &(end, f)) in self.pieces.iter().enumerate() {
856            let end = if i == last { 1.0 } else { moved(end, &self.family, to) };
857            // A piece that rounds to nothing has no probability either way.
858            if end > pieces.last().map_or(0.0, |p| p.0) {
859                pieces.push((end, f));
860            }
861        }
862        match pieces.last_mut() {
863            Some(p) => p.0 = 1.0,
864            None => pieces.push((1.0, self.pieces[last].1)),
865        }
866        Self {
867            id: self.id,
868            family: *to,
869            domain: self.domain.moved(&self.family, to),
870            pieces,
871            int: self.int,
872            form: self.form.clone(),
873            weight: self.weight.clone(),
874        }
875    }
876    /// Another function of the same draw, a float.
877    fn with(&self, pieces: Pieces) -> Self {
878        Self {
879            id: self.id,
880            family: self.family,
881            domain: self.domain.clone(),
882            pieces,
883            int: false,
884            form: self.form.clone(),
885            weight: self.weight.clone(),
886        }
887    }
888    /// Another function of the same draw, an int if `int`.
889    fn with_type(&self, pieces: Pieces, int: bool) -> Self {
890        Self {
891            int,
892            ..self.with(pieces)
893        }
894    }
895    /// The function, when it's affine everywhere.
896    pub fn affine(&self) -> Option<Affine> {
897        match self.pieces[..] {
898            [(_, f)] => f.as_affine(),
899            _ => None,
900        }
901    }
902    /// Whether it's affine on each piece: then it has a density where it
903    /// isn't constant.
904    pub fn piecewise_affine(&self) -> bool {
905        self.pieces.iter().all(|(_, f)| f.as_affine().is_some())
906    }
907    pub fn key(&self) -> impl Ord + std::hash::Hash + use<> {
908        (
909            self.id,
910            family_key(&self.family),
911            self.domain.key(),
912            self.pieces
913                .iter()
914                .map(|(end, f)| (float_key(*end), f.key()))
915                .collect::<Vec<_>>(),
916            self.int,
917            self.form.as_ref().map(|f| f.key()),
918            self.weight.as_ref().map(|w| w.key()),
919        )
920    }
921    pub fn value(mut self) -> OpResult<Value> {
922        if !self.pieces.iter().all(|(_, f)| f.is_finite()) {
923            return Err(unsupported("non-finite affine coefficients"));
924        }
925        self.simplify();
926        match self.pieces[..] {
927            [(_, f)] if f.as_constant().is_some() => Ok(number(f.outer.offset, self.int)),
928            _ => Ok(Value::Analytic(Arc::new(self))),
929        }
930    }
931    /// Whether each of its values has positive probability: it's a constant
932    /// on each part of the domain, like `floor(x)`.
933    pub fn is_discrete(&self) -> bool {
934        self.pieces.len() > 1 && self.segments().iter().all(|(_, _, f)| f.as_constant().is_some())
935    }
936    /// Each value of a discrete outcome, with where the draw is when it's
937    /// that value, and that part's share of the domain's probability.
938    pub fn atoms(&self) -> Vec<(Value, Domain, f64)> {
939        let mass = self.measure(&self.domain);
940        let mut atoms: Vec<(f64, Vec<(f64, f64)>)> = Vec::new();
941        for (lo, hi, f) in self.segments() {
942            let c = f.outer.offset;
943            match atoms.iter_mut().find(|(d, _)| float_key(*d) == float_key(c)) {
944                Some((_, parts)) => parts.push((lo, hi)),
945                None => atoms.push((c, vec![(lo, hi)])),
946            }
947        }
948        atoms.sort_by(|a, b| a.0.total_cmp(&b.0));
949        atoms
950            .into_iter()
951            .map(|(c, parts)| {
952                let mut region = Vec::with_capacity(parts.len());
953                for (lo, hi) in parts {
954                    push(&mut region, lo, hi);
955                }
956                let region = Domain(region);
957                let share = self.measure(&region) / mass;
958                (number(c, self.int), region, share)
959            })
960            .collect()
961    }
962    /// Gives each piece with no probability in the domain to a neighbour,
963    /// since nothing outside the domain is read, then joins neighbours that
964    /// are the same function. Only an affine function, monotone everywhere,
965    /// stretches over another piece: one whose kernel turns stays on its
966    /// side.
967    fn simplify(&mut self) {
968        if self.pieces.len() == 1 {
969            return;
970        }
971        let affine = |g: Fun| g.as_affine().is_some();
972        let d = &self.domain.0;
973        let mut out: Pieces = Vec::with_capacity(self.pieces.len());
974        let mut any_live = false;
975        let (mut start, mut j) = (0.0, 0);
976        for &(end, f) in &self.pieces {
977            while j < d.len() && d[j].1 <= start {
978                j += 1;
979            }
980            let live = j < d.len() && d[j].0 < end;
981            start = end;
982            if !live {
983                match out.last_mut() {
984                    Some(last) if affine(last.1) => last.0 = end,
985                    _ => out.push((end, f)),
986                }
987                continue;
988            }
989            if !any_live && affine(f) {
990                // A first live piece that's affine starts at 0.
991                out.clear();
992            }
993            any_live = true;
994            match out.last_mut() {
995                Some(last) if last.1 == f && !f.kernel.turns() => last.0 = end,
996                _ => out.push((end, f)),
997            }
998        }
999        if any_live {
1000            self.pieces = out;
1001        }
1002    }
1003    /// Each part of the domain, with the piece of the function there, as
1004    /// `(lo, hi, f)`.
1005    pub fn segments(&self) -> Vec<(f64, f64, Fun)> {
1006        let mut out = Vec::with_capacity(self.domain.0.len() + self.pieces.len() - 1);
1007        let (mut k, mut start) = (0, 0.0);
1008        for &(a, b) in &self.domain.0 {
1009            loop {
1010                let (end, f) = self.pieces[k];
1011                let (lo, hi) = (a.max(start), b.min(end));
1012                if lo < hi {
1013                    out.push((lo, hi, f));
1014                }
1015                if end >= b || k + 1 == self.pieces.len() {
1016                    break;
1017                }
1018                start = end;
1019                k += 1;
1020            }
1021        }
1022        out
1023    }
1024    /// Whether a single value has positive probability: a constant piece.
1025    pub fn has_atoms(&self) -> bool {
1026        self.pieces.len() > 1 && self.segments().iter().any(|(_, _, f)| f.as_constant().is_some())
1027    }
1028    /// The measure of part of its coordinates: its length, or with a
1029    /// likelihood, the likelihood's integral there.
1030    pub fn measure(&self, d: &Domain) -> f64 {
1031        match &self.weight {
1032            None => d.mass(),
1033            Some(l) => l.mass(d),
1034        }
1035    }
1036    /// The probability of `part` of its coordinates, given its domain.
1037    pub fn share(&self, part: &Domain) -> f64 {
1038        (self.measure(&self.domain.intersect(part)) / self.measure(&self.domain)).clamp(0.0, 1.0)
1039    }
1040    pub fn cdf(&self, x: f64) -> f64 {
1041        let affine = self.affine().filter(|_| self.weight.is_none());
1042        let Some(f) = affine else {
1043            return match below(&self.family, &self.pieces, x, false) {
1044                Ok(at_most) => self.share(&at_most),
1045                Err(_) => f64::NAN,
1046            };
1047        };
1048        let p = self.family.cdf((x - f.offset) / f.scale);
1049        let below = self.domain.intersect(&Domain::below(p)).mass() / self.domain.mass();
1050        (if f.scale > 0.0 { below } else { 1.0 - below }).clamp(0.0, 1.0)
1051    }
1052    /// The density at `x`, of a piecewise affine outcome: the caller rules
1053    /// out atoms and other functions.
1054    pub fn pdf(&self, x: f64) -> f64 {
1055        let mass = self.measure(&self.domain);
1056        // The posterior density, in the coordinates, relative to the prior.
1057        let weight = |u: f64| self.weight.as_ref().map_or(1.0, |l| l.at(u));
1058        let Some(f) = self.affine() else {
1059            return self
1060                .segments()
1061                .into_iter()
1062                .filter_map(|(lo, hi, f)| Some((lo, hi, f.as_affine().filter(|f| f.scale != 0.0)?)))
1063                .map(|(lo, hi, f)| {
1064                    let x = (x - f.offset) / f.scale;
1065                    let u = self.family.cdf(x);
1066                    if lo <= u && u <= hi {
1067                        self.family.pdf(x) * weight(u) / (f.scale.abs() * mass)
1068                    } else {
1069                        0.0
1070                    }
1071                })
1072                .sum();
1073        };
1074        let x = (x - f.offset) / f.scale;
1075        let u = self.family.cdf(x);
1076        if self.domain.0.iter().any(|(a, b)| *a <= u && u <= *b) {
1077            self.family.pdf(x) * weight(u) / (f.scale.abs() * mass)
1078        } else {
1079            0.0
1080        }
1081    }
1082    pub fn quantile(&self, q: f64) -> f64 {
1083        if let Some(l) = &self.weight {
1084            return self.weighted_quantile(l, q);
1085        }
1086        let Some(f) = self.affine() else {
1087            return self.piecewise_quantile(q);
1088        };
1089        let mut left = q * self.domain.mass();
1090        let n = self.domain.0.len();
1091        for i in 0..n {
1092            let (a, b) = self.domain.0[if f.scale > 0.0 { i } else { n - 1 - i }];
1093            if left <= b - a || i + 1 == n {
1094                let u = if f.scale > 0.0 { a + left } else { b - left };
1095                return f.at(self.family.quantile(u.clamp(a, b)));
1096            }
1097            left -= b - a;
1098        }
1099        f64::NAN
1100    }
1101    /// The quantile of the mixture of the segments lies between the lowest
1102    /// and the highest of theirs.
1103    fn piecewise_quantile(&self, q: f64) -> f64 {
1104        let family = &self.family;
1105        let segments = self.segments();
1106        // A segment's value `q` of the way through its probability.
1107        let within = |lo: f64, hi: f64, f: Fun, q: f64| {
1108            if let Some(c) = f.as_constant() {
1109                return c;
1110            }
1111            let up = f.increasing(family, lo, hi);
1112            if q <= 0.0 || q >= 1.0 {
1113                let (a, b) = f.ends(family, lo, hi);
1114                return if (q <= 0.0) == up { a } else { b };
1115            }
1116            let u = if up { lo + q * (hi - lo) } else { hi - q * (hi - lo) };
1117            f.at(family.quantile(u.clamp(lo, hi)))
1118        };
1119        let each = |q: f64| segments.iter().map(move |&(lo, hi, f)| within(lo, hi, f, q));
1120        let lo = each(q).fold(f64::INFINITY, f64::min);
1121        let hi = each(q).fold(f64::NEG_INFINITY, f64::max);
1122        if q <= 0.0 {
1123            return lo;
1124        }
1125        if q >= 1.0 {
1126            return hi;
1127        }
1128        let (lo, hi) = (lo.max(-f64::MAX), hi.min(f64::MAX));
1129        if lo >= hi || self.cdf(lo) >= q {
1130            return lo;
1131        }
1132        let t = crate::continuous::invert(|t| self.cdf(t), q, lo, hi);
1133        // Where a single segment takes values around `t`, its own quantile
1134        // is exact: bisection only gets as close as the CDF rounds.
1135        let mut needed = q * self.domain.mass();
1136        let mut around = None;
1137        for &(lo, hi, f) in &segments {
1138            let (a, b) = f.ends(family, lo, hi);
1139            if a.max(b) <= t {
1140                needed -= hi - lo;
1141            } else if a.min(b) < t {
1142                if around.is_some() {
1143                    return t;
1144                }
1145                around = Some((lo, hi, f));
1146            }
1147        }
1148        let Some((lo, hi, f)) = around else { return t };
1149        let exact = within(lo, hi, f, (needed / (hi - lo)).clamp(0.0, 1.0));
1150        if (exact - t).abs() <= 1e-9 * t.abs().max(1.0) {
1151            exact
1152        } else {
1153            t
1154        }
1155    }
1156    /// The quantile under a likelihood: each segment's from the inverse of
1157    /// its integral, and between the lowest and the highest of theirs, the
1158    /// inverse of the CDF, which inherits its tolerance.
1159    fn weighted_quantile(&self, l: &Likelihood, q: f64) -> f64 {
1160        let family = &self.family;
1161        let segments = self.segments();
1162        let within = |lo: f64, hi: f64, f: Fun| {
1163            if let Some(c) = f.as_constant() {
1164                return c;
1165            }
1166            let up = f.increasing(family, lo, hi);
1167            if q <= 0.0 || q >= 1.0 {
1168                let (a, b) = f.ends(family, lo, hi);
1169                return if (q <= 0.0) == up { a } else { b };
1170            }
1171            let u = l.inverse(&Domain(vec![(lo, hi)]), if up { q } else { 1.0 - q });
1172            f.at(family.quantile(u.clamp(lo, hi)))
1173        };
1174        let each = || segments.iter().map(|&(lo, hi, f)| within(lo, hi, f));
1175        let lo = each().fold(f64::INFINITY, f64::min);
1176        let hi = each().fold(f64::NEG_INFINITY, f64::max);
1177        if q <= 0.0 || segments.len() == 1 {
1178            return lo;
1179        }
1180        if q >= 1.0 {
1181            return hi;
1182        }
1183        let (lo, hi) = (lo.max(-f64::MAX), hi.min(f64::MAX));
1184        if lo >= hi || self.cdf(lo) >= q {
1185            return lo;
1186        }
1187        crate::continuous::invert(|t| self.cdf(t), q, lo, hi)
1188    }
1189    /// The mean and variance under a likelihood, by integration, the
1190    /// variance about the mean: `None` if an integral doesn't converge.
1191    fn weighted_moments(&self, l: &Likelihood) -> Option<(f64, f64)> {
1192        let key = (
1193            self.pieces.iter().map(|(end, f)| (float_key(*end), f.key())).collect(),
1194            self.domain.key(),
1195        );
1196        l.moments(key, || self.integrate_moments(l))
1197    }
1198    fn integrate_moments(&self, l: &Likelihood) -> Option<(f64, f64)> {
1199        let total = l.mass(&self.domain);
1200        let segments = self.segments();
1201        let integrate = |g: &dyn Fn(f64) -> f64| -> Option<f64> {
1202            let mut sum = 0.0;
1203            for &(lo, hi, f) in &segments {
1204                sum += match f.as_constant() {
1205                    Some(c) => g(c) * l.mass(&Domain(vec![(lo, hi)])),
1206                    None => l.integral(|u| g(f.at(self.family.quantile(u))), lo, hi)?,
1207                };
1208            }
1209            Some(sum / total)
1210        };
1211        let mean = integrate(&|y| y)?;
1212        let variance = integrate(&|y| (y - mean) * (y - mean))?;
1213        (mean.is_finite() && variance.is_finite()).then_some((mean, variance))
1214    }
1215    /// The mean and variance by numerical integration, to `tolerance`, for
1216    /// `@integrate` where there's no formula: `None` if an integral doesn't
1217    /// converge.
1218    pub fn integrated_moments(&self, tolerance: f64) -> Option<(f64, f64)> {
1219        if let Some(l) = &self.weight {
1220            return self.weighted_moments(l);
1221        }
1222        let total = self.domain.mass();
1223        let segments = self.segments();
1224        let integrate = |g: &dyn Fn(f64) -> f64| -> Option<f64> {
1225            let mut sum = 0.0;
1226            for &(lo, hi, f) in &segments {
1227                sum += match f.as_constant() {
1228                    Some(c) => g(c) * (hi - lo),
1229                    None => crate::quadrature::moment(|u| g(f.at(self.family.quantile(u))), lo, hi, tolerance)?,
1230                };
1231            }
1232            Some(sum / total)
1233        };
1234        let mean = integrate(&|y| y)?;
1235        let variance = integrate(&|y| (y - mean) * (y - mean))?;
1236        (mean.is_finite() && variance.is_finite()).then_some((mean, variance))
1237    }
1238    /// The mean and variance, from formulas: `None` where a piece's function
1239    /// has none for the draw's family.
1240    pub fn moments(&self) -> Option<(f64, f64)> {
1241        if let Some(l) = &self.weight {
1242            return self.weighted_moments(l);
1243        }
1244        let Some(f) = self.affine() else {
1245            let mass = self.domain.mass();
1246            let parts: Vec<_> = self
1247                .segments()
1248                .into_iter()
1249                .map(|(lo, hi, f)| Some((segment_moments(&self.family, f, lo, hi)?, (hi - lo) / mass)))
1250                .collect::<Option<_>>()?;
1251            let mean = parts.iter().map(|((m, _), p)| m * p).sum::<f64>();
1252            let variance = parts.iter().map(|((m, v), p)| p * (v + (m - mean).powi(2))).sum();
1253            return Some((mean, variance));
1254        };
1255        let (mean, variance) = if self.domain == Domain::full() {
1256            (self.family.mean(), self.family.variance())
1257        } else {
1258            let parts: Vec<_> = self
1259                .domain
1260                .0
1261                .iter()
1262                .map(|&(a, b)| {
1263                    let (m, v) = self
1264                        .family
1265                        .interval_moments(self.family.quantile(a), self.family.quantile(b));
1266                    (m, v, (b - a) / self.domain.mass())
1267                })
1268                .collect();
1269            let mean = parts.iter().map(|(m, _, p)| m * p).sum::<f64>();
1270            let variance = parts.iter().map(|(m, v, p)| p * (v + (m - mean).powi(2))).sum();
1271            (mean, variance)
1272        };
1273        Some((f.at(mean), f.scale * f.scale * variance))
1274    }
1275
1276    pub fn sd(&self) -> Option<f64> {
1277        match self.affine() {
1278            Some(f) if self.domain == Domain::full() && self.weight.is_none() => Some(f.scale.abs() * self.family.sd()),
1279            _ => self.moments().map(|(_, v)| libm::sqrt(v)),
1280        }
1281    }
1282}
1283
1284/// The mean and variance of a piece's function of the draw, where the draw
1285/// is between the coordinates `lo` and `hi`: `None` where there's no
1286/// formula for that function and the draw's family.
1287fn segment_moments(family: &Family, f: Fun, lo: f64, hi: f64) -> Option<(f64, f64)> {
1288    if let Some(c) = f.as_constant() {
1289        return Some((c, 0.0));
1290    }
1291    let (first, second) = match f.kernel {
1292        Kernel::Id => {
1293            let (m, v) = family.interval_moments(family.quantile(lo), family.quantile(hi));
1294            return Some((f.outer.at(m), f.outer.scale * f.outer.scale * v));
1295        }
1296        Kernel::Square => (
1297            moments::power(family, f.inner, 2.0, lo, hi)?,
1298            moments::power(family, f.inner, 4.0, lo, hi)?,
1299        ),
1300        Kernel::Sqrt => (
1301            moments::power(family, f.inner, 0.5, lo, hi)?,
1302            moments::power(family, f.inner, 1.0, lo, hi)?,
1303        ),
1304        Kernel::Recip => (
1305            moments::power(family, f.inner, -1.0, lo, hi)?,
1306            moments::power(family, f.inner, -2.0, lo, hi)?,
1307        ),
1308        Kernel::Exp => (
1309            moments::exp(family, f.inner, 1.0, lo, hi)?,
1310            moments::exp(family, f.inner, 2.0, lo, hi)?,
1311        ),
1312        Kernel::Ln => {
1313            return moments::ln(family, f.inner, lo, hi).map(|(m, v)| (f.outer.at(m), f.outer.scale.powi(2) * v));
1314        }
1315    };
1316    let variance = (second - first * first).max(0.0);
1317    let (mean, variance) = (f.outer.at(first), f.outer.scale * f.outer.scale * variance);
1318    (mean.is_finite() && variance.is_finite()).then_some((mean, variance))
1319}
1320
1321/// Moments of functions of a draw restricted to an interval of its CDF
1322/// coordinates, `lo` to `hi`, from formulas. `None` where a family has none
1323/// for that function, or the moment isn't finite.
1324mod moments {
1325    use super::Affine;
1326    use crate::continuous::{Family, beta_cdf, gamma_cdf, ln_beta, std_normal_cdf, std_normal_quantile};
1327
1328    fn finite(x: f64) -> Option<f64> {
1329        x.is_finite().then_some(x)
1330    }
1331
1332    fn std_normal_pdf(z: f64) -> f64 {
1333        if z.is_infinite() {
1334            0.0
1335        } else {
1336            (-0.5 * z * z).exp() / (2.0 * std::f64::consts::PI).sqrt()
1337        }
1338    }
1339
1340    /// E[Y^k], k = 0 to `n`, for a normal Y with mean `m` and sd `t`
1341    /// between `a` and `b`, whose probability there is `mass`: by the
1342    /// recurrence E[Y^k] = m E[Y^(k−1)] + (k−1) t² E[Y^(k−2)]
1343    /// − t (b^(k−1) φ(β) − a^(k−1) φ(α)) / mass.
1344    fn normal_powers(m: f64, t: f64, a: f64, b: f64, mass: f64, n: usize) -> Vec<f64> {
1345        let (alpha, beta) = ((a - m) / t, (b - m) / t);
1346        let (fa, fb) = (std_normal_pdf(alpha), std_normal_pdf(beta));
1347        // An infinite end contributes nothing: φ falls faster than any power.
1348        let edge = |x: f64, f: f64, k: usize| if f == 0.0 { 0.0 } else { x.powi(k as i32) * f };
1349        let mut out = vec![1.0];
1350        for k in 1..=n {
1351            let previous = out[k - 1];
1352            let before = if k >= 2 { out[k - 2] } else { 0.0 };
1353            let boundary = edge(b, fb, k - 1) - edge(a, fa, k - 1);
1354            out.push(m * previous + (k as f64 - 1.0) * t * t * before - t * boundary / mass);
1355        }
1356        out
1357    }
1358
1359    /// E[(s X + o)^p] for the draw X.
1360    pub fn power(family: &Family, inner: Affine, p: f64, lo: f64, hi: f64) -> Option<f64> {
1361        let (a, b) = (family.quantile(lo), family.quantile(hi));
1362        let mass = hi - lo;
1363        let Affine { scale: s, offset: o } = inner;
1364        match *family {
1365            Family::Uniform { .. } => {
1366                // Y = s X + o is uniform between its values at the ends.
1367                let (ya, yb) = (inner.at(a), inner.at(b));
1368                let (y0, y1) = (ya.min(yb), ya.max(yb));
1369                uniform_power(y0, y1, p)
1370            }
1371            Family::Normal { mean, sd } if p.fract() == 0.0 && p >= 0.0 => {
1372                // Y = s X + o is normal; truncated where X is.
1373                let (ya, yb) = (inner.at(a), inner.at(b));
1374                let (y0, y1) = (ya.min(yb), ya.max(yb));
1375                let powers = normal_powers(s * mean + o, s.abs() * sd, y0, y1, mass, p as usize);
1376                finite(powers[p as usize])
1377            }
1378            _ if o == 0.0 && s > 0.0 => finite(s.powf(p) * positive_power(family, p, lo, hi)?),
1379            // A whole power of s X + o, from those of X.
1380            _ if p.fract() == 0.0 && p >= 0.0 => {
1381                let n = p as i32;
1382                let mut total = 0.0;
1383                for j in 0..=n {
1384                    let choose = (0..j).fold(1.0, |c, i| c * (n - i) as f64 / (i + 1) as f64);
1385                    let xj = if j == 0 {
1386                        1.0
1387                    } else {
1388                        positive_power(family, j as f64, lo, hi)?
1389                    };
1390                    total += choose * s.powi(j) * o.powi(n - j) * xj;
1391                }
1392                finite(total)
1393            }
1394            _ => None,
1395        }
1396    }
1397
1398    /// E[Y^p] for Y uniform from `y0` to `y1`.
1399    fn uniform_power(y0: f64, y1: f64, p: f64) -> Option<f64> {
1400        // A real power is only taken where Y ≥ 0, which a coordinate's
1401        // rounding can put a hair below.
1402        let y0 = if p.fract() != 0.0 { y0.max(0.0) } else { y0 };
1403        if y1 <= y0 {
1404            return finite(y0.powf(p));
1405        }
1406        // Real powers need Y ≥ 0; negative ones, Y away from 0.
1407        if p.fract() != 0.0 && y0 < 0.0 || p < 0.0 && y0 <= 0.0 && y1 >= 0.0 {
1408            return None;
1409        }
1410        let integral = if p == -1.0 {
1411            libm::log(y1.abs()) - libm::log(y0.abs())
1412        } else {
1413            (y1.powf(p + 1.0) - y0.powf(p + 1.0)) / (p + 1.0)
1414        };
1415        finite(integral / (y1 - y0))
1416    }
1417
1418    /// E[X^p] for the draw X of a family on positive numbers.
1419    fn positive_power(family: &Family, p: f64, lo: f64, hi: f64) -> Option<f64> {
1420        let mass = hi - lo;
1421        let (a, b) = (family.quantile(lo), family.quantile(hi));
1422        let value = match *family {
1423            Family::Lognormal { mu, sigma } => {
1424                // X^p is lognormal: e^(pμ + p²σ²/2), on the shifted interval.
1425                let (alpha, beta) = (std_normal_quantile(lo), std_normal_quantile(hi));
1426                let share = std_normal_cdf(beta - p * sigma) - std_normal_cdf(alpha - p * sigma);
1427                (p * mu + 0.5 * p * p * sigma * sigma).exp() * share / mass
1428            }
1429            Family::Gamma { shape, scale } if shape + p > 0.0 => {
1430                let ratio = (libm::lgamma(shape + p) - libm::lgamma(shape)).exp() * scale.powf(p);
1431                let share = gamma_cdf(shape + p, b / scale) - gamma_cdf(shape + p, a / scale);
1432                ratio * share / mass
1433            }
1434            Family::Exponential { rate } => {
1435                return positive_power(
1436                    &Family::Gamma {
1437                        shape: 1.0,
1438                        scale: 1.0 / rate,
1439                    },
1440                    p,
1441                    lo,
1442                    hi,
1443                );
1444            }
1445            Family::Beta { a: alpha, b: beta } if alpha + p > 0.0 => {
1446                let ratio = (ln_beta(alpha + p, beta) - ln_beta(alpha, beta)).exp();
1447                let share = beta_cdf(alpha + p, beta, b) - beta_cdf(alpha + p, beta, a);
1448                ratio * share / mass
1449            }
1450            _ => return None,
1451        };
1452        finite(value)
1453    }
1454
1455    /// E[e^(c (s X + o))] for the draw X.
1456    pub fn exp(family: &Family, inner: Affine, c: f64, lo: f64, hi: f64) -> Option<f64> {
1457        let (a, b) = (family.quantile(lo), family.quantile(hi));
1458        let mass = hi - lo;
1459        let Affine { scale: s, offset: o } = inner;
1460        let k = c * s;
1461        let value = match *family {
1462            Family::Uniform { .. } => {
1463                let (ya, yb) = (inner.at(a), inner.at(b));
1464                let (y0, y1) = (ya.min(yb), ya.max(yb));
1465                let width = c * (y1 - y0);
1466                if width == 0.0 {
1467                    (c * y0).exp()
1468                } else {
1469                    (c * y0).exp() * libm::expm1(width) / width
1470                }
1471            }
1472            Family::Normal { mean, sd } => {
1473                // e^(kX) of a normal X, on the interval shifted by kσ.
1474                let (alpha, beta) = ((a - mean) / sd, (b - mean) / sd);
1475                let share = std_normal_cdf(beta - k * sd) - std_normal_cdf(alpha - k * sd);
1476                (c * o + k * mean + 0.5 * k * k * sd * sd).exp() * share / mass
1477            }
1478            Family::Gamma { shape, scale } if k * scale < 1.0 => {
1479                let r = 1.0 - k * scale;
1480                let share = gamma_cdf(shape, b * r / scale) - gamma_cdf(shape, a * r / scale);
1481                (c * o).exp() * r.powf(-shape) * share / mass
1482            }
1483            Family::Exponential { rate } => {
1484                return exp(
1485                    &Family::Gamma {
1486                        shape: 1.0,
1487                        scale: 1.0 / rate,
1488                    },
1489                    inner,
1490                    c,
1491                    lo,
1492                    hi,
1493                );
1494            }
1495            _ => return None,
1496        };
1497        finite(value)
1498    }
1499
1500    /// The mean and variance of ln(s X + o) for the draw X.
1501    pub fn ln(family: &Family, inner: Affine, lo: f64, hi: f64) -> Option<(f64, f64)> {
1502        let (a, b) = (family.quantile(lo), family.quantile(hi));
1503        let Affine { scale: s, offset: o } = inner;
1504        match *family {
1505            Family::Uniform { .. } => {
1506                let (ya, yb) = (inner.at(a), inner.at(b));
1507                // Y > 0 where `ln` is taken, but for a coordinate's rounding.
1508                let (y0, y1) = (ya.min(yb).max(0.0), ya.max(yb));
1509                if y1 <= y0 {
1510                    let l = libm::log(y0);
1511                    return Some((l, 0.0)).filter(|(l, _)| l.is_finite());
1512                }
1513                // ∫ ln y = y ln y − y, ∫ ln² y = y (ln² y − 2 ln y + 2),
1514                // both 0 at y = 0.
1515                let first = |y: f64| if y == 0.0 { 0.0 } else { y * libm::log(y) - y };
1516                let second = |y: f64| {
1517                    if y == 0.0 {
1518                        return 0.0;
1519                    }
1520                    let l = libm::log(y);
1521                    y * (l * l - 2.0 * l + 2.0)
1522                };
1523                let m = (first(y1) - first(y0)) / (y1 - y0);
1524                let m2 = (second(y1) - second(y0)) / (y1 - y0);
1525                Some((m, (m2 - m * m).max(0.0)))
1526            }
1527            // ln(sX) = ln s + μ + σZ, with Z normal where X is.
1528            Family::Lognormal { mu, sigma } if o == 0.0 && s > 0.0 => {
1529                let z = Family::Normal { mean: 0.0, sd: 1.0 };
1530                let (m, v) = z.interval_moments(std_normal_quantile(lo), std_normal_quantile(hi));
1531                Some((libm::log(s) + mu + sigma * m, sigma * sigma * v))
1532            }
1533            _ => None,
1534        }
1535    }
1536}
1537
1538/// Where independent axes are: a union of disjoint boxes, each a domain of
1539/// each axis in its CDF coordinates.
1540#[derive(Clone, Debug)]
1541pub struct Event {
1542    /// Each axis itself, with its domain where the event was made: no two
1543    /// share a latent.
1544    pub draws: Vec<Analytic>,
1545    pub boxes: Vec<Vec<Domain>>,
1546}
1547
1548/// The most boxes an event can have: each `or` of events of different
1549/// draws can multiply them.
1550const MAX_BOXES: usize = 10_000;
1551
1552/// The common parts of two unions of disjoint boxes.
1553fn intersect_boxes(a: &[Vec<Domain>], b: &[Vec<Domain>]) -> OpResult<Vec<Vec<Domain>>> {
1554    let mut out = Vec::with_capacity(a.len() * b.len());
1555    for x in a {
1556        for y in b {
1557            let both: Vec<Domain> = x.iter().zip(y).map(|(x, y)| x.intersect(y)).collect();
1558            if both.iter().all(|d| !d.0.is_empty()) {
1559                out.push(both);
1560            }
1561        }
1562        if out.len() > MAX_BOXES {
1563            return Err(OpError::limit(format!(
1564                "an event with more than {MAX_BOXES} boxes is over the limit"
1565            )));
1566        }
1567    }
1568    Ok(out)
1569}
1570
1571/// Where none of a union of disjoint boxes of `n` axes is: the box minus
1572/// each box in turn, each difference as disjoint boxes (outside the first
1573/// axis's part; inside it and outside the second's; …).
1574fn complement_boxes(boxes: &[Vec<Domain>], n: usize) -> OpResult<Vec<Vec<Domain>>> {
1575    let mut out = vec![vec![Domain::full(); n]];
1576    for b in boxes {
1577        let mut outside = Vec::with_capacity(n);
1578        for i in 0..n {
1579            let mut part: Vec<Domain> = b[..i].to_vec();
1580            part.push(b[i].complement());
1581            part.extend(std::iter::repeat_n(Domain::full(), n - i - 1));
1582            if part.iter().all(|d| !d.0.is_empty()) {
1583                outside.push(part);
1584            }
1585        }
1586        out = intersect_boxes(&out, &outside)?;
1587    }
1588    Ok(out)
1589}
1590
1591impl Event {
1592    pub fn single(draw: Analytic, yes: Domain) -> Self {
1593        Self {
1594            draws: vec![draw],
1595            boxes: vec![vec![yes]],
1596        }
1597    }
1598    pub fn key(&self) -> impl Ord + std::hash::Hash + use<> {
1599        (
1600            self.draws.iter().map(Analytic::key).collect::<Vec<_>>(),
1601            self.boxes
1602                .iter()
1603                .map(|b| b.iter().map(Domain::key).collect::<Vec<_>>())
1604                .collect::<Vec<_>>(),
1605        )
1606    }
1607    pub fn probability(&self) -> f64 {
1608        let p: f64 = self
1609            .boxes
1610            .iter()
1611            .map(|b| self.draws.iter().zip(b).map(|(d, yes)| d.share(yes)).product::<f64>())
1612            .sum();
1613        p.clamp(0.0, 1.0)
1614    }
1615    pub fn value(self) -> Value {
1616        let p = self.probability();
1617        if p == 0.0 || p == 1.0 {
1618            Value::Bool(p == 1.0)
1619        } else {
1620            Value::Event(Arc::new(self))
1621        }
1622    }
1623    /// Where it doesn't hold.
1624    pub fn complement(&self) -> OpResult<Event> {
1625        Ok(Event {
1626            draws: self.draws.clone(),
1627            boxes: self.side(false)?,
1628        })
1629    }
1630    /// The boxes where it holds, or where it doesn't. Of one axis, that's
1631    /// one box.
1632    pub fn side(&self, yes: bool) -> OpResult<Vec<Vec<Domain>>> {
1633        match (&self.boxes[..], self.draws.len()) {
1634            ([b], 1) => Ok(vec![vec![if yes { b[0].clone() } else { b[0].complement() }]]),
1635            _ if yes => Ok(self.boxes.clone()),
1636            _ => complement_boxes(&self.boxes, self.draws.len()),
1637        }
1638    }
1639    /// How many intervals it holds, which operations on it cost.
1640    pub fn size(&self) -> usize {
1641        let boxes: usize = self
1642            .boxes
1643            .iter()
1644            .map(|b| b.iter().map(|d| d.0.len()).sum::<usize>())
1645            .sum();
1646        boxes + self.draws.iter().map(|d| d.domain.0.len()).sum::<usize>()
1647    }
1648    /// Whether worlds can be restricted to it: each axis is a latent, not
1649    /// a form of several. Restricting a form would make its latents a
1650    /// truncated joint distribution.
1651    pub fn restrictable(&self) -> bool {
1652        self.draws.iter().all(|d| d.form.is_none())
1653    }
1654    /// Restricts a world to one of its boxes: the box's probability given
1655    /// what the world knew.
1656    pub fn restrict(&self, part: &[Domain], context: &mut Constraints) -> f64 {
1657        let mut p = 1.0;
1658        for (draw, yes) in self.draws.iter().zip(part) {
1659            let prior = &draw.domain;
1660            let domain = prior.intersect(yes);
1661            p *= draw.measure(&domain) / draw.measure(prior);
1662            Arc::make_mut(context).insert(
1663                draw.id,
1664                Latent {
1665                    family: draw.family,
1666                    domain,
1667                    form: None,
1668                    weight: draw.weight.clone(),
1669                },
1670            );
1671        }
1672        p.clamp(0.0, 1.0)
1673    }
1674    /// Both events, either, or the same outcome for both, by `op`
1675    /// (`And`, `Or`, or `Eq`). Their axes must be the same or independent.
1676    pub fn combine(&self, other: &Event, op: BinOp) -> OpResult<Event> {
1677        let mut draws = self.draws.clone();
1678        let mut at = Vec::with_capacity(other.draws.len());
1679        for d in &other.draws {
1680            match draws.iter().position(|x| x.same_axis(d)) {
1681                Some(i) => at.push(i),
1682                None => {
1683                    if draws.iter().any(|x| x.shares_latents(d)) {
1684                        return Err(unsupported("combining events of draws that depend on each other"));
1685                    }
1686                    at.push(draws.len());
1687                    draws.push(d.clone());
1688                }
1689            }
1690        }
1691        let n = draws.len();
1692        let mine: Vec<Vec<Domain>> = self
1693            .boxes
1694            .iter()
1695            .map(|b| {
1696                let mut part = b.clone();
1697                part.resize(n, Domain::full());
1698                part
1699            })
1700            .collect();
1701        let theirs: Vec<Vec<Domain>> = other
1702            .boxes
1703            .iter()
1704            .map(|b| {
1705                let mut part = vec![Domain::full(); n];
1706                for (d, &i) in b.iter().zip(&at) {
1707                    part[i] = d.clone();
1708                }
1709                part
1710            })
1711            .collect();
1712        let boxes = match op {
1713            BinOp::And => intersect_boxes(&mine, &theirs)?,
1714            BinOp::Or => {
1715                let mut boxes = mine.clone();
1716                boxes.extend(intersect_boxes(&theirs, &complement_boxes(&mine, n)?)?);
1717                boxes
1718            }
1719            _ => {
1720                let mut boxes = intersect_boxes(&mine, &theirs)?;
1721                let (not_mine, not_theirs) = (complement_boxes(&mine, n)?, complement_boxes(&theirs, n)?);
1722                boxes.extend(intersect_boxes(&not_mine, &not_theirs)?);
1723                boxes
1724            }
1725        };
1726        Ok(Event { draws, boxes }.simplified())
1727    }
1728    /// Of one axis, its boxes as one.
1729    fn simplified(mut self) -> Event {
1730        if self.draws.len() == 1 && self.boxes.len() != 1 {
1731            let union = self.boxes.iter().fold(Domain(Vec::new()), |all, b| all.union(&b[0]));
1732            self.boxes = vec![vec![union]];
1733        }
1734        self
1735    }
1736}
1737
1738/// An exact update of a draw `x`, read in the world it's updated in, from
1739/// observing `seen` with `x` itself as the parameter: the logarithm of the
1740/// observation's probability (a density for a normal), and what the world
1741/// knows about the draw after it. `None` if the draw's distribution isn't
1742/// the prior of a conjugate pair with this observation.
1743///
1744/// A draw restricted to part of its range stays restricted: its posterior
1745/// is the updated family on the same values, and the probability includes
1746/// the ratio of the updated and the current family's mass there.
1747pub fn update(x: &Analytic, seen: crate::conjugate::Seen) -> OpResult<Option<(f64, Latent)>> {
1748    use crate::conjugate::Seen;
1749    if x.affine() != Some(Affine::IDENTITY) || x.form.is_some() || x.weight.is_some() {
1750        return Ok(None);
1751    }
1752    // A uniform probability is a beta(1, 1) on part of its range, and an
1753    // exponential rate a gamma with shape 1.
1754    let prior = match (x.family, seen) {
1755        (Family::Uniform { lo, hi }, Seen::Binomial { .. } | Seen::Bernoulli(_)) if lo >= 0.0 && hi <= 1.0 => {
1756            Family::Beta { a: 1.0, b: 1.0 }
1757        }
1758        (Family::Exponential { rate }, Seen::Poisson { .. }) => Family::Gamma {
1759            shape: 1.0,
1760            scale: 1.0 / rate,
1761        },
1762        (family, _) => family,
1763    };
1764    let domain = x.domain.moved(&x.family, &prior);
1765    let Some((ln, posterior)) = crate::conjugate::update(&prior, seen) else {
1766        return Ok(None);
1767    };
1768    if ln == f64::NEG_INFINITY {
1769        return Ok(Some((
1770            ln,
1771            Latent {
1772                family: prior,
1773                domain,
1774                form: None,
1775                weight: None,
1776            },
1777        )));
1778    }
1779    let after = domain.moved(&prior, &posterior);
1780    if after.mass() == 0.0 {
1781        // Not impossible: beyond what CDF coordinates can tell apart.
1782        return Err(OpError::unsupported(
1783            "this update leaves the draw where its updated distribution has too little probability to represent",
1784        )
1785        .help("the draw is restricted to a region far into the tail of its distribution after this observation"));
1786    }
1787    let ln = if domain == Domain::full() && after == Domain::full() {
1788        ln
1789    } else {
1790        ln + libm::log(after.mass()) - libm::log(domain.mass())
1791    };
1792    Ok(Some((
1793        ln,
1794        Latent {
1795            family: posterior,
1796            domain: after,
1797            form: None,
1798            weight: None,
1799        },
1800    )))
1801}
1802
1803/// Where an outcome of a draw, by its `pieces`, is outside the range from
1804/// `lo` to `hi`, in the draw's coordinates within its domain. Each bound is
1805/// included when its flag is.
1806pub fn outside(
1807    x: &Analytic,
1808    pieces: &[(f64, Fun)],
1809    lo: Option<(f64, bool)>,
1810    hi: Option<(f64, bool)>,
1811) -> OpResult<Domain> {
1812    let pieces = split_turns(&x.family, pieces.to_vec());
1813    let mut out = Domain(Vec::new());
1814    if let Some((t, included)) = lo {
1815        out = out.union(&below(&x.family, &pieces, t, included)?);
1816    }
1817    if let Some((t, included)) = hi {
1818        out = out.union(&below(&x.family, &pieces, t, !included)?.complement());
1819    }
1820    Ok(x.domain.intersect(&out))
1821}
1822
1823/// A draw from `normal(mean, sd)`, with `mean` an outcome of normal draws:
1824/// the mean plus a fresh normal latent `id` with standard deviation `sd`.
1825/// `None` if the mean isn't an unrestricted linear form of normal draws.
1826pub fn normal_draw(mean: &Analytic, sd: f64, id: u64) -> OpResult<Option<Value>> {
1827    let Some(m) = mean.linear() else {
1828        return Ok(None);
1829    };
1830    let noise = Linear {
1831        constant: 0.0,
1832        terms: vec![Term {
1833            id,
1834            coef: 1.0,
1835            family: Family::Normal { mean: 0.0, sd },
1836        }],
1837    };
1838    Ok(Some(m.plus(&noise, 1.0)?.value()?))
1839}
1840
1841/// An observation's log density, and what the world knows about each latent
1842/// after it.
1843pub type Observed = (f64, Vec<(u64, Latent)>);
1844
1845/// Observing `y` from `normal(mean, sd)`, with `mean` an outcome of normal
1846/// draws: the logarithm of its density, and what the world knows about
1847/// each latent after it. `None` if the mean isn't affine in a normal latent
1848/// or a linear form of several.
1849///
1850/// With one latent, it's the conjugate update, by what `y` says about the
1851/// latent itself. With several, which are unrestricted, their posterior is
1852/// normal with a covariance. Each becomes its posterior mean plus a linear
1853/// form of fresh standard normal latents, numbered after `next`: with
1854/// prior standard deviations `s`, coefficients `a` and `S` the variance of
1855/// `y`, the covariance is D^½ (I − c cᵀ) D^½ for D = diag(s²) and
1856/// c = s a / √S, and its square root D^½ (I − γ c cᵀ), γ = 1 / (1 + sd / √S).
1857pub fn observe_normal(mean: &Analytic, y: f64, sd: f64, next: &mut u64) -> OpResult<Option<Observed>> {
1858    if mean.weight.is_some() {
1859        return Ok(None);
1860    }
1861    if mean.form.is_none() && matches!(mean.family, Family::Normal { .. }) {
1862        let Some(a) = mean.affine() else {
1863            return Ok(None);
1864        };
1865        // y ~ normal(s x + o, sd) is (y − o) / s ~ normal(x, sd / |s|), with
1866        // the density divided by |s|.
1867        let seen = crate::conjugate::Seen::Normal {
1868            y: (y - a.offset) / a.scale,
1869            sd: sd / a.scale.abs(),
1870        };
1871        let updated = update(&mean.latent(), seen)?;
1872        return Ok(updated.map(|(ln, latent)| (ln - libm::log(a.scale.abs()), vec![(mean.id, latent)])));
1873    }
1874    let Some(form) = mean.linear() else {
1875        return Ok(None);
1876    };
1877    let (m, variance) = form.moments();
1878    let total = variance + sd * sd;
1879    let r = y - m;
1880    let ln = -0.5 * libm::log(2.0 * std::f64::consts::PI * total) - r * r / (2.0 * total);
1881    let root = libm::sqrt(total);
1882    let gamma = 1.0 / (1.0 + sd / root);
1883    let c: Vec<f64> = form.terms.iter().map(|t| t.family.sd() * t.coef / root).collect();
1884    let fresh: Vec<u64> = form
1885        .terms
1886        .iter()
1887        .map(|_| {
1888            *next += 1;
1889            *next
1890        })
1891        .collect();
1892    let standard = Family::Normal { mean: 0.0, sd: 1.0 };
1893    let mut known = Vec::with_capacity(form.terms.len());
1894    for (i, t) in form.terms.iter().enumerate() {
1895        let s = t.family.sd();
1896        let constant = t.family.mean() + s * s * t.coef * r / total;
1897        let terms = fresh
1898            .iter()
1899            .zip(&c)
1900            .enumerate()
1901            .map(|(j, (&id, &cj))| Term {
1902                id,
1903                coef: s * (f64::from(u8::from(i == j)) - gamma * c[i] * cj),
1904                family: standard,
1905            })
1906            .collect();
1907        let linear = Linear::from_terms(constant, terms)?;
1908        let latent = Latent {
1909            family: linear.normal()?,
1910            domain: Domain::full(),
1911            form: Some(Arc::new(linear)),
1912            weight: None,
1913        };
1914        known.push((t.id, latent));
1915    }
1916    Ok(Some((ln, known)))
1917}
1918
1919pub fn unsupported(what: &str) -> OpError {
1920    OpError::unsupported(format!("{what} isn't supported for analytic continuous draws yet"))
1921        .help("use `@mode sample(runs: 10_000)` for this operation")
1922}
1923
1924/// Restricting worlds by an outcome of several normal draws together would
1925/// make them a truncated joint distribution.
1926pub fn joint_condition(what: &str) -> OpError {
1927    OpError::unsupported(format!("{what} isn't supported when enumerating yet")).help(
1928        "it would make the draws a truncated joint distribution: report it instead, or use `@mode sample(runs: 10_000)`",
1929    )
1930}
1931
1932/// An operand of an operation on outcomes of one draw.
1933enum Operand<'a> {
1934    Pieces(&'a [(f64, Fun)]),
1935    Constant(f64),
1936}
1937
1938fn operand<'a>(v: &'a Value, x: &Analytic) -> OpResult<Operand<'a>> {
1939    match v {
1940        Value::Analytic(y) if y.same_axis(x) => Ok(Operand::Pieces(&y.pieces)),
1941        Value::Analytic(_) => Err(different_draws()),
1942        v => v
1943            .as_f64()
1944            .filter(|v| v.is_finite())
1945            .map(Operand::Constant)
1946            .ok_or_else(|| unsupported("this operand")),
1947    }
1948}
1949
1950/// `f(a, b)` on each part where both operands have one function.
1951fn zip(a: &Operand, b: &Operand, f: impl Fn(Fun, Fun) -> OpResult<Fun>) -> OpResult<Pieces> {
1952    match (a, b) {
1953        (Operand::Pieces(a), Operand::Constant(c)) => {
1954            a.iter().map(|&(end, a)| Ok((end, f(a, Fun::constant(*c))?))).collect()
1955        }
1956        (Operand::Constant(c), Operand::Pieces(b)) => {
1957            b.iter().map(|&(end, b)| Ok((end, f(Fun::constant(*c), b)?))).collect()
1958        }
1959        (Operand::Pieces(a), Operand::Pieces(b)) => overlay(a, b)
1960            .into_iter()
1961            .map(|(end, (a, b))| Ok((end, f(a, b)?)))
1962            .collect(),
1963        (Operand::Constant(_), Operand::Constant(_)) => unreachable!("an operand is analytic"),
1964    }
1965}
1966
1967fn nonlinear() -> OpError {
1968    unsupported("nonlinear arithmetic")
1969}
1970
1971/// `a + b`, or `a − b` with `sign` −1. Affine functions add as they did;
1972/// polynomials of degree 2 add as polynomials; the same kernel of the same
1973/// inner function adds its outer ones.
1974fn sum(a: Fun, b: Fun, sign: f64) -> OpResult<Fun> {
1975    if let (Some(x), Some(y)) = (a.as_affine(), b.as_affine()) {
1976        return Ok(Fun::affine(if sign > 0.0 {
1977            Affine {
1978                scale: x.scale + y.scale,
1979                offset: x.offset + y.offset,
1980            }
1981        } else {
1982            Affine {
1983                scale: x.scale - y.scale,
1984                offset: x.offset - y.offset,
1985            }
1986        }));
1987    }
1988    if let Some(c) = b.as_constant() {
1989        return Ok(a.then(1.0, sign * c));
1990    }
1991    if let Some(c) = a.as_constant() {
1992        return Ok(b.then(sign, c));
1993    }
1994    if let (Some(p), Some(q)) = (a.polynomial(), b.polynomial()) {
1995        return Ok(Fun::from_polynomial([0, 1, 2].map(|i| p[i] + sign * q[i])));
1996    }
1997    if a.kernel == b.kernel && a.inner == b.inner {
1998        let outer = Affine {
1999            scale: a.outer.scale + sign * b.outer.scale,
2000            offset: a.outer.offset + sign * b.outer.offset,
2001        };
2002        return Ok(if outer.scale == 0.0 {
2003            Fun::constant(outer.offset)
2004        } else {
2005            Fun { outer, ..a }
2006        });
2007    }
2008    Err(nonlinear())
2009}
2010
2011fn difference(a: Fun, b: Fun) -> OpResult<Fun> {
2012    sum(a, b, -1.0)
2013}
2014
2015/// `a × b`: a constant scales; affine functions multiply into a square.
2016fn product(a: Fun, b: Fun) -> OpResult<Fun> {
2017    if let (Some(x), Some(y)) = (a.as_affine(), b.as_affine()) {
2018        if y.scale == 0.0 {
2019            return Ok(Fun::affine(Affine {
2020                scale: x.scale * y.offset,
2021                offset: x.offset * y.offset,
2022            }));
2023        }
2024        if x.scale == 0.0 {
2025            return Ok(Fun::affine(Affine {
2026                scale: y.scale * x.offset,
2027                offset: y.offset * x.offset,
2028            }));
2029        }
2030    }
2031    if let Some(c) = b.as_constant() {
2032        return Ok(scaled(a, c));
2033    }
2034    if let Some(c) = a.as_constant() {
2035        return Ok(scaled(b, c));
2036    }
2037    match (a.polynomial(), b.polynomial()) {
2038        (Some([p0, p1, 0.0]), Some([q0, q1, 0.0])) => Ok(Fun::from_polynomial([p0 * q0, p0 * q1 + p1 * q0, p1 * q1])),
2039        _ => Err(nonlinear()),
2040    }
2041}
2042
2043/// `c × f`, a constant again when `c` is 0.
2044fn scaled(f: Fun, c: f64) -> Fun {
2045    if c == 0.0 { Fun::constant(0.0) } else { f.then(c, 0.0) }
2046}
2047
2048/// `a ÷ b`: by a constant, or a constant by an affine function.
2049fn quotient(a: Fun, b: Fun, plain: bool) -> OpResult<Fun> {
2050    if let Some(c) = b.as_constant() {
2051        if c == 0.0 {
2052            return Err(if plain {
2053                OpError::fault(Fault::DivisionByZero, "division by zero")
2054            } else {
2055                unsupported("dividing by a value that's zero on part of its domain")
2056            });
2057        }
2058        return Ok(match a.as_affine() {
2059            Some(x) => Fun::affine(Affine {
2060                scale: x.scale / c,
2061                offset: x.offset / c,
2062            }),
2063            None => Fun {
2064                outer: Affine {
2065                    scale: a.outer.scale / c,
2066                    offset: a.outer.offset / c,
2067                },
2068                ..a
2069            },
2070        });
2071    }
2072    // c / (s y + o) = (c / s) / (y + o / s).
2073    if let (Some(c), Some(y)) = (a.as_constant(), b.as_affine()) {
2074        return Ok(Fun {
2075            outer: Affine {
2076                scale: c / y.scale,
2077                offset: 0.0,
2078            },
2079            kernel: Kernel::Recip,
2080            inner: Affine {
2081                scale: 1.0,
2082                offset: y.offset / y.scale,
2083            },
2084        });
2085    }
2086    Err(nonlinear())
2087}
2088
2089/// The most pieces an outcome can have. Repeated folds like
2090/// `y = abs(2 * y - 1)` double them each time.
2091const MAX_PIECES: usize = 100_000;
2092
2093/// Pieces cost work, and their number is limited.
2094fn spend(pieces: &[(f64, Fun)], budget: &mut Budget) -> OpResult<()> {
2095    if pieces.len() > MAX_PIECES {
2096        return Err(OpError::limit(format!(
2097            "a continuous outcome with more than {MAX_PIECES} pieces is over the limit"
2098        )));
2099    }
2100    if pieces.len() > 1 {
2101        budget.work(pieces.len() as u64)?;
2102    }
2103    Ok(())
2104}
2105
2106pub fn binary(op: BinOp, a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
2107    use BinOp::*;
2108    if matches!(a, Value::Joint(_)) || matches!(b, Value::Joint(_)) {
2109        return crate::joint::binary(op, a, b, budget);
2110    }
2111    let x = match (a, b) {
2112        (Value::Analytic(x), _) | (_, Value::Analytic(x)) => x,
2113        _ => return Err(unsupported("this operation")),
2114    };
2115    if op == Pow {
2116        return power(a, b, budget);
2117    }
2118    if let (Value::Analytic(p), Value::Analytic(q)) = (a, b) {
2119        if !p.same_axis(q) {
2120            return combined(op, p, q, budget);
2121        }
2122    }
2123    let (p, q) = (operand(a, x)?, operand(b, x)?);
2124    let plain = matches!(q, Operand::Constant(_));
2125    let pieces = match op {
2126        Add => zip(&p, &q, |a, b| sum(a, b, 1.0))?,
2127        Sub => zip(&p, &q, difference)?,
2128        Mul => zip(&p, &q, product)?,
2129        Div => zip(&p, &q, |a, b| quotient(a, b, plain))?,
2130        Eq | Ne | Lt | Le | Gt | Ge => {
2131            let d = split_turns(&x.family, zip(&p, &q, difference)?);
2132            spend(&d, budget)?;
2133            if !d.iter().all(|(_, f)| f.is_finite()) {
2134                return Err(unsupported("non-finite affine coefficients"));
2135            }
2136            return Ok(compare(x, op, &d)?.value());
2137        }
2138        _ => return Err(nonlinear()),
2139    };
2140    let pieces = split_turns(&x.family, pieces);
2141    spend(&pieces, budget)?;
2142    // As with numbers, sums, differences and products of ints are ints.
2143    let int = matches!(op, Add | Sub | Mul) && is_int(a) && is_int(b);
2144    x.with_type(pieces, int).value()
2145}
2146
2147/// Combining outcomes of different axes: only sums and differences of
2148/// normal latents, which are linear forms of them.
2149fn different_draws() -> OpError {
2150    OpError::unsupported(
2151        "combining different continuous draws isn't supported when enumerating, except to add or subtract normal ones without restrictions",
2152    )
2153    .help("use `@mode sample(runs: 10_000)` for this operation")
2154}
2155
2156/// An operation on outcomes of different axes: a sum or difference of
2157/// normal latents is a linear form of them, and a comparison an event of
2158/// the difference.
2159fn combined(op: BinOp, p: &Analytic, q: &Analytic, budget: &mut Budget) -> OpResult<Value> {
2160    use BinOp::*;
2161    let sign = match op {
2162        Add => 1.0,
2163        Sub | Eq | Ne | Lt | Le | Gt | Ge => -1.0,
2164        _ => return Err(different_draws()),
2165    };
2166    let (Some(a), Some(b)) = (p.linear(), q.linear()) else {
2167        return Err(different_draws());
2168    };
2169    budget.work((a.terms.len() + b.terms.len()) as u64)?;
2170    let d = a.plus(&b, sign)?.value()?;
2171    match op {
2172        Add | Sub => Ok(d),
2173        _ => crate::ops::binary(op, &d, &Value::Float(0.0), budget),
2174    }
2175}
2176
2177/// `a ^ b` with an outcome of a draw: a square, a square root or a
2178/// reciprocal of one, or a positive number to the power of an affine one.
2179fn power(a: &Value, b: &Value, budget: &mut Budget) -> OpResult<Value> {
2180    match (a, b.as_f64()) {
2181        (Value::Analytic(_), Some(2.0)) => binary(BinOp::Mul, a, a, budget),
2182        (Value::Analytic(_), Some(1.0)) => Ok(a.clone()),
2183        (Value::Analytic(y), Some(0.5)) => transformed(y, Kernel::Sqrt, budget),
2184        (Value::Analytic(_), Some(-1.0)) => binary(BinOp::Div, &Value::Float(1.0), a, budget),
2185        (Value::Analytic(_), _) => Err(unsupported("this power")),
2186        // c^y = e^(y ln c).
2187        (base, _) => match (base.as_f64(), b) {
2188            (Some(c), Value::Analytic(_)) if c > 0.0 => {
2189                let exponent = binary(BinOp::Mul, b, &Value::Float(libm::log(c)), budget)?;
2190                match exponent {
2191                    Value::Analytic(e) => transformed(&e, Kernel::Exp, budget),
2192                    v => crate::ops::binary(BinOp::Pow, base, &v, budget),
2193                }
2194            }
2195            _ => Err(unsupported("this power")),
2196        },
2197    }
2198}
2199
2200/// The function a kernel makes of `f`, if it makes one this can hold.
2201fn compose(kernel: Kernel, f: Fun, family: &Family, start: f64, end: f64) -> OpResult<Fun> {
2202    if let Some(c) = f.as_constant() {
2203        // Outside the domain, the caller has checked, a constant can't be
2204        // there with any probability.
2205        let y = kernel.apply(c);
2206        return Ok(Fun::constant(if y.is_finite() { y } else { 0.0 }));
2207    }
2208    if let Some(inner) = f.as_affine() {
2209        return Ok(Fun::of(kernel, inner));
2210    }
2211    let Affine { scale: s, offset: o } = f.outer;
2212    match (kernel, f.kernel) {
2213        // √(s y²) = √s |y|, on the side of 0 the piece is on.
2214        (Kernel::Sqrt, Kernel::Square) if o == 0.0 && s > 0.0 => {
2215            let sign = if f.side(family, start, end) < 0.0 { -1.0 } else { 1.0 };
2216            let k = sign * s.sqrt();
2217            Ok(Fun::affine(Affine {
2218                scale: k * f.inner.scale,
2219                offset: k * f.inner.offset,
2220            }))
2221        }
2222        // ln(s e^y) = ln s + y.
2223        (Kernel::Ln, Kernel::Exp) if o == 0.0 && s > 0.0 => Ok(Fun::affine(Affine {
2224            scale: f.inner.scale,
2225            offset: f.inner.offset + libm::log(s),
2226        })),
2227        // e^(ln y + o) = e^o y.
2228        (Kernel::Exp, Kernel::Ln) if s == 1.0 => {
2229            let k = o.exp();
2230            Ok(Fun::affine(Affine {
2231                scale: k * f.inner.scale,
2232                offset: k * f.inner.offset,
2233            }))
2234        }
2235        _ => Err(unsupported("this composition of continuous transforms")),
2236    }
2237}
2238
2239/// `sqrt`, `exp` or `ln` of an outcome of a draw. Where the outcome is
2240/// outside what the function is defined for, with positive probability, it
2241/// faults there: everywhere is a fault, as for a number; part of the
2242/// domain isn't supported yet.
2243pub fn transformed(x: &Analytic, kernel: Kernel, budget: &mut Budget) -> OpResult<Value> {
2244    let (name, outside) = match kernel {
2245        Kernel::Sqrt => ("sqrt", below(&x.family, &x.pieces, 0.0, true)?),
2246        Kernel::Ln => ("ln", below(&x.family, &x.pieces, 0.0, false)?),
2247        // Beyond this, e^y isn't a finite float.
2248        Kernel::Exp => (
2249            "exp",
2250            below(&x.family, &x.pieces, 709.782712893384, false)?.complement(),
2251        ),
2252        _ => unreachable!("a transform built in"),
2253    };
2254    let faulting = x.domain.intersect(&outside).mass();
2255    if faulting > 0.0 {
2256        return Err(if faulting >= x.domain.mass() {
2257            OpError::fault(Fault::DomainError, format!("`{name}` isn't defined for this outcome"))
2258        } else {
2259            unsupported(&format!(
2260                "`{name}` of an outcome outside its domain on part of its range"
2261            ))
2262        });
2263    }
2264    let mut pieces = Vec::with_capacity(x.pieces.len());
2265    let mut start = 0.0;
2266    for &(end, f) in &x.pieces {
2267        pieces.push((end, compose(kernel, f, &x.family, start, end)?));
2268        start = end;
2269    }
2270    spend(&pieces, budget)?;
2271    x.with(pieces).value()
2272}
2273
2274/// The event that a difference `d` of outcomes of `x`'s draw compares with
2275/// zero by `op`.
2276fn compare(x: &Analytic, op: BinOp, d: &[(f64, Fun)]) -> OpResult<Event> {
2277    use BinOp::*;
2278    let yes = match op {
2279        Lt => below(&x.family, d, 0.0, true)?,
2280        Le => below(&x.family, d, 0.0, false)?,
2281        Gt => below(&x.family, d, 0.0, false)?.complement(),
2282        Ge => below(&x.family, d, 0.0, true)?.complement(),
2283        _ => {
2284            // Only a constant piece can equal zero with positive probability.
2285            let mut equal = Vec::new();
2286            let mut start = 0.0;
2287            for &(end, f) in d {
2288                if f.as_constant() == Some(0.0) {
2289                    push(&mut equal, start, end);
2290                }
2291                start = end;
2292            }
2293            let equal = Domain(equal);
2294            if op == Eq { equal } else { equal.complement() }
2295        }
2296    };
2297    Ok(Event::single(x.latent(), yes))
2298}
2299
2300pub fn negate(x: &Analytic) -> OpResult<Value> {
2301    x.with_type(x.pieces.iter().map(|&(end, f)| (end, f.negated())).collect(), x.int)
2302        .value()
2303}
2304
2305pub fn abs(x: &Analytic, budget: &mut Budget) -> OpResult<Value> {
2306    let negative = below(&x.family, &x.pieces, 0.0, true)?;
2307    let negated: Pieces = x.pieces.iter().map(|&(end, f)| (end, f.negated())).collect();
2308    let pieces = select(&negative, &negated, &x.pieces);
2309    spend(&pieces, budget)?;
2310    x.with_type(pieces, x.int).value()
2311}
2312
2313/// A number as the language types it: an int, or a float.
2314fn number(c: f64, int: bool) -> Value {
2315    match probl_number::Integer::from_f64(c).filter(|_| int) {
2316        Some(n) => Value::Int(n),
2317        None => Value::Float(c),
2318    }
2319}
2320
2321/// Whether a value is an int, or an outcome that is one.
2322fn is_int(v: &Value) -> bool {
2323    match v {
2324        Value::Int(_) => true,
2325        Value::Analytic(x) => x.int,
2326        _ => false,
2327    }
2328}
2329
2330/// How `floor`, `ceil`, `trunc` and `round` make an int of a number.
2331#[derive(Clone, Copy, Debug)]
2332pub enum Rounding {
2333    Floor,
2334    Ceil,
2335    Trunc,
2336    Round,
2337}
2338impl Rounding {
2339    fn apply(self, v: f64) -> f64 {
2340        match self {
2341            Rounding::Floor => v.floor(),
2342            Rounding::Ceil => v.ceil(),
2343            Rounding::Trunc => libm::trunc(v),
2344            // Half away from zero, as `round` does.
2345            Rounding::Round => v.round(),
2346        }
2347    }
2348    /// Where it can change, strictly between `lo` and `hi`: the whole
2349    /// numbers, or the halves for `round`.
2350    fn steps(self, lo: f64, hi: f64) -> impl Iterator<Item = f64> {
2351        let half = if matches!(self, Rounding::Round) { 0.5 } else { 0.0 };
2352        let first = (lo - half).floor() + 1.0;
2353        let last = (hi - half).ceil() - 1.0;
2354        (0..)
2355            .map(move |i| first + i as f64)
2356            .take_while(move |k| *k <= last)
2357            .map(move |k| k + half)
2358    }
2359}
2360
2361/// `floor(x)`, `ceil(x)`, `trunc(x)` or `round(x)`: an int that is constant
2362/// on each part of the draw where `x` is between two of the rounding's
2363/// steps, with that part's probability. Where `x` is constant, it rounds as
2364/// a number does, ties included; elsewhere a step has no probability.
2365pub fn rounded(x: &Analytic, how: Rounding, budget: &mut Budget) -> OpResult<Value> {
2366    let family = &x.family;
2367    let mut pieces: Pieces = Vec::new();
2368    let mut at = 0.0;
2369    for (lo, hi, f) in x.segments() {
2370        if lo > at {
2371            // Outside the domain: anything, which `value` gives away.
2372            pieces.push((lo, Fun::constant(0.0)));
2373        }
2374        at = hi;
2375        if let Some(c) = f.as_constant() {
2376            pieces.push((hi, Fun::constant(how.apply(c))));
2377            continue;
2378        }
2379        let (a, b) = f.ends(family, lo, hi);
2380        if !a.is_finite() || !b.is_finite() {
2381            return Err(unsupported("rounding an outcome whose values have no bound"));
2382        }
2383        let (vmin, vmax) = (a.min(b), a.max(b));
2384        let steps: Vec<f64> = how.steps(vmin, vmax).take(MAX_PIECES + 1).collect();
2385        if pieces.len() + steps.len() > MAX_PIECES {
2386            return Err(OpError::limit(format!(
2387                "a continuous outcome with more than {MAX_PIECES} pieces is over the limit"
2388            )));
2389        }
2390        budget.work(steps.len() as u64 + 1)?;
2391        // Each step, as a coordinate of the draw, in increasing order.
2392        let up = f.increasing(family, lo, hi);
2393        let mut bounds: Vec<(f64, f64)> = steps
2394            .iter()
2395            .map(|&t| (f.coordinate(family, lo, hi, t).clamp(lo, hi), t))
2396            .collect();
2397        if !up {
2398            bounds.reverse();
2399        }
2400        let mut from = if up { vmin } else { vmax };
2401        for (u, t) in bounds {
2402            pieces.push((u, Fun::constant(how.apply((from + t) / 2.0))));
2403            from = t;
2404        }
2405        let to = if up { vmax } else { vmin };
2406        pieces.push((hi, Fun::constant(how.apply((from + to) / 2.0))));
2407    }
2408    if at < 1.0 {
2409        pieces.push((1.0, Fun::constant(0.0)));
2410    }
2411    // Ends that rounded together leave parts with no probability.
2412    let mut tiling: Pieces = Vec::with_capacity(pieces.len());
2413    for (end, f) in pieces {
2414        match tiling.last_mut() {
2415            Some(last) if end <= last.0 => {}
2416            _ => tiling.push((end, f)),
2417        }
2418    }
2419    if let Some(last) = tiling.last_mut() {
2420        last.0 = 1.0;
2421    }
2422    spend(&tiling, budget)?;
2423    x.with_type(tiling, true).value()
2424}
2425
2426/// `min` or `max` of numbers and outcomes of one draw. As with numbers, each
2427/// argument replaces the best so far where it's strictly better.
2428pub fn extreme(args: &[Value], want_max: bool, budget: &mut Budget) -> OpResult<Value> {
2429    let x = args
2430        .iter()
2431        .find_map(|v| match v {
2432            Value::Analytic(x) => Some(x),
2433            _ => None,
2434        })
2435        .expect("an analytic argument");
2436    let tiling = |v: &Value| -> OpResult<Pieces> {
2437        Ok(match operand(v, x)? {
2438            Operand::Pieces(p) => p.to_vec(),
2439            Operand::Constant(c) => vec![(1.0, Fun::constant(c))],
2440        })
2441    };
2442    let mut best = tiling(&args[0])?;
2443    for v in &args[1..] {
2444        let v = tiling(v)?;
2445        let d = split_turns(
2446            &x.family,
2447            zip(&Operand::Pieces(&v), &Operand::Pieces(&best), difference)?,
2448        );
2449        let better = if want_max {
2450            below(&x.family, &d, 0.0, false)?.complement()
2451        } else {
2452            below(&x.family, &d, 0.0, true)?
2453        };
2454        best = select(&better, &v, &best);
2455        spend(&best, budget)?;
2456    }
2457    x.with_type(best, args.iter().all(is_int)).value()
2458}
2459
2460/// `clamp(x, lo, hi)` with numbers for bounds, which the caller has checked
2461/// are in order. It's an int if `x` and the bounds are.
2462pub fn clamp(x: &Analytic, lo: &Value, hi: &Value, budget: &mut Budget) -> OpResult<Value> {
2463    let int = x.int && is_int(lo) && is_int(hi);
2464    let (Some(lo), Some(hi)) = (lo.as_f64(), hi.as_f64()) else {
2465        return Err(unsupported("`clamp` with these bounds"));
2466    };
2467    let under = below(&x.family, &x.pieces, lo, true)?;
2468    let over = below(&x.family, &x.pieces, hi, false)?.complement();
2469    let inside = select(&over, &[(1.0, Fun::constant(hi))], &x.pieces);
2470    let pieces = select(&under, &[(1.0, Fun::constant(lo))], &inside);
2471    spend(&pieces, budget)?;
2472    x.with_type(pieces, int).value()
2473}
2474
2475pub fn logic(and: bool, a: &Value, b: &Value) -> OpResult<Value> {
2476    let (event, other) = match (a, b) {
2477        (Value::Event(e), b) | (b, Value::Event(e)) => (e, b),
2478        _ => return Err(unsupported("this logical operation")),
2479    };
2480    match other {
2481        Value::Bool(x) if *x != and => Ok(Value::Bool(*x)),
2482        Value::Bool(_) => Ok(Value::Event(event.clone())),
2483        Value::Event(other) => Ok(event.combine(other, if and { BinOp::And } else { BinOp::Or })?.value()),
2484        _ => Err(unsupported("combining this event with a probability")),
2485    }
2486}
2487
2488/// Scan through aggregate values too: symbolic numbers must never become
2489/// categorical keys or be compared by their internal identity as user data.
2490pub fn contains(v: &Value) -> bool {
2491    match v {
2492        Value::Analytic(_) | Value::Event(_) => return true,
2493        Value::List(_) | Value::Map(_) | Value::Bag(_) | Value::Record(_) | Value::Dist(_) | Value::Closure(_) => {}
2494        _ => return false,
2495    }
2496    let mut pending = vec![v];
2497    while let Some(v) = pending.pop() {
2498        match v {
2499            Value::Analytic(_) | Value::Event(_) => return true,
2500            Value::List(v) => pending.extend(v.iter()),
2501            Value::Map(v) => pending.extend(v.iter().flat_map(|(k, v)| [k, v])),
2502            Value::Bag(v) => pending.extend(v.keys()),
2503            Value::Record(v) => pending.extend(v.fields.iter().map(|(_, v)| v)),
2504            Value::Dist(d) => pending.extend(d.outcomes.iter().map(|(v, _)| v)),
2505            Value::Closure(c) => pending.extend(c.captured.iter()),
2506            _ => {}
2507        }
2508    }
2509    false
2510}
2511
2512/// Whether a value holds an outcome of a draw, other than events.
2513pub fn has_outcomes(v: &Value) -> bool {
2514    let mut pending = vec![v];
2515    while let Some(v) = pending.pop() {
2516        match v {
2517            Value::Analytic(_) => return true,
2518            Value::List(v) => pending.extend(v.iter()),
2519            Value::Map(v) => pending.extend(v.values()),
2520            Value::Record(v) => pending.extend(v.fields.iter().map(|(_, v)| v)),
2521            Value::Dist(d) => pending.extend(d.outcomes.iter().map(|(v, _)| v)),
2522            Value::Closure(c) => pending.extend(c.captured.iter()),
2523            _ => {}
2524        }
2525    }
2526    false
2527}
2528
2529/// A value as reports keep it: each outcome of a draw by its marginal alone,
2530/// without the identity of its draws, so that equal laws merge.
2531pub fn marginal(v: &Value) -> Value {
2532    if !contains(v) {
2533        return v.clone();
2534    }
2535    let bare = |a: &Analytic| Analytic {
2536        id: 0,
2537        form: None,
2538        ..a.clone()
2539    };
2540    match v {
2541        Value::Analytic(a) => Value::Analytic(Arc::new(bare(a))),
2542        Value::Event(e) => Value::Event(Arc::new(Event {
2543            draws: e.draws.iter().map(bare).collect(),
2544            boxes: e.boxes.clone(),
2545        })),
2546        Value::List(xs) => Value::list(xs.iter().map(marginal).collect()),
2547        Value::Record(r) => crate::ops::make_record(
2548            r.ty.clone(),
2549            r.fields.iter().map(|(k, v)| (k.clone(), marginal(v))).collect(),
2550        ),
2551        Value::Map(xs) => Value::map(xs.iter().map(|(k, v)| (k.clone(), marginal(v))).collect()),
2552        Value::Dist(d) => {
2553            crate::dist::Dist::from_pairs(d.outcomes.iter().map(|(v, p)| (marginal(v), *p)).collect(), d.missing)
2554                .into_value()
2555        }
2556        _ => v.clone(),
2557    }
2558}
2559
2560/// The same value with its latents renamed by `map`.
2561pub fn renamed(v: &Value, map: &BTreeMap<u64, u64>) -> Value {
2562    if !contains(v) {
2563        return v.clone();
2564    }
2565    match v {
2566        Value::Analytic(a) => Value::Analytic(Arc::new(a.renamed(map))),
2567        Value::Event(e) => Value::Event(Arc::new(Event {
2568            draws: e.draws.iter().map(|d| d.renamed(map)).collect(),
2569            boxes: e.boxes.clone(),
2570        })),
2571        Value::List(xs) => Value::list(xs.iter().map(|v| renamed(v, map)).collect()),
2572        Value::Record(r) => crate::ops::make_record(
2573            r.ty.clone(),
2574            r.fields.iter().map(|(k, v)| (k.clone(), renamed(v, map))).collect(),
2575        ),
2576        Value::Map(xs) => Value::map(xs.iter().map(|(k, v)| (k.clone(), renamed(v, map))).collect()),
2577        Value::Closure(f) => Value::Closure(Arc::new(Closure {
2578            func: f.func,
2579            captured: f.captured.iter().map(|v| renamed(v, map)).collect(),
2580        })),
2581        Value::Dist(d) => crate::dist::Dist::from_pairs(
2582            d.outcomes.iter().map(|(v, p)| (renamed(v, map), *p)).collect(),
2583            d.missing,
2584        )
2585        .into_value(),
2586        _ => v.clone(),
2587    }
2588}
2589
2590/// Latents reachable through an immutable value (including closure captures).
2591pub fn collect_ids(v: &Value, ids: &mut std::collections::BTreeSet<u64>) {
2592    let mut pending = vec![v];
2593    while let Some(v) = pending.pop() {
2594        match v {
2595            Value::Analytic(a) => a.collect_ids(ids),
2596            Value::Event(e) => e.draws.iter().for_each(|d| d.collect_ids(ids)),
2597            Value::List(v) => pending.extend(v.iter()),
2598            Value::Map(v) => pending.extend(v.iter().flat_map(|(k, v)| [k, v])),
2599            Value::Bag(v) => pending.extend(v.keys()),
2600            Value::Record(v) => pending.extend(v.fields.iter().map(|(_, v)| v)),
2601            Value::Dist(d) => pending.extend(d.outcomes.iter().map(|(v, _)| v)),
2602            Value::Closure(c) => pending.extend(c.captured.iter()),
2603            _ => {}
2604        }
2605    }
2606}
2607
2608/// Adds the latents that the forms of those in `ids` are made of, which
2609/// the world must keep knowing about too.
2610pub fn collect_form_ids(context: &Constraints, ids: &mut std::collections::BTreeSet<u64>) {
2611    let mut pending: Vec<u64> = ids.iter().copied().collect();
2612    while let Some(id) = pending.pop() {
2613        if let Some(Latent { form: Some(f), .. }) = context.get(&id) {
2614            for t in &f.terms {
2615                if ids.insert(t.id) {
2616                    pending.push(t.id);
2617                }
2618            }
2619        }
2620    }
2621}
2622
2623/// Snapshot a value under this world's posterior without mutating aliases.
2624pub fn resolve(v: &Value, context: &Constraints, budget: &mut Budget) -> OpResult<Value> {
2625    if context.is_empty() || !contains(v) {
2626        return Ok(v.clone());
2627    }
2628    resolve_at(v, context, budget, 0)
2629}
2630
2631/// The linear form of independent latents that `form` is in a world: each
2632/// latent that an observation of several replaced by its form, and each
2633/// one's distribution there. `None` if that's `form` itself. Unsupported
2634/// where a latent is restricted, since the form isn't normal then.
2635fn expand(form: &Linear, c: &Constraints, b: &mut Budget, depth: usize) -> OpResult<Option<Linear>> {
2636    if depth > 64 {
2637        return Err(OpError::limit("analytic value nesting exceeds the limit of 64"));
2638    }
2639    b.work(form.terms.len() as u64)?;
2640    let mut changed = false;
2641    let mut constant = form.constant;
2642    let mut terms = Vec::with_capacity(form.terms.len());
2643    for t in &form.terms {
2644        match c.get(&t.id) {
2645            None => terms.push(*t),
2646            Some(Latent { form: Some(sub), .. }) => {
2647                changed = true;
2648                let sub = match expand(sub, c, b, depth + 1)? {
2649                    Some(e) => e,
2650                    None => (**sub).clone(),
2651                };
2652                constant += t.coef * sub.constant;
2653                terms.extend(sub.terms.iter().map(|u| Term {
2654                    coef: t.coef * u.coef,
2655                    ..*u
2656                }));
2657            }
2658            Some(l) if l.domain.is_full() && matches!(l.family, Family::Normal { .. }) => {
2659                changed |= l.family != t.family;
2660                terms.push(Term { family: l.family, ..*t });
2661            }
2662            Some(_) => return Err(different_draws()),
2663        }
2664    }
2665    if !changed {
2666        return Ok(None);
2667    }
2668    Linear::from_terms(constant, terms).map(Some)
2669}
2670
2671/// `x`, whose axis the world now knows to be `e`, as a function of `e`'s
2672/// axis, and how its coordinates moved.
2673fn onto(x: &Analytic, e: Linear) -> OpResult<(Analytic, Transport)> {
2674    let (id, form, to, map) = match e.terms[..] {
2675        [] => return Err(unsupported("a draw that an observation made constant")),
2676        [t] => (
2677            t.id,
2678            None,
2679            t.family,
2680            Affine {
2681                scale: t.coef,
2682                offset: e.constant,
2683            },
2684        ),
2685        _ => {
2686            let (axis, map) = e.axis();
2687            let to = axis.normal()?;
2688            (0, Some(axis), to, map)
2689        }
2690    };
2691    let t = Transport {
2692        from: x.family,
2693        to,
2694        map,
2695    };
2696    Ok((x.rebased(id, form, &t), t))
2697}
2698
2699/// An outcome as a world reads it, and how its coordinates moved, if they
2700/// did.
2701fn read(x: &Analytic, c: &Constraints, b: &mut Budget) -> OpResult<(Analytic, Option<Transport>)> {
2702    if let Some(form) = &x.form {
2703        return Ok(match expand(form, c, b, 0)? {
2704            None => (x.clone(), None),
2705            Some(e) => {
2706                let (x, t) = onto(x, e)?;
2707                (x, Some(t))
2708            }
2709        });
2710    }
2711    Ok(match c.get(&x.id) {
2712        Some(Latent { form: Some(sub), .. }) => {
2713            let e = expand(sub, c, b, 0)?.unwrap_or_else(|| (**sub).clone());
2714            let (x, t) = onto(x, e)?;
2715            (x, Some(t))
2716        }
2717        Some(l) => {
2718            let mut y = x.moved(&l.family);
2719            y.domain = y.domain.intersect(&l.domain);
2720            y.weight = l.weight.clone();
2721            let t = (l.family != x.family).then(|| Transport::moved(x.family, l.family));
2722            (y, t)
2723        }
2724        None => (x.clone(), None),
2725    })
2726}
2727
2728fn resolve_at(v: &Value, c: &Constraints, b: &mut Budget, depth: usize) -> OpResult<Value> {
2729    b.work(1)?;
2730    if depth > 64 {
2731        return Err(OpError::limit("analytic value nesting exceeds the limit of 64"));
2732    }
2733    Ok(match v {
2734        Value::Analytic(x) => read(x, c, b)?.0.value()?,
2735        Value::Event(e) => {
2736            let mut draws = Vec::with_capacity(e.draws.len());
2737            let mut boxes = e.boxes.clone();
2738            let mut moved = false;
2739            for (i, d) in e.draws.iter().enumerate() {
2740                let (d, t) = read(d, c, b)?;
2741                if let Some(t) = t {
2742                    moved = true;
2743                    for part in &mut boxes {
2744                        part[i] = t.domain(&part[i]);
2745                    }
2746                }
2747                draws.push(d);
2748            }
2749            // An observation can make independent draws depend on each other.
2750            if moved && draws.len() > 1 {
2751                for (i, d) in draws.iter().enumerate() {
2752                    if draws[..i].iter().any(|x| x.same_axis(d) || x.shares_latents(d)) {
2753                        return Err(unsupported(
2754                            "an event of draws that an observation made depend on each other",
2755                        ));
2756                    }
2757                }
2758            }
2759            Event { draws, boxes }.simplified().value()
2760        }
2761        Value::List(xs) => Value::list(
2762            xs.iter()
2763                .map(|v| resolve_at(v, c, b, depth + 1))
2764                .collect::<OpResult<_>>()?,
2765        ),
2766        Value::Record(r) => crate::ops::make_record(
2767            r.ty.clone(),
2768            r.fields
2769                .iter()
2770                .map(|(k, v)| Ok((k.clone(), resolve_at(v, c, b, depth + 1)?)))
2771                .collect::<OpResult<_>>()?,
2772        ),
2773        Value::Map(xs) => Value::map(
2774            xs.iter()
2775                .map(|(k, v)| Ok((k.clone(), resolve_at(v, c, b, depth + 1)?)))
2776                .collect::<OpResult<_>>()?,
2777        ),
2778        Value::Closure(f) => Value::Closure(Arc::new(Closure {
2779            func: f.func,
2780            captured: f
2781                .captured
2782                .iter()
2783                .map(|v| resolve_at(v, c, b, depth + 1))
2784                .collect::<OpResult<_>>()?,
2785        })),
2786        Value::Dist(d) => {
2787            let pairs = d
2788                .outcomes
2789                .iter()
2790                .map(|(v, p)| Ok((resolve_at(v, c, b, depth + 1)?, *p)))
2791                .collect::<OpResult<_>>()?;
2792            crate::ops::combine(pairs, d.missing, b)?
2793        }
2794        // Analytic keys and bags are rejected when constructed.
2795        _ => v.clone(),
2796    })
2797}
2798
2799#[cfg(test)]
2800mod tests {
2801    use super::*;
2802
2803    fn uniform(lo: f64, hi: f64) -> Analytic {
2804        Analytic::new(1, Family::Uniform { lo, hi })
2805    }
2806    fn budget() -> Budget {
2807        Budget::unlimited()
2808    }
2809    fn analytic(v: Value) -> Analytic {
2810        match v {
2811            Value::Analytic(a) => (*a).clone(),
2812            v => panic!("expected an analytic value, found {v:?}"),
2813        }
2814    }
2815    fn close(a: f64, b: f64) {
2816        assert!((a - b).abs() < 1e-12, "{a} vs {b}");
2817    }
2818
2819    #[test]
2820    fn overlay_refines_both_tilings() {
2821        let a = [(0.25, 'a'), (1.0, 'b')];
2822        let b = [(0.5, 1), (0.75, 2), (1.0, 3)];
2823        assert_eq!(
2824            overlay(&a, &b),
2825            vec![(0.25, ('a', 1)), (0.5, ('b', 1)), (0.75, ('b', 2)), (1.0, ('b', 3))]
2826        );
2827    }
2828
2829    #[test]
2830    fn abs_of_a_symmetric_uniform_is_uniform() {
2831        // x ~ uniform(-1, 1): |x| ~ uniform(0, 1).
2832        let y = analytic(abs(&uniform(-1.0, 1.0), &mut budget()).unwrap());
2833        assert_eq!(y.pieces.len(), 2);
2834        close(y.cdf(0.3), 0.3);
2835        close(y.quantile(0.05), 0.05);
2836        close(y.quantile(0.5), 0.5);
2837        let (mean, variance) = y.moments().unwrap();
2838        close(mean, 0.5);
2839        close(variance, 1.0 / 12.0);
2840        close(y.pdf(0.4), 1.0);
2841        assert!(!y.has_atoms());
2842    }
2843
2844    #[test]
2845    fn min_with_a_number_has_an_atom() {
2846        // x ~ uniform(0, 2): min(x, 1) is x below 1, and 1 with probability 1/2.
2847        let x = Value::Analytic(Arc::new(uniform(0.0, 2.0)));
2848        let y = analytic(extreme(&[x, Value::Float(1.0)], false, &mut budget()).unwrap());
2849        assert!(y.has_atoms());
2850        close(y.cdf(0.5), 0.25);
2851        close(y.cdf(1.0), 1.0);
2852        close(y.quantile(0.25), 0.5);
2853        assert_eq!(y.quantile(0.5), 1.0);
2854        assert_eq!(y.quantile(0.95), 1.0);
2855        let (mean, variance) = y.moments().unwrap();
2856        close(mean, 0.75);
2857        close(variance, 2.0 / 3.0 - 0.75 * 0.75);
2858    }
2859
2860    #[test]
2861    fn clamp_has_two_atoms() {
2862        // x ~ uniform(-1, 2): clamp(x, 0, 1) is 0 and 1 with 1/3 each.
2863        let y = analytic(
2864            clamp(
2865                &uniform(-1.0, 2.0),
2866                &Value::Float(0.0),
2867                &Value::Float(1.0),
2868                &mut budget(),
2869            )
2870            .unwrap(),
2871        );
2872        assert_eq!(y.pieces.len(), 3);
2873        close(y.cdf(0.0), 1.0 / 3.0);
2874        close(y.cdf(0.5), 0.5);
2875        close(y.moments().unwrap().0, 0.5);
2876        assert_eq!(y.quantile(0.2), 0.0);
2877        assert_eq!(y.quantile(0.9), 1.0);
2878    }
2879
2880    #[test]
2881    fn comparisons_of_pieces_are_events_of_the_latent() {
2882        // |x| < 0.5 for x ~ uniform(-1, 1) is -0.5 < x < 0.5.
2883        let x = uniform(-1.0, 1.0);
2884        let y = abs(&x, &mut budget()).unwrap();
2885        let Value::Event(e) = binary(BinOp::Lt, &y, &Value::Float(0.5), &mut budget()).unwrap() else {
2886            panic!("expected an event");
2887        };
2888        assert_eq!(e.boxes, vec![vec![Domain(vec![(0.25, 0.75)])]]);
2889        assert_eq!(e.draws[0].affine(), Some(Affine::IDENTITY));
2890        // An atom equals its value with its probability.
2891        let m = extreme(&[Value::Analytic(Arc::new(x)), Value::Float(0.0)], true, &mut budget()).unwrap();
2892        let Value::Event(e) = binary(BinOp::Eq, &m, &Value::Float(0.0), &mut budget()).unwrap() else {
2893            panic!("expected an event");
2894        };
2895        close(e.probability(), 0.5);
2896    }
2897
2898    #[test]
2899    fn pieces_outside_the_domain_go() {
2900        // Once x > 0, |x| is x.
2901        let mut x = uniform(-1.0, 1.0);
2902        x.domain = Domain(vec![(0.5, 1.0)]);
2903        let y = analytic(abs(&x, &mut budget()).unwrap());
2904        assert_eq!(y.pieces, vec![(1.0, Fun::IDENTITY)]);
2905        // And |x| - x is 0.
2906        let d = binary(
2907            BinOp::Sub,
2908            &Value::Analytic(Arc::new(y)),
2909            &Value::Analytic(Arc::new(x)),
2910            &mut budget(),
2911        )
2912        .unwrap();
2913        assert!(matches!(d, Value::Float(z) if z == 0.0), "{d:?}");
2914    }
2915
2916    #[test]
2917    fn the_same_draw_combines_piece_by_piece() {
2918        // |x| + x is 0 below zero and 2x above.
2919        let x = uniform(-1.0, 1.0);
2920        let y = abs(&x, &mut budget()).unwrap();
2921        let s = analytic(binary(BinOp::Add, &y, &Value::Analytic(Arc::new(x)), &mut budget()).unwrap());
2922        assert_eq!(
2923            s.pieces,
2924            vec![
2925                (0.5, Fun::constant(0.0)),
2926                (
2927                    1.0,
2928                    Fun::affine(Affine {
2929                        scale: 2.0,
2930                        offset: 0.0
2931                    })
2932                )
2933            ]
2934        );
2935        close(s.moments().unwrap().0, 0.5);
2936        // A constant piece times anything is affine; elsewhere it isn't.
2937        let product = binary(
2938            BinOp::Mul,
2939            &Value::Analytic(Arc::new(s.clone())),
2940            &Value::Float(3.0),
2941            &mut budget(),
2942        );
2943        assert!(product.is_ok());
2944        // Affine pieces multiply into squares: 0 below zero, 4x² above.
2945        let s = Value::Analytic(Arc::new(s));
2946        let square = analytic(binary(BinOp::Mul, &s, &s, &mut budget()).unwrap());
2947        close(square.moments().unwrap().0, 2.0 / 3.0);
2948        // A cube isn't one of the functions a piece can be.
2949        let cube = binary(BinOp::Mul, &Value::Analytic(Arc::new(square)), &s, &mut budget());
2950        assert!(cube.unwrap_err().message.contains("nonlinear arithmetic"));
2951    }
2952}