minor cleanup/edits to srm predictor co-routines

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/LIFSRM.py

@@ -70,6 +70,7 @@ def __init__(
 
         ## analytical SRM state compartments/variables (NOTE: designed to avoid maintaining explicit spike history tensors)
         self.t_last_spike = Compartment(jnp.full((self.batch_size, self.n_units), -1.0))
+        self.next_spike_t = Compartment(jnp.full((self.batch_size, self.n_units), -1.0))
         self.j_lowpass = Compartment(jnp.zeros((self.batch_size, self.n_units))) ## integrated input trace
         self.last_t_eval = Compartment(jnp.zeros((self.batch_size, self.n_units))) ## tracks clock index (when evaluated)
 
@@ -98,9 +99,23 @@ def reset(self):
         self.v.set(jnp.full((self.batch_size, self.n_units), self.v_rest))
         self.s.set(jnp.zeros((self.batch_size, self.n_units)))
         self.t_last_spike.set(jnp.full((self.batch_size, self.n_units), -1.0))
+        self.next_spike_t.set(jnp.full((self.batch_size, self.n_units), -1.0))
         self.j_lowpass.set(jnp.zeros((self.batch_size, self.n_units)))
         self.last_t_eval.set(jnp.zeros((self.batch_size, self.n_units)))
 
+    @compilable
+    def predict_next_spike(self, t_start):
+        next_spike_t = LIFSRM._predict_next_spike(  ## call next-spike-time predictor
+            t_start,
+            self.t_last_spike.get(),
+            self.j_lowpass.get(),
+            self.tau_m,
+            self.v_rest,
+            self.v_reset,
+            self.thr
+        )
+        self.next_spike_t.set(next_spike_t) ## store predicted next spike time(s)
+
     @staticmethod
     def _evaluate_SRM_filter( ## kernel co-routine
         t, 
@@ -129,11 +144,11 @@ def _evaluate_SRM_filter( ## kernel co-routine
         return v_total, new_j_trace
 
     @staticmethod
-    def predict_next_spike( ## next-spike-time prediction co-routine
+    def _predict_next_spike( ## next-spike-time prediction co-routine
         t_start, t_last_spike, j_lowpass, tau_m, v_rest, v_reset, thr
     ):
-        ## co-routine to predict future next spike, takes advantage of an LIF-SRM's 
-        ## closed-form setup; specificaly, this function calculates a future (global)
+        ## co-routine predicts future next spike, taking particular advantage of an LIF-SRM's
+        ## closed-form setup; specifically, this function calculates a future (global)
         ## clock time-stamp as to when this LIF model's decaying voltage would 
         ## cross a firing threshold (note, this does not require step-wise numerical integration)
 
@@ -160,4 +175,3 @@ def predict_next_spike( ## next-spike-time prediction co-routine
         ## safety check: if total driving force is insufficient to cross,
         ## then flag output as "un-triggered" (i.e.,-1.0)
         return jnp.where(can_reach_thr, predicted_t, -1.0) # predicted spike time(s)
-

ngclearn/components/neurons/spiking/RAFSRM.py

@@ -1,5 +1,5 @@
 from ngclearn.components.jaxComponent import JaxComponent
-from jax import numpy as jnp, jit
+from jax import numpy as jnp, jit, lax
 from ngclearn import compilable, Compartment
 
 
@@ -82,6 +82,7 @@ def __init__(
         self.s = Compartment(jnp.zeros((self.batch_size, self.n_units)))
         ## set up analytical state filters (to match RAFCell construction)
         self.t_last_spike = Compartment(jnp.full((self.batch_size, self.n_units), -1.0))
+        self.next_spike_t = Compartment(jnp.full((self.batch_size, self.n_units), -1.0))
         self.j_v = Compartment(jnp.zeros((self.batch_size, self.n_units)))
         self.j_w = Compartment(jnp.zeros((self.batch_size, self.n_units)))
         self.last_t_eval = Compartment(jnp.zeros((self.batch_size, self.n_units)))
@@ -111,10 +112,26 @@ def reset(self): #
         self.v.set(jnp.zeros((self.batch_size, self.n_units)))
         self.s.set(jnp.zeros((self.batch_size, self.n_units)))
         self.t_last_spike.set(jnp.full((self.batch_size, self.n_units), -1.0))
+        self.next_spike_t.set(jnp.full((self.batch_size, self.n_units), -1.0))
         self.j_v.set(jnp.zeros((self.batch_size, self.n_units)))
         self.j_w.set(jnp.zeros((self.batch_size, self.n_units)))
         self.last_t_eval.set(jnp.zeros((self.batch_size, self.n_units)))
 
+    @compilable
+    def predict_next_spike(self, t_start):
+        next_spike_t = RAFSRM._predict_next_spike(  ## next-spike-time predictor
+            t_start,
+            self.j_v.get(),
+            self.j_w.get(),
+            self.tau_v,
+            self.tau_w,
+            self.omega,
+            self.dampen_factor,
+            self.thr,
+            max_iters=10
+        )
+        self.next_spike_t.set(next_spike_t)  ## store predicted next spike time(s)
+
     @staticmethod
     def _evaluate_SRM_filter( ## kernel co-routine
         t, 
@@ -176,11 +193,18 @@ def _evaluate_SRM_filter( ## kernel co-routine
         return v_total, new_v, new_w
 
     @staticmethod
-    def predict_next_spike(
-        t_start, j_v, j_w, tau_v, tau_w, omega, dampen_factor, thr, max_iters=10
+    def _predict_next_spike( ## co-routine used to compute when a future spike will occur
+        t_start,
+        j_v,
+        j_w,
+        tau_v,
+        tau_w,
+        omega,
+        dampen_factor,
+        thr,
+        max_iters=10
     ):
-        ## co-routine used to compute when a future spike will occur
-        ## based on a Newton-Raphson root finding process to predict when exactly
+        ## NOTE: this co-routine is based on a Newton-Raphson root finding process to predict when exactly
         ## continuous sub-threshold wave equation will cross threshold line
 
         ## standardize SRM system parameters
@@ -220,7 +244,7 @@ def scan_body(carry, _): ## set up scanner for this search process
             return s_next, None
 
         ## execute unrolled iteration passes completely w/in JAX/JIT compiler context
-        final_s, _ = jax.lax.scan(scan_body, s_guess, None, length=max_iters)
+        final_s, _ = lax.scan(scan_body, s_guess, None, length=max_iters)
         predicted_t = t_start + final_s ## compute absolute (executed) global clock projection coordinate
         ## employ a safety check:
         ### if initial slope is negative (or diverging), flag it as un-triggered (i.e., -1.0)