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`], which picks the target site uniformly
7//! over **all** sites and dispatches the proposal by value type (Gaussian /
8//! log-space walk for `F64`, flip for `Bool`, reflected discrete walk for
9//! `U64`, prior-resample for `Usize`, integer walk for `I64`), including the
10//! reversible-jump corrections for structure-changing models.
11
12use std::collections::HashMap;
13
14use fugue::inference::mcmc_utils::DiminishingAdaptation;
15use fugue::{
16    adaptive_mcmc_chain_with_overrides, adaptive_single_site_mh, Address, SiteProposal, Trace,
17};
18use rand::Rng;
19
20use super::likelihood::GenomeLikelihood;
21use super::model::EvolutionModel;
22use super::prior::GenomePrior;
23
24/// An MH chain over the fixed-β Boltzmann target `π_β ∝ p(x)·exp(β·f(x))`.
25pub struct EvolutionChain<P, L>
26where
27    P: GenomePrior,
28    L: GenomeLikelihood<P::Genome>,
29{
30    model: EvolutionModel<P, L>,
31    adaptation: DiminishingAdaptation,
32    overrides: HashMap<Address, SiteProposal>,
33}
34
35impl<P, L> EvolutionChain<P, L>
36where
37    P: GenomePrior,
38    L: GenomeLikelihood<P::Genome>,
39{
40    /// Create a chain over the model's fixed-β target.
41    pub fn new(model: EvolutionModel<P, L>) -> Self {
42        Self {
43            model,
44            adaptation: DiminishingAdaptation::new(0.44, 0.7),
45            overrides: HashMap::new(),
46        }
47    }
48
49    /// Set the adaptation's target acceptance rate (default 0.44).
50    pub fn target_rate(mut self, rate: f64) -> Self {
51        self.adaptation = DiminishingAdaptation::new(rate, 0.7);
52        self
53    }
54
55    /// Force a specific `f64` proposal for one address (e.g.
56    /// `SiteProposal::Reflect { lower, upper }` for a bounded coordinate).
57    pub fn override_site(mut self, addr: Address, proposal: SiteProposal) -> Self {
58        self.overrides.insert(addr, proposal);
59        self
60    }
61
62    /// The underlying model.
63    pub fn model(&self) -> &EvolutionModel<P, L> {
64        &self.model
65    }
66
67    /// Draw an initial state: a prior sample's fully-scored trace.
68    pub fn init<R: Rng>(&self, rng: &mut R) -> Trace {
69        use fugue::runtime::handler::run;
70        use fugue::runtime::interpreters::PriorHandler;
71        let (_g, trace) = run(
72            PriorHandler {
73                rng,
74                trace: Trace::default(),
75            },
76            (self.model.target_model())(),
77        );
78        trace
79    }
80
81    /// Warm-start the chain from a given genome: encode it under the model's
82    /// prior ([`GenomePrior::trace_of`]) and score it through the target
83    /// program. Works for any prior — including grammar priors over trees —
84    /// so a classic GA/GP result can seed an inference chain. Returns `None`
85    /// if the genome is outside the prior's support (its target density is
86    /// `−∞`, which can never be left by an MH chain).
87    pub fn init_from(&self, genome: &P::Genome) -> Option<Trace> {
88        let (_g, trace) = self.model.score(genome);
89        if trace.total_log_weight().is_finite() {
90            Some(trace)
91        } else {
92            None
93        }
94    }
95
96    /// One π_β-invariant transition. Moves ANY site type.
97    pub fn step<R: Rng>(&mut self, rng: &mut R, current: &Trace) -> (P::Genome, Trace) {
98        adaptive_single_site_mh(
99            rng,
100            self.model.target_model(),
101            current,
102            &mut self.adaptation,
103        )
104    }
105
106    /// Full warmup-then-frozen chain: `warmup` adaptive iterations are
107    /// discarded, then `n` samples are collected from the frozen kernel.
108    /// Returns decoded genomes with their traces.
109    pub fn run_chain<R: Rng>(&self, rng: &mut R, n: usize, warmup: usize) -> Vec<(P::Genome, Trace)>
110    where
111        P::Genome: Clone,
112    {
113        adaptive_mcmc_chain_with_overrides(
114            rng,
115            self.model.target_model(),
116            n,
117            warmup,
118            &self.overrides,
119        )
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use crate::fitness::traits::Fitness;
127    use crate::genome::bounds::{Bounds, MultiBounds};
128    use crate::genome::real_vector::RealVector;
129    use crate::genome::traits::{BinaryGenome, PermutationGenome, RealValuedGenome};
130    use crate::inference::model::tests::PtrFitness;
131    use crate::inference::prior::{BitStringPrior, PermutationPrior, UniformBoxPrior};
132    use rand::rngs::StdRng;
133    use rand::SeedableRng;
134
135    fn linear_x0(g: &RealVector) -> f64 {
136        g.genes()[0]
137    }
138
139    /// Regression: EV-90 — no MH sample may escape the uniform-prior bounds,
140    /// and the boundary is not over-weighted: on [-2, 2] with f(x) = x the
141    /// β=1 Boltzmann posterior is ∝ e^x truncated to [-2, 2], with analytic
142    /// mean (e² + 3e⁻²)/(e² − e⁻²) ≈ 1.0746. Re-driven through the fugue
143    /// kernel instead of the deleted hand-rolled one.
144    #[test]
145    fn test_mh_respects_bounds() {
146        let prior = UniformBoxPrior::new(MultiBounds::new(vec![Bounds::new(-2.0, 2.0)]));
147        let model = EvolutionModel::new(prior, PtrFitness(linear_x0)).with_beta(1.0);
148        let mut chain = EvolutionChain::new(model);
149
150        let mut rng = StdRng::seed_from_u64(20260710);
151        let mut current = chain.init(&mut rng);
152        let mut samples = Vec::new();
153        for i in 0..40_000 {
154            let (g, t) = chain.step(&mut rng, &current);
155            current = t;
156            let x = g.genes()[0];
157            assert!((-2.0..=2.0).contains(&x), "MH sample escaped bounds: {}", x);
158            if i >= 5_000 {
159                samples.push(x);
160            }
161        }
162        let mean = samples.iter().sum::<f64>() / samples.len() as f64;
163        let analytic = {
164            let e2 = 2.0_f64.exp();
165            let em2 = (-2.0_f64).exp();
166            (e2 + 3.0 * em2) / (e2 - em2)
167        };
168        assert!(
169            (mean - analytic).abs() < 0.1,
170            "posterior mean {} deviates from truncated-exponential analytic {}",
171            mean,
172            analytic
173        );
174    }
175
176    /// New regression (dead-chain fix): a BitString chain must actually move.
177    /// The old `EvolutionStep::propose` cloned every non-F64 choice unchanged,
178    /// making this exact scenario a silent no-op forever.
179    #[test]
180    fn test_bitstring_chain_moves() {
181        #[derive(Clone, Copy)]
182        struct OnesCount;
183        impl Fitness for OnesCount {
184            type Genome = crate::genome::bit_string::BitString;
185            type Value = f64;
186            fn evaluate(&self, g: &Self::Genome) -> f64 {
187                g.bits().iter().filter(|&&b| b).count() as f64
188            }
189        }
190
191        let model = EvolutionModel::new(BitStringPrior::uniform(8), OnesCount).with_beta(1.0);
192        let mut chain = EvolutionChain::new(model);
193        let mut rng = StdRng::seed_from_u64(11);
194        let init = chain.init(&mut rng);
195        let init_bits: Vec<Option<bool>> = (0..8)
196            .map(|i| init.get_bool(&fugue::addr!("bit", i)))
197            .collect();
198
199        let mut current = init.clone();
200        let mut moved = false;
201        for _ in 0..200 {
202            let (_g, t) = chain.step(&mut rng, &current);
203            current = t;
204            let bits: Vec<Option<bool>> = (0..8)
205                .map(|i| current.get_bool(&fugue::addr!("bit", i)))
206                .collect();
207            if bits != init_bits {
208                moved = true;
209                break;
210            }
211        }
212        assert!(moved, "BitString chain never moved (dead-chain regression)");
213    }
214
215    /// New regression (dead-chain fix): a Permutation chain must move AND stay
216    /// inside the permutation support (the sequential categorical prior gives
217    /// colliding proposals probability zero).
218    #[test]
219    fn test_permutation_chain_moves() {
220        #[derive(Clone, Copy)]
221        struct SortedNess;
222        impl Fitness for SortedNess {
223            type Genome = crate::genome::permutation::Permutation;
224            type Value = f64;
225            fn evaluate(&self, g: &Self::Genome) -> f64 {
226                // Rewards ascending order.
227                g.permutation().windows(2).filter(|w| w[0] < w[1]).count() as f64
228            }
229        }
230
231        let model = EvolutionModel::new(PermutationPrior::new(5), SortedNess).with_beta(1.0);
232        let mut chain = EvolutionChain::new(model);
233        let mut rng = StdRng::seed_from_u64(17);
234        let init = chain.init(&mut rng);
235        let read_perm = |t: &Trace| -> Vec<usize> {
236            (0..5)
237                .map(|i| t.get_usize(&fugue::addr!("perm", i)).unwrap())
238                .collect()
239        };
240        let init_perm = read_perm(&init);
241
242        let mut current = init;
243        let mut moved = false;
244        for _ in 0..500 {
245            let (g, t) = chain.step(&mut rng, &current);
246            current = t;
247            assert!(
248                g.is_valid_permutation(),
249                "chain left the permutation support: {:?}",
250                g.permutation()
251            );
252            if read_perm(&current) != init_perm {
253                moved = true;
254            }
255        }
256        assert!(
257            moved,
258            "Permutation chain never moved (dead-chain regression)"
259        );
260    }
261}
fugue_evo/inference/mh.rs

fugue_evo/inference/
mh.rs
