Improve type annotations

Author: maxencefaldor

examples/31_boids.ipynb

@@ -87,14 +87,15 @@
     "num_boids = 256\n",
     "dt = 0.01\n",
     "\n",
+    "acceleration_max = jnp.inf\n",
+    "acceleration_scale = 1.0\n",
+    "perception = 0.1\n",
+    "separation_distance = 0.025\n",
+    "\n",
     "separation_weight = 4.5\n",
     "alignment_weight = 0.65\n",
     "cohesion_weight = 0.75\n",
-    "perception = 0.1\n",
-    "separation_distance = 0.025\n",
-    "acceleration_scale = 1.0\n",
     "noise_scale = 0.1\n",
-    "acceleration_max = jnp.inf\n",
     "\n",
     "key = jax.random.key(seed)\n",
     "rngs = nnx.Rngs(seed)"
@@ -114,14 +115,14 @@
    "outputs": [],
    "source": [
     "boid_policy = BoidPolicy(\n",
+    "\tacceleration_max=acceleration_max,\n",
+    "\tacceleration_scale=acceleration_scale,\n",
+    "\tperception=perception,\n",
+    "\tseparation_distance=separation_distance,\n",
     "\tseparation_weight=separation_weight,\n",
     "\talignment_weight=alignment_weight,\n",
     "\tcohesion_weight=cohesion_weight,\n",
-    "\tperception=perception,\n",
-    "\tseparation_distance=separation_distance,\n",
-    "\tacceleration_scale=acceleration_scale,\n",
     "\tnoise_scale=noise_scale,\n",
-    "\tacceleration_max=acceleration_max,\n",
     "\trngs=rngs,\n",
     ")\n",
     "\n",
@@ -226,20 +227,21 @@
     "\n",
     "\tdef __init__(\n",
     "\t\tself,\n",
-    "\t\t*,\n",
     "\t\tnum_neighbors: int = 16,  # Number of neighbors to consider\n",
-    "\t\tperception: float = 0.1,  # Perception radius\n",
     "\t\thidden_features: int = 8,  # Hidden layer size from reference\n",
+    "\t\t*,\n",
+    "\t\tacceleration_max: float = jnp.inf,\n",
     "\t\tacceleration_scale: float = 10.0,  # Scaling factor from reference\n",
+    "\t\tperception: float = 0.1,  # Perception radius\n",
     "\t\trngs: nnx.Rngs,\n",
     "\t):\n",
-    "\t\t\"\"\"Initialize the boid policy.\"\"\"\n",
-    "\t\tself.rngs = rngs\n",
+    "\t\t\"\"\"Initialize boid policy.\"\"\"\n",
     "\t\tself.num_neighbors = num_neighbors\n",
-    "\t\tself.perception = perception\n",
+    "\t\tself.acceleration_max = acceleration_max\n",
     "\t\tself.acceleration_scale = acceleration_scale\n",
+    "\t\tself.perception = perception\n",
+    "\t\tself.rngs = rngs\n",
     "\n",
-    "\t\t# Define the neural network layers similar to BoidNetwork\n",
     "\t\tself.dense1 = nnx.Linear(4, hidden_features, rngs=rngs)\n",
     "\t\tself.dense2 = nnx.Linear(hidden_features, hidden_features, rngs=rngs)\n",
     "\t\tself.dense3 = nnx.Linear(hidden_features, hidden_features, rngs=rngs)\n",
@@ -269,6 +271,11 @@
     "\n",
     "\t\treturn global2local, local2global, global2local_rot, local2global_rot\n",
     "\n",
+    "\tdef _clip_by_norm(self, vector: jax.Array, max_val: float) -> jax.Array:\n",
+    "\t\t\"\"\"Limit the magnitude of a vector.\"\"\"\n",
+    "\t\tnorm = jnp.linalg.norm(vector)\n",
+    "\t\treturn jnp.where(norm > max_val, vector * max_val / norm, vector)\n",
+    "\n",
     "\tdef __call__(self, state: BoidsState, boid_idx: int) -> jax.Array:\n",
     "\t\t\"\"\"Compute acceleration for a boid based on its neighbors.\n",
     "\n",
@@ -341,9 +348,12 @@
     "\n",
     "\t\t# Transform back to global frame\n",
     "\t\tdv_hom = jnp.concatenate([dv_local, jnp.zeros(1)], axis=-1)  # 3D homogeneous\n",
-    "\t\tdv_global = (l2gr @ dv_hom[:, None])[:2, 0]  # Back to 2D global coords\n",
+    "\t\tacceleration = (l2gr @ dv_hom[:, None])[:2, 0]  # Back to 2D global coords\n",
+    "\n",
+    "\t\t# Limit acceleration\n",
+    "\t\tacceleration = self._clip_by_norm(acceleration, self.acceleration_max)\n",
     "\n",
-    "\t\treturn dv_global"
+    "\t\treturn acceleration"
    ]
   },
   {

src/cax/core/__init__.py

@@ -1,5 +1,9 @@
 """Core Module."""
 
-from .cs import ComplexSystem
+from .cs import ComplexSystem, Input, State
 
-__all__ = ["ComplexSystem"]
+__all__ = [
+	"ComplexSystem",
+	"Input",
+	"State",
+]

src/cax/core/perceive/__init__.py

@@ -13,14 +13,15 @@
 from .conv_perceive import ConvPerceive
 from .kernels import grad2_kernel, grad_kernel, identity_kernel, neighbors_kernel
 from .moore_perceive import MoorePerceive
-from .perceive import Perceive
+from .perceive import Perceive, Perception
 from .von_neumann_perceive import VonNeumannPerceive
 
 __all__ = [
-	"ConvPerceive",
-	"MoorePerceive",
 	"Perceive",
+	"Perception",
+	"MoorePerceive",
 	"VonNeumannPerceive",
+	"ConvPerceive",
 	"identity_kernel",
 	"neighbors_kernel",
 	"grad_kernel",

src/cax/core/perceive/conv_perceive.py

@@ -4,9 +4,9 @@
 
 from flax import nnx
 
-from cax.types import Perception, State
+from cax.core import State
 
-from .perceive import Perceive
+from .perceive import Perceive, Perception
 
 
 class ConvPerceive(Perceive):

src/cax/core/perceive/moore_perceive.py

@@ -4,9 +4,9 @@
 
 import jax.numpy as jnp
 
-from cax.types import Perception, State
+from cax.core import State
 
-from .perceive import Perceive
+from .perceive import Perceive, Perception
 
 
 class MoorePerceive(Perceive):

src/cax/core/perceive/perceive.py

@@ -6,9 +6,14 @@
 that process the current state of the environment or system and return a perception.
 """
 
-from flax import nnx
+from flax import nnx, struct
 
-from cax.types import Perception, State
+from cax.core import State
+
+
+@struct.dataclass
+class Perception:
+	"""Perception class."""
 
 
 class Perceive(nnx.Module):

src/cax/core/perceive/von_neumann_perceive.py

@@ -4,9 +4,9 @@
 
 import jax.numpy as jnp
 
-from cax.types import Perception, State
+from cax.core import State
 
-from .perceive import Perceive
+from .perceive import Perceive, Perception
 
 
 class VonNeumannPerceive(Perceive):

src/cax/core/update/__init__.py

@@ -16,4 +16,9 @@
 from .residual_update import ResidualUpdate
 from .update import Update
 
-__all__ = ["Update", "MLPUpdate", "NCAUpdate", "ResidualUpdate"]
+__all__ = [
+	"Update",
+	"MLPUpdate",
+	"ResidualUpdate",
+	"NCAUpdate",
+]

src/cax/core/update/mlp_update.py

@@ -7,8 +7,10 @@
 from flax.nnx.nn import initializers
 from flax.nnx.nn.linear import default_kernel_init
 
-from cax.core.update.update import Update
-from cax.types import Input, Perception, State
+from cax.core import Input, State
+from cax.core.perceive import Perception
+
+from .update import Update
 
 
 class MLPUpdate(Update):

src/cax/core/update/nca_update.py

@@ -6,8 +6,10 @@
 from flax import nnx
 from jax import Array
 
-from cax.core.update.residual_update import ResidualUpdate
-from cax.types import Input, Perception, State
+from cax.core import Input, State
+from cax.core.perceive import Perception
+
+from .residual_update import ResidualUpdate
 
 
 class NCAUpdate(ResidualUpdate):