fugue-evo 0.2.0

//! Tempered SMC over the Boltzmann posterior, built on fugue's SMC engine
//!
//! The old `EvolutionarySMC` hand-rolled the whole tempering loop (linear β
//! ladder, weight normalization, ESS, systematic resampling, MH sweeps) and
//! carried a weight-model bug: it reweighted by `dβ · f(x)` on top of a
//! `β·f(x)` factor, double-counting β. This rebuild deletes all of it. The
//! driver is [`fugue::adaptive_smc_with_kernel`] run against the model's
//! **untempered** target (`factor(f)`, see
//! [`EvolutionModel::smc_model`](super::model::EvolutionModel::smc_model)):
//! fugue supplies β by likelihood-tempering with an adaptive ESS-driven
//! ladder, applies it exactly once, and returns an unbiased log-evidence
//! estimate for free.
//!
//! Crossover is fugue's [`CrossoverKernel`] — a population-coupled Metropolis
//! move on the product target — driven by an address mask supplied here
//! (genome knowledge stays downstream; trace-space mechanics live upstream).

use std::marker::PhantomData;

use fugue::runtime::handler::run;
use fugue::runtime::interpreters::ScoreGivenTrace;
use fugue::{
    adaptive_smc, adaptive_smc_with_kernel, decode_particle, CrossoverKernel, Model, Particle,
    ResamplingMethod, SMCConfig, Trace,
};
use rand::Rng;

use super::model::EvolutionModel;
use super::prior::GenomePrior;
use crate::fitness::traits::Fitness;
use crate::genome::trace_genome::{gene_address, TraceGenome};

/// Configuration of the crossover population kernel.
#[derive(Clone, Debug)]
pub struct CrossoverConfig {
    /// Number of (pair, swap) proposals per sweep.
    pub n_pairs: usize,
    /// Per-address probability that a site joins the swap mask.
    pub swap_probability: f64,
}

impl Default for CrossoverConfig {
    fn default() -> Self {
        Self {
            n_pairs: 32,
            swap_probability: 0.5,
        }
    }
}

/// Configuration for [`EvolutionSMC::run`].
pub struct EvoSmcConfig {
    /// Number of particles.
    pub num_particles: usize,
    /// ESS threshold fraction driving both the adaptive β ladder and
    /// resampling (fugue `SMCConfig::ess_threshold`).
    pub ess_threshold: f64,
    /// Resampling algorithm.
    pub resampling: ResamplingMethod,
    /// Per-particle MH rejuvenation sweeps per tempering step.
    pub rejuvenation_steps: usize,
    /// Population crossover kernel; `None` = per-particle rejuvenation only.
    pub crossover: Option<CrossoverConfig>,
}

impl Default for EvoSmcConfig {
    fn default() -> Self {
        Self {
            num_particles: 500,
            ess_threshold: 0.5,
            resampling: ResamplingMethod::Systematic,
            rejuvenation_steps: 3,
            crossover: Some(CrossoverConfig::default()),
        }
    }
}

/// The result of a tempered-SMC evolution run: fugue particles (traces +
/// normalized weights) approximating the Boltzmann posterior `π ∝ p·exp(f)`,
/// plus the log-evidence estimate.
///
/// Genomes are not cached on particles; they are recovered by **decode-replay**
/// (replaying the particle's trace through the prior/target program, whose
/// return value *is* the decoded genome).
pub struct EvolutionPosterior<G: TraceGenome> {
    /// Final weighted particle population (fugue particles).
    pub particles: Vec<Particle>,
    /// Unbiased estimate of the log normalizing constant
    /// `log Σ_x p(x)·exp(f(x))` — the Bayesian model score.
    pub log_evidence: f64,
    _g: PhantomData<G>,
}

impl<G: TraceGenome> EvolutionPosterior<G> {
    /// Recover the genome of one particle by replaying its trace.
    pub fn genome(&self, particle: &Particle, model_fn: &impl Fn() -> Model<G>) -> G {
        decode_particle(particle, model_fn)
    }

