Merge branch 'main' into ek/mathavan-cushion

pooltool/ani/__init__.py

@@ -34,7 +34,7 @@
 # max backstroke length, as fraction of cue stick length
 backstroke_fraction = 0.5
 # max masse angle
-max_elevate = 80
+max_elevate = 89.9
 # number of seconds that camera rotation is disabled when shot is being calculated
 rotate_downtime = 0.3
 

pooltool/ani/collision.py

@@ -147,12 +147,14 @@ def process_ball_collision(self, entry):
         phi = ((self.avoid_nodes["cue_stick_focus"].getH() + 180) % 360) * np.pi / 180
         c = np.array([np.cos(phi), np.sin(phi), 0])
         gamma = np.arccos(np.dot(n, c))
-        AB = (ball.params.R + multisystem.active.cue.specs.tip_radius) * np.cos(gamma)
+        AB = (
+            ball.params.R + multisystem.active.cue.specs.shaft_radius_at_tip
+        ) * np.cos(gamma)
 
         # Center of blocking ball transect
         Ax, Ay, _ = entry.getSurfacePoint(scene)
-        Ax -= (AB + multisystem.active.cue.specs.tip_radius) * np.cos(phi)
-        Ay -= (AB + multisystem.active.cue.specs.tip_radius) * np.sin(phi)
+        Ax -= (AB + multisystem.active.cue.specs.shaft_radius_at_tip) * np.cos(phi)
+        Ay -= (AB + multisystem.active.cue.specs.shaft_radius_at_tip) * np.sin(phi)
         Az = ball.params.R
 
         # Center of aim, leveled to ball height
@@ -191,8 +193,8 @@ def get_cue_radius(self, length):
         bounds = visual.cue.get_node("cue_stick").get_tight_bounds()
         L = bounds[1][0] - bounds[0][0]  # cue length
 
-        r = multisystem.active.cue.specs.tip_radius
-        R = multisystem.active.cue.specs.butt_radius
+        r = multisystem.active.cue.specs.shaft_radius_at_tip
+        R = multisystem.active.cue.specs.shaft_radius_at_butt
 
         m = (R - r) / L  # rise/run
         b = r  # intercept

pooltool/ani/hud.py

@@ -2,8 +2,9 @@
 
 from abc import ABC, abstractmethod
 from collections import deque
-from typing import List
+from typing import List, Optional
 
+import numpy as np
 from direct.gui.OnscreenImage import OnscreenImage
 from direct.interval.LerpInterval import LerpFunc
 from panda3d.core import CardMaker, NodePath, TextNode, TransparencyAttrib
@@ -12,7 +13,9 @@
 import pooltool.ani.tasks as tasks
 import pooltool.ani.utils as autils
 from pooltool.ani.globals import Global
-from pooltool.objects.cue.datatypes import Cue
+from pooltool.objects.ball.datatypes import Ball, BallParams
+from pooltool.objects.cue.datatypes import Cue, CueSpecs
+from pooltool.ptmath.utils import tip_center_offset
 from pooltool.utils import panda_path
 from pooltool.utils.strenum import StrEnum, auto
 
@@ -69,7 +72,7 @@ def toggle_help(self):
             return
         self.elements[HUDElement.help_text].toggle()
 
