improve sts & orb_squared plotting functionality

example.ipynb

@@ -92,7 +92,7 @@
     "mfh_model.pp.plot_mo_eigenvector(mo_index=index, spin=spin)\n",
     "\n",
     "# Visualize the orbital squared\n",
-    "mfh_model.pp.plot_orb_squared_map(plt.gca(), evec)"
+    "mfh_model.pp.plot_orb_squared_map(evec)"
    ]
   },
   {
@@ -217,14 +217,14 @@
     "    fig, axs = plt.subplots(nrows=1, ncols=7, figsize=(7 * mfh_model.figure_size[0], mfh_model.figure_size[1]))\n",
     "    \n",
     "    mfh_model.pp.plot_no_eigenvector(i_mo, ax=axs[0])\n",
-    "    mfh_model.pp.plot_orb_squared_map(axs[1], mfh_model.pp.no_evecs[i_mo], h=h)\n",
+    "    mfh_model.pp.plot_orb_squared_map(mfh_model.pp.no_evecs[i_mo], h=h, ax=axs[1])\n",
     "    \n",
     "    mfh_model.pp.plot_mo_eigenvector(i_mo, spin=0, ax=axs[2])\n",
     "    mfh_model.pp.plot_mo_eigenvector(i_mo, spin=1, ax=axs[3])\n",
-    "    mfh_model.pp.plot_sts_map(axs[4], mfh_model.evals[0, i_mo], h=h)\n",
+    "    mfh_model.pp.plot_sts_map(mfh_model.evals[0, i_mo], h=h, ax=axs[4])\n",
     "    \n",
     "    tb_model.pp.plot_mo_eigenvector(i_mo, spin=0, ax=axs[5])\n",
-    "    tb_model.pp.plot_sts_map(axs[6], tb_model.evals[0, i_mo], h=h)\n",
+    "    tb_model.pp.plot_sts_map(tb_model.evals[0, i_mo], h=h, ax=axs[6])\n",
     "    \n",
     "    plt.subplots_adjust(wspace=0.0, hspace=0)\n",
     "    plt.show()"

tb_mean_field_hubbard/mfh_pp.py

@@ -76,11 +76,20 @@ def calc_orb_map(self, evec, h=10.0, edge_space=5.0, dx=0.1, z_eff=3.25):
 
         return orb_map, extent
 
-    def plot_orb_squared_map(self, ax, evec, h=10.0, edge_space=5.0, dx=0.1, title=None, cmap='seismic', z_eff=3.25):
+    def plot_orb_squared_map(self, evec, h=10.0, edge_space=5.0, dx=0.1, title=None, cmap='seismic', z_eff=3.25, ax=None):
         orb_map, extent = self.calc_orb_map(evec, h, edge_space, dx, z_eff)
-        ax.imshow((np.abs(orb_map)**2).T, origin='lower', cmap=cmap, extent=extent)
-        ax.axis('off')
-        ax.set_title(title)
+
+        use_ax = ax
+        if ax is None:
+            plt.figure(figsize=self.mfh.figure_size)
+            use_ax = plt.gca()
+
+        use_ax.imshow((np.abs(orb_map)**2).T, origin='lower', cmap=cmap, extent=extent)
+        use_ax.axis('off')
+        use_ax.set_title(title)
+
+        if ax is None:
+            plt.show()
 
     def calc_sts_map(self, energy, broadening=0.05, h=10.0, edge_space=5.0, dx=0.1, z_eff=3.25):
 
@@ -101,21 +110,30 @@ def calc_sts_map(self, energy, broadening=0.05, h=10.0, edge_space=5.0, dx=0.1,
         return final_map, extent
 
     def plot_sts_map(self,
-                     ax,
                      energy,
                      broadening=0.05,
                      h=10.0,
                      edge_space=5.0,
                      dx=0.1,
                      title=None,
                      cmap='seismic',
-                     z_eff=3.25):
+                     z_eff=3.25,
+                     ax=None):
 
         final_map, extent = self.calc_sts_map(energy, broadening, h, edge_space, dx, z_eff)
 
-        ax.imshow(final_map.T, origin='lower', cmap=cmap, extent=extent)
-        ax.axis('off')
-        ax.set_title(title)
+        use_ax = ax
+        if ax is None:
+            plt.figure(figsize=self.mfh.figure_size)
+            use_ax = plt.gca()
+
+        use_ax.imshow(final_map.T, origin='lower', cmap=cmap, extent=extent)
+        use_ax.axis('off')
+        use_ax.set_title(title)
+
+        if ax is None:
+            plt.show()
+
 
     def plot_eigenvector(self, ax, evec, title=None):
         utils.make_evec_plot(ax, self.mfh.ase_geom, self.mfh.neighbor_list, evec, title=title)
@@ -174,10 +192,10 @@ def report(self, num_orb=2, sts_h=10.0, sts_broad=0.05):
             self.plot_mo_eigenvector(i_mo, spin=1, ax=axs[1])
 
             title1 = "sts h=%.1f, en: %.2f" % (sts_h, self.mfh.evals[0][i_mo])
-            self.plot_sts_map(axs[2], self.mfh.evals[0][i_mo], broadening=sts_broad, h=sts_h, title=title1)
+            self.plot_sts_map(self.mfh.evals[0][i_mo], broadening=sts_broad, h=sts_h, title=title1, ax=axs[2])
 
             title2 = "sts h=%.1f, en: %.2f" % (sts_h, self.mfh.evals[1][i_mo])
-            self.plot_sts_map(axs[3], self.mfh.evals[1][i_mo], broadening=sts_broad, h=sts_h, title=title2)
+            self.plot_sts_map(self.mfh.evals[1][i_mo], broadening=sts_broad, h=sts_h, title=title2, ax=axs[3])
 
             plt.show()