    /// Decode the whole population as `(genome, normalized_weight)` pairs.
    pub fn genomes(&self, model_fn: &impl Fn() -> Model<G>) -> Vec<(G, f64)> {
        self.particles
            .iter()
            .map(|p| (decode_particle(p, model_fn), p.weight))
            .collect()
    }

    /// Self-normalised weighted posterior mean of coordinate `gene#coord`.
    pub fn weighted_mean(&self, coord: usize) -> f64 {
        let addr = gene_address(G::trace_prefix(), coord);
        let mut total_w = 0.0;
        let mut mean = 0.0;
        for p in &self.particles {
            if let Some(x) = p.trace.get_f64(&addr) {
                mean += p.weight * x;
                total_w += p.weight;
            }
        }
        if total_w > 0.0 {
            mean / total_w
        } else {
            0.0
        }
    }

    /// Self-normalised weighted posterior variance of coordinate `gene#coord`.
    pub fn weighted_variance(&self, coord: usize) -> f64 {
        let addr = gene_address(G::trace_prefix(), coord);
        let mean = self.weighted_mean(coord);
        let mut total_w = 0.0;
        let mut var = 0.0;
        for p in &self.particles {
            if let Some(x) = p.trace.get_f64(&addr) {
                var += p.weight * (x - mean).powi(2);
                total_w += p.weight;
            }
        }
        if total_w > 0.0 {
            var / total_w
        } else {
            0.0
        }
    }

    /// The decoded genome with the highest fitness, and that fitness —
    /// the optimizer-mode readout for benchmarking against the classic layer.
    pub fn best<F>(&self, fitness: &F, model_fn: &impl Fn() -> Model<G>) -> Option<(G, f64)>
    where
        F: Fitness<Genome = G, Value = f64>,
    {
        self.particles
            .iter()
            .map(|p| {
                let g = decode_particle(p, model_fn);
                let f = fitness.evaluate(&g);
                (g, f)
            })
            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
    }
}

/// The tempered-SMC evolution driver.
pub struct EvolutionSMC;

impl EvolutionSMC {
    /// Run tempered SMC targeting the Boltzmann posterior
    /// `π ∝ p(x)·exp(f(x))` of `model` (β is supplied by fugue's adaptive
    /// tempering; `model`'s own β setting is ignored here by construction).
    pub fn run<P, F, R>(
        rng: &mut R,
        model: &EvolutionModel<P, F>,
        cfg: EvoSmcConfig,
    ) -> EvolutionPosterior<P::Genome>
    where
        P: GenomePrior,
        F: Fitness<Genome = P::Genome, Value = f64> + Clone + Send + Sync + 'static,
        R: Rng,
    {
        let model_fn = model.smc_model();
        let smc_cfg = SMCConfig {
            resampling_method: cfg.resampling,
            ess_threshold: cfg.ess_threshold,
            rejuvenation_steps: cfg.rejuvenation_steps,
        };
        let result = match cfg.crossover {
            None => adaptive_smc(rng, cfg.num_particles, &model_fn, smc_cfg),
            Some(xcfg) => {
                let p_swap = xcfg.swap_probability.clamp(0.0, 1.0);
                let mut kernel = CrossoverKernel {
                    n_pairs: xcfg.n_pairs,
                    // Value-independent, pair-symmetric mask: each address of
                    // the first parent joins the swap independently.
                    mask: Box::new(move |a: &Trace, _b: &Trace, rng: &mut dyn rand::RngCore| {
                        a.choices
                            .keys()
                            .filter(|_| rand::Rng::gen::<f64>(rng) < p_swap)
                            .cloned()
                            .collect()
                    }),
                };
                adaptive_smc_with_kernel(rng, cfg.num_particles, &model_fn, smc_cfg, &mut kernel)
            }
        };
        EvolutionPosterior {
            particles: result.particles,
            log_evidence: result.log_evidence,
            _g: PhantomData,
        }
    }
}

