[ENH] add features to plot_mod_heatmap

Author: liuzhenqi77

netneurotools/plotting.py

@@ -75,6 +75,7 @@ def sort_communities(consensus, communities):
 def plot_mod_heatmap(data, communities, *, inds=None, edgecolor='black',
                      ax=None, figsize=(6.4, 4.8), xlabels=None, ylabels=None,
                      xlabelrotation=90, ylabelrotation=0, cbar=True,
+                     square=True, ticklabels=None,
                      mask_diagonal=True, **kwargs):
     """
     Plots `data` as heatmap with borders drawn around `communities`
@@ -102,6 +103,10 @@ def plot_mod_heatmap(data, communities, *, inds=None, edgecolor='black',
     {x,y}labelrotation : float, optional
         Angle of the rotation of the labels. Available only if `{x,y}labels`
         provided. Default : xlabelrotation: 90, ylabelrotation: 0
+    square : bool, optional
+        Setting the matrix with equal aspect. Default: True
+    ticklabels : list, optional
+        List of labels for each area. Default: None
     cbar : bool, optional
         Whether to plot colorbar. Default: True
     mask_diagonal : bool, optional
@@ -132,6 +137,10 @@ def plot_mod_heatmap(data, communities, *, inds=None, edgecolor='black',
     coll = ax.pcolormesh(plot_data, edgecolor='none', **kwargs)
     ax.set(xlim=(0, plot_data.shape[1]), ylim=(0, plot_data.shape[0]))
 
+    # set equal aspect
+    if square:
+        ax.set_aspect('equal')
+
     for side in ['top', 'right', 'left', 'bottom']:
         ax.spines[side].set_visible(False)
 
@@ -179,6 +188,13 @@ def plot_mod_heatmap(data, communities, *, inds=None, edgecolor='black',
                 ax.set_yticklabels(labels=ylabels, rotation=ylabelrotation)
                 ax.tick_params(left=False, bottom=False)
 
+    if ticklabels is not None:
+        labels_ind = [ticklabels[i] for i in inds]
+        ax.set_xticks(np.arange(len(labels_ind))+0.5)
+        ax.set_yticks(np.arange(len(labels_ind))+0.5)
+        ax.set_xticklabels(labels_ind, rotation=90)
+        ax.set_yticklabels(labels_ind)
+
     return ax