rustytorch/demos/rtx-cardiosim-demo/src/monodomain.rs

//! Monodomain equation solver for cardiac electrophysiology.
//!
//! Solves the monodomain equation:
//! ∂V/∂t = ∇·(D∇V) + Iion(V, w)
//!
//! where V is transmembrane voltage, D is diffusion tensor, and Iion is ionic current.

use cardiosim_shared::{
    ConductivityParams, HeartMesh, IonicModel, SimulationConfig, StimulationProtocol, Vector3D,
    VoltageField,
};

use crate::CardioSimError;
use crate::ionic_models::IonicState;

/// Monodomain solver.
#[derive(Debug)]
pub struct MonodomainSolver {
    ionic_model: IonicModel,
    dt: f32,
    total_time: f32,
    output_interval: f32,
    conductivity: ConductivityParams,
}

impl MonodomainSolver {
    /// Create a new monodomain solver.
    #[must_use]
    pub fn new(config: &SimulationConfig) -> Self {
        Self {
            ionic_model: config.ionic_model,
            dt: config.dt,
            total_time: config.total_time,
            output_interval: config.output_interval,
            conductivity: config.conductivity.clone(),
        }
    }

    /// Solve monodomain equation on heart mesh.
    pub fn solve(
        &self,
        mesh: &HeartMesh,
        protocol: &StimulationProtocol,
    ) -> Result<Vec<VoltageField>, CardioSimError> {
        let n_vertices = mesh.vertices.len();
        let n_steps = (self.total_time / self.dt) as usize;
        let output_steps = (self.output_interval / self.dt) as usize;

        // Initialize state
        let mut voltage = vec![-85.0_f32; n_vertices];
        let mut ionic_states: Vec<IonicState> = (0..n_vertices)
            .map(|_| IonicState::new(self.ionic_model))
            .collect();

        // Precompute diffusion operator
        let diffusion = self.compute_diffusion_operator(mesh);

        let mut output = Vec::new();

        for step in 0..n_steps {
            let t = step as f32 * self.dt;

            // Apply stimulation
            let stim = self.compute_stimulus(mesh, protocol, t);

            // Diffusion step (implicit or semi-implicit)
            let dv_diffusion = self.apply_diffusion(&voltage, &diffusion);

            // Ionic current step
            let (dv_ionic, dw) = self.compute_ionic_currents(&voltage, &ionic_states);

            // Update voltage
            for i in 0..n_vertices {
                voltage[i] += self.dt * (dv_diffusion[i] + dv_ionic[i] + stim[i]);
            }

            // Update ionic state
            for (i, state) in ionic_states.iter_mut().enumerate() {
                state.recovery += self.dt * dw[i];
            }

            // Output
            if step % output_steps == 0 {
                output.push(VoltageField {
                    time: t,
                    voltages: voltage.clone(),
                });
            }
        }

        Ok(output)
    }

    fn compute_diffusion_operator(&self, mesh: &HeartMesh) -> DiffusionOperator {
        // Build sparse Laplacian with fiber anisotropy
        let n = mesh.vertices.len();
        let mut neighbors: Vec<Vec<(usize, f32)>> = vec![Vec::new(); n];

        // Build neighbor list from triangles
        for tri in &mesh.triangles {
            for i in 0..3 {
                let v1 = tri[i];
                let v2 = tri[(i + 1) % 3];

                let p1 = &mesh.vertices[v1];
                let p2 = &mesh.vertices[v2];
                let dist = p1.distance_to(p2);

                // Get fiber-aligned conductivity
                let fiber = &mesh.fibers[v1];
                let edge = Vector3D::new(p2.x - p1.x, p2.y - p1.y, p2.z - p1.z);
                let cos_angle =
                    (fiber.x * edge.x + fiber.y * edge.y + fiber.z * edge.z) / (dist + 1e-8);

                // Interpolate between longitudinal and transverse conductivity
                let sigma = self.conductivity.sigma_l * cos_angle.abs()
                    + self.conductivity.sigma_t * (1.0 - cos_angle.abs());

                let weight = sigma / (dist * dist + 1e-8);

                neighbors[v1].push((v2, weight));
            }
        }

        DiffusionOperator { neighbors }
    }

