fugue_evo/inference/
prior.rs

1//! Priors over genomes as probabilistic programs
2//!
3//! A [`GenomePrior`] is *the* load-bearing abstraction of the inference layer:
4//! instead of a closed enum of built-in priors, the prior over genomes is an
5//! arbitrary fugue [`Model`] written by the user (or one of the constructors
6//! below). Running it under a `PriorHandler` both draws `p(x)` and accumulates
7//! `log_prior`; scoring an existing genome's trace against it recovers the
8//! genuine prior density — there is no hand-written density code anywhere in
9//! this layer.
10//!
11//! The model returns the **decoded genome** `G`, not a bare vector: the model's
12//! return value *is* the decode, which is what lets the SMC layer recover a
13//! genome from a bare particle trace by replay (see
14//! [`crate::inference::smc::EvolutionPosterior`]).
15
16use fugue::{addr, plate, sample, Bernoulli, Categorical, Model, ModelExt, Normal, Uniform};
17
18use crate::error::GenomeError;
19use crate::genome::bit_string::BitString;
20use crate::genome::bounds::MultiBounds;
21use crate::genome::permutation::Permutation;
22use crate::genome::real_vector::RealVector;
23use crate::genome::trace_genome::TraceGenome;
24use crate::genome::traits::{BinaryGenome, PermutationGenome, RealValuedGenome};
25
26/// A prior distribution over genomes, expressed as a probabilistic program.
27///
28/// The program must sample the genome's canonical trace sites (the same
29/// addresses [`TraceGenome::to_trace`] writes — `gene#i`, `bit#i`, `perm#i`,
30/// …) and return the assembled genome. Anything expressible as a fugue
31/// `Model` is a valid prior: correlated coordinates, hierarchical scales,
32/// variable-length genomes (a length site followed by that many coordinate
33/// sites — fugue's MH treats the resulting births/deaths as reversible-jump
34/// moves with no extra code here).
35pub trait GenomePrior: Clone + Send + Sync + 'static {
36    /// The genome type this prior generates.
37    type Genome: TraceGenome;
38
39    /// The generative program `p(x)`.
40    fn model(&self) -> Model<Self::Genome>;
41
42    /// Encode a genome as a trace **under this prior's address scheme** — the
43    /// inverse direction of running [`Self::model`]. The default delegates to
44    /// the genome's canonical [`TraceGenome::to_trace`] encoding, which is
45    /// correct whenever the prior samples exactly the canonical sites (all the
46    /// vector priors here). Priors with their own generative scheme — e.g.
47    /// [`ArithmeticGrammarPrior`](super::grammar::ArithmeticGrammarPrior)'s
48    /// tree-path grammar — override this so that scoring, weighted traces, and
49    /// chain warm-starts work for any genome the prior can express.
50    fn trace_of(&self, genome: &Self::Genome) -> fugue::Trace {
51        genome.to_trace()
52    }
53
54    /// Check that `genome` has the shape this prior generates, **before** its
55    /// encoding is replayed through a program. The vector priors here report
56    /// [`GenomeError::DimensionMismatch`] for a genome whose length differs
57    /// from the prior's dimension — a shorter one would leave the program
58    /// visiting sites the encoding lacks, a longer one would carry sites the
59    /// program never visits (previously replay panicked on the former and
60    /// silently truncated the latter). The default accepts everything; priors
61    /// whose `trace_of` is total over their genome type (the grammar prior)
62    /// need nothing more, because structural mismatches are caught by the
63    /// replay itself in [`EvolutionModel::score`](super::model::EvolutionModel::score).
64    fn validate(&self, genome: &Self::Genome) -> Result<(), GenomeError> {
65        let _ = genome;
66        Ok(())
67    }
68}
69
70fn check_dimension(expected: usize, actual: usize) -> Result<(), GenomeError> {
71    if expected == actual {
72        Ok(())
73    } else {
74        Err(GenomeError::DimensionMismatch { expected, actual })
75    }
76}
77
78/// Independent uniform prior over a bounded box (per-dimension `[min, max]`).
79///
80/// The support behavior formerly hand-coded in the old `Prior::UniformBounds`
81/// match (`−∞` outside the box) now falls out of scoring `Uniform` sites under
82/// replay: an out-of-bounds value has `log_prob = −∞` and any MH move onto it
83/// is rejected.
84#[derive(Clone, Debug)]
85pub struct UniformBoxPrior {
86    bounds: MultiBounds,
87}
88
89impl UniformBoxPrior {
90    /// Uniform prior over the given per-dimension bounds.
91    pub fn new(bounds: MultiBounds) -> Self {
92        Self { bounds }
93    }
94
95    /// The bounds of the box.
96    pub fn bounds(&self) -> &MultiBounds {
97        &self.bounds
98    }
99}
100
101impl GenomePrior for UniformBoxPrior {
102    type Genome = RealVector;
103
104    fn model(&self) -> Model<RealVector> {
105        let bounds = self.bounds.clone();
106        plate!(i in 0..bounds.dimension().max(1) => {
107            let (lo, hi) = match bounds.get(i) {
108                Some(b) if b.max > b.min => (b.min, b.max),
109                Some(b) => (b.min - 1e-9, b.min + 1e-9),
110                None => (-1.0, 1.0),
111            };
112            sample(addr!("gene", i), Uniform::new(lo, hi).expect("valid uniform prior bounds"))
113        })
114        .map(|genes| RealVector::from_genes(genes).expect("plate produced genes"))
115    }
116
117    fn validate(&self, genome: &RealVector) -> Result<(), GenomeError> {
118        check_dimension(self.bounds.dimension().max(1), genome.genes().len())
119    }
120}
121
122/// Independent Gaussian `N(mean, std²)` prior on every real coordinate.
123#[derive(Clone, Debug)]
124pub struct GaussianPrior {
125    mean: f64,
126    std: f64,
127    dim: usize,
128}
129
130impl GaussianPrior {
131    /// I.i.d. Gaussian prior with the given per-coordinate mean/std over `dim`
132    /// coordinates. `std` must be positive.
133    pub fn new(mean: f64, std: f64, dim: usize) -> Self {
134        assert!(std > 0.0, "Gaussian prior std must be > 0");
135        Self { mean, std, dim }
136    }
137}
138
139impl GenomePrior for GaussianPrior {
140    type Genome = RealVector;
141
142    fn model(&self) -> Model<RealVector> {
143        let (mean, std, dim) = (self.mean, self.std, self.dim.max(1));
144        plate!(i in 0..dim => {
145            sample(addr!("gene", i), Normal::new(mean, std).expect("valid Gaussian prior"))
146        })
147        .map(|genes| RealVector::from_genes(genes).expect("plate produced genes"))
148    }
149
150    fn validate(&self, genome: &RealVector) -> Result<(), GenomeError> {
151        check_dimension(self.dim.max(1), genome.genes().len())
152    }
153}
154
155/// Independent `Bernoulli(p)` prior on every bit of a [`BitString`].
156#[derive(Clone, Debug)]
157pub struct BitStringPrior {
158    p_one: f64,
159    len: usize,
160}
161
162impl BitStringPrior {
163    /// Prior over `len`-bit strings with per-bit probability `p_one` of a set
164    /// bit. `p_one` must lie in `(0, 1)` so every string has support.
165    pub fn new(p_one: f64, len: usize) -> Self {
166        assert!(
167            p_one > 0.0 && p_one < 1.0,
168            "BitStringPrior p_one must be in (0, 1)"
169        );
170        Self { p_one, len }
171    }
172
173    /// Uniform prior over `len`-bit strings (`p = 1/2` per bit).
174    pub fn uniform(len: usize) -> Self {
175        Self::new(0.5, len)
176    }
177}
178
179impl GenomePrior for BitStringPrior {
180    type Genome = BitString;
181
182    fn model(&self) -> Model<BitString> {
183        let (p, len) = (self.p_one, self.len.max(1));
184        plate!(i in 0..len => {
185            sample(addr!("bit", i), Bernoulli::new(p).expect("valid Bernoulli prior"))
186        })
187        .map(|bits| BitString::from_bits(bits).expect("plate produced bits"))
188    }
189
190    fn validate(&self, genome: &BitString) -> Result<(), GenomeError> {
191        check_dimension(self.len.max(1), genome.bits().len())
192    }
193}
194
195/// Fisher–Yates / Lehmer-code uniform prior over permutations of `0..n`.
196///
197/// Position `i` samples a **rank** at `perm#i`, uniform over the `n−i` values
198/// not yet used (a `Usize` in `0..n−i`); the rank sequence decodes to a
199/// permutation against the shrinking available-value list. This coincides
200/// site-for-site with [`Permutation::to_trace`]'s Lehmer encoding, so scoring
201/// an existing genome's trace under replay finds every site with matching
202/// semantics, and:
203///
204/// - every model execution decodes to a valid permutation (density exactly
205///   `1/n!`), and
206/// - under single-site MH, resampling one rank always decodes to a *different
207///   valid* permutation — the rank encoding is what makes single-site moves
208///   live (a raw value encoding would turn every single-site change into a
209///   duplicate and freeze the chain).
210#[derive(Clone, Debug)]
211pub struct PermutationPrior {
212    n: usize,
213}
214
215impl PermutationPrior {
216    /// Uniform prior over permutations of `0..n`.
217    pub fn new(n: usize) -> Self {
218        assert!(n > 0, "PermutationPrior needs n > 0");
219        Self { n }
220    }
221}
222
223impl GenomePrior for PermutationPrior {
224    type Genome = Permutation;
225
226    fn model(&self) -> Model<Permutation> {
227        let n = self.n;
228        fn rank_model(n: usize, i: usize, ranks: Vec<usize>) -> Model<Vec<usize>> {
229            if i == n {
230                return fugue::pure(ranks);
231            }
232            let k = n - i;
233            let probs = vec![1.0 / k as f64; k];
234            sample(
235                addr!("perm", i),
236                Categorical::new(probs).expect("valid categorical prior"),
237            )
238            .bind(move |r| {
239                let mut ranks = ranks;
240                ranks.push(r);
241                rank_model(n, i + 1, ranks)
242            })
243        }
244        rank_model(n, 0, Vec::with_capacity(n)).map(move |ranks| {
245            let mut available: Vec<usize> = (0..n).collect();
246            let perm: Vec<usize> = ranks.into_iter().map(|r| available.remove(r)).collect();
247            Permutation::from_permutation(perm).expect("Lehmer decode produced a permutation")
248        })
249    }
250
251    fn validate(&self, genome: &Permutation) -> Result<(), GenomeError> {
252        check_dimension(self.n, genome.permutation().len())
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259    use fugue::runtime::handler::run;
260    use fugue::runtime::interpreters::{PriorHandler, ScoreGivenTrace};
261    use fugue::Trace;
262    use rand::rngs::StdRng;
263    use rand::SeedableRng;
264
265    #[test]
266    fn test_gaussian_prior_draws_match_moments() {
267        let prior = GaussianPrior::new(0.0, 2.0, 1);
268        let mut rng = StdRng::seed_from_u64(99);
269        let xs: Vec<f64> = (0..5000)
270            .map(|_| {
271                let (g, _) = run(
272                    PriorHandler {
273                        rng: &mut rng,
274                        trace: Trace::default(),
275                    },
276                    prior.model(),
277                );
278                g.genes()[0]
279            })
280            .collect();
281        let mean = xs.iter().sum::<f64>() / xs.len() as f64;
282        let var = xs.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / xs.len() as f64;
283        assert!(mean.abs() < 0.2, "prior mean {}", mean);
284        assert!((var.sqrt() - 2.0).abs() < 0.2, "prior std {}", var.sqrt());
285    }
286
287    /// Replacement anchor for the deleted hand-written `log_prior_density`:
288    /// scoring a genome's trace under the prior model reproduces the analytic
289    /// Gaussian log-density.
290    #[test]
291    fn test_prior_model_log_prior_matches_analytic() {
292        let prior = GaussianPrior::new(1.0, 2.0, 3);
293        let g = RealVector::new(vec![0.5, 1.5, -2.0]);
294        let (_, scored) = run(
295            ScoreGivenTrace {
296                base: g.to_trace(),
297                trace: Trace::default(),
298            },
299            prior.model(),
300        );
301        let normal = Normal::new(1.0, 2.0).unwrap();
302        let analytic: f64 = g
303            .genes()
304            .iter()
305            .map(|x| fugue::Distribution::log_prob(&normal, x))
306            .sum();
307        assert!((scored.log_prior - analytic).abs() < 1e-12);
308    }
309
310    #[test]
311    fn test_uniform_prior_out_of_box_scores_neg_inf() {
312        let prior = UniformBoxPrior::new(MultiBounds::symmetric(1.0, 2));
313        let g = RealVector::new(vec![0.5, 5.0]); // second coordinate outside
314        let (_, scored) = run(
315            ScoreGivenTrace {
316                base: g.to_trace(),
317                trace: Trace::default(),
318            },
319            prior.model(),
320        );
321        assert_eq!(scored.log_prior, f64::NEG_INFINITY);
322    }
323
324    /// EV-N3: the vector priors reject genomes of the wrong dimension before
325    /// any replay, with the exact expected/actual pair.
326    #[test]
327    fn test_vector_priors_validate_dimension() {
328        let g = GaussianPrior::new(0.0, 1.0, 3);
329        assert_eq!(g.validate(&RealVector::new(vec![0.0; 3])), Ok(()));
330        assert_eq!(
331            g.validate(&RealVector::new(vec![0.0; 2])),
332            Err(GenomeError::DimensionMismatch {
333                expected: 3,
334                actual: 2
335            })
336        );
337        let u = UniformBoxPrior::new(MultiBounds::symmetric(1.0, 2));
338        assert_eq!(
339            u.validate(&RealVector::new(vec![0.0; 5])),
340            Err(GenomeError::DimensionMismatch {
341                expected: 2,
342                actual: 5
343            })
344        );
345        let b = BitStringPrior::uniform(4);
346        assert!(b.validate(&BitString::new(vec![true; 4])).is_ok());
347        assert!(b.validate(&BitString::new(vec![true; 3])).is_err());
348        let p = PermutationPrior::new(4);
349        assert!(p.validate(&Permutation::new(vec![0, 1, 2, 3])).is_ok());
350        assert!(p.validate(&Permutation::new(vec![0, 1, 2])).is_err());
351    }
352
353    #[test]
354    fn test_permutation_prior_generates_valid_permutations() {
355        let prior = PermutationPrior::new(6);
356        let mut rng = StdRng::seed_from_u64(7);
357        for _ in 0..50 {
358            let (p, trace) = run(
359                PriorHandler {
360                    rng: &mut rng,
361                    trace: Trace::default(),
362                },
363                prior.model(),
364            );
365            assert!(p.is_valid_permutation());
366            // Density of any permutation is 1/n!.
367            let expected = -(720.0f64).ln(); // ln(1/6!)
368            assert!((trace.log_prior - expected).abs() < 1e-9);
369            // The trace encoding coincides with Permutation::to_trace.
370            let canonical = p.to_trace();
371            for (addr, choice) in &canonical.choices {
372                assert_eq!(trace.choices[addr].value, choice.value);
373            }
374        }
375    }
376
377    #[test]
378    fn test_bitstring_prior_scores_canonical_trace() {
379        let prior = BitStringPrior::uniform(4);
380        let g = BitString::from_bits(vec![true, false, true, true]).unwrap();
381        let (decoded, scored) = run(
382            ScoreGivenTrace {
383                base: g.to_trace(),
384                trace: Trace::default(),
385            },
386            prior.model(),
387        );
388        assert_eq!(decoded.bits(), g.bits());
389        assert!((scored.log_prior - 4.0 * (0.5f64).ln()).abs() < 1e-12);
390    }
391}
fugue_evo/inference/prior.rs

fugue_evo/inference/
prior.rs
