Merge branch 'major_release_update' of github.com:NACLab/ngc-learn into major_release_update

Author: Alexander Ororbia

ngclearn/components/neurons/graded/gaussianErrorCell.py

@@ -2,6 +2,7 @@
 from ngclearn.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, jit
 from ngclearn.utils import tensorstats
+from ngcsimlib.compilers.process import transition
 
 class GaussianErrorCell(JaxComponent): ## Rate-coded/real-valued error unit/cell
     """
@@ -64,8 +65,9 @@ def __init__(self, name, n_units, batch_size=1, sigma=1., shape=None, **kwargs):
         self.modulator = Compartment(restVals + 1.0) # to be set/consumed
         self.mask = Compartment(restVals + 1.0)
 
+    @transition(output_compartments=["dmu", "dtarget", "dSigma", "L", "mask"])
     @staticmethod
-    def _advance_state(dt, mu, target, Sigma, modulator, mask): ## compute Gaussian error cell output
+    def advance_state(dt, mu, target, Sigma, modulator, mask): ## compute Gaussian error cell output
         # Moves Gaussian cell dynamics one step forward. Specifically, this routine emulates the error unit
         # behavior of the local cost functional:
         # FIXME: Currently, below does: L(targ, mu) = -(1/(2*sigma)) * ||targ - mu||^2_2
@@ -83,16 +85,9 @@ def _advance_state(dt, mu, target, Sigma, modulator, mask): ## compute Gaussian
         mask = mask * 0. + 1. ## "eat" the mask as it should only apply at time t
         return dmu, dtarget, dSigma, jnp.squeeze(L), mask
 
-    @resolver(_advance_state)
-    def advance_state(self, dmu, dtarget, dSigma, L, mask):
-        self.dmu.set(dmu)
-        self.dtarget.set(dtarget)
-        self.dSigma.set(dSigma)
-        self.L.set(L)
-        self.mask.set(mask)
-
+    @transition(output_compartments=["dmu", "dtarget", "dSigma", "target", "mu", "modulator", "L", "mask"])
     @staticmethod
-    def _reset(batch_size, shape, sigma_shape): ## reset core components/statistics
+    def reset(batch_size, shape, sigma_shape): ## reset core components/statistics
         _shape = (batch_size, shape[0])
         if len(shape) > 1:
             _shape = (batch_size, shape[0], shape[1], shape[2])
@@ -107,17 +102,6 @@ def _reset(batch_size, shape, sigma_shape): ## reset core components/statistics
         mask = jnp.ones(_shape)
         return dmu, dtarget, dSigma, target, mu, modulator, L, mask
 
-    @resolver(_reset)
-    def reset(self, dmu, dtarget, dSigma, target, mu, modulator, L, mask):
-        self.dmu.set(dmu)
-        self.dtarget.set(dtarget)
-        self.dSigma.set(dSigma)
-        self.target.set(target)
-        self.mu.set(mu)
-        self.modulator.set(modulator)
-        self.L.set(L)
-        self.mask.set(mask)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

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)
+  print(outs)
+  ## output should equal input
+  # assert_array_equal(outs, y_seq, tol=1e-3)
+  np.testing.assert_allclose(outs, y_seq, atol=1e-3)
+

tests/components/neurons/graded/test_bernoulliErrorCell.py

@@ -0,0 +1,66 @@
+# %%
+
+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 BernoulliErrorCell
+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_bernoulliErrorCell():
+  np.random.seed(42)
+  name = "bernoulli_error_ctx"
+  dkey = random.PRNGKey(42)
+  dkey, *subkeys = random.split(dkey, 100)
+  dt = 1.  # ms
+  with Context(name) as ctx:
+    a = BernoulliErrorCell(
+      name="a", n_units=1, batch_size=1, input_logits=False, shape=None
+    )
+    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.p.set(x)
+
+    @Context.dynamicCommand
+    def clamp_target(x):
+      a.target.set(x)
+
+  ## input spike train
+  x_seq = jnp.asarray(np.random.randn(1, 10))
+  target_seq = (jnp.arange(10)[None] - 5.0) / 2.0
+  ## desired output/epsp pulses
+  y_seq = jnp.asarray([[-2.8193381, -4976.9263, -2.1224928, -2939.0425, -1233.3916, -0.24662945, -708.30042, 0.28213939, 3550.8477, 1.3651246]], 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)
+      target_xt = jnp.array([[target_seq[0, ts]]])
+      ctx.clamp_target(target_xt)
+      ctx.run(t=ts * 1., dt=dt)
+      outs.append(a.dp.value)
+  outs = jnp.concatenate(outs, axis=1)
+  # print(outs)
+  ## output should equal input
+  np.testing.assert_allclose(outs, y_seq, atol=1e-7)
+
+# test_bernoulliErrorCell()

tests/components/neurons/graded/test_gaussianErrorCell.py

@@ -0,0 +1,77 @@
+# %%
+
+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 GaussianErrorCell
+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_gaussianErrorCell():
+  np.random.seed(42)
+  name = "gaussian_error_ctx"
+  dkey = random.PRNGKey(42)
+  dkey, *subkeys = random.split(dkey, 100)
+  dt = 1.  # ms
+  with Context(name) as ctx:
+    a = GaussianErrorCell(
+      name="a", n_units=1, batch_size=1, sigma=1.0, shape=None
+    )
+    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_mu(x):
+      a.mu.set(x)
+
+    @Context.dynamicCommand
+    def clamp_target(x):
+      a.target.set(x)
+
+  ## input sequence
+  mu_seq = jnp.asarray(np.random.randn(1, 10))
+  target_seq = (jnp.arange(10)[None] - 5.0) / 2.0
+  ## expected output based on the Gaussian error cell formula
+  ## L = -0.5 * (target - mu)^2 / sigma, dmu = (target - mu) / sigma
+  expected_dmu = (target_seq - mu_seq) / 1.0  # sigma = 1.0
+  expected_L = -0.5 * jnp.square(target_seq - mu_seq) / 1.0
+
+  dmu_outs = []
+  L_outs = []
+  ctx.reset()
+  for ts in range(mu_seq.shape[1]):
+      mu_t = jnp.array([[mu_seq[0, ts]]])  ## get data at time t
+      ctx.clamp_mu(mu_t)
+      target_t = jnp.array([[target_seq[0, ts]]])
+      ctx.clamp_target(target_t)
+      ctx.run(t=ts * 1., dt=dt)
+      dmu_outs.append(a.dmu.value)
+      L_outs.append(a.L.value)
+
+  dmu_outs = jnp.concatenate(dmu_outs, axis=1)
+  L_outs = jnp.array(L_outs)[None] # (1, 10)
+  # print(dmu_outs.shape)
+  # print(L_outs.shape)
+  # print(expected_dmu.shape)
+  # print(expected_L.shape)
+
+  ## verify outputs match expected values
+  np.testing.assert_allclose(dmu_outs, expected_dmu, atol=1e-5)
+  np.testing.assert_allclose(L_outs, expected_L, atol=1e-5)
+
+# test_gaussianErrorCell()