Merge pull request #21 from WMD-group/Minor_Improvements

Minor improvements

script/get_eig.py

@@ -20,8 +20,6 @@
 import matplotlib.pyplot as plt
 
 
-E_MIN = -3
-E_MAX = 3
 dq = 0.01
 
 # Copied from pymatgen
@@ -83,16 +81,16 @@ def plot_eigs(eigval, q):
 
             # occpied
             indx_occ = np.where(occ > 0.9)
-            plt.plot([q+float(spin)*dq]*len(y[indx_occ]), y[indx_occ] - efermi, color=c, lw=0, marker='o')
+            plt.plot([q+float(spin)*dq]*len(y[indx_occ]), y[indx_occ] - efermi, color="C0", lw=0, marker='o')
 
             # unoccpied
             indx_unocc = np.where(occ < 0.3)
-            plt.plot([q+float(spin)*dq]*len(y[indx_unocc]), y[indx_unocc] - efermi, color=c, lw=0, marker='o', fillstyle='none')
+            plt.plot([q+float(spin)*dq]*len(y[indx_unocc]), y[indx_unocc] - efermi, color="C1", lw=0, marker='o', fillstyle='none')
 
             # halfoccpied
             indx_hocc = (0.3 <= occ) & (occ <= 0.9)
             if len(indx_hocc) > 0:
-                plt.plot([q+float(spin)*dq]*len(y[indx_hocc]), y[indx_hocc] - efermi, color=c, lw=0, marker='o', fillstyle='bottom')
+                plt.plot([q+float(spin)*dq]*len(y[indx_hocc]), y[indx_hocc] - efermi, color="C2", lw=0, marker='o', fillstyle='bottom')
 
 
 def read_poscar(i_path, l_get_sorted_symbols=False):
@@ -129,6 +127,7 @@ def plot(qs, eigvals, e_min, e_max):
     for q, eigval in zip(qs, eigvals):
         plot_eigs(eigval, q)
     plt.ylim((e_min, e_max))
+    plt.savefig("Q_eigs.pdf", bbox_inches="tight")
     plt.show()
 
 def read_data(paths, pivot_path):

src/Plotter.jl

@@ -14,7 +14,10 @@ using Plots, LaTeXStrings
 
 Plots `Potential` with its' wave functions if `lplt_wf` is `true`.
 """
-function plot_pot!(pot::Potential; lplt_wf = false, plt = nothing, color = Nothing, label = "", scale_factor = 2e-2)
+function plot_pot!(
+    pot::Potential; lplt_wf = false, plt = nothing, color = Nothing, label = "", scale_factor = 2e-2,
+    wf_sampling_rate = 5, alpha = 0.5, linealpha=1
+    )
     if plt == nothing
         plt = plot()
     end
@@ -24,19 +27,19 @@ function plot_pot!(pot::Potential; lplt_wf = false, plt = nothing, color = Nothi
     # plot wave functions
     if lplt_wf
         ϵ = pot.ϵ; χ = pot.χ
-        for i = 1:length(ϵ)
+        for i = 1:wf_sampling_rate:length(ϵ)
             plot!(plt, pot.Q, χ[i, :] * scale_factor .+ ϵ[i], lw = 0,
                   fillrange = χ[i, :] * -scale_factor .+ ϵ[i],
-                  color = color, alpha = 0.5, label = "")
+                  color = color, alpha = alpha, label = "")
         end
     end
 
     # plot function
-    plot!(plt, pot.Q, pot.E, lw = 4, color = color, label = label)
+    plot!(plt, pot.Q, pot.E, lw = 4, color = color, label = label, alpha = linealpha)
 
     # plot data
     if size(pot.QE_data)[1] > 1 
-        scatter!(plt, pot.QE_data.Q, pot.QE_data.E, color = color, label = "")
+        scatter!(plt, pot.QE_data.Q, pot.QE_data.E, color = color, label = "", alpha = linealpha)
     end
     xaxis!(L"\ Q (amu^{1\/2} Å) \ (^{}$$"); yaxis!(L"\ Energy  (eV) \ (^{}$$")