-    def update_cue(self, cue: Cue):
+    def update_cue(self, cue: Cue, cue_ball: Optional[Ball] = None):
         """Update HUD to reflect english, jack, and power of cue
 
         Returns silently if HUD is not initialized.
@@ -78,6 +81,15 @@ def update_cue(self, cue: Cue):
         if not self.initialized:
             return
 
+        if cue_ball is not None:
+            tip_offset_a, tip_offset_b = tip_center_offset(
+                np.array([cue.a, cue.b]), cue.specs.tip_radius, cue_ball.params.R
+            )
+            self.elements[HUDElement.english].set_tip_center(tip_offset_a, tip_offset_b)
+            self.elements[HUDElement.english].set_shaft_to_ball_diameter_ratio(
+                cue.specs.shaft_radius_at_tip / cue_ball.params.R
+            )
+
         self.elements[HUDElement.english].set(cue.a, cue.b)
         self.elements[HUDElement.jack].set(cue.theta)
         self.elements[HUDElement.power].set(cue.V0)
@@ -334,17 +346,25 @@ def __init__(self):
         BaseHUDElement.__init__(self)
         self.dir = ani.model_dir / "hud" / "english"
         self.text_scale = 0.2
+        self.ball_scale = 0.15
         self.text_color = (1, 1, 1, 1)
 
         self.circle = OnscreenImage(
             image=panda_path(self.dir / "circle.png"),
             parent=self.dummy_right,
-            scale=0.15,
+            scale=self.ball_scale,
         )
         self.circle.setTransparency(TransparencyAttrib.MAlpha)
         autils.alignTo(self.circle, self.dummy_right, autils.CL, autils.C)
         self.circle.setZ(-0.65)
 
+        self.tip_circle = OnscreenImage(
+            image=panda_path(self.dir / "tip-outline.png"),
+            parent=self.circle,
+            scale=CueSpecs.default().shaft_radius_at_tip / BallParams.default().R,
+        )
+        self.tip_circle.setTransparency(TransparencyAttrib.MAlpha)
+
         self.crosshairs = OnscreenImage(
             image=panda_path(self.dir / "crosshairs.png"),
             pos=(0, 0, 0),
@@ -367,6 +387,12 @@ def set(self, a, b):
         self.crosshairs.setPos(-a, 0, b)
         self.text.setText(f"({a:.3f},{b:.3f})")
 
+    def set_shaft_to_ball_diameter_ratio(self, ratio):
+        self.tip_circle.setScale(ratio)
+
+    def set_tip_center(self, a, b):
+        self.tip_circle.setPos(-a, 0, b)
+
     def init(self):
         self.show()
 

pooltool/ani/image/interface.py

@@ -58,7 +58,7 @@ def image_stack(
     if show_hud:
         hud.init()
         hud.elements[HUDElement.help_text].help_hint.hide()
-        hud.update_cue(system.cue)
+        hud.update_cue(system.cue, system.balls[system.cue.cue_ball_id])
     else:
         hud.destroy()
 

pooltool/ani/modes/aim.py

@@ -12,6 +12,7 @@
 from pooltool.ani.modes.datatypes import BaseMode, Mode
 from pooltool.ani.modes.shot import ShotMode
 from pooltool.ani.mouse import MouseMode, mouse
+from pooltool.ptmath.utils import norm2d, tip_contact_offset
 from pooltool.system.datatypes import multisystem
 from pooltool.system.render import visual
 
@@ -172,9 +173,11 @@ def cue_avoidance(self):
 
         if (theta < cue_avoid.min_theta) or self.magnet_theta:
             theta = cue_avoid.min_theta
-            multisystem.active.cue.set_state(theta=theta)
+            system_cue = multisystem.active.cue
+            system_cue.set_state(theta=theta)
+            system_cue_ball = multisystem.active.balls[system_cue.cue_ball_id]
             visual.cue.set_render_state_as_object_state()
-            hud.update_cue(multisystem.active.cue)
+            hud.update_cue(system_cue, system_cue_ball)
 
         if cam.theta < theta + ani.min_camera:
             cam.rotate(theta=theta + ani.min_camera)
@@ -194,8 +197,9 @@ def aim_apply_power(self):
         if V0 > ani.max_stroke_speed:
             V0 = ani.max_stroke_speed
 
-        multisystem.active.cue.set_state(V0=V0)
-        hud.update_cue(multisystem.active.cue)
+        system_cue = multisystem.active.cue
+        system_cue.set_state(V0=V0)
+        hud.update_cue(system_cue, multisystem.active.balls[system_cue.cue_ball_id])
 
     def aim_elevate_cue(self):
         cue = visual.cue.get_node("cue_stick_focus")
@@ -219,8 +223,9 @@ def aim_elevate_cue(self):
         if cam.theta < (new_elevation + ani.min_camera):
             cam.rotate(theta=new_elevation + ani.min_camera)
 
-        multisystem.active.cue.set_state(theta=new_elevation)
-        hud.update_cue(multisystem.active.cue)
+        system_cue = multisystem.active.cue
+        system_cue.set_state(theta=new_elevation)
+        hud.update_cue(system_cue, multisystem.active.balls[system_cue.cue_ball_id])
 
     def apply_english(self):
         with mouse:
@@ -237,10 +242,20 @@ def apply_english(self):
         new_y = cue.getY() + delta_y
         new_z = cue.getZ() + delta_z
 
-        norm = np.sqrt(new_y**2 + new_z**2)
-        if norm > ani.max_english * R:
-            new_y *= ani.max_english * R / norm
-            new_z *= ani.max_english * R / norm
+        cue_axis_offset = (
+            np.array([-new_y, new_z]) / R
+        )  # components normalized to ball radius
+        contact_point_offset = tip_contact_offset(
+            cue_axis_offset, multisystem.active.cue.specs.tip_radius, R
+        )
+
+        norm = norm2d(contact_point_offset)
+        if norm > ani.max_english:
+            limit_scaling_factor = ani.max_english / norm
+            new_y *= limit_scaling_factor
+            new_z *= limit_scaling_factor
+            cue_axis_offset *= limit_scaling_factor
+            contact_point_offset *= limit_scaling_factor
 
         cue.setY(new_y)
         cue.setZ(new_z)
@@ -257,9 +272,10 @@ def apply_english(self):
             cam.rotate(theta=new_theta)
 
         multisystem.active.cue.set_state(
-            a=-new_y / R,
-            b=new_z / R,
+            a=contact_point_offset[0],
+            b=contact_point_offset[1],
             theta=-cue_focus.getR(),
         )
 
-        hud.update_cue(multisystem.active.cue)
+        system_cue = multisystem.active.cue
+        hud.update_cue(system_cue, multisystem.active.balls[system_cue.cue_ball_id])

pooltool/ani/modes/shot.py

@@ -63,7 +63,8 @@ def enter(
         if playback_mode is not None:
             visual.animate(playback_mode)
 
-        hud.update_cue(multisystem.active.cue)
+        system_cue = multisystem.active.cue
+        hud.update_cue(system_cue, multisystem.active.balls[system_cue.cue_ball_id])
 
         tasks.register_event("space", visual.toggle_pause)
         tasks.register_event("arrow_up", visual.speed_up)
@@ -166,7 +167,8 @@ def exit(self, key="soft"):
             visual.cue.set_render_state_as_object_state()
 
             # Set the HUD
-            hud.update_cue(multisystem.active.cue)
+            system_cue = multisystem.active.cue
+            hud.update_cue(system_cue, multisystem.active.balls[system_cue.cue_ball_id])
 
         elif key == "reset":
             if multisystem.active_index != len(multisystem) - 1:
@@ -298,4 +300,5 @@ def change_animation(shot_index):
         cue_avoid.init_collisions()
 
         # Set the HUD
-        hud.update_cue(multisystem.active.cue)
+        system_cue = multisystem.active.cue
+        hud.update_cue(system_cue, multisystem.active.balls[system_cue.cue_ball_id])

pooltool/ani/modes/view.py

@@ -12,6 +12,7 @@
 from pooltool.ani.modes.datatypes import BaseMode, Mode
 from pooltool.ani.mouse import MouseMode, mouse
 from pooltool.objects.ball.datatypes import Ball
+from pooltool.ptmath.utils import norm2d, tip_contact_offset
 from pooltool.system.datatypes import multisystem
 from pooltool.system.render import PlaybackMode, visual
 
@@ -159,8 +160,9 @@ def view_apply_power(self):
         if V0 > ani.max_stroke_speed:
             V0 = ani.max_stroke_speed
 
-        multisystem.active.cue.set_state(V0=V0)
-        hud.update_cue(multisystem.active.cue)
+        system_cue = multisystem.active.cue
+        system_cue.set_state(V0=V0)
+        hud.update_cue(system_cue, multisystem.active.balls[system_cue.cue_ball_id])
 
     def view_elevate_cue(self):
         visual.cue.show_nodes(ignore=("cue_cseg",))
@@ -183,8 +185,9 @@ def view_elevate_cue(self):
 
         cue.setR(-new_elevation)
 
-        multisystem.active.cue.set_state(theta=new_elevation)
-        hud.update_cue(multisystem.active.cue)
+        system_cue = multisystem.active.cue
+        system_cue.set_state(theta=new_elevation)
+        hud.update_cue(system_cue, multisystem.active.balls[system_cue.cue_ball_id])
 
     def view_apply_english(self):
         visual.cue.show_nodes(ignore=("cue_cseg",))
@@ -202,10 +205,21 @@ def view_apply_english(self):
         new_y = cue.getY() + delta_y
         new_z = cue.getZ() + delta_z
 
-        norm = np.sqrt(new_y**2 + new_z**2)
-        if norm > ani.max_english * R:
-            new_y *= ani.max_english * R / norm
-            new_z *= ani.max_english * R / norm
+        # y corresponds to side spin, z to top/bottom spin
+        cue_axis_offset = (
+            np.array([-new_y, new_z]) / R
+        )  # components normalized to ball radius
+        contact_point_offset = tip_contact_offset(
+            cue_axis_offset, multisystem.active.cue.specs.tip_radius, R
+        )
+
+        norm = norm2d(contact_point_offset)
+        if norm > ani.max_english:
+            limit_scaling_factor = ani.max_english / norm
+            new_y *= limit_scaling_factor
+            new_z *= limit_scaling_factor
+            cue_axis_offset *= limit_scaling_factor
+            contact_point_offset *= limit_scaling_factor
 
         cue.setY(new_y)
         cue.setZ(new_z)
@@ -219,12 +233,13 @@ def view_apply_english(self):
             cue_focus.setR(-cue_avoid.min_theta)
 
         multisystem.active.cue.set_state(
-            a=-new_y / R,
-            b=new_z / R,
+            a=contact_point_offset[0],
+            b=contact_point_offset[1],
             theta=-visual.cue.get_node("cue_stick_focus").getR(),
         )
 
-        hud.update_cue(multisystem.active.cue)
+        system_cue = multisystem.active.cue
+        hud.update_cue(system_cue, multisystem.active.balls[system_cue.cue_ball_id])
 
     def change_animation(self, shot_index):
         """Switch to a different system in the system collection"""
@@ -256,4 +271,5 @@ def change_animation(self, shot_index):
         cue_avoid.init_collisions()
 
         # Set the HUD
-        hud.update_cue(multisystem.active.cue)
+        system_cue = multisystem.active.cue
+        hud.update_cue(system_cue, multisystem.active.balls[system_cue.cue_ball_id])

pooltool/objects/ball/render.py

@@ -117,7 +117,7 @@ def init_collision(self, cue: Cue):
             CollisionNode(f"ball_csphere_{self._ball.id}")
         )
         collision_node.node().addSolid(
-            CollisionCapsule(0, 0, -R, 0, 0, R, cue.specs.tip_radius + R)
+            CollisionCapsule(0, 0, -R, 0, 0, R, cue.specs.shaft_radius_at_tip + R)
         )
         if ani.settings["graphics"]["debug"]:
             collision_node.show()