from copy import copy from pprint import pprint from typing import Tuple, Optional import numpy as np from numpy import linalg, sqrt from rebound import Simulation, Particle from rebound.simulation import POINTER_REB_SIM, reb_collision from scipy.constants import astronomical_unit, G from extradata import ExtraData, ParticleData, CollisionMeta, Input from massloss import RbfMassloss, Massloss, LeiZhouMassloss, SimpleNNMassloss from massloss.perfect_merging import PerfectMerging from utils import unique_hash, clamp, PlanetaryRadius massloss_estimator: Optional[Massloss] = None # global waterloss estimator cache def get_mass_fractions(input_data: Input) -> Tuple[float, float, float, CollisionMeta]: global massloss_estimator print("v_esc", input_data.escape_velocity) print("v_orig,v_si", input_data.velocity_original, input_data.velocity_si) print("v/v_esc", input_data.velocity_esc) data = copy(input_data) if data.gamma > 1: data.gamma = 1 / data.gamma data.alpha = clamp(data.alpha, 0, 60) data.velocity_esc = clamp(data.velocity_esc, 1, 5) m_ceres = 9.393e+20 m_earth = 5.9722e+24 data.projectile_mass = clamp(data.projectile_mass, 2 * m_ceres, 2 * m_earth) data.gamma = clamp(data.gamma, 1 / 10, 1) water_retention, mantle_retention, core_retention = \ massloss_estimator.estimate(data.alpha, data.velocity_esc, data.projectile_mass, data.gamma, ) metadata = CollisionMeta() metadata.interpolation_input = [data.alpha, data.velocity_esc, data.projectile_mass, data.gamma] metadata.input = input_data metadata.adjusted_input = data metadata.raw_water_retention = water_retention metadata.raw_mantle_retention = mantle_retention metadata.raw_core_retention = core_retention water_retention = clamp(water_retention, 0, 1) mantle_retention = clamp(mantle_retention, 0, 1) core_retention = clamp(core_retention, 0, 1) metadata.water_retention = water_retention metadata.mantle_retention = mantle_retention metadata.core_retention = core_retention return water_retention, mantle_retention, core_retention, metadata def merge_particles(sim_p: POINTER_REB_SIM, collision: reb_collision, ed: ExtraData): global massloss_estimator print("--------------") print("colliding") sim: Simulation = sim_p.contents print("current time step", sim.dt) print(f"p1 is {collision.p1}") print(f"p2 is {collision.p2}") # the assignment to cp1 or cp2 seems to be random # also look at a copy instead of the original particles # to avoid issues after they have been modified cp1: Particle = sim.particles[collision.p1].copy() cp2: Particle = sim.particles[collision.p2].copy() # just calling the more massive one the target to keep its type/name # Sun<->Protoplanet -> Sun # and to keep collsions mostly reproducable if cp1.m > cp2.m: target = cp1 projectile = cp2 else: # also when masses are the same target = cp2 projectile = cp1 if collision.p1 > collision.p2: lower_index_particle_index = collision.p2 else: lower_index_particle_index = collision.p1 print(f"colliding {target.hash.value} ({ed.pd(target).type}) " f"with {projectile.hash.value} ({ed.pd(projectile).type})") projectile_wmf = ed.pd(projectile).water_mass_fraction projectile_cmf = ed.pd(projectile).core_mass_fraction target_wmf = ed.pd(target).water_mass_fraction target_cmf = ed.pd(target).core_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(target.vxyz) v2 = np.array(projectile.vxyz) r1 = np.array(target.xyz) r2 = np.array(projectile.xyz) vdiff = v2 - v1 rdiff = r2 - r1 vdiff_n = linalg.norm(vdiff) rdiff_n = linalg.norm(rdiff) print("dt", sim.dt) # during a collision ias15 should always be used, otherwise something weird has happend assert sim.ri_mercurius.mode == 1 print("rdiff", rdiff) print("vdiff", vdiff) print("sum_radii", target.r + projectile.r) print("rdiff_n", rdiff_n) print("vdiff_n", vdiff_n) ang = float(np.degrees(np.arccos(np.dot(rdiff, vdiff) / (rdiff_n * vdiff_n)))) if ang > 90: ang = 180 - ang print("angle_deg", ang) print() # get mass fraction gamma = projectile.m / target.m # calculate mutual escape velocity (for norming the velocities in the interpolation) in SI units escape_velocity = sqrt(2 * G * (target.m + projectile.m) / ((target.r + projectile.r) * astronomical_unit)) print("interpolating") if not massloss_estimator: methods = [RbfMassloss, LeiZhouMassloss, PerfectMerging, SimpleNNMassloss] per_name = {} for method in methods: per_name[method.name] = method try: estimator_class = per_name[ed.meta.massloss_method] except KeyError: print("invalid mass loss estimation method") print("please use one of these:") print(per_name) raise massloss_estimator = estimator_class() # 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 input_data = Input( alpha=ang, velocity_original=vdiff_n, escape_velocity=escape_velocity, gamma=gamma, projectile_mass=projectile.m, target_water_fraction=target_wmf, projectile_water_fraction=projectile_wmf, ) water_ret, mantle_ret, core_ret, meta = get_mass_fractions(input_data) print("mass retentions:", water_ret, mantle_ret, core_ret) meta.collision_velocities = (v1.tolist(), v2.tolist()) meta.collision_positions = (target.xyz, projectile.xyz) meta.collision_radii = (target.r, projectile.r) hash = unique_hash(ed) # hash for newly created particle # handle loss of water and core mass water_mass = target.m * target_wmf + projectile.m * projectile_wmf core_mass = target.m * target_cmf + projectile.m * projectile_cmf mantle_mass = target.m + projectile.m - water_mass - core_mass water_mass *= water_ret mantle_mass *= mantle_ret core_mass *= core_ret total_mass = water_mass + mantle_mass + core_mass final_wmf = water_mass / total_mass final_cmf = core_mass / total_mass print(final_wmf) # create new object preserving momentum merged_planet = (target * target.m + projectile * projectile.m) / total_mass merged_planet.m = total_mass merged_planet.hash = hash merged_planet.r = PlanetaryRadius(merged_planet.m, final_wmf, final_cmf).total_radius / astronomical_unit ed.pdata[hash.value] = ParticleData( water_mass_fraction=final_wmf, core_mass_fraction=final_cmf, type=ed.pd(target).type, 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 pprint(meta) ed.tree.add(target, projectile, merged_planet, meta) sim.particles[lower_index_particle_index] = merged_planet sim.move_to_com() sim.integrator_synchronize() sim.ri_mercurius.recalculate_coordinates_this_timestep = 1 sim.ri_mercurius.recalculate_dcrit_this_timestep = 1 print("collision finished") print("--------------") # from rebound docs: # A return value of 0 indicates that both particles remain in the simulation. # A return value of 1 (2) indicates that particle 1 (2) should be removed from the simulation. # A return value of 3 indicates that both particles should be removed from the simulation. # always keep lower index particle and delete other one # this keeps the N_active working if lower_index_particle_index == collision.p1: print("deleting p2") return 2 elif lower_index_particle_index == collision.p2: print("deleting p1") return 1 else: raise ValueError("invalid index")