#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Module for Structure-of-Arrays (SoA) container for simulation bodies and a
proxy for AoS-like access.
The `Bodies` class uses an SoA layout (e.g., all masses in one array, all
positions in another) for cache-efficient computations. The `BodyProxy` class
provides a convenient AoS-like interface (`body.x`, `body.v`) for a single
body, combining the best of both worlds.
:module: body
:author: Le Bars, Yoann
"""
# Defines the public names of this module.
__all__ = ['Bodies', 'BodyProxy']
import numpy as np
from scipy.spatial.transform import Rotation
# pylint: disable=too-many-instance-attributes
[docs]
class Bodies:
"""
Structure-of-Arrays (SoA) container for all bodies in the simulation.
This layout improves cache performance by storing related data contiguously.
:ivar np.ndarray m: Mass of each body.
:ivar np.ndarray inv_m: Inverse mass of each body.
:ivar np.ndarray r: Radius of each body.
:ivar np.ndarray i_inv: Inverse moment of inertia of each body.
:ivar np.ndarray x: Position of each body.
:ivar np.ndarray v: Velocity of each body.
:ivar np.ndarray a: Acceleration of each body.
:ivar np.ndarray q: Orientation of each body.
:ivar np.ndarray omega: Angular velocity of each body.
:ivar np.ndarray alpha: Angular acceleration of each body.
:ivar int _size: Number of bodies.
"""
m: np.ndarray
inv_m: np.ndarray
r: np.ndarray
i_inv: np.ndarray
x: np.ndarray
v: np.ndarray
a: np.ndarray
q: np.ndarray
omega: np.ndarray
alpha: np.ndarray
_size: int
def __init__(self, n: int):
"""
Class initialiser.
:param int n: Number of bodies.
"""
self.m = np.zeros(n)
self.inv_m = np.zeros(n)
self.r = np.zeros(n)
self.i_inv = np.zeros(n)
self.x = np.zeros((n, 3))
self.v = np.zeros((n, 3))
self.a = np.zeros((n, 3))
self.q = np.array([Rotation.identity()] * n)
self.omega = np.zeros((n, 3))
self.alpha = np.zeros((n, 3))
self._size = 0
[docs]
def emplace_back(self, m: float, r: float, x0: np.ndarray, v0: np.ndarray) -> None:
"""
Adds a new body to the data structure.
:param float m: Body mass.
:param float r: Body radius.
:param np.ndarray x0: Body position.
:param np.ndarray v0: Body velocity.
"""
if self._size >= len(self.m):
raise IndexError("Bodies container is full.")
# Index of the new body.
idx = self._size
self.m[idx] = m
self.inv_m[idx] = 1.0 / m if m > 0 else 0.0
self.r[idx] = r
# The moment of inertia for a solid sphere is I = (2/5) * m * r^2.
if self.inv_m[idx] > 0 and r > 0:
inertia = 0.4 * m * r * r
self.i_inv[idx] = 1.0 / inertia
else:
self.i_inv[idx] = 0.0
self.x[idx] = x0
self.v[idx] = v0
# a, q, omega, alpha are already initialized to zero/identity.
self._size += 1
def _get_quaternion_derivative(self, omega: np.ndarray, q_coeffs: np.ndarray) -> np.ndarray:
"""
Calculates the time derivative of a batch of quaternions.
The rate of change of an orientation quaternion `q` due to an angular velocity `ω`
is given by the equation: dq/dt = 0.5 * ω_q * q, where `ω_q` is the angular
velocity expressed as a pure quaternion [0, ω_x, ω_y, ω_z].
This function returns `2 * dq/dt` for convenience.
:param np.ndarray omega: Array of angular velocities.
:param np.ndarray q_coeffs: Array of quaternion coefficients [x, y, z, w].
:returns: The quaternion derivative multiplied by 2.
:rtype: np.ndarray
"""
q_vec = q_coeffs[:, :3]
q_w = q_coeffs[:, 3, np.newaxis]
# Vector part of the derivative product: w*ω + ω x v
delta_q_vec = q_w * omega + np.cross(omega, q_vec)
# Scalar part of the derivative product: -ω · v
delta_q_w = -np.sum(omega * q_vec, axis=1)
return np.hstack((delta_q_vec, delta_q_w[:, np.newaxis]))
def _update_rotations(self, dt: float) -> None:
"""
Vectorized update of orientations for all bodies using quaternion integration.
:param float dt: Time step.
"""
# This is an optimization: only update bodies that are actually rotating.
active_mask = np.sum(self.omega**2, axis=1) > 1e-12
if not np.any(active_mask):
return
# --- 1. Filter data for active bodies ---
active_omega = self.omega[active_mask]
# Extract quaternion coefficients into a NumPy array for vectorized processing.
# This is much faster than creating a new vectorized Rotation object.
q_coeffs = np.array([q.as_quat() for q in self.q[active_mask]])
# --- 2. Integrate using 2nd-order Runge-Kutta (RK2 / Midpoint Method) ---
# This is more stable than Forward Euler.
# k1 = f(q_old)
# The derivative returns 2 * dq/dt, so we use 0.25*dt instead of 0.5*dt.
k1 = self._get_quaternion_derivative(active_omega, q_coeffs)
# q_mid = q_old + 0.5 * dt * (0.5 * k1)
q_mid = q_coeffs + 0.25 * dt * k1
# Normalize the midpoint quaternion to prevent drift.
q_mid /= np.linalg.norm(q_mid, axis=1, keepdims=True)
# k2 = f(q_mid)
# The derivative returns 2 * dq/dt, so we use 0.5*dt instead of dt.
k2 = self._get_quaternion_derivative(active_omega, q_mid)
# q_new = q_old + dt * (0.5 * k2)
new_q_coeffs = q_coeffs + 0.5 * dt * k2
# --- 3. Update the original array of Rotation objects ---
# Normalizing is crucial to counteract numerical drift that can cause the quaternion to
# lose its unit length. We normalize the raw coefficients and create new Rotation objects.
norms = np.linalg.norm(new_q_coeffs, axis=1, keepdims=True)
# Avoid division by zero for null quaternions
valid_norms = norms.flatten() > 1e-9
new_q_coeffs[valid_norms] /= norms[valid_norms]
self.q[active_mask] = [Rotation.from_quat(q) for q in new_q_coeffs]
[docs]
def integrate_part1(self, dt: float, dt_pos: float) -> None:
"""
Velocity Verlet: Part 1 (Vectorized)
Updates positions and half-updates velocities for all bodies.
:param float dt: Time step.
:param float dt_pos: Time step for position update.
"""
# Update positions using the full time step from the previous frame.
self.x += self.v * dt_pos + 0.5 * self.a * dt_pos * dt_pos
# Perform the first half-step update for velocities using the new dt.
self.omega += 0.5 * self.alpha * dt
self.v += 0.5 * self.a * dt
# Update rotations and angular velocities.
self._update_rotations(dt)
[docs]
def integrate_part2(self, dt: float) -> None:
"""
Velocity Verlet: Part 2 (Vectorized)
Completes the velocity updates for all bodies.
:param float dt: Time step.
"""
self.v += 0.5 * self.a * dt
self.omega += 0.5 * self.alpha * dt
def __len__(self) -> int:
"""
Access to the number of bodies.
:returns: Number of bodies.
:rtype: int
"""
return self._size
def __getitem__(self, index: int) -> 'BodyProxy':
"""
Access to a given body.
:param int index: Body index.
:returns: Body proxy.
:rtype: BodyProxy
"""
return BodyProxy(self, index)
def __iter__(self):
"""
Allows for correct iteration over the bodies.
"""
for i in range(self._size):
yield self[i]
[docs]
class BodyProxy:
"""
A proxy object that provides an AoS-like interface to a body
stored in the `Bodies` SoA container.
The `Bodies` class uses a Structure-of-Arrays (SoA) layout for performance, where all
positions are in one array, all velocities in another, etc. This is ideal for vectorized
physics calculations.
However, it can be inconvenient to work with a single body's properties. This `BodyProxy`
class solves that problem by implementing the Proxy design pattern. It provides a classic
Array-of-Structures (AoS) interface (e.g., `body.x`, `body.v`) for a single body, while
the underlying data remains in the efficient SoA format. This gives us the best of both
worlds: high-performance vectorized operations on the `Bodies` container and intuitive,
object-oriented access to individual bodies.
:ivar Bodies _bodies: The main SoA container.
:ivar int _index: The index of the body this proxy refers to.
"""
_bodies: Bodies
_index: int
def __init__(self, bodies: Bodies, index: int):
"""
Class initialiser.
:param Bodies bodies: The main SoA container.
:param int index: The index of the body this proxy refers to.
"""
self._bodies = bodies
self._index = index
# --- Getters ---
@property
def m(self) -> float:
"""Body mass."""
return self._bodies.m[self._index]
@property
def inv_m(self) -> float:
"""Inverse of the body mass."""
return self._bodies.inv_m[self._index]
@property
def r(self) -> float:
"""Body radius."""
return self._bodies.r[self._index]
@property
def i_inv(self) -> float:
"""Inverse of the moment of inertia."""
return self._bodies.i_inv[self._index]
@property
def x(self) -> np.ndarray:
"""Body position."""
return self._bodies.x[self._index]
@property
def v(self) -> np.ndarray:
"""Body velocity."""
return self._bodies.v[self._index]
@property
def a(self) -> np.ndarray:
"""Body acceleration."""
return self._bodies.a[self._index]
@property
def q(self) -> Rotation:
"""Body orientation."""
return self._bodies.q[self._index]
@property
def omega(self) -> np.ndarray:
"""Body angular velocity."""
return self._bodies.omega[self._index]
@property
def alpha(self) -> np.ndarray:
"""Body angular acceleration."""
return self._bodies.alpha[self._index]
# --- Setters for mutable properties ---
[docs]
def set_acceleration(self, new_a: np.ndarray) -> None:
"""
Set the acceleration.
:param np.ndarray new_a: New acceleration vector.
"""
self._bodies.a[self._index] = new_a
[docs]
def accumulate_acceleration(self, accel: np.ndarray) -> None:
"""
Accumulate an acceleration.
:param np.ndarray accel: Acceleration vector.
"""
self._bodies.a[self._index] += accel
[docs]
def accumulate_angular_acceleration(self, angular_accel: np.ndarray) -> None:
"""
Accumulate an angular acceleration.
:param np.ndarray angular_accel: Angular acceleration vector.
"""
self._bodies.alpha[self._index] += angular_accel
[docs]
def apply_impulse(self, j_total: np.ndarray, r_vec: np.ndarray, epsilon: float) -> None:
"""
Apply an impulse to the body, updating its linear and angular velocity based on the
impulse-momentum theorem.
:param np.ndarray j_total: Total impulse vector.
:param np.ndarray r_vec: Vector from the body's centre of mass to the point of impulse
application.
:param float epsilon: Small value to avoid division by zero.
"""
# --- Linear Impulse ---
# The linear impulse-momentum theorem states that an impulse J applied to an object
# changes its linear momentum (p = m*v) by Δp = J.
# Therefore, the change in velocity is Δv = J / m = J * inv_m.
self._bodies.v[self._index] += j_total * self.inv_m
# --- Angular Impulse ---
if self.i_inv > epsilon:
# The angular impulse-momentum theorem states that an angular impulse changes the
# angular momentum (L = I*ω). The angular impulse is the torque created by the
# linear impulse, which is τ = r x J.
# The change in angular velocity is Δω = I⁻¹ * (r x J).
self._bodies.omega[self._index] += self.i_inv * \
np.cross(r_vec, j_total)