Improve plotting

ex2.py

@@ -30,8 +30,8 @@ class bandstructure:
         """
         self.mat = material
 
-        max_G = np.sqrt(11) * 2*np.pi/si.a
-        self.lat_vecs = si.gen_lat_vecs(max_G)
+        max_G = np.sqrt(11) * 2*np.pi/self.mat.a
+        self.lat_vecs = self.mat.gen_lat_vecs(max_G)
         self.assemble_hamiltonian_wo_diag()
 
         self.E_vals = None
@@ -87,7 +87,7 @@ class bandstructure:
         ham = self.hamiltonian(k)
         return np.linalg.eigvalsh(ham)
 
-    def compute_bandstructure(self, num_points: int = 10) -> None:
+    def compute_bandstructure(self, num_points: int = 20) -> None:
         """
         Generates all Energy Eigenvalues along the L-Gamma-X-U-Gamma
         path of the irreducible wedge
@@ -96,7 +96,10 @@ class bandstructure:
             num_points (int): the number of k-vectors between two high-symmetry points
         """
         E_vals = []
-        k_vecs = []
+
+        self.tick_positions = []
+        self.tick_labels = []
+        i = 0
 
         k0 = None
         high_symmetry_points = ["L", "Gamma", "X", "U", "Gamma"]
@@ -107,18 +110,22 @@ class bandstructure:
             if k0 is not None:
                 vecs = self.mat.gen_k_vecs(k0, k1, num_points)
                 for k in vecs:
-                    k_vecs.append(k)
                     E_vals.append(self.get_E(k))
+                    i += 1
 
-            k_vecs.append(k1)
             E_vals.append(self.get_E(k1))
 
+            # One tick for every high symmetry point
+            self.tick_positions.append(i)
+            if point == "Gamma":
+                point = r"$\Gamma$"
+            self.tick_labels.append(point) 
+
             k0 = np.copy(k1)
 
         self.E_vals = np.array(E_vals)
-        self.k_vecs = np.array(k_vecs)
 
-    def plot_bandstructure(self, filename: Optional[str] = None, num_bands: int = 3) -> None:
+    def plot_bandstructure(self, filename: Optional[str] = None, num_bands: int = 6) -> None:
         """
         Plots the num_bands lowest bands
         and saves the plot to filename
@@ -135,14 +142,12 @@ class bandstructure:
                                layout="constrained",
                                )
 
-        tick_positions = [0, 11, 22, 33, 44]
-        tick_labels = ["L", "Gamma", "X", "U", "Gamma"]
-
         for band in range(num_bands):
             ax.plot(self.E_vals[:,band] / c.e)
 
-        ax.set_xticks(tick_positions)
-        ax.set_xticklabels(tick_labels)
+        ax.set_xticks(self.tick_positions)
+        ax.set_xticklabels(self.tick_labels)
+        ax.grid()
 
         ax.set_ylabel("Energy in eV", fontsize=12)
 
@@ -150,10 +155,11 @@ class bandstructure:
 
 
 if __name__ == "__main__":
-    si = fcc("Si")
-    band = bandstructure(si)
+    for mat_name in ["Si", "Ge", "ZnSe"]: 
+        mat = fcc(mat_name)
 
-    band.compute_bandstructure()
-    band.plot_bandstructure()
+        band = bandstructure(mat)
+        band.compute_bandstructure()
+        band.plot_bandstructure()