consider precision of input signal, need to merge lava first

src/lava/proc/s4d/models.py

@@ -23,6 +23,7 @@ class S4dModel(PyLoihiProcessModel):
     a_in = LavaPyType(PyInPort.VEC_DENSE, float)
     s_out = LavaPyType(PyOutPort.VEC_DENSE, float)
     s4_exp: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=3)
+    inp_exp: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=3)
 
     # S4 variables
     s4_state: np.ndarray = LavaPyType(np.ndarray, complex)

src/lava/proc/s4d/process.py

@@ -18,7 +18,8 @@ def __init__(
             b: float,
             c: float,
             s4_state: ty.Optional[int] = 0,
-            s4_exp: ty.Optional[int] = 0) -> None:
+            s4_exp: ty.Optional[int] = 0,
+            inp_exp: ty.Optional[int] = 0) -> None:
         """
         Neuron process that implements S4D (described by
         Gu et al., 2022) dynamics.
@@ -55,14 +56,19 @@ def __init__(
             Scaling exponent with base 2 for the S4 state variables.
             Note: This should only be used for nc models.
             Default is 0.
+        inp_exp: int
+            Bit precision of the input signal.
+            Note: This should only be used for nc models.
+            Default is 0.
         """
 
         super().__init__(shape=shape,
                          a=a,
                          b=b,
                          c=c,
                          s4_state=s4_state,
-                         s4_exp=s4_exp)
+                         s4_exp=s4_exp,
+                         inp_exp=inp_exp)
         # Ports
         self.a_in = InPort(shape=shape)
         self.s_out = OutPort(shape=shape)
@@ -71,8 +77,9 @@ def __init__(
         self.a = Var(shape=shape, init=a)
         self.b = Var(shape=shape, init=b)
         self.c = Var(shape=shape, init=c)
-        self.s4_state = Var(shape=shape, init=0)
+        self.s4_state = Var(shape=shape, init=s4_state)
         self.s4_exp = Var(shape=(1,), init=s4_exp)
+        self.inp_exp = Var(shape=(1,), init=inp_exp)
 
     @property
     def shape(self) -> ty.Tuple[int, ...]:

tests/lava/proc/s4d/test_process.py

@@ -14,17 +14,20 @@ def test_init(self) -> None:
         """Tests instantiation of S4d"""
         shape = 10
         s4_exp = 12
+        inp_exp = 8
         a = np.ones(shape) * 0.5
         b = np.ones(shape) * 0.8
         c = np.ones(shape) * 0.9
         s4d = S4d(shape=(shape,),
                   s4_exp=s4_exp,
+                  inp_exp=inp_exp,
                   a=a,
                   b=b,
                   c=c)
 
         self.assertEqual(s4d.shape, (shape,))
         self.assertEqual(s4d.s4_exp.init, s4_exp)
+        self.assertEqual(s4d.inp_exp.init, inp_exp)
         np.testing.assert_array_equal(s4d.a.init, a)
         np.testing.assert_array_equal(s4d.b.init, b)
         np.testing.assert_array_equal(s4d.c.init, c)

tests/lava/proc/s4d/utils.py

@@ -24,107 +24,6 @@ def get_coefficients(
     return s4d_A, s4d_B, s4d_C
 
 
-def run_bit_acc_model(
-        inp: np.ndarray,
-        num_steps: int,
-        model_dim: int,
-        d_states: int,
-        a: np.ndarray,
-        b: np.ndarray,
-        c: np.ndarray,
-        s4d_exp: int,
-        is_real: bool) -> Tuple[np.ndarray]:
-    """
-    Run original S4d model in fixed precision.
-
-    This function simulates the behavior of a linear time-invariant system
-    with diagonalized state-space representation. (S4D)
-    The state-space equations are given by:
-    s4_state_{k+1} = A * s4_state_k + B * input_k
-    out_k = C * s4_state_k
-
-    where:
-    - s4_state_k is the state vector at time step k,
-    - input_k is the input vector at time step k,
-    - out_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.
-
-    The function computes the output of the system for the given input signal
-    over num_steps time steps.
-
-    Parameters
-    ----------
-    inp: np.ndarray
-        Input signal to the model.
-    num_steps: int
-        Number of time steps to simulate the model.
-    model_dim: int
-        Dimensionality of the model.
-    d_states: int
-        Number of model states.
-    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.
-    s4d_exp: int
-        Bit precision of a, b, c and the s4_state.
-    is_real: bool
-        Whether a, b, c and the s4_state are complex or real valued.
-
-    Returns
-    -------
-    Tuple[np.ndarray]
-        Tuple containing the output of the fixed precision model simulation.
-    """
-
-    a = a[:model_dim * d_states]
-    b = b[:model_dim * d_states]
-    c = c[:model_dim * d_states]
-    out = np.zeros((model_dim * d_states, num_steps))
-    expansion_weights = np.kron(np.eye(model_dim), np.ones(d_states))
-    expanded_inp = np.matmul(expansion_weights.T, inp)
-
-    if is_real:
-        a = (a * 2**s4d_exp).astype(int)
-        b = (b * 2**s4d_exp).astype(int)
-        c = (c * 2**s4d_exp).astype(int)
-        s4_state = np.zeros((model_dim * d_states,)).flatten()
-        for idx, data_in in enumerate(expanded_inp.T):
-            s4_state = (s4_state * a * 2**-s4d_exp).astype(int)
-            + (data_in * b * 2**-s4d_exp).astype(int)
-            out[:, idx] = (c * s4_state * 2**-s4d_exp).astype(int) * 2
-    else:
-        s4_state_real = np.zeros((1, model_dim * d_states)).astype(int)
-        s4_state_imag = np.zeros((1, model_dim * d_states)).astype(int)
-        a_imag = (a.imag * 2**s4d_exp).astype(int)
-        a_real = (a.real * 2**s4d_exp).astype(int)
-        b_imag = (b.imag * 2**s4d_exp).astype(int)
-        b_real = (b.real * 2**s4d_exp).astype(int)
-        c_real = (c.real * 2**s4d_exp).astype(int)
-        c_imag = (c.imag * 2**s4d_exp).astype(int)
-
-        for idx, data_in in enumerate(expanded_inp.T):
-            s4_state_real_copy = np.copy(s4_state_real)
-            s4_state_imag_copy = np.copy(s4_state_imag)
-            s4_state_real = (s4_state_real * a_real * 2**-s4d_exp).astype(int)
-            - (s4_state_imag_copy * a_imag * 2**-s4d_exp).astype(int)
-            + (data_in * b_real * 2**-s4d_exp).astype(int)
-            s4_state_imag = ((s4_state_imag * a_real) * 2**-s4d_exp).astype(int)
-            + ((s4_state_real_copy * a_imag) * 2**-s4d_exp).astype(int)
-            + ((data_in * b_imag) * 2**-s4d_exp).astype(int)
-            out_val = ((c_real * s4_state_real) * 2**-s4d_exp).astype(int)
-            - ((c_imag * s4_state_imag) * 2**-s4d_exp).astype(int)
-            out[:, idx] = 2 * out_val
-    return out
-
-
 def run_original_model(
         inp: np.ndarray,
         num_steps: int,