bug fix

Author: yohsama

__init__.py

@@ -1,3 +1,3 @@
 # __init__.py
-from . import constant, plotting, readhead, readvasp, ultils ,velocity_acf
-__all__ = ['readhead', 'ultils','constant','plotting','readvasp','velocity_acf']
+from . import constant, plotting, readhead, readvasp, utilities ,velocity_acf
+__all__ = ['readhead', 'utilities','constant','plotting','readvasp','velocity_acf']

cvasp.py

@@ -2,7 +2,7 @@
 import sys
 import matplotlib.pyplot as plt
 from readvasp import *
-import plotting2
+from plotting import plot_band,plot_dos,plot_band_dos
 import yaml
 import os
 
@@ -34,7 +34,7 @@ def from_input(config):
     config["plotset"]["plot_type"] = type
     config["groupinfo"] = [None]
     if type > 0:
-        config["Nbandpicture"] = int(input("How many bands?\n"))
+        config["Nbandpicture"] = 1 if config["plotset"]["object"] == "DOS" else int(input("How many groups to plot the band?\n"))
         config["groupinfo"] = []
         for N in range(config["Nbandpicture"]):
             print("next group\n")
@@ -51,7 +51,7 @@ def from_input(config):
                 if tmp:
                     groups.append(tmp)
             config["groupinfo"].append(groups)
-    tmp = input('Energy range, e.g -2 5  default is [-2 5]\n') or "-2 5"
+    tmp = input('Energy range, e.g. -2 5; default is [-2 5]\n') or "-2 5"
     config["plotset"]['Elim'] = [float(i) for i in tmp.split()]
     config["plotset"]['fermi'] = input("Set 0eV to fermi? (VBM,CBM,band No. or any Energy)\n"
                                        "F or 0 will change nothing.\n"
@@ -65,38 +65,38 @@ def ini_config():
     plotset = {"plot_type": 0, "object": "band",
                 "scale": 50,
                "Elim": [-2, 5], "fermi": "vbm", "int": 100}
-    figset = {"figsize": (7, 6), "dpi": 150}
+    figset = {"figsize": (7, 6), "dpi": 300,"format":"png"}
     fontset = {"fontsize": 16, "font": "Arial"}
     fileset = {"eigfile": ["EIGENVAL"], "dosfile": "DOSCAR", "prosfile": [
         "PROCAR"], "kptfile": ["KPOINTS"], "posfile": "POSCAR"}
+    paras = {"dos_paras":{},"band_paras":{},"dos_legend_paras":{},"band_legend_paras":{"loc":"upper right"}}
     config = {"method": method,
               "plotset": plotset,
               "fileset": fileset,
               "figset": figset,
-              "fontset": fontset}
+              "fontset": fontset,
+              "paras":paras}
     return config
 
 
 def read_config(file='cvasp.yml'):
     with open(file) as f:
-        data = yaml.load(f, Loader=yaml.FullLoader)
-    # print(data)
-    return data
+        config = yaml.load(f, Loader=yaml.FullLoader)
+    return config
 
 
 def cvasp(config, axs=None):
-    print(config["plotset"]["object"].lower())
     if config["plotset"]["object"].lower() == "band":
         KPT = from_kpoints.get_kpoints(config["fileset"]["kptfile"])
         if config["plotset"]["plot_type"] > 0:
             eigdata = from_procar.get_procar(config["fileset"]["prosfile"])
         else:
             eigdata = from_eigenval.get_eigenvalue(config["fileset"]["eigfile"])
-        plotting2.plot_band(KPT, eigdata, config, axs=axs)
+        plot_band(KPT, eigdata, config, axs=axs)
 
     if config["plotset"]["object"].lower() == "dos":
         dosdata = from_doscar.get_doscar(config["fileset"]["dosfile"])
-        plotting2.plot_dos(dosdata, config, axs=axs)
+        plot_dos(dosdata, config, ax=axs)
 
     if "band" in config["plotset"]["object"].lower() and \
        "dos" in config["plotset"]["object"].lower():
@@ -106,15 +106,19 @@ def cvasp(config, axs=None):
         else:
             eigdata = from_eigenval.get_eigenvalue(config["fileset"]["eigfile"])
         dosdata = from_doscar.get_doscar(config["fileset"]["dosfile"])
-        plotting2.plot_band_dos(KPT, eigdata, dosdata, config, axs=axs)
+        plot_band_dos(KPT, eigdata, dosdata, config, axs=axs)
 
 
 
 if __name__ == "__main__":
+    import matplotlib as mpl
+    # mpl.use("TkAgg")
     config = ini_config()
     try:
         data = read_config()
         config.update(data)
+        with open('cvasp.yml', "w", encoding='utf-8') as f:
+            yaml.dump(config, f)
     except Exception as e:
         config = from_input(config)
         with open('cvasp.yml', "w", encoding='utf-8') as f:
@@ -123,7 +127,6 @@ def cvasp(config, axs=None):
            ["font"], size=config["fontset"]["fontsize"])
     plt.rcParams['font.serif'] = [config["fontset"]
                                   ["font"]] + plt.rcParams['font.serif']
-    print(config["plotset"]["object"].lower())
     ispin=1
     if "band" in config["plotset"]["object"].lower():
         if config["plotset"]["plot_type"] > 0:
@@ -135,12 +138,12 @@ def cvasp(config, axs=None):
         fig, axs = plt.subplots(1, len(config["groupinfo"])*ispin,
                                 figsize=config["figset"]["figsize"], sharey=True)
     elif config["plotset"]["object"].lower() == "dos":
-        fig, ax = plt.subplots(1, 1,
+        fig, axs = plt.subplots(1, 1,
                                figsize=config["figset"]["figsize"])
     elif "band" in config["plotset"]["object"].lower() and \
         "dos" in config["plotset"]["object"].lower():
         fig, axs = plt.subplots(1, len(config["groupinfo"])*ispin + 1,
                                 figsize=config["figset"]["figsize"], sharey=True)
     cvasp(config, axs=axs)
     plt.show()
-    fig.savefig('band_dos.png', dpi=300)
+    fig.savefig(f'{config["plotset"]["object"]}.{config["figset"]["format"]}',format=config["figset"]["format"], dpi=config["figset"]["dpi"])