add tests test_apply_scissor_insulator and test_apply_scissor_spin_polarized (#3082)

janosh · web-flow · commit c01fae1fd23b · 2023-06-20T10:56:37.000-07:00
diff --git a/pymatgen/electronic_structure/bandstructure.py b/pymatgen/electronic_structure/bandstructure.py
@@ -335,8 +335,10 @@ def is_metal(self, efermi_tol=1e-4) -> bool:
             True if a metal, False if not
         """
         for values in self.bands.values():
-            for i in range(self.nb_bands):
-                if np.any(values[i, :] - self.efermi < -efermi_tol) and np.any(values[i, :] - self.efermi > efermi_tol):
+            for idx in range(self.nb_bands):
+                if np.any(values[idx, :] - self.efermi < -efermi_tol) and np.any(
+                    values[idx, :] - self.efermi > efermi_tol
+                ):
                     return True
         return False
 
@@ -690,7 +692,7 @@ def from_dict(cls, dct):
     def from_old_dict(cls, dct):
         """
         Args:
-            dct (dict): A dict with all data for a band structure symm line object.
+            dct (dict): A dict with all data for a band structure symmetry line object.
 
         Returns:
             A BandStructureSymmLine object
@@ -863,57 +865,56 @@ def get_branch(self, index):
             branches
         """
         to_return = []
-        for i in self.get_equivalent_kpoints(index):
+        for idx in self.get_equivalent_kpoints(index):
             for b in self.branches:
-                if b["start_index"] <= i <= b["end_index"]:
+                if b["start_index"] <= idx <= b["end_index"]:
                     to_return.append(
                         {
                             "name": b["name"],
                             "start_index": b["start_index"],
                             "end_index": b["end_index"],
-                            "index": i,
+                            "index": idx,
                         }
                     )
         return to_return
 
     def apply_scissor(self, new_band_gap):
         """
         Apply a scissor operator (shift of the CBM) to fit the given band gap.
-        If it's a metal. We look for the band crossing the fermi level
+        If it's a metal, we look for the band crossing the Fermi level
         and shift this one up. This will not work all the time for metals!
 
         Args:
             new_band_gap: the band gap the scissor band structure need to have.
 
         Returns:
-            a BandStructureSymmLine object with the applied scissor shift
+            BandStructureSymmLine: with the applied scissor shift
         """
         if self.is_metal():
-            # moves then the highest index band crossing the fermi level
-            # find this band...
+            # moves then the highest index band crossing the fermi level find this band...
             max_index = -1000
             # spin_index = None
-            for i in range(self.nb_bands):
+            for idx in range(self.nb_bands):
                 below = False
                 above = False
                 for j in range(len(self.kpoints)):
-                    if self.bands[Spin.up][i][j] < self.efermi:
+                    if self.bands[Spin.up][idx][j] < self.efermi:
                         below = True
-                    if self.bands[Spin.up][i][j] > self.efermi:
+                    if self.bands[Spin.up][idx][j] > self.efermi:
                         above = True
-                if above and below and i > max_index:
-                    max_index = i
+                if above and below and idx > max_index:
+                    max_index = idx
                     # spin_index = Spin.up
                 if self.is_spin_polarized:
                     below = False
                     above = False
                     for j in range(len(self.kpoints)):
-                        if self.bands[Spin.down][i][j] < self.efermi:
+                        if self.bands[Spin.down][idx][j] < self.efermi:
                             below = True
-                        if self.bands[Spin.down][i][j] > self.efermi:
+                        if self.bands[Spin.down][idx][j] > self.efermi:
                             above = True
-                    if above and below and i > max_index:
-                        max_index = i
+                    if above and below and idx > max_index:
+                        max_index = idx
                         # spin_index = Spin.down
             old_dict = self.as_dict()
             shift = new_band_gap
@@ -937,9 +938,9 @@ def as_dict(self):
         """
         JSON-serializable dict representation of BandStructureSymmLine.
         """
-        d = super().as_dict()
-        d["branches"] = self.branches
-        return d
+        dct = super().as_dict()
+        dct["branches"] = self.branches
+        return dct
 
 
 class LobsterBandStructureSymmLine(BandStructureSymmLine):
@@ -951,7 +952,7 @@ def as_dict(self):
         """
         JSON-serializable dict representation of BandStructureSymmLine.
         """
-        d = {
+        dct = {
             "@module": type(self).__module__,
             "@class": type(self).__name__,
             "lattice_rec": self.lattice_rec.as_dict(),
@@ -961,62 +962,62 @@ def as_dict(self):
         # kpoints are not kpoint objects dicts but are frac coords (this makes
         # the dict smaller and avoids the repetition of the lattice
         for k in self.kpoints:
-            d["kpoints"].append(k.as_dict()["fcoords"])
-        d["branches"] = self.branches
-        d["bands"] = {str(int(spin)): self.bands[spin].tolist() for spin in self.bands}
-        d["is_metal"] = self.is_metal()
+            dct["kpoints"].append(k.as_dict()["fcoords"])
+        dct["branches"] = self.branches
+        dct["bands"] = {str(int(spin)): self.bands[spin].tolist() for spin in self.bands}
+        dct["is_metal"] = self.is_metal()
         vbm = self.get_vbm()
-        d["vbm"] = {
+        dct["vbm"] = {
             "energy": vbm["energy"],
             "kpoint_index": [int(x) for x in vbm["kpoint_index"]],
             "band_index": {str(int(spin)): vbm["band_index"][spin] for spin in vbm["band_index"]},
             "projections": {str(spin): v for spin, v in vbm["projections"].items()},
         }
         cbm = self.get_cbm()
-        d["cbm"] = {
+        dct["cbm"] = {
             "energy": cbm["energy"],
             "kpoint_index": [int(x) for x in cbm["kpoint_index"]],
             "band_index": {str(int(spin)): cbm["band_index"][spin] for spin in cbm["band_index"]},
             "projections": {str(spin): v for spin, v in cbm["projections"].items()},
         }
-        d["band_gap"] = self.get_band_gap()
-        d["labels_dict"] = {}
-        d["is_spin_polarized"] = self.is_spin_polarized
+        dct["band_gap"] = self.get_band_gap()
+        dct["labels_dict"] = {}
+        dct["is_spin_polarized"] = self.is_spin_polarized
         # MongoDB does not accept keys starting with $. Add a blank space to fix the problem
         for c, label in self.labels_dict.items():
             mongo_key = c if not c.startswith("$") else " " + c
-            d["labels_dict"][mongo_key] = label.as_dict()["fcoords"]
+            dct["labels_dict"][mongo_key] = label.as_dict()["fcoords"]
         if len(self.projections) != 0:
-            d["structure"] = self.structure.as_dict()
-            d["projections"] = {str(int(spin)): np.array(v).tolist() for spin, v in self.projections.items()}
-        return d
+            dct["structure"] = self.structure.as_dict()
+            dct["projections"] = {str(int(spin)): np.array(v).tolist() for spin, v in self.projections.items()}
+        return dct
 
     @classmethod
-    def from_dict(cls, d):
+    def from_dict(cls, dct):
         """
         Args:
-            d (dict): A dict with all data for a band structure symm line
+            dct (dict): A dict with all data for a band structure symmetry line
                 object.
 
         Returns:
             A BandStructureSymmLine object
         """
         try:
             # Strip the label to recover initial string (see trick used in as_dict to handle $ chars)
-            labels_dict = {k.strip(): v for k, v in d["labels_dict"].items()}
+            labels_dict = {k.strip(): v for k, v in dct["labels_dict"].items()}
             projections = {}
             structure = None
-            if d.get("projections"):
-                if isinstance(d["projections"]["1"][0][0], dict):
+            if dct.get("projections"):
+                if isinstance(dct["projections"]["1"][0][0], dict):
                     raise ValueError("Old band structure dict format detected!")
-                structure = Structure.from_dict(d["structure"])
-                projections = {Spin(int(spin)): np.array(v) for spin, v in d["projections"].items()}
+                structure = Structure.from_dict(dct["structure"])
+                projections = {Spin(int(spin)): np.array(v) for spin, v in dct["projections"].items()}
 
             return LobsterBandStructureSymmLine(
-                d["kpoints"],
-                {Spin(int(k)): d["bands"][k] for k in d["bands"]},
-                Lattice(d["lattice_rec"]["matrix"]),
-                d["efermi"],
+                dct["kpoints"],
+                {Spin(int(k)): dct["bands"][k] for k in dct["bands"]},
+                Lattice(dct["lattice_rec"]["matrix"]),
+                dct["efermi"],
                 labels_dict,
                 structure=structure,
                 projections=projections,
@@ -1028,39 +1029,39 @@ def from_dict(cls, d):
                 "format. The old format will be retired in pymatgen "
                 "5.0."
             )
-            return LobsterBandStructureSymmLine.from_old_dict(d)
+            return LobsterBandStructureSymmLine.from_old_dict(dct)
 
     @classmethod
-    def from_old_dict(cls, d):
+    def from_old_dict(cls, dct):
         """
         Args:
-            d (dict): A dict with all data for a band structure symm line
+            dct (dict): A dict with all data for a band structure symmetry line
                 object.
 
         Returns:
             A BandStructureSymmLine object
         """
         # Strip the label to recover initial string (see trick used in as_dict to handle $ chars)
-        labels_dict = {k.strip(): v for k, v in d["labels_dict"].items()}
+        labels_dict = {k.strip(): v for k, v in dct["labels_dict"].items()}
         projections = {}
         structure = None
-        if "projections" in d and len(d["projections"]) != 0:
-            structure = Structure.from_dict(d["structure"])
+        if "projections" in dct and len(dct["projections"]) != 0:
+            structure = Structure.from_dict(dct["structure"])
             projections = {}
-            for spin in d["projections"]:
+            for spin in dct["projections"]:
                 dd = []
-                for i in range(len(d["projections"][spin])):
+                for i in range(len(dct["projections"][spin])):
                     ddd = []
-                    for j in range(len(d["projections"][spin][i])):
-                        ddd.append(d["projections"][spin][i][j])
+                    for j in range(len(dct["projections"][spin][i])):
+                        ddd.append(dct["projections"][spin][i][j])
                     dd.append(np.array(ddd))
                 projections[Spin(int(spin))] = np.array(dd)
 
         return LobsterBandStructureSymmLine(
-            d["kpoints"],
-            {Spin(int(k)): d["bands"][k] for k in d["bands"]},
-            Lattice(d["lattice_rec"]["matrix"]),
-            d["efermi"],
+            dct["kpoints"],
+            {Spin(int(k)): dct["bands"][k] for k in dct["bands"]},
+            Lattice(dct["lattice_rec"]["matrix"]),
+            dct["efermi"],
             labels_dict,
             structure=structure,
             projections=projections,
diff --git a/pymatgen/electronic_structure/tests/test_bandstructure.py b/pymatgen/electronic_structure/tests/test_bandstructure.py
@@ -61,12 +61,14 @@ def test_from_dict(self):
 
 class BandStructureSymmLineTest(PymatgenTest):
     def setUp(self):
-        self.bs = loadfn(os.path.join(PymatgenTest.TEST_FILES_DIR, "Cu2O_361_bandstructure.json"))
-        self.bs2 = loadfn(os.path.join(PymatgenTest.TEST_FILES_DIR, "CaO_2605_bandstructure.json"))
-        self.bs_spin = loadfn(os.path.join(PymatgenTest.TEST_FILES_DIR, "NiO_19009_bandstructure.json"))
-        self.bs_cbm0 = loadfn(os.path.join(PymatgenTest.TEST_FILES_DIR, "InN_22205_bandstructure.json"))
-        self.bs_cu = loadfn(os.path.join(PymatgenTest.TEST_FILES_DIR, "Cu_30_bandstructure.json"))
-        self.bs_diff_spins = loadfn(os.path.join(PymatgenTest.TEST_FILES_DIR, "VBr2_971787_bandstructure.json"))
+        self.bs: BandStructureSymmLine = loadfn(f"{PymatgenTest.TEST_FILES_DIR}/Cu2O_361_bandstructure.json")
+        self.bs2: BandStructureSymmLine = loadfn(f"{PymatgenTest.TEST_FILES_DIR}/CaO_2605_bandstructure.json")
+        self.bs_spin: BandStructureSymmLine = loadfn(f"{PymatgenTest.TEST_FILES_DIR}/NiO_19009_bandstructure.json")
+        self.bs_cbm0: BandStructureSymmLine = loadfn(f"{PymatgenTest.TEST_FILES_DIR}/InN_22205_bandstructure.json")
+        self.bs_cu: BandStructureSymmLine = loadfn(f"{PymatgenTest.TEST_FILES_DIR}/Cu_30_bandstructure.json")
+        self.bs_diff_spins: BandStructureSymmLine = loadfn(
+            f"{PymatgenTest.TEST_FILES_DIR}/VBr2_971787_bandstructure.json"
+        )
         warnings.simplefilter("ignore")
 
     def tearDown(self):
@@ -188,7 +190,7 @@ def test_get_sym_eq_kpoints_and_degeneracy(self):
         cbm_k = bs.get_cbm()["kpoint"].frac_coords
         vbm_k = bs.get_vbm()["kpoint"].frac_coords
         assert bs.get_kpoint_degeneracy(cbm_k) is None
-        bs.structure = loadfn(os.path.join(PymatgenTest.TEST_FILES_DIR, "CaO_2605_structure.json"))
+        bs.structure: BandStructureSymmLine = loadfn(f"{PymatgenTest.TEST_FILES_DIR}/CaO_2605_structure.json")
         assert bs.get_kpoint_degeneracy(cbm_k) == 3
         assert bs.get_kpoint_degeneracy(vbm_k) == 1
         cbm_eqs = bs.get_sym_eq_kpoints(cbm_k)
@@ -229,11 +231,31 @@ def test_old_format_load(self):
             bs_old = BandStructureSymmLine.from_dict(d)
             assert bs_old.get_projection_on_elements()[Spin.up][0][0]["Zn"] == 0.0971
 
+    def test_apply_scissor_insulator(self):
+        # test applying a scissor operator to a metal
+        for scissor in (1, 3):
+            bs_scissored = self.bs.apply_scissor(scissor)
+            assert not bs_scissored.is_metal()
+            assert bs_scissored.nb_bands == 48
+            assert bs_scissored.efermi == approx(3.75640309 + scissor)
+            orig_efermi = self.bs_spin.efermi
+            assert bs_scissored.efermi != approx(orig_efermi)
+
+    def test_apply_scissor_spin_polarized(self):
+        # test applying a scissor operator to a spin-polarized system
+        bs_scissored = self.bs_spin.apply_scissor(1.0)
+        assert bs_scissored.is_metal()
+        assert bs_scissored.nb_bands == 27
+        assert {*bs_scissored.bands} == {Spin.up, Spin.down}
+        assert bs_scissored.efermi == approx(4.64005999)
+        orig_efermi = self.bs_spin.efermi
+        assert bs_scissored.efermi != approx(orig_efermi)
+
 
 class ReconstructBandStructureTest(PymatgenTest):
     def setUp(self):
-        self.bs_cu = loadfn(os.path.join(PymatgenTest.TEST_FILES_DIR, "Cu_30_bandstructure.json"))
-        self.bs_cu2 = loadfn(os.path.join(PymatgenTest.TEST_FILES_DIR, "Cu_30_bandstructure.json"))
+        self.bs_cu: BandStructureSymmLine = loadfn(f"{PymatgenTest.TEST_FILES_DIR}/Cu_30_bandstructure.json")
+        self.bs_cu2: BandStructureSymmLine = loadfn(f"{PymatgenTest.TEST_FILES_DIR}/Cu_30_bandstructure.json")
         warnings.simplefilter("ignore")
 
     def tearDown(self):
