henad-models 0.3.0

//! The SIR epidemic as a [`GpuGridModel`], with the rules of [`crate::sir`] and the grid in a storage buffer.
//!
//! The CPU model draws from one `xorshift64` stream along each row, and WGSL has no 64-bit integers to run it.
//! Each cell instead keeps a `pcg_hash` state in a second buffer, `rng`, and the step advances it one round beside the
//! cell. The engine ping-pongs both buffers together, and display and reduce read `state` alone.
//!
//! [`seed_cells`] repeats the draws of the CPU model's `init`, and the two backends agree at tick 0. Their random
//! streams differ after tick 0.

use henad_compute::cpu::grid_engine::{GRID_INIT_SEED, grid_init_rng};
use henad_core::action::ActionDescriptor;
use henad_core::authoring::model::binding::BindingDecl;
use henad_core::authoring::model::gpu_grid_model::{GpuGridAction, GpuGridModel};
use henad_core::authoring::primitives::rng::{mix_seed, pcg_hash, xorshift64};
use henad_core::helpers::{extract_f32, extract_u32, f32_param, u32_param};
use henad_core::params::{ParamDescriptor, ParamValue};
use henad_core::view::{StatDescriptor, StatValue};

use crate::shader_bindings::gpu_sir::seed_outbreak::Params as ActionParams;
use crate::shader_bindings::gpu_sir::step::Params as StepParams;
use crate::sir::PALETTE;

/// Domain separator of the `rng` buffer's seed, so the buffer and the cell seeding stream do not start correlated.
const RNG_INIT_SEED: u64 = GRID_INIT_SEED ^ 0x5EED_5EED_5EED_5EED;

// The list repeats the ids and order the CPU engine composes for `SirGridModel`, so one parameter vector drives
// either backend.
henad_core::params! {
    const PARAM_WIDTH = u32_param("grid_width", "Grid Width", DEFAULT_DIM, 1, 16_384);
    const PARAM_HEIGHT = u32_param("grid_height", "Grid Height", DEFAULT_DIM, 1, 16_384);
    const PARAM_INFECTION_RATE =
        f32_param("infection_rate", "Infection Rate", DEFAULT_INFECTION_RATE, 0.0, 1.0, Some(0.01));
    const PARAM_RECOVERY_RATE =
        f32_param("recovery_rate", "Recovery Rate", DEFAULT_RECOVERY_RATE, 0.0, 1.0, Some(0.01));
    const PARAM_INITIAL_INFECTED_PCT = f32_param(
        "initial_infected_pct",
        "Initial Infected",
        DEFAULT_INITIAL_INFECTED_PCT,
        0.0,
        1.0,
        Some(0.001),
    )
    .percent();
}

const DEFAULT_DIM: u32 = 1024;
const DEFAULT_INFECTION_RATE: f32 = 0.3;
const DEFAULT_RECOVERY_RATE: f32 = 0.05;
const DEFAULT_INITIAL_INFECTED_PCT: f32 = 0.01;

// Cell states, taken from the shader that defines them rather than mirrored.
use crate::shader_bindings::gpu_sir::step::{I as CELL_I, R as CELL_R, S as CELL_S};

/// Returns the initial cell states, one `u32` per cell, each infected with probability `initial_infected_pct`.
///
/// The cells are drawn from `rng` as `SirGridModel::init` draws them. With `rng` from `grid_init_rng`, the grid
/// matches the CPU model's bit for bit.
pub fn seed_cells(width: u32, height: u32, initial_infected_pct: f32, mut rng: u64) -> Vec<u32> {
    let threshold = (initial_infected_pct * u32::MAX as f32) as u32;
    let mut cells = vec![0u32; (width as usize) * (height as usize)];
    for cell in &mut cells {
        rng = xorshift64(rng);
        *cell = u32::from(((rng >> 32) as u32) < threshold);
    }
    cells
}

