ReNunney/rust/track2-core/examples/example.rs

//! Example usage of Track 2 components.
//!
//! This example demonstrates how to use the Track 2 components together
//! to run simulations, estimate extinction probabilities, and find thresholds.

use track2_core::{
    Config, NunneySimulationKernel,
    BinarySearchStrategy, SweepStrategy,
    Track2Job, Track2Result,
    extinction_trait::ExtinctionEstimator,
};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 1. Create a simulation configuration
    let config = Config {
        K: 1000.0,
        N0: 500.0,
        R: 10.0,
        T: 50.0,
        n: 3,
        u: 0.001,
        epochs: 5,
        target_extinction_probability: 0.5,
        seed: None,
    };

    println!("Config: {}", config);
    println!("Mutation supply M = 2*K*u = {}", config.mutation_supply());

    // 2. Create a simulation kernel
    let kernel = NunneySimulationKernel::new(None);

    // 3. Run a single simulation
    println!("\nRunning a single simulation...");
    let result = kernel.run(&config);
    println!("  Extinct: {}, Prob: {}", result.extinct, result.extinction_probability());
    println!(
        "  Final N: {}, target: {:.3}, mismatch: {:.3}",
        result.final_state.N,
        result.final_state.optimum,
        result.final_state.mean_mismatch
    );
    println!(
        "  Mean allele: {:.3}, tracking gap: {:.3}, births: {}, survivors: {}",
        result.final_state.allele_means.first().copied().unwrap_or(0.0),
        result.final_state.mean_tracking_gap,
        result.final_state.birth_count,
        result.final_state.surviving_offspring_count
    );
    println!(
        "  First/last nonzero allele t: {:?} / {:?}",
        result.first_nonzero_allele_t,
        result.last_nonzero_allele_t
    );

    // 4. Estimate extinction probability with multiple runs
    println!("\nEstimating extinction probability with 5 runs...");
    let num_runs = 5;
    let prob = kernel.estimate_extinction_probability(&config, num_runs);
    println!("  Extinction probability: {}", prob);

    // 5. Search for threshold using binary search
    println!("\nSearching for threshold with binary search...");
    let t_values = vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0];
    let _strategy = BinarySearchStrategy::new();
    
    // Note: This requires a real ExtinctionEstimator implementation
    // For now, we'll just show the structure
    // let estimate = strategy.search_threshold(&estimator, &config, &t_values, 0.5);
    // if let Some(estimate) = estimate {
    //     println!("  Estimated threshold: {}", estimate.estimated_t);
    //     println!("  Lower: {}, Upper: {}", estimate.lower_t, estimate.upper_t);
    // }

    // 6. Generate sweep data for plotting
    println!("\nGenerating sweep data...");
    let sweep_strategy = SweepStrategy::with_default();
    println!("  Testing T values: {:?}", t_values);
    println!("  Strategy: {} runs per T", sweep_strategy.runs_per_T);

    // 7. Create a Track 2 job manifest
    println!("\nCreating Track 2 job...");
    let job = Track2Job::new(
        "example-job-1".to_string(),
        42,
        config,
        0.5,
        t_values.clone(),
    );
    println!("  Job ID: {}", job.job_id);
    println!("  Track: {}", job.track);
    println!("  Job kind: {}", job.job_kind);

    // 8. Create a result summary
    let summary = track2_core::Track2Summary {
        extinction_probability: prob,
        estimated_t: 0.0, // Will be set after threshold search
        num_runs,
        extinct_count: (prob * num_runs as f64) as u32,
        tested_t_values: t_values.clone(),
        search_strategy: "sweep".to_string(),
    };

    let result_manifest = Track2Result::success(job.job_id, summary, vec![]);
    println!("  Status: {}", result_manifest.status);
    println!("  Exit code: {}", result_manifest.exit_code);

    println!("\n✓ Example completed successfully!");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn example_simulation() {
        let config = Config {
            K: 100.0,
            N0: 50.0,
            R: 10.0,
            T: 20.0,
            n: 1,
            u: 0.001,
            epochs: 1,
            target_extinction_probability: 0.5,
            seed: Some(123),
        };

        let kernel = NunneySimulationKernel::new(None);
        let result = kernel.run(&config);

        assert_eq!(result.config.K, 100.0);
        assert_eq!(result.config.n, 1);
        assert!(result.generations() > 0);
    }

    #[test]
    fn config_derived_values() {
        let config = Config {
            K: 1000.0,
            N0: 500.0,
            R: 10.0,
            T: 50.0,
            n: 3,
            u: 0.001,
            epochs: 5,
            target_extinction_probability: 0.5,
            seed: None,
        };

        assert_eq!(config.mutation_supply(), 2.0 * 1000.0 * 0.001);
        assert_eq!(config.generations_per_epoch(), 250);
    }
}