def test_element_recovery(verbose: bool = False) -> bool:
"""Test recovery of initial orbital elements for selected asteroids"""
# List of asteroids to test: first 25
asteroid_names = [
'Ceres', 'Pallas', 'Juno', 'Vesta', 'Astraea',
'Hebe', 'Iris', 'Flora', 'Metis', 'Hygiea',
'Parthenope', 'Victoria', 'Egeria', 'Irene', 'Eunomia',
'Psyche', 'Thetis', 'Melpomene', 'Fortuna', 'Massalia',
'Lutetia', 'Kalliope', 'Thalia', 'Phocaea']
# Load asteroid data as DataFrame
ast_elt = load_data()
# Get the epoch from the DataFrame
epoch_mjd: float = ast_elt.epoch_mjd[1]
epoch: datetime = mjd_to_datetime(epoch_mjd)
# Rebound simulation of the planets and moons on this date
sim_base = make_sim_planets(epoch=epoch)
# Add selected asteroids
sim_ast, asteroid_names_out = make_sim_asteroids(sim_base=sim_base, ast_elt=ast_elt, n0=1, n1=31)
# Create the reference simulation
sim_hrzn = make_sim_asteroids_horizons(asteroid_names=asteroid_names, epoch=epoch)
# Report the difference
object_names = ['Earth'] + asteroid_names
pos_err, ang_err = \
report_sim_difference(sim0=sim_hrzn, sim1=sim_ast, object_names=object_names, verbose=True)
# Report details of one specific asteroid
report_one_asteroid(sim=sim_ast, asteroid_name='Ceres', epoch=epoch, verbose=True)
# Threshold for pass
pos_tol: float = 1.0E-5
ang_tol: float = 2.0
# Test result
isOK: bool = all(pos_err < pos_tol) and (ang_err < ang_tol)
msg: str = 'PASS' if isOK else 'FAIL'
print(f'\n***** {msg} *****')
return isOK
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 main():
"""Integrate the orbits of the planets and major moons"""
# Process command line arguments
parser = argparse.ArgumentParser(
description='Integrate the orbits of planets and moons.')
parser.add_argument(
'type',
nargs='?',
metavar='type',
type=str,
default='A',
help='type of integration: p- planets; d- dwarfs; m-moons; a-all, '
'A-run all 4 strategies')
parser.add_argument(
'steps_per_day',
nargs='?',
metavar='SPD',
type=int,
default=-1,
help='the (max) number of steps per day taken by the integrator')
parser.add_argument('--test',
default=False,
action='store_true',
help='run in test mode')
args = parser.parse_args()
# Unpack command line arguments
integration_type = args.type
steps_per_day: int = args.steps_per_day
# Flags for planets and moons
run_planets: bool = integration_type in ('p', 'A')
run_moons: bool = integration_type in ('m', 'A')
run_dwarfs: bool = integration_type in ('d', 'A')
run_all: bool = integration_type in ('a', 'A')
# Reference epoch for asteroids file
epoch_mjd: float = 58600.0
# Convert to a datetime
epoch: datetime = mjd_to_datetime(epoch_mjd)
# epoch_dt: datetime = datetime(2019,4,27)
# Start and end times of simulation
dt0: datetime = datetime(2000, 1, 1)
dt1: datetime = datetime(2040, 12, 31)
# Integrator choices
integrator_planets: str = 'ias15'
integrator_moons: str = 'ias15'
integrator_dwarfs: str = 'ias15'
integrator_all: str = 'ias15'
# Integrator time step
steps_per_day_planets: int = steps_per_day if steps_per_day >= 0 else 16
steps_per_day_moons: int = steps_per_day if steps_per_day >= 0 else 16
steps_per_day_dwarfs: int = steps_per_day if steps_per_day >= 0 else 16
steps_per_day_all: int = steps_per_day if steps_per_day >= 0 else 16
# Initial configuration of planets, moons, and dwarfs
sim_planets = make_sim_planets(epoch=epoch,
integrator=integrator_planets,
steps_per_day=steps_per_day_planets)
sim_moons = make_sim_moons(epoch=epoch,
integrator=integrator_moons,
steps_per_day=steps_per_day_moons)
sim_dwarfs = make_sim_dwarfs(epoch=epoch,
integrator=integrator_dwarfs,
steps_per_day=steps_per_day_dwarfs)
sim_all = make_sim_all(epoch=epoch,
integrator=integrator_all,
steps_per_day=steps_per_day_all)
# If we are in test mode, add the asteroids to the base simulations
if args.test:
add_test_asteroids(sim=sim_planets,
object_names=object_names_planets,
test_asteroids=test_asteroids,
epoch=epoch)
add_test_asteroids(sim=sim_moons,
object_names=object_names_moons,
test_asteroids=test_asteroids,
epoch=epoch)
add_test_asteroids(sim=sim_dwarfs,
object_names=object_names_dwarfs,
test_asteroids=test_asteroids,
epoch=epoch)
add_test_asteroids(sim=sim_all,
object_names=object_names_all,
test_asteroids=test_asteroids,
epoch=epoch)
# Shared time_step and save_step
time_step: int = 1
save_step: int = 1
# Integrate the planets from dt0 to dt1
if run_planets:
sim_name = 'planets' if not args.test else 'planets_test'
test_name = 'planets_com'
sim_long_name = 'sun and 8 planets'
# Run the planets simulation
sa_planets = \
run_one_sim(sim_name=sim_name, sim_epoch=sim_planets, object_names=object_names_planets,
epoch=epoch, dt0=dt0, dt1=dt1,
time_step=time_step, save_step=save_step, steps_per_day=steps_per_day_planets)
# Test the planets simulation
if run_planets and args.test:
pos_err_planets, ang_err_planets = \
test_one_sim(sa=sa_planets, test_objects=test_objects_planets,
sim_name=sim_name, test_name=test_name, sim_long_name=sim_long_name,
integrator=integrator_planets, steps_per_day=steps_per_day_planets,
make_plot_internal=False, verbose_internal=False,
make_plot_asteroids=False, verbose_asteroids=False)
# Integrate the planets and moons from dt0 to dt1
if run_moons:
sim_name = 'moons' if not args.test else 'moons_test'
test_name = 'planets'
sim_long_name = f'{len(object_names_moons)} objects in solar system: sun, planets, moons'
# Run the moons simulation
sa_moons = \
run_one_sim(sim_name=sim_name, sim_epoch=sim_moons, object_names=object_names_moons,
epoch=epoch, dt0=dt0, dt1=dt1,
time_step=time_step, save_step=save_step, steps_per_day=steps_per_day_moons)
# Test the moons simulation
if run_moons and args.test:
pos_err_moons, ang_err_moons = \
test_one_sim(sa=sa_moons, test_objects=test_objects_moons,
sim_name=sim_name, test_name=test_name, sim_long_name=sim_long_name,
integrator=integrator_moons, steps_per_day=steps_per_day_moons,
make_plot_internal=True, verbose_internal=False,
make_plot_asteroids=True, verbose_asteroids=False)
# Integrate the planets and dwarf planets from dt0 to dt1
if run_dwarfs:
sim_name = 'dwarfs' if not args.test else 'dwarfs_test'
test_name = 'planets'
sim_long_name = f'{len(object_names_moons)} objects in solar system: sun, planets, dwarf planets'
# Run the dwarfs simulation
sa_dwarfs = \
run_one_sim(sim_name=sim_name, sim_epoch=sim_dwarfs, object_names=object_names_dwarfs,
epoch=epoch, dt0=dt0, dt1=dt1,
time_step=time_step, save_step=save_step, steps_per_day=steps_per_day_dwarfs)
# Test the dwarfs simulation
if run_dwarfs and args.test:
pos_err_dwarfs, ang_err_dwarfs = \
test_one_sim(sa=sa_dwarfs, test_objects=test_objects_dwarfs,
sim_name=sim_name, test_name=test_name, sim_long_name=sim_long_name,
integrator=integrator_dwarfs, steps_per_day=steps_per_day_dwarfs,
make_plot_internal=True, verbose_internal=False,
make_plot_asteroids=True, verbose_asteroids=False)
# Integrate all the bodies from dt0 to dt1
if run_all:
sim_name = 'all' if not args.test else 'all_test'
test_name = 'planets'
sim_long_name = f'all {len(object_names_all)} heavy objects in solar system: ' \
f'sun, planets, moons, dwarf planets'
# Run the all simulation
sa_all = \
run_one_sim(sim_name=sim_name, sim_epoch=sim_all, object_names=object_names_all,
epoch=epoch, dt0=dt0, dt1=dt1,
time_step=time_step, save_step=save_step, steps_per_day=steps_per_day_all)
# Test the all simulation
if run_all and args.test:
pos_err_all, ang_err_all = \
test_one_sim(sa=sa_all, test_objects=test_objects_all,
sim_name=sim_name, test_name=test_name, sim_long_name=sim_long_name,
integrator=integrator_all, steps_per_day=steps_per_day_all,
make_plot_internal=True, verbose_internal=False,
make_plot_asteroids=True, verbose_asteroids=False)
if args.test:
# Plot position error
fig, ax = plt.subplots(figsize=[16, 10])
ax.set_title(f'Position Error on 10 Test Asteroids')
ax.set_ylabel('RMS Position Error in AU')
ax.ticklabel_format(axis='y', style='sci', scilimits=(
0,
0,
))
test_years = list(range(2000, 2041))
if run_planets:
ax.plot(test_years,
pos_err_planets,
label='planets',
marker='+',
color='blue')
if run_moons:
ax.plot(test_years,
pos_err_moons,
label='moons',
marker='x',
color='green')
if run_dwarfs:
ax.plot(test_years,
pos_err_dwarfs,
label='dwarfs',
marker='o',
color='red')
ax.grid()
ax.legend()
fig.savefig(
fname=f'../figs/integration_test/planets/sim_ang_error_comp.png',
bbox_inches='tight')
# Plot angle error
fig, ax = plt.subplots(figsize=[16, 10])
ax.set_title(f'Angle Error on 10 Test Asteroids')
ax.set_ylabel('RMS Angle Error vs. Earth in Arcseconds')
test_years = list(range(2000, 2041))
if run_planets:
ax.plot(test_years,
ang_err_planets,
label='planets',
marker='+',
color='blue')
if run_moons:
ax.plot(test_years,
ang_err_moons,
label='moons',
marker='x',
color='green')
if run_dwarfs:
ax.plot(test_years,
ang_err_dwarfs,
label='dwarfs',
marker='o',
color='red')
ax.grid()
ax.legend()
fig.savefig(
fname=f'../figs/integration_test/planets/sim_pos_error_comp.png',
bbox_inches='tight')
def main():
"""Main routine for integrating the orbits of known asteroids"""
# Process command line arguments
parser = argparse.ArgumentParser(description='Integrate the orbits of known asteroids from JPL ephemeris file.')
parser.add_argument('n0', nargs='?', metavar='n0', type=int, default=0,
help='the first asteroid number to process')
parser.add_argument('n_ast', nargs='?', metavar='B', type=int, default=1000,
help='the number of asteroids to process in this batch')
parser.add_argument('--progress', default=False, action='store_true',
help='display progress bar')
parser.add_argument('--test', default=False, action='store_true',
help='run in test mode')
args = parser.parse_args()
# If run in test mode, run tests without processing any asteroid trajectories
if args.test:
# Test that initial orbital elements recovered from the JPL file
print_header(f'Testing recovery of initial orbital elements with JPL text file vs. Horizons')
test_element_recovery(verbose=True)
# Test the integration vs. Horizons
print_header(f'Testing asteroid integration vs. Horizons')
test_asteroid_sim(verbose=True, make_plot=True)
# Test numpy arrays
print_header(f'Testing Numpy array vs. simulation archive:')
test_numpy(verbose=True)
# Quit early in test mode: don't want to do any integrations
print()
exit()
# Unpack command line arguments
n0: int = args.n0
n1: int = n0 + args.n_ast
progbar: bool = args.progress
# Load asteroid data as DataFrame
ast_elt: pd.DataFrame = load_data()
# Get the epoch from the DataFrame
epoch_mjd: float = ast_elt.epoch_mjd[1]
epoch: datetime = mjd_to_datetime(epoch_mjd)
# Start and end times of simulation
dt0: datetime = datetime(2000, 1, 1)
dt1: datetime = datetime(2040,12,31)
# Rebound simulation of the planets on this date
integrator: str = 'ias15'
steps_per_day: int = 16
sim_base: rebound.Simulation = make_sim_planets(epoch=epoch, integrator=integrator, steps_per_day=steps_per_day)
# Add selected asteroids
sim: rebound.Simulation
asteroid_names: List[str]
sim, asteroid_names = make_sim_asteroids(sim_base=sim_base, ast_elt=ast_elt, n0=n0, n1=n1)
# The list of object names corresponding to this simulation
object_names: List[str] = object_names_planets + asteroid_names
# Integrate the asteroids from dt0 to dt1 with a time step of 1 day
fname: str = f'../data/asteroids/sim_asteroids_n_{n0:06}_{n1:06}.bin'
time_step: int = 1
save_step: int = 32
save_elements: bool = True
print(f'Processing asteroid trajectories for asteroid numbers {n0} to {n1}...')
make_archive(fname_archive=fname, sim_epoch=sim, object_names=object_names,
epoch=epoch, dt0=dt0, dt1=dt1,
time_step=time_step, save_step=save_step,
save_elements=save_elements, progbar=progbar)
def ast_data_add_calc_elements(ast_elt) -> pd.DataFrame:
"""Add the true anomaly and other calculated orbital elements to the asteroid DataFrame"""
# Number of asteroids
N: int = len(ast_elt)
# Initialize empty arrays for computed orbital elements
f = np.zeros(N)
P = np.zeros(N)
mean_motion = np.zeros(N)
long = np.zeros(N)
theta = np.zeros(N)
pomega = np.zeros(N)
T_peri = np.zeros(N)
# Get the epoch from the DataFrame
epoch_mjd: float = ast_elt.epoch_mjd[1]
epoch: datetime = mjd_to_datetime(epoch_mjd)
# Rebound simulation of the planets and moons on this date
sim_base = make_sim_planets(epoch=epoch)
# Make a gigantic simulation with all these asteroids
n0: int = np.min(ast_elt.Num)
n1 = np.max(ast_elt.Num) + 1
print(f'Making big simulation with all {n1} asteroids...')
sim_ast, asteroid_names = make_sim_asteroids(sim_base=sim_base, ast_elt=ast_elt, n0=n0, n1=n1, progbar=True)
# Calculate orbital elements for all particles; must specify primary = Sun!!!
print(f'Computing orbital elements...')
orbits = sim_ast.calculate_orbits(primary=sim_ast.particles[0])
# Slice orbits so it skips the planets and only includes the asteroids
orbits = orbits[sim_base.N-1:]
# Iterate over all the asteroids in the simulation
print(f'Copying additional orbital elements to DataFrame...')
mask = (n0 <= ast_elt.Num) & (ast_elt.Num < n1)
iters = list(enumerate(ast_elt.Num[mask]))
for i, num in tqdm_console(iters):
# Look up the orbit of asteroid i
orb = orbits[i]
# Unpack the additional (calculated) orbital elements
f[i] = orb.f
P[i] = orb.P
mean_motion[i] = orb.n
long[i] = orb.l
theta[i] = orb.theta
pomega[i] = orb.pomega
T_peri[i] = orb.T
# Save computed orbital elements to the DataFrame
ast_elt['f'] = f
ast_elt['P'] = P
ast_elt['n'] = mean_motion
ast_elt['long'] = long
ast_elt['theta'] = theta
ast_elt['pomega'] = pomega
ast_elt['T_peri'] = T_peri
# Return the updated DataFrame
return ast_elt