Multisize Ball-Ball Collision

pooltool/physics/resolve/ball_ball/frictional_inelastic/__init__.py

@@ -1,5 +1,6 @@
 import attrs
 import numpy as np
+import quaternion
 from numba import jit
 
 import pooltool.constants as const
@@ -13,23 +14,28 @@
 from pooltool.physics.resolve.models import BallBallModel
 
 
-@jit(nopython=True, cache=const.use_numba_cache)
-def _resolve_ball_ball(rvw1, rvw2, R, u_b, e_b):
+def _resolve_ball_ball(rvw1, m1, R1, rvw2, m2, R2, u_b, e_b):
     unit_x = np.array([1.0, 0.0, 0.0])
-
-    # rotate the x-axis to be in line with the line of centers
     delta_centers = rvw2[0] - rvw1[0]
-    # FIXME3D: this should use quaternion rotation in 3D
-    theta = ptmath.angle(delta_centers, unit_x)
-    rvw1[1] = ptmath.coordinate_rotation(rvw1[1], -theta)
-    rvw1[2] = ptmath.coordinate_rotation(rvw1[2], -theta)
-    rvw2[1] = ptmath.coordinate_rotation(rvw2[1], -theta)
-    rvw2[2] = ptmath.coordinate_rotation(rvw2[2], -theta)
+    frame_rotation = ptmath.quaternion_from_vector_to_vector(delta_centers, unit_x)
+    rvw1 = quaternion.rotate_vectors(frame_rotation, rvw1)
+    rvw2 = quaternion.rotate_vectors(frame_rotation, rvw2)
+    rvw1, rvw2 = _resolve_ball_ball_x_normal(rvw1, m1, R1, rvw2, m2, R2, u_b, e_b)
+    rvw1 = quaternion.rotate_vectors(frame_rotation.conjugate(), rvw1)
+    rvw2 = quaternion.rotate_vectors(frame_rotation.conjugate(), rvw2)
+    return rvw1, rvw2
+
+
+@jit(nopython=True, cache=const.use_numba_cache)
+def _resolve_ball_ball_x_normal(rvw1, m1, R1, rvw2, m2, R2, u_b, e_b):
+    unit_x = np.array([1.0, 0.0, 0.0])
 
     # velocity normal component, same for both slip and no-slip after collison cases
-    v1_n_f = 0.5 * ((1.0 - e_b) * rvw1[1][0] + (1.0 + e_b) * rvw2[1][0])
-    v2_n_f = 0.5 * ((1.0 + e_b) * rvw1[1][0] + (1.0 - e_b) * rvw2[1][0])
-    D_v_n_magnitude = abs(v2_n_f - v1_n_f)
+    v_12_n = rvw1[1][0] - rvw2[1][0]
+    D_v1_n = -(1 + e_b) / (1 + m1 / m2) * v_12_n
+    D_v2_n = -(m1 / m2) * D_v1_n
+    v1_n_f = rvw1[1][0] + D_v1_n
+    v2_n_f = rvw2[1][0] + D_v2_n
 
     # angular velocity normal component, unchanged
     w1_n_f = rvw1[2][0]
@@ -43,8 +49,8 @@ def _resolve_ball_ball(rvw1, rvw2, R, u_b, e_b):
     rvw1_f = rvw1.copy()
     rvw2_f = rvw2.copy()
 
-    v1_c = ptmath.surface_velocity(rvw1, unit_x, R)
-    v2_c = ptmath.surface_velocity(rvw2, -unit_x, R)
+    v1_c = ptmath.surface_velocity(rvw1, unit_x, R1)
+    v2_c = ptmath.surface_velocity(rvw2, -unit_x, R2)
     v12_c = v1_c - v2_c
     has_relative_velocity = ptmath.norm3d(v12_c) > const.EPS
 
@@ -53,61 +59,50 @@ def _resolve_ball_ball(rvw1, rvw2, R, u_b, e_b):
     if has_relative_velocity:
         # tangent components for slip condition
         v12_c_hat = ptmath.unit_vector(v12_c)
-        D_v1_t = u_b * D_v_n_magnitude * -v12_c_hat
-        D_w1 = 2.5 / R * ptmath.cross(unit_x, D_v1_t)
+        D_v1_t = u_b * abs(D_v1_n) * -v12_c_hat
+        D_w1 = 2.5 / R1 * ptmath.cross(unit_x, D_v1_t)
+        D_v2_t = -(m1 / m2) * D_v1_t
+        D_w2 = 2.5 / R2 * ptmath.cross(-unit_x, D_v2_t)
         rvw1_f[1] = rvw1[1] + D_v1_t
         rvw1_f[2] = rvw1[2] + D_w1
