rustsim_core/

soa.rs

//! SoA (Structure-of-Arrays) extraction for GPU-friendly data layout.
//!
//! ABM agents are stored as AoS (Array-of-Structures) in `AgentStore`.
//! For GPU kernels, we need flat contiguous arrays of each field.
//! `SoaExtractable` lets agent types define how to extract/write-back
//! their data to/from flat `f32` buffers suitable for GPU upload.
//!
//! For workloads that require `f64` precision (scientific simulations,
//! long time horizons, or ill-conditioned dynamics), implement
//! [`SoaExtractableF64`] instead of or alongside [`SoaExtractable`].
//! The two traits are independent: an agent type may implement either,
//! both, or neither. CUDA kernels that target `f64` must declare `double`
//! parameters.

use crate::agent::Agent;
use crate::store::AgentStore;
use crate::types::AgentId;

/// Trait for agents whose numeric fields can be extracted into SoA buffers
/// and written back from SoA buffers after GPU computation.
///
/// Each "column" is a `Vec<f32>` representing one field across all agents.
/// The order of agents in the buffers matches the order of IDs returned
/// by `extract`.
///
/// # Example
///
/// ```ignore
/// impl SoaExtractable for Particle {
///     fn num_columns() -> usize { 2 } // x, vx
///     fn column_names() -> Vec<&'static str> { vec!["x", "vx"] }
///     fn extract_row(&self, columns: &mut [Vec<f32>]) {
///         columns[0].push(self.x);
///         columns[1].push(self.vx);
///     }
///     fn write_back_row(&mut self, columns: &[&[f32]], row: usize) {
///         self.x = columns[0][row];
///         self.vx = columns[1][row];
///     }
/// }
/// ```
pub trait SoaExtractable: Agent {
    /// Number of f32 columns to extract.
    fn num_columns() -> usize;

    /// Human-readable names for each column (for debugging / PTX variable naming).
    fn column_names() -> Vec<&'static str>;

    /// Push this agent's values into the column vectors.
    fn extract_row(&self, columns: &mut [Vec<f32>]);

    /// Read this agent's values back from the column slices at `row`.
    fn write_back_row(&mut self, columns: &[&[f32]], row: usize);
}

/// Extract SoA buffers from an `AgentStore`.
///
/// Returns `(ids, columns)` where:
/// - `ids[i]` is the `AgentId` for row `i`
/// - `columns[c][i]` is column `c` for agent `i`
pub fn extract_soa<A, S>(store: &S) -> (Vec<AgentId>, Vec<Vec<f32>>)
where
    A: SoaExtractable,
    S: AgentStore<A>,
{
    let ids = store.iter_ids();
    let n = ids.len();
    let nc = A::num_columns();
    let mut columns: Vec<Vec<f32>> = (0..nc).map(|_| Vec::with_capacity(n)).collect();

    for &id in &ids {
        if let Some(agent) = store.get(id) {
            agent.extract_row(&mut columns);
        }
    }

    (ids, columns)
}

/// Write SoA buffers back into an `AgentStore`.
///
/// `ids` and `columns` must have the same row count and match the
/// order returned by `extract_soa`.
pub fn write_back_soa<A, S>(store: &S, ids: &[AgentId], columns: &[Vec<f32>])
where
    A: SoaExtractable,
    S: AgentStore<A>,
{
    let col_refs: Vec<&[f32]> = columns.iter().map(|c| c.as_slice()).collect();
    for (row, &id) in ids.iter().enumerate() {
        if let Some(mut agent) = store.get_mut(id) {
            agent.write_back_row(&col_refs, row);
        }
    }
}

// -----------------------------------------------------------------------
// f64 path
// -----------------------------------------------------------------------

/// Like [`SoaExtractable`] but using `f64` columns.
///
/// Implement this trait when your kernel needs double precision — e.g.
/// scientific workloads, long-horizon integrators, or ill-conditioned
/// dynamics. `f32` remains the default for parity with the CUDA batch path
/// and most ABM workloads.
///
/// An agent type may implement both [`SoaExtractable`] and
/// [`SoaExtractableF64`] independently; the two extraction paths do not
/// interact. Use [`cast_columns_f64_to_f32`] as a convenience when you want
/// to downcast an `f64` extraction result to `f32` for an existing `f32`
/// kernel.
pub trait SoaExtractableF64: Agent {
    /// Number of `f64` columns to extract.
    fn num_columns() -> usize;

