use core::fmt;
use crate::threading::try_for_each_pair_chunk_mut;
use crate::engines::soa::phys_obj::{AttrsError, PhysObj};
use crate::models::particles::attrs::{ATTR_A, ATTR_R, ATTR_V, ParticleSelection};
use crate::models::particles::state::{
BorrowedMasks, ParticleStateError, mask_labels, validate_vector_attr_f64,
};
#[derive(Debug, Clone, PartialEq)]
pub enum IntegratorError {
InvalidDt {
dt: f64,
},
State(ParticleStateError),
}
impl From<AttrsError> for IntegratorError {
fn from(value: AttrsError) -> Self {
Self::State(value.into())
}
}
impl From<ParticleStateError> for IntegratorError {
fn from(value: ParticleStateError) -> Self {
Self::State(value)
}
}
impl fmt::Display for IntegratorError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InvalidDt { dt } => {
write!(
f,
"integration time step must be finite and positive; got {dt}"
)
}
Self::State(error) => write!(f, "invalid particle state: {error}"),
}
}
}
impl std::error::Error for IntegratorError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::State(error) => Some(error),
_ => None,
}
}
}
pub trait Integrator: Send {
fn apply(&mut self, objects: &mut PhysObj, dt: f64) -> Result<(), IntegratorError>;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct ExplicitEuler;
#[derive(Debug, Clone, Copy, Default)]
pub struct SemiImplicitEuler;
fn apply_euler(objects: &mut PhysObj, dt: f64, semi_implicit: bool) -> Result<(), IntegratorError> {
if !dt.is_finite() || dt <= 0.0 {
return Err(IntegratorError::InvalidDt { dt });
}
let v = objects.core.get::<f64>(ATTR_V)?;
let (dim, n) = (v.dim(), v.num_vectors());
validate_vector_attr_f64(objects, ATTR_A, dim, n)?;
validate_vector_attr_f64(objects, ATTR_R, dim, n)?;
let flags = mask_labels(objects, ParticleSelection::AliveOnly);
let ([r, v, a], flags) = objects
.core
.get_mixed_mut::<f64, u8, 3, 2>([ATTR_R, ATTR_V, ATTR_A], flags)?;
let masks = BorrowedMasks::new(flags, n)?;
let acceleration = a.borrow_values();
r.edit_values(|positions| {
v.edit_values(|velocities| {
try_for_each_pair_chunk_mut(
positions,
velocities,
dim,
|start, positions, velocities| {
if masks.all_active_in(start / dim, (start + positions.len()) / dim) {
crate::math::kernels::euler_f64(
positions,
velocities,
&acceleration[start..start + positions.len()],
dt,
semi_implicit,
);
return Ok(());
}
for (row, (r, v)) in positions
.chunks_exact_mut(dim)
.zip(velocities.chunks_exact_mut(dim))
.enumerate()
{
let index = start / dim + row;
if masks.should_skip(index) {
continue;
}
let a = &acceleration[index * dim..][..dim];
for ((r, v), &a) in r.iter_mut().zip(v).zip(a) {
let old = *v;
*v += a * dt;
*r += if semi_implicit { *v * dt } else { old * dt };
}
}
Ok::<_, std::convert::Infallible>(())
},
)
.expect("infallible validated integration");
})
});
Ok(())
}
impl Integrator for ExplicitEuler {
fn apply(&mut self, objects: &mut PhysObj, dt: f64) -> Result<(), IntegratorError> {
apply_euler(objects, dt, false)
}
}
impl Integrator for SemiImplicitEuler {
fn apply(&mut self, objects: &mut PhysObj, dt: f64) -> Result<(), IntegratorError> {
apply_euler(objects, dt, true)
}
}