/// Returns the initial `pcg_hash` state of every cell, hashed from `seed` and the cell's index.
fn seed_rng_states(width: u32, height: u32, seed: u64) -> Vec<u32> {
    let seed32 = (seed ^ (seed >> 32)) as u32;
    (0..(width as usize) * (height as usize))
        .map(|idx| pcg_hash(seed32 ^ idx as u32))
        .collect()
}

/// The SIR epidemic as a [`GpuGridModel`], with a random number state per cell.
#[derive(Debug)]
pub struct GpuSir;

impl GpuGridModel for GpuSir {
    const NAME: &'static str = "SIR Epidemic (GPU)";
    const ID: &'static str = "gpu_sir";
    const DESCRIPTION: &'static str =
        "Classic SIR compartmental model on a toroidal grid with Moore neighborhood, stepped entirely on the GPU";
    const PALETTE: &'static [[u8; 4]] = &PALETTE;
    const STATS: &'static [StatDescriptor] = &[
        StatDescriptor::new("Susceptible", PALETTE[0]),
        StatDescriptor::new("Infected", PALETTE[1]),
        StatDescriptor::new("Recovered", PALETTE[2]),
    ];

    const BUFFERS: &'static [&'static str] = &["state", "rng"];

    const STEP_BINDINGS: &'static [BindingDecl] = crate::binding_decls::bindings::GPU_SIR_STEP;
    const DISPLAY_BINDINGS: &'static [BindingDecl] = crate::binding_decls::bindings::GPU_SIR_DISPLAY;
    const REDUCE_BINDINGS: &'static [BindingDecl] = crate::binding_decls::bindings::GPU_SIR_REDUCE;

    const STEP_SHADER: &'static str = crate::shader_bindings::gpu_sir::step::SHADER_STRING;
    const DISPLAY_SHADER: &'static str = crate::shader_bindings::gpu_sir::display::SHADER_STRING;
    const REDUCE_SHADER: &'static str = crate::shader_bindings::gpu_sir::reduce::SHADER_STRING;

    const ACTIONS: &'static [GpuGridAction] = &[GpuGridAction {
        desc: ActionDescriptor::new("seed_outbreak", "Seed outbreak"),
        shader: crate::shader_bindings::gpu_sir::seed_outbreak::SHADER_STRING,
        bindings: crate::binding_decls::bindings::GPU_SIR_SEED_OUTBREAK,
    }];

    fn param_descriptors() -> Vec<ParamDescriptor> {
        descriptors()
    }

    fn dims(params: &[ParamValue]) -> (u32, u32) {
        (
            extract_u32(params, PARAM_WIDTH, DEFAULT_DIM),
            extract_u32(params, PARAM_HEIGHT, DEFAULT_DIM),
        )
    }

    fn seed_buffers(width: u32, height: u32, params: &[ParamValue], seed: Option<u64>) -> Vec<Vec<u32>> {
        let initial_infected_pct = extract_f32(params, PARAM_INITIAL_INFECTED_PCT, DEFAULT_INITIAL_INFECTED_PCT);
        let cells = grid_init_rng(seed);
        let rng = seed.map_or(RNG_INIT_SEED, |s| mix_seed(s ^ RNG_INIT_SEED));
        vec![
            seed_cells(width, height, initial_infected_pct, cells),
            seed_rng_states(width, height, rng),
        ]
    }

    fn step_params_bytes(width: u32, height: u32, params: &[ParamValue]) -> Vec<u8> {
        bytemuck::bytes_of(&StepParams {
            width,
            height,
            infection_rate: extract_f32(params, PARAM_INFECTION_RATE, DEFAULT_INFECTION_RATE),
            recovery_rate: extract_f32(params, PARAM_RECOVERY_RATE, DEFAULT_RECOVERY_RATE),
        })
        .to_vec()
    }

    fn action_params_bytes(_action: usize, width: u32, height: u32, params: &[ParamValue], seed: u32) -> Vec<u8> {
        let pct = extract_f32(params, PARAM_INITIAL_INFECTED_PCT, DEFAULT_INITIAL_INFECTED_PCT);
        bytemuck::bytes_of(&ActionParams {
            width,
            height,
            threshold: (pct * u32::MAX as f32) as u32,
            seed,
        })
        .to_vec()
    }

    fn stats(counts: &[u32]) -> Vec<StatValue> {
        vec![
            StatValue::Scalar(f64::from(counts[CELL_S as usize])),
            StatValue::Scalar(f64::from(counts[CELL_I as usize])),
            StatValue::Scalar(f64::from(counts[CELL_R as usize])),
        ]
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use henad_compute::cpu::grid_engine::{GRID_PARAM_BASE, GridModelState};
    use henad_compute::gpu::GpuContext;
    use henad_compute::gpu::grid_engine::GpuGridState;
    use henad_compute::gpu::sim_thread::GpuSimState as _;
    use henad_core::authoring::model::grid_model::GridModel as _;
    use henad_core::grid::Grid2D;
    use henad_core::model::SimState as _;
    use henad_core::view::StatEntry;
    use henad_explore::testing::{TestDeviceRequest, headless_test_device};

    use crate::sir::SirGridModel;

    type State = GpuGridState<GpuSir>;

    pub(super) fn headless_context() -> Option<GpuContext> {
        headless_test_device(&TestDeviceRequest::baseline())
    }

    pub(super) fn params(
        width: u32,
        height: u32,
        infection_rate: f32,
        recovery_rate: f32,
        initial_infected_pct: f32,
    ) -> Vec<ParamValue> {
        vec![
            ParamValue::U32(width),
            ParamValue::U32(height),
            ParamValue::F32(infection_rate),
            ParamValue::F32(recovery_rate),
            ParamValue::F32(initial_infected_pct),
        ]
    }

    fn sir_counts(state: &State) -> (u64, u64, u64) {
        let stats = state.stats();
        let scalar = |entry: &StatEntry| match &entry.value {
            StatValue::Scalar(v) => *v as u64,
            other => panic!("expected a scalar stat, got {other:?}"),
        };
        (scalar(&stats[0]), scalar(&stats[1]), scalar(&stats[2]))
    }

    /// Runs the display, reduce and readback passes as the sim thread's one-shot snapshot does.
    fn refresh_stats(ctx: &GpuContext, state: &mut State) {
        let mut encoder = ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        state.encode_snapshot_passes(&mut encoder);
        ctx.queue.submit(Some(encoder.finish()));
        state.begin_stats_readback();
        state.poll_stats_readback(&ctx.device, true);
    }

    fn step_once(ctx: &GpuContext, state: &mut State) {
        let mut encoder = ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        state.encode_steps(&mut encoder, 1, None);
        ctx.queue.submit(Some(encoder.finish()));
    }

    /// The total of S, I and R holds exactly at every tick, whatever the per-cell RNG streams do.
    #[test]
    fn population_conserved_over_many_ticks() {
        let Some(ctx) = headless_context() else {
            log::warn!("skipping population_conserved_over_many_ticks: no adapter");
            return;
        };

        let (width, height) = (64u32, 64u32);
        let mut state = State::new(&ctx, &params(width, height, 0.3, 0.05, 0.1));
        let total = u64::from(width) * u64::from(height);

        for tick in 0..50 {
            refresh_stats(&ctx, &mut state);
            let (s, i, r) = sir_counts(&state);
            assert_eq!(s + i + r, total, "S+I+R must equal total population at tick {tick}");
            step_once(&ctx, &mut state);
        }
    }

    /// With `infection_rate == 0`, S can never lose a member.
    #[test]
    fn zero_infection_rate_freezes_susceptible_count() {
        let Some(ctx) = headless_context() else {
            log::warn!("skipping zero_infection_rate_freezes_susceptible_count: no adapter");
            return;
        };

        let (width, height) = (64u32, 64u32);
        let mut state = State::new(&ctx, &params(width, height, 0.0, 0.05, 0.2));

        refresh_stats(&ctx, &mut state);
        let (initial_s, _, _) = sir_counts(&state);

        for tick in 0..20 {
            step_once(&ctx, &mut state);
            refresh_stats(&ctx, &mut state);
            let (s, _, _) = sir_counts(&state);
            assert_eq!(
                s, initial_s,
                "susceptible count must not change at tick {tick} when infection_rate is 0"
            );
        }
    }

    /// With `recovery_rate == 0`, I can never lose a member, so at any positive infection rate I never decreases.
    #[test]
    fn zero_recovery_rate_keeps_infected_non_decreasing() {
        let Some(ctx) = headless_context() else {
            log::warn!("skipping zero_recovery_rate_keeps_infected_non_decreasing: no adapter");
            return;
        };

        let (width, height) = (64u32, 64u32);
        let mut state = State::new(&ctx, &params(width, height, 0.5, 0.0, 0.1));

        refresh_stats(&ctx, &mut state);
        let (_, mut prev_i, _) = sir_counts(&state);

        for tick in 0..20 {
            step_once(&ctx, &mut state);
            refresh_stats(&ctx, &mut state);
            let (_, i, r) = sir_counts(&state);
            assert!(
                i >= prev_i,
                "infected count must not decrease at tick {tick} when recovery_rate is 0"
            );
            assert_eq!(r, 0, "no cell can reach R when recovery_rate is 0 (tick {tick})");
            prev_i = i;
        }
    }

    /// The port repeats `SirGridModel::init` rather than calling it, so every cell is compared.
    #[test]
    fn the_seeded_grid_matches_the_cpu_init() {
        let (width, height) = (37u32, 23u32);
        let p = params(width, height, 0.3, 0.05, 0.2);
        for seed in [None, Some(7)] {
            let mut grid = Grid2D::new(width, height);
            SirGridModel::init(&mut grid, &p[GRID_PARAM_BASE..], &mut grid_init_rng(seed));
            let cpu: Vec<u32> = grid.current().iter().copied().map(u32::from).collect();
            assert_eq!(
                GpuSir::seed_buffers(width, height, &p, seed)[0],
                cpu,
                "the seeded state differs from the CPU grid for seed {seed:?}"
            );
        }
    }

    /// Initial seeding uses the same PRNG, traversal order, and threshold as the CPU model, so the
    /// tick-0 compartment counts (before any RNG-dependent step) must match exactly.
    #[test]
    fn initial_seed_matches_cpu_model() {
        let Some(ctx) = headless_context() else {
            log::warn!("skipping initial_seed_matches_cpu_model: no adapter");
            return;
        };

        let (width, height) = (64u32, 64u32);
        let p = params(width, height, 0.3, 0.05, 0.2);

        let mut gpu = State::new(&ctx, &p);
        let cpu = GridModelState::<SirGridModel>::from_params(&p);

        refresh_stats(&ctx, &mut gpu);
        let (gpu_s, gpu_i, gpu_r) = sir_counts(&gpu);

        let cpu_scalar = |entry: &StatEntry| match &entry.value {
            StatValue::Scalar(v) => *v as u64,
            other => panic!("expected a scalar stat, got {other:?}"),
        };
        let cpu_stats = cpu.stats();
        assert_eq!(
            gpu_s,
            cpu_scalar(&cpu_stats[0]),
            "initial susceptible count must match the CPU model"
        );
        assert_eq!(
            gpu_i,
            cpu_scalar(&cpu_stats[1]),
            "initial infected count must match the CPU model"
        );
        assert_eq!(
            gpu_r,
            cpu_scalar(&cpu_stats[2]),
            "initial recovered count must be 0 on both backends"
        );
    }

    #[test]
    fn population_is_total_cells() {
        let Some(ctx) = headless_context() else {
            log::warn!("skipping population_is_total_cells: no adapter");
            return;
        };

        let (width, height) = (32u32, 16u32);
        let state = State::new(&ctx, &params(width, height, 0.3, 0.05, 0.1));
        assert_eq!(state.population(), u64::from(width) * u64::from(height));
    }
}

