patched ode_utils backend wrt jax, cleaned up unit-tests, added disable flag for phasor-cell

Author: Alexander Ororbia

ngclearn/components/input_encoders/phasorCell.py

@@ -34,8 +34,8 @@ class PhasorCell(JaxComponent):
     """
 
     # Define Functions
-    def __init__(self, name, n_units, target_freq=63.75, batch_size=1,
-                 **kwargs):
+    def __init__(
+            self, name, n_units, target_freq=63.75, batch_size=1, disable_phasor=False, **kwargs):
         super().__init__(name, **kwargs)
 
         ## Phasor meta-parameters
@@ -63,6 +63,7 @@ def __init__(self, name, n_units, target_freq=63.75, batch_size=1,
         # beta = random.poisson(subkey, lam=target_freq, shape=self.angles.value.shape) / target_freq
 
         self.base_scale = random.poisson(subkey[0], lam=target_freq, shape=self.angles.value.shape) / target_freq
+        self.disable_phasor = disable_phasor
 
     def validate(self, dt=None, **validation_kwargs):
         valid = super().validate(**validation_kwargs)
@@ -87,7 +88,7 @@ def validate(self, dt=None, **validation_kwargs):
 
     @transition(output_compartments=["outputs", "tols", "key", "angles"])
     @staticmethod
-    def advance_state(t, dt, target_freq, key, inputs, angles, tols, base_scale):
+    def advance_state(t, dt, target_freq, key, inputs, angles, tols, base_scale, disable_phasor):
         ms_per_second = 1000  # ms/s
         events_per_ms = target_freq / ms_per_second  # e/s s/ms -> e/ms
         ms_per_event = 1 / events_per_ms  # ms/e
@@ -108,6 +109,8 @@ def advance_state(t, dt, target_freq, key, inputs, angles, tols, base_scale):
         updated_angles = jnp.where(updated_angles > angle_per_event,
                                    updated_angles - angle_per_event,
                                    updated_angles)
+        if disable_phasor:
+            outputs = inputs + 0
         tols = tols * (1. - outputs) + t * outputs
 
         return outputs, tols, key, updated_angles

ngclearn/components/neurons/spiking/LIFCell.py

@@ -14,7 +14,7 @@
 #from ngcsimlib.component import Component
 from ngcsimlib.compartment import Compartment
 
-@jit
+#@jit
 def _dfv_internal(j, v, rfr, tau_m, refract_T, v_rest, v_decay=1.): ## raw voltage dynamics
     mask = (rfr >= refract_T).astype(jnp.float32) # get refractory mask
     ## update voltage / membrane potential
@@ -188,16 +188,16 @@ def advance_state(
         j = j * resist_m
         ############################################################################
         ### Runs leaky integrator (leaky integrate-and-fire; LIF) neuronal dynamics.
-        _v_thr = thr_theta + thr #v_theta + v_thr ## calc present voltage threshold
+        _v_thr = thr_theta + thr ## calc present voltage threshold
         #mask = (rfr >= refract_T).astype(jnp.float32) # get refractory mask
         ## update voltage / membrane potential
         v_params = (j, rfr, tau_m, refract_T, v_rest, v_decay)
         if intgFlag == 1:
             _, _v = step_rk2(0., v, _dfv, dt, v_params)
-        else: #_v = v + (v_rest - v) * (dt/tau_m) + (j * mask)
+        else:
             _, _v = step_euler(0., v, _dfv, dt, v_params)
-        ## obtain action potentials/spikes
-        s = (_v > _v_thr).astype(jnp.float32)
+        ## obtain action potentials/spikes/pulses
+        s = (_v > _v_thr) * 1.
         ## update refractory variables
         _rfr = (rfr + dt) * (1. - s)
         ## perform hyper-polarization of neuronal cells

ngclearn/utils/diffeq/ode_utils.py

@@ -45,17 +45,17 @@ def get_integrator_code(integrationType): ## integrator type decoding routine
 @jit
 def _sum_combine(*args, **kwargs): ## fast co-routine for simple addition/summation
     _sum = 0
-    for arg, val in zip(args, kwargs.values()):
+    for arg, val in zip(args, kwargs.values()): ## Sigma^I_{i=1} a_i
         _sum = _sum + val * arg
     return _sum
 
 @jit
 def _step_forward(t, x, dx_dt, dt, x_scale): ## internal step co-routine
-    _t = t + dt
-    _x = x * x_scale + dx_dt * dt
+    _t = t + dt ## advance time forward by dt (denominator)
+    _x = x * x_scale + dx_dt * dt ## advance variable(s) forward by dt (numerator)
     return _t, _x
 
-@partial(jit, static_argnums=(2, 3, 5)) #(2, 3, 4, 5)
+@partial(jit, static_argnums=(2))
 def step_euler(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via the Euler method, i.e., a
@@ -84,7 +84,7 @@ def step_euler(t, x, dfx, dt, params, x_scale=1.):
     _t, _x = next_state
     return _t, _x
 
-@partial(jit, static_argnums=(1, 2, 4)) #(1, 2, 3, 4))
+@partial(jit, static_argnums=(1))
 def _euler(carry, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via the Euler method, i.e., a
@@ -112,7 +112,7 @@ def _euler(carry, dfx, dt, params, x_scale=1.):
     new_carry = (_t, _x)
     return new_carry, (new_carry, carry)
 
-@partial(jit, static_argnums=(2, 3, 5)) #(2, 3, 4, 5))
+@partial(jit, static_argnums=(2))
 def step_heun(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via Heun's method, i.e., a
@@ -150,7 +150,7 @@ def step_heun(t, x, dfx, dt, params, x_scale=1.):
     _t, _x = next_state
     return _t, _x
 
-@partial(jit, static_argnums=(1, 2, 4)) #(1, 2, 3, 4, ))
+@partial(jit, static_argnums=(1))
 def _heun(carry, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via Heun's method, i.e., a
@@ -189,7 +189,7 @@ def _heun(carry, dfx, dt, params, x_scale=1.):
     new_carry = (_t, _x)
     return new_carry, (new_carry, carry)
 
-@partial(jit, static_argnums=(2, 3, 5)) #(2, 3, 4, 5))
+@partial(jit, static_argnums=(2))
 def step_rk2(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via the midpoint method, i.e., a
@@ -224,7 +224,7 @@ def step_rk2(t, x, dfx, dt, params, x_scale=1.):
     _t, _x = next_state
     return _t, _x
 
-@partial(jit, static_argnums=(1, 2, 4)) #(1, 2, 3, 4, ))
+@partial(jit, static_argnums=(1))
 def _rk2(carry, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via the midpoint method, i.e., a
@@ -260,7 +260,7 @@ def _rk2(carry, dfx, dt, params, x_scale=1.):
     new_carry = (_t, _x)
     return new_carry, (new_carry, carry)
 
-@partial(jit, static_argnums=(2, 3, 5)) #(2, 3, 4, 5))
+@partial(jit, static_argnums=(2))
 def step_rk4(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via the midpoint method, i.e., a
@@ -295,7 +295,7 @@ def step_rk4(t, x, dfx, dt, params, x_scale=1.):
     _t, _x = next_state
     return _t, _x
 
-@partial(jit, static_argnums=(1, 2, 4)) #(1, 2, 3, 4, ))
+@partial(jit, static_argnums=(1))
 def _rk4(carry, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via the midpoint method, i.e., a
@@ -338,7 +338,7 @@ def _rk4(carry, dfx, dt, params, x_scale=1.):
     new_carry = (_t, _x)
     return new_carry, (new_carry, carry)
 
-@partial(jit, static_argnums=(2, 3, 5)) #(2, 3, 4, 5))
+@partial(jit, static_argnums=(2))
 def step_ralston(t, x, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via Ralston's method, i.e., a
@@ -375,7 +375,7 @@ def step_ralston(t, x, dfx, dt, params, x_scale=1.):
     _t, _x = next_state
     return _t, _x
 
-@partial(jit, static_argnums=(1, 2, 4)) #(1, 2, 3, 4,))
+@partial(jit, static_argnums=(1))
 def _ralston(carry, dfx, dt, params, x_scale=1.):
     """
     Iteratively integrates one step forward via Ralston's method, i.e., a
@@ -416,7 +416,6 @@ def _ralston(carry, dfx, dt, params, x_scale=1.):
     new_carry = (_t, _x)
     return new_carry, (new_carry, carry)
 
-
 @partial(jit, static_argnums=(0, 3, 4, 5, 6, 7, 8))
 def solve_ode(method_name, t0, x0, T, dfx, dt, params=None, x_scale=1., sols_only=True):
     if method_name =='euler':

tests/components/input_encoders/test_phasorCell.py

@@ -22,7 +22,7 @@ def test_phasorCell1():
     # T = 300  # ms
     # ---- build a simple Poisson cell system ----
     with Context(name) as ctx:
-        a = PhasorCell(name="a", n_units=1, target_freq=1000., key=subkeys[0])
+        a = PhasorCell(name="a", n_units=1, target_freq=1000., disable_phasor=True, key=subkeys[0])
 
         advance_process = (Process()
                            >> a.advance_state)
@@ -49,6 +49,7 @@ def clamp(x):
         outs.append(a.outputs.value)
         #print(a.outputs.value)
     outs = jnp.concatenate(outs, axis=1)
+    #print(outs)
 
     ## output should equal input
     assert_array_equal(outs, x_seq)