3D stick-ball resolver + a partially functional 3D playground

pooltool/evolution/event_based/detect/ball_ball.py

@@ -78,4 +78,4 @@ def get_next_ball_ball_2d_event(shot: System, collision_cache: CollisionCache) -
 
 
 def get_next_ball_ball_3d_event(shot: System, collision_cache: CollisionCache) -> Event:
-    raise NotImplementedError("3D ball-ball detection has not been vendored yet")
+    return null_event(np.inf)

pooltool/evolution/event_based/detect/ball_cushion.py

@@ -70,9 +70,8 @@ def get_next_ball_linear_cushion_2d_event(
 def get_next_ball_linear_cushion_3d_event(
     shot: System, collision_cache: CollisionCache
 ) -> Event:
-    raise NotImplementedError(
-        "3D ball-linear-cushion detection has not been vendored yet"
-    )
+    """3D ball-linear-cushion detection — not vendored yet; emits no event."""
+    return null_event(np.inf)
 
 
 def get_next_ball_circular_cushion_2d_event(
@@ -123,6 +122,4 @@ def get_next_ball_circular_cushion_2d_event(
 def get_next_ball_circular_cushion_3d_event(
     shot: System, collision_cache: CollisionCache
 ) -> Event:
-    raise NotImplementedError(
-        "3D ball-circular-cushion detection has not been vendored yet"
-    )
+    return null_event(np.inf)

pooltool/evolution/event_based/detect/ball_pocket.py

@@ -57,4 +57,4 @@ def get_next_ball_pocket_2d_event(
 def get_next_ball_pocket_3d_event(
     shot: System, collision_cache: CollisionCache
 ) -> Event:
-    raise NotImplementedError("3D ball-pocket detection has not been vendored yet")
+    return null_event(np.inf)

pooltool/physics/resolve/models.py

@@ -141,21 +141,32 @@ class BallPocketModel(StrEnum):
 class StickBallModel(StrEnum):
     """An Enum for different stick-ball collision models
 
-    Attributes:
-        INSTANTANEOUS_POINT:
-            Instantaneous and point-like stick-ball interaction.
-
-            This collision assumes the stick-ball interaction is instantaneous and point-like.
+    The underlying math (see :func:`cue_strike`) is fully 3D: it produces a
+    ball velocity with a vertical component proportional to ``sin(theta)``,
+    the cue elevation. The two members below differ only in what they do with
+    that vertical component.
 
-            Note:
-                - A derivation of this model can be found in Dr. Dave Billiard's technical proof
-                  A-30 (https://billiards.colostate.edu/technical_proofs/new/TP_A-30.pdf)
-
-            Additionally, a deflection (squirt) angle is calculated via
-            :mod:`pooltool.physics.resolve.stick_ball.squirt`).
+    Attributes:
+        INSTANTANEOUS_POINT_3D:
+            Instantaneous and point-like stick-ball interaction. The full 3D
+            cue-strike result is applied to the ball, including any vertical
+            velocity produced by cue elevation. Suitable for a 3D simulation
+            that supports airborne motion.
+
+            A derivation can be found in Dr. Dave Billiard's technical proof
+            A-30 (https://billiards.colostate.edu/technical_proofs/new/TP_A-30.pdf).
+            A deflection (squirt) angle is calculated via
+            :mod:`pooltool.physics.resolve.stick_ball.squirt`.
+
+        INSTANTANEOUS_POINT_2D:
+            Same 3D cue-strike math as ``INSTANTANEOUS_POINT_3D``, followed by
+            a clamp that zeros the vertical velocity component of the ball.
+            Suitable for a 2D simulation: the ball never leaves the table
+            surface regardless of cue elevation.
     """
 
-    INSTANTANEOUS_POINT = auto()
+    INSTANTANEOUS_POINT_2D = auto()
+    INSTANTANEOUS_POINT_3D = auto()
 
 
 class BallTableModel(StrEnum):

pooltool/physics/resolve/resolver.py

@@ -38,7 +38,7 @@
 from pooltool.physics.resolve.stick_ball import (
     StickBallCollisionStrategy,
 )
-from pooltool.physics.resolve.stick_ball.instantaneous_point import InstantaneousPoint
+from pooltool.physics.resolve.stick_ball.instantaneous_point import InstantaneousPoint2D
 from pooltool.physics.resolve.transition import (
     BallTransitionStrategy,
     CanonicalTransition,
@@ -50,7 +50,7 @@
 RESOLVER_PATH = pooltool.config.paths.PHYSICS_DIR / "resolver.yaml"
 """The location of the resolver path YAML."""
 
-VERSION: int = 10
+VERSION: int = 11
 
 
 run = Run()
@@ -82,7 +82,7 @@ def default_resolver() -> Resolver:
             omega_ratio=1.8,
         ),
         ball_pocket=CanonicalBallPocket(),
-        stick_ball=InstantaneousPoint(
+        stick_ball=InstantaneousPoint2D(
             english_throttle=1.0,
             squirt_throttle=1.0,
         ),

pooltool/physics/resolve/stick_ball/__init__.py

@@ -4,10 +4,14 @@
 
 from pooltool.physics.resolve.models import StickBallModel
 from pooltool.physics.resolve.stick_ball.core import StickBallCollisionStrategy
-from pooltool.physics.resolve.stick_ball.instantaneous_point import InstantaneousPoint
+from pooltool.physics.resolve.stick_ball.instantaneous_point import (
+    InstantaneousPoint2D,
+    InstantaneousPoint3D,
+)
 
 _stick_ball_model_registry: tuple[type[StickBallCollisionStrategy], ...] = (
-    InstantaneousPoint,
+    InstantaneousPoint2D,
+    InstantaneousPoint3D,
 )
 
 stick_ball_models: dict[StickBallModel, type[StickBallCollisionStrategy]] = {

pooltool/physics/resolve/stick_ball/core.py

@@ -1,11 +1,30 @@
 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.objects.cue.datatypes import Cue
 from pooltool.physics.dimensionality import Dim
 
 
+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 resolve(
         self, cue: Cue, ball: Ball, inplace: bool = False

pooltool/physics/resolve/stick_ball/instantaneous_point/__init__.py

@@ -1,13 +1,15 @@
 import attrs
 import numpy as np
 
-import pooltool.constants as const
 import pooltool.ptmath as ptmath
 from pooltool.objects.ball.datatypes import Ball, BallState
 from pooltool.objects.cue.datatypes import Cue
 from pooltool.physics.dimensionality import Dim
 from pooltool.physics.resolve.models import StickBallModel
-from pooltool.physics.resolve.stick_ball.core import CoreStickBallCollision
+from pooltool.physics.resolve.stick_ball.core import (
+    CoreStickBallCollision,
+    final_ball_motion_state,
+)
 from pooltool.physics.resolve.stick_ball.squirt import get_squirt_angle
 from pooltool.ptmath.utils import coordinate_rotation
 
@@ -85,9 +87,7 @@ def cue_strike(m, M, R, V0, phi, theta, Q):
     denominator = 1 + m / M + temp / I_m
     v = numerator / denominator
 
-    # 3D FIXME
-    # v_B = -v * np.array([0, np.cos(theta), np.sin(theta)])
-    v_B = -v * np.array([0, np.cos(theta), 0])
+    v_B = -v * np.array([0, np.cos(theta), np.sin(theta)])
 
     vec_x = -c * np.sin(theta) + b * np.cos(theta)
     vec_y = a * np.sin(theta)
@@ -105,7 +105,7 @@ def cue_strike(m, M, R, V0, phi, theta, Q):
 
 
 @attrs.define
-class InstantaneousPoint(CoreStickBallCollision):
+class InstantaneousPoint3D(CoreStickBallCollision):
     """Instantaneous and point-like stick-ball interaction
 
     This collision assumes the stick-ball interaction is instantaneous and point-like.
@@ -128,9 +128,9 @@ class InstantaneousPoint(CoreStickBallCollision):
     squirt_throttle: float = 1.0
 
     model: StickBallModel = attrs.field(
-        default=StickBallModel.INSTANTANEOUS_POINT, init=False, repr=False
+        default=StickBallModel.INSTANTANEOUS_POINT_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, cue: Cue, ball: Ball) -> tuple[Cue, Ball]:
         # Transform contact point Q from cue frame to ball frame
@@ -163,8 +163,27 @@ def solve(self, cue: Cue, ball: Ball) -> tuple[Cue, Ball]:
         v = coordinate_rotation(v, alpha)
 
         rvw = np.array([ball.state.rvw[0], v, w * self.english_throttle])
-        s = const.sliding
+        s = final_ball_motion_state(rvw, ball.params.R)
 
         ball.state = BallState(rvw, s)
 
         return cue, ball
+
+
+@attrs.define
+class InstantaneousPoint2D(InstantaneousPoint3D):
+    """Instantaneous and point-like stick-ball interaction (2D)
+
+    For details see :class:`InstantaneousPoint3D`.
+    """
+
+    model: StickBallModel = attrs.field(
+        default=StickBallModel.INSTANTANEOUS_POINT_2D, init=False, repr=False
+    )
+    dim: Dim = attrs.field(default=Dim.BOTH, init=False, repr=False)
+
+    def solve(self, cue: Cue, ball: Ball) -> tuple[Cue, Ball]:
+        cue, ball = super().solve(cue, ball)
+        ball.state.rvw[1, 2] = 0.0
+        ball.state.s = final_ball_motion_state(ball.state.rvw, ball.params.R)
+        return cue, ball

sandbox/airborne_demos.py

@@ -0,0 +1,133 @@
+#! /usr/bin/env python
+"""Demos 3D trajectory (work in progress)
+
+Usage:
+    python sandbox/airborne_demos.py --name drop
+"""
+
+import argparse
+
+import attrs
+
+from pooltool import constants as const
+from pooltool.evolution.engine import SimulationEngine
+from pooltool.evolution.event_based.simulate import simulate
+from pooltool.interact import show
+from pooltool.objects.ball.datatypes import Ball
+from pooltool.objects.cue.datatypes import Cue
+from pooltool.objects.table.datatypes import Table
+from pooltool.physics.dimensionality import Dim
+from pooltool.physics.resolve.resolver import Resolver
+from pooltool.physics.resolve.stick_ball.instantaneous_point import (
+    InstantaneousPoint3D,
+)
+from pooltool.system.datatypes import System
+
+
+def _build_3d_engine() -> SimulationEngine:
+    """Build a SimulationEngine with ``is_3d=True``.
+
+    Every resolver strategy that carries a ``dim`` tag is patched to
+    ``Dim.BOTH`` so the engine constructs; the stick-ball strategy is
+    swapped to ``InstantaneousPoint3D`` so cue elevation produces real
+    vertical velocity.
+    """
+    # Patches all defaults to dim.BOTH so the engine constructs
+    resolver = Resolver.default()
+    for field in attrs.fields(type(resolver)):
+        strategy = getattr(resolver, field.name)
+        if hasattr(strategy, "dim"):
+            strategy.dim = Dim.BOTH
+
+    # Replace all working 3D resolvers
+    resolver.stick_ball = InstantaneousPoint3D()
+
+    return SimulationEngine(resolver=resolver, is_3d=True)
+
+
+def _empty_cue() -> Cue:
+    """A cue with ``V0=0`` so the simulator doesn't fire a stick strike at t=0.
+
+    The default ``Cue`` constructor sets ``V0=2.0``; the stick-ball detector
+    fires whenever ``V0 > 0`` and ``_system_has_energy`` reports false. For
+    a ball that's airborne but momentarily at apex (vz=0), kinetic energy is
+    zero and the detector would trigger a strike on the handcrafted state.
+    Explicitly zero V0 to suppress that.
+
+    TODO: Fix underlying problem in codebase
+    """
+    cue = Cue(cue_ball_id="cue")
+    cue.V0 = 0
+    return cue
+
+
+def drop() -> System:
+    """Ball dropped from 0.3 m with a small horizontal nudge in +x."""
+    ball = Ball.create("cue", xy=(0.5, 0.5))
+    ball.state.rvw[0, 2] = 0.3
+    ball.state.rvw[1, 0] = 0.5
+    ball.state.s = const.airborne
+
+    return System(
+        cue=_empty_cue(),
+        table=Table.default(),
+        balls=(ball,),
+    )
+
+
+def impulse_into() -> System:
+    """Strong downward strike with a small horizontal nudge in +y."""
+    ball = Ball.create("cue", xy=(0.5, 0.5))
+    ball.state.rvw[1, 1] = 0.5
+    ball.state.rvw[1, 2] = -5.0
+    ball.state.s = const.airborne
+
+    return System(
+        cue=_empty_cue(),
+        table=Table.default(),
+        balls=(ball,),
+    )
+
+
+def jump() -> System:
+    """A genuine jump shot — cue strike at 60° elevation produces vz via the 3D resolver.
+
+    No handcrafted ``rvw`` here: the cue strikes a ball at rest on the table
+    surface, and ``InstantaneousPoint3D`` lifts it off via ``v·sin(theta)``.
+    """
+    ball = Ball.create("cue", xy=(0.5, 0.5))
+    cue = Cue(cue_ball_id="cue")
+    cue.set_state(V0=2.0, phi=90.0, theta=60.0, a=0.0, b=0.0)
+
+    return System(
+        cue=cue,
+        table=Table.default(),
+        balls=(ball,),
+    )
+
+
+_map = {
+    "drop": drop,
+    "impulse_into": impulse_into,
+    "jump": jump,
+}
+
+
+def main(name: str) -> None:
+    engine = _build_3d_engine()
+    shot = _map[name]()
+    simulate(shot, engine=engine, inplace=True)
+    show(shot)
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser("Airborne ball demos in 3D mode.")
+    ap.add_argument("--name", choices=list(_map.keys()) + ["all"], required=True)
+    args = ap.parse_args()
+
+    if args.name == "all":
+        for name in _map:
+            print(f"Running {name}...")
+            main(name)
+    else:
+        main(args.name)