Implementation of LIF reset models in CPU backend

src/lava/proc/io/encoder.py

@@ -129,7 +129,8 @@ def __init__(self, proc_params: Optional[Dict] = None):
 
     def run_spk(self):
         self.s_out.send(self.s_out_buf)
-        a_in_data = np.left_shift(self.a_in.recv(), self.spike_exp)
+        a_in_data = np.left_shift(self.a_in.recv().astype(int),
+                                  self.spike_exp)
         self.s_out_buf = self.encode_delta(a_in_data)
 
 
@@ -235,7 +236,8 @@ def encode_delta_sparse_8(self, s_out):
     def run_spk(self):
         self.s_out.send(self.data, self.idx)
         # Receive synaptic input
-        a_in_data = np.left_shift(self.a_in.recv(), self.spike_exp)
+        a_in_data = np.left_shift(self.a_in.recv().astype(int),
+                                  self.spike_exp)
         s_out = self.encode_delta(a_in_data)
         if self.compression == Compression.SPARSE:
             self.data, self.idx = self.encode_sparse(s_out)

src/lava/proc/lif/models.py

@@ -8,7 +8,7 @@
 from lava.magma.core.resources import CPU
 from lava.magma.core.decorator import implements, requires, tag
 from lava.magma.core.model.py.model import PyLoihiProcessModel
-from lava.proc.lif.process import LIF, TernaryLIF
+from lava.proc.lif.process import LIF, LIFReset, TernaryLIF
 
 
 class AbstractPyLifModelFloat(PyLoihiProcessModel):
@@ -320,3 +320,110 @@ def reset_voltage(self, spike_vector: np.ndarray):
         """Reset voltage of all spiking neurons to 0.
         """
         self.v[spike_vector != 0] = 0  # Reset voltage to 0 wherever we spiked
+
+
+@implements(proc=LIFReset, protocol=LoihiProtocol)
+@requires(CPU)
+@tag('floating_pt')
+class PyLifResetModelFloat(AbstractPyLifModelFloat):
+    """Implementation of Leaky-Integrate-and-Fire neural process with reset
+    in floating point precision.
+    """
+    s_out: PyOutPort = LavaPyType(PyOutPort.VEC_DENSE, float)
+    vth: float = LavaPyType(float, float)
+
+    def __init__(self, proc_params):
+        super(PyLifResetModelFloat, self).__init__(proc_params)
+        self.reset_interval = proc_params['reset_interval']
+        self.reset_offset = (proc_params['reset_offset']) % self.reset_interval
+
+    def spiking_activation(self):
+        """Spiking activation function for LIF.
+        """
+        return self.v > self.vth
+
+    def run_spk(self):
+        """The run function that performs the actual computation during
+        execution orchestrated by a PyLoihiProcessModel using the
+        LoihiProtocol.
+        """
+        # Receive synaptic input
+        a_in_data = self.a_in.recv()
+
+        if (self.time_step % self.reset_interval) == self.reset_offset:
+            self.u *= 0
+            self.v *= 0
+
+        self.subthr_dynamics(activation_in=a_in_data)
+
+        s_out = self.spiking_activation()
+
+        # Reset voltage of spiked neurons to 0
+        self.reset_voltage(spike_vector=s_out)
+        self.s_out.send(s_out)
+
+
+@implements(proc=LIFReset, protocol=LoihiProtocol)
+@requires(CPU)
+@tag('bit_accurate_loihi', 'fixed_pt')
+class PyLifResetModelBitAcc(AbstractPyLifModelFixed):
+    """Implementation of Leaky-Integrate-and-Fire neural process with reset
+    bit-accurate with Loihi's hardware LIF dynamics, which means, it mimics
+    Loihi behaviour.
+
+    Precisions of state variables
+
+    - du: unsigned 12-bit integer (0 to 4095)
+    - dv: unsigned 12-bit integer (0 to 4095)
+    - bias_mant: signed 13-bit integer (-4096 to 4095). Mantissa part of neuron
+      bias.
+    - bias_exp: unsigned 3-bit integer (0 to 7). Exponent part of neuron bias.
+    - vth: unsigned 17-bit integer (0 to 131071).
+
+    """
+    s_out: PyOutPort = LavaPyType(PyOutPort.VEC_DENSE, np.int32, precision=24)
+    vth: int = LavaPyType(int, np.int32, precision=17)
+
+    def __init__(self, proc_params):
+        super(PyLifResetModelBitAcc, self).__init__(proc_params)
+        self.effective_vth = 0
+        self.reset_interval = proc_params['reset_interval']
+        self.reset_offset = (proc_params['reset_offset']) % self.reset_interval
+
+    def scale_threshold(self):
+        """Scale threshold according to the way Loihi hardware scales it. In
+        Loihi hardware, threshold is left-shifted by 6-bits to MSB-align it
+        with other state variables of higher precision.
+        """
+        self.effective_vth = np.left_shift(self.vth, self.vth_shift)
+        self.isthrscaled = True
+
+    def spiking_activation(self):
+        """Spike when voltage exceeds threshold.
+        """
+        return self.v > self.effective_vth
+
+    def run_spk(self):
+        """The run function that performs the actual computation during
+        execution orchestrated by a PyLoihiProcessModel using the
+        LoihiProtocol.
+        """
+        # Receive synaptic input
+        a_in_data = self.a_in.recv()
+
+        if (self.time_step % self.reset_interval) == self.reset_offset:
+            self.u *= 0
+            self.v *= 0
+
+        self.scale_bias()
+
+        if not self.isthrscaled:
+            self.scale_threshold()
+
+        self.subthr_dynamics(activation_in=a_in_data)
+
+        s_out = self.spiking_activation()
+
+        # Reset voltage of spiked neurons to 0
+        self.reset_voltage(spike_vector=s_out)
+        self.s_out.send(s_out)