/// Tests that drive the real [`henad_compute::gpu::GpuSimThread`], similar to
/// `gpu_game_of_life`'s `runner_tests`.
#[cfg(test)]
#[cfg(not(target_arch = "wasm32"))]
mod runner_tests {
    use std::time::{Duration, Instant};

    use henad_compute::gpu::grid_engine::GpuGridState;
    use henad_compute::gpu::sim_thread::{GpuBatchSettings, GpuSimThread};
    use henad_compute::snapshot::{Snapshot, SnapshotView};

    use super::GpuSir;
    use super::tests::{headless_context, params};
    use henad_compute::entry::ModelState;
    use henad_core::view::StatValue;

    fn wait_for(thread: &mut GpuSimThread, timeout: Duration, pred: impl Fn(&Snapshot) -> bool) -> Option<Snapshot> {
        let deadline = Instant::now() + timeout;
        while Instant::now() < deadline {
            if let Some(snap) = thread.take_snapshot()
                && pred(&snap)
            {
                return Some(snap);
            }
            std::thread::sleep(Duration::from_millis(2));
        }
        None
    }

    fn sir_counts(snap: &Snapshot) -> (u64, u64, u64) {
        let scalar = |entry: &henad_core::view::StatEntry| match &entry.value {
            StatValue::Scalar(v) => *v as u64,
            other => panic!("expected a scalar stat, got {other:?}"),
        };
        (scalar(&snap.stats[0]), scalar(&snap.stats[1]), scalar(&snap.stats[2]))
    }

