initial commit

.gitignore

@@ -0,0 +1,9 @@
+initcon
+bac
+.idea
+*.json
+*.bin
+super.in
+__pycache__/
+*.gv
+*.pdf

analyze.old.py

@@ -0,0 +1,28 @@
+import json
+import matplotlib.pyplot as plt
+
+import numpy as np
+from rebound import SimulationArchive, OrbitPlot
+
+sa = SimulationArchive("water.bin")
+with open("water.meta.json") as f:
+    meta = json.load(f)
+    tmax = meta["tmax"]
+    savesteps = meta["savesteps"]
+    max_n = meta["max_n"]
+
+times = np.linspace(0., tmax, savesteps)
+data = np.zeros((3, savesteps * max_n))
+j = 0
+
+for i, t in enumerate(times):
+    sim = sa.getSimulation(t=t)
+    for particle in range(1, sim.N):
+        data[0, j] = i
+        data[1, j] = particle
+        data[2, j] = sim.particles[particle].a
+        j += 1
+plt.scatter(data[0, :], data[2, :], c=data[1, :], cmap="tab10", s=.4)
+plt.colorbar()
+# OrbitPlot(sim, slices=1)
+plt.show()

analyze.py

@@ -0,0 +1,34 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from rebound import SimulationArchive, Particle
+
+from extradata import ExtraData
+from utils import filename_from_argv
+
+fn = filename_from_argv()
+sa = SimulationArchive(str(fn.with_suffix(".bin")))
+ed = ExtraData.load(fn.with_suffix(".extra.json"))
+
+times = np.linspace(0., ed.meta.current_time, ed.meta.current_steps)
+
+data = {}
+for i, t in enumerate(times):
+    sim = sa.getSimulation(t=t)
+    for pn in range(1, sim.N):
+        part: Particle = sim.particles[pn]
+        hash = part.hash.value
+        if hash not in data:
+            data[hash] = ([], [])
+        data[hash][0].append(part.x)
+        data[hash][1].append(part.y)
+
+for name, d in data.items():
+    times, values = d
+    print(list(map(len, [times, values])))
+    if True:
+        plt.scatter(times, values, label=name, s=.9)
+    else:
+        plt.plot(times, values, label=name)
+# plt.legend()
+# OrbitPlot(sim, slices=1)
+plt.show()

extradata.py

@@ -0,0 +1,73 @@
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict
+
+from rebound import Particle
+
+
+class ExtraData:
+
+    def __init__(self):
+        self.tree = CollisionTree()
+        self.pdata: Dict[int, ParticleData] = {}
+        self.meta = Meta()
+
+    def save(self, filename: Path):
+        pdata = {}
+        for k, v in self.pdata.items():
+            pdata[k] = v.__dict__
+
+        with filename.open("w") as f:
+            json.dump({
+                "meta": self.meta.save(),
+                "pdata": pdata,
+                "tree": self.tree.save()
+            }, f, indent=2)
+
+    @classmethod
+    def load(cls, filename: Path):
+        with filename.open() as f:
+            data = json.load(f)
+        self = cls()
+        self.meta = Meta(**data["meta"])
+
+        self.tree.load(data["tree"])
+
+        for k, v in data["pdata"].items():
+            self.pdata[k] = ParticleData(**v)
+
+        return self
+
+
+@dataclass
+class ParticleData:
+    water_mass_fraction: float
+
+
+@dataclass
+class Meta:
+    tmax: float = None
+    savesteps: int = None
+    max_n: int = None
+    walltime: int = None  # seconds
+    cputime: int = None  # seconds
+    current_time: float = None
+    current_steps: float = None
+
+    def save(self):
+        return self.__dict__
+
+
+class CollisionTree:
+    def __init__(self):
+        self._tree = {}
+
+    def add(self, source1: Particle, source2: Particle, to: Particle):
+        self._tree[to.hash.value] = [source1.hash.value, source2.hash.value]
+
+    def save(self):
+        return self._tree
+
+    def load(self, tree):
+        self._tree = tree

graph.py

@@ -0,0 +1,13 @@
+from graphviz import Digraph
+
+from extradata import ExtraData
+from utils import filename_from_argv
+
+ed = ExtraData.load(filename_from_argv().with_suffix(".extra.json"))
+
+dot = Digraph(comment='Collisions')
+for merged, originals in ed.tree._tree.items():
+    for parent in originals:
+        dot.edge(str(parent), str(merged))
+
+dot.render('graph.gv', view=True)

main.py

