fugue_evo/inference/
mh.rs

1//! Metropolis–Hastings over the Boltzmann target, delegated to fugue
2//!
3//! The old `EvolutionStep` hand-rolled its proposal (and only ever perturbed
4//! `F64` choices, so BitString/Permutation chains silently never moved). This
5//! wrapper deletes all of that: one transition is one call into
6//! [`fugue::adaptive_single_site_mh_cached`], which picks the target site
7//! uniformly over **all** sites and dispatches the proposal by value type
8//! (for `F64` from the site's [`fugue::Support`]: Gaussian walk on the reals,
9//! log-space walk on the positives, reflected walk on a bounded interval;
10//! flip for `Bool`, reflected discrete walk for `U64`, prior-resample for
11//! `Usize`, integer walk for `I64`), including the reversible-jump
12//! corrections for structure-changing models. Per-address
13//! [`SiteProposal`] overrides registered with
14//! [`EvolutionChain::override_site`] are honoured by every entry point —
15//! [`EvolutionChain::step`], [`EvolutionChain::step_scored`] and
16//! [`EvolutionChain::run_chain`] alike.
17//!
18//! A transition costs **one** model execution (the proposal): the current
19//! state's log-density and per-site densities are read from the scored trace
20//! the caller holds, which is why `current` must be a trace produced by
21//! [`EvolutionChain::init`], [`EvolutionChain::init_from`] or a previous
22//! `step` (see [`EvolutionChain::step`]).
23
24use std::collections::HashMap;
25
26use fugue::inference::mcmc_utils::DiminishingAdaptation;
27use fugue::runtime::handler::run;
28use fugue::runtime::interpreters::{PriorHandler, ScoreGivenTrace};
29use fugue::{
30    adaptive_mcmc_chain_with_overrides, adaptive_single_site_mh_cached, Address, SiteProposal,
31    Trace,
32};
33use rand::Rng;
34
35use super::likelihood::GenomeLikelihood;
36use super::model::EvolutionModel;
37use super::prior::GenomePrior;
38use crate::error::GenomeError;
39
40fn finite_state(trace: Trace) -> Result<Trace, GenomeError> {
41    if trace.total_log_weight().is_finite() {
42        Ok(trace)
43    } else {
44        Err(GenomeError::ConstraintViolation(
45            "genome is outside the prior's support (target log-density is not finite)".to_string(),
46        ))
47    }
48}
49
50/// An MH chain over the fixed-β Boltzmann target `π_β ∝ p(x)·exp(β·f(x))`.
51pub struct EvolutionChain<P, L>
52where
53    P: GenomePrior,
54    L: GenomeLikelihood<P::Genome>,
55{
56    model: EvolutionModel<P, L>,
57    adaptation: DiminishingAdaptation,
58    overrides: HashMap<Address, SiteProposal>,
59}
60
61impl<P, L> EvolutionChain<P, L>
62where
63    P: GenomePrior,
64    L: GenomeLikelihood<P::Genome>,
65{
66    /// Create a chain over the model's fixed-β target.
67    pub fn new(model: EvolutionModel<P, L>) -> Self {
68        Self {
69            model,
70            adaptation: DiminishingAdaptation::new(0.44, 0.7),
71            overrides: HashMap::new(),
72        }
73    }
74
75    /// Set the adaptation's target acceptance rate (default 0.44).
76    pub fn target_rate(mut self, rate: f64) -> Self {
77        self.adaptation = DiminishingAdaptation::new(rate, 0.7);
78        self
79    }
80
81    /// Force a specific `f64` proposal for one address (e.g.
82    /// `SiteProposal::Reflect { lower, upper }` for a bounded coordinate, or
83    /// `SiteProposal::PriorResample` for an independence move).
84    ///
85    /// Honoured by [`Self::step`], [`Self::step_scored`] and
86    /// [`Self::run_chain`]. Since fugue selects the default `f64` proposal
87    /// from the site's declared [`fugue::Support`] — a `Uniform` site already
88    /// gets a reflected walk at its own bounds — an override is only needed
89    /// to *change* that default (a narrower reflection interval, a log-space
90    /// walk on a `Normal` site known to be positive, …).
91    pub fn override_site(mut self, addr: Address, proposal: SiteProposal) -> Self {
92        self.overrides.insert(addr, proposal);
93        self
94    }
95
96    /// The registered per-address proposal overrides.
97    pub fn overrides(&self) -> &HashMap<Address, SiteProposal> {
98        &self.overrides
99    }
100
101    /// The underlying model.
102    pub fn model(&self) -> &EvolutionModel<P, L> {
103        &self.model
104    }
105
106    /// Draw an initial state: a prior sample's fully-scored trace (latent
107    /// likelihood sites included).
108    pub fn init<R: Rng>(&self, rng: &mut R) -> Trace {
109        let (_g, trace) = run(
110            PriorHandler {
111                rng,
112                trace: Trace::default(),
113            },
114            (self.model.target_model())(),
115        );
116        trace
117    }
118
119    /// Warm-start the chain from a given genome: encode it under the model's
120    /// prior ([`GenomePrior::trace_of`]) and score it through the target
121    /// program. Works for any prior — including grammar priors over trees —
122    /// so a classic GA/GP result can seed an inference chain. Returns `None`
123    /// if the genome is outside the prior's support (its target density is
124    /// `−∞`, which can never be left by an MH chain) **or** cannot be scored
125    /// from its encoding alone: wrong dimension for the prior, or a likelihood
126    /// with latent nuisance sites (see [`Self::try_init_from`] for the reason
127    /// and [`Self::init_from_with_latents`] to draw them). Never panics
128    /// (EV-N3).
129    pub fn init_from(&self, genome: &P::Genome) -> Option<Trace> {
130        self.try_init_from(genome).ok()
131    }
132
133    /// [`Self::init_from`] with the reason on failure:
134    /// [`EvolutionModel::score`]'s errors for a structural mismatch, or
135    /// [`GenomeError::ConstraintViolation`] for a genome outside the prior's
136    /// support.
137    pub fn try_init_from(&self, genome: &P::Genome) -> Result<Trace, GenomeError> {
138        let (_g, trace) = self.model.score(genome)?;
139        finite_state(trace)
140    }
141
142    /// Warm-start from a genome when the likelihood has **latent nuisance
143    /// sites** (an inferred noise scale, a Pareto weight): the genome's sites
144    /// come from its encoding, the latent ones are drawn from their priors
145    /// with `rng`, and the result is a complete, fully scored state. Same
146    /// errors as [`Self::try_init_from`].
147    pub fn init_from_with_latents<R: Rng>(
148        &self,
149        rng: &mut R,
150        genome: &P::Genome,
151    ) -> Result<Trace, GenomeError> {
152        let (_g, trace) = self.model.score_with_latents(rng, genome)?;
153        finite_state(trace)
154    }
155
156    /// One π_β-invariant transition. Moves ANY site type; honours
157    /// [`Self::override_site`]. Returns the decoded genome and the new state
158    /// (the freshly scored proposal on acceptance, a copy of `current` on
159    /// rejection).
160    ///
161    /// Costs exactly one model execution — the proposal — plus, on rejection,
162    /// a replay of the **prior program only** (no likelihood / fitness
163    /// evaluation) to decode the genome of the unchanged state. Callers who
164    /// keep their own decoded genome can use [`Self::step_scored`] and skip
165    /// even that.
166    ///
167    /// # Contract on `current`
168    ///
169    /// `current` must be a fully scored trace of this chain's target: one
170    /// returned by [`Self::init`], [`Self::init_from`] /
171    /// [`Self::init_from_with_latents`], or a previous `step` /
172    /// `step_scored`. Its accumulators and per-site densities are trusted as
173    /// the current state's log-density and as the reverse-move densities of
174    /// sites a proposal makes vanish. A trace assembled by hand —
175    /// [`TraceGenome::to_trace`](crate::genome::trace_genome::TraceGenome::to_trace)
176    /// or [`GenomePrior::trace_of`], whose per-site `logp` is 0 — violates
177    /// this and over-accepts structure-shrinking moves until the first
178    /// acceptance; route it through `init_from` first.
179    pub fn step<R: Rng>(&mut self, rng: &mut R, current: &Trace) -> (P::Genome, Trace) {
180        match self.step_scored(rng, current) {
181            Some((g, t, _log_weight)) => (g, t),
182            None => (self.decode(current), current.clone()),
183        }
184    }
185
186    /// One π_β-invariant transition from a scored state, at the cost of a
187    /// single model execution: `Some((genome, scored_trace, log_weight))` on
188    /// acceptance — `log_weight == scored_trace.total_log_weight()` — or
189    /// `None` on rejection, in which case the caller keeps `current`. Same
190    /// contract on `current` as [`Self::step`].
191    pub fn step_scored<R: Rng>(
192        &mut self,
193        rng: &mut R,
194        current: &Trace,
195    ) -> Option<(P::Genome, Trace, f64)> {
196        adaptive_single_site_mh_cached(
197            rng,
198            self.model.target_model(),
199            current,
200            &mut self.adaptation,
201            &self.overrides,
202            true,
203        )
204    }
205
206    /// Decode the genome of a chain state by replaying the **prior program**
207    /// over it (the prior's return value *is* the decoded genome). No
208    /// likelihood or fitness is evaluated. `state` must be a complete
209    /// assignment for the prior — every trace this chain hands out is.
210    pub fn decode(&self, state: &Trace) -> P::Genome {
211        let (g, _) = run(
212            ScoreGivenTrace {
213                base: state.clone(),
214                trace: Trace::default(),
215            },
216            self.model.prior().model(),
217        );
218        g
219    }
220
221    /// Full warmup-then-frozen chain: `warmup` adaptive iterations are
222    /// discarded, then `n` samples are collected from the frozen kernel.
223    /// Returns decoded genomes with their traces.
224    pub fn run_chain<R: Rng>(&self, rng: &mut R, n: usize, warmup: usize) -> Vec<(P::Genome, Trace)>
225    where
226        P::Genome: Clone,
227    {
228        adaptive_mcmc_chain_with_overrides(
229            rng,
230            self.model.target_model(),
231            n,
232            warmup,
233            &self.overrides,
234        )
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use crate::fitness::traits::Fitness;
242    use crate::genome::bounds::{Bounds, MultiBounds};
243    use crate::genome::real_vector::RealVector;
244    use crate::genome::traits::{BinaryGenome, PermutationGenome, RealValuedGenome};
245    use crate::inference::model::tests::PtrFitness;
246    use crate::inference::prior::{BitStringPrior, PermutationPrior, UniformBoxPrior};
247    use rand::rngs::StdRng;
248    use rand::SeedableRng;
249
250    fn linear_x0(g: &RealVector) -> f64 {
251        g.genes()[0]
252    }
253
254    /// Regression: EV-90 — no MH sample may escape the uniform-prior bounds,
255    /// and the boundary is not over-weighted: on [-2, 2] with f(x) = x the
256    /// β=1 Boltzmann posterior is ∝ e^x truncated to [-2, 2], with analytic
257    /// mean (e² + 3e⁻²)/(e² − e⁻²) ≈ 1.0746. Re-driven through the fugue
258    /// kernel instead of the deleted hand-rolled one.
259    #[test]
260    fn test_mh_respects_bounds() {
261        let prior = UniformBoxPrior::new(MultiBounds::new(vec![Bounds::new(-2.0, 2.0)]));
262        let model = EvolutionModel::new(prior, PtrFitness(linear_x0)).with_beta(1.0);
263        let mut chain = EvolutionChain::new(model);
264
265        let mut rng = StdRng::seed_from_u64(20260710);
266        let mut current = chain.init(&mut rng);
267        let mut samples = Vec::new();
268        for i in 0..40_000 {
269            let (g, t) = chain.step(&mut rng, &current);
270            current = t;
271            let x = g.genes()[0];
272            assert!((-2.0..=2.0).contains(&x), "MH sample escaped bounds: {}", x);
273            if i >= 5_000 {
274                samples.push(x);
275            }
276        }
277        let mean = samples.iter().sum::<f64>() / samples.len() as f64;
278        let analytic = {
279            let e2 = 2.0_f64.exp();
280            let em2 = (-2.0_f64).exp();
281            (e2 + 3.0 * em2) / (e2 - em2)
282        };
283        assert!(
284            (mean - analytic).abs() < 0.1,
285            "posterior mean {} deviates from truncated-exponential analytic {}",
286            mean,
287            analytic
288        );
289    }
290
291    /// Regression: FG-N1 downstream — the same truncated-exponential anchor on
292    /// `[-0.5, 0.5]`. Under the pre-fix fugue proposal selector a `Uniform`
293    /// site whose support excludes `-1` but contains negatives was put on a
294    /// log-space walk whenever its first draw was positive, so the chain
295    /// inherited the sign of its initial state and could never cross zero
296    /// (posterior mean ≈ +0.27 or ≈ −0.23 instead of the analytic
297    /// `a·coth(a) − 1 = 0.0820` for `ρ ∝ eˣ` on `[−a, a]`, `a = 1/2`). fugue
298    /// now selects the proposal from `Distribution::support()` (Reflect for a
299    /// bounded site), so the chain driven through `EvolutionChain::step`
300    /// visits both signs with the analytic mass `P(x > 0) = 0.6225`.
301    #[test]
302    fn test_mh_bounded_prior_containing_negatives_mixes_across_zero() {
303        let prior = UniformBoxPrior::new(MultiBounds::new(vec![Bounds::new(-0.5, 0.5)]));
304        let model = EvolutionModel::new(prior, PtrFitness(linear_x0)).with_beta(1.0);
305        let analytic_mean = 0.5 / (0.5f64).tanh() - 1.0;
306        let analytic_p_pos = (0.5f64.exp() - 1.0) / (0.5f64.exp() - (-0.5f64).exp());
307        for seed in [1u64, 2, 3, 20260710] {
308            let mut chain = EvolutionChain::new(model.clone());
309            let mut rng = StdRng::seed_from_u64(seed);
310            let mut current = chain.init(&mut rng);
311            let mut samples = Vec::new();
312            for i in 0..40_000 {
313                let (g, t) = chain.step(&mut rng, &current);
314                current = t;
315                let x = g.genes()[0];
316                assert!((-0.5..=0.5).contains(&x), "MH sample escaped bounds: {}", x);
317                if i >= 5_000 {
318                    samples.push(x);
319                }
320            }
321            let n = samples.len() as f64;
322            let mean = samples.iter().sum::<f64>() / n;
323            let p_pos = samples.iter().filter(|&&x| x > 0.0).count() as f64 / n;
324            assert!(
325                (mean - analytic_mean).abs() < 0.04,
326                "seed {seed}: posterior mean {mean} deviates from analytic {analytic_mean}"
327            );
328            assert!(
329                (p_pos - analytic_p_pos).abs() < 0.08,
330                "seed {seed}: P(x > 0) = {p_pos} vs analytic {analytic_p_pos} — chain stuck on one sign"
331            );
332        }
333    }
334
335    /// EV-N2 / X-5(a): `step` honours `override_site`. A `Reflect` override
336    /// narrower than the prior box confines a chain started inside it — the
337    /// reflected walk can never propose outside `[lower, upper]` — whereas
338    /// the same chain without the override wanders over the whole box.
339    #[test]
340    fn test_step_honours_override_site() {
341        let prior = || UniformBoxPrior::new(MultiBounds::new(vec![Bounds::new(-2.0, 2.0)]));
342        let start = RealVector::new(vec![0.0]);
343
344        let mut confined = EvolutionChain::new(EvolutionModel::new(prior(), PtrFitness(linear_x0)))
345            .override_site(
346                fugue::addr!("gene", 0),
347                SiteProposal::Reflect {
348                    lower: -0.5,
349                    upper: 0.5,
350                },
351            );
352        let mut rng = StdRng::seed_from_u64(3);
353        let mut current = confined.init_from(&start).expect("in support");
354        let mut accepted = 0;
355        for _ in 0..5_000 {
356            if let Some((g, t, _)) = confined.step_scored(&mut rng, &current) {
357                accepted += 1;
358                current = t;
359                let x = g.genes()[0];
360                assert!(
361                    (-0.5..=0.5).contains(&x),
362                    "override ignored: reflected chain left [-0.5, 0.5] at {x}"
363                );
364            }
365        }
366        assert!(
367            accepted > 500,
368            "confined chain barely moved ({accepted} acceptances)"
369        );
370
371        let mut free = EvolutionChain::new(EvolutionModel::new(prior(), PtrFitness(linear_x0)));
372        let mut rng = StdRng::seed_from_u64(3);
373        let mut current = free.init_from(&start).expect("in support");
374        let mut escaped = false;
375        for _ in 0..5_000 {
376            let (g, t) = free.step(&mut rng, &current);
377            current = t;
378            if g.genes()[0].abs() > 0.5 {
379                escaped = true;
380                break;
381            }
382        }
383        assert!(
384            escaped,
385            "without the override the chain must explore the whole box"
386        );
387    }
388
389    /// EV-N2 / X-5(b): a transition costs one model execution. `init_from`
390    /// evaluates the fitness once (the scoring replay); each `step` evaluates
391    /// it exactly once more (the proposal), whether accepted or rejected —
392    /// the rejected path decodes the genome from the prior program alone.
393    #[test]
394    fn test_step_costs_one_fitness_evaluation() {
395        use std::sync::atomic::{AtomicUsize, Ordering};
396        use std::sync::Arc;
397
398        #[derive(Clone)]
399        struct Counting(Arc<AtomicUsize>);
400        impl Fitness for Counting {
401            type Genome = RealVector;
402            type Value = f64;
403            fn evaluate(&self, g: &RealVector) -> f64 {
404                self.0.fetch_add(1, Ordering::SeqCst);
405                -0.5 * g.genes()[0].powi(2)
406            }
407        }
408
409        let counter = Arc::new(AtomicUsize::new(0));
410        let prior = UniformBoxPrior::new(MultiBounds::new(vec![Bounds::new(-3.0, 3.0)]));
411        let mut chain = EvolutionChain::new(EvolutionModel::new(prior, Counting(counter.clone())));
412        let mut rng = StdRng::seed_from_u64(5);
413        let mut current = chain
414            .init_from(&RealVector::new(vec![0.3]))
415            .expect("in support");
416        assert_eq!(counter.load(Ordering::SeqCst), 1, "init_from scores once");
417
418        let n = 2_000;
419        let mut rejections = 0;
420        for _ in 0..n {
421            let before = counter.load(Ordering::SeqCst);
422            let (g, t) = chain.step(&mut rng, &current);
423            assert_eq!(
424                counter.load(Ordering::SeqCst) - before,
425                1,
426                "a step must evaluate the fitness exactly once"
427            );
428            let x = t.get_f64(&fugue::addr!("gene", 0)).unwrap();
429            if x == current.get_f64(&fugue::addr!("gene", 0)).unwrap() {
430                rejections += 1;
431            }
432            assert_eq!(g.genes()[0], x);
433            current = t;
434        }
435        assert!(
436            rejections > 0,
437            "some proposals must be rejected for the test to bite"
438        );
439        assert_eq!(counter.load(Ordering::SeqCst), 1 + n);
440    }
441
442    /// New regression (dead-chain fix): a BitString chain must actually move.
443    /// The old `EvolutionStep::propose` cloned every non-F64 choice unchanged,
444    /// making this exact scenario a silent no-op forever.
445    #[test]
446    fn test_bitstring_chain_moves() {
447        #[derive(Clone, Copy)]
448        struct OnesCount;
449        impl Fitness for OnesCount {
450            type Genome = crate::genome::bit_string::BitString;
451            type Value = f64;
452            fn evaluate(&self, g: &Self::Genome) -> f64 {
453                g.bits().iter().filter(|&&b| b).count() as f64
454            }
455        }
456
457        let model = EvolutionModel::new(BitStringPrior::uniform(8), OnesCount).with_beta(1.0);
458        let mut chain = EvolutionChain::new(model);
459        let mut rng = StdRng::seed_from_u64(11);
460        let init = chain.init(&mut rng);
461        let init_bits: Vec<Option<bool>> = (0..8)
462            .map(|i| init.get_bool(&fugue::addr!("bit", i)))
463            .collect();
464
465        let mut current = init.clone();
466        let mut moved = false;
467        for _ in 0..200 {
468            let (_g, t) = chain.step(&mut rng, &current);
469            current = t;
470            let bits: Vec<Option<bool>> = (0..8)
471                .map(|i| current.get_bool(&fugue::addr!("bit", i)))
472                .collect();
473            if bits != init_bits {
474                moved = true;
475                break;
476            }
477        }
478        assert!(moved, "BitString chain never moved (dead-chain regression)");
479    }
480
481    /// New regression (dead-chain fix): a Permutation chain must move AND stay
482    /// inside the permutation support (the sequential categorical prior gives
483    /// colliding proposals probability zero).
484    #[test]
485    fn test_permutation_chain_moves() {
486        #[derive(Clone, Copy)]
487        struct SortedNess;
488        impl Fitness for SortedNess {
489            type Genome = crate::genome::permutation::Permutation;
490            type Value = f64;
491            fn evaluate(&self, g: &Self::Genome) -> f64 {
492                // Rewards ascending order.
493                g.permutation().windows(2).filter(|w| w[0] < w[1]).count() as f64
494            }
495        }
496
497        let model = EvolutionModel::new(PermutationPrior::new(5), SortedNess).with_beta(1.0);
498        let mut chain = EvolutionChain::new(model);
499        let mut rng = StdRng::seed_from_u64(17);
500        let init = chain.init(&mut rng);
501        let read_perm = |t: &Trace| -> Vec<usize> {
502            (0..5)
503                .map(|i| t.get_usize(&fugue::addr!("perm", i)).unwrap())
504                .collect()
505        };
506        let init_perm = read_perm(&init);
507
508        let mut current = init;
509        let mut moved = false;
510        for _ in 0..500 {
511            let (g, t) = chain.step(&mut rng, &current);
512            current = t;
513            assert!(
514                g.is_valid_permutation(),
515                "chain left the permutation support: {:?}",
516                g.permutation()
517            );
518            if read_perm(&current) != init_perm {
519                moved = true;
520            }
521        }
522        assert!(
523            moved,
524            "Permutation chain never moved (dead-chain regression)"
525        );
526    }
527}
fugue_evo/inference/mh.rs

fugue_evo/inference/
mh.rs