impl EvolutionSMC {
    /// Like [`EvolutionSMC::run`], but with an explicit population kernel
    /// (e.g. a [`CrossoverKernel`] with a
    /// [`subtree_crossover_mask`](super::grammar::subtree_crossover_mask) for
    /// grammar-driven tree genomes). `cfg.crossover` is ignored.
    pub fn run_with_kernel<P, F, R, K>(
        rng: &mut R,
        model: &EvolutionModel<P, F>,
        cfg: EvoSmcConfig,
        kernel: &mut K,
    ) -> EvolutionPosterior<P::Genome>
    where
        P: GenomePrior,
        F: Fitness<Genome = P::Genome, Value = f64> + Clone + Send + Sync + 'static,
        R: Rng,
        K: fugue::PopulationKernel<P::Genome>,
    {
        let model_fn = model.smc_model();
        let smc_cfg = SMCConfig {
            resampling_method: cfg.resampling,
            ess_threshold: cfg.ess_threshold,
            rejuvenation_steps: cfg.rejuvenation_steps,
        };
        let result = adaptive_smc_with_kernel(rng, cfg.num_particles, &model_fn, smc_cfg, kernel);
        EvolutionPosterior {
            particles: result.particles,
            log_evidence: result.log_evidence,
            _g: PhantomData,
        }
    }
}