@@ -0,0 +1,43 @@
+import json
+from math import radians
+
+import numpy as np
+from rebound import Particle, Simulation
+from reboundx import Extras
+
+data = [
+    [5.202887e0, 0.04838624e0, 1.30439695e0, 274.2545707e0, 100.4739091e0, 19.66796068e0, 0.000954792e0, "JUPITER"],
+    [9.53667594e0, 0.05386179e0, 2.48599187e0, 338.9364538e0, 113.6624245e0, 317.3553659e0, 0.000285886e0, "SATURN"],
+    [19.18916464e0, 0.04725744e0, 0.77263783e0, 96.93735127e0, 74.01692503e0, 142.2838282e0, 4.36624e-05, "URANUS"],
+    [30.06992276e0, 0.00859048e0, 1.77004347e0, 273.1805365e0, 131.7842257e0, 259.915208e0, 5.15139e-05, "NEPTUNE"]
+]
+
+
+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)
+
+
+if __name__ == '__main__':
+    main()

merge.py

@@ -0,0 +1,131 @@
+import sys
+from typing import List
+
+import numpy as np
+from numpy import linalg, sqrt
+from rebound import Simulation, Particle
+from scipy.constants import astronomical_unit, G, year
+
+from extradata import ExtraData, ParticleData
+from radius_utils import radius
+from utils import unique_hash, clamp
+
+sys.path.append("./bac")
+
+from bac.simulation_list import SimulationList
+from bac.CustomScaler import CustomScaler
+from bac.interpolators.rbf import RbfInterpolator
+
+simulations = SimulationList.jsonlines_load()
+
+scaler = CustomScaler()
+scaler.fit(simulations.X)
+
+scaled_data = scaler.transform_data(simulations.X)
+water_interpolator = RbfInterpolator(scaled_data, simulations.Y_water)
+mass_interpolator = RbfInterpolator(scaled_data, simulations.Y_mass)
+
+
+def get_mass_fractions(alpha, velocity_original, escape_velocity, gamma, projectile_mass, target_water_fraction,
+                       projectile_water_fraction):
+    velocity_si = velocity_original * astronomical_unit / year
+    print("v_esc", escape_velocity)
+    velocity = velocity_si / escape_velocity
+    print("v_orig,v_si", velocity_original, velocity_si)
+    print("v", velocity)
+    if alpha > 90:
+        alpha = 180 - alpha
+    if gamma > 1:
+        gamma = 1 / gamma
+    alpha = clamp(alpha, 0, 60)
+    velocity = clamp(velocity, 1, 5)
+
+    m_ceres = 9.393e+20
+    m_earth = 5.9722e+24
+    projectile_mass = clamp(projectile_mass, 2 * m_ceres, 2 * m_earth)
+    gamma = clamp(gamma, 1 / 10, 1)
+
+    testinput = [alpha, velocity, projectile_mass, gamma,
+                 target_water_fraction, projectile_water_fraction]
+
+    print("# alpha velocity projectile_mass gamma target_water_fraction projectile_water_fraction\n")
+    print(" ".join(map(str, testinput)))
+
+    scaled_input = list(scaler.transform_parameters(testinput))
+    water_retention = water_interpolator.interpolate(*scaled_input)
+    mass_retention = mass_interpolator.interpolate(*scaled_input)
+
+    water_retention = clamp(water_retention, 0, 1)
+    mass_retention = clamp(mass_retention, 0, 1)
+
+    return water_retention, mass_retention
+
+
+def merge_particles(sim: Simulation, ed: ExtraData):
+    print("colliding")
+    collided: List[Particle] = []
+    p: Particle
+    for p in sim.particles:
+        # print(p.lastcollision, sim.t)
+        # if p.lastcollision == sim.t:
+        if p.lastcollision >= sim.t - 1:
+            collided.append(p)
+    # if not collided:
+    #     print("empty collision")
+    #     return
+    print(collided)
+    assert len(collided) == 2, "More or fewer than 2 objects collided with each other"
+    cp1: Particle  # projectile
+    cp2: Particle  # target
+    cp1, cp2 = collided
+
+    projectile_wmf = ed.pdata[cp1.hash.value].water_mass_fraction
+    target_wmf = ed.pdata[cp2.hash.value].water_mass_fraction
+
+    v1 = np.array(cp1.vxyz)
+    v2 = np.array(cp2.vxyz)
+    vdiff = linalg.norm(v2 - v1)  # AU/year
+    v1_u = v1 / linalg.norm(v1)
+    v2_u = v2 / linalg.norm(v2)
+    # https://stackoverflow.com/a/13849249/4398037
+    ang = np.degrees(np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0)))
+
+    gamma = cp1.m / cp2.m
+
+    escape_velocity = sqrt(2 * G * (cp1.m + cp2.m) / ((cp1.r + cp2.r) * astronomical_unit))
+
+    print("interpolating")
+    water_ret, stone_ret = get_mass_fractions(
+        alpha=ang, velocity_original=vdiff, escape_velocity=escape_velocity, gamma=gamma, projectile_mass=cp1.m,
+        target_water_fraction=target_wmf, projectile_water_fraction=projectile_wmf)
+    print(water_ret, stone_ret)
+    hash = unique_hash()
+
+    water_mass = cp1.m * projectile_wmf + cp2.m * target_wmf
+    stone_mass = cp1.m + cp2.m - water_mass
+
+    water_mass *= water_ret
+    stone_ret *= stone_ret
+
+    total_mass = water_mass + stone_mass
+    final_wmf = water_mass / total_mass
+    print(final_wmf)
+
+    merged_planet = (cp1 * cp1.m + cp2 * cp2.m) / (cp1.m + cp2.m)
+    merged_planet.m = total_mass
+    merged_planet.hash = hash
+
+    merged_planet.r = radius(merged_planet.m, final_wmf) / astronomical_unit
+    ed.pdata[hash.value] = ParticleData(water_mass_fraction=final_wmf)
+
+    ed.tree.add(cp1, cp2, merged_planet)
+
+    cp1_hash=cp1.hash
+    cp2_hash=cp2.hash
+
+    # don't use cp1 and cp2 from now on as they will change
+
+    print("removing",cp1_hash.value,cp2_hash.value)
+    sim.remove(hash=cp1_hash)
+    sim.remove(hash=cp2_hash)
+    sim.add(merged_planet)

