Merge branch 'main' into ek/recursive-modelargs

docs/examples/30_degree_rule.pct.py

@@ -75,7 +75,7 @@
 #
 # We'll start with a table. Since we don't want collisions with cushions to interfere with our trajectory, let's make an unrealistically large $10\text{m} \times 10\text{m}$  [Table](../autoapi/pooltool/index.rst#pooltool.Table).
 
-# %% trusted=true
+# %%
 import pooltool as pt
 
 table_specs = pt.objects.BilliardTableSpecs(l=10, w=10)
@@ -85,20 +85,20 @@
 # %% [markdown]
 # Next, we'll create two [Ball](../autoapi/pooltool/index.rst#pooltool.Ball) objects.
 
-# %% trusted=true
+# %%
 cue_ball = pt.Ball.create("cue", xy=(2.5, 1.5))
 obj_ball = pt.Ball.create("obj", xy=(2.5, 3.0))
 
 # %% [markdown]
 # Next, we'll need a [Cue](../autoapi/pooltool/index.rst#pooltool.Cue).
 
-# %% trusted=true
+# %%
 cue = pt.Cue(cue_ball_id="cue")
 
 # %% [markdown]
 # Finally, we'll need to wrap these objects up into a [System](../autoapi/pooltool/index.rst#pooltool.System). We'll call this our system *template*, with the intention of reusing it for many different shots.
 
