use crate::math::{Fix128, Vec3Fix};
use crate::solver::{BodyType, PhysicsWorld, RigidBody};
use crate::world_participant::{
ObservationSink, Participant, ParticipantFault, ParticipantKind, StateError, SubstepCtx,
};
#[cfg(not(feature = "std"))]
use alloc::{vec, vec::Vec};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum NBodyError {
NonPositiveGravitationalConstant,
NegativeSoftening,
NegativeMass,
LengthMismatch,
}
impl core::fmt::Display for NBodyError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let msg = match self {
Self::NonPositiveGravitationalConstant => "the gravitational constant must be positive",
Self::NegativeSoftening => "the softening length must not be negative",
Self::NegativeMass => "a mass must not be negative",
Self::LengthMismatch => "positions, velocities and masses must have the same length",
};
f.write_str(msg)
}
}
#[cfg(feature = "std")]
impl std::error::Error for NBodyError {}
fn check_masses(masses: &[Fix128]) -> Result<(), NBodyError> {
if masses.iter().any(|m| m.is_negative()) {
Err(NBodyError::NegativeMass)
} else {
Ok(())
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct DirectSum {
gravitational_constant: Fix128,
softening: Fix128,
}
impl DirectSum {
pub fn new(gravitational_constant: Fix128, softening: Fix128) -> Result<Self, NBodyError> {
if gravitational_constant <= Fix128::ZERO {
return Err(NBodyError::NonPositiveGravitationalConstant);
}
if softening.is_negative() {
return Err(NBodyError::NegativeSoftening);
}
Ok(Self {
gravitational_constant,
softening,
})
}
#[must_use]
pub fn gravitational_constant(&self) -> Fix128 {
self.gravitational_constant
}
#[must_use]
pub fn softening(&self) -> Fix128 {
self.softening
}
pub fn accelerations(
&self,
positions: &[Vec3Fix],
masses: &[Fix128],
out: &mut Vec<Vec3Fix>,
) -> Result<(), NBodyError> {
if positions.len() != masses.len() {
return Err(NBodyError::LengthMismatch);
}
check_masses(masses)?;
let n = positions.len();
out.clear();
out.resize(n, Vec3Fix::ZERO);
let eps2 = self.softening * self.softening;
for i in 0..n {
for j in (i + 1)..n {
let d = positions[j] - positions[i];
let r2 = d.length_squared() + eps2;
if r2.is_zero() {
continue;
}
let inv_r = Fix128::ONE / r2.sqrt();
let pull = d * (self.gravitational_constant * inv_r * inv_r * inv_r);
out[i] = out[i] + pull * masses[j];
out[j] = out[j] - pull * masses[i];
}
}
Ok(())
}
pub fn potential_energy(
&self,
positions: &[Vec3Fix],
masses: &[Fix128],
) -> Result<Fix128, NBodyError> {
if positions.len() != masses.len() {
return Err(NBodyError::LengthMismatch);
}
check_masses(masses)?;
let eps2 = self.softening * self.softening;
let mut total = Fix128::ZERO;
for i in 0..positions.len() {
for j in (i + 1)..positions.len() {
let r2 = (positions[j] - positions[i]).length_squared() + eps2;
if r2.is_zero() {
continue;
}
total = total - self.gravitational_constant * masses[i] * masses[j] / r2.sqrt();
}
}
Ok(total)
}
pub fn kick_bodies(&self, bodies: &mut [RigidBody], dt: Fix128) {
let _ = self.kick_dynamic(bodies, dt);
}
fn kick_dynamic(&self, bodies: &mut [RigidBody], dt: Fix128) -> usize {
let index: Vec<usize> = (0..bodies.len())
.filter(|&i| {
bodies[i].body_type == BodyType::Dynamic && bodies[i].inv_mass > Fix128::ZERO
})
.collect();
if index.len() < 2 {
return index.len();
}
let positions: Vec<Vec3Fix> = index.iter().map(|&i| bodies[i].position).collect();
let masses: Vec<Fix128> = index
.iter()
.map(|&i| Fix128::ONE / bodies[i].inv_mass)
.collect();
let mut acc = Vec::new();
if self.accelerations(&positions, &masses, &mut acc).is_err() {
return 0;
}
for (k, &i) in index.iter().enumerate() {
bodies[i].velocity = bodies[i].velocity + acc[k] * dt;
}
index.len()
}
pub fn step_world(&self, world: &mut PhysicsWorld, dt: Fix128) {
if dt <= Fix128::ZERO {
return;
}
let half = dt.half();
if self.kick_dynamic(&mut world.bodies, half) >= 2 {
for i in 0..world.bodies.len() {
if world.bodies[i].body_type == BodyType::Dynamic {
world.wake_body(i);
}
}
}
world.step(dt);
let _ = self.kick_dynamic(&mut world.bodies, half);
}
}
pub const NBODY_PARTICIPANT_KIND: ParticipantKind =
ParticipantKind::new(u32::from_be_bytes(*b"NBDY"));
pub const NBODY_OBS_POTENTIAL: u32 = 0;
pub const NBODY_OBS_BODIES: u32 = 1;
const NBODY_STATE_VERSION: u32 = 1;
const NBODY_STATE_LEN: usize = 4 + 8 + 16 + 8;
#[derive(Clone, Debug)]
pub struct NBodyParticipant {
law: DirectSum,
digest: u64,
potential: Fix128,
bodies: u64,
}
impl NBodyParticipant {
#[must_use]
pub fn new(law: DirectSum) -> Self {
let digest = fnv_fix(
fnv_fix(FNV_OFFSET, law.gravitational_constant),
law.softening,
);
Self {
law,
digest,
potential: Fix128::ZERO,
bodies: 0,
}
}
#[must_use]
pub fn law(&self) -> &DirectSum {
&self.law
}
fn pair_forces(&self, bodies: &[RigidBody], index: &[usize]) -> Option<(Vec<Vec3Fix>, Fix128)> {
let n = index.len();
let mut forces = vec![Vec3Fix::ZERO; n];
let mut potential = Fix128::ZERO;
if n < 2 {
return Some((forces, potential));
}
let masses = index
.iter()
.map(|&i| Fix128::ONE.checked_div(bodies[i].inv_mass))
.collect::<Option<Vec<Fix128>>>()?;
let g = self.law.gravitational_constant;
let eps2 = self.law.softening.checked_mul(self.law.softening)?;
for a in 0..n {
for b in (a + 1)..n {
let d = checked_sub(bodies[index[b]].position, bodies[index[a]].position)?;
let r2 = d.checked_length_squared()?.checked_add(eps2)?;
if r2.is_zero() {
continue;
}
let r = r2.sqrt();
let inv_r = Fix128::ONE.checked_div(r)?;
let s = g
.checked_mul(inv_r)?
.checked_mul(inv_r)?
.checked_mul(inv_r)?;
let force = d
.checked_scale(s)?
.checked_scale(masses[a])?
.checked_scale(masses[b])?;
forces[a] = forces[a].checked_add(force)?;
forces[b] = checked_sub(forces[b], force)?;
let pair = g
.checked_mul(masses[a])?
.checked_mul(masses[b])?
.checked_div(r)?;
potential = potential.checked_add(pair.checked_mul(Fix128::NEG_ONE)?)?;
}
}
Some((forces, potential))
}
}
fn checked_sub(a: Vec3Fix, b: Vec3Fix) -> Option<Vec3Fix> {
a.checked_add(Vec3Fix::new(
b.x.checked_mul(Fix128::NEG_ONE)?,
b.y.checked_mul(Fix128::NEG_ONE)?,
b.z.checked_mul(Fix128::NEG_ONE)?,
))
}
const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
fn fnv(mut h: u64, bytes: &[u8]) -> u64 {
for &b in bytes {
h ^= u64::from(b);
h = h.wrapping_mul(0x0000_0100_0000_01b3);
}
h
}
fn fnv_fix(h: u64, f: Fix128) -> u64 {
fnv(fnv(h, &f.hi.to_le_bytes()), &f.lo.to_le_bytes())
}
fn put_fix(out: &mut Vec<u8>, f: Fix128) {
out.extend_from_slice(&f.hi.to_le_bytes());
out.extend_from_slice(&f.lo.to_le_bytes());
}
fn get_fix(bytes: &[u8], at: usize) -> Fix128 {
let mut hi = [0u8; 8];
let mut lo = [0u8; 8];
hi.copy_from_slice(&bytes[at..at + 8]);
lo.copy_from_slice(&bytes[at + 8..at + 16]);
Fix128 {
hi: i64::from_le_bytes(hi),
lo: u64::from_le_bytes(lo),
}
}
fn get_u64(bytes: &[u8], at: usize) -> u64 {
let mut b = [0u8; 8];
b.copy_from_slice(&bytes[at..at + 8]);
u64::from_le_bytes(b)
}
impl Participant for NBodyParticipant {
fn kind(&self) -> ParticipantKind {
NBODY_PARTICIPANT_KIND
}
fn substep(&mut self, ctx: &mut SubstepCtx<'_>, _h: Fix128) -> Result<(), ParticipantFault> {
let index: Vec<usize> = ctx
.bodies()
.iter()
.enumerate()
.filter(|(_, b)| b.body_type == BodyType::Dynamic && b.inv_mass > Fix128::ZERO)
.map(|(i, _)| i)
.collect();
let (forces, potential) = self
.pair_forces(ctx.bodies(), &index)
.ok_or(ParticipantFault::OutOfRange)?;
for (&i, &f) in index.iter().zip(&forces) {
ctx.add_force(i, f)
.map_err(|_| ParticipantFault::InvalidState)?;
}
self.potential = potential;
self.bodies = index.len() as u64;
Ok(())
}
fn observe(&self, out: &mut ObservationSink) {
out.push(NBODY_OBS_POTENTIAL, self.potential);
out.push(NBODY_OBS_BODIES, Fix128::from_int(self.bodies as i64));
}
fn write_state(&self, out: &mut Vec<u8>) {
out.extend_from_slice(&NBODY_STATE_VERSION.to_le_bytes());
out.extend_from_slice(&self.digest.to_le_bytes());
put_fix(out, self.potential);
out.extend_from_slice(&self.bodies.to_le_bytes());
}
fn check_state(&self, bytes: &[u8]) -> Result<(), StateError> {
if bytes.len() != NBODY_STATE_LEN {
return Err(StateError::Length {
expected: NBODY_STATE_LEN,
found: bytes.len(),
});
}
let mut version = [0u8; 4];
version.copy_from_slice(&bytes[0..4]);
if u32::from_le_bytes(version) != NBODY_STATE_VERSION || get_u64(bytes, 4) != self.digest {
return Err(StateError::InvalidValue);
}
Ok(())
}
fn read_state(&mut self, bytes: &[u8]) {
self.potential = get_fix(bytes, 12);
self.bodies = get_u64(bytes, 28);
}
}
pub fn kinetic_energy(velocities: &[Vec3Fix], masses: &[Fix128]) -> Result<Fix128, NBodyError> {
if velocities.len() != masses.len() {
return Err(NBodyError::LengthMismatch);
}
check_masses(masses)?;
Ok(velocities
.iter()
.zip(masses)
.fold(Fix128::ZERO, |acc, (v, &m)| {
acc + (m * v.length_squared()).half()
}))
}
pub fn total_momentum(velocities: &[Vec3Fix], masses: &[Fix128]) -> Result<Vec3Fix, NBodyError> {
if velocities.len() != masses.len() {
return Err(NBodyError::LengthMismatch);
}
check_masses(masses)?;
Ok(velocities
.iter()
.zip(masses)
.fold(Vec3Fix::ZERO, |acc, (&v, &m)| acc + v * m))
}
#[derive(Clone, Debug, Default)]
pub struct VelocityVerlet {
acc: Vec<Vec3Fix>,
cached_positions: Vec<Vec3Fix>,
cached_masses: Vec<Fix128>,
cached_law: Option<DirectSum>,
}
impl VelocityVerlet {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn step(
&mut self,
law: &DirectSum,
positions: &mut [Vec3Fix],
velocities: &mut [Vec3Fix],
masses: &[Fix128],
dt: Fix128,
) -> Result<(), NBodyError> {
if positions.len() != velocities.len() || positions.len() != masses.len() {
return Err(NBodyError::LengthMismatch);
}
check_masses(masses)?;
let fresh = self.cached_law == Some(*law)
&& self.cached_positions.as_slice() == &*positions
&& self.cached_masses.as_slice() == masses;
if !fresh {
law.accelerations(positions, masses, &mut self.acc)?;
}
let half = dt.half();
for ((x, v), a) in positions
.iter_mut()
.zip(velocities.iter_mut())
.zip(&self.acc)
{
*v = *v + *a * half;
*x = *x + *v * dt;
}
law.accelerations(positions, masses, &mut self.acc)?;
for (v, a) in velocities.iter_mut().zip(&self.acc) {
*v = *v + *a * half;
}
self.cached_positions.clear();
self.cached_positions.extend_from_slice(positions);
self.cached_masses.clear();
self.cached_masses.extend_from_slice(masses);
self.cached_law = Some(*law);
Ok(())
}
}
#[cfg(test)]
mod participant_tests {
use super::*;
#[cfg(not(feature = "std"))]
use alloc::vec;
fn resident(g: i64, eps: Fix128) -> NBodyParticipant {
NBodyParticipant::new(DirectSum::new(Fix128::from_int(g), eps).expect("law"))
}
#[test]
fn digest_depends_on_the_law() {
let a = resident(1, Fix128::ZERO);
assert_ne!(a.digest, resident(2, Fix128::ZERO).digest);
assert_ne!(a.digest, resident(1, Fix128::from_ratio(1, 8)).digest);
assert_eq!(a.kind(), NBODY_PARTICIPANT_KIND);
assert_eq!(a.kind().get(), 0x4e42_4459);
}
#[test]
fn payload_round_trip_and_refusals() {
let mut a = resident(1, Fix128::ZERO);
a.potential = Fix128::from_ratio(-3, 4);
a.bodies = 7;
let mut bytes = Vec::new();
a.write_state(&mut bytes);
assert_eq!(bytes.len(), NBODY_STATE_LEN);
let mut b = resident(1, Fix128::ZERO);
assert_eq!(b.check_state(&bytes), Ok(()));
b.read_state(&bytes);
assert_eq!((b.potential, b.bodies), (a.potential, 7));
assert_eq!(
b.check_state(&bytes[..35]),
Err(StateError::Length {
expected: 36,
found: 35
})
);
let mut version = bytes.clone();
version[0] = 2;
assert_eq!(b.check_state(&version), Err(StateError::InvalidValue));
assert_eq!(
resident(3, Fix128::ZERO).check_state(&bytes),
Err(StateError::InvalidValue)
);
}
#[test]
fn checked_sub_refuses_a_wrap() {
let one = Vec3Fix::new(Fix128::ONE, Fix128::ZERO, Fix128::ZERO);
assert_eq!(checked_sub(one, one), Some(Vec3Fix::ZERO));
let low = Vec3Fix::new(Fix128::from_int(i64::MIN / 2), Fix128::ZERO, Fix128::ZERO);
let high = Vec3Fix::new(
Fix128::from_int(i64::MAX / 2 + 2),
Fix128::ZERO,
Fix128::ZERO,
);
assert_eq!(checked_sub(low, high), None);
}
#[test]
fn pair_forces_check_the_range() {
let far = [
RigidBody::new_dynamic(Vec3Fix::ZERO, Fix128::ONE),
RigidBody::new_dynamic(
Vec3Fix::new(Fix128::from_int(2), Fix128::ZERO, Fix128::ZERO),
Fix128::ONE,
),
];
let (f, u) = resident(1, Fix128::ZERO)
.pair_forces(&far, &[0, 1])
.expect("in range");
assert_eq!(
f,
vec![
Vec3Fix::new(Fix128::from_ratio(1, 4), Fix128::ZERO, Fix128::ZERO),
Vec3Fix::new(Fix128::from_ratio(-1, 4), Fix128::ZERO, Fix128::ZERO),
]
);
assert_eq!(u, Fix128::from_ratio(-1, 2));
let near = [
RigidBody::new_dynamic(Vec3Fix::ZERO, Fix128::ONE),
RigidBody::new_dynamic(
Vec3Fix::new(Fix128::from_ratio(1, 1024), Fix128::ZERO, Fix128::ZERO),
Fix128::ONE,
),
];
assert!(resident(1 << 60, Fix128::ZERO)
.pair_forces(&near, &[0, 1])
.is_none());
}
}