Clean up and comment
This commit is contained in:
parent
8bb6ba19d5
commit
6822df0432
SSH key fingerprint:
SHA256:cj/YZ5VT+QOKncqSkx+ibKTIn0Obg7OIzwzl9BL8EO8
5 changed files with 249 additions and 152 deletions
185
data.py
185
data.py
|
|
@ -1,93 +1,148 @@
|
|||
from pathlib import Path
|
||||
|
||||
import matplotlib.animation as animation
|
||||
import matplotlib.cm as cm
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import physics
|
||||
from matplotlib.collections import PathCollection
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from physics import n_body_matrix
|
||||
|
||||
|
||||
def parse_csv(filename: str, dimensions=2):
|
||||
if not (1 < dimensions < 4):
|
||||
raise ValueError(f"Can only show 2or 3 dimensional scenes, not {dimensions}")
|
||||
with open(filename, 'r') as file:
|
||||
lines = file.read().strip().splitlines()
|
||||
pos = np.zeros((len(lines), dimensions))
|
||||
vel = np.zeros((len(lines), dimensions))
|
||||
rad = np.zeros((len(lines), 1))
|
||||
for i, values in enumerate(map(lambda l: map(float, l.split(',')), lines)):
|
||||
def parse_csv(file: Path) -> tuple[NDArray, NDArray, NDArray]:
|
||||
"""
|
||||
Reads the starting conditions of a simulation from a CSV
|
||||
:param file: The CSV file
|
||||
:return: The position, velocity, and radius matrices, in 2 or 3 dimensions
|
||||
"""
|
||||
# Get text from file
|
||||
lines = file.read_text().strip().splitlines()
|
||||
field_count = len(lines[0].split(","))
|
||||
if field_count not in (5, 7):
|
||||
raise RuntimeError("CSV format not recognized, can only show scenes in 2 or 3 dimensions")
|
||||
|
||||
# Allocate matrices
|
||||
dimensions = (field_count - 1) // 2
|
||||
pos = np.zeros((len(lines), dimensions), dtype=np.float64)
|
||||
vel = np.zeros((len(lines), dimensions), dtype=np.float64)
|
||||
rad = np.zeros((len(lines), 1), dtype=np.float64)
|
||||
|
||||
# Parse CSV lines
|
||||
for i, line in enumerate(lines):
|
||||
values = tuple(float(field) for field in line.split(","))
|
||||
if dimensions == 2:
|
||||
[x, y, vx, vy, r] = values
|
||||
pos[i] = [x, y]
|
||||
vel[i] = [vx, vy]
|
||||
x, y, vx, vy, r = values
|
||||
pos[i] = (x, y)
|
||||
vel[i] = (vx, vy)
|
||||
rad[i] = r
|
||||
elif dimensions == 3:
|
||||
[x, y, z, vx, vy, vz, r] = values
|
||||
pos[i] = [x, y, z]
|
||||
vel[i] = [vx, vy, vz]
|
||||
x, y, z, vx, vy, vz, r = values
|
||||
pos[i] = (x, y, z)
|
||||
vel[i] = (vx, vy, vz)
|
||||
rad[i] = r
|
||||
|
||||
return pos, vel, rad
|
||||
|
||||
|
||||
class Animator:
|
||||
def __init__(self, pos: np.ndarray, vel: np.ndarray, rad: np.ndarray):
|
||||
self.pos = pos
|
||||
self.vel = vel
|
||||
self.rad = rad
|
||||
self.mass = np.pi * 4 / 3 * rad ** 3
|
||||
"""Runs the simulation and displays it in a plot"""
|
||||
def __init__(self, pos: NDArray, vel: NDArray, rad: NDArray, gravity: float) -> None:
|
||||
"""
|
||||
Sets up the simulation using the given data
|
||||
:param pos: Start positions of the objects
|
||||
:param vel: Start velocities of the objects
|
||||
:param rad: Radii of the objects
|
||||
:param gravity: Strength of gravity in this simulation
|
||||
"""
|
||||
# We update our arrays in-place, so we'll make out own copies to be safe
|
||||
self._pos = pos.copy()
|
||||
self._vel = vel.copy()
|
||||
self._rad = rad.copy()
|
||||
# Calculate volume of a sphere of this radius
|
||||
self._mass = np.pi * 4 / 3 * rad ** 3
|
||||
|
||||
n, d = self.pos.shape
|
||||
self._gravity = gravity
|
||||
|
||||
self.scat: plt.PathCollection = None
|
||||
self.colours = cm.rainbow(
|
||||
object_count, dimensions = self._pos.shape
|
||||
|
||||
# Objects will be represented using a scatter plot
|
||||
self._scatter_plot: plt.PathCollection | None = None
|
||||
# We'll give each objects a random colour to better differentiate them
|
||||
self._object_colours: NDArray = cm.rainbow(
|
||||
np.random.random(
|
||||
(n,)
|
||||
(object_count,)
|
||||
)
|
||||
)
|
||||
|
||||
self.fig = plt.figure()
|
||||
if d == 2:
|
||||
self.ax = self.fig.add_subplot()
|
||||
# Create plot in an appropriate number of dimensions
|
||||
self._fig = plt.figure()
|
||||
if dimensions == 2:
|
||||
self._ax = self._fig.add_subplot()
|
||||
else:
|
||||
self.ax = self.fig.add_subplot(projection="3d")
|
||||
self.ani = animation.FuncAnimation(
|
||||
self.fig,
|
||||
self._ax = self._fig.add_subplot(projection="3d")
|
||||
|
||||
# Set up animation loop
|
||||
# This attribute never gets used again, but we'll keep a reference so that it doesn't get garbage collected
|
||||
self._animation = animation.FuncAnimation(
|
||||
self._fig,
|
||||
self.update,
|
||||
interval=1000 / (15 * 2 ** 4),
|
||||
init_func=self.setup_plot,
|
||||
blit=True,
|
||||
cache_frame_data=False
|
||||
cache_frame_data=False,
|
||||
)
|
||||
|
||||
def setup_plot(self):
|
||||
_n, d = self.pos.shape
|
||||
if d == 2:
|
||||
self.scat = self.ax.scatter(
|
||||
self.pos[:, 0],
|
||||
self.pos[:, 1],
|
||||
c=self.colours,
|
||||
s=self.rad * 10
|
||||
)
|
||||
self.ax.axis([-950, 950, -500, 500])
|
||||
else:
|
||||
self.scat = self.ax.scatter(
|
||||
self.pos[:, 0],
|
||||
self.pos[:, 1],
|
||||
self.pos[:, 2],
|
||||
c=self.colours,
|
||||
s=self.rad * 10,
|
||||
depthshade=False
|
||||
)
|
||||
self.ax.axis([-500, 500, -500, 500, -500, 500])
|
||||
return self.scat,
|
||||
|
||||
def update(self, *_args, **_kwargs):
|
||||
_n, d = self.pos.shape
|
||||
physics.n_body_matrix(self.pos, self.vel, self.mass, constrain=2.)
|
||||
self.scat.set_offsets(self.pos[:, :2])
|
||||
if d == 3:
|
||||
self.scat.set_3d_properties(self.pos[:, 2], 'z')
|
||||
self.scat.set_sizes(self.rad[:, 0] * 10)
|
||||
self.fig.canvas.draw()
|
||||
return self.scat,
|
||||
|
||||
def show(self):
|
||||
# Display the animation
|
||||
plt.show()
|
||||
|
||||
def setup_plot(self) -> tuple[PathCollection]:
|
||||
"""
|
||||
This is a FuncAnimation initialization function in the form matplotlib expects
|
||||
:return: The single scatter plot we're using
|
||||
"""
|
||||
_n, dimensions = self._pos.shape
|
||||
# Set up the scatter plot in 2 or 3 dimensions
|
||||
if dimensions == 2:
|
||||
self._scatter_plot = self._ax.scatter(
|
||||
self._pos[:, 0],
|
||||
self._pos[:, 1],
|
||||
c=self._object_colours,
|
||||
s=self._rad * 10, # To make the objects more visible
|
||||
)
|
||||
# These values work nicely for a landscape window
|
||||
self._ax.axis([-950, 950, -500, 500])
|
||||
else:
|
||||
self._scatter_plot = self._ax.scatter(
|
||||
self._pos[:, 0],
|
||||
self._pos[:, 1],
|
||||
self._pos[:, 2],
|
||||
c=self._object_colours,
|
||||
s=self._rad * 10, # To make the objects more visible
|
||||
depthshade=False, # I find it confusing, YMMV
|
||||
)
|
||||
# These values work nicely for a square window
|
||||
self._ax.axis([-500, 500, -500, 500, -500, 500])
|
||||
return self._scatter_plot,
|
||||
|
||||
def update(self, *_args, **_kwargs) -> tuple[PathCollection]:
|
||||
"""
|
||||
This is a FuncAnimation update function in the form matplotlib expects
|
||||
:arg _args: We don't need any of matplotlib's information for our implementation
|
||||
"arg _kwargs: Again, not necessary for us
|
||||
:return: The single scatter plot we're using
|
||||
"""
|
||||
_n, dimensions = self._pos.shape
|
||||
# Update the state of our simulation
|
||||
n_body_matrix(self._pos, self._vel, self._mass, self._gravity)
|
||||
|
||||
# Set the X and Y values of the objects
|
||||
self._scatter_plot.set_offsets(self._pos[:, :2])
|
||||
if dimensions == 3:
|
||||
# Update the Z value if in 3D
|
||||
self._scatter_plot.set_3d_properties(self._pos[:, 2], 'z')
|
||||
# Use radius to represent mass
|
||||
self._scatter_plot.set_sizes(self._rad[:, 0] * 10)
|
||||
# Redraw the plot
|
||||
self._fig.canvas.draw()
|
||||
return self._scatter_plot,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue