Merge branch 'main' into tfaour-patch-1
This commit is contained in:
commit
ba8348b8b3
Binary file not shown.
Binary file not shown.
Binary file not shown.
@ -1,57 +0,0 @@
|
||||
from enum import Enum
|
||||
import numpy as np
|
||||
|
||||
from ..units import *
|
||||
from .calc import format_sig_figs
|
||||
|
||||
class Body:
|
||||
"""
|
||||
Base class for any orbital body
|
||||
"""
|
||||
def __init__(self, X: Position, V: Velocity, m: Mass, name: str = ""):
|
||||
"""
|
||||
x (position) and v (velocity)
|
||||
"""
|
||||
self.X = X
|
||||
self.V = V
|
||||
self.A = Acceleration([0,0,0])
|
||||
self.m = m
|
||||
self.name = name
|
||||
|
||||
def save(self):
|
||||
return (self.X, self.V, self.m)
|
||||
|
||||
@classmethod
|
||||
def load(cls, tup):
|
||||
return cls(tup[0], tup[1], tup[2])
|
||||
|
||||
def step(self, step_size: float):
|
||||
self.X += step_size*self.V
|
||||
self.V += step_size*self.A
|
||||
self.A = HighPrecisionVector([0,0,0])
|
||||
|
||||
def __str__(self):
|
||||
pos = 'm '.join([format_sig_figs(real_pos(x), 3) for x in self.X])
|
||||
vel = 'm/s '.join([format_sig_figs(real_vel(v), 3) for v in self.V])
|
||||
return str(f"{self.name}: X = {pos}, V = {vel}")
|
||||
|
||||
def E(self):
|
||||
return self.ke() + self.pe()
|
||||
|
||||
def pe(self):
|
||||
return -self.m/self.dist_from_o()
|
||||
|
||||
def dist_from_o(self):
|
||||
return sum([x**2 for x in self.X]).sqrt()
|
||||
|
||||
|
||||
def ke(self):
|
||||
return Decimal(0.5)*self.m*(self._speed()**2)
|
||||
|
||||
def _speed(self):
|
||||
return sum([v**2 for v in self.V]).sqrt()
|
||||
|
||||
def speed(self):
|
||||
return str(f"{format_sig_figs(real_vel(self._speed),5)}")
|
||||
|
||||
|
||||
@ -1,84 +0,0 @@
|
||||
import numpy as np
|
||||
import sys
|
||||
from decimal import Decimal
|
||||
from time import time
|
||||
|
||||
from ..units import HighPrecisionMatrix
|
||||
|
||||
def calculate_distances(positions):
|
||||
N = len(positions)
|
||||
dists = HighPrecisionMatrix(N,N)
|
||||
for i in range(N):
|
||||
dists[i][i] = Decimal(0)
|
||||
for j in range(i+1, N):
|
||||
d = sum([x**2 for x in (positions[i]-positions[j])]).sqrt()
|
||||
dists[i][j] = Decimal(d)
|
||||
dists[j][i] = Decimal(d)
|
||||
return dists
|
||||
|
||||
def calculate_distances_np(positions):
|
||||
positions = np.array(positions)
|
||||
diffs = positions[:, np.newaxis] - positions[np.newaxis, :]
|
||||
dists = np.linalg.norm(diffs, axis=-1)
|
||||
np.fill_diagonal(dists, 0)
|
||||
return dists
|
||||
|
||||
def print_progress_bar(iteration, total, start_time, length=50):
|
||||
"""Prints a progress bar to the console."""
|
||||
percent = (iteration / total) * 100
|
||||
filled_length = int(length * iteration // total)
|
||||
bar = '#' * filled_length + '-' * (length - filled_length)
|
||||
sys.stdout.write(f'\r[{bar}] {percent:.2f}% {int(iteration/(time()-start_time))} steps/s')
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def format_sig_figs(value, sig_figs):
|
||||
"""Format a number to a specified number of significant figures for printing."""
|
||||
if value == 0:
|
||||
return "0"
|
||||
return f"{value:.{sig_figs-1}e}"
|
||||
|
||||
|
||||
|
||||
def plot_points_terminal(vectors, stdscr, scale=500000, grid_width=30, grid_height=30):
|
||||
"""Plots multiple points in the terminal, scaled and centered at (0,0)."""
|
||||
stdscr.clear()
|
||||
|
||||
if not vectors:
|
||||
stdscr.addstr(0, 0, "No vectors provided.")
|
||||
stdscr.refresh()
|
||||
return
|
||||
|
||||
# Apply scaling
|
||||
scaled_vectors = {(round(vec[0] / scale), round(vec[1] / scale)) for vec in vectors}
|
||||
|
||||
# Find min and max coordinates
|
||||
min_x = min(vec[0] for vec in scaled_vectors)
|
||||
max_x = max(vec[0] for vec in scaled_vectors)
|
||||
min_y = min(vec[1] for vec in scaled_vectors)
|
||||
max_y = max(vec[1] for vec in scaled_vectors)
|
||||
|
||||
# Center offsets to keep (0,0) in middle
|
||||
center_x = (grid_width // 2) - min_x
|
||||
center_y = (grid_height // 2) - min_y
|
||||
|
||||
# Adjust coordinates for plotting
|
||||
adjusted_vectors = {(vec[0] + center_x, vec[1] + center_y) for vec in scaled_vectors}
|
||||
|
||||
# Ensure grid boundaries
|
||||
max_terminal_y, max_terminal_x = stdscr.getmaxyx()
|
||||
max_x = min(grid_width, max_terminal_x - 5)
|
||||
max_y = min(grid_height, max_terminal_y - 5)
|
||||
|
||||
# Draw grid with points
|
||||
for i in range(grid_height, -1, -1):
|
||||
row = f"{i - center_y:2} | "
|
||||
for j in range(grid_width + 1):
|
||||
row += "● " if (j, i) in adjusted_vectors else ". "
|
||||
stdscr.addstr(max_y - i, 0, row[:max_terminal_x - 1]) # Ensure no overflow
|
||||
|
||||
# Print X-axis labels
|
||||
x_labels = " " + " ".join(f"{j - center_x:2}" for j in range(max_x + 1))
|
||||
stdscr.addstr(max_y + 1, 0, x_labels[:max_terminal_x - 1]) # Avoid out-of-bounds error
|
||||
|
||||
stdscr.refresh()
|
||||
@ -1,136 +0,0 @@
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
from decimal import InvalidOperation, DivisionByZero
|
||||
from time import time
|
||||
|
||||
from .body import Body
|
||||
from .calc import *
|
||||
from ..units import *
|
||||
|
||||
|
||||
|
||||
class Simulator:
|
||||
"""
|
||||
Everything a simulator needs to run:
|
||||
- a step size
|
||||
- bodies with initial positions and momenta
|
||||
- number of steps to take
|
||||
- a reset mechanism
|
||||
- a checkpoint mechanism
|
||||
- how often to save?
|
||||
- overwrite output file if exists? default false
|
||||
if output file is a saved checkpoint file, then use the
|
||||
from_checkpoint class method to create the class.
|
||||
|
||||
otherwise if the output file exists, class will not start.
|
||||
|
||||
output is as text:
|
||||
first line of file is list of masses of bodies
|
||||
each subsequent line is list of positions and velocities of each bodies:
|
||||
body1x, body1y, body1vx, body1vy, body2x, body2y, body2vx etc
|
||||
|
||||
For some of these, it should come with sensible defaults with
|
||||
the ability to change
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
bodies: list[Body],
|
||||
step_size: float,
|
||||
steps_per_save: int,
|
||||
output_file: str,
|
||||
current_step: int = 0,
|
||||
overwrite_output: bool = False
|
||||
):
|
||||
self.output_file = Path(output_file)
|
||||
if output_file.exists() and not overwrite_output:
|
||||
raise FileExistsError(f"File {output_file} exists and overwrite flag not given.")
|
||||
|
||||
self.bodies = bodies
|
||||
self.step_size = step_size
|
||||
self.steps_per_save = steps_per_save
|
||||
self.current_step = current_step
|
||||
|
||||
if output_file.exists() and overwrite_output:
|
||||
print(f"Warning! Overwriting file: {output_file}")
|
||||
|
||||
self._checkpoint()
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_checkpoint(cls, output_file: Path):
|
||||
data = np.load("last_checkpoint.npz")
|
||||
positions = data["positions"]
|
||||
velocities = data["velocities"]
|
||||
masses = data["masses"]
|
||||
step_size = data["steps"][0]
|
||||
current_step = data["steps"][1]
|
||||
steps_per_save = data["steps"][2]
|
||||
|
||||
bodies = [
|
||||
Body(val[0], val[1], val[2]) for val in zip(
|
||||
positions, velocities, masses
|
||||
)
|
||||
]
|
||||
return cls(
|
||||
bodies,
|
||||
step_size,
|
||||
steps_per_save,
|
||||
output_file,
|
||||
current_step,
|
||||
)
|
||||
|
||||
|
||||
def _checkpoint(self):
|
||||
"""
|
||||
Two things - save high precision last checkpoint for resuming
|
||||
then save lower precision text for trajectories
|
||||
"""
|
||||
body_X_np = np.array([
|
||||
body.X for body in self.bodies
|
||||
])
|
||||
body_V_np = np.array([
|
||||
body.V for body in self.bodies
|
||||
])
|
||||
body_m_np = np.array([
|
||||
body.m for body in self.bodies
|
||||
])
|
||||
stepsz_n_np = np.array([
|
||||
self.step_size,
|
||||
self.current_step,
|
||||
self.steps_per_save
|
||||
])
|
||||
|
||||
np.savez("last_checkpoint.npz",
|
||||
positions=body_X_np,
|
||||
velocities=body_V_np,
|
||||
masses=body_m_np,
|
||||
steps=stepsz_n_np)
|
||||
|
||||
def run(self, steps):
|
||||
time_start = time()
|
||||
for i in range(steps):
|
||||
self.calculate_forces()
|
||||
self.move_bodies()
|
||||
self.current_step += 1
|
||||
if (self.current_step % self.steps_per_save == 0):
|
||||
print_progress_bar(i, steps, time_start)
|
||||
#self._checkpoint()
|
||||
|
||||
def calculate_forces(self):
|
||||
positions = [
|
||||
body.X for body in self.bodies
|
||||
]
|
||||
dists = calculate_distances_np(positions)
|
||||
for i in range(len(self.bodies)):
|
||||
for j in range(i, len(self.bodies)):
|
||||
if i == j:
|
||||
continue
|
||||
vec = self.bodies[i].X - self.bodies[j].X
|
||||
f = vec/(dists[i][j]**3)
|
||||
self.bodies[i].A += f*self.bodies[j].m
|
||||
self.bodies[j].A += f*self.bodies[i].m
|
||||
|
||||
|
||||
def move_bodies(self):
|
||||
for body in self.bodies:
|
||||
body.step(self.step_size)
|
||||
@ -1,74 +0,0 @@
|
||||
import numpy as np
|
||||
from decimal import Decimal
|
||||
|
||||
#Declare these as units to make code clearer
|
||||
class HighPrecisionVector(np.ndarray):
|
||||
def __new__(cls, input_array, *args, **kwargs):
|
||||
decimal_array = [Decimal(i) for i in input_array]
|
||||
obj = np.asarray(decimal_array).view(cls)
|
||||
return obj
|
||||
|
||||
class HighPrecisionMatrix(np.ndarray):
|
||||
def __new__(cls, dim1, dim2, *args, **kwargs):
|
||||
decimal_array = [Decimal(0) for _ in range(dim1)]
|
||||
decimal_matrix = [decimal_array for _ in range(dim2)]
|
||||
obj = np.asarray(decimal_matrix).view(cls)
|
||||
return obj
|
||||
|
||||
class Position(HighPrecisionVector):
|
||||
pass
|
||||
class Velocity(HighPrecisionVector):
|
||||
pass
|
||||
class Acceleration(HighPrecisionVector):
|
||||
pass
|
||||
|
||||
Mass = Decimal
|
||||
|
||||
####################
|
||||
# USEFUL CONSTANTS
|
||||
####################
|
||||
EARTH_MASS = Decimal(5972 * 10**21) #kg
|
||||
EARTH_RADIUS = Decimal(6378 * 10**3) #meters
|
||||
EARTH_ORBITAL_VELOCITY = Decimal(29780) # m/s
|
||||
AU = Decimal(149_597_870_700) #meters
|
||||
|
||||
MOON_MASS = Decimal(734767309 * 10**14)
|
||||
MOON_ORBITAL_VELOCITY = Decimal(1022) #m/s relative to earth
|
||||
|
||||
SUN_MASS = Decimal(1989 * 10**27) #kg
|
||||
SUN_RADIUS = Decimal(6957 * 10**5) #meters
|
||||
|
||||
pi_approx = Decimal("3.14159265358979323846264338327950288419716939937510")
|
||||
|
||||
|
||||
#NORMALIZING CONSTANTS
|
||||
G = Decimal(6.67430e-11)
|
||||
r_0 = Decimal(EARTH_RADIUS) #1.496e11
|
||||
m_0 = Decimal(5.972e24)
|
||||
t_0 = np.sqrt((r_0**3) / (G*m_0))
|
||||
|
||||
def norm_pos(pos):
|
||||
return Decimal(pos) / Decimal(r_0)
|
||||
|
||||
def real_pos(pos):
|
||||
return pos * Decimal(r_0)
|
||||
|
||||
def norm_mass(mass):
|
||||
return Decimal(mass) / Decimal(m_0)
|
||||
|
||||
def real_mass(mass):
|
||||
return Decimal(mass) * Decimal(m_0)
|
||||
|
||||
def norm_time(time):
|
||||
return Decimal(time) / Decimal(t_0)
|
||||
|
||||
def real_time(time):
|
||||
return Decimal(time) * Decimal(t_0)
|
||||
|
||||
def norm_vel(vel):
|
||||
return vel / Decimal(r_0/t_0)
|
||||
|
||||
def real_vel(vel):
|
||||
return vel * Decimal(r_0/t_0)
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user