    /// Human-readable names for each column.
    fn column_names() -> Vec<&'static str>;

    /// Push this agent's values into the column vectors.
    fn extract_row(&self, columns: &mut [Vec<f64>]);

    /// Read this agent's values back from the column slices at `row`.
    fn write_back_row(&mut self, columns: &[&[f64]], row: usize);
}

/// Extract `f64` SoA buffers from an `AgentStore`.
pub fn extract_soa_f64<A, S>(store: &S) -> (Vec<AgentId>, Vec<Vec<f64>>)
where
    A: SoaExtractableF64,
    S: AgentStore<A>,
{
    let ids = store.iter_ids();
    let n = ids.len();
    let nc = <A as SoaExtractableF64>::num_columns();
    let mut columns: Vec<Vec<f64>> = (0..nc).map(|_| Vec::with_capacity(n)).collect();

    for &id in &ids {
        if let Some(agent) = store.get(id) {
            <A as SoaExtractableF64>::extract_row(&agent, &mut columns);
        }
    }

    (ids, columns)
}

/// Write `f64` SoA buffers back into an `AgentStore`.
pub fn write_back_soa_f64<A, S>(store: &S, ids: &[AgentId], columns: &[Vec<f64>])
where
    A: SoaExtractableF64,
    S: AgentStore<A>,
{
    let col_refs: Vec<&[f64]> = columns.iter().map(|c| c.as_slice()).collect();
    for (row, &id) in ids.iter().enumerate() {
        if let Some(mut agent) = store.get_mut(id) {
            <A as SoaExtractableF64>::write_back_row(&mut agent, &col_refs, row);
        }
    }
}

/// Convenience helper: downcast `f64` columns to `f32` columns.
///
/// Useful when you have an `f64` extraction path but want to feed an
/// existing `f32` kernel. Precision loss is explicit and at the call site.
pub fn cast_columns_f64_to_f32(columns: &[Vec<f64>]) -> Vec<Vec<f32>> {
    columns
        .iter()
        .map(|c| c.iter().map(|v| *v as f32).collect())
        .collect()
}

/// Convenience helper: upcast `f32` columns to `f64` columns.
pub fn cast_columns_f32_to_f64(columns: &[Vec<f32>]) -> Vec<Vec<f64>> {
    columns
        .iter()
        .map(|c| c.iter().map(|v| *v as f64).collect())
        .collect()
}

#[cfg(test)]
mod tests_f64 {
    use super::*;
    use crate::store::{AgentStore, HashMapStore};

    #[derive(Clone, Debug)]
    struct P {
        id: AgentId,
        x: f64,
        vx: f64,
    }

    impl Agent for P {
        fn id(&self) -> AgentId {
            self.id
        }
    }

    impl SoaExtractableF64 for P {
        fn num_columns() -> usize {
            2
        }
        fn column_names() -> Vec<&'static str> {
            vec!["x", "vx"]
        }
        fn extract_row(&self, columns: &mut [Vec<f64>]) {
            columns[0].push(self.x);
            columns[1].push(self.vx);
        }
        fn write_back_row(&mut self, columns: &[&[f64]], row: usize) {
            self.x = columns[0][row];
            self.vx = columns[1][row];
        }
    }

    #[test]
    fn extract_and_write_back_preserve_f64_precision() {
        let mut store: HashMapStore<P> = HashMapStore::new();
        // A value that loses precision when round-tripped through f32.
        let precise = 1.0_f64 + 1.0e-10_f64;
        store.insert(P {
            id: 1,
            x: precise,
            vx: 2.0,
        });

        let (ids, mut cols) = extract_soa_f64::<P, _>(&store);
        assert_eq!(ids.len(), 1);
        assert_eq!(cols[0][0], precise);

        cols[0][0] = precise * 2.0;
        write_back_soa_f64::<P, _>(&store, &ids, &cols);

        let a = store.get(1).unwrap();
        assert_eq!(a.x, precise * 2.0);
    }

    #[test]
    fn f32_f64_cast_round_trip_loses_low_bits() {
        let f64_cols = vec![vec![1.0_f64 + 1.0e-10_f64]];
        let f32_cols = cast_columns_f64_to_f32(&f64_cols);
        let back = cast_columns_f32_to_f64(&f32_cols);
        // f32 cannot represent 1.0 + 1e-10, so the round-trip rounds to 1.0.
        assert_eq!(back[0][0], 1.0);
    }
}