incorporate reviews

Author: smm-ncl

src/lava/proc/s4d/models.py

@@ -1,21 +1,21 @@
-# Copyright (C) 2022 Intel Corporation
+# Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: BSD-3-Clause
 # See: https://spdx.org/licenses/
 
-from lava.proc.sdn.models import AbstractSigmaDeltaModel
+import numpy as np
 from typing import Any, Dict
+from lava.proc.sdn.models import AbstractSigmaDeltaModel
 from lava.magma.core.decorator import implements, requires, tag
-import numpy as np
 from lava.magma.core.sync.protocols.loihi_protocol import LoihiProtocol
-from lava.proc.s4d.process import SigmaS4Delta, SigmaS4DeltaLayer
+from lava.proc.s4d.process import SigmaS4dDelta, SigmaS4dDeltaLayer
 from lava.magma.core.resources import CPU
 from lava.magma.core.model.py.ports import PyInPort, PyOutPort
 from lava.magma.core.model.py.type import LavaPyType
 from lava.magma.core.model.sub.model import AbstractSubProcessModel
 from lava.proc.sparse.process import Sparse
 
 
-class AbstractSigmaS4DeltaModel(AbstractSigmaDeltaModel):
+class AbstractSigmaS4dDeltaModel(AbstractSigmaDeltaModel):
     a_in = None
     s_out = None
 
@@ -29,62 +29,87 @@ class AbstractSigmaS4DeltaModel(AbstractSigmaDeltaModel):
     bias = None
 
     # S4 Variables
-    A = None
-    B = None
-    C = None
-    S4state = None
-    S4_exp = None
+    a = None
+    b = None
+    c = None
+    s4_state = None
+    s4_exp = None
 
     def __init__(self, proc_params: Dict[str, Any]) -> None:
-        super().__init__(proc_params)
-        self.A = self.proc_params['A']
-        self.B = self.proc_params['B']
-        self.C = self.proc_params['C']
-        self.S4state = self.proc_params['S4state']
-
-    def activation_dynamics(self, sigma_data: np.ndarray) -> np.ndarray:
-        """Sigma Delta activation dynamics. Performs S4D dynamics.
+        """
+        Sigma delta neuron model that implements S4D
+        (as described by Gu et al., 2022) dynamics as its activation function.
 
         Parameters
         ----------
-        sigma_data: np.ndarray
-            sigma decoded data
+        shape: Tuple
+            Shape of the sigma process.
+        vth: int or float
+            Threshold of the delta encoder.
+        a: np.ndarray
+            Diagonal elements of the state matrix of the S4D model.
+        b: np.ndarray
+            Diagonal elements of the input matrix of the S4D model.
+        c: np.ndarray
+            Diagonal elements of the output matrix of the S4D model.
+        state_exp: int
+            Scaling exponent with base 2 for the reconstructed sigma variables.
+            Note: This should only be used for nc models.
+            Default is 0.
+        s4_exp: int
+            Scaling exponent with base 2 for the S4 state variables.
+            Note: This should only be used for nc models.
+            Default is 0.
+        """
+        super().__init__(proc_params)
+        self.a = self.proc_params['a']
+        self.b = self.proc_params['b']
+        self.c = self.proc_params['c']
+        self.s4_state = self.proc_params['s4_state']
 
-        Returns
-        -------
-        np.ndarray
-            activation output
+    def activation_dynamics(self, sigma_data: np.ndarray) -> np.ndarray:
+        """Sigma Delta activation dynamics. Performs S4D dynamics.
 
-        Notes
-        -----
         This function simulates the behavior of a linear time-invariant system
-        with diagonalized state-space representation. (S4D)
+        with diagonalized state-space representation.
+        (For reference see Gu et al., 2022)
+
         The state-space equations are given by:
-        x_{k+1} = A * x_k + B * u_k
-        y_k = C * x_k
+        s4_state_{k+1} = A * s4_state_k + B * input_k
+        act_k = C * s4_state_k
 
         where:
-        - x_k is the state vector at time step k,
-        - u_k is the input vector at time step k,
-        - y_k is the output vector at time step k,
+        - s4_state_k is the state vector at time step k,
+        - input_k is the input vector at time step k,
+        - act_k is the output vector at time step k,
         - A is the diagonal state matrix,
         - B is the diagonal input matrix,
         - C is the diagonal output matrix.
 
         The function computes the next output step of the
         system for the given input signal.
