Refactor weight loading in dcm_app.py to use JSON instead of pickle, adding a new function for restoring arrays from loaded data.

dcm_app.py

@@ -2,7 +2,7 @@ import ase
 import jax
 import jax.numpy as jnp
 import numpy as np
-import pickle
+import json
 from rdkit import Chem
 from rdkit.Chem import AllChem
 from rdkit.Chem import Draw
@@ -75,14 +75,29 @@ def get_grid_points(coordinates):
     return grid_points
 
 
-def load_pickle_dict(path):
-    with open(path, "rb") as handle:
-        return pickle.load(handle)
+def restore_arrays(obj):
+    if isinstance(obj, dict):
+        return {key: restore_arrays(value) for key, value in obj.items()}
+    if isinstance(obj, list):
+        restored = [restore_arrays(value) for value in obj]
+        if any(isinstance(value, dict) for value in restored):
+            return restored
+        try:
+            return np.array(restored)
+        except Exception:
+            return restored
+    return obj
+
+
+def load_json_dict(path):
+    with open(path, "r", encoding="utf-8") as handle:
+        payload = json.load(handle)
+    return restore_arrays(payload)
 
 
 def load_weights():
-    dcm1_weights = load_pickle_dict("wbs/best_0.0_params_dict.pkl")
-    dcm2_weights = load_pickle_dict("wbs/dcm2-best_1000.0_params_dict.pkl")
+    dcm1_weights = load_json_dict("wbs/best_0.0_params.json")
+    dcm2_weights = load_json_dict("wbs/dcm2-best_1000.0_params.json")
     return dcm1_weights, dcm2_weights
 
 

scripts/convert_pkls_to_dicts.py

@@ -1,9 +1,20 @@
 #!/usr/bin/env python3
 import argparse
+import json
 import pickle
 from collections.abc import Mapping
 from pathlib import Path
 
+try:
+    import numpy as np
+except Exception:  # pragma: no cover - optional dependency for conversion
+    np = None
+
+try:
+    import jax.numpy as jnp
+except Exception:  # pragma: no cover - optional dependency for conversion
+    jnp = None
+
 
 def as_dict(obj):
     if isinstance(obj, dict):
@@ -15,18 +26,36 @@ def as_dict(obj):
     return obj
 
 
+def to_jsonable(obj):
+    if isinstance(obj, dict):
+        return {key: to_jsonable(value) for key, value in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [to_jsonable(value) for value in obj]
+    if np is not None:
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        if isinstance(obj, (np.integer, np.floating)):
+            return obj.item()
+    if jnp is not None and isinstance(obj, jnp.ndarray):
+        return np.array(obj).tolist() if np is not None else obj.tolist()
+    if hasattr(obj, "tolist"):
+        return obj.tolist()
+    return obj
+
+
 def convert_file(source_path, output_path):
     with source_path.open("rb") as handle:
         payload = pickle.load(handle)
     dict_payload = as_dict(payload)
-    with output_path.open("wb") as handle:
-        pickle.dump(dict_payload, handle, protocol=pickle.HIGHEST_PROTOCOL)
+    json_payload = to_jsonable(dict_payload)
+    with output_path.open("w", encoding="utf-8") as handle:
+        json.dump(json_payload, handle, indent=2, sort_keys=True)
 
 
 def build_output_path(source_path, output_dir, suffix):
     if output_dir is not None:
-        return output_dir / source_path.name
-    return source_path.with_name(f"{source_path.stem}{suffix}{source_path.suffix}")
+        return output_dir / f"{source_path.stem}{suffix}.json"
+    return source_path.with_name(f"{source_path.stem}{suffix}.json")
 
 
 def main():