Solver migration

README.md

@@ -6,7 +6,7 @@
 ![GitHub Actions Workflow Status](https://img.shields.io/github/actions/workflow/status/ekiefl/pooltool/test.yml)
 
 ![PyPI - Version](https://img.shields.io/pypi/v/pooltool-billiards)
-![PyPI - Python Version](https://img.shields.io/pypi/pyversions/pooltool-billiards)
+[![Python Version](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2Fekiefl%2Fpooltool%2Fmain%2Fpyproject.toml)](https://pypi.org/project/pooltool-billiards/)
 [![codecov](https://codecov.io/gh/ekiefl/pooltool/graph/badge.svg?flag=service-no-ani)](https://codecov.io/gh/ekiefl/pooltool)
 
 [![Discord](https://img.shields.io/badge/Discord-Join%20Server-7289da?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/8Y8qUgzZhz)

docs/include_exclude.py

@@ -29,7 +29,6 @@
         "pooltool.ruleset.snooker",
         # API: pooltool.physics
         "pooltool.physics.evolve",
-        "pooltool.physics.motion",
         "pooltool.physics.resolve",
     ],
     "module": [

pooltool/evolution/event_based/detect/__init__.py

@@ -3,14 +3,11 @@
     get_next_ball_ball_3d_event,
 )
 from pooltool.evolution.event_based.detect.ball_cushion import (
-    get_next_ball_circular_cushion_2d_event,
-    get_next_ball_circular_cushion_3d_event,
-    get_next_ball_linear_cushion_2d_event,
-    get_next_ball_linear_cushion_3d_event,
+    get_next_ball_circular_cushion_event,
+    get_next_ball_linear_cushion_event,
 )
 from pooltool.evolution.event_based.detect.ball_pocket import (
-    get_next_ball_pocket_2d_event,
-    get_next_ball_pocket_3d_event,
+    get_next_ball_pocket_event,
 )
 from pooltool.evolution.event_based.detect.ball_table import (
     get_next_ball_table_event,
@@ -24,12 +21,9 @@
     "EventDetector",
     "get_next_ball_ball_2d_event",
     "get_next_ball_ball_3d_event",
-    "get_next_ball_circular_cushion_2d_event",
-    "get_next_ball_circular_cushion_3d_event",
-    "get_next_ball_linear_cushion_2d_event",
-    "get_next_ball_linear_cushion_3d_event",
-    "get_next_ball_pocket_2d_event",
-    "get_next_ball_pocket_3d_event",
+    "get_next_ball_circular_cushion_event",
+    "get_next_ball_linear_cushion_event",
+    "get_next_ball_pocket_event",
     "get_next_ball_table_event",
     "get_next_stick_ball_event",
 ]

pooltool/evolution/event_based/detect/ball_ball.py

@@ -3,15 +3,84 @@
 from itertools import combinations
 
 import numpy as np
+from numba import jit
+from numpy.typing import NDArray
 
 import pooltool.constants as const
 import pooltool.ptmath as ptmath
 from pooltool.events import Event, EventType, ball_ball_collision, null_event
 from pooltool.evolution.event_based.cache import CollisionCache
-from pooltool.physics.motion.solve import ball_ball_collision_time
+from pooltool.physics.utils import get_u_vec
+from pooltool.ptmath.roots import quartic
+from pooltool.ptmath.roots.core import get_real_positive_smallest_root
 from pooltool.system.datatypes import System
 
 
+@jit(nopython=True, cache=const.use_numba_cache)
+def ball_ball_collision_time(
+    rvw1: NDArray[np.float64],
+    rvw2: NDArray[np.float64],
+    s1: int,
+    s2: int,
+    mu1: float,
+    mu2: float,
+    m1: float,
+    m2: float,
+    g1: float,
+    g2: float,
+    R: float,
+) -> float:
+    """Get the time until collision between 2 balls."""
+    c1x, c1y = rvw1[0, 0], rvw1[0, 1]
+    c2x, c2y = rvw2[0, 0], rvw2[0, 1]
+
+    if s1 == const.spinning or s1 == const.pocketed or s1 == const.stationary:
+        a1x, a1y, b1x, b1y = 0, 0, 0, 0
+    else:
+        phi1 = ptmath.angle(rvw1[1])
+        v1 = ptmath.norm3d(rvw1[1])
+
+        u1 = get_u_vec(rvw1, R, phi1, s1)
+
+        K1 = -0.5 * mu1 * g1
+        cos_phi1 = np.cos(phi1)
+        sin_phi1 = np.sin(phi1)
+
+        a1x = K1 * (u1[0] * cos_phi1 - u1[1] * sin_phi1)
+        a1y = K1 * (u1[0] * sin_phi1 + u1[1] * cos_phi1)
+        b1x = v1 * cos_phi1
+        b1y = v1 * sin_phi1
+
+    if s2 == const.spinning or s2 == const.pocketed or s2 == const.stationary:
+        a2x, a2y, b2x, b2y = 0.0, 0.0, 0.0, 0.0
+    else:
+        phi2 = ptmath.angle(rvw2[1])
+        v2 = ptmath.norm3d(rvw2[1])
+
+        u2 = get_u_vec(rvw2, R, phi2, s2)
+
+        K2 = -0.5 * mu2 * g2
+        cos_phi2 = np.cos(phi2)
+        sin_phi2 = np.sin(phi2)
+
+        a2x = K2 * (u2[0] * cos_phi2 - u2[1] * sin_phi2)
+        a2y = K2 * (u2[0] * sin_phi2 + u2[1] * cos_phi2)
+        b2x = v2 * cos_phi2
+        b2y = v2 * sin_phi2
+
+    Ax, Ay = a2x - a1x, a2y - a1y
+    Bx, By = b2x - b1x, b2y - b1y
+    Cx, Cy = c2x - c1x, c2y - c1y
+
+    a = Ax * Ax + Ay * Ay
+    b = 2 * Ax * Bx + 2 * Ay * By
+    c = Bx * Bx + 2 * Ax * Cx + 2 * Ay * Cy + By * By
+    d = 2 * Bx * Cx + 2 * By * Cy
+    e = Cx * Cx + Cy * Cy - 4 * R * R
+
+    return get_real_positive_smallest_root(quartic.solve(a, b, c, d, e))
+
+
 def get_next_ball_ball_2d_event(shot: System, collision_cache: CollisionCache) -> Event:
     """Detect the next ball-ball collision in 2D mode."""
     cache = collision_cache.times.setdefault(EventType.BALL_BALL, {})

pooltool/evolution/event_based/detect/ball_cushion.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
 import numpy as np
+from numba import jit
+from numpy.typing import NDArray
 
 import pooltool.constants as const
+import pooltool.physics.evolve as evolve
+import pooltool.ptmath as ptmath
 from pooltool.events import (
     Event,
     EventType,
@@ -11,14 +15,142 @@
     null_event,
 )
 from pooltool.evolution.event_based.cache import CollisionCache
-from pooltool.physics.motion.solve import (
-    ball_circular_cushion_collision_time,
-    ball_linear_cushion_collision_time,
-)
+from pooltool.physics.utils import get_u_vec
+from pooltool.ptmath.roots import get_real_positive_smallest_root, quartic
 from pooltool.system.datatypes import System
 
 
-def get_next_ball_linear_cushion_2d_event(
+@jit(nopython=True, cache=const.use_numba_cache)
+def ball_vertical_plane_collision_time(
+    rvw: NDArray[np.float64],
+    s: int,
+    lx: float,
+    ly: float,
+    l0: float,
+    p1: NDArray[np.float64],
+    p2: NDArray[np.float64],
+    direction: int,
+    mu: float,
+    m: float,
+    g: float,
+    R: float,
+) -> float:
+    """Get time until collision between a ball and a vertical plane.
+
+    For ball trajectories limited to the playing surface, this suffices for
+    detecting ball collisions with linear cushion segments.
+
+    Note:
+        - This is broken for airborne balls.
+    """
+    if s == const.spinning or s == const.pocketed or s == const.stationary:
+        return np.inf
+
+    phi = ptmath.angle(rvw[1])
+    v = ptmath.norm3d(rvw[1])
+
+    u = get_u_vec(rvw, R, phi, s)
+
+    K = -0.5 * mu * g
+    cos_phi = np.cos(phi)
+    sin_phi = np.sin(phi)
+
+    ax = K * (u[0] * cos_phi - u[1] * sin_phi)
+    ay = K * (u[0] * sin_phi + u[1] * cos_phi)
+    bx, by = v * cos_phi, v * sin_phi
+    cx, cy = rvw[0, 0], rvw[0, 1]
+
+    A = lx * ax + ly * ay
+    B = lx * bx + ly * by
+
+    if direction == 0:
+        C = l0 + lx * cx + ly * cy + R * np.sqrt(lx * lx + ly * ly)
+        root1, root2 = ptmath.roots.quadratic.solve(A, B, C)
+        roots = [root1, root2]
+    elif direction == 1:
+        C = l0 + lx * cx + ly * cy - R * np.sqrt(lx * lx + ly * ly)
+        root1, root2 = ptmath.roots.quadratic.solve(A, B, C)
+        roots = [root1, root2]
+    else:
+        C1 = l0 + lx * cx + ly * cy + R * np.sqrt(lx * lx + ly * ly)
+        C2 = l0 + lx * cx + ly * cy - R * np.sqrt(lx * lx + ly * ly)
+        root1, root2 = ptmath.roots.quadratic.solve(A, B, C1)
+        root3, root4 = ptmath.roots.quadratic.solve(A, B, C2)
+        roots = [root1, root2, root3, root4]
+
+    min_time = np.inf
+    for root in roots:
+        if np.isnan(root):
+            continue
+
+        if np.abs(root.imag) > const.EPS:
+            continue
+
+        if root.real <= const.EPS:
+            continue
+
+        rvw_dtau, _ = evolve.evolve_ball_motion(s, rvw, R, m, mu, 1, mu, g, root.real)
+        s_score = -np.dot(p1 - rvw_dtau[0], p2 - p1) / np.dot(p2 - p1, p2 - p1)
+
+        if not (0 <= s_score <= 1):
+            continue
+
+        if root.real < min_time:
+            min_time = root.real
+
+    return min_time
+
+
+@jit(nopython=True, cache=const.use_numba_cache)
+def ball_vertical_cylinder_collision_time(
+    rvw: NDArray[np.float64],
+    s: int,
+    a: float,
+    b: float,
+    r: float,
+    mu: float,
+    m: float,
+    g: float,
+    R: float,
+) -> float:
+    """Get the time until collision between a ball and a vertical cylinder.
+
+    For ball trajectories limited to the playing surface, this suffices for
+    detecting ball collisions with circular cushion segments.
+
+    Note:
+        - This is broken for airborne balls.
+    """
+
+    if s == const.spinning or s == const.pocketed or s == const.stationary:
+        return np.inf
+
+    phi = ptmath.angle(rvw[1])
+    v = ptmath.norm3d(rvw[1])
+
+    u = get_u_vec(rvw, R, phi, s)
+
+    K = -0.5 * mu * g
+    cos_phi = np.cos(phi)
+    sin_phi = np.sin(phi)
+
+    ax = K * (u[0] * cos_phi - u[1] * sin_phi)
+    ay = K * (u[0] * sin_phi + u[1] * cos_phi)
+    bx, by = v * cos_phi, v * sin_phi
+    cx, cy = rvw[0, 0], rvw[0, 1]
+
+    A = 0.5 * (ax * ax + ay * ay)
+    B = ax * bx + ay * by
+    C = ax * (cx - a) + ay * (cy - b) + 0.5 * (bx * bx + by * by)
+    D = bx * (cx - a) + by * (cy - b)
+    E = 0.5 * (a * a + b * b + cx * cx + cy * cy - (r + R) * (r + R)) - (
+        cx * a + cy * b
+    )
+
+    return get_real_positive_smallest_root(quartic.solve(A, B, C, D, E))
+
+
+def get_next_ball_linear_cushion_event(
     shot: System, collision_cache: CollisionCache
 ) -> Event:
     """Detect the next ball-vs-linear-cushion collision in 2D mode."""
@@ -41,20 +173,24 @@ def get_next_ball_linear_cushion_2d_event(
                 cache[obj_ids] = np.inf
                 continue
 
-            dtau_E = ball_linear_cushion_collision_time(
-                rvw=state.rvw,
-                s=state.s,
-                lx=cushion.lx,
-                ly=cushion.ly,
-                l0=cushion.l0,
-                p1=cushion.p1,
-                p2=cushion.p2,
-                direction=cushion.direction,
-                mu=(params.u_s if state.s == const.sliding else params.u_r),
-                m=params.m,
-                g=params.g,
-                R=params.R,
-            )
+            if ball.state.s == const.airborne:
+                # TODO
+                dtau_E = np.inf
+            else:
+                dtau_E = ball_vertical_plane_collision_time(
+                    rvw=state.rvw,
+                    s=state.s,
+                    lx=cushion.lx,
+                    ly=cushion.ly,
+                    l0=cushion.l0,
+                    p1=cushion.p1,
+                    p2=cushion.p2,
+                    direction=cushion.direction,
+                    mu=(params.u_s if state.s == const.sliding else params.u_r),
+                    m=params.m,
+                    g=params.g,
+                    R=params.R,
+                )
 
             cache[obj_ids] = shot.t + dtau_E
 
@@ -67,14 +203,7 @@ def get_next_ball_linear_cushion_2d_event(
     )
 
 
-def get_next_ball_linear_cushion_3d_event(
-    shot: System, collision_cache: CollisionCache
-) -> Event:
-    """3D ball-linear-cushion detection — not vendored yet; emits no event."""
-    return null_event(np.inf)
-
-
-def get_next_ball_circular_cushion_2d_event(
+def get_next_ball_circular_cushion_event(
     shot: System, collision_cache: CollisionCache
 ) -> Event:
     """Detect the next ball-vs-circular-cushion collision in 2D mode."""
@@ -97,17 +226,22 @@ def get_next_ball_circular_cushion_2d_event(
                 cache[obj_ids] = np.inf
                 continue
 
-            dtau_E = ball_circular_cushion_collision_time(
-                rvw=state.rvw,
-                s=state.s,
-                a=cushion.a,
-                b=cushion.b,
-                r=cushion.radius,
-                mu=(params.u_s if state.s == const.sliding else params.u_r),
-                m=params.m,
-                g=params.g,
-                R=params.R,
-            )
+            if ball.state.s == const.airborne:
+                # TODO
+                dtau_E = np.inf
+            else:
+                dtau_E = ball_vertical_cylinder_collision_time(
+                    rvw=state.rvw,
+                    s=state.s,
+                    a=cushion.a,
+                    b=cushion.b,
+                    r=cushion.radius,
+                    mu=(params.u_s if state.s == const.sliding else params.u_r),
+                    m=params.m,
+                    g=params.g,
+                    R=params.R,
+                )
+
             cache[obj_ids] = shot.t + dtau_E
 
     ball_id, cushion_id = min(cache, key=lambda k: cache[k])
@@ -117,9 +251,3 @@ def get_next_ball_circular_cushion_2d_event(
         cushion=shot.table.cushion_segments.circular[cushion_id],
         time=cache[(ball_id, cushion_id)],
     )
-
-
-def get_next_ball_circular_cushion_3d_event(
-    shot: System, collision_cache: CollisionCache
-) -> Event:
-    return null_event(np.inf)