Add ball-table resolver strategies

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} "

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."""

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)

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

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()

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",
+]

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

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

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

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