fugue_evo/checkpoint/
state.rs

1//! Checkpoint state structures
2//!
3//! Complete evolution state for checkpointing and recovery.
4
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7
8use crate::diagnostics::GenerationStats;
9use crate::genome::traits::EvolutionaryGenome;
10use crate::population::individual::Individual;
11
12/// Current checkpoint format version
13pub const CHECKPOINT_VERSION: u32 = 1;
14
15/// Complete evolution state for checkpointing
16#[derive(Clone, Debug, Serialize, Deserialize)]
17#[serde(bound = "")]
18pub struct Checkpoint<G>
19where
20    G: Clone + Serialize + EvolutionaryGenome,
21{
22    /// Schema version for forward compatibility
23    pub version: u32,
24    /// Current generation
25    pub generation: usize,
26    /// Total fitness evaluations
27    pub evaluations: usize,
28    /// Population with fitness values
29    pub population: Vec<Individual<G>>,
30    /// RNG state for reproducibility (serialized bytes)
31    pub rng_state: Option<Vec<u8>>,
32    /// Best individual found so far
33    pub best: Option<Individual<G>>,
34    /// Algorithm-specific state
35    pub algorithm_state: AlgorithmState,
36    /// Hyperparameter state if using adaptive learning
37    pub hyperparameter_state: Option<HyperparameterState>,
38    /// Statistics history
39    pub statistics: Vec<GenerationStats>,
40    /// Custom metadata
41    pub metadata: HashMap<String, String>,
42}
43
44impl<G> Checkpoint<G>
45where
46    G: Clone + Serialize + EvolutionaryGenome,
47{
48    /// Create a new checkpoint
49    pub fn new(generation: usize, population: Vec<Individual<G>>) -> Self {
50        Self {
51            version: CHECKPOINT_VERSION,
52            generation,
53            evaluations: 0,
54            population,
55            rng_state: None,
56            best: None,
57            algorithm_state: AlgorithmState::SimpleGA,
58            hyperparameter_state: None,
59            statistics: Vec::new(),
60            metadata: HashMap::new(),
61        }
62    }
63
64    /// Set the number of evaluations
65    pub fn with_evaluations(mut self, evaluations: usize) -> Self {
66        self.evaluations = evaluations;
67        self
68    }
69
70    /// Set the best individual
71    pub fn with_best(mut self, best: Individual<G>) -> Self {
72        self.best = Some(best);
73        self
74    }
75
76    /// Set algorithm-specific state
77    pub fn with_algorithm_state(mut self, state: AlgorithmState) -> Self {
78        self.algorithm_state = state;
79        self
80    }
81
82    /// Set hyperparameter state
83    pub fn with_hyperparameter_state(mut self, state: HyperparameterState) -> Self {
84        self.hyperparameter_state = Some(state);
85        self
86    }
87
88    /// Add statistics history
89    pub fn with_statistics(mut self, stats: Vec<GenerationStats>) -> Self {
90        self.statistics = stats;
91        self
92    }
93
94    /// Add custom metadata
95    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
96        self.metadata.insert(key.into(), value.into());
97        self
98    }
99
100    /// Set RNG state from raw serialized bytes.
101    ///
102    /// Prefer [`Checkpoint::with_rng`] for a type-checked capture of a
103    /// [`SnapshotRng`](crate::checkpoint::rng::SnapshotRng).
104    pub fn with_rng_state(mut self, state: Vec<u8>) -> Self {
105        self.rng_state = Some(state);
106        self
107    }
108
109    /// Capture a [`SnapshotRng`](crate::checkpoint::rng::SnapshotRng)'s full state into this checkpoint (EV-02).
110    ///
111    /// Restoring the checkpoint's RNG via [`Checkpoint::restore_rng`] and
112    /// continuing evolution then reproduces the exact stochastic trajectory a
113    /// continuous (never-interrupted) run would have taken. Reproducible resume
114    /// requires a ChaCha RNG (`ChaCha8Rng`/`ChaCha12Rng`/`ChaCha20Rng`); generic
115    /// generators such as `StdRng`/`ThreadRng` cannot be snapshotted.
116    pub fn with_rng<R: crate::checkpoint::rng::SnapshotRng>(
117        mut self,
118        rng: &R,
119    ) -> Result<Self, crate::error::CheckpointError> {
120        self.rng_state = Some(rng.capture()?);
121        Ok(self)
122    }
123
124    /// Restore a [`SnapshotRng`](crate::checkpoint::rng::SnapshotRng) previously captured via [`Checkpoint::with_rng`].
125    ///
126    /// Returns `Ok(None)` if no RNG state was stored in this checkpoint.
127    pub fn restore_rng<R: crate::checkpoint::rng::SnapshotRng>(
128        &self,
129    ) -> Result<Option<R>, crate::error::CheckpointError> {
130        match &self.rng_state {
131            Some(bytes) => Ok(Some(R::restore(bytes)?)),
132            None => Ok(None),
133        }
134    }
135
136    /// Check if checkpoint is compatible with current version
137    pub fn is_compatible(&self) -> bool {
138        self.version <= CHECKPOINT_VERSION
139    }
140
141    /// Get the checkpoint version
142    pub fn version(&self) -> u32 {
143        self.version
144    }
145}
146
147/// Algorithm-specific state variants
148#[derive(Clone, Debug, Serialize, Deserialize)]
149pub enum AlgorithmState {
150    /// Simple generational GA (no additional state)
151    SimpleGA,
152    /// Steady-state GA
153    SteadyState { replacement_count: usize },
154    /// CMA-ES state
155    CmaEs(CmaEsCheckpointState),
156    /// NSGA-II state
157    Nsga2 { pareto_front_indices: Vec<usize> },
158    /// HBGA state
159    Hbga {
160        population_params: Vec<f64>,
161        temperature: f64,
162    },
163    /// Island model state
164    Island {
165        island_populations: Vec<Vec<usize>>,
166        migration_count: usize,
167    },
168    /// Interactive GA state
169    Interactive {
170        /// Serialized aggregator state (JSON)
171        aggregator_state: String,
172        /// Number of pending evaluations
173        pending_evaluations: usize,
174        /// Evaluation mode
175        evaluation_mode: String,
176    },
177    /// Custom algorithm state (JSON serialized)
178    Custom(String),
179}
180
181/// CMA-ES checkpoint state
182#[derive(Clone, Debug, Serialize, Deserialize)]
183pub struct CmaEsCheckpointState {
184    /// Mean vector
185    pub mean: Vec<f64>,
186    /// Global step size
187    pub sigma: f64,
188    /// Covariance matrix (flattened row-major)
189    pub covariance: Vec<f64>,
190    /// Evolution path for sigma
191    pub path_sigma: Vec<f64>,
192    /// Evolution path for covariance
193    pub path_c: Vec<f64>,
194    /// Dimension
195    pub dimension: usize,
196}
197
198/// Hyperparameter learning state
199#[derive(Clone, Debug, Serialize, Deserialize)]
200pub struct HyperparameterState {
201    /// Mutation rate posterior (alpha, beta for Beta distribution)
202    pub mutation_rate_posterior: Option<(f64, f64)>,
203    /// Crossover probability posterior
204    pub crossover_prob_posterior: Option<(f64, f64)>,
205    /// Selection temperature posterior (shape, rate for Gamma)
206    pub temperature_posterior: Option<(f64, f64)>,
207    /// Step size posteriors (mu, sigma_sq for LogNormal)
208    pub step_size_posteriors: Vec<(f64, f64)>,
209    /// Operator selection weights
210    pub operator_weights: Vec<f64>,
211    /// History window for learning
212    pub history_size: usize,
213}
214
215impl Default for HyperparameterState {
216    fn default() -> Self {
217        Self {
218            mutation_rate_posterior: None,
219            crossover_prob_posterior: None,
220            temperature_posterior: None,
221            step_size_posteriors: Vec::new(),
222            operator_weights: Vec::new(),
223            history_size: 100,
224        }
225    }
226}
227
228/// Builder for creating checkpoints
229pub struct CheckpointBuilder<G>
230where
231    G: Clone + Serialize + EvolutionaryGenome,
232{
233    checkpoint: Checkpoint<G>,
234}
235
236impl<G> CheckpointBuilder<G>
237where
238    G: Clone + Serialize + EvolutionaryGenome,
239{
240    /// Create a new checkpoint builder
241    pub fn new(generation: usize, population: Vec<Individual<G>>) -> Self {
242        Self {
243            checkpoint: Checkpoint::new(generation, population),
244        }
245    }
246
247    /// Set evaluations count
248    pub fn evaluations(mut self, count: usize) -> Self {
249        self.checkpoint.evaluations = count;
250        self
251    }
252
253    /// Set best individual
254    pub fn best(mut self, individual: Individual<G>) -> Self {
255        self.checkpoint.best = Some(individual);
256        self
257    }
258
259    /// Set algorithm state
260    pub fn algorithm_state(mut self, state: AlgorithmState) -> Self {
261        self.checkpoint.algorithm_state = state;
262        self
263    }
264
265    /// Set hyperparameter state
266    pub fn hyperparameters(mut self, state: HyperparameterState) -> Self {
267        self.checkpoint.hyperparameter_state = Some(state);
268        self
269    }
270
271    /// Set statistics
272    pub fn statistics(mut self, stats: Vec<GenerationStats>) -> Self {
273        self.checkpoint.statistics = stats;
274        self
275    }
276
277    /// Add metadata
278    pub fn metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
279        self.checkpoint.metadata.insert(key.into(), value.into());
280        self
281    }
282
283    /// Set RNG state
284    pub fn rng_state(mut self, state: Vec<u8>) -> Self {
285        self.checkpoint.rng_state = Some(state);
286        self
287    }
288
289    /// Build the checkpoint
290    pub fn build(self) -> Checkpoint<G> {
291        self.checkpoint
292    }
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298    use crate::genome::real_vector::RealVector;
299
300    #[test]
301    fn test_checkpoint_creation() {
302        let population: Vec<Individual<RealVector>> = vec![
303            Individual::new(RealVector::new(vec![1.0, 2.0])),
304            Individual::new(RealVector::new(vec![3.0, 4.0])),
305        ];
306
307        let checkpoint = Checkpoint::new(10, population.clone());
308
309        assert_eq!(checkpoint.version, CHECKPOINT_VERSION);
310        assert_eq!(checkpoint.generation, 10);
311        assert_eq!(checkpoint.population.len(), 2);
312    }
313
314    #[test]
315    fn test_checkpoint_builder() {
316        let population: Vec<Individual<RealVector>> =
317            vec![Individual::new(RealVector::new(vec![1.0]))];
318
319        let checkpoint = CheckpointBuilder::new(5, population)
320            .evaluations(1000)
321            .algorithm_state(AlgorithmState::SimpleGA)
322            .metadata("experiment", "test_run")
323            .build();
324
325        assert_eq!(checkpoint.generation, 5);
326        assert_eq!(checkpoint.evaluations, 1000);
327        assert_eq!(
328            checkpoint.metadata.get("experiment"),
329            Some(&"test_run".to_string())
330        );
331    }
332
333    #[test]
334    fn test_checkpoint_compatibility() {
335        let population: Vec<Individual<RealVector>> = vec![];
336        let checkpoint = Checkpoint::new(0, population);
337
338        assert!(checkpoint.is_compatible());
339    }
340
341    #[test]
342    fn test_cmaes_checkpoint_state() {
343        let state = CmaEsCheckpointState {
344            mean: vec![0.0, 0.0],
345            sigma: 1.0,
346            covariance: vec![1.0, 0.0, 0.0, 1.0],
347            path_sigma: vec![0.0, 0.0],
348            path_c: vec![0.0, 0.0],
349            dimension: 2,
350        };
351
352        let alg_state = AlgorithmState::CmaEs(state);
353        if let AlgorithmState::CmaEs(s) = alg_state {
354            assert_eq!(s.dimension, 2);
355            assert_eq!(s.sigma, 1.0);
356        } else {
357            panic!("Expected CmaEs state");
358        }
359    }
360
361    #[test]
362    fn test_checkpoint_with_methods() {
363        let population: Vec<Individual<RealVector>> =
364            vec![Individual::new(RealVector::new(vec![1.0, 2.0]))];
365        let best = Individual::with_fitness(RealVector::new(vec![0.5, 0.5]), 10.0);
366
367        let checkpoint = Checkpoint::new(5, population)
368            .with_evaluations(500)
369            .with_best(best.clone())
370            .with_algorithm_state(AlgorithmState::SteadyState {
371                replacement_count: 10,
372            })
373            .with_metadata("run_id", "test123")
374            .with_rng_state(vec![1, 2, 3, 4]);
375
376        assert_eq!(checkpoint.evaluations, 500);
377        assert!(checkpoint.best.is_some());
378        assert_eq!(
379            checkpoint.metadata.get("run_id"),
380            Some(&"test123".to_string())
381        );
382        assert!(checkpoint.rng_state.is_some());
383    }
384
385    #[test]
386    fn test_checkpoint_with_hyperparameters() {
387        let population: Vec<Individual<RealVector>> = vec![];
388        let hp_state = HyperparameterState {
389            mutation_rate_posterior: Some((2.0, 8.0)),
390            crossover_prob_posterior: Some((5.0, 5.0)),
391            ..Default::default()
392        };
393
394        let checkpoint = Checkpoint::new(0, population).with_hyperparameter_state(hp_state);
395
396        assert!(checkpoint.hyperparameter_state.is_some());
397        let hp = checkpoint.hyperparameter_state.unwrap();
398        assert_eq!(hp.mutation_rate_posterior, Some((2.0, 8.0)));
399    }
400
401    #[test]
402    fn test_checkpoint_with_statistics() {
403        use crate::diagnostics::{GenerationStats, TimingStats};
404
405        let population: Vec<Individual<RealVector>> = vec![];
406        let stats = vec![
407            GenerationStats {
408                generation: 0,
409                evaluations: 100,
410                best_fitness: 10.0,
411                worst_fitness: 1.0,
412                mean_fitness: 5.0,
413                median_fitness: 5.0,
414                fitness_std: 2.0,
415                diversity: 0.5,
416                timing: TimingStats::default(),
417            },
418            GenerationStats {
419                generation: 1,
420                evaluations: 200,
421                best_fitness: 15.0,
422                worst_fitness: 2.0,
423                mean_fitness: 7.0,
424                median_fitness: 7.0,
425                fitness_std: 1.5,
426                diversity: 0.4,
427                timing: TimingStats::default(),
428            },
429        ];
430
431        let checkpoint = Checkpoint::new(2, population).with_statistics(stats.clone());
432
433        assert_eq!(checkpoint.statistics.len(), 2);
434    }
435
436    #[test]
437    fn test_checkpoint_version() {
438        let population: Vec<Individual<RealVector>> = vec![];
439        let checkpoint = Checkpoint::new(0, population);
440
441        assert_eq!(checkpoint.version(), CHECKPOINT_VERSION);
442    }
443
444    #[test]
445    fn test_checkpoint_builder_full() {
446        use crate::diagnostics::{GenerationStats, TimingStats};
447
448        let population: Vec<Individual<RealVector>> =
449            vec![Individual::new(RealVector::new(vec![1.0]))];
450        let best = Individual::with_fitness(RealVector::new(vec![0.0]), 100.0);
451        let hp_state = HyperparameterState::default();
452        let stats = vec![GenerationStats {
453            generation: 0,
454            evaluations: 100,
455            best_fitness: 100.0,
456            worst_fitness: 10.0,
457            mean_fitness: 50.0,
458            median_fitness: 50.0,
459            fitness_std: 10.0,
460            diversity: 0.5,
461            timing: TimingStats::default(),
462        }];
463
464        let checkpoint = CheckpointBuilder::new(10, population)
465            .evaluations(5000)
466            .best(best)
467            .algorithm_state(AlgorithmState::Nsga2 {
468                pareto_front_indices: vec![0, 1, 2],
469            })
470            .hyperparameters(hp_state)
471            .statistics(stats)
472            .metadata("version", "1.0")
473            .rng_state(vec![0, 1, 2, 3])
474            .build();
475
476        assert_eq!(checkpoint.generation, 10);
477        assert_eq!(checkpoint.evaluations, 5000);
478        assert!(checkpoint.best.is_some());
479        assert!(checkpoint.hyperparameter_state.is_some());
480        assert_eq!(checkpoint.statistics.len(), 1);
481        assert!(checkpoint.rng_state.is_some());
482    }
483
484    #[test]
485    fn test_algorithm_state_variants() {
486        // Test all algorithm state variants
487        let simple_ga = AlgorithmState::SimpleGA;
488        let steady_state = AlgorithmState::SteadyState {
489            replacement_count: 5,
490        };
491        let nsga2 = AlgorithmState::Nsga2 {
492            pareto_front_indices: vec![0, 1],
493        };
494        let hbga = AlgorithmState::Hbga {
495            population_params: vec![1.0, 2.0],
496            temperature: 1.5,
497        };
498        let island = AlgorithmState::Island {
499            island_populations: vec![vec![0, 1], vec![2, 3]],
500            migration_count: 3,
501        };
502        let interactive = AlgorithmState::Interactive {
503            aggregator_state: "{}".to_string(),
504            pending_evaluations: 10,
505            evaluation_mode: "pairwise".to_string(),
506        };
507        let custom = AlgorithmState::Custom("custom_state".to_string());
508
509        // Verify they're all different variants (pattern matching)
510        assert!(matches!(simple_ga, AlgorithmState::SimpleGA));
511        assert!(matches!(steady_state, AlgorithmState::SteadyState { .. }));
512        assert!(matches!(nsga2, AlgorithmState::Nsga2 { .. }));
513        assert!(matches!(hbga, AlgorithmState::Hbga { .. }));
514        assert!(matches!(island, AlgorithmState::Island { .. }));
515        assert!(matches!(interactive, AlgorithmState::Interactive { .. }));
516        assert!(matches!(custom, AlgorithmState::Custom(_)));
517    }
518
519    #[test]
520    fn test_checkpoint_rng_capture_restore_round_trip() {
521        // regression: EV-02 - with_rng must capture a ChaCha RNG's state such
522        // that restore_rng reproduces an identical draw sequence.
523        use rand::{Rng, SeedableRng};
524        use rand_chacha::ChaCha8Rng;
525
526        let mut rng = ChaCha8Rng::seed_from_u64(1234);
527        for _ in 0..11 {
528            let _: u64 = rng.gen();
529        }
530
531        let population: Vec<Individual<RealVector>> = vec![];
532        let checkpoint = Checkpoint::new(0, population).with_rng(&rng).unwrap();
533
534        let mut restored: ChaCha8Rng = checkpoint
535            .restore_rng()
536            .unwrap()
537            .expect("rng state must be present");
538
539        for _ in 0..500 {
540            assert_eq!(rng.gen::<u64>(), restored.gen::<u64>());
541        }
542    }
543
544    #[test]
545    fn test_checkpoint_restore_rng_absent() {
546        use rand_chacha::ChaCha8Rng;
547        let population: Vec<Individual<RealVector>> = vec![];
548        let checkpoint = Checkpoint::new(0, population);
549        let restored: Option<ChaCha8Rng> = checkpoint.restore_rng().unwrap();
550        assert!(restored.is_none());
551    }
552
553    #[test]
554    fn test_hyperparameter_state_default() {
555        let hp = HyperparameterState::default();
556
557        assert!(hp.mutation_rate_posterior.is_none());
558        assert!(hp.crossover_prob_posterior.is_none());
559        assert!(hp.temperature_posterior.is_none());
560        assert!(hp.step_size_posteriors.is_empty());
561        assert!(hp.operator_weights.is_empty());
562        assert_eq!(hp.history_size, 100);
563    }
564}
fugue_evo/checkpoint/state.rs

fugue_evo/checkpoint/
state.rs
