SigmaS4Delta Neuronmodel and Layer with Unittests

src/lava/proc/s4d/models.py

@@ -0,0 +1,154 @@
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: BSD-3-Clause
+# See: https://spdx.org/licenses/
+
+from lava.proc.sdn.models import AbstractSigmaDeltaModel
+from typing import Any, Dict
+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.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):
+    a_in = None
+    s_out = None
+
+    # SigmaDelta Variables
+    vth = None
+    sigma = None
+    act = None
+    residue = None
+    error = None
+    state_exp = None
+    bias = None
+
+    # S4 Variables
+    A = None
+    B = None
+    C = None
+    S4state = 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.
+
+        Parameters
+        ----------
+        sigma_data: np.ndarray
+            sigma decoded data
+
+        Returns
+        -------
+        np.ndarray
+            activation output
+
+        Notes
+        -----
+        This function simulates the behavior of a linear time-invariant system
+        with diagonalized state-space representation. (S4D)
+        The state-space equations are given by:
+        x_{k+1} = A * x_k + B * u_k
+        y_k = C * x_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,
+        - 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.
+        """
+
+        self.S4state = self.S4state * self.A + sigma_data * self.B
+        act = self.C * self.S4state * 2
+        return act
+
+
+@implements(proc=SigmaS4Delta, protocol=LoihiProtocol)
+@requires(CPU)
+@tag('floating_pt')
+class PySigmaS4DeltaModelFloat(AbstractSigmaS4DeltaModel):
+    """Floating point implementation of SigmaS4Delta neuron."""
+    a_in = LavaPyType(PyInPort.VEC_DENSE, float)
+    s_out = LavaPyType(PyOutPort.VEC_DENSE, float)
+
+    vth: np.ndarray = LavaPyType(np.ndarray, float)
+    sigma: np.ndarray = LavaPyType(np.ndarray, float)
+    act: np.ndarray = LavaPyType(np.ndarray, float)
+    residue: np.ndarray = LavaPyType(np.ndarray, float)
+    error: np.ndarray = LavaPyType(np.ndarray, float)
+    bias: np.ndarray = LavaPyType(np.ndarray, float)
+
+    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 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)
+
+    def run_spk(self) -> None:
+        # Receive synaptic input
+        a_in_data = self.a_in.recv()
+        s_out = self.dynamics(a_in_data)
+        self.s_out.send(s_out)
+
+
+@implements(proc=SigmaS4DeltaLayer, 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")
+        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")
+        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.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.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)

src/lava/proc/s4d/process.py

@@ -0,0 +1,185 @@
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: BSD-3-Clause
+# See: https://spdx.org/licenses/
+
+from lava.magma.core.process.process import AbstractProcess
+from lava.magma.core.process.variable import Var
+from lava.magma.core.process.ports.ports import InPort, OutPort
+from lava.proc.sdn.process import ActivationMode, SigmaDelta
+import typing as ty
+import numpy as np
+
+
+class SigmaS4Delta(SigmaDelta, AbstractProcess):
+    def __init__(
+            self,
+            shape: ty.Tuple[int, ...],
+            vth: ty.Union[int, float],
+            A: float,
+            B: float,
+            C: float,
+            state_exp: ty.Optional[int] = 0,
+            S4_exp: ty.Optional[int] = 0) -> None:
+        """
+        Sigma delta neuron process that implements S4D dynamics as its
+        activation function.
+
+        Parameters
+        ----------
+        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.
+
+        Notes
+        -----
+        This process simulates the behavior of a linear time-invariant system
+        with diagonal state-space representation.
+        The state-space equations are given by:
+        x_{k+1} = A * x_k + B * u_k
+        y_k = C * x_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,
+        - A is the diagonal state matrix,
+        - B is the diagonal input matrix,
+        - C is the diagonal output matrix.
+        """
+
+        super().__init__(shape=shape,
+                         vth=vth,
+                         A=A,
+                         B=B,
+                         C=C,
+                         S4state=0,
+                         state_exp=state_exp,
+                         S4_exp=S4_exp)
+
+        # Variables for S4
+        self.A = Var(shape=shape, init=A)
+        self.B = Var(shape=shape, init=B)
+        self.C = Var(shape=shape, init=C)
+        self.S4state = Var(shape=shape, init=0)
+        self.S4_exp = Var(shape=(1,), init=S4_exp)
+
+
+class SigmaS4DeltaLayer(AbstractProcess):
+    def __init__(
+            self,
+            shape: ty.Tuple[int, ...],
+            vth: ty.Union[int, float],
+            A: float,
+            B: float,
+            C: float,
+            d_states: ty.Optional[int] = 1,
+            S4_exp: ty.Optional[int] = 0,
+            num_message_bits: ty.Optional[int] = 24,
+            state_exp: ty.Optional[int] = 0) -> None:
+        """
+        Combines S4D neuron with Sparse Processes that allow for multiple
+        d_states.
+
+        Parameters
+        ----------
+        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.
+        d_states: int
+            Number of hidden states of the S4D model.
+            Default is 1.
+        state_exp: int
+            Scaling exponent with base 2 for the reconstructed sigma variables.
+            Note: Only relevant for nc model.
+            Default is 0.
+        num_message_bits: int
+            Number of message bits to be used in Sparse connection processes.
+            Note: Only relevant for nc model.
+        S4_exp: int
+            Scaling exponent with base 2 for the S4 state variables.
+            Note: Only relevant for nc model.
+            Default is 0.
+
+        Notes
+        -----
+        Connectivity: Sparse -> SigmaS4Delta -> Sparse.
+        Relieves user from computing required connection weights for multiple
+        d_states.
+
+        This process simulates the behavior of a linear time-invariant system
+        with diagonalized state-space representation. (S4D)
+        The state-space equations are given by:
+        x_{k+1} = A * x_k + B * u_k
+        y_k = C * x_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,
+        - A is the diagonal state matrix,
+        - B is the diagonal input matrix,
+        - C is the diagonal output matrix.
+        """
+
+        # Automatically takes care of expansion and reduction of dimensionality
+        # for multiple hidden states (d_states)
+        conn_weights = np.kron(np.eye(shape[0]), np.ones(d_states))
+        S4state = 0
+        super().__init__(shape=shape,
+                         vth=vth,
+                         A=A,
+                         B=B,
+                         C=C,
+                         S4_exp=S4_exp,
+                         S4state=S4state,
+                         conn_weights=conn_weights,
+                         num_message_bits=num_message_bits,
+                         d_states=d_states,
+                         state_exp=state_exp,
+                         act_mode=ActivationMode.UNIT)
+
+        # Ports
+        self.s_in = InPort(shape=shape)
+        self.a_out = OutPort(shape=shape)
+
+        # General variables
+        self.state_exp = Var(shape=(1,), init=state_exp)
+
+        # Variables for S4
+        self.A = Var(shape=(shape[0] * d_states,), init=A)
+        self.B = Var(shape=(shape[0] * d_states,), init=B)
+        self.C = Var(shape=(shape[0] * d_states,), init=C)
+        self.S4state = Var(shape=(shape[0] * d_states,), init=0)
+        self.S4_exp = Var(shape=(1,), init=S4_exp)
+
+        # Variables for connecting Dense processes
+        # Project input_dim to input_dim * d_states
+        self.conn_weights = Var(shape=shape, init=conn_weights)
+        self.num_message_bits = Var(shape=(1,), init=num_message_bits)
+
+    @property
+    def shape(self) -> ty.Tuple[int, ...]:
+        """Return shape of the Process."""
+        return self.proc_params['shape']

src/lava/proc/sdn/process.py

@@ -126,7 +126,8 @@ def __init__(
             act_mode: ty.Optional[ActivationMode] = ActivationMode.RELU,
             cum_error: ty.Optional[bool] = False,
             spike_exp: ty.Optional[int] = 0,
-            state_exp: ty.Optional[int] = 0) -> None:
+            state_exp: ty.Optional[int] = 0,
+            **kwargs) -> None:
         """Sigma delta neuron process. At the moment only ReLu activation is
         supported. Spike mechanism based on accumulated error is also supported.
 
@@ -173,7 +174,7 @@ def __init__(
         """
         super().__init__(shape=shape, vth=vth, bias=bias,
                          act_mode=act_mode, cum_error=cum_error,
-                         spike_exp=spike_exp, state_exp=state_exp)
+                         spike_exp=spike_exp, state_exp=state_exp, **kwargs)
         # scaling factor for fixed precision scaling
         vth = vth * (1 << (spike_exp + state_exp))
         bias = bias * (1 << (spike_exp + state_exp))