tests/lava/proc/lif/test_models.py

@@ -12,10 +12,11 @@
 from lava.magma.core.process.process import AbstractProcess
 from lava.magma.core.process.variable import Var
 from lava.magma.core.resources import CPU
-from lava.magma.core.run_configs import RunConfig
+from lava.magma.core.run_configs import Loihi1SimCfg, Loihi2SimCfg, RunConfig
 from lava.magma.core.run_conditions import RunSteps
 from lava.magma.core.sync.protocols.loihi_protocol import LoihiProtocol
-from lava.proc.lif.process import LIF, TernaryLIF
+from lava.proc.lif.process import LIF, LIFReset, TernaryLIF
+from lava.proc import io
 
 
 class LifRunConfig(RunConfig):
@@ -732,3 +733,93 @@ def test_fixed_pm_neg_impulse_dv(self):
         lif_v_float = np.right_shift(np.array(lif_v), 6)
         self.assertListEqual(expected_v_timeseries, lif_v)
         self.assertListEqual(expected_float_v, lif_v_float.tolist())
+
+
+class TestTLIFReset(unittest.TestCase):
+    """Test LIF reset process models"""
+
+    def test_float_model(self):
+        """Test float model"""
+        num_neurons = 10
+        num_steps = 16
+        reset_interval = 4
+        reset_offset = 3
+
+        lif_reset = LIFReset(shape=(num_neurons,),
+                             u=np.arange(num_neurons),
+                             du=0,
+                             dv=0,
+                             vth=100,
+                             bias_mant=np.arange(num_neurons) + 1,
+                             reset_interval=reset_interval,
+                             reset_offset=reset_offset)
+
+        u_logger = io.sink.Read(buffer=num_steps)
+        v_logger = io.sink.Read(buffer=num_steps)
+
+        u_logger.connect_var(lif_reset.u)
+        v_logger.connect_var(lif_reset.v)
+
+        lif_reset.run(condition=RunSteps(num_steps),
+                      run_cfg=Loihi2SimCfg())
+        u = u_logger.data.get()
+        v = v_logger.data.get()
+        lif_reset.stop()
+
+        # Lava timesteps start from t=0. So the first reset offset is missed.
+        u_gt_pre = np.vstack([np.arange(num_neurons)] * 2).T
+        u_gt_post = np.zeros((num_neurons, num_steps - reset_offset + 1))
+
+        dt = (1 + np.arange(reset_offset - 1)).reshape(1, -1)
+        v_gt_pre = np.arange(num_neurons).reshape(-1, 1) * dt \
+            + (1 + np.arange(num_neurons)).reshape(-1, 1) * dt
+        dt = (1 + np.arange(num_steps - reset_offset + 1) % 4).reshape(1, -1)
+        v_gt_post = (1 + np.arange(num_neurons)).reshape(-1, 1) * dt
+
+        self.assertTrue(np.array_equal(u[:, :reset_offset - 1], u_gt_pre))
+        self.assertTrue(np.array_equal(u[:, reset_offset - 1:], u_gt_post))
+        self.assertTrue(np.array_equal(v[:, :reset_offset - 1], v_gt_pre))
+        self.assertTrue(np.array_equal(v[:, reset_offset - 1:], v_gt_post))
+
+    def test_fixed_model(self):
+        """Test fixed model"""
+        num_neurons = 10
+        num_steps = 16
+        reset_interval = 4
+        reset_offset = 3
+
+        lif_reset = LIFReset(shape=(num_neurons,),
+                             u=np.arange(num_neurons),
+                             du=-1,
+                             dv=0,
+                             vth=100,
+                             bias_mant=np.arange(num_neurons) + 1,
+                             reset_interval=reset_interval,
+                             reset_offset=reset_offset)
+
+        u_logger = io.sink.Read(buffer=num_steps)
+        v_logger = io.sink.Read(buffer=num_steps)
+
+        u_logger.connect_var(lif_reset.u)
+        v_logger.connect_var(lif_reset.v)
+
+        lif_reset.run(condition=RunSteps(num_steps),
+                      run_cfg=Loihi2SimCfg(select_tag='fixed_pt'))
+        u = u_logger.data.get()
+        v = v_logger.data.get()
+        lif_reset.stop()
+
+        # Lava timesteps start from t=0. So the first reset offset is missed.
+        u_gt_pre = np.vstack([np.arange(num_neurons)] * 2).T
+        u_gt_post = np.zeros((num_neurons, num_steps - reset_offset + 1))
+
+        dt = (1 + np.arange(reset_offset - 1)).reshape(1, -1)
+        v_gt_pre = np.arange(num_neurons).reshape(-1, 1) * dt \
+            + (1 + np.arange(num_neurons)).reshape(-1, 1) * dt
+        dt = (1 + np.arange(num_steps - reset_offset + 1) % 4).reshape(1, -1)
+        v_gt_post = (1 + np.arange(num_neurons)).reshape(-1, 1) * dt
+
+        self.assertTrue(np.array_equal(u[:, :reset_offset - 1], u_gt_pre))
+        self.assertTrue(np.array_equal(u[:, reset_offset - 1:], u_gt_post))
+        self.assertTrue(np.array_equal(v[:, :reset_offset - 1], v_gt_pre))
+        self.assertTrue(np.array_equal(v[:, reset_offset - 1:], v_gt_post))