refactor: improve status reporting structure and format detection

- Wrapped folder status results in a dict with "summary" (showing percent of each status) and "details" (per-folder info).
- Added computation of status percentages (PENDING, DONE, NOT_CONVERGED) for the summary.
- Updated `write_status_report` to determine output format (JSON/YAML) automatically from the file extension.
- Ensured `forces_sum` is always included and uses `np.nan` for missing values, with conversion to JSON/YAML-safe output.
- Improved robustness of serialization for nested structures and NaN values.
- Updated CLI to reflect the new API and output handling.

Author: singularitti

src/vasp_snake/report_status.py

@@ -1,13 +1,11 @@
 import json
 import os
+from collections import Counter
 from enum import Enum
 
+import click
 import numpy as np
-
-try:
-    import yaml
-except ImportError:
-    yaml = None
+import yaml
 
 from vasp_snake.force import parse_forces_and_check_zero
 
@@ -21,20 +19,20 @@ class JobStatus(Enum):
 
 
 def classify_folders(root=".", atol=1e-6):
-    status = {}
+    details = {}
     for folder in sorted(os.listdir(root)):
         folder_path = os.path.join(root, folder)
         if not os.path.isdir(folder_path) or folder.startswith("."):
             continue
         outcar = os.path.join(folder_path, "OUTCAR")
         if not os.path.exists(outcar):
-            forces_sum = [np.nan] * 3
+            forces_sum = [np.nan, np.nan, np.nan]
             job_status = JobStatus.PENDING
             reason = "OUTCAR missing"
         else:
             forces_sum, is_converged = parse_forces_and_check_zero(outcar, atol=atol)
             if forces_sum is None:
-                forces_sum = [np.nan] * 3
+                forces_sum = [np.nan, np.nan, np.nan]
                 job_status = JobStatus.NOT_CONVERGED
                 reason = "No force block found"
             elif is_converged:
@@ -49,52 +47,71 @@ def classify_folders(root=".", atol=1e-6):
                 float(f)
                 for f in (forces_sum if forces_sum is not None else [np.nan] * 3)
             ]
-        status[folder] = {
+        details[folder] = {
             "status": job_status.value,
             "forces_sum": forces_sum,
             "reason": reason,
         }
-    return status
+    # Compute summary percentages
+    status_list = [v["status"] for v in details.values()]
+    total = len(status_list)
+    counter = Counter(status_list)
+    summary = {
+        status: counter.get(status, 0) / total if total else 0.0
+        for status in [s.value for s in JobStatus]
+    }
+    return {"summary": summary, "details": details}
+
 
+def write_status_report(status_dict, filename):
+    # Numpy floats/nans don't serialize well with json/yaml, so convert
+    def convert(obj):
+        if isinstance(obj, float) and np.isnan(obj):
+            return None  # Or "NaN" if you want the string
+        if isinstance(obj, (np.generic, np.ndarray)):
+            return obj.tolist()
+        return obj
+
+    # Recursively convert NaN in all nested structures
+    def recursive_convert(o):
+        if isinstance(o, dict):
+            return {k: recursive_convert(v) for k, v in o.items()}
+        elif isinstance(o, list):
+            return [recursive_convert(x) for x in o]
+        else:
+            return convert(o)
 
-def write_status_report(status_dict, filename, out_format="json"):
-    if out_format == "json":
+    output = recursive_convert(status_dict)
+    ext = os.path.splitext(filename)[-1].lower()
+    if ext in [".json"]:
         with open(filename, "w") as f:
-            json.dump(status_dict, f, indent=2)
-    elif out_format == "yaml":
-        if yaml is None:
-            raise ImportError("pyyaml is required for YAML output.")
+            json.dump(output, f, indent=2)
+    elif ext in [".yaml", ".yml"]:
         with open(filename, "w") as f:
-            yaml.dump(status_dict, f, sort_keys=False)
+            yaml.dump(output, f, sort_keys=False)
     else:
-        raise ValueError("Unsupported format: {}".format(out_format))
+        raise ValueError(f"Unknown file extension: {ext}")
 
 
-# Optional: CLI interface
-if __name__ == "__main__":
-    import click
+@click.command()
+@click.option(
+    "--output",
+    default="vasp_status.json",
+    help="Output file name (default: vasp_status.json)",
+)
+@click.option("--folders", default=".", help="Root directory containing folders")
+@click.option(
+    "--atol",
+    default=1e-6,
+    type=float,
+    show_default=True,
+    help="Convergence tolerance for force sum norm",
+)
+def main(output, folders, atol):
+    status = classify_folders(folders, atol=atol)
+    write_status_report(status, output)
 
-    @click.command()
-    @click.option(
-        "--format",
-        "out_format",
-        type=click.Choice(["json", "yaml"], case_sensitive=False),
-        default="json",
-        help="Output format (json or yaml)",
-    )
-    @click.option("--output", default=None, help="Output file name (optional)")
-    @click.option("--folders", default=".", help="Root directory containing folders")
-    @click.option(
-        "--atol",
-        default=1e-6,
-        type=float,
-        show_default=True,
-        help="Convergence tolerance for force sum norm",
-    )
-    def main(out_format, output, folders, atol):
-        status = classify_folders(folders, atol=atol)
-        if output is None:
-            output = f"vasp_status.{out_format}"
-        write_status_report(status, output, out_format)
 
+# Optional: CLI interface
+if __name__ == "__main__":
     main()