3D Ball-Ball Resolution

pooltool/evolution/event_based/detect/ball_ball.py

@@ -18,10 +18,30 @@
 )
 from pooltool.physics.utils import get_u_vec
 from pooltool.ptmath.roots import quadratic, quartic
-from pooltool.ptmath.roots.core import get_real_positive_smallest_root
+from pooltool.ptmath.roots.core import (
+    get_real_positive_smallest_root,
+    is_real_number,
+)
 from pooltool.system.datatypes import Ball, System
 
 
+# @jit(nopython=True, cache=const.use_numba_cache)
+def select_ball_ball_collision_root(
+    sorted_real_positive_roots: NDArray[np.float64], p12: NDArray[np.float64]
+):
+    """Smallest positive real root for which the two balls are moving towards each other"""
+
+    v12: NDArray[np.float64] = np.array([p12[1], 0.5 * p12[2]])
+
+    for t in sorted_real_positive_roots:
+        p12_collision = p12[0] + p12[1] * t + p12[2] * t * t
+        v12_collision = v12[0] + v12[1] * t
+        if np.dot(p12_collision, v12_collision) > 0:
+            continue
+        return t
+    return np.inf
+
+
 def ball_ball_collision_time(
     ball1: Ball,
     ball2: Ball,
@@ -54,9 +74,15 @@ def ball_ball_collision_time(
     if C[4] == 0.0:
         # C[3] must also be 0.0, and this is a quadratic
         assert C[3] == 0.0
-        return get_real_positive_smallest_root(quadratic.solve(C[2], C[1], C[0]))
+        roots = quadratic.solve(C[2], C[1], C[0])
+    else:
+        roots = quartic.solve(C[4], C[3], C[2], C[1], C[0])
+
+    sorted_real_positive_roots = np.array(
+        sorted(root.real for root in roots if is_real_number(root) and root.real > 0)
+    )
 
-    return get_real_positive_smallest_root(quartic.solve(C[4], C[3], C[2], C[1], C[0]))
+    return select_ball_ball_collision_root(sorted_real_positive_roots, p12)
 
 
 @jit(nopython=True, cache=const.use_numba_cache)

pooltool/physics/resolve/ball_ball/__init__.py

@@ -5,15 +5,19 @@
 import attrs
 
 from pooltool.physics.resolve.ball_ball.core import BallBallCollisionStrategy
-from pooltool.physics.resolve.ball_ball.frictional_inelastic import FrictionalInelastic
+from pooltool.physics.resolve.ball_ball.frictional_inelastic import (
+    FrictionalInelastic2D,
+    FrictionalInelastic3D,
+)
 from pooltool.physics.resolve.ball_ball.frictional_mathavan import FrictionalMathavan
 from pooltool.physics.resolve.ball_ball.frictionless_elastic import FrictionlessElastic
 from pooltool.physics.resolve.models import BallBallModel
 
 _ball_ball_model_registry: tuple[type[BallBallCollisionStrategy], ...] = (
     FrictionlessElastic,
     FrictionalMathavan,
-    FrictionalInelastic,
+    FrictionalInelastic2D,
+    FrictionalInelastic3D,
 )
 
 ball_ball_models: dict[BallBallModel, type[BallBallCollisionStrategy]] = {

pooltool/physics/resolve/ball_ball/core.py

@@ -2,13 +2,31 @@
 from typing import Protocol
 
 import numpy as np
+from numpy.typing import NDArray
 
 import pooltool.constants as const
 import pooltool.ptmath as ptmath
 from pooltool.objects.ball.datatypes import Ball
 from pooltool.physics.dimensionality import Dim
 
 
+# stolen from stick_ball/core.py
+# TODO: move to common place
+def final_ball_motion_state(rvw: NDArray[np.float64], R: float) -> int:
+    """Return the final (post-strike) motion state label.
+
+    If the z-velocity is non-zero the ball is considered airborne, otherwise
+    it is sliding (a struck ball is always kinetic).
+
+    Notes:
+        - A universal ``final_ball_motion_state`` fn could be a good idea.
+    """
+    if rvw[1, 2] != 0.0:
+        return const.airborne
+
+    return const.sliding
+
+
 class _BaseStrategy(Protocol):
     def make_kiss(self, ball1: Ball, ball2: Ball) -> tuple[Ball, Ball]: ...
 

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

@@ -1,12 +1,16 @@
 import attrs
 import numpy as np
+import quaternion
 from numba import jit
 
 import pooltool.constants as const
 import pooltool.ptmath as ptmath
-from pooltool.objects.ball.datatypes import Ball, BallState
+from pooltool.objects.ball.datatypes import Ball
 from pooltool.physics.dimensionality import Dim
-from pooltool.physics.resolve.ball_ball.core import CoreBallBallCollision
+from pooltool.physics.resolve.ball_ball.core import (
+    CoreBallBallCollision,
+    final_ball_motion_state,
+)
 from pooltool.physics.resolve.ball_ball.friction import (
     AlciatoreBallBallFriction,
     BallBallFrictionStrategy,
@@ -15,18 +19,21 @@
 from pooltool.physics.utils import surface_velocity
 
 
-@jit(nopython=True, cache=const.use_numba_cache)
 def _resolve_ball_ball(rvw1, rvw2, R, 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, rvw2, R, 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, rvw2, R, 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])
@@ -88,22 +95,11 @@ def _resolve_ball_ball(rvw1, rvw2, R, u_b, e_b):
     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):
+class FrictionalInelastic3D(CoreBallBallCollision):
     """A simple ball-ball collision model including ball-ball friction, and coefficient of restitution for equal-mass balls
 
     Largely inspired by Dr. David Alciatore's technical proofs
@@ -115,9 +111,9 @@ class FrictionalInelastic(CoreBallBallCollision):
     friction: BallBallFrictionStrategy = AlciatoreBallBallFriction()
 
     model: BallBallModel = attrs.field(
-        default=BallBallModel.FRICTIONAL_INELASTIC, init=False, repr=False
+        default=BallBallModel.FRICTIONAL_INELASTIC_3D, init=False, repr=False
     )
-    dim: Dim = attrs.field(default=Dim.TWO, init=False, repr=False)
+    dim: Dim = attrs.field(default=Dim.THREE, init=False, repr=False)
 
     def solve(self, ball1: Ball, ball2: Ball) -> tuple[Ball, Ball]:
         """Resolves the collision."""
@@ -130,7 +126,35 @@ def solve(self, ball1: Ball, ball2: Ball) -> tuple[Ball, Ball]:
             e_b=(ball1.params.e_b + ball2.params.e_b) / 2,
         )
 
-        ball1.state = BallState(rvw1, const.sliding)
-        ball2.state = BallState(rvw2, const.sliding)
+        ball1.state.rvw = rvw1
+        ball2.state.rvw = rvw2
+
+        ball1.state.s = final_ball_motion_state(rvw1, ball1.params.R)
+        ball2.state.s = final_ball_motion_state(rvw2, ball2.params.R)
+
+        return ball1, ball2
+
+
+@attrs.define
+class FrictionalInelastic2D(FrictionalInelastic3D):
+    """A simple ball-ball collision model including ball-ball friction, and coefficient of restitution for equal-mass balls
+
+    For details see :class:`FrictionalInelastic3D`.
+    """
+
+    model: BallBallModel = attrs.field(
+        default=BallBallModel.FRICTIONAL_INELASTIC_2D, init=False, repr=False
+    )
+    dim: Dim = attrs.field(default=Dim.TWO, init=False, repr=False)
+
+    def solve(self, ball1: Ball, ball2: Ball) -> tuple[Ball, Ball]:
+        """Resolves the collision."""
+        ball1, ball2 = super().solve(ball1, ball2)
+
+        # remove any z velocity components for 2D
+        ball1.state.rvw[1, 2] = 0.0
+        ball1.state.rvw[1, 2] = 0.0
+        ball1.state.s = const.sliding
+        ball2.state.s = const.sliding
 
         return ball1, ball2

pooltool/physics/resolve/models.py

@@ -44,7 +44,8 @@ class BallBallModel(StrEnum):
     """
 
     FRICTIONLESS_ELASTIC = auto()
-    FRICTIONAL_INELASTIC = auto()
+    FRICTIONAL_INELASTIC_2D = auto()
+    FRICTIONAL_INELASTIC_3D = auto()
     FRICTIONAL_MATHAVAN = auto()
 
 

pooltool/physics/resolve/resolver.py

@@ -17,7 +17,9 @@
 from pooltool.physics.resolve.ball_ball.friction import (
     AlciatoreBallBallFriction,
 )
-from pooltool.physics.resolve.ball_ball.frictional_inelastic import FrictionalInelastic
+from pooltool.physics.resolve.ball_ball.frictional_inelastic import (
+    FrictionalInelastic2D,
+)
 from pooltool.physics.resolve.ball_cushion import (
     BallCCushionCollisionStrategy,
     BallLCushionCollisionStrategy,
@@ -68,7 +70,7 @@ def default_resolver() -> Resolver:
     The resolver YAML is found at `RESOLVER_PATH`.
     """
     return Resolver(
-        ball_ball=FrictionalInelastic(
+        ball_ball=FrictionalInelastic2D(
             friction=AlciatoreBallBallFriction(
                 a=0.009951,
                 b=0.108,

tests/physics/resolve/ball_ball/test_ball_ball.py

@@ -7,7 +7,10 @@
 from pooltool import ptmath
 from pooltool.objects.ball.datatypes import Ball
 from pooltool.physics.resolve.ball_ball.core import BallBallCollisionStrategy
-from pooltool.physics.resolve.ball_ball.frictional_inelastic import FrictionalInelastic
+from pooltool.physics.resolve.ball_ball.frictional_inelastic import (
+    FrictionalInelastic2D,
+    FrictionalInelastic3D,
+)
 from pooltool.physics.resolve.ball_ball.frictional_mathavan import FrictionalMathavan
 from pooltool.physics.resolve.ball_ball.frictionless_elastic import FrictionlessElastic
 from pooltool.physics.utils import tangent_surface_velocity
@@ -87,7 +90,12 @@ def test_head_on_zero_spin(model: BallBallCollisionStrategy):
 
 
 @pytest.mark.parametrize(
-    "model", [FrictionalInelastic(), FrictionalMathavan(num_iterations=int(1e6))]
+    "model",
+    [
+        FrictionalInelastic2D(),
+        FrictionalInelastic3D(),
+        FrictionalMathavan(num_iterations=int(1e6)),
+    ],
 )
 @pytest.mark.parametrize("e_b", [0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
 def test_head_on_zero_spin_inelastic(model: BallBallCollisionStrategy, e_b: float):
@@ -117,7 +125,9 @@ def test_head_on_zero_spin_inelastic(model: BallBallCollisionStrategy, e_b: floa
         assert cb_f.state.rvw[1][0] > 0
 
 
-@pytest.mark.parametrize("model", [FrictionalInelastic(), FrictionalMathavan()])
+@pytest.mark.parametrize(
+    "model", [FrictionalInelastic2D(), FrictionalInelastic3D(), FrictionalMathavan()]
+)
 @pytest.mark.parametrize("e_b", [0.6, 0.8, 1.0])
 def test_translating_head_on_zero_spin_inelastic(
     model: BallBallCollisionStrategy, e_b: float
@@ -134,7 +144,9 @@ def test_translating_head_on_zero_spin_inelastic(
     assert abs(cb_f.vel[1] - ob_f.vel[1]) < 1e-10
 
 
-@pytest.mark.parametrize("model", [FrictionalInelastic(), FrictionalMathavan()])
+@pytest.mark.parametrize(
+    "model", [FrictionalInelastic2D(), FrictionalInelastic3D(), FrictionalMathavan()]
+)
 @pytest.mark.parametrize("cb_wz_i", [0.1, 1, 10, 100])
 def test_head_on_z_spin(model: BallBallCollisionStrategy, cb_wz_i: float):
     """Cue ball has positive z-spin (e.g. hitting right-hand-side of cue ball)"""
@@ -155,7 +167,12 @@ def test_head_on_z_spin(model: BallBallCollisionStrategy, cb_wz_i: float):
 
 
 @pytest.mark.parametrize(
-    "model", [FrictionalInelastic(), FrictionalMathavan(num_iterations=int(1e5))]
+    "model",
+    [
+        FrictionalInelastic2D(),
+        FrictionalInelastic3D(),
+        FrictionalMathavan(num_iterations=int(1e5)),
+    ],
 )
 @pytest.mark.parametrize("speed", np.logspace(-1, 1, 4))
 @pytest.mark.parametrize(
@@ -201,7 +218,7 @@ def test_gearing_z_spin(
     assert abs(ob_f.avel[2]) < 5e-3, "Gearing english shouldn't cause induced side-spin"
 
 
-@pytest.mark.parametrize("model", [FrictionalInelastic()])
+@pytest.mark.parametrize("model", [FrictionalInelastic2D(), FrictionalInelastic3D()])
 @pytest.mark.parametrize("speed", np.logspace(0, 1, 4))
 @pytest.mark.parametrize(
     "line_of_centers_angle_radians", np.linspace(0, 2.0 * math.pi, 6, endpoint=False)