-        rvw2_f[1] = rvw2[1] - D_v1_t
-        rvw2_f[2] = rvw2[2] + D_w1
+        rvw2_f[1] = rvw2[1] + D_v2_t
+        rvw2_f[2] = rvw2[2] + D_w2
 
         # calculate new relative contact velocity
-        v1_c_slip = ptmath.surface_velocity(rvw1_f, unit_x, R)
-        v2_c_slip = ptmath.surface_velocity(rvw2_f, -unit_x, R)
+        v1_c_slip = ptmath.surface_velocity(rvw1_f, unit_x, R1)
+        v2_c_slip = ptmath.surface_velocity(rvw2_f, -unit_x, R2)
         v12_c_slip = v1_c_slip - v2_c_slip
 
     # if there was no relative velocity to begin with, or if slip changed directions,
     # then slip condition is invalid so we need to calculate no-slip condition
     if not has_relative_velocity or np.dot(v12_c, v12_c_slip) <= 0:  # type: ignore
         # velocity tangent component for no-slip condition
-        D_v1_t = -(1.0 / 7.0) * (
-            rvw1[1] - rvw2[1] + R * ptmath.cross(rvw1[2] + rvw2[2], unit_x)
-        )
-        D_w1 = -(5.0 / 14.0) * (
-            ptmath.cross(unit_x, rvw1[1] - rvw2[1]) / R + rvw1[2] + rvw2[2]
-        )
+        D_v1_t = -(2.0 / 7.0) * v12_c / (1 + m1 / m2)
+        D_w1 = 2.5 / R1 * ptmath.cross(unit_x, D_v1_t)
+        D_v2_t = -(m1 / m2) * D_v1_t
+        D_w2 = 2.5 / R2 * ptmath.cross(-unit_x, D_v2_t)
         rvw1_f[1] = rvw1[1] + D_v1_t
         rvw1_f[2] = rvw1[2] + D_w1
-        rvw2_f[1] = rvw2[1] - D_v1_t
-        rvw2_f[2] = rvw2[2] + D_w1
+        rvw2_f[1] = rvw2[1] + D_v2_t
+        rvw2_f[2] = rvw2[2] + D_w2
 
     # reintroduce the final normal components
     rvw1_f[1][0] = v1_n_f
     rvw2_f[1][0] = v2_n_f
     rvw1_f[2][0] = w1_n_f
     rvw2_f[2][0] = w2_n_f
 
-    # rotate everything back to the original frame
-    rvw1_f[1] = ptmath.coordinate_rotation(rvw1_f[1], theta)
-    rvw1_f[2] = ptmath.coordinate_rotation(rvw1_f[2], theta)
-    rvw2_f[1] = ptmath.coordinate_rotation(rvw2_f[1], theta)
-    rvw2_f[2] = ptmath.coordinate_rotation(rvw2_f[2], theta)
-
-    # FIXME3D: include z velocity components
-    # remove any z velocity components from spin-induced throw
-    rvw1_f[1][2] = 0.0
-    rvw2_f[1][2] = 0.0
-
     return rvw1_f, rvw2_f
 
 
 @attrs.define
 class FrictionalInelastic(CoreBallBallCollision):
-    """A simple ball-ball collision model including ball-ball friction, and coefficient of restitution for equal-mass balls
+    """A simple ball-ball collision model including ball-ball friction, and coefficient of restitution
 
     Largely inspired by Dr. David Alciatore's technical proofs
