Add nbands argument to get_band_occupancy_info function

wladerer · Mar 19, 2024 · 85158f9 · 85158f9
1 parent 34f7f48
commit 85158f9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/vsh/scripts/__init__.py b/vsh/scripts/__init__.py
@@ -637,6 +637,7 @@ def wavecar(subparsers):
     info_choices.add_argument(
         "--occ", help="Print bands near the Fermi level to stdout", action="store_true"
     )
+    info_choices.add_argument("--nbands", help="Number of bands to print for --occ", type=int, default=19)
 
     subp_wavecar.add_argument("input", help="WAVECAR or cube file")
     subp_wavecar.add_argument("-S", "--structure", help="Structure file")

diff --git a/vsh/scripts/wavecar.py b/vsh/scripts/wavecar.py
@@ -153,7 +153,7 @@ def print_fermi_energy(args):
     print(f"{wave.efermi}")
 
 
-def get_band_occupancy_info(wavecar: Wavecar):
+def get_band_occupancy_info(wavecar: Wavecar, nbands=19):
     """Gets the band occupancy information from a WAVECAR file."""
     import pandas as pd
 
@@ -181,12 +181,12 @@ def get_band_occupancy_info(wavecar: Wavecar):
     # Remove duplicate bands
     df = df.drop_duplicates()
 
-    return df[df["relative_occupancy"] > 0].tail(19)
+    return df[df["relative_occupancy"] > 0].tail(nbands)
 
 
 def print_band_occupancy_info(args):
     wave = Wavecar(args.input)
-    df = get_band_occupancy_info(wave)
+    df = get_band_occupancy_info(wave, args.nbands)
     # print dataframe without index
 
     print(df.to_string(index=False))