1use 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#[derive(Clone, Debug, PartialEq)]
32pub struct Latent {
33 pub family: Family,
34 pub domain: Domain,
35 pub form: Option<Arc<Linear>>,
36 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#[derive(Clone, Copy, Debug, PartialEq)]
53pub struct Term {
54 pub id: u64,
55 pub coef: f64,
56 pub family: Family,
57}
58
59#[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 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 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 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 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 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 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 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
195fn moved(c: f64, from: &Family, to: &Family) -> f64 {
197 to.cdf(from.quantile(c)).clamp(0.0, 1.0)
198}
199
200#[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 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 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
297fn 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#[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 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 fn inverse(self, y: f64) -> f64 {
341 (y - self.offset) / self.scale
342 }
343 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#[derive(Clone, Copy, Debug)]
356struct Transport {
357 from: Family,
358 to: Family,
359 map: Affine,
360}
361impl Transport {
362 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 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 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 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#[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 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 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 fn increasing(self, side: f64) -> bool {
465 match self {
466 Kernel::Square => side >= 0.0,
467 Kernel::Recip => false,
468 _ => true,
469 }
470 }
471 fn turns(self) -> bool {
473 matches!(self, Kernel::Square | Kernel::Recip)
474 }
475}
476
477#[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 fn of(kernel: Kernel, inner: Affine) -> Fun {
501 Fun {
502 outer: Affine::IDENTITY,
503 kernel,
504 inner,
505 }
506 }
507 pub fn as_affine(self) -> Option<Affine> {
509 (self.kernel == Kernel::Id).then_some(self.outer)
510 }
511 pub fn as_constant(self) -> Option<f64> {
513 (self.outer.scale == 0.0).then_some(self.outer.offset)
514 }
515 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 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 fn side(self, family: &Family, start: f64, end: f64) -> f64 {
552 self.inner.at(family.quantile((start + end) / 2.0))
553 }
554 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 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 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 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 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 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
666pub type Pieces = Vec<(f64, Fun)>;
670
671fn 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 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
687fn 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
696fn 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
709fn 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 pub int: bool,
739 pub form: Option<Arc<Linear>>,
742 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 pub(crate) fn latent(&self) -> Self {
759 self.with(vec![(1.0, Fun::IDENTITY)])
760 }
761 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 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 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 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 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 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 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 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 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 fn with_type(&self, pieces: Pieces, int: bool) -> Self {
890 Self {
891 int,
892 ..self.with(pieces)
893 }
894 }
895 pub fn affine(&self) -> Option<Affine> {
897 match self.pieces[..] {
898 [(_, f)] => f.as_affine(),
899 _ => None,
900 }
901 }
902 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 pub fn is_discrete(&self) -> bool {
934 self.pieces.len() > 1 && self.segments().iter().all(|(_, _, f)| f.as_constant().is_some())
935 }
936 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(®ion) / mass;
958 (number(c, self.int), region, share)
959 })
960 .collect()
961 }
962 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 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 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 pub fn has_atoms(&self) -> bool {
1026 self.pieces.len() > 1 && self.segments().iter().any(|(_, _, f)| f.as_constant().is_some())
1027 }
1028 pub fn measure(&self, d: &Domain) -> f64 {
1031 match &self.weight {
1032 None => d.mass(),
1033 Some(l) => l.mass(d),
1034 }
1035 }
1036 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 pub fn pdf(&self, x: f64) -> f64 {
1055 let mass = self.measure(&self.domain);
1056 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 fn piecewise_quantile(&self, q: f64) -> f64 {
1104 let family = &self.family;
1105 let segments = self.segments();
1106 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 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 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 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 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 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
1284fn 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
1321mod 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 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 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 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 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 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 _ 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 fn uniform_power(y0: f64, y1: f64, p: f64) -> Option<f64> {
1400 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 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 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 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 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 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 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 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 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 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#[derive(Clone, Debug)]
1541pub struct Event {
1542 pub draws: Vec<Analytic>,
1545 pub boxes: Vec<Vec<Domain>>,
1546}
1547
1548const MAX_BOXES: usize = 10_000;
1551
1552fn 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
1571fn 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 pub fn complement(&self) -> OpResult<Event> {
1625 Ok(Event {
1626 draws: self.draws.clone(),
1627 boxes: self.side(false)?,
1628 })
1629 }
1630 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 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 pub fn restrictable(&self) -> bool {
1652 self.draws.iter().all(|d| d.form.is_none())
1653 }
1654 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 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(¬_mine, ¬_theirs)?);
1723 boxes
1724 }
1725 };
1726 Ok(Event { draws, boxes }.simplified())
1727 }
1728 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
1738pub 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 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 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
1803pub 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
1823pub 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
1841pub type Observed = (f64, Vec<(u64, Latent)>);
1844
1845pub 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 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
1924pub 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
1932enum 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
1950fn 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
1971fn 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
2015fn 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
2043fn scaled(f: Fun, c: f64) -> Fun {
2045 if c == 0.0 { Fun::constant(0.0) } else { f.then(c, 0.0) }
2046}
2047
2048fn 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 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
2089const MAX_PIECES: usize = 100_000;
2092
2093fn 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 let int = matches!(op, Add | Sub | Mul) && is_int(a) && is_int(b);
2144 x.with_type(pieces, int).value()
2145}
2146
2147fn 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
2156fn 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
2177fn 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 (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
2200fn compose(kernel: Kernel, f: Fun, family: &Family, start: f64, end: f64) -> OpResult<Fun> {
2202 if let Some(c) = f.as_constant() {
2203 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 (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 (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 (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
2239pub 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 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
2274fn 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 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
2313fn 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
2321fn is_int(v: &Value) -> bool {
2323 match v {
2324 Value::Int(_) => true,
2325 Value::Analytic(x) => x.int,
2326 _ => false,
2327 }
2328}
2329
2330#[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 Rounding::Round => v.round(),
2346 }
2347 }
2348 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
2361pub 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 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 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 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
2426pub 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
2460pub 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
2488pub 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
2512pub 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
2529pub 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
2560pub 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
2590pub 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
2608pub 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
2623pub 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
2631fn 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
2671fn 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
2699fn 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 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 _ => 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 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 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 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 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 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 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 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 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 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 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 let cube = binary(BinOp::Mul, &Value::Analytic(Arc::new(square)), &s, &mut budget());
2950 assert!(cube.unwrap_err().message.contains("nonlinear arithmetic"));
2951 }
2952}