Restructure transforms (#25)

docs/source/examples/evo_boids.rst

@@ -40,6 +40,7 @@ First we import JAX, Chex, Flax, and Esquilax
 .. testsetup:: evo_boids
 
    from typing import Callable, Tuple
+   from functools import partial
 
 .. testcode:: evo_boids
 
@@ -85,10 +86,11 @@ add aggregates position and velocity data from neighbours in-range
 
 .. testcode:: evo_boids
 
-   @esquilax.transforms.spatial(
-       10,
-       (jnp.add, jnp.add, jnp.add, jnp.add),
-       (0, jnp.zeros(2), 0.0, 0.0),
+   @partial(
+       esquilax.transforms.spatial,
+       n_bins=10,
+       reduction=(jnp.add, jnp.add, jnp.add, jnp.add),
+       default=(0, jnp.zeros(2), 0.0, 0.0),
        include_self=False,
    )
    def observe(_k: chex.PRNGKey, _params: Params, _a: Boid, b: Boid):
@@ -180,8 +182,12 @@ to calculate reward contributions
 
 .. testcode:: evo_boids
 
-   @esquilax.transforms.spatial(
-       5, jnp.add, 0.0, include_self=False,
+   @partial(
+       esquilax.transforms.spatial,
+       n_bins=5,
+       reduction=jnp.add,
+       default=0.0,
+       include_self=False,
    )
    def reward(_k: chex.PRNGKey, params: Params, a: chex.Array, b: chex.Array):
        d = esquilax.utils.shortest_distance(a, b, norm=True)

docs/source/examples/hard_coded_boids.rst

@@ -15,6 +15,7 @@ We first import JAX, `Chex <https://chex.readthedocs.io/en/latest/>`_, and Esqui
 
 .. testcode:: hard_coded_boids
 
+   from functools import partial
    import chex
    import jax
    import jax.numpy as jnp
@@ -54,10 +55,11 @@ Firstly agents observe the state of neighbours within a given range
 
 .. testcode:: hard_coded_boids
 
-   @esquilax.transforms.spatial(
-       5,
-       (jnp.add, jnp.add, jnp.add, jnp.add),
-       (0, jnp.zeros(2), jnp.zeros(2), jnp.zeros(2)),
+   @partial(
+       esquilax.transforms.spatial,
+       n_bins=5,
+       reduction=(jnp.add, jnp.add, jnp.add, jnp.add),
+       default=(0, jnp.zeros(2), jnp.zeros(2), jnp.zeros(2)),
        include_self=False,
    )
    def observe(_key: chex.PRNGKey, params: Params, a: Boid, b: Boid):

docs/source/examples/rl_boids.rst

@@ -14,6 +14,7 @@ example, but wrap them up in an environment class
 .. testsetup:: rl_boids
 
    from typing import Callable, Tuple, Any
+   from functools import partial
    import chex
    import evosax
    import flax
@@ -48,10 +49,11 @@ example, but wrap them up in an environment class
            x = flax.linen.tanh(x)
            return x
 
-   @esquilax.transforms.spatial(
-       10,
-       (jnp.add, jnp.add, jnp.add, jnp.add),
-       (0, jnp.zeros(2), 0.0, 0.0),
+   @partial(
+       esquilax.transforms.spatial,
+       n_bins=10,
+       reduction=(jnp.add, jnp.add, jnp.add, jnp.add),
+       default=(0, jnp.zeros(2), 0.0, 0.0),
        include_self=False,
    )
    def observe(_k: chex.PRNGKey, _params: Params, _a: Boid, b: Boid):
@@ -110,8 +112,12 @@ example, but wrap them up in an environment class
        )
        return (pos + d_pos) % 1.0
 
-   @esquilax.transforms.spatial(
-       5, jnp.add, 0.0, include_self=False,
+   @partial(
+       esquilax.transforms.spatial,
+       n_bins=5,
+       reduction=jnp.add,
+       default=0.0,
+       include_self=False,
    )
    def reward(_k: chex.PRNGKey, params: Params, a: chex.Array, b: chex.Array):
        d = esquilax.utils.shortest_distance(a, b, norm=True)

examples/boids/hard_coded.py

@@ -1,3 +1,5 @@
+from functools import partial
+
 import chex
 import jax
 import jax.numpy as jnp
@@ -21,10 +23,11 @@ class Params:
     close_range: float
 
 
-@esquilax.transforms.spatial(
-    5,
-    (jnp.add, jnp.add, jnp.add, jnp.add),
-    (0, jnp.zeros(2), jnp.zeros(2), jnp.zeros(2)),
+@partial(
+    esquilax.transforms.spatial,
+    n_bins=5,
+    reduction=(jnp.add, jnp.add, jnp.add, jnp.add),
+    default=(0, jnp.zeros(2), jnp.zeros(2), jnp.zeros(2)),
     include_self=False,
 )
 def observe(_key: chex.PRNGKey, params: Params, a: Boids, b: Boids):

examples/boids/updates.py

@@ -1,3 +1,4 @@
+from functools import partial
 from typing import Tuple
 
 import chex
@@ -25,10 +26,11 @@ class Params:
     collision_penalty: float = 0.1
 
 
-@esquilax.transforms.spatial(
-    10,
-    (jnp.add, jnp.add, jnp.add, jnp.add),
-    (0, jnp.zeros(2), 0.0, 0.0),
+@partial(
+    esquilax.transforms.spatial,
+    n_bins=10,
+    reduction=(jnp.add, jnp.add, jnp.add, jnp.add),
+    default=(0, jnp.zeros(2), 0.0, 0.0),
     include_self=False,
 )
 def observe(_k: chex.PRNGKey, _params: Params, a: Boid, b: Boid):
@@ -87,10 +89,11 @@ def move(_key: chex.PRNGKey, _params: Params, x):
     return (pos + d_pos) % 1.0
 
 
-@esquilax.transforms.spatial(
-    5,
-    jnp.add,
-    0.0,
+@partial(
+    esquilax.transforms.spatial,
+    n_bins=5,
+    reduction=jnp.add,
+    default=0.0,
     include_self=False,
 )
 def rewards(_k: chex.PRNGKey, params: Params, a: chex.Array, b: chex.Array):

examples/opinion_dynamics.py

@@ -1,3 +1,5 @@
+from functools import partial
+
 import chex
 import jax
 import jax.numpy as jnp
@@ -18,7 +20,9 @@ class SimState:
     weights: chex.Array
 
 
-@esquilax.transforms.graph_reduce((jnp.add, jnp.add), (0, 0.0))
+@partial(
+    esquilax.transforms.graph_reduce, reduction=(jnp.add, jnp.add), default=(0, 0.0)
+)
 def collect_opinions(_, params: Params, my_opinion, your_opinion, weight):
     d = jnp.abs(my_opinion - your_opinion)
     w = params.strength * weight