def add_one_object_hrzn(sim: rebound.Simulation, object_name: str, epoch: datetime, hrzn: Dict):
"""Add one object to a simulation with data fromm horizons (cache or API)."""
# Identifiers for this object
object_id = name_to_id[object_name]
key = (epoch, object_id)
try:
# Try to look up the object on the horizons cache
p: horizon_entry = hrzn[key]
sim.add(m=p.m, x=p.x, y=p.y, z=p.z, vx=p.vx, vy=p.vy, vz=p.vz, hash=rebound.hash(object_name))
except KeyError:
# Search string for the horizon API
horizon_name = object_to_horizon_name[object_name]
# Convert epoch to a horizon date string
horizon_date: str = datetime_to_horizons(epoch)
print(f'Searching Horizons as of {horizon_date}')
# Add the particle
sim.add(horizon_name, date=horizon_date)
# Set the mass and hash of this particle
p: rebound.Particle = sim.particles[-1]
# p.m = mass_tbl[object_name]
p.m = mass_tbl.get(object_name, 0.0)
p.hash = rebound.hash(object_name)
# Create an entry for this particle on the cache
entry: horizon_entry = horizon_entry(m=p.m, x=p.x, y=p.y, z=p.z, vx=p.vx, vy=p.vy, vz=p.vz,
object_name=object_name, object_id=object_id, horizon_name=horizon_name)
hrzn[key] = entry
# Save the revised cache
save_horizons_cache(hrzn)
def test_removing_particles(self):
self.sim.add(a=30.)
self.sim.remove(index=0)
self.sim.remove(index=0)
self.sim.remove(hash=rebound.hash("earth"), keepSorted=0)
self.assertEqual(self.sim.get_particle_by_hash("jupiter").hash, rebound.hash("jupiter"))
self.assertEqual(self.sim.N, 2)
with self.assertRaises(rebound.ParticleNotFound):
self.sim.get_particle_by_hash("earth")
self.sim.remove(hash="jupiter")
with self.assertRaises(rebound.ParticleNotFound):
self.sim.get_particle_by_hash("jupiter")
def test_removing_particles(self):
self.sim.add(a=30.)
self.sim.remove(0)
self.sim.remove(0)
self.sim.remove(hash=rebound.hash("earth"), keepSorted=0)
self.assertEqual(self.sim.particles["jupiter"].hash.value, rebound.hash("jupiter").value)
self.assertEqual(self.sim.N, 2)
with self.assertRaises(rebound.ParticleNotFound):
self.sim.particles["earth"]
self.sim.remove(hash="jupiter")
with self.assertRaises(rebound.ParticleNotFound):
self.sim.particles["jupiter"]
def test_removing_particles(self):
self.sim.add(a=30.)
self.sim.remove(0)
self.sim.remove(0)
self.sim.remove(hash=rebound.hash("earth"), keepSorted=0)
self.assertEqual(self.sim.particles["jupiter"].hash.value,
rebound.hash("jupiter").value)
self.assertEqual(self.sim.N, 2)
with self.assertRaises(rebound.ParticleNotFound):
self.sim.particles["earth"]
self.sim.remove(hash="jupiter")
with self.assertRaises(rebound.ParticleNotFound):
self.sim.particles["jupiter"]
def test_removing_particles(self):
self.sim.add(a=30.)
self.sim.remove(index=0)
self.sim.remove(index=0)
self.sim.remove(hash=rebound.hash("earth"), keepSorted=0)
self.assertEqual(
self.sim.get_particle_by_hash("jupiter").hash,
rebound.hash("jupiter"))
self.assertEqual(self.sim.N, 2)
with self.assertRaises(rebound.ParticleNotFound):
self.sim.get_particle_by_hash("earth")
self.sim.remove(hash="jupiter")
with self.assertRaises(rebound.ParticleNotFound):
self.sim.get_particle_by_hash("jupiter")
def make_sim(sim_name: str, object_names: List[str], epoch: datetime,
integrator: str, steps_per_day: int,
save_file: bool) -> rebound.Simulation:
"""Create or load simulation with the specified objects at the specified time"""
# Filename for archive
file_date: str = epoch.strftime('%Y-%m-%d_%H-%M')
fname_archive: str = f'../data/planets/{sim_name}_{file_date}.bin'
# If this file already exists, load it and check for both extra and missing bodies
sim: rebound.Simulation
try:
# Attempt to load the named file
sim = rebound.Simulation(fname_archive)
# print(f'Loaded {fname_archive}')
# Generate list of missing object names
objects_missing: List[str] = [
nm for nm in object_names if nm not in sim.particles
]
# Extend the simulation and save it with the augmented bodies
if objects_missing:
print(f'Found missing objects in {fname_archive}:')
print(objects_missing)
extend_sim_horizons(sim, object_names=objects_missing, epoch=epoch)
# Sets of named and input object hashes
hashes_sim: Set[int] = set(p.hash.value for p in sim.particles)
hashes_input: Set[str] = set(
rebound.hash(nm).value for nm in object_names)
# Filter the simulation so only the named objects are included
hashes_remove: List[str] = [
h for h in hashes_sim if h not in hashes_input
]
for h in hashes_remove:
sim.remove(hash=h)
except:
# Initialize simulation
sim = make_sim_horizons(object_names=object_names, epoch=epoch)
# Move to center of momentum
sim.move_to_com()
# Set integrator and time step
sim.integrator = integrator
dt: float = 1.0 / steps_per_day if steps_per_day > 0 else 0
sim.dt = dt
if integrator == 'ias15':
ias15 = sim.ri_ias15
ias15.min_dt = dt
# Save a snapshot to the archive file if requested
if save_file:
sim.simulationarchive_snapshot(filename=fname_archive, deletefile=True)
# Return the simulation
return sim
def extend_sim_horizons(sim: rebound.Simulation, object_names: List[str], epoch: datetime) -> None:
"""Extend a rebound simulation with initial data from the NASA Horizons system"""
# Generate list of missing object names
hashes_present: Set[int] = set(p.hash.value for p in sim.particles)
objects_missing: List[str] = [nm for nm in object_names if rebound.hash(nm).value not in hashes_present]
# Add missing objects one at a time if not already present
for object_name in objects_missing:
add_one_object_hrzn(sim=sim, object_name=object_name, epoch=epoch, hrzn=hrzn)
# Move to center of mass
sim.move_to_com()
def hash(self, value):
PY3 = sys.version_info[0] == 3
if PY3:
string_types = str,
int_types = int,
else:
string_types = basestring,
int_types = int, long,
if isinstance(value, int_types):
self._hash = value
elif isinstance(value, string_types):
self._hash = rebound.hash(value)
else:
raise AttributeError("Expecting integer or string as argument")
def extend_sim(sim: rebound.Simulation, object_names_new: List[str],
epoch: datetime):
"""Extend an existing simulation"""
# Generate list of missing object names
hashes_present: Set[int] = set(p.hash.value for p in sim.particles)
objects_missing: List[str] = [
nm for nm in object_names_new
if rebound.hash(nm).value not in hashes_present
]
# Extend the simulation and save it with the augmented bodies
if objects_missing:
extend_sim_horizons(sim, object_names=objects_missing, epoch=epoch)
return sim
def hash(self, value):
PY3 = sys.version_info[0] == 3
hash_types = c_uint32, c_uint, c_ulong
if PY3:
string_types = str,
int_types = int,
else:
string_types = basestring,
int_types = int, long,
if isinstance(value, hash_types):
self._hash = value.value
elif isinstance(value, string_types):
self._hash = rebound.hash(value).value
elif isinstance(value, int_types):
self._hash = value
else:
raise AttributeError("Hash must be set to an integer, a ctypes.c_uint32 or a string. See UniquelyIdentifyingParticles.ipynb ipython_example.")
def hash(self, value):
PY3 = sys.version_info[0] == 3
hash_types = c_uint32, c_uint, c_ulong
if PY3:
string_types = str,
int_types = int,
else:
string_types = basestring,
int_types = int, long,
if isinstance(value, hash_types):
self._hash = value.value
elif isinstance(value, string_types):
self._hash = rebound.hash(value).value
elif isinstance(value, int_types):
self._hash = value
else:
raise AttributeError("Hash must be set to an integer, a ctypes.c_uint32 or a string. See UniquelyIdentifyingParticles.ipynb ipython_example.")
def test_hash(self):
self.case(hashes=[0,1,2,3])
self.case(hashes=[0,None,2,3])
self.case(hashes=[2,3,1,4])
self.case(hashes=[2,4,None,3])
self.case(hashes=[2,0,4,3])
self.case(hashes=[2,0,4,None])
self.case(hashes=[2,3,4,0])
self.case(hashes=[None,4,3,0])
e = rebound.hash("earth").value
self.case(hashes=[0,1,2,e])
self.case(hashes=[0,None,e,3])
self.case(hashes=[2,3,e,4])
self.case(hashes=[2,e,None,3])
self.case(hashes=[2,0,e,3])
self.case(hashes=[e,0,4,None])
self.case(hashes=[2,e,4,0])
self.case(hashes=[None,4,e,0])
def test_hash(self):
self.case(hashes=[0, 1, 2, 3])
self.case(hashes=[0, None, 2, 3])
self.case(hashes=[2, 3, 1, 4])
self.case(hashes=[2, 4, None, 3])
self.case(hashes=[2, 0, 4, 3])
self.case(hashes=[2, 0, 4, None])
self.case(hashes=[2, 3, 4, 0])
self.case(hashes=[None, 4, 3, 0])
e = rebound.hash("earth").value
self.case(hashes=[0, 1, 2, e])
self.case(hashes=[0, None, e, 3])
self.case(hashes=[2, 3, e, 4])
self.case(hashes=[2, e, None, 3])
self.case(hashes=[2, 0, e, 3])
self.case(hashes=[e, 0, 4, None])
self.case(hashes=[2, e, 4, 0])
self.case(hashes=[None, 4, e, 0])
def make_sim_asteroids(sim_base: rebound.Simulation,
ast_elt: pd.DataFrame,
n0: int, n1: int,
progbar: bool = False) -> Tuple[rebound.Simulation, List[str]]:
"""
Create a simulation with the selected asteroids by their ID numbers.
INPUTS:
sim_base: the base simulation, with e.g. the sun, planets, and selected moons
ast_elt: the DataFrame with asteroid orbital elements at the specified epoch
n0: the first asteroid number to add, inclusive
n1: the last asteroid number to add, exclusive
"""
# Start with a copy of the base simulation
sim = sim_base.copy()
# Set the number of active particles to the base simulation
# https://rebound.readthedocs.io/en/latest/ipython/Testparticles.html
sim.N_active = sim_base.N
# Add the specified asteroids one at a time
mask = (n0 <= ast_elt.Num) & (ast_elt.Num < n1)
nums = ast_elt.index[mask]
iters = tqdm_console(nums) if progbar else nums
for num in iters:
# Unpack the orbital elements
a = ast_elt.a[num]
e = ast_elt.e[num]
inc = ast_elt.inc[num]
Omega = ast_elt.Omega[num]
omega = ast_elt.omega[num]
M = ast_elt.M[num]
name = ast_elt.Name[num]
# Set the primary to the sun (NOT the solar system barycenter!)
primary = sim.particles['Sun']
# Add the new asteroid
sim.add(m=0.0, a=a, e=e, inc=inc, Omega=Omega, omega=omega, M=M, primary=primary)
# Set the hash to the asteroid's name
sim.particles[-1].hash = rebound.hash(name)
# The corresponding list of asteroid names
asteroid_names = [name for name in ast_elt.Name[nums]]
# Return the new simulation including the asteroids
return sim, asteroid_names
def make_sim_horizons(object_names: List[str], horizon_date: str):
"""Create a new rebound simulation with initial data from the NASA Horizons system"""
# Create a simulation
sim = rebound.Simulation()
# Set units
sim.units = ('yr', 'AU', 'Msun')
# Add these objects from Horizons
sim.add(object_names, date=horizon_date)
# Add hashes for the object names
for i, particle in enumerate(sim.particles):
particle.hash = rebound.hash(object_names[i])
# Move to center of mass
sim.move_to_com()
return sim
def make_sim_from_elts(elts: Dict[str, np.array], epoch: float):
"""
Create a simulation for planets and asteroids parameterized by their orbital elements
INPUTS:
elts: dictionary with keys 'a', 'e', etc. for 6 orbital elements. values arrays of shape (N,)
epoch: MJD as of which these orbital elements apply
"""
# Convert epoch to a datetime
epoch_dt = mjd_to_datetime(epoch)
# Base Rebound simulation of the planets and moons on this date
sim = make_sim_planets(epoch=epoch_dt)
# Set the number of active particles to the base simulation
sim.N_active = sim.N
# Unpack the orbital elements
a = elts['a']
e = elts['e']
inc = elts['inc']
Omega = elts['Omega']
omega = elts['omega']
f = elts['f']
# Get the number of asteroids in the data set
n: int = a.shape[0]
# Add the specified asteroids one at a time
for i in range(n):
# Set the primary to the sun (NOT the solar system barycenter!)
# Put this inside the loop b/c not guaranteed to remain constant as particles are added
primary = sim.particles['Sun']
# Add the new asteroid
sim.add(m=0.0, a=a[i], e=e[i], inc=inc[i], Omega=Omega[i], omega=omega[i], f=f[i], primary=primary)
# Set the hash to the asteroid's number in this batch
sim.particles[-1].hash = rebound.hash(f'{i}')
# Return the new simulation including the asteroids
return sim
def get_wd_hash():
wd_hash = rebound.hash('wd').value
return wd_hash
def test_add(self):
self.sim.add(a=7., hash = "planet 9")
self.assertEqual(rebound.hash("planet 9"), self.sim.get_particle_by_hash("planet 9").hash)
def test_rebound_hash(self):
self.assertNotEqual(rebound.hash("foo"), rebound.hash("bar"))
self.assertEqual(rebound.hash("earth"), 1424801690)
def test_add(self):
self.sim.add(a=7., hash="planet 9")
self.assertEqual(
rebound.hash("planet 9").value,
self.sim.particles["planet 9"].hash.value)
def propagate(self, epochs, units=Units(), **kwargs):
'''
Numerically integrate all bodies to the desired date.
This routine synchronizes the epochs.
'''
if kwargs.get('func') is not None:
f = True
func = kwargs.get('func')
else:
f = False
units.timescale = 'tdb'
epochs = self.detect_timescale(np.atleast_1d(epochs), units.timescale)
self.move_to_com()
#for time in track(np.sort(epochs.tdb.jd)):
# self.integrate(time, exact_finish_time=1)
for time in np.sort(epochs.tdb.jd):
self.integrate(time, exact_finish_time=1)
if f == True:
func(self)
for p in self.particles:
h = p.hash.value
arr = [time] + p.xyz + p.vxyz
self.simdata[h].append(arr)
pepoch = []
px = []
py = []
pz = []
pvx = []
pvy = []
pvz = []
pname = []
for n in self.perturber_names:
ts, xs, ys, zs, vxs, vys, vzs = np.array(
self.simdata[rebound.hash(n).value]).T
pepoch.append(ts.tolist())
px.append(xs.tolist())
py.append(ys.tolist())
pz.append(zs.tolist())
pvx.append(vxs.tolist())
pvy.append(vys.tolist())
pvz.append(vzs.tolist())
pname.append([n for _ in range(len(ts))])
pepoch = np.hstack(pepoch)
px = np.hstack(px)
py = np.hstack(py)
pz = np.hstack(pz)
pvx = np.hstack(pvx)
pvy = np.hstack(pvy)
pvz = np.hstack(pvz)
pname = np.hstack(pname)
planets = SpaceRock(x=px,
y=py,
z=pz,
vx=pvx,
vy=pvy,
vz=pvz,
name=pname,
epoch=pepoch,
origin='ssb',
units=units)
if hasattr(self, 'testparticle_names'):
epoch = []
x = []
y = []
z = []
vx = []
vy = []
vz = []
name = []
for n in self.testparticle_names:
ts, xs, ys, zs, vxs, vys, vzs = np.array(
self.simdata[rebound.hash(n).value]).T
epoch.append(ts.tolist())
x.append(xs.tolist())
y.append(ys.tolist())
z.append(zs.tolist())
vx.append(vxs.tolist())
vy.append(vys.tolist())
vz.append(vzs.tolist())
name.append([n for _ in range(len(ts))])
epoch = np.hstack(epoch)
x = np.hstack(x)
y = np.hstack(y)
z = np.hstack(z)
vx = np.hstack(vx)
vy = np.hstack(vy)
vz = np.hstack(vz)
name = np.hstack(name)
rocks = SpaceRock(x=x,
y=y,
z=z,
vx=vx,
vy=vy,
vz=vz,
name=name,
epoch=epoch,
origin='ssb',
units=units)
else:
rocks = ()
return rocks, planets, self
'Varda': ao + 174567,
'2007 OR10': ao + 225088,
'229762': ao + 229762,
# Miscellaneous heavy non-numbered asteroids
'Vanth': uo + 1,
'Hiʻiaka': uo + 2,
'2001 QC298': uo + 3,
}
# *********************************************************************************************************************
# Dictionary mapping from ID to name
id_to_name: Dict[int, str] = {name_to_id[name]: name for name in name_to_id}
# Dictionary mapping from rebound hash to name
hash_to_name: Dict[int, str] = {
rebound.hash(name_to_id[name]).value: name
for name in name_to_id
}
name_to_hash: Dict[str, int] = {
name: rebound.hash(name_to_id[name]).value
for name in name_to_id
}
# *********************************************************************************************************************
# Mass of all solar system objects available on Wikipedia
# https://en.wikipedia.org/wiki/List_of_Solar_System_objects_by_size
mass_tbl = {
'Sun': 1.000000000000E+00,
# Inner planets
'Mercury': 1.660506399135E-07,
'Venus': 2.448266324709E-06,
def test_add(self):
self.sim.add(a=7., hash="planet 9")
self.assertEqual(rebound.hash("planet 9"),
self.sim.get_particle_by_hash("planet 9").hash)
def calc_ast_pos(elts: Dict[str, np.array], epoch: float, ts: np.array) -> np.array:
"""
Calculate asteroid positions from the given elements on the fly
INPUTS:
elts: dictionary with keys 'a', 'e', etc. for 6 orbital elements. values arrays of shape (N,)
epoch: MJD as of which these orbital elements apply
ts: array of MJDs as of which
Outputs:
q_earth: position of earth at input times; shape (traj_size, 3,)
q_ast: positions of asteroids at input times; shape (num_ast, traj_size, 3,)
Note: this function has some overlap with make_archive_impl in rebount_utils.py
It's different though because it's intended to generate arrays on the fly, not save them to disk.
"""
# Build the simulation at the epoch
sim_epoch = make_sim_from_elts(elts, epoch)
# Create copies of the simulation to integrate forward and backward
sim_fwd: rebound.Simulation = sim_epoch.copy()
sim_back: rebound.Simulation = sim_epoch.copy()
# Set the time counter on both simulation copies to the epoch time
sim_fwd.t = epoch
sim_back.t = epoch
# Set the times for snapshots in both directions;
idx: int = np.searchsorted(ts, epoch, side='left')
ts_fwd: np.array = ts[idx:]
ts_back: np.array = ts[:idx][::-1]
# Number of snapshots
M_back: int = len(ts_back)
M_fwd: int = len(ts_fwd)
M: int = M_back + M_fwd
# Number of particles
N = sim_epoch.N
# Number of asteroids (smaller than total number of particles; not including sun and planets)
N_ast: int = elts['a'].shape[0]
# Number of heavy bodies; saved before first asteroid in output array
N_heavy: int = N - N_ast
# The sun is the primary and is always in column 0
sun_idx = 0
# Column index of Earth by searching for hash of earth in all hashes of the simulation
hash_earth = rebound.hash('Earth')
hashes = np.zeros(N, dtype=np.uint32)
sim_epoch.serialize_particle_data(hash=hashes)
earth_idx = np.argmax(hashes==hash_earth)
# Initialize array for the positions of all particles in the simulation
# Match the order required by rebound for efficient serialization!
# This must be permuted at the end
q: np.array = np.zeros(shape=(M, N, 3,), dtype=np.float64)
# Subfunction: process one row of the loop
def process_row(sim: rebound.Simulation, t: float, row: int):
# Integrate to the current time step with an exact finish time
sim.integrate(t, exact_finish_time=1)
# Serialize the positions of earth and all the bodies
sim.serialize_particle_data(xyz=q[row])
# process times in the forward direction
for i, t in enumerate(ts_fwd):
# Row index for position data
row: int = M_back + i
# Process this row
process_row(sim=sim_fwd, t=t, row=row)
# process times in the backward direction
for i, t in enumerate(ts_back):
# Row index for position data
row: int = M_back - (i+1)
# Process this row
process_row(sim=sim_back, t=t, row=row)
# Position of sun is the relevant column
q_sun = q[:, sun_idx, 0:3]
# Position of earth is the relevant column
q_earth = q[:, earth_idx, 0:3]
# The asteroid positions are in the right-hand slice of q_all
# Transpose so resulting array has size (N_ast, traj_size, 3)
q_ast = q[:, N_heavy:N, 0:3].transpose((1,0,2))
return q_ast, q_sun, q_earth
def test_add(self):
self.sim.add(a=7., hash = "planet 9")
self.assertEqual(rebound.hash("planet 9").value, self.sim.particles["planet 9"].hash.value)
def test_rebound_hash(self):
self.assertNotEqual(
rebound.hash("foo").value,
rebound.hash("bar").value)
self.assertEqual(rebound.hash("earth").value, 1424801690)