+
+        Parameters
+        ----------
+        sigma_data: np.ndarray
+            sigma decoded data
+
+        Returns
+        -------
+        np.ndarray
+            activation output
         """
 
-        self.S4state = self.S4state * self.A + sigma_data * self.B
-        act = self.C * self.S4state * 2
+        self.s4_state = self.s4_state * self.a + sigma_data * self.b
+        act = self.c * self.s4_state * 2
         return act
 
 
-@implements(proc=SigmaS4Delta, protocol=LoihiProtocol)
+@implements(proc=SigmaS4dDelta, protocol=LoihiProtocol)
 @requires(CPU)
 @tag('floating_pt')
-class PySigmaS4DeltaModelFloat(AbstractSigmaS4DeltaModel):
-    """Floating point implementation of SigmaS4Delta neuron."""
+class PySigmaS4dDeltaModelFloat(AbstractSigmaS4dDeltaModel):
+    """Floating point implementation of SigmaS4dDelta neuron."""
     a_in = LavaPyType(PyInPort.VEC_DENSE, float)
     s_out = LavaPyType(PyOutPort.VEC_DENSE, float)
 
@@ -98,13 +123,13 @@ class PySigmaS4DeltaModelFloat(AbstractSigmaS4DeltaModel):
     state_exp: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=3)
     cum_error: np.ndarray = LavaPyType(np.ndarray, bool, precision=1)
     spike_exp: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=3)
-    S4_exp: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=3)
+    s4_exp: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=3)
 
-    # S4 stuff
-    S4state: np.ndarray = LavaPyType(np.ndarray, float)
-    A: np.ndarray = LavaPyType(np.ndarray, float)
-    B: np.ndarray = LavaPyType(np.ndarray, float)
-    C: np.ndarray = LavaPyType(np.ndarray, float)
+    # S4 vaiables
+    s4_state: np.ndarray = LavaPyType(np.ndarray, float)
+    a: np.ndarray = LavaPyType(np.ndarray, float)
+    b: np.ndarray = LavaPyType(np.ndarray, float)
+    c: np.ndarray = LavaPyType(np.ndarray, float)
 
     def run_spk(self) -> None:
         # Receive synaptic input
@@ -113,42 +138,42 @@ def run_spk(self) -> None:
         self.s_out.send(s_out)
 
 
-@implements(proc=SigmaS4DeltaLayer, protocol=LoihiProtocol)
+@implements(proc=SigmaS4dDeltaLayer, protocol=LoihiProtocol)
 class SubDenseLayerModel(AbstractSubProcessModel):
     def __init__(self, proc):
         """Builds (Sparse -> S4D -> Sparse) connection of the process."""
         conn_weights = proc.proc_params.get("conn_weights")
         shape = proc.proc_params.get("shape")
         state_exp = proc.proc_params.get("state_exp")
         num_message_bits = proc.proc_params.get("num_message_bits")
-        S4_exp = proc.proc_params.get("S4_exp")
+        s4_exp = proc.proc_params.get("s4_exp")
         d_states = proc.proc_params.get("d_states")
-        A = proc.proc_params.get("A")
-        B = proc.proc_params.get("B")
-        C = proc.proc_params.get("C")
+        a = proc.proc_params.get("a")
+        b = proc.proc_params.get("b")
+        c = proc.proc_params.get("c")
         vth = proc.proc_params.get("vth")
 
         # Instantiate processes
         self.sparse1 = Sparse(weights=conn_weights.T, weight_exp=state_exp,
                               num_message_bits=num_message_bits)
-        self.sigmaS4delta = SigmaS4Delta(shape=(shape[0] * d_states,),
-                                         vth=vth,
-                                         state_exp=state_exp,
-                                         S4_exp=S4_exp,
-                                         A=A,
-                                         B=B,
-                                         C=C)
+        self.sigma_S4d_delta = SigmaS4dDelta(shape=(shape[0] * d_states,),
+                                             vth=vth,
+                                             state_exp=state_exp,
+                                             s4_exp=s4_exp,
+                                             a=a,
+                                             b=b,
+                                             c=c)
         self.sparse2 = Sparse(weights=conn_weights, weight_exp=state_exp,
                               num_message_bits=num_message_bits)
 
         # Make connections Sparse -> SigmaS4Delta -> Sparse
         proc.in_ports.s_in.connect(self.sparse1.in_ports.s_in)
-        self.sparse1.out_ports.a_out.connect(self.sigmaS4delta.in_ports.a_in)
-        self.sigmaS4delta.out_ports.s_out.connect(self.sparse2.s_in)
+        self.sparse1.out_ports.a_out.connect(self.sigma_S4d_delta.in_ports.a_in)
+        self.sigma_S4d_delta.out_ports.s_out.connect(self.sparse2.s_in)
         self.sparse2.out_ports.a_out.connect(proc.out_ports.a_out)
 
-        # Set aliasses
-        proc.vars.A.alias(self.sigmaS4delta.vars.A)
-        proc.vars.B.alias(self.sigmaS4delta.vars.B)
-        proc.vars.C.alias(self.sigmaS4delta.vars.C)
-        proc.vars.S4state.alias(self.sigmaS4delta.vars.S4state)
+        # Set aliases
+        proc.vars.a.alias(self.sigma_S4d_delta.vars.a)
+        proc.vars.b.alias(self.sigma_S4d_delta.vars.b)
+        proc.vars.c.alias(self.sigma_S4d_delta.vars.c)
+        proc.vars.s4_state.alias(self.sigma_S4d_delta.vars.s4_state)