Implement the logic part 2

pooltool/physics/resolve/ball_cushion/unrealistic/__init__.py

@@ -12,9 +12,15 @@
 
 
 def _solve(
-    ball: Ball, cushion: LinearCushionSegment
+    ball: Ball, cushion: LinearCushionSegment, restitution: bool = True
 ) -> Tuple[Ball, LinearCushionSegment]:
-    """Given ball and cushion, unrealistically reflect the ball's momentum"""
+    """Given ball and cushion, unrealistically reflect the ball's momentum
+
+    Args:
+        restitution:
+            By default, the ball's momentum is reflected without loss. Set this to true
+            if the ball's restitution coefficient should dampen the outgoing velocity.
+    """
     rvw = ball.state.rvw
 
     # Two things about the normal:
@@ -39,9 +45,9 @@ def _solve(
     rvw_R = math.coordinate_rotation(rvw.T, -psi).T
 
     # Reverse velocity component lying in normal direction
-    rvw_R[1, 0] *= -1
+    rvw_R[1, 0] *= -1 * (1 if not restitution else ball.params.e_c)
 
-    # Rotate frame of refernce back to the table frame
+    # Rotate frame of reference back to the table frame
     rvw = math.coordinate_rotation(rvw_R.T, psi).T
 
     # Set the ball's rvw
@@ -54,7 +60,10 @@ def _solve(
 
 
 class UnrealisticLinear(CoreBallLCushionCollision):
+    def __init__(self, restitution: bool = True) -> None:
+        self.restitution = restitution
+
     def solve(
         self, ball: Ball, cushion: LinearCushionSegment
     ) -> Tuple[Ball, LinearCushionSegment]:
-        return _solve(ball, cushion)
+        return _solve(ball, cushion, self.restitution)