Refactor Resolver serialization + ball ball friction sub models

pooltool/physics/__init__.py

@@ -5,37 +5,29 @@
 from pooltool.physics.resolve.ball_ball import (
     BallBallCollisionStrategy,
     BallBallModel,
-    get_ball_ball_model,
 )
 from pooltool.physics.resolve.ball_cushion import (
     BallCCushionCollisionStrategy,
     BallCCushionModel,
     BallLCushionCollisionStrategy,
     BallLCushionModel,
-    get_ball_circ_cushion_model,
-    get_ball_lin_cushion_model,
 )
 from pooltool.physics.resolve.ball_pocket import (
     BallPocketModel,
     BallPocketStrategy,
-    get_ball_pocket_model,
 )
 from pooltool.physics.resolve.resolver import (
-    RESOLVER_CONFIG_PATH,
+    RESOLVER_PATH,
     Resolver,
-    ResolverConfig,
 )
 from pooltool.physics.resolve.stick_ball import (
     StickBallCollisionStrategy,
     StickBallModel,
-    get_stick_ball_model,
 )
 from pooltool.physics.resolve.transition import (
     BallTransitionModel,
     BallTransitionStrategy,
-    get_transition_model,
 )
-from pooltool.physics.resolve.types import ArgType, ModelArgs
 
 __all__ = [
     "BallBallCollisionStrategy",
@@ -48,20 +40,11 @@
     "evolve",
     "resolve",
     "Resolver",
-    "RESOLVER_CONFIG_PATH",
-    "ResolverConfig",
+    "RESOLVER_PATH",
     "BallBallModel",
-    "get_ball_ball_model",
     "BallCCushionModel",
     "BallLCushionModel",
-    "get_ball_circ_cushion_model",
-    "get_ball_lin_cushion_model",
     "BallPocketModel",
-    "get_ball_pocket_model",
     "StickBallModel",
-    "get_stick_ball_model",
     "BallTransitionModel",
-    "get_transition_model",
-    "ArgType",
-    "ModelArgs",
 ]

pooltool/physics/resolve/__init__.py

@@ -3,57 +3,43 @@
 from pooltool.physics.resolve.ball_ball import (
     BallBallCollisionStrategy,
     BallBallModel,
-    get_ball_ball_model,
 )
 from pooltool.physics.resolve.ball_cushion import (
     BallCCushionCollisionStrategy,
     BallCCushionModel,
     BallLCushionCollisionStrategy,
     BallLCushionModel,
-    get_ball_circ_cushion_model,
-    get_ball_lin_cushion_model,
 )
 from pooltool.physics.resolve.ball_pocket import (
     BallPocketModel,
     BallPocketStrategy,
-    get_ball_pocket_model,
 )
 from pooltool.physics.resolve.resolver import (
-    RESOLVER_CONFIG_PATH,
+    RESOLVER_PATH,
     Resolver,
-    ResolverConfig,
 )
 from pooltool.physics.resolve.stick_ball import (
     StickBallCollisionStrategy,
     StickBallModel,
-    get_stick_ball_model,
 )
 from pooltool.physics.resolve.transition import (
     BallTransitionModel,
     BallTransitionStrategy,
-    get_transition_model,
 )
 
 __all__ = [
     "Resolver",
-    "RESOLVER_CONFIG_PATH",
-    "ResolverConfig",
+    "RESOLVER_PATH",
     "BallBallCollisionStrategy",
     "BallBallModel",
-    "get_ball_ball_model",
     "BallCCushionCollisionStrategy",
     "BallCCushionModel",
     "BallLCushionCollisionStrategy",
     "BallLCushionModel",
-    "get_ball_circ_cushion_model",
-    "get_ball_lin_cushion_model",
     "BallPocketModel",
     "BallPocketStrategy",
-    "get_ball_pocket_model",
     "StickBallCollisionStrategy",
     "StickBallModel",
-    "get_stick_ball_model",
     "BallTransitionModel",
     "BallTransitionStrategy",
-    "get_transition_model",
 ]

pooltool/physics/resolve/ball_ball/__init__.py

@@ -1,63 +1,24 @@
 """Models for ball-ball collisions."""
 
-from typing import Dict, Optional, Type
+from typing import Dict, Type
 
 from pooltool.physics.resolve.ball_ball.core import BallBallCollisionStrategy
 from pooltool.physics.resolve.ball_ball.frictional_inelastic import FrictionalInelastic
 from pooltool.physics.resolve.ball_ball.frictional_mathavan import FrictionalMathavan
 from pooltool.physics.resolve.ball_ball.frictionless_elastic import FrictionlessElastic
-from pooltool.physics.resolve.types import ModelArgs
-from pooltool.utils.strenum import StrEnum, auto
+from pooltool.physics.resolve.models import BallBallModel
 
