Dumb lunar landing (#52)

cotix/_abstract_shapes.py

@@ -15,7 +15,7 @@
 SupportFn = Callable[[Float[Array, "2"]], Float[Array, "2"]]
 
 
-class AbstractShape(eqx.Module, strict=True):
+class AbstractShape(eqx.Module):
     "blah blah"
 
     @abc.abstractmethod
@@ -53,5 +53,5 @@ def get_center(self):
         raise NotImplementedError
 
 
-class AbstractConvexShape(AbstractShape, strict=True):
+class AbstractConvexShape(AbstractShape):
     ...

cotix/_bodies.py

@@ -119,6 +119,9 @@ def load(self, other):
             .set_friction_coefficient(other.friction_coefficient)
         )
 
+    def draw(self, painter):
+        self.shape.draw(painter)
+
 
 class Ball(AbstractBody, strict=True):
     """
@@ -207,17 +210,17 @@ class DynamicBody(eqx.Module, strict=True):
     friction_coefficient: Float[Array, ""]
 
     def __init__(self, other: AbstractBody):
-        self.mass = other.mass
-        self.inertia = other.inertia
+        self.mass = other.mass.astype(jnp.float32)
+        self.inertia = other.inertia.astype(jnp.float32)
 
-        self.position = other.position
-        self.velocity = other.velocity
+        self.position = other.position.astype(jnp.float32)
+        self.velocity = other.velocity.astype(jnp.float32)
 
-        self.angle = other.angle
-        self.angular_velocity = other.angular_velocity
+        self.angle = other.angle.astype(jnp.float32)
+        self.angular_velocity = other.angular_velocity.astype(jnp.float32)
 
-        self.elasticity = other.elasticity
-        self.friction_coefficient = other.friction_coefficient
+        self.elasticity = other.elasticity.astype(jnp.float32)
+        self.friction_coefficient = other.friction_coefficient.astype(jnp.float32)
 
     def set_mass(self, mass: Float[Array, ""]):
         """Sets the mass (a scalar) of the body."""
@@ -294,17 +297,17 @@ def __init__(
         friction_coefficient=jnp.array(1.0),
         shape=None,
     ):
-        self.mass = mass
-        self.inertia = inertia
+        self.mass = mass.astype(jnp.float32)
+        self.inertia = inertia.astype(jnp.float32)
 
-        self.position = position
-        self.velocity = velocity
+        self.position = position.astype(jnp.float32)
+        self.velocity = velocity.astype(jnp.float32)
 
-        self.angle = angle
-        self.angular_velocity = angular_velocity
+        self.angle = angle.astype(jnp.float32)
+        self.angular_velocity = angular_velocity.astype(jnp.float32)
 
-        self.elasticity = elasticity
-        self.friction_coefficient = friction_coefficient
+        self.elasticity = elasticity.astype(jnp.float32)
+        self.friction_coefficient = friction_coefficient.astype(jnp.float32)
 
         self.shape = shape
         self.shape = self.shape.update_transform(