From d71a91343e60c1012622a27fba63be669bb45f7c Mon Sep 17 00:00:00 2001
From: Lukas Winkler <git@lw1.at>
Date: Thu, 21 Jan 2021 16:45:31 +0100
Subject: [PATCH] better collision statistics

---
 collisionstats.py      | 45 ++++++++++++++++++++++++++++
 extradata.py           |  2 ++
 merge.py               |  7 +++++
 visualize_collision.py | 68 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 122 insertions(+)
 create mode 100644 collisionstats.py
 create mode 100644 visualize_collision.py

diff --git a/collisionstats.py b/collisionstats.py
new file mode 100644
index 0000000..8d00924
--- /dev/null
+++ b/collisionstats.py
@@ -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()
diff --git a/extradata.py b/extradata.py
index c013c8e..bb2cb26 100644
--- a/extradata.py
+++ b/extradata.py
@@ -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
diff --git a/merge.py b/merge.py
index 9e61729..7171a8e 100644
--- a/merge.py
+++ b/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.
diff --git a/visualize_collision.py b/visualize_collision.py
new file mode 100644
index 0000000..342d8db
--- /dev/null
+++ b/visualize_collision.py
@@ -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()