-# %% trusted=true
+# %%
 system_template = pt.System(
     table=table,
     cue=cue,
@@ -112,7 +112,7 @@
 #
 # So in the function call below, `pt.aim.at_ball(system, "obj", cut=30)` returns the angle `phi` that the cue ball should be directed at such that a cut angle of 30 degrees with the object ball is achieved.
 
-# %% trusted=true
+# %%
 # Creates a deep copy of the template
 system = system_template.copy()
 
@@ -144,7 +144,7 @@
 #
 # Since that can't be embedded into the documentation, we'll instead plot the trajectory of the cue ball and object ball by accessing ther historical states.
 
-# %% trusted=true
+# %%
 cue_ball = system.balls["cue"]
 obj_ball = system.balls["obj"]
 cue_history = cue_ball.history_cts
@@ -154,7 +154,7 @@
 # %% [markdown]
 # The [BallHistory](../autoapi/pooltool/objects/index.rst#pooltool.objects.BallHistory) holds the ball's historical states, each stored as a [BallState](../autoapi/pooltool/objects/index.rst#pooltool.objects.BallState) object. Each attribute of the ball states can be concatenated into numpy arrays with the [BallHistory.vectorize](../autoapi/pooltool/objects/index.rst#pooltool.objects.BallHistory.vectorize) method.
 
-# %% trusted=true
+# %%
 rvw_cue, s_cue, t_cue = cue_history.vectorize()
 rvw_obj, s_obj, t_obj = obj_history.vectorize()
 
@@ -165,12 +165,12 @@
 # %% [markdown]
 # We can grab the xy-coordinates from the `rvw` array by with the following.
 
-# %% trusted=true
+# %%
 coords_cue = rvw_cue[:, 0, :2]
 coords_obj = rvw_obj[:, 0, :2]
 coords_cue.shape
 
-# %% trusted=true editable=true slideshow={"slide_type": ""} tags=[]
+# %% editable=true slideshow={"slide_type": ""} tags=[]
 import plotly.graph_objects as go
 import plotly.io as pio
 
@@ -203,28 +203,28 @@
 #
 # As mentioned before, the carom angle is the angle between the cue ball velocity right before collision, and the cue ball velocity post-collision, once the ball has stopped sliding on the cloth. Hidden somewhere in the system **event list** one can find the events corresponding to these precise moments in time:
 
-# %% trusted=true
+# %%
 system.events[:6]
 
 # %% [markdown]
 # Programatically, we can pick out these two events of interest with event selection syntax.
 #
 # Since there is only one ball-ball collision, it's easy to select with [filter_type](../autoapi/pooltool/events/index.rst#pooltool.events.filter_type):
 
-# %% trusted=true
+# %%
 collision = pt.events.filter_type(system.events, pt.EventType.BALL_BALL)[0]
 collision
 
 # %% [markdown]
 # To get the event when the cue ball stops sliding, we can similarly try filtering by the sliding to rolling transition event:
 
-# %% trusted=true
+# %%
 pt.events.filter_type(system.events, pt.EventType.SLIDING_ROLLING)
 
 # %% [markdown]
 # But there are many sliding to rolling transition events, and to make matters worse, they are shared by both the cue ball and the object ball. What we need is the **first** **sliding to rolling** transition that the **cue ball** undergoes **after** the **ball-ball** collision. We can achieve this multi-criteria query with [filter_events](../autoapi/pooltool/events/index.rst#pooltool.events.filter_events):
 
-# %% trusted=true
+# %%
 transition = pt.events.filter_events(
     system.events,
     pt.events.by_time(t=collision.time, after=True),
@@ -236,16 +236,13 @@
 # %% [markdown]
 # Now, we can dive into these two events and pull out the cue ball velocities we need to calculate the carom angle.
 
-# %% trusted=true
+# %%
 # Velocity prior to impact
-for agent in collision.agents:
-    if agent.id == "cue":
-        # agent.initial is a copy of the Ball before resolving the collision
-        velocity_initial = agent.initial.state.rvw[1, :2]
+velocity_initial = collision.get_ball("cue", initial=True).vel[:2]
 
 # Velocity post sliding
-# We choose `final` here for posterity, but the velocity is the same both before and after resolving the transition.
-velocity_final = transition.agents[0].final.state.rvw[1, :2]
+# We choose the "final" here for posterity, but the velocity is the same both before and after resolving the transition.
+velocity_final = transition.get_ball("cue", initial=False).vel[:2]
 
 carom_angle = pt.ptmath.utils.angle_between_vectors(velocity_final, velocity_initial)
 
@@ -259,7 +256,7 @@
 # We calculated the carom angle for a single cut angle, 30 degrees. Let's write a function called `get_carom_angle` so we can do that repeatedly for different cut angles.
 
 
-# %% trusted=true
+# %%
 def get_carom_angle(system: pt.System) -> float:
     assert system.simulated
 
@@ -283,7 +280,7 @@ def get_carom_angle(system: pt.System) -> float:
 # `get_carom_angle` assumes the passed system has already been simulated, so we'll need another function to take care of that. We'll cue stick speed and cut angle as parameters.
 
 
-# %% trusted=true
+# %%
 def simulate_experiment(V0: float, cut_angle: float) -> pt.System:
     system = system_template.copy()
     phi = pt.aim.at_ball(system, "obj", cut=cut_angle)
@@ -296,7 +293,7 @@ def simulate_experiment(V0: float, cut_angle: float) -> pt.System:
 # We'll also want the ball hit fraction:
 
 
-# %% trusted=true
+# %%
 import numpy as np
 
 def get_ball_hit_fraction(cut_angle: float) -> float:
@@ -306,7 +303,7 @@ def get_ball_hit_fraction(cut_angle: float) -> float:
 # %% [markdown]
 # With these functions, we are ready to simulate how carom angle varies as a function of cut angle.
 
-# %% trusted=true
+# %%
 import pandas as pd
 
 data = {
@@ -329,7 +326,7 @@ def get_ball_hit_fraction(cut_angle: float) -> float:
 # %% [markdown]
 # From this dataframe we can make some plots. On top of the ball-hit fraction, plot, I'll create a box between a $1/4$ ball hit and a $3/4$ ball hit, since this is the carom angle range that the 30-degree rule is defined with respect to.
 
-# %% trusted=true editable=true slideshow={"slide_type": ""} tags=["nbsphinx-thumbnail"]
+# %% editable=true slideshow={"slide_type": ""} tags=["nbsphinx-thumbnail"]
 import matplotlib.pyplot as plt
 
 x_min = 0.25
@@ -354,7 +351,7 @@ def get_ball_hit_fraction(cut_angle: float) -> float:
 #
 # For your reference, here is the same plot but with cut angle $\phi$ as the x-axis:
 
-# %% trusted=true
+# %%
 fig, ax = plt.subplots()
 ax.scatter(frame['phi'], frame['theta'], color='#1f77b4')
 ax.set_title('Carom Angle vs Cut Angle', fontsize=20)
@@ -378,7 +375,7 @@ def get_ball_hit_fraction(cut_angle: float) -> float:
 # Since pooltool's baseline physics engine makes the same assumptions, we should expect the results to be the same. Let's directly compare:
 
 
-# %% trusted=true
+# %%
 def get_theoretical_carom_angle(phi) -> float:
     return np.atan2(np.sin(phi) * np.cos(phi), (np.sin(phi) ** 2 + 2 / 5))
 
@@ -425,7 +422,7 @@ def get_theoretical_carom_angle(phi) -> float:
 #
 # Interestingly, the carom angle is independent of the speed:
 
-# %% trusted=true
+# %%
 for V0 in np.linspace(1, 4, 20):
     system = simulate_experiment(V0, 30)
     carom_angle = get_carom_angle(system)
@@ -434,7 +431,7 @@ def get_theoretical_carom_angle(phi) -> float:
 # %% [markdown]
 # This doesn't mean that the trajectories are the same though. Here are the trajectories:
 
-# %% trusted=true
+# %%
 import numpy as np
 import plotly.graph_objects as go
 

docs/examples/straight_shot.pct.py

@@ -115,7 +115,7 @@ def create_system(d, D):
     obj_ball = create_object_ball(cue_ball, d)
 
     return pt.System(
-        cue=pt.Cue.default(),
+        cue=pt.Cue(cue_ball_id="CB"),
         balls=(cue_ball, obj_ball),
         table=table,
     )

pooltool/ani/animate.py

@@ -506,7 +506,8 @@ def create_system(self):
 
         game = get_ruleset(game_type)()
         game.players = [
-            Player("Player"),
+            Player("Player 1"),
+            Player("Player 2"),
         ]
 
         table = Table.from_game_type(game_type)

pooltool/events/datatypes.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import Any, Dict, Optional, Tuple, Type, Union
+from typing import Any, Dict, Optional, Tuple, Type, Union, cast
 
 from attrs import define, evolve, field
 from cattrs.converters import Converter
@@ -75,6 +75,35 @@ def is_transition(self) -> bool:
             EventType.SLIDING_ROLLING,
         }
 
+    def has_ball(self) -> bool:
+        """Returns True if this event type can involve a Ball."""
+        return (
+            self
+            in {
+                EventType.BALL_BALL,
+                EventType.BALL_LINEAR_CUSHION,
+                EventType.BALL_CIRCULAR_CUSHION,
+                EventType.BALL_POCKET,
+                EventType.STICK_BALL,
+            }
+            or self.is_transition()
+        )
+
+    def has_cushion(self) -> bool:
+        """Returns True if this event type can involve a cushion (linear or circular)."""
+        return self in {
+            EventType.BALL_LINEAR_CUSHION,
+            EventType.BALL_CIRCULAR_CUSHION,
+        }
+
+    def has_pocket(self) -> bool:
+        """Returns True if this event type can involve a Pocket."""
+        return self == EventType.BALL_POCKET
+
+    def has_stick(self) -> bool:
+        """Returns True if this event type can involve a CueStick."""
+        return self == EventType.STICK_BALL
+
 
 Object = Union[
     NullObject,
@@ -185,22 +214,6 @@ def set_final(self, obj: Object) -> None:
         else:
             self.final = obj.copy()
 
-    def matches(self, obj: Object) -> bool:
-        """Determines if the given object matches the agent.
-
-        It checks if the object is of the correct class type and if the IDs match.
-
-        Args:
-            obj: The object to compare with the agent.
-
-        Returns:
-            bool:
-                True if the object's class type and ID match the agent's type and ID,
-                False otherwise.
-        """
-        correct_class = _class_to_type[type(obj)] == self.agent_type
-        return correct_class and obj.id == self.id
-
     @staticmethod
     def from_object(obj: Object, set_initial: bool = False) -> Agent:
         """Creates an agent instance from an object.
@@ -228,6 +241,17 @@ def copy(self) -> Agent:
         """Create a copy."""
         return evolve(self)
 
+    def _get_state(self, initial: bool) -> Object:
+        """Return either the initial or final state of the given agent.
+
+        Raises ValueError if that state is None.
+        """
+        obj = self.initial if initial else self.final
+        if obj is None:
+            which = "initial" if initial else "final"
+            raise ValueError(f"Agent '{self.id}' has no {which} state in this event.")
+        return obj
+
 
 def _disambiguate_agent_structuring(
     uo: Dict[str, Any], _: Type[Agent], con: Converter
@@ -329,3 +353,81 @@ def copy(self) -> Event:
         """Create a copy."""
         # NOTE is this deep-ish copy?
         return evolve(self)
+
+    def _find_agent(self, agent_type: AgentType, agent_id: str) -> Agent:
+        """Return the Agent with the specified agent_type and ID.
+
+        Raises:
+            ValueError if not found.
+        """
+        for agent in self.agents:
+            if agent.agent_type == agent_type and agent.id == agent_id:
+                return agent
+        raise ValueError(
+            f"No agent of type {agent_type} with ID '{agent_id}' found in this event."
+        )
+
+    def get_ball(self, ball_id: str, initial: bool = True) -> Ball:
+        """Return the Ball object with the given ID, either final or initial.
+
+        Args:
+            ball_id: The ID of the ball to retrieve.
+            initial: If True, return the ball's initial state; otherwise final state.
+
+        Raises:
+            ValueError: If the event does not involve a ball or if no matching ball is found.
+        """
+        if not self.event_type.has_ball():
+            raise ValueError(
+                f"Event of type {self.event_type} does not involve a Ball."
+            )
+
+        agent = self._find_agent(AgentType.BALL, ball_id)
+        obj = agent._get_state(initial)
+        return cast(Ball, obj)
+
+    def get_pocket(self, pocket_id: str, initial: bool = True) -> Pocket:
+        """Return the Pocket object with the given ID, either final or initial."""
+        if not self.event_type.has_pocket():
+            raise ValueError(
+                f"Event of type {self.event_type} does not involve a Pocket."
+            )
+
+        agent = self._find_agent(AgentType.POCKET, pocket_id)
+        obj = agent._get_state(initial)
+        return cast(Pocket, obj)
+
+    def get_cushion(
+        self, cushion_id: str
+    ) -> Union[LinearCushionSegment, CircularCushionSegment]:
+        """Return the cushion segment with the given ID."""
+        if not self.event_type.has_cushion():
+            raise ValueError(
+                f"Event of type {self.event_type} does not involve a cushion."
+            )
+
+        try:
+            agent = self._find_agent(AgentType.LINEAR_CUSHION_SEGMENT, cushion_id)
+            return cast(LinearCushionSegment, agent.initial)
+        except ValueError:
+            pass
+
+        try:
+            agent = self._find_agent(AgentType.CIRCULAR_CUSHION_SEGMENT, cushion_id)
+            return cast(CircularCushionSegment, agent.initial)
+        except ValueError:
+            pass
+
+        raise ValueError(
+            f"No agent of linear/circular cushion with ID '{cushion_id}' found in this event."
+        )
+
+    def get_stick(self, stick_id: str) -> Pocket:
+        """Return the cue stick with the given ID."""
+        if not self.event_type.has_pocket():
+            raise ValueError(
+                f"Event of type {self.event_type} does not involve a Pocket."
+            )
+
+        agent = self._find_agent(AgentType.POCKET, stick_id)
+        return cast(Pocket, agent.initial)