rebound-collisions/merge.py

import sys
from pathlib import Path
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(Path("./bac/save.jsonl"))

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 interpolate(alpha, velocity, projectile_mass, gamma):
    hard_coded_water_mass_fraction = 0.15  # workaround to get proper results for water poor collisions
    testinput = [alpha, velocity, projectile_mass, gamma,
                 hard_coded_water_mass_fraction, hard_coded_water_mass_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)
    return float(water_retention), float(mass_retention)


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)
    orig_velocity = velocity
    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)

    water_retention, mass_retention = interpolate(alpha, velocity, projectile_mass, gamma)

    water_retention = clamp(water_retention, 0, 1)
    mass_retention = clamp(mass_retention, 0, 1)

    metadata = {"water_retention": water_retention, "mass_retention": mass_retention,
                "testinput": [alpha, velocity, projectile_mass, gamma],
                "velocity_si": velocity_si, "escape_velocity": escape_velocity, "orig_velocity": orig_velocity}

    return water_retention, mass_retention, metadata


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 - sim.dt:
            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"

    # the assignment to cp1 or cp2 is mostly random
    # (cp1 is the one with a lower index in sim.particles)
    # naming them projectile or target is therefore also arbitrary
    cp1: Particle  # projectile
    cp2: Particle  # target
    cp1, cp2 = collided

    # just called the more massive one the main particle to keep its type/name
    # Sun<->Protoplanet -> Sun
    main_particle = cp1.m if cp1.m > cp2.m else cp2

    projectile_wmf = ed.pd(cp1).water_mass_fraction
    target_wmf = ed.pd(cp2).water_mass_fraction

    # get the velocities, velocity differences and unit vector as numpy arrays
    # all units are in sytem units (so AU/year)
    v1 = np.array(cp1.vxyz)
    v2 = np.array(cp2.vxyz)
    vdiff = linalg.norm(v2 - v1)
    v1_u = v1 / linalg.norm(v1)
    v2_u = v2 / linalg.norm(v2)
    # get angle between the two velocities as degrees
    # https://stackoverflow.com/a/13849249/4398037
    ang = np.degrees(np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0)))

    # get mass fraction
    # if it is >1 it will be inverted during interpolation
    gamma = cp1.m / cp2.m

    # calculate mutual escape velocity (for norming the velocities in the interpolation)
    escape_velocity = sqrt(2 * G * (cp1.m + cp2.m) / ((cp1.r + cp2.r) * astronomical_unit))

    print("interpolating")

    # let interpolation calculate water and mass retention fraction
    # meta is just a bunch of intermediate results that will be logged to help
    # understand the collisions better
    water_ret, stone_ret, meta = 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("interpolation finished")
    print(water_ret, stone_ret)

    hash = unique_hash()  # hash for newly created particle

    # handle loss of water and core mass
    water_mass = cp1.m * projectile_wmf + cp2.m * target_wmf
    stone_mass = cp1.m + cp2.m - water_mass

    water_mass *= water_ret
    stone_mass *= stone_ret

    total_mass = water_mass + stone_mass
    final_wmf = water_mass / total_mass
    print(final_wmf)
    # create new object preserving momentum
    merged_planet = (cp1 * cp1.m + cp2 * cp2.m) / total_mass
    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,
        type=ed.pd(main_particle).type
    )

    meta["total_mass"] = total_mass
    meta["final_wmf"] = final_wmf
    meta["final_radius"] = merged_planet.r
    meta["target_wmf"] = target_wmf
    meta["projectile_wmf"] = projectile_wmf
    meta["time"] = sim.t
    ed.tree.add(cp1, cp2, merged_planet, meta)

    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)

    sim.move_to_com()
    sim.ri_whfast.recalculate_coordinates_this_timestep = 1
    sim.integrator_synchronize()


def handle_escape(sim: Simulation, ed: ExtraData):
    h = 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:
            h = p.hash
    if not h:
        raise RuntimeError("Escape without escaping particle")
    sim.remove(hash=h)

    sim.move_to_com()
    sim.ri_whfast.recalculate_coordinates_this_timestep = 1
    sim.integrator_synchronize()