tree_d_3d.gravity

Barnes-Hut N-body solver for fast gravitational force calculations.

This module implements the Barnes-Hut algorithm for O(N log N) computation of gravitational forces in N-body simulations. It uses quadtree (2D) or octree (3D) spatial decomposition to approximate distant bodies.

References

Barnes, J., & Hut, P. (1986). A hierarchical O(N log N) force-calculation algorithm. Nature, 324(6096), 446-449.

Classes

Name	Description
BarnesHutSimulation	N-body simulation using Barnes-Hut algorithm.
BarnesHutTree	Quadtree-based Barnes-Hut N-body tree.
Body	Represents a body (particle) in N-body simulation.
Node	Represents a node in the quadtree.

BarnesHutSimulationclass#

View source

class BarnesHutSimulation(bodies: Sequence[Body], bounds: Quad, theta: float = 0.5, G: float = 1.0, softening: float = 0.1, dt: float = 0.01)

N-body simulation using Barnes-Hut algorithm.

This class manages particles and implements the main simulation loop using leapfrog integration with force calculation via the Barnes-Hut tree algorithm.

Gattributeinstance attribute#

View source

G: float = G

bodiesattributeinstance attribute#

View source

bodies: list[Body] = list(bodies)

boundsattributeinstance attribute#

View source

bounds: Quad = bounds

dtattributeinstance attribute#

View source

dt: float = dt

softeningattributeinstance attribute#

View source

softening: float = softening

thetaattributeinstance attribute#

View source

theta: float = theta

build_treemethod#

View source

def build_tree() -> BarnesHutTree

Build a new Barnes-Hut tree from current particle positions.

Returns

BarnesHutTree

New BarnesHutTree

get_center_of_massmethod#

View source

def get_center_of_mass() -> np.ndarray

Calculate center of mass of all bodies.

Returns

np.ndarray

Position vector of center of mass

get_kinetic_energymethod#

View source

def get_kinetic_energy() -> float

Calculate total kinetic energy of all bodies.

Returns

float

Total kinetic energy

stepmethod#

View source

def step() -> BarnesHutTree

Advance simulation by one timestep using leapfrog integration.

Returns

BarnesHutTree

The updated Barnes-Hut tree (for visualization/diagnostics)

BarnesHutTreeclass#

View source

class BarnesHutTree(root_quad: Quad)

Quadtree-based Barnes-Hut N-body tree.

rootattributeinstance attribute#

View source

root: Node = Node(root_quad)

collect_nodesmethod#

View source

def collect_nodes() -> list[Node]

Collect all nodes in the tree (for visualization).

Returns

list[Node]

List of all non-empty nodes

compute_force_onmethod#

View source

def compute_force_on(body: Body, theta: float = 0.5, G: float = 1.0, softening: float = 0.1) -> None

Compute force on a body from all other bodies in tree.

Parameters

Name	Type	Default	Description
body	Body	required	Body to compute force on
theta	float	0.5	Opening angle (trade-off accuracy vs speed, typically 0.3-0.9)
G	float	1.0	Gravitational constant
softening	float	0.1	Softening length to prevent force singularities

insertmethod#

View source

def insert(body: Body) -> None

Insert a body into the tree.

Parameters

Name	Type	Description
body	Body	Body to insert

Bodyclass#

View source

class Body(x: float, y: float, vx: float = 0.0, vy: float = 0.0, mass: float = 1.0)

Represents a body (particle) in N-body simulation.

forceattributeinstance attribute#

View source

force: np.ndarray = np.zeros(2, dtype=np.float64)

massattributeinstance attribute#

View source

mass: float = float(mass)

posattributeinstance attribute#

View source

pos: np.ndarray = np.array([x, y], dtype=np.float64)

velattributeinstance attribute#

View source

vel: np.ndarray = np.array([vx, vy], dtype=np.float64)

add_forcemethod#

View source

def add_force(fx: float, fy: float) -> None

Add force component to the body.

Parameters

Name	Description
fx, fy	Force components to add

reset_forcemethod#

View source

def reset_force() -> None

Reset force to zero.

Nodeclass#

View source

class Node(quad: Quad)

Represents a node in the quadtree.

NEattributeinstance attribute#

View source

NE: Node | None = None

NWattributeinstance attribute#

View source

NW: Node | None = None

SEattributeinstance attribute#

View source

SE: Node | None = None

SWattributeinstance attribute#

View source

SW: Node | None = None

bodyattributeinstance attribute#

View source

body: Body | None = None

center_of_massattributeinstance attribute#

View source

center_of_mass: np.ndarray = np.array([0.0, 0.0], dtype=np.float64)

massattributeinstance attribute#

View source

mass: float = 0.0

quadattributeinstance attribute#

View source

quad: Quad = quad

is_externalmethod#

View source

def is_external() -> bool

Check if this is a leaf node (no children).