16 angular momentum and friction

cotix/_bodies.py

@@ -29,6 +29,9 @@ class AbstractBody(eqx.Module, strict=True):
     angle: AbstractVar[Float[Array, ""]]
     angular_velocity: AbstractVar[Float[Array, ""]]
 
+    elasticity: AbstractVar[Float[Array, ""]]
+    friction_coefficient: AbstractVar[Float[Array, ""]]
+
     shape: AbstractVar[UniversalShape]
 
     def update_transform(self):
@@ -58,15 +61,17 @@ def set_inertia(self, inertia: Float[Array, ""]):
 
     def set_position(self, position: Float[Array, "2"]):
         """Sets position of the center of mass of the body."""
-        return eqx.tree_at(lambda x: x.position, self, position)
+        tmp = eqx.tree_at(lambda x: x.position, self, position)
+        return tmp.update_transform()
 
     def set_velocity(self, velocity: Float[Array, "2"]):
         """Sets velocity of the center of mass of the body."""
         return eqx.tree_at(lambda x: x.velocity, self, velocity)
 
     def set_angle(self, angle: Float[Array, ""]):
         """Sets the angle of rotation around its center of mass."""
-        return eqx.tree_at(lambda x: x.angle, self, angle)
+        tmp = eqx.tree_at(lambda x: x.angle, self, angle)
+        return tmp.update_transform()
 
     def set_angular_velocity(self, angular_velocity: Float[Array, ""]):
         """Sets the rate of change of bodys angle."""
@@ -76,6 +81,13 @@ def set_shape(self, shape: UniversalShape):
         """Replaces the shape with any other shape: not recommended to use."""
         return eqx.tree_at(lambda x: x.shape, self, shape)
 
+    def get_center_of_mass(self):
+        return self.position
+
+    def get_mass_matrix(self):
+        # it is a scalar since we are in 2d, so there is 1 axis of rotation
+        return self.inertia
+
     def __invariant__(self):
         return (
             # Checks for nans
@@ -109,19 +121,34 @@ class Ball(AbstractBody, strict=True):
     angle: Float[Array, ""]
     angular_velocity: Float[Array, ""]
 
+    elasticity: Float[Array, ""]
+    friction_coefficient: Float[Array, ""]
+
     shape: UniversalShape
 
-    def __init__(self, mass, velocity, shape):
+    def __init__(self, mass, position, velocity, shape):
+        # check that the shape is a circle
+        if not (isinstance(shape.parts[0], Circle) and len(shape.parts) == 1):
+            raise ValueError("Ball universal shape must be a circle")
+
         self.mass = mass
-        self.inertia = mass
+        self.inertia = (
+            2 * (mass * shape.parts[0].radius ** 2) / 5
+        )  # inertia of a solid ball
 
-        self.position = jnp.zeros((2,))
+        self.position = position
         self.velocity = velocity
 
         self.angle = jnp.array(0.0)
         self.angular_velocity = jnp.array(0.0)
 
+        self.elasticity = jnp.array(1.0)
+        self.friction_coefficient = jnp.array(1.0)
+
         self.shape = shape
+        self.shape = self.shape.update_transform(
+            angle=self.angle, position=self.position
+        )
 
     @staticmethod
     def make_default():
@@ -131,6 +158,7 @@ def make_default():
         ball = Ball(
             jnp.array(1.0),
             jnp.zeros((2,)),
+            jnp.zeros((2,)),
             UniversalShape(Circle(jnp.array(0.05), jnp.zeros((2,)))),
         )
         ball = (
@@ -143,3 +171,52 @@ def make_default():
             .set_shape(UniversalShape(Circle(jnp.array(0.05), jnp.zeros((2,)))))
         )
         return ball
+
+
+class AnyBody(AbstractBody, strict=True):
+    """
+    A body with any shape. Useful for tests.
+    """
+
+    mass: Float[Array, ""]
+    inertia: Float[Array, ""]
+
+    position: Float[Array, "2"]
+    velocity: Float[Array, "2"]
+
+    angle: Float[Array, ""]
+    angular_velocity: Float[Array, ""]
+
+    elasticity: Float[Array, ""]
+    friction_coefficient: Float[Array, ""]
+
+    shape: UniversalShape
+
+    def __init__(
+        self,
+        mass,
+        inertia,
+        position,
+        velocity,
+        angle,
+        angular_velocity,
+        elasticity,
+        friction_coefficient,
+        shape,
+    ):
+        self.mass = mass
+        self.inertia = inertia
+
+        self.position = position
+        self.velocity = velocity
+
+        self.angle = angle
+        self.angular_velocity = angular_velocity
+
+        self.elasticity = elasticity
+        self.friction_coefficient = friction_coefficient
+
+        self.shape = shape
+        self.shape = self.shape.update_transform(
+            angle=self.angle, position=self.position
+        )