def labels(row):
try:
sa = rebound.SimulationArchive(pathtosa+'sa'+row['runstring'])
sim = sa[0]
P1 = sim.particles[1].P # Need initial orbital period for TTVsystems, where P1 != 1
try: # Needed for old integrations (random and Naireen) because no snapshot on end
sim = rebound.Simulation(pathtosa+'../../final_conditions/runs/fc'+row['runstring'])
row['instability_time'] = sim.t/P1
try:
ssim = rebound.Simulation(pathtossa+'../../final_conditions/shadowruns/fc'+row['runstring'])
row['shadow_instability_time'] = ssim.t/P1
except:
print('No shadow for random {0}'.format(row['runstring']))
except: # New runs (resonant and Ari) have snapshots at collision
sim = sa[-1]
if sim.t > 9.99e3 and sim.t < 1.0001e4: # catch all stable integrations accidentally thrown out. Have checked this catches all thrown out, and doesn't catch any with actual instability times in this range
row['instability_time'] = 1e9
row['shadow_instability_time'] = 1e9
else:
row['instability_time'] = sim.t/P1
ssa = rebound.SimulationArchive(pathtossa+'sa'+row['runstring'])
ssim = ssa[-1]
row['shadow_instability_time'] = ssim.t/P1
row['Stable'] = row['instability_time'] > 9.99e8
except:
print('Error', pathtosa+'sa'+row['runstring'])
return row
def test_append_to_corrupt_snapshot_by_2bytes(self):
sim = rebound.Simulation()
sim.add(m=1.)
sim.add(m=1e-3, a=1.)
sim.add(m=5e-3, a=2.25)
sim.automateSimulationArchive("simulationarchive.bin",
interval=1000,
deletefile=True)
sim.integrate(3001)
s1 = os.path.getsize('simulationarchive.bin')
with open('simulationarchive.bin', 'r+b') as f:
f.seek(0, os.SEEK_END)
f.seek(f.tell() - 2, os.SEEK_SET)
f.truncate() # truncate by 2 bytes
s2 = os.path.getsize('simulationarchive.bin')
self.assertEqual(s1, s2 + 2)
sa = rebound.SimulationArchive("simulationarchive.bin")
sim = sa[-1]
sim.automateSimulationArchive("simulationarchive.bin", interval=1000)
sim.integrate(7001)
sa = rebound.SimulationArchive("simulationarchive.bin")
self.assertEqual(sa.nblobs, 8)
self.assertAlmostEqual(sa[-1].t, 7000, places=0)
def es(system, Nshadows, tmax=1.e4):
distpath = 'hussain2019data/resonant_distributions/'
folder = distpath + "Res_sys_{0}_1e8/simulation_archives/".format(
system) #ic{1:0=7d}.bin".format(system, 0)
root, dirs, files = next(os.walk(folder))
Nsys = 0
for file in files:
try:
sim = rebound.SimulationArchive(folder + file)[0]
Nsys += 1
except:
print('Didnt load')
Nout = 1000
data = np.zeros((Nsys + 1, Nout))
for j, file in enumerate(files[:Nshadows]):
sim = rebound.SimulationArchive(folder + file)[0]
sim.collision_resolve = collision
sim.exit_max_distance = 100.
ps = sim.particles
times = np.logspace(0, np.log10(tmax), Nout)
for i, time in enumerate(times):
try:
sim.integrate(time)
data[j, i] = ps[2].e
except:
break
data[-1, :] = times
return data
def labels(row):
try:
sa = rebound.SimulationArchive(pathtosa + 'sa' + row['runstring'])
sim = sa[0]
P1 = sim.particles[
1].P # Need initial orbital period for TTVsystems, where P1 != 1
try: # Needed for old integrations (random and Naireen) because no snapshot on end
sim = rebound.Simulation(pathtosa +
'../../final_conditions/runs/fc' +
row['runstring'])
except: # New runs (resonant and Ari) have snapshots at collision
sa = rebound.SimulationArchive(pathtosa + 'sa' + row['runstring'])
sim = sa[-1]
row['instability_time'] = sim.t / P1
try:
ssim = rebound.Simulation(pathtossa +
'../../final_conditions/shadowruns/fc' +
row['runstring'])
except:
ssa = rebound.SimulationArchive(pathtossa + 'sa' +
row['runstring'])
ssim = ssa[-1]
row['shadow_instability_time'] = ssim.t / P1
row['Stable'] = row['instability_time'] > 9.99e8
except:
print(pathtosa + 'sa' + row['runstring'])
return row
def get_features(system, dir_SA):
try:
df = pd.read_csv('systems/%s_features.csv' % system)
print('***Loaded predictions for system %s.' % system)
except:
SAs = glob.glob("%s/*_SA_inc.bin" % dir_SA)
# make data frame
print(
'***Couldnt retrieve predictions for system %s, generating from scratch***'
% system)
mf = list(model_features) # copy
mf += [
'name', 'id', 'Stable', 'instability_time', 'Rel_Eerr', 'P1',
'sim_time'
]
df = pd.DataFrame(columns=mf)
for index, dir_sim in enumerate(SAs):
try:
basename = os.path.basename(dir_sim.split('_SA_inc.bin')[0])
#dir_final = dir_sim.split('_SA.bin')[0] + '_final.bin'
P1 = rebound.SimulationArchive(dir_sim)[0].particles[1].P
#sim_time = rebound.SimulationArchive(dir_final)[-1].t
sim_time = rebound.SimulationArchive(dir_sim)[-1].t
features = gen.system(dir_sim, sim_time, P1, index)
features['name'] = basename
features['id'] = basename.split('_')[-1]
df = pd.concat([df, features[mf]])
except:
pass
df.to_csv('systems/%s_features.csv' % system)
return df
def test_asteroid_sim(make_plot: bool = False, verbose: bool=False) -> bool:
"""Test the integration of the asteroids against Horizons"""
# Load the simulation archive for the first 1000 asteroids
n0: int = 0
n1: int = 1000
fname: str = f'../data/asteroids/sim_asteroids_n_{n0:06}_{n1:06}.bin'
sa: rebound.SimulationArchive = rebound.SimulationArchive(fname)
# List of objects to test: Earth and the first 25 asteroids
test_objects: List[str] = [
'Sun', 'Earth',
'Ceres', 'Pallas', 'Juno', 'Vesta', 'Astraea',
'Hebe', 'Iris', 'Flora', 'Metis', 'Hygiea',
'Parthenope', 'Victoria', 'Egeria', 'Irene', 'Eunomia',
'Psyche', 'Thetis', 'Melpomene', 'Fortuna', 'Massalia',
'Lutetia', 'Kalliope', 'Thalia', 'Phocaea']
# Other args to test_integration
sim_name: str = 'planets'
# Test against the asteroid test set
pos_err, ang_err = \
test_integration(sa=sa, test_objects=test_objects,
sim_name=sim_name, test_name='asteroids',
make_plot=make_plot, verbose=verbose)
# Threshold for pass
pos_tol: float = 1.0E-5
ang_tol: float = 2.0
# Test result
isOK: bool = (max(pos_err) < pos_tol) and (max(ang_err) < ang_tol)
msg: str = 'PASS' if isOK else 'FAIL'
print(f'\n***** {msg} *****')
return isOK
def test_sa_whfasthelio_restart(self):
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3,a=1,e=0.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.add(m=1e-3,a=-2,e=1.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.integrator = "whfast"
sim.ri_whfast.coordinates = "democraticheliocentric"
sim.dt = 0.1313
sim.ri_whfast.safe_mode = 0
sim.automateSimulationArchive("test.bin", 10.,deletefile=True)
sim.integrate(42.,exact_finish_time=0)
sim = None
sa = rebound.SimulationArchive("test.bin")
sim = sa[-1]
sim.integrate(80.,exact_finish_time=0)
x1 = sim.particles[1].x
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3,a=1,e=0.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.add(m=1e-3,a=-2,e=1.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.integrator = "whfast"
sim.ri_whfast.coordinates = "democraticheliocentric"
sim.dt = 0.1313
sim.ri_whfast.safe_mode = 0
sim.integrate(80.,exact_finish_time=0)
x0 = sim.particles[1].x
self.assertEqual(x0,x1)
def runSimulation(tmax=40., restart=False, keep_unsynchronized=1, interval=None, safe_mode=True, integrator="ias15",G=1., testparticle=0,simulationarchive_version=2):
if restart:
if keep_unsynchronized==1:
sim = rebound.Simulation("test.bin")
else:
sa = rebound.SimulationArchive("test.bin")
sim = sa.getSimulation(sa.tmax,keep_unsynchronized=0)
else:
sim = rebound.Simulation()
sim.G = G
sim.add(m=1)
sim.add(m=1e-3,a=1,e=0.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.add(m=1e-3,a=-2,e=1.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.integrator = integrator
sim.dt = 0.1313
if simulationarchive_version==1:
sim.simulationarchive_version = 1
if safe_mode==False:
sim.ri_whfast.safe_mode = 1
sim.ri_mercurius.safe_mode = 1
if testparticle>0:
if testparticle==1:
sim.testparticle_type=0
if testparticle==2:
sim.testparticle_type=1
sim.add(m=1e-4,a=1.2,e=0.04,omega=0.21,M=1.41,inc=0.21,Omega=1.1)
sim.N_active = sim.N-1 # one test particle
if interval:
sim.automateSimulationArchive("test.bin", interval, deletefile=True)
sim.integrate(tmax,exact_finish_time=0)
return sim
def test_sa_restart_ias15_walltime(self):
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3, a=1, e=0.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.add(m=1e-3, a=-2, e=1.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.integrator = "ias15"
sim.dt = 0.1313
sim.initSimulationArchive("test.bin", interval_walltime=0.01)
sim.integrate(400., exact_finish_time=0)
sim = None
sa = rebound.SimulationArchive("test.bin")
sim = sa[-1]
self.assertGreater(sim.t, 100.)
sim.integrate(800., exact_finish_time=0)
x1 = sim.particles[1].x
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3, a=1, e=0.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.add(m=1e-3, a=-2, e=1.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.integrator = "ias15"
sim.dt = 0.1313
sim.integrate(800., exact_finish_time=0)
x0 = sim.particles[1].x
self.assertEqual(x0, x1)
tget = 216.123
sim = sa.getSimulation(tget, mode="exact")
self.assertAlmostEqual(sim.t, tget, delta=1e-14)
tget = 116.123
sim = sa.getSimulation(tget, mode="close")
self.assertAlmostEqual(sim.t, tget, delta=sim.dt)
def test_sa_restart(self):
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3, a=1, e=0.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.add(m=1e-3, a=-2, e=1.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.integrator = "whfast"
sim.dt = 0.1313
sim.ri_whfast.safe_mode = 0
sim.initSimulationArchive("test.bin", 10.)
sim.integrate(40., exact_finish_time=0)
sim = None
sa = rebound.SimulationArchive("test.bin")
sim = sa[-1]
sim.integrate(80., exact_finish_time=0)
x1 = sim.particles[1].x
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3, a=1, e=0.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.add(m=1e-3, a=-2, e=1.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.integrator = "whfast"
sim.dt = 0.1313
sim.ri_whfast.safe_mode = 0
sim.integrate(80., exact_finish_time=0)
x0 = sim.particles[1].x
self.assertEqual(x0, x1)
tget = 35.123
sim = sa.getSimulation(tget, mode="exact")
self.assertAlmostEqual(sim.t, tget, delta=1e-14)
tget = 25.123
sim = sa.getSimulation(tget, mode="close")
self.assertAlmostEqual(sim.t, tget, delta=sim.dt)
def get_longitude_of_periastron_stats(row):
print('Processing ', row.simarchive_filename)
sa = rebound.SimulationArchive(row.simarchive_filename,
process_warnings=False)
oa = chaincalc.OrbitArray(sa)
tlook = 1e4 * keplertime #probably 1e2-1e3 times the output cadence
lko = oa.tail_orbitarray(tlook)
mean_dlongitude_of_periastron = []
std_dlongitude_of_periastron = []
for n in range(lko['pomega'].shape[0] - 1):
longstats = [{
'mean':
chaincalc.wrap2pi(lko['pomega'][n + 1, :] - lko['pomega'][n, :] +
a).mean() - a,
'std':
chaincalc.wrap2pi(lko['pomega'][n + 1, :] - lko['pomega'][n, :] +
a).std()
} for a in [0.0, np.pi]]
if longstats[0]['std'] < longstats[1]['std']:
mean_dlongitude_of_periastron.append(longstats[0]['mean'])
std_dlongitude_of_periastron.append(longstats[0]['std'])
else:
mean_dlongitude_of_periastron.append(longstats[1]['mean'])
std_dlongitude_of_periastron.append(longstats[1]['std'])
#print(mean_dlongitude_of_periastron)
#print(std_dlongitude_of_periastron)
#if max(mean_dlongitude_of_periastron) > -np.pi/2:
# print(row.hash)
# print(mean_dlongitude_of_periastron)
return {
'delta_longitude_of_periastron': mean_dlongitude_of_periastron,
'std_delta_longitude_of_periastron': std_dlongitude_of_periastron
}
def test_sa_restart_ias15(self):
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3,a=1,e=0.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.add(m=1e-3,a=-2,e=1.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.integrator = "ias15"
sim.dt = 0.1313
sim.automateSimulationArchive("test.bin", 10.,deletefile=True)
sim.integrate(40.,exact_finish_time=0)
sim = None
sa = rebound.SimulationArchive("test.bin")
sim = sa[-1]
sim.integrate(80.,exact_finish_time=0)
x1 = sim.particles[1].x
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3,a=1,e=0.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.add(m=1e-3,a=-2,e=1.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.integrator = "ias15"
sim.dt = 0.1313
sim.integrate(80.,exact_finish_time=0)
x0 = sim.particles[1].x
self.assertEqual(x0,x1)
tget = 27.123
sim = sa.getSimulation(tget,mode="exact");
self.assertAlmostEqual(sim.t,tget,delta=1e-14)
tget = 25.123
sim = sa.getSimulation(tget,mode="close");
self.assertAlmostEqual(sim.t,tget,delta=sim.dt)
def bin_find(loc):
##Ensure we are in the com frame of the simulation.
t, name = loc
sat = rebound.SimulationArchive(name)
sim = sat.getSimulation(t)
sim.move_to_com()
##Integrate forward for a small time to ensure that the accelerations
##are in sync with the rest of the simulation (this is important for
##calculating tidal forces...
sim.integrate(sim.t + sim.t * 1.0e-14)
ps = sim.particles
##mass of of primary
m0 = ps[0].m
bin_indics = []
for i1, i2 in combinations(
range(1, sim.N),
2): # get all pairs of indices/start at 1 to exclude SMBH
com_d, a_bin, e_bin, p1_com, p2_com, d2, inc, ft = bin_props(
ps[i1], ps[i2])
m1, m2 = ps[i1].m, ps[i2].m
##Hill sphere condition.
inside_hill = (a_bin < ((m1 + m2) / m0)**(1. / 3.) * com_d)
tidal_2 = (m1 * m2 / d2 > ft)
##If the kinetic energy is less than the potential energy
if ((a_bin > 0) and (inside_hill) and (tidal_2)):
#if ((a_bin>0) and (inside_hill)):
bin_indics.append([
sim.t, i1, i2, d2**0.5, a_bin,
a_bin / (((m1 + m2) / m0)**(1. / 3.) * com_d), e_bin
])
return np.array(bin_indics)
def test_sa_mercurius_restart_safemode(self):
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3,a=1,e=0.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.add(m=1e-3,a=-2,e=1.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.integrator = "mercurius"
sim.dt = 0.1313
sim.ri_mercurius.safe_mode = 1
sim.automateSimulationArchive("test.bin", 10.,deletefile=True)
sim.integrate(40.,exact_finish_time=0)
sim = None
sa = rebound.SimulationArchive("test.bin")
sim = sa[-1]
sim.integrate(80.,exact_finish_time=0)
x1 = sim.particles[1].x
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3,a=1,e=0.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.add(m=1e-3,a=-2,e=1.1,omega=0.1,M=0.1,inc=0.1,Omega=0.1)
sim.integrator = "mercurius"
sim.dt = 0.1313
sim.ri_mercurius.safe_mode = 1
sim.integrate(80.,exact_finish_time=0)
x0 = sim.particles[1].x
self.assertEqual(x0,x1)
def test_sa_whfasthelio_restart(self):
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3, a=1, e=0.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.add(m=1e-3, a=-2, e=1.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.integrator = "whfasthelio"
sim.dt = 0.1313
sim.ri_whfasthelio.safe_mode = 0
sim.initSimulationArchive("test.bin", 10.)
sim.integrate(40., exact_finish_time=0)
sim = None
sa = rebound.SimulationArchive("test.bin")
sim = sa[-1]
sim.integrate(80., exact_finish_time=0)
x1 = sim.particles[1].x
sim = rebound.Simulation()
sim.add(m=1)
sim.add(m=1e-3, a=1, e=0.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.add(m=1e-3, a=-2, e=1.1, omega=0.1, M=0.1, inc=0.1, Omega=0.1)
sim.integrator = "whfasthelio"
sim.dt = 0.1313
sim.ri_whfasthelio.safe_mode = 0
sim.integrate(80., exact_finish_time=0)
x0 = sim.particles[1].x
self.assertEqual(x0, x1)
def __init__(self,
file,
time,
t_kepler=25200,
dt=252,
r_earth=6.371e6,
m_earth=5.9722e24,
m_lunar=7.345806e+22,
J_EM=3.5e34,
rho=5.5,
rho0=3.3):
self.file = file
self.time = time
self.sim = rebound.SimulationArchive(file)
self.sim_time = self.sim.getSimulation(time)
self.t_kepler = t_kepler
self.dt = dt
# Planetary constants
self.r_earth = r_earth
self.m_earth = m_earth
self.m_lunar = m_lunar
self.J_EM = J_EM
self.rho = rho
self.rho0 = rho0
self.roche = 2.456 * (self.rho / self.rho0)**(1 / 3) * self.r_earth
def get_simarchive_integration_results(sa,coordinates='jacobi'):
"""
Read a simulation archive and store orbital elements
as arrays in a dictionary.
Arguments
---------
sa : rebound.SimulationArchive or str
The simulation archive to read or the file name of the simulation
archive file. Can also be a reboundx simulation archive.
coordinates : str
The coordinate system to use for calculating orbital elements.
This can be:
- 'jacobi' : Use Jacobi coordinates (including Jacobi masses)
- 'heliocentric' : Use canonical heliocentric elements. The
canonical cooridantes are heliocentric distance vectors.
The conjugate momenta are barycentric momenta.
Returns
-------
sim_results : dict
Dictionary containing time and orbital elements at each
snapshot of the simulation archive.
"""
if type(sa) == str:
sa = rb.SimulationArchive(sa)
if type(sa) == rb.simulationarchive.SimulationArchive:
return _get_rebound_simarchive_integration_results(sa,coordinates)
elif type(sa) == rbx.simulationarchive.SimulationArchive:
return _get_reboundx_simarchive_integration_results(sa,coordinates)
raise TypeError("{} is not a rebound or reboundx simulation archive!".format(sa))
def save_evolution():
fs = np.array(glob.glob('pickles/archived0*'))
for i in range(fs.size):
sa = rebound.SimulationArchive(fs[i])
times, smas, eccs, incs = np.zeros(0), np.zeros((0, 2)), np.zeros(
(0, 2)), np.zeros((0, 2))
ind = 0
for j, sim in enumerate(sa):
times = np.append(times, sim.t)
smas = np.insert(smas,
ind,
np.array([sim.particles[1].a,
sim.particles[2].a]),
axis=0)
eccs = np.insert(eccs,
ind,
np.array([sim.particles[1].e,
sim.particles[2].e]),
axis=0)
incs = np.insert(incs,
ind,
np.rad2deg(
[sim.particles[1].inc, sim.particles[2].inc]),
axis=0)
ind += 1
# Save
fname = fs[i].replace('archived', 'evolution').replace('bin', 'dat')
np.savetxt(fname,
np.array([
times, smas[:, 0], smas[:, 1], eccs[:, 0], eccs[:, 1],
incs[:, 0], incs[:, 1]
]).T,
delimiter='\t',
fmt='%.6f')
def get_k(row):
sim = rebound.SimulationArchive(sim_names + "_sa_%d.bin" % (row[0]))[0]
# print(sim)
p2 = sim.particles[2]
row['h'] = p2.e * np.sin(p2.pomega)
row['k'] = p2.e * np.cos(p2.pomega)
avars = Andoyer.from_Simulation(sim,
a10=sim.particles[1].a,
j=5,
k=1,
i1=1,
i2=2,
average=False)
row['Z12'] = avars.Z
row['Zcom12'] = avars.Zcom
avars = Andoyer.from_Simulation(sim,
a10=sim.particles[1].a,
j=4,
k=1,
i1=2,
i2=3,
average=False)
row['Z23'] = avars.Z
row['Zcom23'] = avars.Zcom
row['e1'] = sim.particles[1].e
row['e2'] = sim.particles[2].e
row['e3'] = sim.particles[3].e
row['m1'] = sim.particles[1].m
row['m2'] = sim.particles[2].m
row['m3'] = sim.particles[3].m
return row
def pred(nsim):
sim = rebound.SimulationArchive(sim_names + "_sa_%d.bin" % nsim)[0]
sim.move_to_com()
sim.integrator = "whfast"
sim.dt = 0.07 * sim.particles[1].P
# prob=model.predict(sim)
prob = model.predict(sim, copy=False)
return prob
def doublecheck(row):
sa = rebound.SimulationArchive(row.simarchive_filename, process_warnings=False)
oa = chaincalc.OrbitArray(sa)
tlook = 1e4*keplertime #probably 1e2-1e3 times the output cadence
lko = oa.tail_orbitarray(tlook)
finding = oa.is_this_a_constant_ratio_chain(oa.compute_tight_angles(lko,pratiocut=0.02))
print(row.hash, finding)
return finding[0]
def make_archive(fname_archive: str,
sim_epoch: rebound.Simulation,
object_names: List[str],
epoch: datetime,
dt0: datetime,
dt1: datetime,
time_step: int,
save_step: int = 1,
save_elements: bool = False,
progbar: bool = False) -> rebound.SimulationArchive:
"""
Load a rebound archive if available; otherwise generate it and save it to disk.
INPUTS:
fname_archive: the file name to save the archive to
sim_epoch: rebound simulation object as of the epoch time; to be integrated in both directions
object_names: the user names of all the objects in the simulation
epoch: a datetime corresponding to sim_epoch
dt0: the earliest datetime to simulate back to
dt1: the latest datetime to simulate forward to
time_step: the time step in days for the simulation
save_step: the interval for saving snapshots to the simulation archive
progbar: flag - whether to display a progress bar
"""
try:
# First try to load the named archive
sa = rebound.SimulationArchive(filename=fname_archive)
except:
# If the archive is not on disk, save it to disk
print(
f'Generating archive {fname_archive}\n'
f'from {dt0} to {dt1}, time_step={time_step}, save_step={save_step}...'
)
make_archive_impl(fname_archive=fname_archive,
sim_epoch=sim_epoch,
object_names=object_names,
epoch=epoch,
dt0=dt0,
dt1=dt1,
time_step=time_step,
save_step=save_step,
save_elements=save_elements,
progbar=progbar)
# Load the new archive into memory
sa = rebound.SimulationArchive(filename=fname_archive)
return sa
def final_time(filename):
try:
sa = rebound.SimulationArchive(trappistdistpath + filename)
sim = sa[0]
P0 = sim.particles[1].P
return sa[-1].t / P0
except Exception as e:
print(e, filename)
return np.nan
def lag_simple_plot(sa_name, i1, i2, extras=[], name='', cols=['r', 'g', 'k'], idx_min=0, idx_max=None, lim=0.1, ms=2, interval=1):
'''
Project particles to co-rotating frame ; this is function is only accurate if the orbital plane is the xy plane.
'''
planes = [['x', 'y'], ['x', 'z'], ['y','z']]
#fig.patch.set_visible(False)
sa = rebound.SimulationArchive(sa_name)
if not idx_max:
idx_max=len(sa)
m0=sa[0].particles[0].m
for ii in range(idx_min, idx_max, interval):
fig,ax=plt.subplots(nrows=1, ncols=3, figsize=(10*len(planes),10))
for kk, plane in enumerate(planes):
ax[kk].set_xlim(-lim, lim)
ax[kk].set_ylim(-lim, lim)
ax[kk].set_xlabel(planes[kk][0])
ax[kk].set_ylabel(planes[kk][1])
sim=sa[ii]
p1,p2=sim.particles[i1],sim.particles[i2]
m1,m2=p1.m,p2.m
com_d, a_bin, e_bin, p1_com, p2_com, d2, inc, ft = bin_props(p1,p2)
theta=-np.arctan2(p2_com.y,p2_com.x)
ct,st=np.cos(theta), np.sin(theta)
p1_xyz=np.dot([[ct, -st,0], [st, ct,0], [0,0,1]], p1_com.xyz)
p2_xyz=np.dot([[ct, -st,0], [st, ct,0], [0,0,1]], p2_com.xyz)
p1_xyz={'x':p1_xyz[0], 'y':p1_xyz[1], 'z':p1_xyz[2]}
p2_xyz={'x':p2_xyz[0], 'y':p2_xyz[1], 'z':p2_xyz[2]}
p1_pos=p1_xyz[plane[0]], p1_xyz[plane[1]]
p2_pos=p2_xyz[plane[0]], p2_xyz[plane[1]]
com=get_com([p1, p2])
ax[kk].plot(p1_pos[0], p1_pos[1], 'o', markersize=ms, color=cols[0])
ax[kk].plot(p2_pos[0], p2_pos[1], 'o', markersize=ms, color=cols[1])
for jj, extra in enumerate(extras):
xyz=np.dot([[ct, -st,0], [st, ct,0], [0,0,1]], (sim.particles[extra]-com).xyz)
xyz={'x':xyz[0], 'y':xyz[1], 'z':xyz[2]}
ax[kk].plot(xyz[plane[0]], xyz[plane[1]], 'o', markersize=ms, color=cols[(2+jj)%len(cols)])
alpha=m2/(m1+m2)
l1=(d2**0.5*(1.-(alpha/3.)**(1./3.)))
l2=(d2**0.5*(1.+(alpha/3.)**(1./3.)))
l3=(-d2**0.5*(1.+5./12.*alpha))
ax[0].plot(l1, 0, 'gv')
ax[0].plot(l2, 0, 'gv')
ax[0].plot(l3, 0, 'gv')
ax[0].plot(d2**0.5*np.cos(np.pi/3.), d2**0.5*np.sin(np.pi/3.), 'gv')
ax[0].plot(d2**0.5*np.cos(np.pi/3.), -d2**0.5*np.sin(np.pi/3.), 'gv')
fig.savefig(sa_name.replace('.bin', '')+name+'_com_{0:03d}.png'.format(ii), bbox_inches='tight', pad_inches=0)
plt.clf()
def massratios(row):
try:
sa = rebound.SimulationArchive(pathtosa+'sa'+row['runstring'])
sim = sa[0]
row['m1'] = sim.particles[1].m/sim.particles[0].m
row['m2'] = sim.particles[2].m/sim.particles[0].m
row['m3'] = sim.particles[3].m/sim.particles[0].m
except:
print(pathtosa+'sa'+row['runstring'])
return row
def get_times(row, args):
systemdir = args
fcpath = systemdir + "/simulation_archives/sa"
try:
sa = rebound.SimulationArchive(fcpath + row["runstring"])
row['t'] = sa[-1].t
del sa
except Exception as e:
row['t'] = np.nan
return row
def training_data(row, safolder, runfunc, args):
try:
sa = rebound.SimulationArchive(safolder + 'sa' + row['runstring'])
sim = sa[0]
except:
print("traininst_data_functions.py Error reading " + safolder + 'sa' +
row['runstring'])
return None
return runfunc(sim, args)
def doTest2(self):
sim1 = runSimulation(40., restart=False, interval=10., **params)
if params['keep_unsynchronized']==1:
sim2 = rebound.Simulation("test.bin")
else:
sa = rebound.SimulationArchive("test.bin")
sim2 = sa.getSimulation(sa.tmax,keep_unsynchronized=0)
compareSim(self,sim1,sim2)
sim1.integrate(sim1.t+12.)
sim2.integrate(sim2.t+12.)
compareSim(self,sim1,sim2)
def system(dir, sim_time, P1, index, max_sim_P1=1e9):
SA = rebound.SimulationArchive(dir)
E0 = SA[0].calculate_energy()
features = generate_features(SA[0], index)
features['Stable'] = 1 if np.isclose(sim_time/P1, max_sim_P1) else 0
features['sim_time'] = sim_time
#print(sim_time/SA[0].particles[1].P, sim_time/P1, features['Stable'].values, dir) #disagreement.. why?
features['instability_time'] = sim_time
features['Rel_Eerr'] = abs((SA[-1].calculate_energy()-E0)/E0)
features['P1'] = P1
return features
def plot_archival_evolution(index, nplanets=2):
sa = rebound.SimulationArchive('pickles/archived%.6d.bin' % index)
times = np.zeros(0)
eccs = np.zeros((0, nplanets))
for i, sim in enumerate(sa):
times = np.append(times, sim.t)
eccs = np.insert(eccs,
i,
np.array([sim.particles[1].e, sim.particles[2].e]),
axis=0)
plt.plot(times, eccs, '-')
plt.show()
