From 83c76db8dd013dc0ccfe888432ec82e91b619bf3 Mon Sep 17 00:00:00 2001 From: Derek McBlane Date: Sun, 24 May 2026 14:33:52 -0400 Subject: [PATCH 1/3] convert ball-ball FrictionalInelastic to FrictionalInelastic3D and FrictionalInelastic2D --- .../physics/resolve/ball_ball/__init__.py | 8 +- pooltool/physics/resolve/ball_ball/core.py | 18 +++++ .../frictional_inelastic/__init__.py | 78 ++++++++++++------- pooltool/physics/resolve/models.py | 3 +- pooltool/physics/resolve/resolver.py | 6 +- .../resolve/ball_ball/test_ball_ball.py | 22 ++++-- 6 files changed, 97 insertions(+), 38 deletions(-) diff --git a/pooltool/physics/resolve/ball_ball/__init__.py b/pooltool/physics/resolve/ball_ball/__init__.py index ebe28ee4..49016cc7 100644 --- a/pooltool/physics/resolve/ball_ball/__init__.py +++ b/pooltool/physics/resolve/ball_ball/__init__.py @@ -5,7 +5,10 @@ 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 @@ -13,7 +16,8 @@ _ball_ball_model_registry: tuple[type[BallBallCollisionStrategy], ...] = ( FrictionlessElastic, FrictionalMathavan, - FrictionalInelastic, + FrictionalInelastic2D, + FrictionalInelastic3D, ) ball_ball_models: dict[BallBallModel, type[BallBallCollisionStrategy]] = { diff --git a/pooltool/physics/resolve/ball_ball/core.py b/pooltool/physics/resolve/ball_ball/core.py index 49f8db67..0a846fee 100644 --- a/pooltool/physics/resolve/ball_ball/core.py +++ b/pooltool/physics/resolve/ball_ball/core.py @@ -2,6 +2,7 @@ from typing import Protocol import numpy as np +from numpy.typing import NDArray import pooltool.constants as const import pooltool.ptmath as ptmath @@ -9,6 +10,23 @@ 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]: ... diff --git a/pooltool/physics/resolve/ball_ball/frictional_inelastic/__init__.py b/pooltool/physics/resolve/ball_ball/frictional_inelastic/__init__.py index 6e3c2035..ac6722c5 100644 --- a/pooltool/physics/resolve/ball_ball/frictional_inelastic/__init__.py +++ b/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 diff --git a/pooltool/physics/resolve/models.py b/pooltool/physics/resolve/models.py index 4eb3174d..91d29ea4 100644 --- a/pooltool/physics/resolve/models.py +++ b/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() diff --git a/pooltool/physics/resolve/resolver.py b/pooltool/physics/resolve/resolver.py index b07847b2..02ebbcf3 100644 --- a/pooltool/physics/resolve/resolver.py +++ b/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, diff --git a/tests/physics/resolve/ball_ball/test_ball_ball.py b/tests/physics/resolve/ball_ball/test_ball_ball.py index 4e30f6db..0c71acf1 100644 --- a/tests/physics/resolve/ball_ball/test_ball_ball.py +++ b/tests/physics/resolve/ball_ball/test_ball_ball.py @@ -7,7 +7,9 @@ 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, +) 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 +89,11 @@ def test_head_on_zero_spin(model: BallBallCollisionStrategy): @pytest.mark.parametrize( - "model", [FrictionalInelastic(), FrictionalMathavan(num_iterations=int(1e6))] + "model", + [ + FrictionalInelastic2D(), + 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 +123,7 @@ 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(), 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 +140,7 @@ 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(), 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 +161,11 @@ def test_head_on_z_spin(model: BallBallCollisionStrategy, cb_wz_i: float): @pytest.mark.parametrize( - "model", [FrictionalInelastic(), FrictionalMathavan(num_iterations=int(1e5))] + "model", + [ + FrictionalInelastic2D(), + FrictionalMathavan(num_iterations=int(1e5)), + ], ) @pytest.mark.parametrize("speed", np.logspace(-1, 1, 4)) @pytest.mark.parametrize( @@ -201,7 +211,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()]) @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) From 1e6d39773246455be4a3ea55b6550c3642696203 Mon Sep 17 00:00:00 2001 From: Derek McBlane Date: Sun, 24 May 2026 15:23:47 -0400 Subject: [PATCH 2/3] add FrictionalInelastic3D to test_ball_ball.py --- tests/physics/resolve/ball_ball/test_ball_ball.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/physics/resolve/ball_ball/test_ball_ball.py b/tests/physics/resolve/ball_ball/test_ball_ball.py index 0c71acf1..9277b871 100644 --- a/tests/physics/resolve/ball_ball/test_ball_ball.py +++ b/tests/physics/resolve/ball_ball/test_ball_ball.py @@ -9,6 +9,7 @@ from pooltool.physics.resolve.ball_ball.core import BallBallCollisionStrategy 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 @@ -92,6 +93,7 @@ def test_head_on_zero_spin(model: BallBallCollisionStrategy): "model", [ FrictionalInelastic2D(), + FrictionalInelastic3D(), FrictionalMathavan(num_iterations=int(1e6)), ], ) @@ -123,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", [FrictionalInelastic2D(), 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 @@ -140,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", [FrictionalInelastic2D(), 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)""" @@ -164,6 +170,7 @@ def test_head_on_z_spin(model: BallBallCollisionStrategy, cb_wz_i: float): "model", [ FrictionalInelastic2D(), + FrictionalInelastic3D(), FrictionalMathavan(num_iterations=int(1e5)), ], ) @@ -211,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", [FrictionalInelastic2D()]) +@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) From 8d0c0ba6ff1efce87090777f01ed69390bc09abe Mon Sep 17 00:00:00 2001 From: Derek McBlane Date: Mon, 15 Jun 2026 22:46:34 -0400 Subject: [PATCH 3/3] fixup ball-ball detection --- .../evolution/event_based/detect/ball_ball.py | 32 +++++++++++++++++-- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/pooltool/evolution/event_based/detect/ball_ball.py b/pooltool/evolution/event_based/detect/ball_ball.py index df782764..b852fc16 100644 --- a/pooltool/evolution/event_based/detect/ball_ball.py +++ b/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)