    #[test]
    fn gpu_thread_publishes_gpu_snapshots_and_conserves_population() {
        let Some(ctx) = headless_context() else {
            log::warn!("skipping gpu_thread_publishes_gpu_snapshots_and_conserves_population: no adapter");
            return;
        };

        let (width, height) = (128u32, 128u32);
        let state = GpuGridState::<GpuSir>::new(&ctx, &params(width, height, 0.3, 0.05, 0.1));
        let mut thread = GpuSimThread::new(ctx, Box::new(state), GpuBatchSettings::default(), None);

        let initial = wait_for(&mut thread, Duration::from_secs(5), |_| true)
            .expect("the GPU thread must publish an initial snapshot before Play");
        assert!(matches!(initial.view, SnapshotView::Gpu(_)));
        let total = initial.population;
        let (s, i, r) = sir_counts(&initial);
        assert_eq!(s + i + r, total, "initial snapshot must already conserve population");

        thread.play();
        let running = wait_for(&mut thread, Duration::from_secs(5), |snap| snap.tick > 0)
            .expect("playing must advance the tick counter and publish fresh snapshots");
        assert!(matches!(running.view, SnapshotView::Gpu(_)));
        let (s, i, r) = sir_counts(&running);
        assert_eq!(s + i + r, total, "population must stay conserved while running");

        thread.pause();
        drop(thread);
    }

    #[test]
    fn registry_with_gpu_context_offers_a_drivable_gpu_sir_model() {
        let Some(ctx) = headless_context() else {
            log::warn!("skipping registry_with_gpu_context_offers_a_drivable_gpu_sir_model: no adapter");
            return;
        };

        let models = crate::example_models();
        let entry = models
            .lookup("gpu_sir", Some(&ctx))
            .expect("a GPU context must make the GPU SIR model selectable");

        let built = entry
            .build(&params(32, 32, 0.3, 0.05, 0.1), None, Some(&ctx))
            .unwrap_or_else(|fault| panic!("the GPU SIR entry's factory failed to build: {fault}"));
        let ModelState::Gpu(mut state) = built else {
            panic!("the GPU SIR entry's factory must yield ModelState::Gpu, not ModelState::Cpu");
        };

        let mut encoder = ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        state.encode_steps(&mut encoder, 4, None);
        ctx.queue.submit(Some(encoder.finish()));
        assert_eq!(state.tick(), 4, "the registry-built state must be steppable");
    }
}