Merge pull request #9 from eimrek/release/2.2.0

Release/2.2.0

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()"

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="tb-mean-field-hubbard",
-    version="2.1.0",
+    version="2.2.0",
     author="Kristjan Eimre",
     author_email="[email protected]",
     description="Package to run tight-binding mean field hubbard calculations",

tb_mean_field_hubbard/__init__.py

@@ -1,4 +1,4 @@
 from .mfh import MeanFieldHubbardModel
 from .utils import create_itx
 
-__version__ = "2.1.0"
+__version__ = "2.2.0"

tb_mean_field_hubbard/mfh.py

@@ -19,6 +19,7 @@ class MeanFieldHubbardModel:
     """
     Class that takes care of running the Mean Field Hubbard calculation
     """
+
     def __init__(self, ase_geom, t_list=[2.7], charge=0, multiplicity='auto', bond_cutoff='auto'):
 
         self.t_list = t_list
@@ -33,6 +34,7 @@ def __init__(self, ase_geom, t_list=[2.7], charge=0, multiplicity='auto', bond_c
 
         self.ase_geom = ase_geom
         self.num_atoms = len(ase_geom)
+        self.onsite_energies = np.zeros(self.num_atoms)  # by default, no on-site energies
 
         self.figure_size = (np.ptp(self.ase_geom.positions, axis=0)[:2] + 1.0) / 4.0
         self.figure_size[0] = max([self.figure_size[0], 3.5])  # to have enough space for titles
@@ -158,6 +160,22 @@ def override_hoppings(self, custom_t_list=[]):
             self.model_a.set_hop(-t, i1, i2, mode='reset')
             self.model_b.set_hop(-t, i1, i2, mode='reset')
 
+    def set_onsite_energies(self, onsites_list):
+        """ Set the on-site energy/potentials for a list of sites
+
+        onsites_list - list of tuples (atom index, on-site energy)
+        """
+
+        for onsite in onsites_list:
+            if len(onsite) != 2:
+                print("Error: please specify (atom index, on-site energy)")
+                return
+            i, e = onsite
+            if not (0 <= i < self.num_atoms):
+                print("Error: index out of range.")
+                return
+            self.onsite_energies[i] = e
+
     def visualize_tb_model(self):
         self.model_a.visualize(0, 1)
         plt.show()
@@ -213,8 +231,8 @@ def mfh_iteration(self, u):
 
         for i_at, d in enumerate(self.spin_resolved_dens):
 
-            self.model_a.set_onsite(u * d[1], i_at, mode="reset")
-            self.model_b.set_onsite(u * d[0], i_at, mode="reset")
+            self.model_a.set_onsite(u * d[1] + self.onsite_energies[i_at], i_at, mode="reset")
+            self.model_b.set_onsite(u * d[0] + self.onsite_energies[i_at], i_at, mode="reset")
 
         # Solve the new TB
         (evals_a, evecs_a) = self.model_a.solve_all(eig_vectors=True)
@@ -309,4 +327,4 @@ def gaps(self, evals):
             return gap_a, gap_b, gap
 
         except IndexError:
-            return np.nan, np.nan, np.nan
+            return np.nan, np.nan, np.nan

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()