#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Provides a Structure-of-Arrays (SoA) container for simulation bodies and a
proxy for Array-of-Structures (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 (n, 3).
: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 np.ndarray torque: Net torque applied to each body in a frame.
: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
torque: 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.zeros((n, 4))
self.q[:, 3] = 1.0 # Identity quaternions (x, y, z, w)
self.omega = np.zeros((n, 3))
self.alpha = np.zeros((n, 3))
self.torque = 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
[docs]
def set_size(self, size: int) -> None:
"""
Sets the number of bodies in the container.
This method should be used with caution and is primarily for initialization
scenarios where bodies are placed in bulk.
:param int size: The new size value.
"""
if size < 0 or size > len(self.m):
raise ValueError(f"Size must be between 0 and {len(self.m)}")
self._size = size
def _get_quaternion_derivative(self, angular_velocities: np.ndarray,
orientations: 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 `omega`
is given by the equation: `dq/dt = 0.5 * omega_q * q`, where `omega_q` is the angular
velocity expressed as a pure quaternion `[omega, 0]`.
This function is a helper for the RK2 integrator in `_update_rotations` and
returns `2 * dq/dt` to simplify the integration steps.
:param np.ndarray angular_velocities: Array of angular velocities (shape: `(k, 3)`).
:param np.ndarray orientations: Array of orientation quaternions [x, y, z, w]
(shape: `(k, 4)`).
:returns: The quaternion derivative multiplied by 2.
:rtype: np.ndarray
"""
# Decompose the orientation quaternions into their vector (q_vec) and scalar (q_w) parts.
q_vec = orientations[:, :3]
# Keep as a column vector for broadcasting
q_w = orientations[:, 3, np.newaxis]
# The product of `omega_q * q` is given by the formula:
# Vector part: q_w * omega + cross(omega, q_vec)
# Scalar part: -dot(omega, q_vec)
# Calculate the vector part of the derivative.
derivative_vec = q_w * angular_velocities + \
np.cross(angular_velocities, q_vec)
# 2. Calculate the scalar part of the derivative.
derivative_w = -np.sum(angular_velocities *
q_vec, axis=1, keepdims=True)
# 3. Combine the vector and scalar parts back into a single array of quaternions.
return np.hstack((derivative_vec, derivative_w))
def _update_rotations(self, dt: float) -> None:
"""
Vectorised update of orientations for all bodies using quaternion integration.
:param float dt: Time step.
"""
# As an optimisation, we create a mask to update only the bodies that are
# actually rotating, avoiding unnecessary calculations.
active_mask = np.sum(self.omega**2, axis=1) > 1e-12
# --- 1. Filter data for active bodies ---
active_omega = self.omega[active_mask]
active_quaternions = 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, active_quaternions)
# q_mid = q_old + 0.5 * dt * (0.5 * k1)
q_mid = active_quaternions + 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 = active_quaternions + 0.5 * dt * k2
# --- 3. Update the original quaternion array ---
# Normalizing is crucial to counteract numerical drift that can cause the quaternion to
# lose its unit length.
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] = new_q_coeffs
[docs]
def integrate_part1(self, dt: float) -> None:
"""
Velocity Verlet: Part 1 (Vectorised).
Updates positions and half-updates velocities for all bodies.
:param float dt: Time step.
"""
self.x += self.v * dt + 0.5 * self.a * dt * dt
# Perform the first half-step update for velocities.
self.v += 0.5 * self.a * dt
# Perform the first half-step update for angular velocity.
self.omega += 0.5 * self.alpha * dt
self._update_rotations(dt)
[docs]
def integrate_part2(self, dt: float) -> None:
"""
Velocity Verlet: Part 2 (Vectorised).
Completes the velocity updates for all bodies.
:param float dt: Time step.
"""
self.v += 0.5 * self.a * dt
# Complete the second half-step for angular velocity.
self.omega += 0.5 * self.alpha * dt
[docs]
def apply_impulses(self, indices: np.ndarray, impulses: np.ndarray, contact_vectors: np.ndarray) -> None:
"""
Applies a batch of impulses to multiple bodies in a fully vectorized manner.
:param np.ndarray indices: The indices of the bodies to apply impulses to.
:param np.ndarray impulses: An array of impulse vectors (J) for each body.
:param np.ndarray contact_vectors: An array of vectors (r) from each body’s
centre of mass to the point of impulse application.
"""
if indices.size == 0:
return
# 1. Vectorized update of linear velocities.
# Δv = J * inv_m
self.v[indices] += impulses * self.inv_m[indices, np.newaxis]
# 2. Vectorized calculation of resulting torques.
# τ = r × J
torques = np.cross(contact_vectors, impulses)
np.add.at(self.torque, indices, torques)
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 vectorised
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 vectorised 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
def __eq__(self, other: object) -> bool:
"""
Compares this proxy with another object for equality.
Two `BodyProxy` instances are considered equal if they refer to the same body
(i.e., the same index within the same `Bodies` container).
:param object other: The object to compare with.
:returns: `True` if the objects are equal, `False` otherwise.
:rtype: bool
"""
if not isinstance(other, BodyProxy):
return NotImplemented
return self._bodies is other._bodies and self._index == other._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."""
# Create a Rotation object on-the-fly from the stored quaternion coefficients.
# This provides a convenient high-level object for consumers (like the Controller)
# while keeping the core simulation data in a high-performance NumPy array.
return Rotation.from_quat(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]
@property
def torque(self) -> np.ndarray:
"""Net torque on the body for the current frame."""
return self._bodies.torque[self._index]
# --- Setters for mutable properties ---
[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 dampen_velocity(self, linear_damping: float, angular_damping: float) -> None:
"""
Applies damping to linear and angular velocities to stabilise the simulation.
:param float linear_damping: Damping factor for linear velocity.
:param float angular_damping: Damping factor for angular velocity.
"""
self._bodies.v[self._index] *= (1.0 - linear_damping)
self._bodies.omega[self._index] *= (1.0 - angular_damping)
[docs]
def accumulate_torque(self, torque_vec: np.ndarray) -> None:
"""
Accumulate a torque vector.
:param np.ndarray torque_vec: Torque vector to add.
"""
self._bodies.torque[self._index] += torque_vec
[docs]
def apply_impulse(self, j_total: np.ndarray, r_vec: np.ndarray) -> np.ndarray:
"""
Applies a linear impulse to the body and returns the resulting torque.
Based on the impulse-momentum theorem, this method updates the body's linear velocity
and calculates the torque generated by the impulse, which will later be used to
update the angular velocity.
: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.
:returns: The torque vector generated by the impulse.
:rtype: np.ndarray
"""
self._bodies.v[self._index] += j_total * self.inv_m
# The torque (τ) generated by an impulse (J) applied at a position (r) relative to
# the centre of mass is given by τ = r × J. This torque will be used by the
# physics engine to calculate the change in angular velocity.
return np.cross(r_vec, j_total)