-    (https://billiards.colostate.edu/technical_proofs), in particular, TP_A-5, TP_A-6,
-    and TP_A-14. These ideas have been extended to include motion of both balls, and a
-    more complete analysis of velocity and angular velocity in their vector forms.
+    (https://billiards.colostate.edu/technical_proofs), in particular, TP_A-5, TP_A-6, and TP_A-14.
+    These ideas have been extended to include motion of both balls, balls of different size and mass,
+    and a more complete analysis of velocity and angular velocity in their vector forms.
     """
 
     friction: BallBallFrictionStrategy = AlciatoreBallBallFriction()
@@ -120,13 +115,20 @@ def solve(self, ball1: Ball, ball2: Ball) -> tuple[Ball, Ball]:
         """Resolves the collision."""
         rvw1, rvw2 = _resolve_ball_ball(
             ball1.state.rvw.copy(),
-            ball2.state.rvw.copy(),
+            ball1.params.m,
             ball1.params.R,
+            ball2.state.rvw.copy(),
+            ball2.params.m,
+            ball2.params.R,
             u_b=self.friction.calculate_friction(ball1, ball2),
             # Average the coefficient of restitution parameters for the two balls
             e_b=(ball1.params.e_b + ball2.params.e_b) / 2,
         )
 
+        # FIXME3D: include z velocity components
+        rvw1[1][2] = 0.0
+        rvw2[1][2] = 0.0
+
         ball1.state = BallState(rvw1, const.sliding)
         ball2.state = BallState(rvw2, const.sliding)
 

pooltool/physics/resolve/ball_ball/frictionless_elastic/__init__.py

@@ -1,5 +1,6 @@
 import attrs
 import numpy as np
+import quaternion
 
 import pooltool.constants as const
 import pooltool.ptmath as ptmath
@@ -8,33 +9,39 @@
 from pooltool.physics.resolve.models import BallBallModel
 
 
-def _resolve_ball_ball(rvw1, rvw2, R):
-    r1, r2 = rvw1[0], rvw2[0]
-    v1, v2 = rvw1[1], rvw2[1]
-
-    n = ptmath.unit_vector(r2 - r1)
-    t = ptmath.coordinate_rotation(n, np.pi / 2)
-
-    v_rel = v1 - v2
-    v_mag = ptmath.norm3d(v_rel)
-
-    beta = ptmath.angle(v_rel, n)
+def _resolve_ball_ball(rvw1, m1, rvw2, m2):
+    unit_x = np.array([1.0, 0.0, 0.0])
+    delta_centers = rvw2[0] - rvw1[0]
+    frame_rotation = ptmath.quaternion_from_vector_to_vector(delta_centers, unit_x)
+    v1 = quaternion.rotate_vectors(frame_rotation, rvw1[1])
+    v2 = quaternion.rotate_vectors(frame_rotation, rvw2[1])
+    v1, v2 = _resolve_ball_ball_x_normal(v1, m1, v2, m2)
+    v1 = quaternion.rotate_vectors(frame_rotation.conjugate(), v1)
+    v2 = quaternion.rotate_vectors(frame_rotation.conjugate(), v2)
+    rvw1[1] = v1
+    rvw2[1] = v2
+    return rvw1, rvw2
 
-    rvw1[1] = t * v_mag * np.sin(beta) + v2
-    rvw2[1] = n * v_mag * np.cos(beta) + v2
 
-    return rvw1, rvw2
+def _resolve_ball_ball_x_normal(v1, m1, v2, m2):
+    v_12_n = v1[0] - v2[0]
+    D_v1_n = -2.0 / (1 + m1 / m2) * v_12_n
+    D_v2_n = -(m1 / m2) * D_v1_n
+    v1[0] += D_v1_n
+    v2[0] += D_v2_n
+    return v1, v2
 
 
 @attrs.define
 class FrictionlessElastic(CoreBallBallCollision):
-    """A frictionless, instantaneous, elastic, equal mass collision resolver.
+    """A frictionless, instantaneous, elastic collision resolver.
 
     This is as simple as it gets.
 
     See Also:
         - This physics of this model is blogged about at
-          https://ekiefl.github.io/2020/04/24/pooltool-theory/#1-elastic-instantaneous-frictionless
+          https://ekiefl.github.io/2020/04/24/pooltool-theory/#1-elastic-instantaneous-frictionless.
+          It's since been modified to handle balls of unequal mass and size.
     """
 
     model: BallBallModel = attrs.field(
@@ -45,11 +52,16 @@ def solve(self, ball1: Ball, ball2: Ball) -> tuple[Ball, Ball]:
         """Resolves the collision."""
         rvw1, rvw2 = _resolve_ball_ball(
             ball1.state.rvw.copy(),
+            ball1.params.m,
             ball2.state.rvw.copy(),
-            ball1.params.R,
+            ball2.params.m,
         )
 
         ball1.state = BallState(rvw1, const.sliding)
         ball2.state = BallState(rvw2, const.sliding)
 
+        # FIXME3D: include z velocity components
+        rvw1[1][2] = 0.0
+        rvw2[1][2] = 0.0
+
         return ball1, ball2