pyproject.toml

@@ -0,0 +1,21 @@
+[tool.poetry]
+name = "reboundtest"
+version = "0.1.0"
+description = ""
+authors = ["Lukas Winkler <git@lw1.at>"]
+
+[tool.poetry.dependencies]
+python = "^3.7"
+rebound = "^3.12.1"
+matplotlib = "^3.2.1"
+pyqt5 = "^5.14.1"
+ipywidgets = "^7.5.1"
+reboundx = "^3.0.4"
+scipy = "^1.4.1"
+graphviz = "^0.13.2"
+
+[tool.poetry.dev-dependencies]
+
+[build-system]
+requires = ["poetry>=0.12"]
+build-backend = "poetry.masonry.api"

radius_utils.py

@@ -0,0 +1,21 @@
+from math import pi
+
+ice_density = 0.917 / 1000 * 100 ** 3  # kg/m^3
+basalt_density = 2.7 / 1000 * 100 ** 3  # kg/m^3
+
+
+def core_radius(total_mass, water_fraction, density):
+    core_mass = total_mass * (1 - water_fraction)
+    return (core_mass / density * 3 / 4 / pi) ** (1 / 3)
+
+
+def total_radius(total_mass, water_fraction, density, inner_radius):
+    mantle_mass = total_mass * water_fraction
+    return (mantle_mass / density * 3 / 4 / pi + inner_radius ** 3) ** (1 / 3)
+
+
+def radius(mass: float, water_fraction: float) -> float:
+    """
+    :return: radius in Meter
+    """
+    return total_radius(mass, water_fraction, ice_density, core_radius(mass, water_fraction, basalt_density))

requirements.txt

@@ -0,0 +1,57 @@
+appnope==0.1.0; sys_platform == "darwin" or platform_system == "Darwin" or python_version >= "3.3" and sys_platform == "darwin"
+attrs==19.3.0
+backcall==0.1.0
+bleach==3.1.3
+colorama==0.4.3; python_version >= "3.3" and sys_platform == "win32" or sys_platform == "win32"
+cycler==0.10.0
+decorator==4.4.2
+defusedxml==0.6.0
+entrypoints==0.3
+graphviz==0.13.2
+importlib-metadata==1.5.0; python_version < "3.8"
+ipykernel==5.1.4
+ipython==7.13.0
+ipython-genutils==0.2.0
+ipywidgets==7.5.1
+jedi==0.16.0
+jinja2==2.11.1
+jsonschema==3.2.0
+jupyter-client==6.1.0
+jupyter-core==4.6.3
+kiwisolver==1.1.0
+markupsafe==1.1.1
+matplotlib==3.2.1
+mistune==0.8.4
+nbconvert==5.6.1
+nbformat==5.0.4
+notebook==6.0.3
+numpy==1.18.2
+pandocfilters==1.4.2
+parso==0.6.2
+pexpect==4.8.0; python_version >= "3.3" and sys_platform != "win32" or sys_platform != "win32"
+pickleshare==0.7.5
+prometheus-client==0.7.1
+prompt-toolkit==3.0.4
+ptyprocess==0.6.0; python_version >= "3.3" and sys_platform != "win32" or sys_platform != "win32" or os_name != "nt" or python_version >= "3.3" and sys_platform != "win32" and (python_version >= "3.3" and sys_platform != "win32" or sys_platform != "win32")
+pygments==2.6.1
+pyparsing==2.4.6
+pyqt5==5.14.1
+pyqt5-sip==12.7.1
+pyrsistent==0.15.7
+python-dateutil==2.8.1
+pywin32==227; sys_platform == "win32"
+pywinpty==0.5.7; os_name == "nt"
+pyzmq==19.0.0
+rebound==3.12.1
+reboundx==3.0.4
+scipy==1.4.1
+send2trash==1.5.0
+six==1.14.0
+terminado==0.8.3
+testpath==0.4.4
+tornado==6.0.4
+traitlets==4.3.3
+wcwidth==0.1.8
+webencodings==0.5.1
+widgetsnbextension==3.5.1
+zipp==3.1.0; python_version < "3.8"

