fix: saving model outputs (#66)

* fix: saving jax arrays

* fix: loading score

* refactor: all file openings to follow same pattern

* fix: saving jax arrays

Author: lwalew

src/mlipaudit/io.py

@@ -49,7 +49,9 @@ def write_benchmark_result_to_disk(
     _output_dir.mkdir(exist_ok=True, parents=True)
     (_output_dir / benchmark_name).mkdir(exist_ok=True)
 
-    with (_output_dir / benchmark_name / RESULT_FILENAME).open("w") as json_file:
+    with open(
+        _output_dir / benchmark_name / RESULT_FILENAME, mode="w", encoding="utf-8"
+    ) as json_file:
         json_as_str = json.loads(result.model_dump_json())  # type: ignore
         json.dump(json_as_str, json_file, indent=2)
 
@@ -71,7 +73,9 @@ def load_benchmark_result_from_disk(
     _results_dir = Path(results_dir)
     benchmark_subdir = _results_dir / benchmark_class.name
 
-    with (benchmark_subdir / RESULT_FILENAME).open("r", encoding="utf-8") as json_file:
+    with open(
+        benchmark_subdir / RESULT_FILENAME, mode="r", encoding="utf-8"
+    ) as json_file:
         json_data = json.load(json_file)
 
     return benchmark_class.result_class(**json_data)  # type: ignore
@@ -142,8 +146,8 @@ def write_scores_to_disk(
     """
     _output_dir = Path(output_dir)
     _output_dir.mkdir(exist_ok=True, parents=True)
-    with (_output_dir / SCORE_FILENAME).open("w") as json_file:
-        json.dump(scores, json_file, indent=2)
+    with open(_output_dir / SCORE_FILENAME, "w", encoding="utf-8") as f:
+        json.dump(scores, f, indent=2)
 
 
 def load_score_from_disk(
@@ -160,8 +164,9 @@ def load_score_from_disk(
         A dictionary of scores where the keys are the
             benchmark names.
     """
-    with (Path(output_dir) / SCORE_FILENAME).open("w") as json_file:
-        scores = json.load(json_file)
+    with open(Path(output_dir) / SCORE_FILENAME, mode="r", encoding="utf-8") as f:
+        scores = json.load(f)
+
     return scores
 
 
@@ -213,7 +218,7 @@ def write_model_output_to_disk(
         json_path = Path(tmpdir) / MODEL_OUTPUT_JSON_FILENAME
         arrays_path = Path(tmpdir) / MODEL_OUTPUT_ARRAYS_FILENAME
 
-        with json_path.open("w") as json_file:
+        with open(json_path, "w", encoding="utf-8") as json_file:
             json.dump(data, json_file)
 
         np.savez(arrays_path, **arrays)

src/mlipaudit/io_helpers.py

@@ -15,6 +15,7 @@
 from dataclasses import fields, is_dataclass
 from typing import Any, ClassVar, Protocol, Type, TypeVar, get_args
 
+import jax.numpy as jnp
 import numpy as np
 import pydantic
 
@@ -63,7 +64,7 @@ def recurse(value):
             return {k: recurse(v) for k, v in value.items()}
         elif isinstance(value, (list, tuple)):
             return [recurse(v) for v in value]
-        elif isinstance(value, np.ndarray):
+        elif isinstance(value, (np.ndarray, jnp.ndarray)):
             key = f"np_{counter[0]}"
             arrays[key] = value
             counter[0] += 1

tests/test_io.py

@@ -17,6 +17,7 @@
 from dataclasses import fields
 from pathlib import Path
 
+import jax.numpy as jnp
 import numpy as np
 from mlip.simulation import SimulationState
 
@@ -78,7 +79,7 @@ def test_model_outputs_io_works(
     # First, set up two different simulation states
     dummy_sim_state_1 = SimulationState(
         atomic_numbers=np.array([1, 8, 6, 1]),
-        positions=np.ones((7, 4, 3)),
+        positions=jnp.ones((7, 4, 3)),
         forces=np.random.rand(7, 4, 3),
         velocities=np.zeros((7, 4, 3)),
         temperature=np.full((7,), 1.23),