Allow np.array as input weights for Sparse

src/lava/proc/sparse/process.py

@@ -3,7 +3,7 @@
 # See: https://spdx.org/licenses/
 
 import numpy as np
-from scipy.sparse import spmatrix
+from scipy.sparse import spmatrix, csr_matrix
 import typing as ty
 
 from lava.magma.core.process.process import AbstractProcess, LogConfig
@@ -21,8 +21,9 @@ class Sparse(AbstractProcess):
 
     Parameters
     ----------
-    weights : scipy.sparse.spmatrix
-        2D connection weight matrix as sparse matrix of form
+    weights : scipy.sparse.spmatrix or np.ndarray
+        both get directly converted to a csr_matrix
+        2D connection weight matrix of form
         (num_flat_output_neurons, num_flat_input_neurons).
 
     weight_exp : int, optional
@@ -57,7 +58,7 @@ class Sparse(AbstractProcess):
         """
     def __init__(self,
                  *,
-                 weights: spmatrix,
+                 weights: ty.Union[spmatrix, np.ndarray],
                  name: ty.Optional[str] = None,
                  num_message_bits: ty.Optional[int] = 0,
                  log_config: ty.Optional[LogConfig] = None,
@@ -69,7 +70,10 @@ def __init__(self,
                          **kwargs)
 
         # Transform weights to csr matrix
-        weights = weights.tocsr()
+        if isinstance(weights, np.ndarray):
+            weights = csr_matrix(weights)
+        else:
+            weights = weights.tocsr()
 
         shape = weights.shape
 
@@ -90,8 +94,9 @@ class LearningSparse(LearningConnectionProcess, Sparse):
 
     Parameters
     ----------
-    weights : scipy.sparse.spmatrix
-        2D connection weight matrix as sparse matrix of form
+    weights : scipy.sparse.spmatrix or np.ndarray
+        both get directly converted to a csr_matrix
+        2D connection weight matrix of form
         (num_flat_output_neurons, num_flat_input_neurons).
 
     weight_exp : int, optional
@@ -150,7 +155,7 @@ class LearningSparse(LearningConnectionProcess, Sparse):
         """
     def __init__(self,
                  *,
-                 weights: spmatrix,
+                 weights: ty.Union[spmatrix, np.ndarray],
                  name: ty.Optional[str] = None,
                  num_message_bits: ty.Optional[int] = 0,
                  log_config: ty.Optional[LogConfig] = None,
@@ -172,7 +177,10 @@ def __init__(self,
                          **kwargs)
 
         # Transform weights to csr matrix
-        weights = weights.tocsr()
+        if isinstance(weights, np.ndarray):
+            weights = csr_matrix(weights)
+        else:
+            weights = weights.tocsr()
 
         shape = weights.shape
 
@@ -189,7 +197,7 @@ def __init__(self,
 class DelaySparse(Sparse):
     def __init__(self,
                  *,
-                 weights: spmatrix,
+                 weights: ty.Union[spmatrix, np.ndarray],
                  delays: ty.Union[spmatrix, int],
                  max_delay: ty.Optional[int] = 0,
                  name: ty.Optional[str] = None,
@@ -201,7 +209,8 @@ def __init__(self,
 
         Parameters
         ----------
-        weights : spmatrix
+        weights : scipy.sparse.spmatrix or np.ndarray
+            both get directly converted to a csr_matrix
             2D connection weight matrix of form (num_flat_output_neurons,
             num_flat_input_neurons) in C-order (row major).
 

tests/lava/proc/sparse/test_process.py

@@ -15,7 +15,6 @@ class TestFunctions(unittest.TestCase):
     """Test helper function for Sparse"""
 
     def test_find_with_explicit_zeros(self):
-
         mat = np.random.randint(-10, 10, (3, 5))
         spmat = csr_matrix(mat)
         spmat.data[0] = 0
@@ -132,3 +131,52 @@ def test_validate_nonzero_delays(self):
                                  DelaySparse,
                                  weights=weights_sparse,
                                  delays=delays_sparse)
+
+    def test_ndarray_gets_converted_into_sparse(self):
+        """Tests if ndarray weights get converted to a spmatrix"""
+
+        shape = (3, 2)
+        weights = np.random.random(shape)
+
+        conn = Sparse(weights=weights)
+
+        self.assertIsInstance(conn.weights.get(), spmatrix)
+        np.testing.assert_array_equal(conn.weights.get().toarray(), weights)
+
+    def test_ndarray_gets_converted_into_sparse_for_learning(self):
+        """Tests if ndarray weights get converted to a spmatrix for
+        LearingSparse"""
+
+        shape = (3, 2)
+        weights = np.random.random(shape)
+
+        learning_rule = STDPLoihi(
+            learning_rate=1,
+            A_plus=1,
+            A_minus=-2,
+            tau_plus=10,
+            tau_minus=10,
+            t_epoch=2,
+        )
+
+        conn = LearningSparse(weights=weights, learning_rule=learning_rule)
+
+        self.assertIsInstance(conn.weights.get(), spmatrix)
+        np.testing.assert_array_equal(conn.weights.get().toarray(), weights)
+
+    def test_ndarray_gets_converted_into_sparse_for_delay(self):
+        """Tests if ndarray weights get converted to a spmatrix for
+        DelaySparse"""
+
+        shape = (3, 2)
+        weights = np.random.random(shape)
+        delays = np.random.randint(0, 3, shape)
+
+        conn = DelaySparse(weights=weights, delays=delays)
+
+        self.assertIsInstance(conn.weights.get(), spmatrix)
+        np.testing.assert_array_equal(conn.weights.get().toarray(), weights)
+
+
+if __name__ == '__main__':
+    unittest.main()