    fn apply_diffusion(&self, voltage: &[f32], op: &DiffusionOperator) -> Vec<f32> {
        let n = voltage.len();
        let mut result = vec![0.0_f32; n];

        for i in 0..n {
            let mut sum = 0.0;
            let mut total_weight = 0.0;

            for &(j, w) in &op.neighbors[i] {
                sum += w * (voltage[j] - voltage[i]);
                total_weight += w;
            }

            // Normalize and scale
            if total_weight > 0.0 {
                result[i] = sum / total_weight.max(1e-8);
            }
        }

        result
    }

    fn compute_ionic_currents(
        &self,
        voltage: &[f32],
        states: &[IonicState],
    ) -> (Vec<f32>, Vec<f32>) {
        let n = voltage.len();
        let mut dv = vec![0.0_f32; n];
        let mut dw = vec![0.0_f32; n];

        for i in 0..n {
            let (ion_dv, ion_dw) = states[i].compute_currents(voltage[i]);
            dv[i] = ion_dv;
            dw[i] = ion_dw;
        }

        (dv, dw)
    }

    fn compute_stimulus(
        &self,
        mesh: &HeartMesh,
        protocol: &StimulationProtocol,
        t: f32,
    ) -> Vec<f32> {
        let n = mesh.vertices.len();
        let mut stim = vec![0.0_f32; n];

        for site in &protocol.sites {
            // Check if we're within a stimulus pulse
            let mut is_active = false;
            for &stim_time in &site.times {
                if t >= stim_time && t < stim_time + site.duration {
                    is_active = true;
                    break;
                }
            }

            if is_active {
                // Apply stimulus to vertices within radius
                for (i, vertex) in mesh.vertices.iter().enumerate() {
                    let dist = vertex.distance_to(&site.center);
                    if dist < site.radius {
                        stim[i] = site.current;
                    }
                }
            }
        }

        stim
    }
}

/// Sparse diffusion operator.
#[derive(Debug)]
struct DiffusionOperator {
    neighbors: Vec<Vec<(usize, f32)>>,
}

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

    #[test]
    fn test_solver_creation() {
        let config = SimulationConfig::default();
        let solver = MonodomainSolver::new(&config);
        assert!(solver.dt > 0.0);
    }

    #[test]
    fn test_solve() {
        let config = SimulationConfig {
            total_time: 10.0,
            output_interval: 5.0,
            ..Default::default()
        };
        let solver = MonodomainSolver::new(&config);

        let mesh = cardiosim_shared::get_sample_heart_mesh();
        let protocol = cardiosim_shared::get_sample_protocol();

        let result = solver.solve(&mesh, &protocol);
        assert!(result.is_ok());

        let fields = result.unwrap();
        assert!(fields.len() >= 2);
    }

    #[test]
    fn test_diffusion_operator() {
        let config = SimulationConfig::default();
        let solver = MonodomainSolver::new(&config);
        let mesh = cardiosim_shared::get_sample_heart_mesh();

        let op = solver.compute_diffusion_operator(&mesh);
        assert_eq!(op.neighbors.len(), mesh.vertices.len());
    }

    #[test]
    fn test_stimulus() {
        let config = SimulationConfig::default();
        let solver = MonodomainSolver::new(&config);
        let mesh = cardiosim_shared::get_sample_heart_mesh();
        let protocol = cardiosim_shared::get_sample_protocol();

        let stim = solver.compute_stimulus(&mesh, &protocol, 0.5);
        assert!(!stim.is_empty());
        // Some vertices should be stimulated
        assert!(stim.iter().any(|&s| s > 0.0));
    }
}