update rate cell

Author: rxng8

ngclearn/components/neurons/graded/rateCell.py

@@ -13,9 +13,15 @@
                                             step_euler, step_rk2, step_rk4
 
 ## rewritten code
-@partial(jit, static_argnums=[5])
+# @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
@@ -31,7 +37,7 @@ def _dfz(t, z, params): ## diff-eq dynamics wrapper
     dz_dt = _dfz_internal(z, j, j_td, tau_m, leak_gamma, priorType)
     return dz_dt
 
-@jit
+# @jit
 def _modulate(j, dfx_val):
     """
     Apply a signal modulator to j (typically of the form of a derivative/dampening function)
@@ -46,6 +52,7 @@ 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):
     """
     Runs leaky rate-coded state dynamics one step in time.
@@ -81,7 +88,7 @@ def _run_cell(dt, j, j_td, z, tau_m, leak_gamma=0., integType=0, priorType=None)
         _, _z = step_euler(0., z, _dfz, dt, params)
     return _z
 
-@jit
+# @jit
 def _run_cell_stateless(j):
     """
     A simplification of running a stateless set of dynamics over j (an identity
@@ -161,7 +168,12 @@ def __init__(self, name, n_units, tau_m, prior=("gaussian", 0.), act_fx="identit
             if tau_m <= 0: ## trigger stateless mode
                 self.is_stateful = False
         priorType, leakRate = prior
-        self.priorType = priorType ## type of scale-shift prior to impose over the leak
+        priorTypeDict = {
+            "laplacian": 0,
+            "cauchy": 1,
+            "exp": 2
+        }
+        self.priorType = priorTypeDict.get(priorType, -1)
         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

tests/components/neurons/graded/test_RateCell.py

@@ -0,0 +1,60 @@
+# %%
+
+from jax import numpy as jnp, random, jit
+from ngcsimlib.context import Context
+import numpy as np
+np.random.seed(42)
+from ngclearn.components import RateCell
+from ngcsimlib.compilers import compile_command, wrap_command
+from numpy.testing import assert_array_equal
+
+from ngcsimlib.compilers.process import Process, transition
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+from ngcsimlib.context import Context
+from ngcsimlib.utils.compartment import Get_Compartment_Batch
+
+
+def test_RateCell1():
+  name = "rate_ctx"
+  dkey = random.PRNGKey(42)
+  dkey, *subkeys = random.split(dkey, 100)
+  dt = 1.  # ms
+  with Context(name) as ctx:
+    a = RateCell(
+      name="a", n_units=1, tau_m=50., prior=("gaussian", 0.), act_fx="identity",
+      threshold=("none", 0.), integration_type="euler",
+      batch_size=1, resist_scale=1., shape=None, is_stateful=True
+    )
+    advance_process = (Process() >> a.advance_state)
+    ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
+    reset_process = (Process() >> a.reset)
+    ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
+
+    # reset_cmd, reset_args = ctx.compile_by_key(a, compile_key="reset")
+    # ctx.add_command(wrap_command(jit(ctx.reset)), name="reset")
+    # advance_cmd, advance_args = ctx.compile_by_key(a, compile_key="advance_state")
+    # ctx.add_command(wrap_command(jit(ctx.advance_state)), name="run")
+
+    @Context.dynamicCommand
+    def clamp(x):
+      a.j.set(x)
+
+  ## input spike train
+  x_seq = jnp.ones((1, 10))
+  ## desired output/epsp pulses
+  y_seq = jnp.asarray([[0.02, 0.04, 0.06, 0.08, 0.09999999999999999, 0.11999999999999998, 0.13999999999999999, 0.15999999999999998, 0.17999999999999998, 0.19999999999999998]], dtype=jnp.float32)
+
+  outs = []
+  ctx.reset()
+  for ts in range(x_seq.shape[1]):
+      x_t = jnp.array([[x_seq[0, ts]]])  ## get data at time t
+      ctx.clamp(x_t)
+      ctx.run(t=ts * 1., dt=dt)
+      outs.append(a.z.value)
+  outs = jnp.concatenate(outs, axis=1)
+
+  ## output should equal input
+  # assert_array_equal(outs, y_seq, tol=1e-3)
+  np.testing.assert_allclose(outs, y_seq, atol=1e-3)
+