Fix handling of tig shape.

pulquero · Jun 5, 2024 · e63c35d · e63c35d
1 parent 3b330bd
commit e63c35d
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 8 deletions.
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="sgw_tools",
-    version="1.105",
+    version="2.0",
     author="Mark Hale",
     license="MIT",
     description="Spectral graph wavelet tools",

diff --git a/sgw_tools/__init__.py b/sgw_tools/__init__.py
@@ -77,18 +77,31 @@ def _tig(g, s, **kwargs):
     if g.G.is_directed() and g.G.q != 0 and ('method' not in kwargs or kwargs['method'] != 'exact'):
         raise Exception("Only method='exact' is currently supported for magnetic Laplacians.")
 
-    if s.ndim == 1:
+    if s.ndim == 1:  # single signal
         s = s[..., np.newaxis]
     assert s.ndim == 2, "Signal shape was {}".format(s.shape)
+
     if g.G.N == 1:
         kwargs['method'] = 'exact'
     tig = g.filter(s[..., np.newaxis], **kwargs)
-    if s.shape[1] == 1:  # single signal
-        tig = tig[..., np.newaxis]
-    if tig.ndim == 1:
-        tig = tig[np.newaxis, ..., np.newaxis]
-    elif tig.ndim == 2:  # single filter
-        tig = tig[..., np.newaxis]
+
+    if tig.ndim == 0:
+        tig = tig[np.newaxis, np.newaxis, np.newaxis]
+    elif tig.ndim == 1:
+        if s.shape == (1,1):  # multiple filters
+            tig = tig[np.newaxis, np.newaxis, ...]
+        elif s.shape[0] == 1:  # single node
+            tig = tig[np.newaxis, ..., np.newaxis]
+        elif s.shape[1] == 1:  # single signal
+            tig = tig[:, np.newaxis, np.newaxis]
+    elif tig.ndim == 2:
+        if s.shape[0] == 1:  # single node
+            tig = tig[np.newaxis, ...]
+        elif s.shape[1] == 1:  # single signal
+            tig = tig[:, np.newaxis, :]
+        else:  # single filter
+            tig = tig[..., np.newaxis]
+
     assert tig.shape == s.shape + (tig.shape[2],), "Tig shape was {}".format(tig.shape)
     return tig
 

diff --git a/tests/test_sgw.py b/tests/test_sgw.py
@@ -20,6 +20,11 @@ def test_tig(self):
         s = np.array([[1,2,3],[-1,-2,-3]])
         sgw._tig(g, s)
 
+        # multiple filters, single signal
+        g = gsp.filters.Heat(G, tau=[1,5])
+        s = np.array([[1],[-1]])
+        sgw._tig(g, s)
+
         # isolated node
         G = sgw.BigGraph([[2]])
         g = gsp.filters.Heat(G)
@@ -33,3 +38,13 @@ def test_tig(self):
         # multiple signals
         s = np.array([[1,2,3]])
         sgw._tig(g, s)
+
+        # multiple filters, single signal
+        g = gsp.filters.Heat(G, tau=[1,5])
+        s = np.array([1])
+        sgw._tig(g, s)
+
+        # multiple filters, multiple signals
+        g = gsp.filters.Heat(G, tau=[1,5])
+        s = np.array([[1,2,3]])
+        sgw._tig(g, s)