mirror of
https://github.com/Findus23/rebound-collisions.git
synced 2024-09-19 15:53:48 +02:00
better collision statistics
This commit is contained in:
parent
f7663a53b3
commit
d71a91343e
4 changed files with 122 additions and 0 deletions
45
collisionstats.py
Normal file
45
collisionstats.py
Normal file
|
@ -0,0 +1,45 @@
|
|||
from matplotlib import pyplot as plt
|
||||
from scipy.constants import mega
|
||||
|
||||
from extradata import ExtraData, CollisionMeta
|
||||
from utils import filename_from_argv, create_figure
|
||||
|
||||
fn = filename_from_argv()
|
||||
ed = ExtraData.load(fn.with_suffix(".extra.json"))
|
||||
|
||||
vs = []
|
||||
angles = []
|
||||
times = []
|
||||
|
||||
for collision in ed.tree.get_tree().values():
|
||||
meta: CollisionMeta = collision["meta"]
|
||||
vs.append(meta.input.velocity_esc)
|
||||
angles.append(meta.input.alpha)
|
||||
times.append(meta.time / mega)
|
||||
|
||||
angle_label = "Impact angle (deg)"
|
||||
v_label = "v/v_esc"
|
||||
time_label = "Time (Myr)"
|
||||
|
||||
fig1, ax1 = create_figure()
|
||||
|
||||
ax1.scatter(angles, vs)
|
||||
ax1.set_xlabel(angle_label)
|
||||
ax1.set_ylabel(v_label)
|
||||
|
||||
fig2, ax2 = create_figure()
|
||||
|
||||
ax2.scatter(times, angles)
|
||||
ax2.set_xlabel(time_label)
|
||||
ax2.set_ylabel(angle_label)
|
||||
ax2.set_xscale("log")
|
||||
|
||||
fig3, ax3 = create_figure()
|
||||
|
||||
ax3.scatter(times, vs)
|
||||
ax3.set_xlabel(time_label)
|
||||
ax3.set_ylabel(v_label)
|
||||
ax3.set_xscale("log")
|
||||
|
||||
|
||||
plt.show()
|
|
@ -37,6 +37,8 @@ class Input:
|
|||
@dataclass
|
||||
class CollisionMeta:
|
||||
collision_velocities: Tuple[List[float], List[float]] = None
|
||||
collision_positions: Tuple[List[float], List[float]] = None
|
||||
collision_radii: Tuple[float, float] = None
|
||||
interpolation_input: List[float] = None
|
||||
raw_water_retention: float = None
|
||||
raw_mass_retention: float = None
|
||||
|
|
7
merge.py
7
merge.py
|
@ -78,9 +78,12 @@ def get_mass_fractions(input_data: Input) -> Tuple[float, float, CollisionMeta]:
|
|||
|
||||
|
||||
def merge_particles(sim_p: POINTER_REB_SIM, collision: reb_collision, ed: ExtraData):
|
||||
print("--------------")
|
||||
print("colliding")
|
||||
sim: Simulation = sim_p.contents
|
||||
collided: List[Particle] = []
|
||||
print("current time step", sim.dt, )
|
||||
print("mode", sim.ri_mercurius.mode)
|
||||
|
||||
# the assignment to cp1 or cp2 is mostly random
|
||||
# (cp1 is the one with a lower index in sim.particles)
|
||||
|
@ -151,6 +154,8 @@ def merge_particles(sim_p: POINTER_REB_SIM, collision: reb_collision, ed: ExtraD
|
|||
print(water_ret, stone_ret)
|
||||
|
||||
meta.collision_velocities = (v1.tolist(), v2.tolist())
|
||||
meta.collision_positions = (cp1.xyz, cp2.xyz)
|
||||
meta.collision_radii = (cp1.r, cp2.r)
|
||||
|
||||
hash = unique_hash() # hash for newly created particle
|
||||
|
||||
|
@ -193,6 +198,8 @@ def merge_particles(sim_p: POINTER_REB_SIM, collision: reb_collision, ed: ExtraD
|
|||
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.
|
||||
|
|
68
visualize_collision.py
Normal file
68
visualize_collision.py
Normal file
|
@ -0,0 +1,68 @@
|
|||
import numpy as np
|
||||
from matplotlib import pyplot as plt
|
||||
from matplotlib.figure import Figure
|
||||
from matplotlib.patches import FancyArrowPatch
|
||||
from mpl_toolkits.mplot3d import Axes3D, proj3d
|
||||
|
||||
from extradata import ExtraData, CollisionMeta
|
||||
from utils import filename_from_argv
|
||||
|
||||
fn = filename_from_argv()
|
||||
ed = ExtraData.load(fn.with_suffix(".extra.json"))
|
||||
|
||||
|
||||
def get_circle(dx, dy, dz, r):
|
||||
u, v = np.mgrid[0:2 * np.pi:40j, 0:np.pi:20j]
|
||||
x = r * np.cos(u) * np.sin(v) + dx
|
||||
y = r * np.sin(u) * np.sin(v) + dy
|
||||
z = r * np.cos(v) + dz
|
||||
return x, y, z
|
||||
|
||||
|
||||
class Arrow3D(FancyArrowPatch):
|
||||
"""
|
||||
https://stackoverflow.com/a/11156353/4398037
|
||||
"""
|
||||
|
||||
def __init__(self, xs, ys, zs, *args, **kwargs):
|
||||
FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs)
|
||||
self._verts3d = xs, ys, zs
|
||||
|
||||
def draw(self, renderer):
|
||||
xs3d, ys3d, zs3d = self._verts3d
|
||||
xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
|
||||
self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
|
||||
FancyArrowPatch.draw(self, renderer)
|
||||
|
||||
|
||||
for collision in ed.tree.get_tree().values():
|
||||
meta: CollisionMeta = collision["meta"]
|
||||
|
||||
fig: Figure = plt.figure()
|
||||
ax: Axes3D = fig.gca(projection='3d')
|
||||
|
||||
i = 0
|
||||
for pos, vel, r in zip(meta.collision_positions, meta.collision_velocities, meta.collision_radii):
|
||||
print(pos, vel, r)
|
||||
circle_coords = get_circle(*pos, r)
|
||||
ax.plot_wireframe(*circle_coords, color=f"C{i}")
|
||||
a = Arrow3D(*[(p, p + v / 100000) for p, v in zip(pos, vel)], mutation_scale=20,
|
||||
lw=1, arrowstyle="-|>", color="k")
|
||||
ax.add_artist(a)
|
||||
i += 1
|
||||
ax.set_title(f"angle={meta.input.alpha:.2f}, v/v_esc={meta.input.velocity_esc:.2f}")
|
||||
xyzlim = np.array([ax.get_xlim3d(), ax.get_ylim3d(), ax.get_zlim3d()]).T
|
||||
diff = min(xyzlim[0] - xyzlim[1])
|
||||
|
||||
print(xyzlim)
|
||||
print(diff)
|
||||
xmin, _ = ax.get_xlim3d()
|
||||
ax.set_xlim3d((xmin, xmin - diff))
|
||||
ymin, _ = ax.get_ylim3d()
|
||||
ax.set_ylim3d((ymin, ymin - diff))
|
||||
zmin, _ = ax.get_zlim3d()
|
||||
ax.set_zlim3d((zmin, zmin - diff))
|
||||
# ax.set_ylim3d(XYZlim)
|
||||
# ax.set_zlim3d(XYZlim * 3/4)
|
||||
|
||||
plt.show()
|
Loading…
Reference in a new issue