2021-01-02 18:23:39 +01:00
|
|
|
import re
|
2020-03-31 15:39:59 +02:00
|
|
|
import time
|
|
|
|
from math import radians
|
2021-01-02 18:23:39 +01:00
|
|
|
from pathlib import Path
|
2021-01-02 16:30:39 +01:00
|
|
|
from shutil import copy
|
2021-02-03 22:04:34 +01:00
|
|
|
from sys import argv
|
2020-03-31 15:39:59 +02:00
|
|
|
|
2021-01-21 14:28:29 +01:00
|
|
|
from rebound import Simulation, Particle, NoParticles, Escape, \
|
2020-12-30 20:41:30 +01:00
|
|
|
SimulationArchive
|
2021-01-21 14:28:29 +01:00
|
|
|
from rebound.simulation import POINTER_REB_SIM, reb_collision
|
2021-01-23 14:41:41 +01:00
|
|
|
from scipy.constants import astronomical_unit, mega
|
2020-03-31 15:39:59 +02:00
|
|
|
|
|
|
|
from extradata import ExtraData, ParticleData
|
2020-12-27 13:47:03 +01:00
|
|
|
from merge import merge_particles, handle_escape
|
2020-03-31 15:39:59 +02:00
|
|
|
from radius_utils import radius
|
2021-01-23 14:41:41 +01:00
|
|
|
from utils import unique_hash, filename_from_argv, innermost_period, total_impulse, process_friendlyness
|
2020-03-31 15:39:59 +02:00
|
|
|
|
2020-12-30 20:41:30 +01:00
|
|
|
MIN_TIMESTEP_PER_ORBIT = 20
|
|
|
|
PERFECT_MERGING = False
|
2021-01-02 18:23:39 +01:00
|
|
|
INITCON_FILE = Path("initcon/conditions_many.input")
|
2020-03-31 15:39:59 +02:00
|
|
|
|
2021-01-24 12:13:21 +01:00
|
|
|
abort = False
|
|
|
|
|
2021-01-08 15:10:38 +01:00
|
|
|
|
2021-02-03 22:04:34 +01:00
|
|
|
def main(fn: Path, testrun=False):
|
2021-01-24 12:13:21 +01:00
|
|
|
global abort
|
2021-01-08 15:10:38 +01:00
|
|
|
start = time.perf_counter()
|
|
|
|
|
|
|
|
if not fn.with_suffix(".bin").exists():
|
|
|
|
# set up a fresh simulation
|
|
|
|
sim = Simulation()
|
|
|
|
|
|
|
|
sim.units = ('yr', 'AU', 'kg')
|
|
|
|
# sim.boundary = "open"
|
|
|
|
# boxsize = 100
|
|
|
|
# sim.configure_box(boxsize)
|
|
|
|
sim.exit_max_distance = 30
|
|
|
|
sim.integrator = "mercurius"
|
|
|
|
# sim.collision = 'line'
|
|
|
|
sim.dt = 1e-2
|
|
|
|
sim.ri_ias15.min_dt = 0.0001 / 365
|
|
|
|
sim.collision = "direct"
|
|
|
|
sim.ri_mercurius.hillfac = 3.
|
2021-01-23 14:41:41 +01:00
|
|
|
tmax = 200 * mega
|
2021-01-08 15:10:38 +01:00
|
|
|
num_savesteps = 20000
|
2021-02-03 22:04:34 +01:00
|
|
|
if testrun:
|
|
|
|
tmax /= 200000
|
|
|
|
num_savesteps /= 1000
|
2021-01-08 15:10:38 +01:00
|
|
|
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.perfect_merging = PERFECT_MERGING
|
|
|
|
|
|
|
|
initcon = INITCON_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()
|
|
|
|
if len(columns) > 7:
|
|
|
|
# print(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 cmf")
|
|
|
|
cmf += diff
|
|
|
|
assert cmf + mmf + wmf - 1 <= 1e-10
|
|
|
|
if i > num_embryos + 3:
|
|
|
|
object_type = "planetesimal"
|
|
|
|
else:
|
|
|
|
object_type = "embryo"
|
2020-12-30 20:41:30 +01:00
|
|
|
else:
|
2021-01-08 15:10:38 +01:00
|
|
|
wmf = mmf = 0
|
|
|
|
cmf = 1
|
|
|
|
if columns[1] == 0:
|
|
|
|
object_type = "sun"
|
|
|
|
else:
|
|
|
|
object_type = "gas giant"
|
|
|
|
extradata.pdata[hash.value] = ParticleData(
|
|
|
|
water_mass_fraction=wmf,
|
2021-02-03 22:04:34 +01:00
|
|
|
type=object_type,
|
|
|
|
total_mass=columns[0]
|
2021-01-02 18:23:39 +01:00
|
|
|
)
|
2021-01-08 15:10:38 +01:00
|
|
|
|
|
|
|
if columns[1] == 0: # that should not be needed, but nevertheless is
|
|
|
|
part = Particle(m=columns[0], hash=hash)
|
|
|
|
else:
|
|
|
|
part = Particle(
|
|
|
|
m=columns[0], a=columns[1], e=columns[2],
|
|
|
|
inc=radians(columns[3]), omega=columns[4],
|
|
|
|
Omega=columns[5], M=columns[6],
|
|
|
|
simulation=sim,
|
|
|
|
hash=hash,
|
|
|
|
r=radius(columns[0], wmf) / astronomical_unit
|
|
|
|
)
|
|
|
|
sim.add(part)
|
|
|
|
i += 1
|
|
|
|
assert sim.N == num_planetesimals + num_embryos + 3
|
|
|
|
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.energy.set_initial_energy(sim.calculate_energy())
|
|
|
|
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.with_suffix(".extra.json"))
|
|
|
|
tmax = extradata.meta.tmax
|
|
|
|
per_savestep = extradata.meta.per_savestep
|
|
|
|
t = extradata.meta.current_time - per_savestep
|
|
|
|
sim = sa.getSimulation(t=t)
|
|
|
|
sim.move_to_com()
|
2021-01-25 15:20:45 +01:00
|
|
|
sim.ri_mercurius.recalculate_coordinates_this_timestep = 1
|
2021-01-08 15:10:38 +01:00
|
|
|
sim.integrator_synchronize()
|
|
|
|
|
|
|
|
num_savesteps = extradata.meta.num_savesteps
|
|
|
|
cputimeoffset = extradata.meta.cputime
|
|
|
|
walltimeoffset = extradata.meta.walltime
|
|
|
|
|
2020-12-30 20:41:30 +01:00
|
|
|
assert sim.dt < innermost_period(sim) / MIN_TIMESTEP_PER_ORBIT
|
2021-01-08 15:10:38 +01:00
|
|
|
|
2021-01-21 14:28:29 +01:00
|
|
|
def collision_resolve_handler(sim_p: POINTER_REB_SIM, collision: reb_collision) -> int:
|
2021-01-24 12:13:21 +01:00
|
|
|
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
|
2021-01-21 14:28:29 +01:00
|
|
|
|
|
|
|
sim.collision_resolve = collision_resolve_handler
|
|
|
|
|
2021-01-08 15:10:38 +01:00
|
|
|
# show_orbits(sim)
|
|
|
|
|
|
|
|
fn.with_suffix(".lock").touch()
|
|
|
|
print("start")
|
|
|
|
|
|
|
|
while t <= tmax:
|
|
|
|
print()
|
|
|
|
print(f"{t / tmax * 100:.2f}%")
|
|
|
|
try:
|
|
|
|
print(f"integrating until {t}")
|
|
|
|
sim.integrate(t, exact_finish_time=0)
|
|
|
|
print("dt", sim.dt)
|
|
|
|
print("t", t)
|
|
|
|
t += per_savestep
|
|
|
|
except Escape:
|
|
|
|
handle_escape(sim, extradata)
|
|
|
|
except NoParticles:
|
|
|
|
print("No Particles left")
|
|
|
|
abort = True
|
|
|
|
print("N", sim.N)
|
|
|
|
print("N_active", sim.N_active)
|
|
|
|
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.meta.current_steps = i
|
|
|
|
extradata.energy.add_energy_value(sim.calculate_energy())
|
|
|
|
print("total impulse", total_impulse(sim))
|
|
|
|
extradata.save(fn.with_suffix(".extra.json"))
|
|
|
|
assert sim.dt < innermost_period(sim) / MIN_TIMESTEP_PER_ORBIT
|
|
|
|
print("fraction", innermost_period(sim) / MIN_TIMESTEP_PER_ORBIT)
|
|
|
|
if abort:
|
|
|
|
exit(1)
|
|
|
|
print("finished")
|
|
|
|
fn.with_suffix(".lock").unlink()
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
fn = filename_from_argv()
|
2021-01-23 14:41:41 +01:00
|
|
|
process_friendlyness(fn)
|
2021-02-03 22:04:34 +01:00
|
|
|
testrun = False
|
|
|
|
if len(argv) > 2 and argv[2] == "test":
|
|
|
|
testrun = True
|
2021-01-08 15:10:38 +01:00
|
|
|
try:
|
2021-02-03 22:04:34 +01:00
|
|
|
main(fn, testrun)
|
2021-01-08 15:10:38 +01:00
|
|
|
except KeyboardInterrupt:
|
|
|
|
print("aborting")
|
|
|
|
lockfile = fn.with_suffix(".lock")
|
|
|
|
print(f"deleting {lockfile}")
|
|
|
|
lockfile.unlink()
|