Update collect_structure_info in src/vasp_snake/collect_info.py

Author: singularitti

src/vasp_snake/collect_info.py

@@ -16,11 +16,21 @@ def collect_structure_info(root=".", atol=1e-6, output="structure_info.json"):
         if info["status"] == JobStatus.DONE.value:
             contcar_path = os.path.join(root, folder, "CONTCAR")
             abs_path = os.path.abspath(contcar_path)
+            outcar_path = os.path.join(root, folder, "OUTCAR")
+            total_magnetization = None
+            if os.path.exists(outcar_path):  # << Only try if OUTCAR exists
+                try:
+                    with open(outcar_path) as f:
+                        outcar_text = f.read()
+                    total_magnetization = parse_total_magnetization(outcar_text)
+                except Exception:
+                    total_magnetization = None
             if not os.path.exists(contcar_path):
                 structure_info[folder] = {
                     "abs_path": abs_path,
                     "volume": np.nan,
                     "composition": None,
+                    "last_total_magnetization": total_magnetization,
                     "reason": "CONTCAR missing",
                 }
                 continue
@@ -32,13 +42,15 @@ def collect_structure_info(root=".", atol=1e-6, output="structure_info.json"):
                     "abs_path": abs_path,
                     "volume": np.nan,
                     "composition": None,
+                    "last_total_magnetization": total_magnetization,
                     "reason": f"Failed to parse CONTCAR: {e}",
                 }
             else:
                 structure_info[folder] = {
                     "abs_path": abs_path,
                     "volume": volume,
                     "composition": composition,
+                    "last_total_magnetization": total_magnetization,
                     "reason": "Success",
                 }
     # Save as JSON