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 process_row(sim: rebound.Simulation, fname: str, t: float, row: int):
# Integrate to the current time step with an exact finish time
sim.integrate(t, exact_finish_time=1)
# Serialize the position and velocity
sim.serialize_particle_data(xyz=q[row])
sim.serialize_particle_data(vxvyvz=v[row])
# Save a snapshot on multiples of save_step or the first / last row
if (i % save_step == 0) or (row in (0, M - 1)):
sim.simulationarchive_snapshot(filename=fname)
# Save the orbital elements if applicable
if save_elements:
# Compute the orbital elements vs. the sun as primary
primary = sim.particles['Sun']
orbits = sim.calculate_orbits(primary=primary, jacobi_masses=False)
# Save the elements on this date as an Nx6 array
# This approach only traverses the (slow) Python list orbits one time
elt_array = np.array([[orb.a, orb.e, orb.inc, orb.Omega, orb.omega, orb.f] \
for orb in orbits])
# Save the elements into the current row of the named orbital elements arrays
# The LHS row mask starts at 1 b/c orbital elements are not computed for the first object (Sun)
orb_a[row, 1:] = elt_array[:, 0]
orb_e[row, 1:] = elt_array[:, 1]
orb_inc[row, 1:] = elt_array[:, 2]
orb_Omega[row, 1:] = elt_array[:, 3]
orb_omega[row, 1:] = elt_array[:, 4]
orb_f[row, 1:] = elt_array[:, 5]
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 main():
args = sys.argv[1:]
from argparse import ArgumentParser
psr = ArgumentParser(
description="Convert Mercury6 output into REBOUND archive.")
psr.add_argument('input_dir', type=str, help='Simulation directory')
psr.add_argument('-o',
'--output_dir',
type=str,
help='Output directory.',
default='./rebsims/')
args = psr.parse_args(args)
indir = args.input_dir
outdir = args.output_dir
sim_name = os.path.basename(os.path.normpath(indir))
# Get surviving planets from element.out
try:
el_out = open(os.path.join(indir, 'element.out'), 'r')
pl_remain = []
plan_re = re.compile(r'^\s*(plan\d+)')
for l in el_out:
match = re.search(plan_re, l)
if match:
pl_remain.append(match.group(1))
except:
print("Could not find element.out. Run element6.for.")
raise
# Create simulation
sim = Simulation()
sim.units = ('day', 'AU', 'Msun')
# Add central star
sim.add(m=1)
# Add planets
for pl in pl_remain:
# Read last line of appropriate file
pl_file = os.path.join(indir, pl + '.aei')
with open(pl_file, 'rb') as f:
f.seek(-2, os.SEEK_END)
while f.read(1) != b'\n':
f.seek(-2, os.SEEK_CUR)
final = f.readline().decode()
# Assume Cartesian positions and velocities.
t, x, y, z, u, v, w, m = [float(s) for s in final.split()]
sim.add(m=m, x=x, y=y, z=z, vx=u, vy=v, vz=w)
outfile = os.path.join(outdir, sim_name + '.bin')
sim.save(outfile)
return
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 reorder_particles(sim: Simulation, ed: ExtraData) -> None:
"""
probably not needed anymore
"""
exit()
particles_by_hash: Dict[int, Particle] = {}
hashes = []
suns = []
gas_giants = []
embryos = []
planetesimals = []
original_N = sim.N
p: Particle
for p in sim.particles:
hash_value = p.hash.value
# save a copy of the particles
particles_by_hash[hash_value] = p.copy()
type = ed.pd(p).type
if type == "sun":
suns.append(hash_value)
# keep sun in the simulation to avoid empty particles list
continue
elif type == "gas giant":
gas_giants.append(hash_value)
elif type == "embryo":
embryos.append(hash_value)
elif type == "planetesimal":
planetesimals.append(hash_value)
else:
raise ValueError(f"unknown type: {type}")
hashes.append(p.hash)
for hash in hashes:
sim.remove(hash=hash)
ordered_particles = gas_giants + embryos + planetesimals
assert len(suns) + len(ordered_particles) == original_N
particle: Particle
for h in ordered_particles:
particle = particles_by_hash[h]
sim.add(particle)
# print(list(sim.particles))
# exit()
sim.N_active = len(suns) + len(ordered_particles) - len(planetesimals)
# mark objects > N_active as testparticle_type = 1
# this means they become semi-active bodies
# TODO: double-check meaning
sim.testparticle_type = 1
assert sim.N == original_N
def main():
sim = Simulation()
rebx = Extras(sim)
rebx.register_param("water", "REBX_TYPE_DOUBLE")
sim.units = ('yr', 'AU', 'Msun')
sim.add(m=1.0) # Sun
for planet in data:
part = Particle(a=planet[0],
e=planet[1],
inc=radians(planet[2]),
omega=radians(planet[3]),
Omega=radians(planet[4]),
M=radians(planet[5]),
m=planet[6],
simulation=sim)
sim.add(part)
max_n = sim.N
sim.particles[1].params["water"] = 0.3
print("start")
tmax = 1e5
savesteps = 1000
times = np.linspace(0., tmax, savesteps)
sim.move_to_com()
sim.automateSimulationArchive("sa.bin", interval=savesteps)
for i, t in enumerate(times):
sim.integrate(t)
print(f"done: {i / 10}% ({sim.N}")
with open("sa.meta.json", "w") as f:
meta = {"tmax": tmax, "savesteps": savesteps, "max_n": max_n}
json.dump(meta, f)
def make_traj_from_sim(sim: rebound.Simulation, n_years: int, sample_freq: int,
debug_print: bool):
"""
Make an array of training data points for the general 3 body problem from a simulation.
Creates one trajectory with the same initial configuration.
INPUTS:
sim: rebound Simulation object
n_years: the number of years of data to simulate, e.g. 2
sample_freq: number of sample points per year, e.g. 365
RETURNS:
input_dict: a dictionary with the input fields t, q0, v0
output_dict: a dictionary with the output fields q, v, a, q0_rec, v0_rec, T, H, H, L
"""
# Number of samples including both start and end
N = n_years * sample_freq + 1
# Unpack the masses and wrap into array m
ps = sim.particles
p0, p1, p2 = ps
m0, m1, m2 = p0.m, p1.m, p2.m
m = np.array([m0, m1, m2], dtype=np.float32)
# Array shapes
num_particles = 3
space_dims = 3
traj_shape = (N, num_particles, space_dims)
# the number of orbital elements is the number of particles minus 1; no elements for primary
elt_shape = (N, num_particles - 1)
mom_shape = (N, 3)
# Initialize cartesian entries to zero vectors
q = np.zeros(traj_shape, dtype=np.float32)
v = np.zeros(traj_shape, dtype=np.float32)
a = np.zeros(traj_shape, dtype=np.float32)
# Initialize orbital elements over time to zero vectors
orb_a = np.zeros(elt_shape, dtype=np.float32)
orb_e = np.zeros(elt_shape, dtype=np.float32)
orb_inc = np.zeros(elt_shape, dtype=np.float32)
orb_Omega = np.zeros(elt_shape, dtype=np.float32)
orb_omega = np.zeros(elt_shape, dtype=np.float32)
orb_f = np.zeros(elt_shape, dtype=np.float32)
# Initialize placeholders for kinetic and potential energy
T = np.zeros(N, dtype=np.float32)
U = np.zeros(N, dtype=np.float32)
# Initialize momentum and angular momentum
P = np.zeros(mom_shape, dtype=np.float32)
L = np.zeros(mom_shape, dtype=np.float32)
# The coefficients for gravitational potential energy on particle pairs
k_01 = sim.G * m0 * m1
k_02 = sim.G * m0 * m2
k_12 = sim.G * m1 * m2
# Set the times for snapshots
ts = np.linspace(0.0, n_years, N, dtype=np.float32)
# The particles for the 3 bodies
p0 = sim.particles[0]
p1 = sim.particles[1]
p2 = sim.particles[2]
# Simulate the orbits
# Start by integrating backward, then forward for a small step
# This allows rebound to correctly initialize the acceleration
sim.integrate(-1E-6, exact_finish_time=1)
sim.integrate(1E-6, exact_finish_time=1)
for i, t in enumerate(ts):
# Integrate to the current time step with an exact finish time
sim.integrate(t, exact_finish_time=1)
# Save the position
q[i] = [[p0.x, p0.y, p0.z], [p1.x, p1.y, p1.z], [p2.x, p2.y, p2.z]]
# Save the velocity
v[i] = [[p0.vx, p0.vy, p0.vz], [p1.vx, p1.vy, p1.vz],
[p2.vx, p2.vy, p2.vz]]
# Save the acceleration
a[i] = [[p0.ax, p0.ay, p0.az], [p1.ax, p1.ay, p1.az],
[p2.ax, p2.ay, p2.az]]
# Calculate the two orbits
orb1 = p1.calculate_orbit()
orb2 = p2.calculate_orbit()
# Save the orbital elements
orb_a[i] = [orb1.a, orb2.a]
orb_e[i] = [orb1.e, orb2.e]
orb_inc[i] = [orb1.inc, orb2.inc]
orb_Omega[i] = [orb1.Omega, orb2.Omega]
orb_omega[i] = [orb1.omega, orb2.omega]
orb_f[i] = [orb1.f, orb2.f]
# Extract the 3 positions
q0 = q[i, 0]
q1 = q[i, 1]
q2 = q[i, 2]
# Extract the 3 velocities
v0 = v[i, 0]
v1 = v[i, 1]
v2 = v[i, 2]
# Kinetic energy
T0 = 0.5 * m0 * np.sum(np.square(v0))
T1 = 0.5 * m1 * np.sum(np.square(v1))
T2 = 0.5 * m2 * np.sum(np.square(v2))
T[i] = T0 + T1 + T2
# Potential energy; 3 pairs of interacting particles (3 choose 2)
r_01 = np.linalg.norm(q1 - q0)
r_02 = np.linalg.norm(q2 - q0)
r_12 = np.linalg.norm(q2 - q1)
U[i] = -(k_01 / r_01 + k_02 / r_02 + k_12 / r_12)
# The momentum vector; should be zero in the COM frame
mv0 = m0 * v0
mv1 = m1 * v1
mv2 = m2 * v2
P[i] = mv0 + mv1 + mv2
# The angular momentum vector; should be constant by conservation of angular momentum
L[i] = np.cross(q0, mv0) + np.cross(q1, mv1) + np.cross(q2, mv2)
# The total energy is the sum of kinetic and potential; should be constant by conservation of energy
H = T + U
# The initial position and velocity
q0 = q[0]
v0 = v[0]
# Assemble the input dict
inputs = {'t': ts, 'q0': q0, 'v0': v0, 'm': m}
# Assemble the output dict
outputs = {
# the trajectory
'q': q,
'v': v,
'a': a,
# the orbital elements over time
'orb_a': orb_a,
'orb_e': orb_e,
'orb_inc': orb_inc,
'orb_Omega': orb_Omega,
'orb_omega': orb_omega,
'orb_f': orb_f,
# the initial conditions, which should be recovered
'q0_rec': q0,
'v0_rec': v0,
# the energy, momentum and angular momentum, which should be conserved
'T': T,
'U': U,
'H': H,
'P': P,
'L': L
}
# Return the dicts
return (inputs, outputs)
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])
def make_archive_impl(fname_archive: str, sim_epoch: rebound.Simulation,
object_names: List[str], epoch: datetime, dt0: datetime,
dt1: datetime, time_step: int, save_step: int,
save_elements: bool, progbar: bool) -> None:
"""
Create a rebound simulation archive 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
save_elements: flag indiciting whether to save orbital elements
progbar: flag - whether to display a progress bar
"""
# Convert epoch, start and end times relative to a base date of the simulation start
# This way, time is indexed from t0=0 to t1 = (dt1-dt0)
epoch_t: float = datetime_to_mjd(epoch, dt0)
t0: float = datetime_to_mjd(dt0, dt0)
t1: float = datetime_to_mjd(dt1, dt0)
# 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_t
sim_back.t = epoch_t
# Set the times for snapshots in both directions;
ts: np.array = np.arange(t0, t1 + time_step, time_step)
idx: int = np.searchsorted(ts, epoch_t, side='left')
ts_fwd: np.array = ts[idx:]
ts_back: np.array = ts[:idx][::-1]
# The epochs corresponding to the times in ts
epochs: List[datetime] = [dt0 + timedelta(t) for t in ts]
# File names for forward and backward integrations
fname_fwd: str = fname_archive.replace('.bin', '_fwd.bin')
fname_back: str = fname_archive.replace('.bin', '_back.bin')
# 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: int = sim_epoch.N
# Initialize arrays for the position and velocity
shape_qv: Tuple[int] = (M, N, 3)
q: np.array = np.zeros(shape_qv, dtype=np.float64)
v: np.array = np.zeros(shape_qv, dtype=np.float64)
# Initialize arrays for orbital elements if applicable
if save_elements:
# Arrays for a, e, inc, Omega, omega, f
shape_elt: Tuple[int] = (M, N)
orb_a: np.array = np.zeros(shape_elt, dtype=np.float64)
orb_e: np.array = np.zeros(shape_elt, dtype=np.float64)
orb_inc: np.array = np.zeros(shape_elt, dtype=np.float64)
orb_Omega: np.array = np.zeros(shape_elt, dtype=np.float64)
orb_omega: np.array = np.zeros(shape_elt, dtype=np.float64)
orb_f: np.array = np.zeros(shape_elt, dtype=np.float64)
# Wrap these into a named tuple
elts: OrbitalElement = \
OrbitalElement(a=orb_a, e=orb_e, inc=orb_inc,
Omega=orb_Omega, omega=orb_omega, f=orb_f)
else:
elts = OrbitalElement()
# Subfunction: process one row of the loop
def process_row(sim: rebound.Simulation, fname: str, t: float, row: int):
# Integrate to the current time step with an exact finish time
sim.integrate(t, exact_finish_time=1)
# Serialize the position and velocity
sim.serialize_particle_data(xyz=q[row])
sim.serialize_particle_data(vxvyvz=v[row])
# Save a snapshot on multiples of save_step or the first / last row
if (i % save_step == 0) or (row in (0, M - 1)):
sim.simulationarchive_snapshot(filename=fname)
# Save the orbital elements if applicable
if save_elements:
# Compute the orbital elements vs. the sun as primary
primary = sim.particles['Sun']
orbits = sim.calculate_orbits(primary=primary, jacobi_masses=False)
# Save the elements on this date as an Nx6 array
# This approach only traverses the (slow) Python list orbits one time
elt_array = np.array([[orb.a, orb.e, orb.inc, orb.Omega, orb.omega, orb.f] \
for orb in orbits])
# Save the elements into the current row of the named orbital elements arrays
# The LHS row mask starts at 1 b/c orbital elements are not computed for the first object (Sun)
orb_a[row, 1:] = elt_array[:, 0]
orb_e[row, 1:] = elt_array[:, 1]
orb_inc[row, 1:] = elt_array[:, 2]
orb_Omega[row, 1:] = elt_array[:, 3]
orb_omega[row, 1:] = elt_array[:, 4]
orb_f[row, 1:] = elt_array[:, 5]
# Integrate the simulation forward in time
idx_fwd: List[Tuple[int, float]] = list(enumerate(ts_fwd))
if progbar:
idx_fwd = tqdm_console(idx_fwd)
for i, t in idx_fwd:
# Row index for position data
row: int = M_back + i
# Process this row
process_row(sim=sim_fwd, fname=fname_fwd, t=t, row=row)
# Integrate the simulation backward in time
idx_back: List[Tuple[int, float]] = list(enumerate(ts_back))
if progbar:
idx_back = tqdm_console(idx_back)
for i, t in idx_back:
# Row index for position data
row: int = M_back - (i + 1)
# Process this row
process_row(sim=sim_back, fname=fname_back, t=t, row=row)
# Load the archives with the forward and backward snapshots
sa_fwd: rebound.SimulationArchive = rebound.SimulationArchive(fname_fwd)
sa_back: rebound.SimulationArchive = rebound.SimulationArchive(fname_back)
# Filename for numpy arrays of position and velocity
fname_np: str = fname_archive.replace('.bin', '.npz')
# Save the epochs as a numpy array
epochs_np: np.array = np.array(epochs)
# Save the object names as a numpy array of strings
object_names_np: np.array = np.array(object_names)
# Save the object name hashes
hashes: np.array = np.zeros(N, dtype=np.uint32)
sim_epoch.serialize_particle_data(hash=hashes)
# Combine the forward and backward archives in forward order from t0 to t1
sims = chain(reversed(sa_back), sa_fwd)
# Process each simulation snapshot in turn
for i, sim in enumerate(sims):
# Save a snapshot on multiples of save_step
sim.simulationarchive_snapshot(fname_archive)
# Save the numpy arrays with the object hashes, position and velocity
np.savez(fname_np,
q=q,
v=v,
elts=elts,
ts=ts,
epochs_np=epochs_np,
hashes=hashes,
object_names_np=object_names_np)
# Delete the forward and backward archives
os.remove(fname_fwd)
os.remove(fname_back)
def add_particles_from_conditions_file(sim: Simulation,
ed: ExtraData,
input_file: str,
testrun=False) -> Tuple[int, int]:
initcon = Path(input_file).read_text()
num_embryos = int(
re.search(r"Generated (\d+) minor bodies", initcon,
re.MULTILINE).group(1))
num_planetesimals = int(
re.search(r"Generated (\d+) small bodies", initcon,
re.MULTILINE).group(1))
sim.N_active = num_embryos + 3
i = 1
for line in initcon.split("\n"):
if line.startswith("#") or line.startswith(
"ERROR") or line == "\n" or not line:
continue
columns = list(map(float, line.split()))
hash = unique_hash(ed)
if len(columns) > 7:
cmf, mmf, wmf = columns[7:]
total_fractions = cmf + mmf + wmf
if total_fractions != 1:
diff = 1 - total_fractions
print(f"fractions don't add up by {diff}")
print("adding rest to mmf")
mmf += diff
assert cmf + mmf + wmf - 1 <= 1e-10
if i > num_embryos + 3:
object_type = "planetesimal"
else:
object_type = "embryo"
else:
wmf = mmf = 0
cmf = 1
if columns[1] == 0:
object_type = "sun"
else:
object_type = "gas giant"
ed.pdata[hash.value] = ParticleData(water_mass_fraction=wmf,
core_mass_fraction=cmf,
type=object_type,
total_mass=columns[0])
if columns[1] == 0: # that should not be needed, but nevertheless is
part = Particle(m=columns[0],
hash=hash,
r=solar_radius / astronomical_unit)
else:
inc = columns[3]
if testrun:
# force particles to collide early when running tests
inc /= 1000
part = Particle(
m=columns[0],
a=columns[1],
e=columns[2],
inc=inc,
omega=columns[4],
Omega=columns[5],
M=columns[6],
simulation=sim,
hash=hash,
r=PlanetaryRadius(columns[0], wmf, cmf).total_radius /
astronomical_unit)
sim.add(part)
i += 1
assert sim.N == num_planetesimals + num_embryos + 3
return num_planetesimals, num_embryos
def main(fn: Path, testrun=False):
global abort
start = time.perf_counter()
if not fn.with_suffix(".bin").exists():
if not testrun:
with open(fn.with_suffix(".yaml")) as f:
parameters = Parameters(**yaml.safe_load(f))
else:
parameters = Parameters(
massloss_method="rbf",
initcon_file="initcon/conditions_many.input")
# set up a fresh simulation
sim = Simulation()
sim.units = ('yr', 'AU', 'kg')
# sim.boundary = "open"
# boxsize = 100
# sim.configure_box(boxsize)
sim.integrator = "mercurius"
sim.dt = 1e-2
sim.ri_ias15.min_dt = 0.0001 / 365
if not parameters.no_merging:
sim.collision = "direct"
sim.ri_mercurius.hillfac = 3.
sim.testparticle_type = 1
tmax = 200 * mega
num_savesteps = 20000
if testrun:
tmax /= 200000
num_savesteps /= 1000
per_savestep = tmax / num_savesteps
extradata = ExtraData()
# times = np.linspace(0., tmax, savesteps)
extradata.meta.tmax = tmax
extradata.meta.per_savestep = per_savestep
extradata.meta.num_savesteps = num_savesteps
extradata.meta.git_hash = git_hash()
extradata.meta.rebound_hash = rebound.__githash__
extradata.meta.massloss_method = parameters.massloss_method
extradata.meta.initcon_file = parameters.initcon_file
extradata.meta.no_merging = parameters.no_merging
num_planetesimals, num_embryos = \
add_particles_from_conditions_file(sim, extradata, parameters.initcon_file, testrun)
sim.move_to_com()
extradata.meta.initial_N = sim.N
extradata.meta.initial_N_planetesimal = num_planetesimals
extradata.meta.initial_N_embryo = num_embryos
extradata.history.append(energy=sim.calculate_energy(),
momentum=total_momentum(sim),
total_mass=total_mass(sim),
time=sim.t,
N=sim.N,
N_active=sim.N_active)
cputimeoffset = walltimeoffset = 0
t = 0
else:
if fn.with_suffix(".lock").exists():
raise FileExistsError(
"Lock file found, is the simulation currently running?")
copy(fn.with_suffix(".bin"), fn.with_suffix(".bak.bin"))
copy(fn.with_suffix(".extra.json"), fn.with_suffix(".extra.bak.json"))
sa = SimulationArchive(str(fn.with_suffix(".bin")))
extradata = ExtraData.load(fn)
tmax = extradata.meta.tmax
per_savestep = extradata.meta.per_savestep
sim = sa[-1]
t = round(sim.t + per_savestep)
print(f"continuing from {t}")
sim.move_to_com()
sim.ri_mercurius.recalculate_coordinates_this_timestep = 1
sim.integrator_synchronize()
if extradata.meta.git_hash != git_hash():
print(
"warning: The saved output was originally run with another version of the code"
)
print(f"original: {extradata.meta.git_hash}")
print(f"current: {git_hash()}")
num_savesteps = extradata.meta.num_savesteps
cputimeoffset = extradata.meta.cputime
walltimeoffset = extradata.meta.walltime
check_heartbeat_needs_recompile()
clibheartbeat = cdll.LoadLibrary("heartbeat/heartbeat.so")
clibheartbeat.init_logfile.argtypes = [c_char_p]
logfile = create_string_buffer(128)
logfile.value = str(fn.with_suffix(".energylog.csv")).encode()
clibheartbeat.init_logfile(logfile)
sim.heartbeat = clibheartbeat.heartbeat
innermost_semimajor_axis = third_kepler_law(
orbital_period=sim.dt * year *
MIN_TIMESTEP_PER_ORBIT) / astronomical_unit * 1.1
print(f"innermost semimajor axis is {innermost_semimajor_axis}")
c_double.in_dll(
clibheartbeat,
"min_distance_from_sun_squared").value = innermost_semimajor_axis**2
c_double.in_dll(clibheartbeat,
"max_distance_from_sun_squared").value = 150**2
assert sim.dt < innermost_period(sim) / MIN_TIMESTEP_PER_ORBIT
def collision_resolve_handler(sim_p: POINTER_REB_SIM,
collision: reb_collision) -> int:
global abort # needed as exceptions don't halt integration
try:
return merge_particles(sim_p, collision, ed=extradata)
except BaseException as exception:
print("exception during collision_resolve")
print(exception)
abort = True
sim_p.contents._status = 1
raise exception
sim.collision_resolve = collision_resolve_handler
# show_orbits(sim)
fn.with_suffix(".lock").touch()
print("start")
while t <= tmax:
print()
print(f"{t / tmax * 100:.2f}%")
set_process_title(fn, t / tmax, sim.N)
try:
print(f"integrating until {t}")
sim.integrate(t, exact_finish_time=0)
print("dt", sim.dt)
print("t", t)
t += per_savestep
except NoParticles:
print("No Particles left")
abort = True
print("N", sim.N)
print("N_active", sim.N_active)
print("fraction", innermost_period(sim) / MIN_TIMESTEP_PER_ORBIT)
assert sim.dt < innermost_period(sim) / MIN_TIMESTEP_PER_ORBIT
escape: hb_event
wide_orbit: hb_event
sun_collision: hb_event
for escape in hb_event_list.in_dll(clibheartbeat, "hb_escapes"):
if not escape.new:
continue
print("escape:", escape.time, escape.hash)
extradata.pdata[escape.hash].escaped = escape.time
escape.new = 0 # make sure to not handle it again
c_int.in_dll(clibheartbeat, "hb_escape_index").value = 0
for sun_collision in hb_event_list.in_dll(clibheartbeat,
"hb_sun_collisions"):
if not sun_collision.new:
continue
print("sun collision:", sun_collision.time, sun_collision.hash)
extradata.pdata[
sun_collision.hash].collided_with_sun = sun_collision.time
sun_collision.new = 0
c_int.in_dll(clibheartbeat, "hb_sun_collision_index").value = 0
for wide_orbit in hb_event_list.in_dll(clibheartbeat,
"hb_wide_orbits"):
if not wide_orbit.new:
continue
print("wide orbit:", wide_orbit.time, wide_orbit.hash)
extradata.pdata[wide_orbit.hash].wide_orbit = wide_orbit.time
wide_orbit.new = 0
c_int.in_dll(clibheartbeat, "hb_sun_collision_index").value = 0
sim.simulationarchive_snapshot(str(fn.with_suffix(".bin")))
extradata.meta.walltime = time.perf_counter() - start + walltimeoffset
extradata.meta.cputime = time.process_time() + cputimeoffset
extradata.meta.current_time = t
extradata.history.append(energy=sim.calculate_energy(),
momentum=total_momentum(sim),
total_mass=total_mass(sim),
time=sim.t,
N=sim.N,
N_active=sim.N_active)
extradata.save(fn)
if abort:
print("aborted")
exit(1)
print("finished")
fn.with_suffix(".lock").unlink()