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,8 @@ class Sparse(AbstractProcess):
 
     Parameters
     ----------
-    weights : scipy.sparse.spmatrix
-        2D connection weight matrix as sparse matrix of form
+    weights : scipy.sparse.spmatrix or np.ndarray
+        2D connection weight matrix of form
         (num_flat_output_neurons, num_flat_input_neurons).
 
     weight_exp : int, optional
@@ -57,7 +57,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 +69,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 +93,8 @@ 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
+        2D connection weight matrix of form
         (num_flat_output_neurons, num_flat_input_neurons).
 
     weight_exp : int, optional
@@ -150,7 +153,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 +175,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 +195,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 +207,7 @@ def __init__(self,
 
         Parameters
         ----------
-        weights : spmatrix
+        weights : scipy.sparse.spmatrix or np.ndarray
             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
@@ -41,7 +40,6 @@ def test_init(self):
 
         conn = Sparse(weights=weights_sparse)
 
-        self.assertIsInstance(conn.weights.init, spmatrix)
         np.testing.assert_array_equal(conn.weights.init.toarray(), weights)
 
 
@@ -72,7 +70,6 @@ def test_init(self):
         conn = LearningSparse(weights=weights_sparse,
                               learning_rule=learning_rule)
 
-        self.assertIsInstance(conn.weights.init, spmatrix)
         np.testing.assert_array_equal(conn.weights.init.toarray(), weights)
 
 
@@ -95,7 +92,6 @@ def test_init(self):
 
         conn = DelaySparse(weights=weights_sparse, delays=delays_sparse)
 
-        self.assertIsInstance(conn.weights.init, spmatrix)
         np.testing.assert_array_equal(conn.weights.init.toarray(), weights)
 
     def test_validate_shapes(self):
@@ -132,3 +128,49 @@ def test_validate_nonzero_delays(self):
                                  DelaySparse,
                                  weights=weights_sparse,
                                  delays=delays_sparse)
+
+    def test_init_of_sparse_with_ndarray(self):
+        """Tests instantiation of Sparse with ndarray as
+         weights"""
+
+        shape = (3, 2)
+        weights = np.random.random(shape)
+
+        conn = Sparse(weights=weights)
+
+        np.testing.assert_array_equal(conn.weights.get().toarray(), weights)
+
+    def test_init_of_learningsparse_with_ndarray(self):
+        """Tests instantiation of LearningSparse with
+        ndarray as weights"""
+
+        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)
+
+        np.testing.assert_array_equal(conn.weights.get().toarray(), weights)
+
+    def test_init_of_delaysparse_with_ndarray(self):
+        """Tests instantiation of DelaySparse with ndarray as weights"""
+
+        shape = (3, 2)
+        weights = np.random.random(shape)
+        delays = np.random.randint(0, 3, shape)
+
+        conn = DelaySparse(weights=weights, delays=delays)
+
+        np.testing.assert_array_equal(conn.weights.get().toarray(), weights)
+
+
+if __name__ == '__main__':
+    unittest.main()