Upload visualization.py

visualization.py

@@ -7,13 +7,43 @@ def plot_fc_matrices(original, reconstructed, generated):
     
     vmin, vmax = -1, 1
     
-    im1 = axes[0].imshow(original, cmap='RdBu_r', vmin=vmin, vmax=vmax)
+    # Convert 1D arrays to 2D matrices if needed
+    def vector_to_matrix(vector):
+        """Convert upper triangular vector to full matrix"""
+        if len(vector.shape) == 1:
+            # Calculate the matrix size based on vector length
+            # For a vector of length n, the matrix size is (-1 + sqrt(1 + 8*n))/2
+            n = len(vector)
+            matrix_size = int((-1 + np.sqrt(1 + 8*n)) / 2)
+            
+            # Create empty matrix
+            matrix = np.zeros((matrix_size, matrix_size))
+            
+            # Fill upper triangle
+            idx = 0
+            for i in range(matrix_size):
+                for j in range(i+1, matrix_size):
+                    matrix[i, j] = vector[idx]
+                    idx += 1
+            
+            # Make symmetric
+            matrix = matrix + matrix.T
+            
+            return matrix
+        return vector
+    
+    # Convert inputs to matrices if needed
+    original_mat = vector_to_matrix(original)
+    reconstructed_mat = vector_to_matrix(reconstructed)
+    generated_mat = vector_to_matrix(generated)
+    
+    im1 = axes[0].imshow(original_mat, cmap='RdBu_r', vmin=vmin, vmax=vmax)
     axes[0].set_title('Original FC')
     
-    im2 = axes[1].imshow(reconstructed, cmap='RdBu_r', vmin=vmin, vmax=vmax)
+    im2 = axes[1].imshow(reconstructed_mat, cmap='RdBu_r', vmin=vmin, vmax=vmax)
     axes[1].set_title('Reconstructed FC')
     
-    im3 = axes[2].imshow(generated, cmap='RdBu_r', vmin=vmin, vmax=vmax)
+    im3 = axes[2].imshow(generated_mat, cmap='RdBu_r', vmin=vmin, vmax=vmax)
     axes[2].set_title('Generated FC')
     
     for ax, im in zip(axes, [im1, im2, im3]):