utils.py

@@ -0,0 +1,28 @@
+from ctypes import c_uint32
+from pathlib import Path
+from random import randint
+from sys import argv
+
+
+def unique_hash() -> c_uint32:
+    """
+    returns a (hopefully) unique 32 bit integer to be used as a particle hash
+
+    when a collision occurs and ruins the output data, please complain to the universe
+    """
+    return c_uint32(randint(0, 2 ** 32 - 1))
+
+
+def clamp(n: float, smallest: float, largest: float) -> float:
+    assert smallest < largest
+    return max(smallest, min(n, largest))
+
+
+def filename_from_argv() -> Path:
+    if len(argv) < 2:
+        raise ValueError("specify filename")
+    fn = argv[1]
+    fn = fn.replace(".bin", "").replace(".meta.json", "")
+    if fn.endswith("."):
+        fn = fn[:-1]
+    return Path(fn.replace(".bin", "").replace(".meta.json", ""))

water_sim.py

@@ -0,0 +1,104 @@
+import time
+from math import radians
+
+import numpy as np
+from rebound import Simulation, Particle, Collision, Escape, reb_simulation_integrator_mercurius
+from scipy.constants import astronomical_unit
+
+from extradata import ExtraData, ParticleData
+from merge import merge_particles
+from radius_utils import radius
+from utils import unique_hash, filename_from_argv
+
+start = time.perf_counter()
+
+sim = Simulation()
+
+fn = filename_from_argv()
+
+sim.units = ('yr', 'AU', 'kg')
+# sim.boundary = "open"
+# boxsize = 20
+# sim.configure_box(boxsize)
+# sim.integrator = "mercurius"
+# sim.dt = 1e-3
+# sim.ri_whfast.safe_mode = 0
+sim.collision = "direct"
+# sim.ri_mercurius.hillfac = 3.
+# sim.collision_resolve = "merge"
+
+extradata = ExtraData()
+
+i = 0
+with open("initcon/conditions.input") as f:
+    for line in f:
+        if line.startswith("#") or line.startswith("ERROR") or line == "\n":
+            continue
+        columns = list(map(float, line.split()))
+        hash = unique_hash()
+        if len(columns) > 7:
+            print(columns[7:])
+            cmf, mmf, wmf = columns[7:]
+            assert cmf + mmf + wmf == 1
+            extradata.pdata[hash.value] = ParticleData(water_mass_fraction=wmf)
+        else:
+            wmf = 0
+        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=radians(columns[4]),
+                            Omega=radians(columns[5]), M=radians(columns[6]),
+                            simulation=sim,
+                            hash=hash,
+                            r=radius(columns[0], wmf) / astronomical_unit
+                            )
+        sim.add(part)
+
+        i += 1
+        if i == 7:  # only use 29 objects to make results comparable with lie-simulation
+            break
+
+print(sim.N)
+# from matplotlib import pyplot as plt
+# OrbitPlot(sim,slices=1,color=True)
+# plt.show()
+
+max_n = sim.N
+print("start")
+tmax = 1e6
+savesteps = 3000
+times = np.linspace(0., tmax, savesteps)
+sim.move_to_com()
+sim.exit_max_distance = 15
+
+extradata.meta.tmax = tmax
+extradata.meta.savesteps = savesteps
+extradata.meta.max_n = max_n
+
+try:
+    fn.with_suffix(".bin").unlink()
+except OSError:
+    pass
+# sim.automateSimulationArchive("water.bin", interval=1, deletefile=True)
+for i, t in enumerate(times, start=1):
+    try:
+        sim.integrate(t)
+        print(sim.dt)
+        merc: reb_simulation_integrator_mercurius = sim.ri_mercurius
+        print(merc._encounterN)
+        print(merc.mode)
+    except Collision:
+        merge_particles(sim, extradata)
+    except Escape:
+        print("something escaped")
+    print(f"{i / savesteps * 100:.2f}% ({sim.N})")
+    sim.simulationarchive_snapshot(str(fn.with_suffix(".bin")))
+    extradata.meta.walltime = time.perf_counter() - start
+    extradata.meta.cputime = time.process_time()
+    extradata.meta.current_time = t
+    extradata.meta.current_steps = i
+    extradata.save(fn.with_suffix(".extra.json"))
+
+# OrbitPlot(sim,slices=1,color=True)
+# plt.show()