Small updates (#20)

* add sessions, return viz, default suffix to csv

* rename test data files

graphutils/graph_io.py

@@ -51,7 +51,7 @@ class NdmgDirectory:
         Send all graph files to a directory of your choosing
     """
 
-    def __init__(self, directory, atlas="", suffix="ssv", delimiter=" "):
+    def __init__(self, directory, atlas="", suffix="csv", delimiter=" "):
         if not isinstance(directory, (str, Path)):
             message = f"Directory must be type str or Path. Instead, it is type {type(directory)}."
             raise TypeError(message)
@@ -215,7 +215,7 @@ def __init__(self, *args, **kwargs):
         self.vertices = self._vertices()
         self.graphs = self._graphs()
         self.subjects = self._parse()[0]
-        self.sessions = self._parse()[0]
+        self.sessions = self._parse()[1]
 
     def __repr__(self):
         return f"NdmgGraphs : {str(self.directory)}"
@@ -229,9 +229,10 @@ def _nx_graphs(self):
         nx_graphs : List[nx.Graph]
             List of networkX graphs corresponding to subjects.
         """
+        fils = [str(name) for name in self.files]
         nx_graphs = [
             nx.read_weighted_edgelist(f, nodetype=int, delimiter=self.delimiter)
-            for f in self.files
+            for f in fils
         ]
         return nx_graphs
 

graphutils/graph_stats.py

@@ -121,6 +121,62 @@ def discriminability(self, PTR=True, **kwargs):
 
         return discr_stat(self.X, self.Y, **kwargs)
 
+    def visualize(self, i, savedir=""):
+        """
+        Visualize the ith graph of self.graphs, passed-to-ranks.
+        
+        Parameters
+        ----------
+        i : int
+            Graph to visualize.
+        savedir : str, optional
+            Directory to save graph into.
+            If left empty, do not save.
+        """
+
+        nmax = np.max(self.graphs)
+
+        if isinstance(i, int):
+            graph = pass_to_ranks(self.graphs[i])
+            sub = self.subjects[i]
+            sesh = self.sessions[i]
+
+        elif isinstance(i, np.ndarray):
+            graph = pass_to_ranks(i)
+            sub = ""
+            sesh = ""
+
+        else:
+            raise TypeError("Passed value must be integer or np.ndarray.")
+
+        viz = heatmap(
+            graph, title=f"sub-{sub}_session-{sesh}", xticklabels=True, yticklabels=True
+        )
+
+        # set color of title
+        viz.set_title(viz.get_title(), color="black")
+
+        # set color of colorbar ticks
+        viz.collections[0].colorbar.ax.yaxis.set_tick_params(color="black")
+
+        # set font size and color of heatmap ticks
+        for item in viz.get_xticklabels() + viz.get_yticklabels():
+            item.set_color("black")
+            item.set_fontsize(7)
+
+        if savedir:
+            p = Path(savedir).resolve()
+            if not p.is_dir():
+                p.mkdir()
+            plt.savefig(
+                p / f"sub-{sub}_sesh-{sesh}.png",
+                facecolor="white",
+                bbox_inches="tight",
+                dpi=300,
+            )
+
+        return viz
+
 
     def visualize(self, i, savedir=""):
         """