rebound-collisions/water_sim.py

import re
import time
from ctypes import Structure, c_uint32, c_double, c_uint, cdll, c_int, create_string_buffer, c_char_p
from dataclasses import dataclass
from pathlib import Path
from shutil import copy
from sys import argv
from typing import Tuple

import rebound
import yaml
from rebound import Simulation, Particle, NoParticles, SimulationArchive
from rebound.simulation import POINTER_REB_SIM, reb_collision
from scipy.constants import astronomical_unit, mega, year

from extradata import ExtraData, ParticleData
from merge import merge_particles
from utils import unique_hash, filename_from_argv, innermost_period, total_momentum, process_friendlyness, total_mass, \
    third_kepler_law, solar_radius, git_hash, check_heartbeat_needs_recompile, PlanetaryRadius, set_process_title

MIN_TIMESTEP_PER_ORBIT = 20

abort = False


class hb_event(Structure):
    _fields_ = [("hash", c_uint32),
                ("time", c_double),
                ("new", c_uint)]
    hash: int
    time: float
    new: int


hb_event_list = hb_event * 500


@dataclass
class Parameters:
    initcon_file: str
    massloss_method: str
    no_merging: bool = False


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()


if __name__ == '__main__':
    fn = filename_from_argv()
    process_friendlyness(fn)
    testrun = False
    if len(argv) > 2 and argv[2] == "test":
        testrun = True
    try:
        main(fn, testrun)
    except KeyboardInterrupt:
        print("aborting")
        lockfile = fn.with_suffix(".lock")
        print(f"deleting {lockfile}")
        lockfile.unlink()