use heartbeat for proper escape and sun collision handling

extradata.py

@@ -13,6 +13,7 @@ class ParticleData:
     water_mass_fraction: float
     type: str
     escaped: float = None
+    collided_with_sun: float = None
     total_mass: float = None
 
     @property

heartbeat/.gitignore

@@ -0,0 +1,4 @@
+.idea
+*.o
+*.so
+rebound.h

heartbeat/build.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+set -x
+set -e
+
+gcc -c -fPIC -O3 -Wall heartbeat.c -o heartbeat.o
+gcc -L. -shared heartbeat.o -o heartbeat.so -lrebound -Wl,-rpath='./heartbeat' -Wall

heartbeat/heartbeat.c

@@ -0,0 +1,59 @@
+#include "rebound.h"
+
+double min_distance_from_sun_squared = 0;
+double max_distance_from_sun_squared = 0;
+
+struct hb_event {
+    uint32_t hash;
+    double time;
+    unsigned int new;
+};
+
+struct hb_event hb_escapes[500];
+struct hb_event hb_sun_collisions[500];
+
+int hb_escape_index = 0;
+int hb_sun_collision_index = 0;
+
+void heartbeat(struct reb_simulation *sim) {
+    if ((sim->steps_done % 100) == 0) {
+        const struct reb_particle *const particles = sim->particles;
+        int N = sim->N - sim->N_var;
+        for (int i = 1; i < N; i++) { // skip sun
+            struct reb_particle p = particles[i];
+            double distance_squared = p.x * p.x + p.y * p.y + p.z * p.z;
+            if (distance_squared > max_distance_from_sun_squared) {
+                printf("max: %f\n", max_distance_from_sun_squared);
+                printf("min: %f\n", min_distance_from_sun_squared);
+                printf("dist: %f\n", distance_squared);
+                printf("px: %f\n", p.x);
+                printf("py: %f\n", p.y);
+                printf("pz: %f\n", p.z);
+                printf("remove %u at t=%f (max)\n", p.hash, sim->t);
+                reb_remove_by_hash(sim, p.hash, 1);
+                hb_escapes[hb_escape_index].hash = p.hash;
+                hb_escapes[hb_escape_index].time = sim->t;
+                hb_escapes[hb_escape_index].new = 1;
+                hb_escape_index++;
+            } else if (distance_squared < min_distance_from_sun_squared) {
+                printf("remove %u at t=%f (min)\n", p.hash, sim->t);
+                reb_remove_by_hash(sim, p.hash, 1);
+                hb_sun_collisions[hb_sun_collision_index].hash = p.hash;
+                hb_sun_collisions[hb_sun_collision_index].time = sim->t;
+                hb_sun_collisions[hb_sun_collision_index].new = 1;
+
+                hb_sun_collision_index++;
+            }
+            if (
+                    (distance_squared > max_distance_from_sun_squared)
+                    || (distance_squared < min_distance_from_sun_squared)
+                    ) {
+                N--;
+                reb_move_to_com(sim);
+                reb_integrator_synchronize(sim);
+                sim->ri_mercurius.recalculate_coordinates_this_timestep = 1;
+                sim->ri_mercurius.recalculate_dcrit_this_timestep = 1;
+            }
+        }
+    }
+}

heartbeat/symlinks.py

@@ -0,0 +1,16 @@
+from pathlib import Path
+
+import rebound
+
+Path("librebound.so").unlink(missing_ok=True)
+Path("rebound.h").unlink(missing_ok=True)
+
+module_dir = Path(rebound.__file__).parent
+rebound_h = module_dir / "rebound.h"
+librebound = Path(rebound.__libpath__)
+
+print(librebound)
+print(rebound_h)
+
+Path("librebound.so").symlink_to(librebound)
+Path("rebound.h").symlink_to(rebound_h)

merge.py