-
-class BallBallModel(StrEnum):
-    """An Enum for different ball-ball collision models
-
-    Attributes:
-        FRICTIONLESS_ELASTIC:
-            Frictionless, instantaneous, elastic, equal mass collision
-            (:class:`FrictionlessElastic`).
-    """
-
-    FRICTIONLESS_ELASTIC = auto()
-    FRICTIONAL_INELASTIC = auto()
-    FRICTIONAL_MATHAVAN = auto()
-
-
-_ball_ball_models: Dict[BallBallModel, Type[BallBallCollisionStrategy]] = {
+ball_ball_models: Dict[BallBallModel, Type[BallBallCollisionStrategy]] = {
     BallBallModel.FRICTIONLESS_ELASTIC: FrictionlessElastic,
     BallBallModel.FRICTIONAL_INELASTIC: FrictionalInelastic,
     BallBallModel.FRICTIONAL_MATHAVAN: FrictionalMathavan,
 }
 
 
-def get_ball_ball_model(
-    model: Optional[BallBallModel] = None, params: ModelArgs = {}
-) -> BallBallCollisionStrategy:
-    """Returns a ball-ball collision model
-
-    Args:
-        model:
-            An Enum specifying the desired model. If not passed,
-            :class:`FrictionalMathavan` is passed with empty params.
-        params:
-            A mapping of parameters accepted by the model.
-
-    Returns:
-        An instantiated model that satisfies the :class:`BallBallCollisionStrategy`
-        protocol.
-    """
-
-    if model is None:
-        return FrictionlessElastic()
-
-    return _ball_ball_models[model](**params)
-
-
 __all__ = [
     "BallBallModel",
-    "get_ball_ball_model",
     "FrictionalMathavan",
     "FrictionalInelastic",
     "FrictionlessElastic",
+    "ball_ball_models",
 ]

pooltool/physics/resolve/ball_ball/friction.py

@@ -0,0 +1,66 @@
+import math
+from typing import Dict, Protocol, Type
+
+import attrs
+import numpy as np
+
+import pooltool.ptmath as ptmath
+from pooltool.objects.ball.datatypes import Ball
+from pooltool.utils.strenum import StrEnum, auto
+
+
+class BallBallFrictionModel(StrEnum):
+    """An Enum for different ball-ball friction models"""
+
+    AVERAGE = auto()
+    ALCIATORE = auto()
+
+
+class BallBallFrictionStrategy(Protocol):
+    """Ball-ball friction models must satisfy this protocol"""
+
+    def calculate_friction(self, ball1: Ball, ball2: Ball) -> float:
+        """This method calculates ball-ball friction"""
+        ...
+
+
+@attrs.define
+class AlciatoreBallBallFriction:
+    """Friction fit curve u_b = a + b * exp(-c * v_rel) used in David Alciatore's TP A-14"""
+
+    a: float = 9.951e-3
+    b: float = 0.108
+    c: float = 1.088
+
+    model: BallBallFrictionModel = attrs.field(
+        default=BallBallFrictionModel.ALCIATORE, init=False
+    )
+
+    def calculate_friction(self, ball1: Ball, ball2: Ball) -> float:
+        unit_x = np.array([1.0, 0.0, 0.0])
+        v1_c = ptmath.surface_velocity(
+            ball1.state.rvw, unit_x, ball1.params.R
+        ) - np.array([ball1.state.rvw[1][0], 0, 0])
+        v2_c = ptmath.surface_velocity(
+            ball2.state.rvw, -unit_x, ball2.params.R
+        ) - np.array([ball2.state.rvw[1][0], 0, 0])
+        relative_surface_speed = ptmath.norm3d(v1_c - v2_c)
+        return self.a + self.b * math.exp(-self.c * relative_surface_speed)
+
+
+@attrs.define
+class AverageBallBallFriction:
+    model: BallBallFrictionModel = attrs.field(
+        default=BallBallFrictionModel.AVERAGE, init=False
+    )
+
+    def calculate_friction(self, ball1: Ball, ball2: Ball) -> float:
+        return (ball1.params.u_b + ball2.params.u_b) / 2
+
+
+ball_ball_friction_models: Dict[
+    BallBallFrictionModel, Type[BallBallFrictionStrategy]
+] = {
+    BallBallFrictionModel.AVERAGE: AverageBallBallFriction,
+    BallBallFrictionModel.ALCIATORE: AlciatoreBallBallFriction,
+}

pooltool/physics/resolve/ball_ball/frictional_inelastic/__init__.py

@@ -1,12 +1,18 @@
 from typing import Tuple
 
+import attrs
 import numpy as np
 from numba import jit
 
 import pooltool.constants as const
 import pooltool.ptmath as ptmath
 from pooltool.objects.ball.datatypes import Ball, BallState
 from pooltool.physics.resolve.ball_ball.core import CoreBallBallCollision
+from pooltool.physics.resolve.ball_ball.friction import (
+    AlciatoreBallBallFriction,
+    BallBallFrictionStrategy,
+)
+from pooltool.physics.resolve.models import BallBallModel
 
 
 @jit(nopython=True, cache=const.use_numba_cache)
@@ -103,6 +109,7 @@ def _resolve_ball_ball(rvw1, rvw2, R, u_b, e_b):
     return rvw1_f, rvw2_f
 
 
+@attrs.define
 class FrictionalInelastic(CoreBallBallCollision):
     """A simple ball-ball collision model including ball-ball friction, and coefficient of restitution for equal-mass balls
 
@@ -111,14 +118,20 @@ class FrictionalInelastic(CoreBallBallCollision):
     and a more complete analysis of velocity and angular velocity in their vector forms.
     """
 
+    friction: BallBallFrictionStrategy = AlciatoreBallBallFriction()
+
+    model: BallBallModel = attrs.field(
+        default=BallBallModel.FRICTIONAL_INELASTIC, init=False
+    )
+
     def solve(self, ball1: Ball, ball2: Ball) -> Tuple[Ball, Ball]:
         """Resolves the collision."""
         rvw1, rvw2 = _resolve_ball_ball(
             ball1.state.rvw.copy(),
             ball2.state.rvw.copy(),
             ball1.params.R,
-            # Assume the interaction coefficients are the average of the two balls
-            u_b=(ball1.params.u_b + ball2.params.u_b) / 2,
+            u_b=self.friction.calculate_friction(ball1, ball2),
+            # Average the coefficient of restitution parameters for the two balls
             e_b=(ball1.params.e_b + ball2.params.e_b) / 2,
         )
 

pooltool/physics/resolve/ball_ball/frictional_mathavan/__init__.py

@@ -1,5 +1,6 @@
 from typing import Optional, Tuple
 
+import attrs
 import numpy as np
 from numba import jit
 from numpy import array, dot, sqrt
@@ -8,6 +9,11 @@
 import pooltool.constants as const
 from pooltool.objects.ball.datatypes import Ball, BallState
 from pooltool.physics.resolve.ball_ball.core import CoreBallBallCollision
+from pooltool.physics.resolve.ball_ball.friction import (
+    AlciatoreBallBallFriction,
+    BallBallFrictionStrategy,
+)
+from pooltool.physics.resolve.models import BallBallModel
 
 INF = float("inf")
 Z_LOC = array([0, 0, 1], dtype=np.float64)
@@ -209,6 +215,7 @@ def _collide_balls(
     return dot(G_T, v_i), dot(G_T, w_i), dot(G_T, v_j), dot(G_T, w_j)
 
 
+@attrs.define
 class FrictionalMathavan(CoreBallBallCollision):
     """Ball-ball collision resolver for the Mathavan et al. (2014) collision model.
 
@@ -229,8 +236,12 @@ class FrictionalMathavan(CoreBallBallCollision):
         https://billiards.colostate.edu/physics_articles/Mathavan_Sports_2014.pdf
     """
 
-    def __init__(self, num_iterations: int = 1000):
-        self.num_iterations = num_iterations
+    friction: BallBallFrictionStrategy = AlciatoreBallBallFriction()
+    num_iterations: int = 1000
+
+    model: BallBallModel = attrs.field(
+        default=BallBallModel.FRICTIONAL_MATHAVAN, init=False
+    )
 
     def solve(self, ball1: Ball, ball2: Ball) -> Tuple[Ball, Ball]:
         """Resolve ball-ball collision via Mathavan et al. (2014).
@@ -257,7 +268,7 @@ def solve(self, ball1: Ball, ball2: Ball) -> Tuple[Ball, Ball]:
             u_s1=ball1.params.u_s,
             u_s2=ball2.params.u_s,
             # Assume the interaction coefficients are the average of the two balls
-            u_b=(ball1.params.u_b + ball2.params.u_b) / 2,
+            u_b=self.friction.calculate_friction(ball1, ball2),
             e_b=(ball1.params.e_b + ball2.params.e_b) / 2,
             N=self.num_iterations,
         )

pooltool/physics/resolve/ball_ball/frictionless_elastic/__init__.py

@@ -1,11 +1,13 @@
 from typing import Tuple
 
+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.physics.resolve.ball_ball.core import CoreBallBallCollision
+from pooltool.physics.resolve.models import BallBallModel
 
 
 def _resolve_ball_ball(rvw1, rvw2, R):
@@ -26,6 +28,7 @@ def _resolve_ball_ball(rvw1, rvw2, R):
     return rvw1, rvw2
 
 
+@attrs.define
 class FrictionlessElastic(CoreBallBallCollision):
     """A frictionless, instantaneous, elastic, equal mass collision resolver.
 
@@ -36,6 +39,10 @@ class FrictionlessElastic(CoreBallBallCollision):
           https://ekiefl.github.io/2020/04/24/pooltool-theory/#1-elastic-instantaneous-frictionless
     """
 
+    model: BallBallModel = attrs.field(
+        default=BallBallModel.FRICTIONLESS_ELASTIC, init=False
+    )
+
     def solve(self, ball1: Ball, ball2: Ball) -> Tuple[Ball, Ball]:
         """Resolves the collision."""
         rvw1, rvw2 = _resolve_ball_ball(