diff --git a/pooltool/evolution/event_based/detect/ball_ball.py b/pooltool/evolution/event_based/detect/ball_ball.py index 2c3c550b..1796f7ea 100644 --- a/pooltool/evolution/event_based/detect/ball_ball.py +++ b/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) diff --git a/pooltool/evolution/event_based/detect/ball_cushion.py b/pooltool/evolution/event_based/detect/ball_cushion.py index d67cbc65..6eaf494e 100644 --- a/pooltool/evolution/event_based/detect/ball_cushion.py +++ b/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) diff --git a/pooltool/evolution/event_based/detect/ball_pocket.py b/pooltool/evolution/event_based/detect/ball_pocket.py index 1424f873..9d6311de 100644 --- a/pooltool/evolution/event_based/detect/ball_pocket.py +++ b/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) diff --git a/pooltool/physics/resolve/models.py b/pooltool/physics/resolve/models.py index 65cf5aa2..4eb3174d 100644 --- a/pooltool/physics/resolve/models.py +++ b/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): diff --git a/pooltool/physics/resolve/resolver.py b/pooltool/physics/resolve/resolver.py index cb1812fb..b07847b2 100644 --- a/pooltool/physics/resolve/resolver.py +++ b/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, ), diff --git a/pooltool/physics/resolve/stick_ball/__init__.py b/pooltool/physics/resolve/stick_ball/__init__.py index 7f7e831e..e45c46ca 100644 --- a/pooltool/physics/resolve/stick_ball/__init__.py +++ b/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]] = { diff --git a/pooltool/physics/resolve/stick_ball/core.py b/pooltool/physics/resolve/stick_ball/core.py index c2af34da..fa4537cc 100644 --- a/pooltool/physics/resolve/stick_ball/core.py +++ b/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 diff --git a/pooltool/physics/resolve/stick_ball/instantaneous_point/__init__.py b/pooltool/physics/resolve/stick_ball/instantaneous_point/__init__.py index f940c7e7..db227a49 100644 --- a/pooltool/physics/resolve/stick_ball/instantaneous_point/__init__.py +++ b/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 diff --git a/sandbox/airborne_demos.py b/sandbox/airborne_demos.py new file mode 100644 index 00000000..18a61967 --- /dev/null +++ b/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)