@@ -214,23 +214,3 @@ def merge_particles(sim_p: POINTER_REB_SIM, collision: reb_collision, ed: ExtraD
         return 2
     else:
         return 1
-
-
-def handle_escape(sim: Simulation, ed: ExtraData):
-    escaped_particle = None
-    p: Particle
-    for p in sim.particles:
-        distance_squared = p.x ** 2 + p.y ** 2 + p.z ** 2
-        if distance_squared > sim.exit_max_distance ** 2:
-            escaped_particle = p
-    if not escaped_particle:
-        raise RuntimeError("Escape without escaping particle")
-    ed.pd(escaped_particle).escaped = sim.t
-    sim.remove(hash=escaped_particle.hash)
-    del escaped_particle  # to avoid using the invalid pointer afterwards
-
-    # reorder_particles(sim, ed)
-    sim.move_to_com()
-    sim.integrator_synchronize()
-    sim.ri_mercurius.recalculate_coordinates_this_timestep = 1
-    sim.ri_mercurius.recalculate_dcrit_this_timestep = 1

water_sim.py

@@ -1,17 +1,17 @@
 import re
 import time
+from ctypes import Structure, c_uint32, c_double, c_uint, cdll, c_int
 from math import radians
 from pathlib import Path
 from shutil import copy
 from sys import argv
 
-from rebound import Simulation, Particle, NoParticles, Escape, \
-    SimulationArchive
+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, handle_escape
+from merge import merge_particles
 from radius_utils import radius
 from utils import unique_hash, filename_from_argv, innermost_period, total_momentum, process_friendlyness, total_mass, \
     third_kepler_law, solar_radius
@@ -23,6 +23,18 @@ INITCON_FILE = Path("initcon/conditions_many.input")
 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
+
+
 def main(fn: Path, testrun=False):
     global abort
     start = time.perf_counter()
@@ -35,7 +47,6 @@ def main(fn: Path, testrun=False):
         # sim.boundary = "open"
         # boxsize = 100
         # sim.configure_box(boxsize)
-        sim.exit_max_distance = 30
         sim.integrator = "mercurius"
         # sim.collision = 'line'
         sim.dt = 1e-2
@@ -55,12 +66,6 @@ def main(fn: Path, testrun=False):
         extradata.meta.num_savesteps = num_savesteps
         extradata.meta.perfect_merging = PERFECT_MERGING
 
-        innermost_semimajor_axis = third_kepler_law(
-            orbital_period=sim.dt * year * MIN_TIMESTEP_PER_ORBIT
-        ) / astronomical_unit
-        print(f"innermost semimajor axis is {innermost_semimajor_axis}")
-
-
         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))
@@ -145,6 +150,16 @@ def main(fn: Path, testrun=False):
         cputimeoffset = extradata.meta.cputime
         walltimeoffset = extradata.meta.walltime
 
+    clibheartbeat = cdll.LoadLibrary("heartbeat/heartbeat.so")
+    sim.heartbeat = clibheartbeat.heartbeat
+    innermost_semimajor_axis = third_kepler_law(
+        orbital_period=sim.dt * year * MIN_TIMESTEP_PER_ORBIT
+    ) / astronomical_unit
+    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 = 30 ** 2
+
     assert sim.dt < innermost_period(sim) / MIN_TIMESTEP_PER_ORBIT
 
     def collision_resolve_handler(sim_p: POINTER_REB_SIM, collision: reb_collision) -> int:
@@ -174,13 +189,29 @@ def main(fn: Path, testrun=False):
             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)
+
+        escape: 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
+        print(c_int.in_dll(clibheartbeat, "hb_sun_collision_index").value, "found")
+        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
@@ -194,8 +225,8 @@ def main(fn: Path, testrun=False):
             N_active=sim.N_active
         )
         extradata.save(fn)
-        assert sim.dt < innermost_period(sim) / MIN_TIMESTEP_PER_ORBIT
         print("fraction", innermost_period(sim) / MIN_TIMESTEP_PER_ORBIT)
+        # assert sim.dt < innermost_period(sim) / MIN_TIMESTEP_PER_ORBIT
         if abort:
             exit(1)
     print("finished")