/// Score a genome's canonical trace under an arbitrary model — convenience
/// used by readouts and tests.
pub fn score_genome<G: TraceGenome, A>(genome: &G, model: Model<A>) -> (A, Trace) {
    run(
        ScoreGivenTrace {
            base: genome.to_trace(),
            trace: Trace::default(),
        },
        model,
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::genome::bounds::{Bounds, MultiBounds};
    use crate::genome::real_vector::RealVector;
    use crate::genome::traits::RealValuedGenome;
    use crate::inference::model::tests::PtrFitness;
    use crate::inference::prior::GaussianPrior;
    use rand::rngs::StdRng;
    use rand::SeedableRng;

    fn quad_k1_c3(g: &RealVector) -> f64 {
        -0.5 * g.genes().iter().map(|x| (x - 3.0).powi(2)).sum::<f64>()
    }

    /// Regression: EV-16 — tempered SMC on a quadratic fitness with a Gaussian
    /// prior reproduces the conjugate Boltzmann posterior.
    ///
    /// Prior N(0, σ0²=4) ⇒ τ0 = 0.25; fitness −0.5(x−3)² ⇒ k = 1, c = 3.
    /// Posterior at β=1: τ = 1.25, mean = 3/1.25 = 2.4, variance = 0.8.
    /// Re-driven through the fugue-backed rebuild — this directly exercises
    /// the β-single-counting fix (fitness enters as `factor(f)`; β only from
    /// tempering).
    #[test]
    fn test_smc_matches_gaussian_conjugate_posterior() {
        let prior = GaussianPrior::new(0.0, 2.0, 1);
        let model = EvolutionModel::new(prior, PtrFitness(quad_k1_c3));
        let mut rng = StdRng::seed_from_u64(42);
        let result = EvolutionSMC::run(
            &mut rng,
            &model,
            EvoSmcConfig {
                num_particles: 4000,
                ess_threshold: 0.5,
                resampling: ResamplingMethod::Systematic,
                rejuvenation_steps: 6,
                crossover: None,
            },
        );

        let mean = result.weighted_mean(0);
        let var = result.weighted_variance(0);
        assert!(
            (mean - 2.4).abs() < 0.15,
            "posterior mean {} vs analytic 2.4",
            mean
        );
        assert!(
            (var - 0.8).abs() < 0.2,
            "posterior variance {} vs analytic 0.8",
            var
        );

        // Weights are self-normalised.
        let total: f64 = result.particles.iter().map(|p| p.weight).sum();
        assert!((total - 1.0).abs() < 1e-6);

        // Analytic evidence check comes for free from the rebuild:
        // Z = ∫ N(x; 0, 4)·e^{-(x-3)²/2} dx = √(2π·0.8)/√(2π·4) · e^{-9/(2·5)}
        let analytic_log_z = 0.5 * ((0.8f64).ln() - (4.0f64).ln()) - 9.0 / (2.0 * 5.0);
        assert!(
            (result.log_evidence - analytic_log_z).abs() < 0.25,
            "log evidence {} vs analytic {}",
            result.log_evidence,
            analytic_log_z
        );
    }

    /// Same conjugate target, with the crossover population kernel enabled —
    /// the kernel must not bias the posterior (product-target invariance) nor
    /// the evidence (FG-58).
    #[test]
    fn test_smc_with_crossover_matches_conjugate_posterior() {
        let prior = GaussianPrior::new(0.0, 2.0, 2);
        // Independent per-coordinate quadratic pull toward 3.
        let model = EvolutionModel::new(prior, PtrFitness(quad_k1_c3));
        let mut rng = StdRng::seed_from_u64(1234);
        let result = EvolutionSMC::run(
            &mut rng,
            &model,
            EvoSmcConfig {
                num_particles: 3000,
                ess_threshold: 0.5,
                resampling: ResamplingMethod::Systematic,
                rejuvenation_steps: 4,
                crossover: Some(CrossoverConfig {
                    n_pairs: 500,
                    swap_probability: 0.5,
                }),
            },
        );
        for coord in 0..2 {
            let mean = result.weighted_mean(coord);
            let var = result.weighted_variance(coord);
            assert!(
                (mean - 2.4).abs() < 0.15,
                "coord {} posterior mean {} vs 2.4",
                coord,
                mean
            );
            assert!(
                (var - 0.8).abs() < 0.25,
                "coord {} posterior variance {} vs 0.8",
                coord,
                var
            );
        }
    }

    /// The rebuilt result exposes decode-replay: recover genomes from bare
    /// particle traces via the prior program's return value.
    #[test]
    fn test_decode_replay_recovers_genomes() {
        let prior = GaussianPrior::new(0.0, 2.0, 1);
        let model = EvolutionModel::new(prior, PtrFitness(quad_k1_c3));
        let mut rng = StdRng::seed_from_u64(5);
        let result = EvolutionSMC::run(
            &mut rng,
            &model,
            EvoSmcConfig {
                num_particles: 100,
                rejuvenation_steps: 2,
                crossover: None,
                ..Default::default()
            },
        );
        let model_fn = model.smc_model();
        let decoded = result.genomes(&model_fn);
        assert_eq!(decoded.len(), 100);
        for (g, _w) in &decoded {
            assert_eq!(g.genes().len(), 1);
        }
        let (best, best_f) = result.best(&PtrFitness(quad_k1_c3), &model_fn).unwrap();
        assert!(best_f.is_finite());
        assert!((quad_k1_c3(&best) - best_f).abs() < 1e-12);
    }

    /// Bounds are respected end-to-end: with a uniform-box prior every
    /// particle stays inside the box (out-of-box scores −∞ and can never
    /// survive).
    #[test]
    fn test_smc_respects_bounds() {
        use crate::inference::prior::UniformBoxPrior;
        let prior = UniformBoxPrior::new(MultiBounds::new(vec![Bounds::new(-2.0, 2.0)]));
        let model = EvolutionModel::new(prior, PtrFitness(|g: &RealVector| g.genes()[0]));
        let mut rng = StdRng::seed_from_u64(9);
        let result = EvolutionSMC::run(
            &mut rng,
            &model,
            EvoSmcConfig {
                num_particles: 300,
                rejuvenation_steps: 3,
                crossover: Some(CrossoverConfig::default()),
                ..Default::default()
            },
        );
        for p in &result.particles {
            let x = p.trace.get_f64(&fugue::addr!("gene", 0)).unwrap();
            assert!((-2.0..=2.0).contains(&x), "particle escaped bounds: {}", x);
        }
    }
}