simulation module
Definition of the Simulation class and the Galaxy constructor.
Source code
"""Definition of the Simulation class and the Galaxy constructor."""
import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
from utils import random_unit_vectors, cascade_round
from distributions import PLUMMER, HERNQUIST, UNIFORM, EXP, NFW
import acceleration
##############################################################################
##############################################################################
class Simulation:
""""Main class for the gravitational simulation.
Attributes:
r_vec (array): position of the particles in the current timestep.
Shape: (number of particles, 3)
rprev_vec (array): position of the particles in the previous timestep.
Shape: (number of particles, 3)
v_vec (array): velocity in the current timestep.
Shape: (number of particles, 3)
a_vec (array): acceleration in the current timestep.
Shape: (number of particles, 3)
mass (array): mass of each particle in the simulation.
Shape: (number of particles,)
type (array): non-unique identifier for each particle.
Shape: (number of particles,)
tracks (array): list of positions through the simulation for central
masses. Shape: (tracked particles, n+1, 3).
CONFIG (array): configuration used to create the simulation.
It will be saved along the state of the simulation.
dt (float): timestep of the simulation
n (int): current timestep. Initialized as n=0.
soft (float): softening length used by the simulation.
verbose (boolean): When True progress statements will be printed.
"""
def __init__(self, dt, soft, verbose, CONFIG, method, **kwargs):
"""Constructor for the Simulation class.
Arguments:
dt (float): timestep of the simulation
n (int): current timestep. Initialized as n=0.
soft (float): softening length used by the simulation.
verbose (bool): When True progress statements will be printed.
CONFIG (dict): configuration file used to create the simulation.
method (string): Optional. Algorithm to use when computing the
gravitational forces. One of 'bruteForce', 'bruteForce_numba',
'bruteForce_numbaopt', 'bruteForce_CPP', 'barnesHut_CPP'.
"""
self.n = 0
self.t = 0
self.dt = dt
self.soft = soft
self.verbose = verbose
self.CONFIG = CONFIG
# Initialize empty arrays for all necessary properties
self.r_vec = np.empty((0, 3))
self.v_vec = np.empty((0, 3))
self.a_vec = np.empty((0, 3))
self.mass = np.empty((0,))
self.type = np.empty((0, 2))
algorithms = {
'bruteForce': acceleration.bruteForce,
'bruteForceNumba': acceleration.bruteForceNumba,
'bruteForceNumbaOptimized': acceleration.bruteForceNumbaOptimized,
'bruteForceCPP': acceleration.bruteForceCPP,
'barnesHutCPP': acceleration.barnesHutCPP
}
try:
self.acceleration = algorithms[method]
except: raise Exception("Method '{}' unknown".format(method))
def add(self, body):
"""Add a body to the simulation. It must expose the public attributes
body.r_vec, body.v_vec, body.a_vec, body.type, body.mass.
Arguments:
body: Object to be added to the simulation (e.g. a Galaxy object)
"""
# Extend all relevant attributes by concatenating the body
for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']:
simattr, bodyattr = getattr(self, name), getattr(body, name)
setattr(self, name, np.concatenate([simattr, bodyattr], axis=0))
# Order based on mass
order = np.argsort(-self.mass)
for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']:
setattr(self, name, getattr(self, name)[order])
# Update the list of objects to keep track of
self.tracks = np.empty((np.sum(self.type[:,0]=='center'), 0, 3))
def step(self):
"""Perform a single step of the simulation.
Makes use of a 4th order Verlet integrator.
"""
# Calculate the acceleration
self.a_vec = self.acceleration(self.r_vec, self.mass, soft=self.soft)
# Update the state using the Verlet algorithm
# (A custom algorithm is written mainly for learning purposes)
self.r_vec, self.rprev_vec = (2*self.r_vec - self.rprev_vec
+ self.a_vec * self.dt**2, self.r_vec)
self.n += 1
# Update tracks
self.tracks = np.concatenate([self.tracks,
self.r_vec[self.type[:,0]=='center'][:,np.newaxis]], axis=1)
def run(self, tmax, saveEvery, outputFolder, **kwargs):
"""Run the galactic simulation.
Attributes:
tmax (float): Time to which the simulation will run to.
This is measured here since the start of the simulation,
not since pericenter.
saveEvery (int): The state is saved every saveEvery steps.
outputFolder (string): It will be saved to /data/outputFolder/
"""
# When the simulation starts, intialize self.rprev_vec
self.rprev_vec = self.r_vec - self.v_vec * self.dt
if self.verbose: print('Simulation starting. Bon voyage!')
while(self.t < tmax):
self.step()
if(self.n % saveEvery == 0):
self.save('data/{}'.format(outputFolder))
print('Simulation complete.')
def save(self, outputFolder):
"""Save the state of the simulation to the outputFolder.
Two files are saved:
sim{self.n}.pickle: serializing the state.
sim{self.n}.png: a simplified 2D plot of x, y.
"""
# Create the output folder if it doesn't exist
if not os.path.exists(outputFolder): os.makedirs(outputFolder)
# Compute some useful quantities
# v_vec is not required by the integrator, but useful
self.v_vec = (self.r_vec - self.rprev_vec)/self.dt
self.t = self.n * self.dt # prevents numerical rounding errors
# Serialize state
file = open(outputFolder+'/data{}.pickle'.format(self.n), "wb")
pickle.dump({'r_vec': self.r_vec, 'v_vec': self.v_vec,
'type': self.type, 'mass': self.mass,
'CONFIG': self.CONFIG, 't': self.t,
'tracks': self.tracks}, file)
# Save simplified plot of the current state.
# Its main use is to detect ill-behaved situations early on.
fig = plt.figure()
plt.xlim(-5, 5); plt.ylim(-5, 5); plt.axis('equal')
# Dark halo is plotted in red, disk in blue, bulge in green
PLTCON = [('dark', 'r', 0.3), ('disk', 'b', 1.0), ('bulge', 'g', 0.5)]
for type_, c, a in PLTCON:
plt.scatter(self.r_vec[self.type[:,0]==type_][:,0],
self.r_vec[self.type[:,0]==type_][:,1], s=0.1, c=c, alpha=a)
# Central mass as a magenta star
plt.scatter(self.r_vec[self.type[:,0]=='center'][:,0],
self.r_vec[self.type[:,0]=='center'][:,1], s=100, marker="*", c='m')
# Save to png file
fig.savefig(outputFolder+'/sim{}.png'.format(self.n), dpi=150)
plt.close(fig)
def project(self, theta, phi, view=0):
"""Projects the 3D simulation onto a plane as viewed from the
direction described by the (theta, phi, view). Angles in radians.
(This is used by the simulated annealing algorithm)
Parameters:
theta (float): polar angle.
phi (float): azimuthal angle.
view (float): rotation along line of sight.
"""
M1 = np.array([[np.cos(phi), np.sin(phi), 0],
[-np.sin(phi), np.cos(phi), 0],
[0, 0, 1]])
M2 = np.array([[1, 0, 0],
[0, np.cos(theta), np.sin(theta)],
[0, -np.sin(theta), np.cos(theta)]])
M3 = np.array([[np.cos(view), np.sin(view), 0],
[-np.sin(view), np.cos(view), 0],
[0, 0, 1]])
M = np.matmul(M1, np.matmul(M2, M3)) # combine rotations
r = np.tensordot(self.r_vec, M, axes=[1, 0])
return r[:,0:2]
def setOrbit(self, galaxy1, galaxy2, e, rmin, R0):
"""Sets the two galaxies galaxy1, galaxy2 in an orbit.
Parameters:
galaxy1 (Galaxy): 1st galaxy (main)
galaxy2 (Galaxy): 2nd galaxy (companion)
e: eccentricity of the orbit
rmin: distance of closest approach
R0: initial separation
"""
m1, m2 = np.sum(galaxy1.mass), np.sum(galaxy2.mass)
# Relevant formulae:
# $r_0 = r (1 + e) \cos(\phi)$, where $r_0$ ($\neq R_0$) is the semi-latus rectum
# $r_0 = r_\textup{min} (1 + e)$
# $v^2 = G M (2/r - 1/a)$, where a is the semimajor axis
# Solve the reduced two-body problem with reduced mass $\mu$ (mu)
M = m1 + m2
r0 = rmin * (1 + e)
try:
phi = np.arccos((r0/R0 - 1) / e) # inverting the orbit equation
phi = -np.abs(phi) # Choose negative (incoming) angle
ainv = (1 - e) / rmin # ainv = $1/a$, as a may be infinite
v = np.sqrt(M * (2/R0 - ainv))
# Finally, calculate the angle of motion. angle = tan(dy/dx)
# $dy/dx = ((dr/d\phi) sin(\phi) + r \cos(\phi))/((dr/d\phi) cos(\phi) - r \sin(\phi))$
dy = R0/r0 * e * np.sin(phi)**2 + np.cos(phi)
dx = R0/r0 * e * np.sin(phi) * np.cos(phi) - np.sin(phi)
vangle = np.arctan2(dy, dx)
except: raise Exception('''The orbital parameters cannot be satisfied.
For elliptical orbits check that R0 is posible (<rmax).''')
# We now need the actual motion of each body
R_vec = np.array([[R0*np.cos(phi), R0*np.sin(phi), 0.]])
V_vec = np.array([[v*np.cos(vangle), v*np.sin(vangle), 0.]])
galaxy1.r_vec += m2/M * R_vec
galaxy1.v_vec += m2/M * V_vec
galaxy2.r_vec += -m1/M * R_vec
galaxy2.v_vec += -m1/M * V_vec
# Explicitely add the galaxies to the simulation
self.add(galaxy1)
self.add(galaxy2)
if self.verbose: print('Galaxies set in orbit.')
##############################################################################
##############################################################################
class Galaxy():
""""Helper class for creating galaxies.
Attributes:
r_vec (array): position of the particles in the current timestep.
Shape: (number of particles, 3)
v_vec (array): velocity in the current timestep.
Shape: (number of particles, 3)
a_vec (array): acceleration in the current timestep.
Shape: (number of particles, 3)
mass (array): mass of each particle in the simulation.
Shape: (number of particles,)
type (array): non-unique identifier for each particle.
Shape: (number of particles,) """
def __init__(self, orientation, centralMass, bulge, disk, halo, sim):
"""Constructor for the Galaxy class.
Parameters:
orientation (tupple): (inclination i, argument of pericenter w)
centralMass (float): mass at the center of the galaxy
bulge (dict): passed to the addBulge method.
disk (dict): passed to the addDisk method.
halo (dict): passed to the addHalo method.
sim (Simulation): simulation object
"""
if sim.verbose: print('Initializing galaxy')
# Build the central mass
self.r_vec = np.zeros((1, 3))
self.v_vec = np.zeros((1, 3))
self.a_vec = np.zeros((1, 3))
self.mass = np.array([centralMass])
self.type = np.array([['center', 0]])
# Build the other components
self.addBulge(**bulge)
if sim.verbose: print('Bulge created.')
self.addDisk(**disk)
if sim.verbose: print('Disk created.')
self.addHalo(**halo)
if sim.verbose: print('Halo created.')
# Correct particles velocities to generate circular orbits
# $a_\textup{centripetal} = v^2/r$
a_vec = sim.acceleration(self.r_vec, self.mass, soft=sim.soft)
a = np.linalg.norm(a_vec, axis=-1, keepdims=True)
r = np.linalg.norm(self.r_vec, axis=-1, keepdims=True)
v = np.linalg.norm(self.v_vec[1:], axis=-1, keepdims=True)
direction_unit = self.v_vec[1:]/v
# Set orbital velocities (except for central mass)
self.v_vec[1:] = np.sqrt(a[1:] * r[1:]) * direction_unit
self.a_vec = np.zeros_like(self.r_vec)
# Rotate the galaxy into its correct orientation
self.rotate(*(np.array(orientation)/360 * 2*np.pi))
if sim.verbose: print('Galaxy set in rotation and oriented.')
def addBulge(self, model, totalMass, particles, l):
"""Adds a bulge to the galaxy.
Parameters:
model (string): parametrization of the bulge.
'plummer' and 'hernquist' are supported.
totalMass (float): total mass of the bulge
particles (int): number of particles in the bulge
l (float): characteristic length scale of the model.
"""
if particles == 0: return None
# Divide the mass equally among all particles
mass = np.ones(particles) * totalMass/particles
self.mass = np.concatenate([self.mass, mass], axis=0)
# Create particles according to the radial distribution from model
if model == 'plummer':
r = PLUMMER.ppf(np.random.rand(particles), scale=l)
elif model == 'hernquist':
r = HERNQUIST.ppf(np.random.rand(particles), scale=l)
else: raise Exception("""Bulge distribution not allowed.
'plummer' and 'hernquist' are the supported values""")
r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles)
self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
# Set them orbitting along random directions normal to r_vec
v_vec = np.cross(r_vec, random_unit_vectors(size=particles))
self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
# Label the particles
type_ = [['bulge', 0]]*particles
self.type = np.concatenate([self.type, type_], axis=0)
def addDisk(self, model, totalMass, particles, l):
"""Adds a disk to the galaxy.
Parameters:
model (string): parametrization of the disk.
'rings', 'uniform' and 'exp' are supported.
totalMass (float): total mass of the bulge
particles (int): number of particles in the bulge
l: fot 'exp' and 'uniform' characteristic length of the
model. For 'rings' tupple of the form (inner radius,
outer radius, number of rings)
"""
if particles == 0: return None
# Create particles according to the radial distribution from model
if model == 'uniform':
r = UNIFORM.ppf(np.random.rand(particles), scale=l)
type_ = [['disk', 0]]*particles
r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D')
self.type = np.concatenate([self.type, type_], axis=0)
elif model == 'rings':
# l = [inner radius, outter radius, number of rings]
distances = np.linspace(*l)
# Aim for roughly constant areal density
# Cascade rounding preserves the total number of particles
perRing = cascade_round(particles * distances / np.sum(distances))
particles = int(np.sum(perRing)) # prevents numerical errors
r_vec = np.empty((0, 3))
for d, n, i in zip(distances, perRing, range(l[2])):
type_ = [['disk', i+1]]*int(n)
self.type = np.concatenate([self.type, type_], axis=0)
phi = np.linspace(0, 2 * np.pi, n, endpoint=False)
ringr = d * np.array([[np.cos(i), np.sin(i), 0] for i in phi])
r_vec = np.concatenate([r_vec, ringr], axis=0)
elif model == 'exp':
r = EXP.ppf(np.random.rand(particles), scale=l)
r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D')
type_ = [['disk', 0]]*particles
self.type = np.concatenate([self.type, type_], axis=0)
else:
raise Exception("""Disk distribution not allowed.
'uniform', 'rings' and 'exp' are the supported values""")
self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
# Divide the mass equally among all particles
mass = np.ones(particles) * totalMass/particles
self.mass = np.concatenate([self.mass, mass], axis=0)
# Set them orbitting along the spin plane
v_vec = np.cross(r_vec, [0, 0, 1])
self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
def addHalo(self, model, totalMass, particles, rs):
"""Adds a halo to the galaxy.
Parameters:
model (string): parametrization of the halo.
Only 'NFW' is supported.
totalMass (float): total mass of the halo
particles (int): number of particles in the halo
rs (float): characteristic length scale of the NFW profile.
"""
if particles == 0: return None
# Divide the mass equally among all particles
mass = np.ones(particles)*totalMass/particles
self.mass = np.concatenate([self.mass, mass], axis=0)
# Create particles according to the radial distribution from model
if model == 'NFW':
r = NFW.ppf(np.random.rand(particles), scale=rs)
else: raise Exception("""Bulge distribution not allowed.
'plummer' and 'hernquist' are the supported values""")
r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles)
self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
# Orbit along random directions normal to the radial vector
v_vec = np.cross(r_vec, random_unit_vectors(size=particles))
self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
# Label the particles
type_ = [['dark'], 0]*particles
self.type = np.concatenate([self.type, type_], axis=0)
def rotate(self, theta, phi):
"""Rotates the galaxy so that its spin is along the (theta, phi)
direction.
Parameters:
theta (float): polar angle.
phi (float): azimuthal angle.
"""
M1 = np.array([[1, 0, 0],
[0, np.cos(theta), np.sin(theta)],
[0, -np.sin(theta), np.cos(theta)]])
M2 = np.array([[np.cos(phi), np.sin(phi), 0],
[-np.sin(phi), np.cos(phi), 0],
[0, 0, 1]])
M = np.matmul(M1, M2) # combine rotations
self.r_vec = np.tensordot(self.r_vec, M, axes=[1, 0])
self.v_vec = np.tensordot(self.v_vec, M, axes=[1, 0])Classes
class Galaxy-
"Helper class for creating galaxies.
Attributes
r_vec:array- position of the particles in the current timestep. Shape: (number of particles, 3)
v_vec:array- velocity in the current timestep. Shape: (number of particles, 3)
a_vec:array- acceleration in the current timestep. Shape: (number of particles, 3)
mass:array- mass of each particle in the simulation. Shape: (number of particles,)
type:array- non-unique identifier for each particle. Shape: (number of particles,)
Source code
class Galaxy(): """"Helper class for creating galaxies. Attributes: r_vec (array): position of the particles in the current timestep. Shape: (number of particles, 3) v_vec (array): velocity in the current timestep. Shape: (number of particles, 3) a_vec (array): acceleration in the current timestep. Shape: (number of particles, 3) mass (array): mass of each particle in the simulation. Shape: (number of particles,) type (array): non-unique identifier for each particle. Shape: (number of particles,) """ def __init__(self, orientation, centralMass, bulge, disk, halo, sim): """Constructor for the Galaxy class. Parameters: orientation (tupple): (inclination i, argument of pericenter w) centralMass (float): mass at the center of the galaxy bulge (dict): passed to the addBulge method. disk (dict): passed to the addDisk method. halo (dict): passed to the addHalo method. sim (Simulation): simulation object """ if sim.verbose: print('Initializing galaxy') # Build the central mass self.r_vec = np.zeros((1, 3)) self.v_vec = np.zeros((1, 3)) self.a_vec = np.zeros((1, 3)) self.mass = np.array([centralMass]) self.type = np.array([['center', 0]]) # Build the other components self.addBulge(**bulge) if sim.verbose: print('Bulge created.') self.addDisk(**disk) if sim.verbose: print('Disk created.') self.addHalo(**halo) if sim.verbose: print('Halo created.') # Correct particles velocities to generate circular orbits # $a_\textup{centripetal} = v^2/r$ a_vec = sim.acceleration(self.r_vec, self.mass, soft=sim.soft) a = np.linalg.norm(a_vec, axis=-1, keepdims=True) r = np.linalg.norm(self.r_vec, axis=-1, keepdims=True) v = np.linalg.norm(self.v_vec[1:], axis=-1, keepdims=True) direction_unit = self.v_vec[1:]/v # Set orbital velocities (except for central mass) self.v_vec[1:] = np.sqrt(a[1:] * r[1:]) * direction_unit self.a_vec = np.zeros_like(self.r_vec) # Rotate the galaxy into its correct orientation self.rotate(*(np.array(orientation)/360 * 2*np.pi)) if sim.verbose: print('Galaxy set in rotation and oriented.') def addBulge(self, model, totalMass, particles, l): """Adds a bulge to the galaxy. Parameters: model (string): parametrization of the bulge. 'plummer' and 'hernquist' are supported. totalMass (float): total mass of the bulge particles (int): number of particles in the bulge l (float): characteristic length scale of the model. """ if particles == 0: return None # Divide the mass equally among all particles mass = np.ones(particles) * totalMass/particles self.mass = np.concatenate([self.mass, mass], axis=0) # Create particles according to the radial distribution from model if model == 'plummer': r = PLUMMER.ppf(np.random.rand(particles), scale=l) elif model == 'hernquist': r = HERNQUIST.ppf(np.random.rand(particles), scale=l) else: raise Exception("""Bulge distribution not allowed. 'plummer' and 'hernquist' are the supported values""") r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles) self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0) # Set them orbitting along random directions normal to r_vec v_vec = np.cross(r_vec, random_unit_vectors(size=particles)) self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0) # Label the particles type_ = [['bulge', 0]]*particles self.type = np.concatenate([self.type, type_], axis=0) def addDisk(self, model, totalMass, particles, l): """Adds a disk to the galaxy. Parameters: model (string): parametrization of the disk. 'rings', 'uniform' and 'exp' are supported. totalMass (float): total mass of the bulge particles (int): number of particles in the bulge l: fot 'exp' and 'uniform' characteristic length of the model. For 'rings' tupple of the form (inner radius, outer radius, number of rings) """ if particles == 0: return None # Create particles according to the radial distribution from model if model == 'uniform': r = UNIFORM.ppf(np.random.rand(particles), scale=l) type_ = [['disk', 0]]*particles r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D') self.type = np.concatenate([self.type, type_], axis=0) elif model == 'rings': # l = [inner radius, outter radius, number of rings] distances = np.linspace(*l) # Aim for roughly constant areal density # Cascade rounding preserves the total number of particles perRing = cascade_round(particles * distances / np.sum(distances)) particles = int(np.sum(perRing)) # prevents numerical errors r_vec = np.empty((0, 3)) for d, n, i in zip(distances, perRing, range(l[2])): type_ = [['disk', i+1]]*int(n) self.type = np.concatenate([self.type, type_], axis=0) phi = np.linspace(0, 2 * np.pi, n, endpoint=False) ringr = d * np.array([[np.cos(i), np.sin(i), 0] for i in phi]) r_vec = np.concatenate([r_vec, ringr], axis=0) elif model == 'exp': r = EXP.ppf(np.random.rand(particles), scale=l) r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D') type_ = [['disk', 0]]*particles self.type = np.concatenate([self.type, type_], axis=0) else: raise Exception("""Disk distribution not allowed. 'uniform', 'rings' and 'exp' are the supported values""") self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0) # Divide the mass equally among all particles mass = np.ones(particles) * totalMass/particles self.mass = np.concatenate([self.mass, mass], axis=0) # Set them orbitting along the spin plane v_vec = np.cross(r_vec, [0, 0, 1]) self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0) def addHalo(self, model, totalMass, particles, rs): """Adds a halo to the galaxy. Parameters: model (string): parametrization of the halo. Only 'NFW' is supported. totalMass (float): total mass of the halo particles (int): number of particles in the halo rs (float): characteristic length scale of the NFW profile. """ if particles == 0: return None # Divide the mass equally among all particles mass = np.ones(particles)*totalMass/particles self.mass = np.concatenate([self.mass, mass], axis=0) # Create particles according to the radial distribution from model if model == 'NFW': r = NFW.ppf(np.random.rand(particles), scale=rs) else: raise Exception("""Bulge distribution not allowed. 'plummer' and 'hernquist' are the supported values""") r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles) self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0) # Orbit along random directions normal to the radial vector v_vec = np.cross(r_vec, random_unit_vectors(size=particles)) self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0) # Label the particles type_ = [['dark'], 0]*particles self.type = np.concatenate([self.type, type_], axis=0) def rotate(self, theta, phi): """Rotates the galaxy so that its spin is along the (theta, phi) direction. Parameters: theta (float): polar angle. phi (float): azimuthal angle. """ M1 = np.array([[1, 0, 0], [0, np.cos(theta), np.sin(theta)], [0, -np.sin(theta), np.cos(theta)]]) M2 = np.array([[np.cos(phi), np.sin(phi), 0], [-np.sin(phi), np.cos(phi), 0], [0, 0, 1]]) M = np.matmul(M1, M2) # combine rotations self.r_vec = np.tensordot(self.r_vec, M, axes=[1, 0]) self.v_vec = np.tensordot(self.v_vec, M, axes=[1, 0])Methods
def __init__(self, orientation, centralMass, bulge, disk, halo, sim)-
Constructor for the Galaxy class.
Parameters
orientation:tupple- (inclination i, argument of pericenter w)
centralMass:float- mass at the center of the galaxy
bulge:dict- passed to the addBulge method.
disk:dict- passed to the addDisk method.
halo:dict- passed to the addHalo method.
sim:Simulation- simulation object
Source code
def __init__(self, orientation, centralMass, bulge, disk, halo, sim): """Constructor for the Galaxy class. Parameters: orientation (tupple): (inclination i, argument of pericenter w) centralMass (float): mass at the center of the galaxy bulge (dict): passed to the addBulge method. disk (dict): passed to the addDisk method. halo (dict): passed to the addHalo method. sim (Simulation): simulation object """ if sim.verbose: print('Initializing galaxy') # Build the central mass self.r_vec = np.zeros((1, 3)) self.v_vec = np.zeros((1, 3)) self.a_vec = np.zeros((1, 3)) self.mass = np.array([centralMass]) self.type = np.array([['center', 0]]) # Build the other components self.addBulge(**bulge) if sim.verbose: print('Bulge created.') self.addDisk(**disk) if sim.verbose: print('Disk created.') self.addHalo(**halo) if sim.verbose: print('Halo created.') # Correct particles velocities to generate circular orbits # $a_\textup{centripetal} = v^2/r$ a_vec = sim.acceleration(self.r_vec, self.mass, soft=sim.soft) a = np.linalg.norm(a_vec, axis=-1, keepdims=True) r = np.linalg.norm(self.r_vec, axis=-1, keepdims=True) v = np.linalg.norm(self.v_vec[1:], axis=-1, keepdims=True) direction_unit = self.v_vec[1:]/v # Set orbital velocities (except for central mass) self.v_vec[1:] = np.sqrt(a[1:] * r[1:]) * direction_unit self.a_vec = np.zeros_like(self.r_vec) # Rotate the galaxy into its correct orientation self.rotate(*(np.array(orientation)/360 * 2*np.pi)) if sim.verbose: print('Galaxy set in rotation and oriented.') def addBulge(self, model, totalMass, particles, l)-
Adds a bulge to the galaxy.
Parameters
model:string- parametrization of the bulge. 'plummer' and 'hernquist' are supported.
totalMass:float- total mass of the bulge
particles:int- number of particles in the bulge
l:float- characteristic length scale of the model.
Source code
def addBulge(self, model, totalMass, particles, l): """Adds a bulge to the galaxy. Parameters: model (string): parametrization of the bulge. 'plummer' and 'hernquist' are supported. totalMass (float): total mass of the bulge particles (int): number of particles in the bulge l (float): characteristic length scale of the model. """ if particles == 0: return None # Divide the mass equally among all particles mass = np.ones(particles) * totalMass/particles self.mass = np.concatenate([self.mass, mass], axis=0) # Create particles according to the radial distribution from model if model == 'plummer': r = PLUMMER.ppf(np.random.rand(particles), scale=l) elif model == 'hernquist': r = HERNQUIST.ppf(np.random.rand(particles), scale=l) else: raise Exception("""Bulge distribution not allowed. 'plummer' and 'hernquist' are the supported values""") r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles) self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0) # Set them orbitting along random directions normal to r_vec v_vec = np.cross(r_vec, random_unit_vectors(size=particles)) self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0) # Label the particles type_ = [['bulge', 0]]*particles self.type = np.concatenate([self.type, type_], axis=0) def addDisk(self, model, totalMass, particles, l)-
Adds a disk to the galaxy.
Parameters
model:string- parametrization of the disk. 'rings', 'uniform' and 'exp' are supported.
totalMass:float- total mass of the bulge
particles:int- number of particles in the bulge
l- fot 'exp' and 'uniform' characteristic length of the model. For 'rings' tupple of the form (inner radius, outer radius, number of rings)
Source code
def addDisk(self, model, totalMass, particles, l): """Adds a disk to the galaxy. Parameters: model (string): parametrization of the disk. 'rings', 'uniform' and 'exp' are supported. totalMass (float): total mass of the bulge particles (int): number of particles in the bulge l: fot 'exp' and 'uniform' characteristic length of the model. For 'rings' tupple of the form (inner radius, outer radius, number of rings) """ if particles == 0: return None # Create particles according to the radial distribution from model if model == 'uniform': r = UNIFORM.ppf(np.random.rand(particles), scale=l) type_ = [['disk', 0]]*particles r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D') self.type = np.concatenate([self.type, type_], axis=0) elif model == 'rings': # l = [inner radius, outter radius, number of rings] distances = np.linspace(*l) # Aim for roughly constant areal density # Cascade rounding preserves the total number of particles perRing = cascade_round(particles * distances / np.sum(distances)) particles = int(np.sum(perRing)) # prevents numerical errors r_vec = np.empty((0, 3)) for d, n, i in zip(distances, perRing, range(l[2])): type_ = [['disk', i+1]]*int(n) self.type = np.concatenate([self.type, type_], axis=0) phi = np.linspace(0, 2 * np.pi, n, endpoint=False) ringr = d * np.array([[np.cos(i), np.sin(i), 0] for i in phi]) r_vec = np.concatenate([r_vec, ringr], axis=0) elif model == 'exp': r = EXP.ppf(np.random.rand(particles), scale=l) r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D') type_ = [['disk', 0]]*particles self.type = np.concatenate([self.type, type_], axis=0) else: raise Exception("""Disk distribution not allowed. 'uniform', 'rings' and 'exp' are the supported values""") self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0) # Divide the mass equally among all particles mass = np.ones(particles) * totalMass/particles self.mass = np.concatenate([self.mass, mass], axis=0) # Set them orbitting along the spin plane v_vec = np.cross(r_vec, [0, 0, 1]) self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0) def addHalo(self, model, totalMass, particles, rs)-
Adds a halo to the galaxy.
Parameters
model:string- parametrization of the halo. Only 'NFW' is supported.
totalMass:float- total mass of the halo
particles:int- number of particles in the halo
rs:float- characteristic length scale of the NFW profile.
Source code
def addHalo(self, model, totalMass, particles, rs): """Adds a halo to the galaxy. Parameters: model (string): parametrization of the halo. Only 'NFW' is supported. totalMass (float): total mass of the halo particles (int): number of particles in the halo rs (float): characteristic length scale of the NFW profile. """ if particles == 0: return None # Divide the mass equally among all particles mass = np.ones(particles)*totalMass/particles self.mass = np.concatenate([self.mass, mass], axis=0) # Create particles according to the radial distribution from model if model == 'NFW': r = NFW.ppf(np.random.rand(particles), scale=rs) else: raise Exception("""Bulge distribution not allowed. 'plummer' and 'hernquist' are the supported values""") r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles) self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0) # Orbit along random directions normal to the radial vector v_vec = np.cross(r_vec, random_unit_vectors(size=particles)) self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0) # Label the particles type_ = [['dark'], 0]*particles self.type = np.concatenate([self.type, type_], axis=0) def rotate(self, theta, phi)-
Rotates the galaxy so that its spin is along the (theta, phi) direction.
Parameters
theta:float- polar angle.
phi:float- azimuthal angle.
Source code
def rotate(self, theta, phi): """Rotates the galaxy so that its spin is along the (theta, phi) direction. Parameters: theta (float): polar angle. phi (float): azimuthal angle. """ M1 = np.array([[1, 0, 0], [0, np.cos(theta), np.sin(theta)], [0, -np.sin(theta), np.cos(theta)]]) M2 = np.array([[np.cos(phi), np.sin(phi), 0], [-np.sin(phi), np.cos(phi), 0], [0, 0, 1]]) M = np.matmul(M1, M2) # combine rotations self.r_vec = np.tensordot(self.r_vec, M, axes=[1, 0]) self.v_vec = np.tensordot(self.v_vec, M, axes=[1, 0])
class Simulation-
"Main class for the gravitational simulation.
Attributes
r_vec:array- position of the particles in the current timestep. Shape: (number of particles, 3)
rprev_vec:array- position of the particles in the previous timestep. Shape: (number of particles, 3)
v_vec:array- velocity in the current timestep. Shape: (number of particles, 3)
a_vec:array- acceleration in the current timestep. Shape: (number of particles, 3)
mass:array- mass of each particle in the simulation. Shape: (number of particles,)
type:array- non-unique identifier for each particle. Shape: (number of particles,)
tracks:array- list of positions through the simulation for central masses. Shape: (tracked particles, n+1, 3).
CONFIG:array- configuration used to create the simulation. It will be saved along the state of the simulation.
dt:float- timestep of the simulation
n:int- current timestep. Initialized as n=0.
soft:float- softening length used by the simulation.
verbose:boolean- When True progress statements will be printed.
Source code
class Simulation: """"Main class for the gravitational simulation. Attributes: r_vec (array): position of the particles in the current timestep. Shape: (number of particles, 3) rprev_vec (array): position of the particles in the previous timestep. Shape: (number of particles, 3) v_vec (array): velocity in the current timestep. Shape: (number of particles, 3) a_vec (array): acceleration in the current timestep. Shape: (number of particles, 3) mass (array): mass of each particle in the simulation. Shape: (number of particles,) type (array): non-unique identifier for each particle. Shape: (number of particles,) tracks (array): list of positions through the simulation for central masses. Shape: (tracked particles, n+1, 3). CONFIG (array): configuration used to create the simulation. It will be saved along the state of the simulation. dt (float): timestep of the simulation n (int): current timestep. Initialized as n=0. soft (float): softening length used by the simulation. verbose (boolean): When True progress statements will be printed. """ def __init__(self, dt, soft, verbose, CONFIG, method, **kwargs): """Constructor for the Simulation class. Arguments: dt (float): timestep of the simulation n (int): current timestep. Initialized as n=0. soft (float): softening length used by the simulation. verbose (bool): When True progress statements will be printed. CONFIG (dict): configuration file used to create the simulation. method (string): Optional. Algorithm to use when computing the gravitational forces. One of 'bruteForce', 'bruteForce_numba', 'bruteForce_numbaopt', 'bruteForce_CPP', 'barnesHut_CPP'. """ self.n = 0 self.t = 0 self.dt = dt self.soft = soft self.verbose = verbose self.CONFIG = CONFIG # Initialize empty arrays for all necessary properties self.r_vec = np.empty((0, 3)) self.v_vec = np.empty((0, 3)) self.a_vec = np.empty((0, 3)) self.mass = np.empty((0,)) self.type = np.empty((0, 2)) algorithms = { 'bruteForce': acceleration.bruteForce, 'bruteForceNumba': acceleration.bruteForceNumba, 'bruteForceNumbaOptimized': acceleration.bruteForceNumbaOptimized, 'bruteForceCPP': acceleration.bruteForceCPP, 'barnesHutCPP': acceleration.barnesHutCPP } try: self.acceleration = algorithms[method] except: raise Exception("Method '{}' unknown".format(method)) def add(self, body): """Add a body to the simulation. It must expose the public attributes body.r_vec, body.v_vec, body.a_vec, body.type, body.mass. Arguments: body: Object to be added to the simulation (e.g. a Galaxy object) """ # Extend all relevant attributes by concatenating the body for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']: simattr, bodyattr = getattr(self, name), getattr(body, name) setattr(self, name, np.concatenate([simattr, bodyattr], axis=0)) # Order based on mass order = np.argsort(-self.mass) for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']: setattr(self, name, getattr(self, name)[order]) # Update the list of objects to keep track of self.tracks = np.empty((np.sum(self.type[:,0]=='center'), 0, 3)) def step(self): """Perform a single step of the simulation. Makes use of a 4th order Verlet integrator. """ # Calculate the acceleration self.a_vec = self.acceleration(self.r_vec, self.mass, soft=self.soft) # Update the state using the Verlet algorithm # (A custom algorithm is written mainly for learning purposes) self.r_vec, self.rprev_vec = (2*self.r_vec - self.rprev_vec + self.a_vec * self.dt**2, self.r_vec) self.n += 1 # Update tracks self.tracks = np.concatenate([self.tracks, self.r_vec[self.type[:,0]=='center'][:,np.newaxis]], axis=1) def run(self, tmax, saveEvery, outputFolder, **kwargs): """Run the galactic simulation. Attributes: tmax (float): Time to which the simulation will run to. This is measured here since the start of the simulation, not since pericenter. saveEvery (int): The state is saved every saveEvery steps. outputFolder (string): It will be saved to /data/outputFolder/ """ # When the simulation starts, intialize self.rprev_vec self.rprev_vec = self.r_vec - self.v_vec * self.dt if self.verbose: print('Simulation starting. Bon voyage!') while(self.t < tmax): self.step() if(self.n % saveEvery == 0): self.save('data/{}'.format(outputFolder)) print('Simulation complete.') def save(self, outputFolder): """Save the state of the simulation to the outputFolder. Two files are saved: sim{self.n}.pickle: serializing the state. sim{self.n}.png: a simplified 2D plot of x, y. """ # Create the output folder if it doesn't exist if not os.path.exists(outputFolder): os.makedirs(outputFolder) # Compute some useful quantities # v_vec is not required by the integrator, but useful self.v_vec = (self.r_vec - self.rprev_vec)/self.dt self.t = self.n * self.dt # prevents numerical rounding errors # Serialize state file = open(outputFolder+'/data{}.pickle'.format(self.n), "wb") pickle.dump({'r_vec': self.r_vec, 'v_vec': self.v_vec, 'type': self.type, 'mass': self.mass, 'CONFIG': self.CONFIG, 't': self.t, 'tracks': self.tracks}, file) # Save simplified plot of the current state. # Its main use is to detect ill-behaved situations early on. fig = plt.figure() plt.xlim(-5, 5); plt.ylim(-5, 5); plt.axis('equal') # Dark halo is plotted in red, disk in blue, bulge in green PLTCON = [('dark', 'r', 0.3), ('disk', 'b', 1.0), ('bulge', 'g', 0.5)] for type_, c, a in PLTCON: plt.scatter(self.r_vec[self.type[:,0]==type_][:,0], self.r_vec[self.type[:,0]==type_][:,1], s=0.1, c=c, alpha=a) # Central mass as a magenta star plt.scatter(self.r_vec[self.type[:,0]=='center'][:,0], self.r_vec[self.type[:,0]=='center'][:,1], s=100, marker="*", c='m') # Save to png file fig.savefig(outputFolder+'/sim{}.png'.format(self.n), dpi=150) plt.close(fig) def project(self, theta, phi, view=0): """Projects the 3D simulation onto a plane as viewed from the direction described by the (theta, phi, view). Angles in radians. (This is used by the simulated annealing algorithm) Parameters: theta (float): polar angle. phi (float): azimuthal angle. view (float): rotation along line of sight. """ M1 = np.array([[np.cos(phi), np.sin(phi), 0], [-np.sin(phi), np.cos(phi), 0], [0, 0, 1]]) M2 = np.array([[1, 0, 0], [0, np.cos(theta), np.sin(theta)], [0, -np.sin(theta), np.cos(theta)]]) M3 = np.array([[np.cos(view), np.sin(view), 0], [-np.sin(view), np.cos(view), 0], [0, 0, 1]]) M = np.matmul(M1, np.matmul(M2, M3)) # combine rotations r = np.tensordot(self.r_vec, M, axes=[1, 0]) return r[:,0:2] def setOrbit(self, galaxy1, galaxy2, e, rmin, R0): """Sets the two galaxies galaxy1, galaxy2 in an orbit. Parameters: galaxy1 (Galaxy): 1st galaxy (main) galaxy2 (Galaxy): 2nd galaxy (companion) e: eccentricity of the orbit rmin: distance of closest approach R0: initial separation """ m1, m2 = np.sum(galaxy1.mass), np.sum(galaxy2.mass) # Relevant formulae: # $r_0 = r (1 + e) \cos(\phi)$, where $r_0$ ($\neq R_0$) is the semi-latus rectum # $r_0 = r_\textup{min} (1 + e)$ # $v^2 = G M (2/r - 1/a)$, where a is the semimajor axis # Solve the reduced two-body problem with reduced mass $\mu$ (mu) M = m1 + m2 r0 = rmin * (1 + e) try: phi = np.arccos((r0/R0 - 1) / e) # inverting the orbit equation phi = -np.abs(phi) # Choose negative (incoming) angle ainv = (1 - e) / rmin # ainv = $1/a$, as a may be infinite v = np.sqrt(M * (2/R0 - ainv)) # Finally, calculate the angle of motion. angle = tan(dy/dx) # $dy/dx = ((dr/d\phi) sin(\phi) + r \cos(\phi))/((dr/d\phi) cos(\phi) - r \sin(\phi))$ dy = R0/r0 * e * np.sin(phi)**2 + np.cos(phi) dx = R0/r0 * e * np.sin(phi) * np.cos(phi) - np.sin(phi) vangle = np.arctan2(dy, dx) except: raise Exception('''The orbital parameters cannot be satisfied. For elliptical orbits check that R0 is posible (<rmax).''') # We now need the actual motion of each body R_vec = np.array([[R0*np.cos(phi), R0*np.sin(phi), 0.]]) V_vec = np.array([[v*np.cos(vangle), v*np.sin(vangle), 0.]]) galaxy1.r_vec += m2/M * R_vec galaxy1.v_vec += m2/M * V_vec galaxy2.r_vec += -m1/M * R_vec galaxy2.v_vec += -m1/M * V_vec # Explicitely add the galaxies to the simulation self.add(galaxy1) self.add(galaxy2) if self.verbose: print('Galaxies set in orbit.')Methods
def __init__(self, dt, soft, verbose, CONFIG, method, **kwargs)-
Constructor for the Simulation class.
Arguments
dt:float- timestep of the simulation
n:int- current timestep. Initialized as n=0.
soft:float- softening length used by the simulation.
verbose:bool- When True progress statements will be printed.
CONFIG:dict- configuration file used to create the simulation.
method:string- Optional. Algorithm to use when computing the gravitational forces. One of 'bruteForce', 'bruteForce_numba', 'bruteForce_numbaopt', 'bruteForce_CPP', 'barnesHut_CPP'.
Source code
def __init__(self, dt, soft, verbose, CONFIG, method, **kwargs): """Constructor for the Simulation class. Arguments: dt (float): timestep of the simulation n (int): current timestep. Initialized as n=0. soft (float): softening length used by the simulation. verbose (bool): When True progress statements will be printed. CONFIG (dict): configuration file used to create the simulation. method (string): Optional. Algorithm to use when computing the gravitational forces. One of 'bruteForce', 'bruteForce_numba', 'bruteForce_numbaopt', 'bruteForce_CPP', 'barnesHut_CPP'. """ self.n = 0 self.t = 0 self.dt = dt self.soft = soft self.verbose = verbose self.CONFIG = CONFIG # Initialize empty arrays for all necessary properties self.r_vec = np.empty((0, 3)) self.v_vec = np.empty((0, 3)) self.a_vec = np.empty((0, 3)) self.mass = np.empty((0,)) self.type = np.empty((0, 2)) algorithms = { 'bruteForce': acceleration.bruteForce, 'bruteForceNumba': acceleration.bruteForceNumba, 'bruteForceNumbaOptimized': acceleration.bruteForceNumbaOptimized, 'bruteForceCPP': acceleration.bruteForceCPP, 'barnesHutCPP': acceleration.barnesHutCPP } try: self.acceleration = algorithms[method] except: raise Exception("Method '{}' unknown".format(method)) def add(self, body)-
Add a body to the simulation. It must expose the public attributes body.r_vec, body.v_vec, body.a_vec, body.type, body.mass.
Arguments
body- Object to be added to the simulation (e.g. a Galaxy object)
Source code
def add(self, body): """Add a body to the simulation. It must expose the public attributes body.r_vec, body.v_vec, body.a_vec, body.type, body.mass. Arguments: body: Object to be added to the simulation (e.g. a Galaxy object) """ # Extend all relevant attributes by concatenating the body for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']: simattr, bodyattr = getattr(self, name), getattr(body, name) setattr(self, name, np.concatenate([simattr, bodyattr], axis=0)) # Order based on mass order = np.argsort(-self.mass) for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']: setattr(self, name, getattr(self, name)[order]) # Update the list of objects to keep track of self.tracks = np.empty((np.sum(self.type[:,0]=='center'), 0, 3)) def project(self, theta, phi, view=0)-
Projects the 3D simulation onto a plane as viewed from the direction described by the (theta, phi, view). Angles in radians. (This is used by the simulated annealing algorithm)
Parameters
theta:float- polar angle.
phi:float- azimuthal angle.
view:float- rotation along line of sight.
Source code
def project(self, theta, phi, view=0): """Projects the 3D simulation onto a plane as viewed from the direction described by the (theta, phi, view). Angles in radians. (This is used by the simulated annealing algorithm) Parameters: theta (float): polar angle. phi (float): azimuthal angle. view (float): rotation along line of sight. """ M1 = np.array([[np.cos(phi), np.sin(phi), 0], [-np.sin(phi), np.cos(phi), 0], [0, 0, 1]]) M2 = np.array([[1, 0, 0], [0, np.cos(theta), np.sin(theta)], [0, -np.sin(theta), np.cos(theta)]]) M3 = np.array([[np.cos(view), np.sin(view), 0], [-np.sin(view), np.cos(view), 0], [0, 0, 1]]) M = np.matmul(M1, np.matmul(M2, M3)) # combine rotations r = np.tensordot(self.r_vec, M, axes=[1, 0]) return r[:,0:2] def run(self, tmax, saveEvery, outputFolder, **kwargs)-
Run the galactic simulation.
Attributes
tmax:float- Time to which the simulation will run to. This is measured here since the start of the simulation, not since pericenter.
saveEvery:int- The state is saved every saveEvery steps.
outputFolder:string- It will be saved to /data/outputFolder/
Source code
def run(self, tmax, saveEvery, outputFolder, **kwargs): """Run the galactic simulation. Attributes: tmax (float): Time to which the simulation will run to. This is measured here since the start of the simulation, not since pericenter. saveEvery (int): The state is saved every saveEvery steps. outputFolder (string): It will be saved to /data/outputFolder/ """ # When the simulation starts, intialize self.rprev_vec self.rprev_vec = self.r_vec - self.v_vec * self.dt if self.verbose: print('Simulation starting. Bon voyage!') while(self.t < tmax): self.step() if(self.n % saveEvery == 0): self.save('data/{}'.format(outputFolder)) print('Simulation complete.') def save(self, outputFolder)-
Save the state of the simulation to the outputFolder. Two files are saved: sim{self.n}.pickle: serializing the state. sim{self.n}.png: a simplified 2D plot of x, y.
Source code
def save(self, outputFolder): """Save the state of the simulation to the outputFolder. Two files are saved: sim{self.n}.pickle: serializing the state. sim{self.n}.png: a simplified 2D plot of x, y. """ # Create the output folder if it doesn't exist if not os.path.exists(outputFolder): os.makedirs(outputFolder) # Compute some useful quantities # v_vec is not required by the integrator, but useful self.v_vec = (self.r_vec - self.rprev_vec)/self.dt self.t = self.n * self.dt # prevents numerical rounding errors # Serialize state file = open(outputFolder+'/data{}.pickle'.format(self.n), "wb") pickle.dump({'r_vec': self.r_vec, 'v_vec': self.v_vec, 'type': self.type, 'mass': self.mass, 'CONFIG': self.CONFIG, 't': self.t, 'tracks': self.tracks}, file) # Save simplified plot of the current state. # Its main use is to detect ill-behaved situations early on. fig = plt.figure() plt.xlim(-5, 5); plt.ylim(-5, 5); plt.axis('equal') # Dark halo is plotted in red, disk in blue, bulge in green PLTCON = [('dark', 'r', 0.3), ('disk', 'b', 1.0), ('bulge', 'g', 0.5)] for type_, c, a in PLTCON: plt.scatter(self.r_vec[self.type[:,0]==type_][:,0], self.r_vec[self.type[:,0]==type_][:,1], s=0.1, c=c, alpha=a) # Central mass as a magenta star plt.scatter(self.r_vec[self.type[:,0]=='center'][:,0], self.r_vec[self.type[:,0]=='center'][:,1], s=100, marker="*", c='m') # Save to png file fig.savefig(outputFolder+'/sim{}.png'.format(self.n), dpi=150) plt.close(fig) def setOrbit(self, galaxy1, galaxy2, e, rmin, R0)-
Sets the two galaxies galaxy1, galaxy2 in an orbit.
Parameters
galaxy1:Galaxy- 1st galaxy (main)
galaxy2:Galaxy- 2nd galaxy (companion)
e- eccentricity of the orbit
rmin- distance of closest approach
R0- initial separation
Source code
def setOrbit(self, galaxy1, galaxy2, e, rmin, R0): """Sets the two galaxies galaxy1, galaxy2 in an orbit. Parameters: galaxy1 (Galaxy): 1st galaxy (main) galaxy2 (Galaxy): 2nd galaxy (companion) e: eccentricity of the orbit rmin: distance of closest approach R0: initial separation """ m1, m2 = np.sum(galaxy1.mass), np.sum(galaxy2.mass) # Relevant formulae: # $r_0 = r (1 + e) \cos(\phi)$, where $r_0$ ($\neq R_0$) is the semi-latus rectum # $r_0 = r_\textup{min} (1 + e)$ # $v^2 = G M (2/r - 1/a)$, where a is the semimajor axis # Solve the reduced two-body problem with reduced mass $\mu$ (mu) M = m1 + m2 r0 = rmin * (1 + e) try: phi = np.arccos((r0/R0 - 1) / e) # inverting the orbit equation phi = -np.abs(phi) # Choose negative (incoming) angle ainv = (1 - e) / rmin # ainv = $1/a$, as a may be infinite v = np.sqrt(M * (2/R0 - ainv)) # Finally, calculate the angle of motion. angle = tan(dy/dx) # $dy/dx = ((dr/d\phi) sin(\phi) + r \cos(\phi))/((dr/d\phi) cos(\phi) - r \sin(\phi))$ dy = R0/r0 * e * np.sin(phi)**2 + np.cos(phi) dx = R0/r0 * e * np.sin(phi) * np.cos(phi) - np.sin(phi) vangle = np.arctan2(dy, dx) except: raise Exception('''The orbital parameters cannot be satisfied. For elliptical orbits check that R0 is posible (<rmax).''') # We now need the actual motion of each body R_vec = np.array([[R0*np.cos(phi), R0*np.sin(phi), 0.]]) V_vec = np.array([[v*np.cos(vangle), v*np.sin(vangle), 0.]]) galaxy1.r_vec += m2/M * R_vec galaxy1.v_vec += m2/M * V_vec galaxy2.r_vec += -m1/M * R_vec galaxy2.v_vec += -m1/M * V_vec # Explicitely add the galaxies to the simulation self.add(galaxy1) self.add(galaxy2) if self.verbose: print('Galaxies set in orbit.') def step(self)-
Perform a single step of the simulation. Makes use of a 4th order Verlet integrator.
Source code
def step(self): """Perform a single step of the simulation. Makes use of a 4th order Verlet integrator. """ # Calculate the acceleration self.a_vec = self.acceleration(self.r_vec, self.mass, soft=self.soft) # Update the state using the Verlet algorithm # (A custom algorithm is written mainly for learning purposes) self.r_vec, self.rprev_vec = (2*self.r_vec - self.rprev_vec + self.a_vec * self.dt**2, self.r_vec) self.n += 1 # Update tracks self.tracks = np.concatenate([self.tracks, self.r_vec[self.type[:,0]=='center'][:,np.newaxis]], axis=1)