update rate cell and fix bug of passing a tuple of (jax Array -- not hashable) to jax jit functions. Basically, simplify the codebase by using a hashmap of functions

Author: rxng8

ngclearn/components/neurons/graded/rateCell.py

@@ -12,29 +12,24 @@
 from ngclearn.utils.diffeq.ode_utils import get_integrator_code, \
                                             step_euler, step_rk2, step_rk4
 
-## rewritten code
-# @partial(jit, static_argnums=[5])
-def _dfz_internal(z, j, j_td, tau_m, leak_gamma, prior_type=None): ## raw dynamics
-    z_leak = z # * 2 ## Default: assume Gaussian
-    prior_type_dict = {
-        0: "laplacian",
-        1: "cauchy",
-        2: "exp"
-    }
-    prior_type = prior_type_dict.get(prior_type, None)
-    if prior_type != None:
-        if prior_type == "laplacian": ## Laplace dist
-            z_leak = jnp.sign(z) ## d/dx of Laplace is signum
-        elif prior_type == "cauchy":  ## Cauchy dist: x ~ (1.0 + tf.math.square(z))
-            z_leak = (z * 2)/(1. + jnp.square(z))
-        elif prior_type == "exp":  ## Exp dist: x ~ -exp(-x^2)
-            z_leak = jnp.exp(-jnp.square(z)) * z * 2
+def _dfz_internal_laplace(z, j, j_td, tau_m, leak_gamma): ## raw dynamics
+    z_leak = jnp.sign(z) ## d/dx of Laplace is signum
+    dz_dt = (-z_leak * leak_gamma + (j + j_td)) * (1./tau_m)
+    return dz_dt
+
+def _dfz_internal_cauchy(z, j, j_td, tau_m, leak_gamma): ## raw dynamics
+    z_leak = (z * 2)/(1. + jnp.square(z))
+    dz_dt = (-z_leak * leak_gamma + (j + j_td)) * (1./tau_m)
+    return dz_dt
+
+def _dfz_internal_exp(z, j, j_td, tau_m, leak_gamma): ## raw dynamics
+    z_leak = jnp.exp(-jnp.square(z)) * z * 2
     dz_dt = (-z_leak * leak_gamma + (j + j_td)) * (1./tau_m)
     return dz_dt
 
-def _dfz(t, z, params): ## diff-eq dynamics wrapper
-    j, j_td, tau_m, leak_gamma, priorType = params
-    dz_dt = _dfz_internal(z, j, j_td, tau_m, leak_gamma, priorType)
+def _dfz_internal_gaussian(z, j, j_td, tau_m, leak_gamma): ## raw dynamics
+    z_leak = z # * 2 ## Default: assume Gaussian
+    dz_dt = (-z_leak * leak_gamma + (j + j_td)) * (1./tau_m)
     return dz_dt
 
 # @jit
@@ -52,8 +47,8 @@ def _modulate(j, dfx_val):
     """
     return j * dfx_val
 
-@partial(jit, static_argnames=["integType", "priorType"])
-def _run_cell(dt, j, j_td, z, tau_m, leak_gamma=0., integType=0, priorType=None):
+# @partial(jit, static_argnames=["integType", "priorType"])
+def _run_cell(dt, j, j_td, z, tau_m, leak_gamma=0., integType=0, priorType=0):
     """
     Runs leaky rate-coded state dynamics one step in time.
 
@@ -77,15 +72,21 @@ def _run_cell(dt, j, j_td, z, tau_m, leak_gamma=0., integType=0, priorType=None)
     Returns:
         New value of membrane/state for next time step
     """
-    if integType == 1:
-        params = (j, j_td, tau_m, leak_gamma, priorType)
-        _, _z = step_rk2(0., z, _dfz, dt, params)
-    elif integType == 2:
-        params = (j, j_td, tau_m, leak_gamma, priorType)
-        _, _z = step_rk4(0., z, _dfz, dt, params)
-    else:
-        params = (j, j_td, tau_m, leak_gamma, priorType)
-        _, _z = step_euler(0., z, _dfz, dt, params)
+    _dfz_fns = {
+        0: lambda t, z, params: _dfz_internal_gaussian(z, *params),
+        1: lambda t, z, params: _dfz_internal_laplace(z, *params),
+        2: lambda t, z, params: _dfz_internal_cauchy(z, *params),
+        3: lambda t, z, params: _dfz_internal_exp(z, *params),
+    }
+    _dfz_fn = _dfz_fns.get(priorType, _dfz_internal_gaussian)
+    _step_fns = {
+        0: step_euler,
+        1: step_rk2,
+        2: step_rk4,
+    }
+    _step_fn = _step_fns.get(integType, step_euler)
+    params = (j, j_td, tau_m, leak_gamma)
+    _, _z = _step_fn(0., z, _dfz_fn, dt, params)
     return _z
 
 # @jit
@@ -169,11 +170,12 @@ def __init__(self, name, n_units, tau_m, prior=("gaussian", 0.), act_fx="identit
                 self.is_stateful = False
         priorType, leakRate = prior
         priorTypeDict = {
-            "laplacian": 0,
-            "cauchy": 1,
-            "exp": 2
+            "gaussian": 0,
+            "laplacian": 1,
+            "cauchy": 2,
+            "exp": 3
         }
-        self.priorType = priorTypeDict.get(priorType, -1)
+        self.priorType = priorTypeDict.get(priorType, 0)
         self.priorLeakRate = leakRate ## degree to which rate neurons leak (according to prior)
         thresholdType, thr_lmbda = threshold
         self.thresholdType = thresholdType ## type of thresholding function to use