diff --git a/pooltool/evolution/engine.py b/pooltool/evolution/engine.py index 09f8ec13..f6327025 100644 --- a/pooltool/evolution/engine.py +++ b/pooltool/evolution/engine.py @@ -4,6 +4,7 @@ import attrs +from pooltool.evolution.event_based.config import DORMANT_IN_2D from pooltool.evolution.event_based.detect import EventDetector from pooltool.physics.dimensionality import Dim from pooltool.physics.resolve import Resolver @@ -42,6 +43,8 @@ def _validate_dimensionality(self) -> None: strategy = getattr(bundle, field.name) if not attrs.has(type(strategy)): continue + if not self.is_3d and field.name in DORMANT_IN_2D: + continue if not hasattr(strategy, "dim"): raise AttributeError( f"{type(bundle).__name__}.{field.name} " diff --git a/pooltool/evolution/event_based/config.py b/pooltool/evolution/event_based/config.py index c3d30f02..288be937 100644 --- a/pooltool/evolution/event_based/config.py +++ b/pooltool/evolution/event_based/config.py @@ -12,3 +12,10 @@ EventType.ROLLING_SPINNING, EventType.SLIDING_ROLLING, } + +DORMANT_IN_2D: frozenset[str] = frozenset({"ball_table"}) +"""Resolver/EventDetector field names that are dormant in 2D mode because the detection +layer doesn't emit their associated event types. In 2D, the ``dim`` of these fields is +not validated against ``SimulationEngine.is_3d`` - any tag is safe because the strategy +will never be invoked. In 3D, they are validated normally so a misdeclared ball-table +strategy is still caught.""" diff --git a/pooltool/objects/ball/params.py b/pooltool/objects/ball/params.py index aaea8c65..9a696f6a 100644 --- a/pooltool/objects/ball/params.py +++ b/pooltool/objects/ball/params.py @@ -45,6 +45,8 @@ class BallParams: The ball-ball coefficient of sliding friction. e_b: The ball-ball coefficient of restitution. + e_t: + The ball-table coefficient of restitution. e_c: The cushion coefficient of restitution. @@ -75,6 +77,7 @@ class BallParams: u_sp_proportionality: float = attrs.field(default=10 * 2 / 5 / 9) u_b: float = attrs.field(default=0.05) e_b: float = attrs.field(default=0.95) + e_t: float = attrs.field(default=0.5) e_c: float = attrs.field(default=0.85) f_c: float = attrs.field(default=0.2) g: float = attrs.field(default=9.81) diff --git a/pooltool/physics/__init__.py b/pooltool/physics/__init__.py index 303b6cf3..c462d4e2 100644 --- a/pooltool/physics/__init__.py +++ b/pooltool/physics/__init__.py @@ -22,6 +22,10 @@ BallPocketModel, ball_pocket_models, ) +from pooltool.physics.resolve.ball_table import ( + BallTableModel, + ball_table_models, +) from pooltool.physics.resolve.resolver import ( RESOLVER_PATH, Resolver, @@ -54,6 +58,7 @@ "BallCCushionModel", "BallLCushionModel", "BallPocketModel", + "BallTableModel", "StickBallModel", "BallTransitionModel", "rel_velocity", @@ -70,6 +75,7 @@ "ball_lcushion_models", "ball_ccushion_models", "ball_pocket_models", + "ball_table_models", "stick_ball_models", "ball_transition_models", # Evolve diff --git a/pooltool/physics/dimensionality.py b/pooltool/physics/dimensionality.py index d72f3783..678ae2ab 100644 --- a/pooltool/physics/dimensionality.py +++ b/pooltool/physics/dimensionality.py @@ -9,16 +9,18 @@ class Dim(StrEnum): construction to validate that the bundled strategies are compatible with its ``is_3d`` setting. - A strategy's ``dim`` is a promise about its behavior, *not* a mode switch. - ``BOTH`` means the strategy behaves identically in either mode; it does not mean - the strategy branches internally based on mode. If a strategy would behave - differently in 2D vs 3D, it should be split into separate ``TWO`` and ``THREE`` - classes. + A strategy's ``dim`` is a promise about *safety*, not behavior. Strategies + don't see ``SimulationEngine.is_3d`` - they see only the inputs handed to + them and must remain correct under the states their mode can produce. + ``BOTH`` is appropriate when the strategy is safe in either mode. It may + still take different code paths depending on the input (e.g. a branch on + ``state == const.airborne`` is dead in 2D and live in 3D), as long as + neither path is incorrect for the mode it runs under. Members: TWO: Safe only when ``SimulationEngine.is_3d`` is ``False``. THREE: Safe only when ``SimulationEngine.is_3d`` is ``True``. - BOTH: Behavior identical in either mode; safe always. + BOTH: Safe in either mode. """ TWO = auto() diff --git a/pooltool/physics/resolve/ball_table/__init__.py b/pooltool/physics/resolve/ball_table/__init__.py new file mode 100644 index 00000000..cf7f469e --- /dev/null +++ b/pooltool/physics/resolve/ball_table/__init__.py @@ -0,0 +1,30 @@ +from typing import cast + +import attrs + +from pooltool.physics.resolve.ball_table.core import BallTableCollisionStrategy +from pooltool.physics.resolve.ball_table.frictional_inelastic import ( + FrictionalInelasticTable, +) +from pooltool.physics.resolve.ball_table.frictionless_inelastic import ( + FrictionlessInelasticTable, +) +from pooltool.physics.resolve.models import BallTableModel + +_ball_table_model_registry: tuple[type[BallTableCollisionStrategy], ...] = ( + FrictionlessInelasticTable, + FrictionalInelasticTable, +) + +ball_table_models: dict[BallTableModel, type[BallTableCollisionStrategy]] = { + cast(BallTableModel, attrs.fields_dict(cls)["model"].default): cls + for cls in _ball_table_model_registry +} + +__all__ = [ + "BallTableCollisionStrategy", + "BallTableModel", + "FrictionalInelasticTable", + "FrictionlessInelasticTable", + "ball_table_models", +] diff --git a/pooltool/physics/resolve/ball_table/core.py b/pooltool/physics/resolve/ball_table/core.py new file mode 100644 index 00000000..2082508c --- /dev/null +++ b/pooltool/physics/resolve/ball_table/core.py @@ -0,0 +1,78 @@ +from abc import ABC, abstractmethod +from typing import Protocol + +import numpy as np +from numpy.typing import NDArray + +import pooltool.constants as const +from pooltool.objects.ball.datatypes import Ball +from pooltool.physics.dimensionality import Dim +from pooltool.physics.utils import on_table, rel_velocity +from pooltool.ptmath.utils import norm2d, norm3d + + +def bounce_height(vz: float, g: float) -> float: + """Return how high a ball with outgoing positive z-velocity will bounce. + + Measured as distance from table to bottom of ball. + """ + return 0.5 * vz**2 / g + + +def final_ball_motion_state(rvw: NDArray[np.float64], R: float) -> int: + """Return the final (post-collision) motion state label for a ball.""" + if rvw[0, 2] < 0: + return const.pocketed + + if rvw[1, 2] != 0.0 or not on_table(rvw, R): + return const.airborne + + if norm3d(rel_velocity(rvw, R)) > const.EPS: + return const.sliding + + if norm2d(rvw[1]) > const.EPS: + return const.rolling + + if rvw[2, 2] != 0.0: + return const.spinning + + return const.stationary + + +class _BaseStrategy(Protocol): + dim: Dim + + def resolve(self, ball: Ball, inplace: bool = False) -> Ball: ... + + def make_kiss(self, ball: Ball) -> Ball: ... + + +class BallTableCollisionStrategy(_BaseStrategy, Protocol): + """Ball-table collision models must satisfy this protocol""" + + def solve(self, ball: Ball) -> Ball: + """Resolves a ball-table collision""" + ... + + +class CoreBallTableCollision(ABC): + """Operations used by every ball-table collision resolver""" + + def make_kiss(self, ball: Ball) -> Ball: + """Translate the ball so its height is exactly its radius. + + If the ball is not at a height R, it is moved vertically such that it is. + """ + ball.state.rvw[0, 2] = ball.params.R + return ball + + def resolve(self, ball: Ball, inplace: bool = False) -> Ball: + if not inplace: + ball = ball.copy() + + ball = self.make_kiss(ball) + return self.solve(ball) + + @abstractmethod + def solve(self, ball: Ball) -> Ball: + pass diff --git a/pooltool/physics/resolve/ball_table/frictional_inelastic/__init__.py b/pooltool/physics/resolve/ball_table/frictional_inelastic/__init__.py new file mode 100644 index 00000000..17e3f9a8 --- /dev/null +++ b/pooltool/physics/resolve/ball_table/frictional_inelastic/__init__.py @@ -0,0 +1,92 @@ +import attrs +import numpy as np +from numba import jit +from numpy.typing import NDArray + +import pooltool.constants as const +import pooltool.physics as physics +import pooltool.ptmath as ptmath +from pooltool.objects.ball.datatypes import Ball, BallState +from pooltool.physics.dimensionality import Dim +from pooltool.physics.resolve.ball_table.core import ( + CoreBallTableCollision, + bounce_height, + final_ball_motion_state, +) +from pooltool.physics.resolve.models import BallTableModel + + +@jit(nopython=True, cache=const.use_numba_cache) +def _resolve_ball_table( + rvw: NDArray[np.float64], R: float, u: float, e: float +) -> NDArray[np.float64]: + rvw_i = rvw.copy() + v_i = rvw_i[1] + w_i = rvw_i[2] + if v_i[2] >= 0: + raise ValueError( + "Ball with non-negative z-velocity can't collide with table surface." + ) + + unit_z = np.array([0.0, 0.0, 1.0]) + + D_v_perpendicular_magnitude = (1 + e) * -v_i[2] + D_v_perpendicular = D_v_perpendicular_magnitude * unit_z + + v_i[2] = 0 + + v_c_i = physics.surface_velocity(rvw_i, -unit_z, R) + has_relative_velocity = ptmath.squared_norm3d(v_c_i) > const.EPS**2 + + if has_relative_velocity: + v_hat_c_i = ptmath.unit_vector(v_c_i) + D_v_parallel_slip = u * D_v_perpendicular_magnitude * -v_hat_c_i + else: + v_hat_c_i = np.zeros(3) + D_v_parallel_slip = np.zeros(3) + + D_v_parallel_no_slip = (2.0 / 7.0) * (R * ptmath.cross(w_i, unit_z) - v_i) + + if not has_relative_velocity or ptmath.squared_norm3d( + D_v_parallel_no_slip + ) <= ptmath.squared_norm3d(D_v_parallel_slip): + rvw[1] = rvw[1] + D_v_perpendicular + D_v_parallel_no_slip + rvw[2] = rvw[2] + (5.0 / 7.0) * (-w_i + ptmath.cross(unit_z, v_i) / R) + else: + rvw[1] = rvw[1] + D_v_perpendicular + D_v_parallel_slip + rvw[2] = rvw[2] + (2.5 / R) * ptmath.norm3d(D_v_parallel_slip) * ptmath.cross( + unit_z, v_hat_c_i + ) + + return rvw + + +@attrs.define +class FrictionalInelasticTable(CoreBallTableCollision): + """A frictional, inelastic ball-table collision. + + Reference: + https://billiards.colostate.edu/technical_proofs/new/TP_A-14.pdf + """ + + min_bounce_height: float = 0.005 + + model: BallTableModel = attrs.field( + default=BallTableModel.FRICTIONAL_INELASTIC, init=False, repr=False + ) + dim: Dim = attrs.field(default=Dim.THREE, init=False, repr=False) + + def solve(self, ball: Ball) -> Ball: + """Resolves the collision.""" + rvw = _resolve_ball_table( + ball.state.rvw.copy(), ball.params.R, ball.params.u_s, ball.params.e_t + ) + + if bounce_height(rvw[1, 2], ball.params.g) < self.min_bounce_height: + rvw[1, 2] = 0 + + state = final_ball_motion_state(rvw, ball.params.R) + + ball.state = BallState(rvw, state) + + return ball diff --git a/pooltool/physics/resolve/ball_table/frictionless_inelastic/__init__.py b/pooltool/physics/resolve/ball_table/frictionless_inelastic/__init__.py new file mode 100644 index 00000000..992efa7c --- /dev/null +++ b/pooltool/physics/resolve/ball_table/frictionless_inelastic/__init__.py @@ -0,0 +1,54 @@ +import attrs + +from pooltool.objects.ball.datatypes import Ball, BallState +from pooltool.physics.dimensionality import Dim +from pooltool.physics.resolve.ball_table.core import ( + CoreBallTableCollision, + bounce_height, + final_ball_motion_state, +) +from pooltool.physics.resolve.models import BallTableModel + + +def _resolve_ball_table(vz0: float, e_t: float) -> float: + if vz0 >= 0: + raise ValueError( + "Ball with non-negative z-velocity can't collide with table surface." + ) + + return -vz0 * e_t + + +@attrs.define +class FrictionlessInelasticTable(CoreBallTableCollision): + """A frictionless, inelastic collision. + + The ball bounces on the table with a coefficient of restitution. There is no + influence of friction, so only the z-component of the velocity is affected. + + To avoid infinite bouncing (the dichotomy paradox), the projected bounce height + is calculated; if it is less than ``min_bounce_height``, the z-component of the + velocity is zeroed and the outgoing ball state is set accordingly. + """ + + min_bounce_height: float = 0.005 + + model: BallTableModel = attrs.field( + default=BallTableModel.FRICTIONLESS_INELASTIC, init=False, repr=False + ) + dim: Dim = attrs.field(default=Dim.THREE, init=False, repr=False) + + def solve(self, ball: Ball) -> Ball: + """Resolves the collision.""" + vz = _resolve_ball_table(ball.state.rvw[1, 2], ball.params.e_t) + + if bounce_height(vz, ball.params.g) < self.min_bounce_height: + vz = 0.0 + + ball.state.rvw[1, 2] = vz + + state = final_ball_motion_state(ball.state.rvw, ball.params.R) + + ball.state = BallState(ball.state.rvw, state) + + return ball diff --git a/pooltool/physics/resolve/models.py b/pooltool/physics/resolve/models.py index daa0ed49..65cf5aa2 100644 --- a/pooltool/physics/resolve/models.py +++ b/pooltool/physics/resolve/models.py @@ -158,6 +158,21 @@ class StickBallModel(StrEnum): INSTANTANEOUS_POINT = auto() +class BallTableModel(StrEnum): + """An Enum for different ball-table collision models + + Attributes: + FRICTIONLESS_INELASTIC: + Frictionless, instantaneous, inelastic collision. + FRICTIONAL_INELASTIC: + Frictional, inelastic + (https://billiards.colostate.edu/technical_proofs/new/TP_A-14.pdf). + """ + + FRICTIONLESS_INELASTIC = auto() + FRICTIONAL_INELASTIC = auto() + + class BallTransitionModel(StrEnum): """An Enum for different transition models diff --git a/pooltool/physics/resolve/resolver.py b/pooltool/physics/resolve/resolver.py index d08fa82e..cb1812fb 100644 --- a/pooltool/physics/resolve/resolver.py +++ b/pooltool/physics/resolve/resolver.py @@ -30,6 +30,10 @@ BallPocketStrategy, CanonicalBallPocket, ) +from pooltool.physics.resolve.ball_table import ( + BallTableCollisionStrategy, + FrictionalInelasticTable, +) from pooltool.physics.resolve.serialize import register_serialize_hooks from pooltool.physics.resolve.stick_ball import ( StickBallCollisionStrategy, @@ -46,7 +50,7 @@ RESOLVER_PATH = pooltool.config.paths.PHYSICS_DIR / "resolver.yaml" """The location of the resolver path YAML.""" -VERSION: int = 9 +VERSION: int = 10 run = Run() @@ -82,6 +86,9 @@ def default_resolver() -> Resolver: english_throttle=1.0, squirt_throttle=1.0, ), + ball_table=FrictionalInelasticTable( + min_bounce_height=0.005, + ), transition=CanonicalTransition(), version=VERSION, ) @@ -101,6 +108,7 @@ class Resolver: ball_circular_cushion: BallCCushionCollisionStrategy ball_pocket: BallPocketStrategy stick_ball: StickBallCollisionStrategy + ball_table: BallTableCollisionStrategy transition: BallTransitionStrategy version: int | None = None @@ -142,6 +150,10 @@ def resolve(self, shot: System, event: Event) -> None: ball = shot.balls[ids[1]] self.stick_ball.resolve(cue, ball, inplace=True) ball.state.t = event.time + elif event.event_type == EventType.BALL_TABLE: + ball = shot.balls[ids[0]] + self.ball_table.resolve(ball, inplace=True) + ball.state.t = event.time _snapshot_final(shot, event) diff --git a/pooltool/physics/resolve/serialize.py b/pooltool/physics/resolve/serialize.py index ca6d2463..e3018183 100644 --- a/pooltool/physics/resolve/serialize.py +++ b/pooltool/physics/resolve/serialize.py @@ -23,6 +23,10 @@ BallPocketStrategy, ball_pocket_models, ) +from pooltool.physics.resolve.ball_table import ( + BallTableCollisionStrategy, + ball_table_models, +) from pooltool.physics.resolve.stick_ball import ( StickBallCollisionStrategy, stick_ball_models, @@ -42,6 +46,7 @@ StickBallCollisionStrategy: stick_ball_models, BallTransitionStrategy: ball_transition_models, BallBallFrictionStrategy: ball_ball_friction_models, + BallTableCollisionStrategy: ball_table_models, } diff --git a/pooltool/physics/utils.py b/pooltool/physics/utils.py index dcbef489..3ee5b725 100644 --- a/pooltool/physics/utils.py +++ b/pooltool/physics/utils.py @@ -50,6 +50,12 @@ def get_u_vec( return coordinate_rotation(unit_vector(rel_vel), -phi) +@jit(nopython=True, cache=const.use_numba_cache) +def on_table(rvw: NDArray[np.float64], R: float) -> bool: + """True when the ball's center is at the table-plane height (z == R).""" + return rvw[0, 2] == R + + @jit(nopython=True, cache=const.use_numba_cache) def get_airborne_time(rvw: NDArray[np.float64], R: float, g: float) -> float: """Time until an airborne ball's bottom touches the table plane (z = R). diff --git a/tests/evolution/test_engine.py b/tests/evolution/test_engine.py index 5351dd6f..23b755de 100644 --- a/tests/evolution/test_engine.py +++ b/tests/evolution/test_engine.py @@ -85,3 +85,22 @@ def test_dim_both_strategy_accepted_in_3d(engine_3d: SimulationEngine): detector=engine_3d.detector, is_3d=True, ) + + +def test_dormant_ball_table_skipped_in_2d(): + """In 2D, ball_table's dim is not validated (the slot is dormant).""" + resolver = SimulationEngine().resolver + assert resolver.ball_table.dim == Dim.THREE + SimulationEngine(resolver=resolver, is_3d=False) + + +def test_misdeclared_ball_table_still_caught_in_3d(engine_3d: SimulationEngine): + """In 3D, ball_table's dim *is* validated — a Dim.TWO ball_table is a bug.""" + engine_3d.resolver.ball_table.dim = Dim.TWO + + with pytest.raises(ValueError, match=r"ball_table.*incompatible with is_3d=True"): + SimulationEngine( + resolver=engine_3d.resolver, + detector=engine_3d.detector, + is_3d=True, + ) diff --git a/tests/physics/resolve/ball_table/__init__.py b/tests/physics/resolve/ball_table/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/physics/resolve/ball_table/test_ball_table.py b/tests/physics/resolve/ball_table/test_ball_table.py new file mode 100644 index 00000000..aac13def --- /dev/null +++ b/tests/physics/resolve/ball_table/test_ball_table.py @@ -0,0 +1,100 @@ +from math import isclose + +import numpy as np +import pytest + +from pooltool import constants +from pooltool.objects.ball.datatypes import Ball +from pooltool.physics.resolve.ball_table.core import ( + BallTableCollisionStrategy, + bounce_height, +) +from pooltool.physics.resolve.ball_table.frictional_inelastic import ( + FrictionalInelasticTable, +) +from pooltool.physics.resolve.ball_table.frictionless_inelastic import ( + FrictionlessInelasticTable, +) + + +@pytest.mark.parametrize( + "vz,g,expected", + [ + (4.0, 9.8, 0.5 * 16 / 9.8), + (0.0, 9.8, 0.0), + (10.0, 9.8, 0.5 * 100 / 9.8), + (2.5, 10.0, 0.5 * 6.25 / 10.0), + ], +) +def test_bounce_height(vz, g, expected): + result = bounce_height(vz, g) + assert isclose(result, expected, rel_tol=1e-7) + + +def test_bounce_height_negative_vz(): + vz = -3.0 + g = 9.8 + expected = 0.5 * (vz**2) / g + assert isclose(bounce_height(vz, g), expected, rel_tol=1e-7) + + +models = [FrictionlessInelasticTable(), FrictionalInelasticTable()] + + +def example() -> Ball: + return Ball.create("cue", xy=(0, 0)) + + +@pytest.mark.parametrize("model", models) +@pytest.mark.parametrize("vz0", -np.logspace(-5, 4, 10, base=10)) +def test_non_negative_output_velocity(model: BallTableCollisionStrategy, vz0: float): + ball = example() + ball.state.rvw[1, 2] = vz0 + + model.resolve(ball, inplace=True) + + assert ball.state.rvw[1, 2] >= 0 + + +@pytest.mark.parametrize("model", models) +@pytest.mark.parametrize("vz0", -np.logspace(-5, 4, 10, base=10)) +def test_decaying_velocity(model: BallTableCollisionStrategy, vz0: float): + ball = example() + ball.state.rvw[1, 2] = vz0 + + model.resolve(ball, inplace=True) + + assert ball.state.rvw[1, 2] <= -vz0 + + +@pytest.mark.parametrize("model", models) +def test_positive_incoming_velocity_fails(model: BallTableCollisionStrategy): + ball = example() + ball.state.rvw[1, 2] = +1 + + with pytest.raises(ValueError, match="can't collide with table surface"): + model.resolve(ball, inplace=True) + + +@pytest.mark.parametrize("model", models) +def test_zero_incoming_velocity_fails(model: BallTableCollisionStrategy): + ball = example() + ball.state.rvw[1, 2] = 0.0 + + with pytest.raises(ValueError, match="can't collide with table surface"): + model.resolve(ball, inplace=True) + + +@pytest.mark.parametrize("model", models) +def test_non_airborne_outgoing_state(model: BallTableCollisionStrategy): + """A very small incoming velocity leads to a non-airborne outgoing state. + + This avoids perpetual bouncing (the dichotomy paradox). + """ + ball = example() + ball.state.rvw[1, 2] = -0.001 + + model.resolve(ball, inplace=True) + + assert ball.state.s != constants.airborne + assert ball.state.rvw[1, 2] == 0.0 diff --git a/tests/physics/resolve/ball_table/test_frictionless_inelastic.py b/tests/physics/resolve/ball_table/test_frictionless_inelastic.py new file mode 100644 index 00000000..960f4c66 --- /dev/null +++ b/tests/physics/resolve/ball_table/test_frictionless_inelastic.py @@ -0,0 +1,26 @@ +from math import isclose + +import pytest + +from pooltool.physics.resolve.ball_table.frictionless_inelastic import ( + _resolve_ball_table, +) + + +@pytest.mark.parametrize( + "vz0,e_t,expected", + [ + (-5.0, 0.9, 4.5), + (-1.0, 1.0, 1.0), + (-2.0, 0.5, 1.0), + (-3.5, 0.75, 2.625), + ], +) +def test_resolve_ball_table_valid(vz0, e_t, expected): + assert isclose(_resolve_ball_table(vz0, e_t), expected) + + +@pytest.mark.parametrize("vz0,e_t", [(0.0, 1.0), (1.0, 0.5), (2.0, 0.9)]) +def test_resolve_ball_table_invalid(vz0, e_t): + with pytest.raises(ValueError, match="can't collide with table surface"): + _resolve_ball_table(vz0, e_t)