Refactor weight loading in dcm_app.py to use pickle for parameter files, replacing pandas with a dedicated function for loading dictionaries.

dcm_app.py

@@ -2,7 +2,7 @@ import ase
 import jax
 import jax.numpy as jnp
 import numpy as np
-import pandas as pd
+import pickle
 from rdkit import Chem
 from rdkit.Chem import AllChem
 from rdkit.Chem import Draw
@@ -75,9 +75,14 @@ def get_grid_points(coordinates):
     return grid_points
 
 
+def load_pickle_dict(path):
+    with open(path, "rb") as handle:
+        return pickle.load(handle)
+
+
 def load_weights():
-    dcm1_weights = pd.read_pickle("wbs/best_0.0_params.pkl")
-    dcm2_weights = pd.read_pickle("wbs/dcm2-best_1000.0_params.pkl")
+    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")
     return dcm1_weights, dcm2_weights
 
 

scripts/convert_pkls_to_dicts.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+import argparse
+import pickle
+from collections.abc import Mapping
+from pathlib import Path
+
+
+def as_dict(obj):
+    if isinstance(obj, dict):
+        return {key: as_dict(value) for key, value in obj.items()}
+    if isinstance(obj, Mapping):
+        return {key: as_dict(value) for key, value in obj.items()}
+    if hasattr(obj, "to_dict"):
+        return as_dict(obj.to_dict())
+    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)
+
+
+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}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert pickled weights into plain dictionaries."
+    )
+    parser.add_argument(
+        "--input-dir",
+        default="wbs",
+        help="Directory containing .pkl files (default: wbs).",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default=None,
+        help="Directory to write dictionary pickles. Defaults next to input.",
+    )
+    parser.add_argument(
+        "--suffix",
+        default="_dict",
+        help="Suffix to append before extension (default: _dict).",
+    )
+    args = parser.parse_args()
+
+    input_dir = Path(args.input_dir)
+    if not input_dir.is_dir():
+        raise SystemExit(f"Input dir not found: {input_dir}")
+
+    output_dir = Path(args.output_dir) if args.output_dir else None
+    if output_dir is not None:
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+    pkls = sorted(input_dir.glob("*.pkl"))
+    if not pkls:
+        raise SystemExit(f"No .pkl files found in {input_dir}")
+
+    for source_path in pkls:
+        if source_path.stem.endswith(args.suffix):
+            continue
+        output_path = build_output_path(source_path, output_dir, args.suffix)
+        convert_file(source_path, output_path)
+        print(f"{source_path} -> {output_path}")
+
+
+if __name__ == "__main__":
+    main()

wbs/best_0.0_params_dict.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:580eff73f9c8bb5e3d67766bf7d923d8855943ccf0b25fb463f0fc0b65e162a1
+size 21694

wbs/best_1000.0_params_dict.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f859c8c20019e87081e319fb03771c22415dca8b33c5a15163b3a2830249e933
+size 22663

wbs/dcm2-best_1000.0_params_dict.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be8f1a6f7648fc4549a06d32a588a7684f03a6f54dfcc785eefcb6c588f32319
+size 22179

wbs/dcm3-best_1000.0_params_dict.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f859c8c20019e87081e319fb03771c22415dca8b33c5a15163b3a2830249e933
+size 22663

wbs/dcm4-best_1000.0_params_dict.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68ac932b26f93df7f99a7c911e5ed111f6818ddb65d32a56b0110e5f82199933
+size 23158