Version 0.3.0 changes

apax/md/ase_calc.py

@@ -112,6 +112,7 @@ def __init__(
         self.model_config, self.params = restore_parameters(model_dir)
         self.n_models = check_for_ensemble(self.params)
         self.padding_factor = padding_factor
+        self.padded_length = 0
 
         if self.model_config.model.calc_stress:
             self.implemented_properties.append("stress")
@@ -148,6 +149,10 @@ def initialize(self, atoms):
         self.step = get_step_fn(model, atoms, self.neigbor_from_jax)
         self.neighbor_fn = neighbor_fn
 
+        if self.neigbor_from_jax:
+            positions = jnp.asarray(atoms.positions, dtype=jnp.float64)
+            self.neighbors = self.neighbor_fn.allocate(positions)
+
     def set_neighbours_and_offsets(self, atoms, box):
         idxs_i, idxs_j, offsets = neighbour_list("ijS", atoms, self.r_max)
 

apax/md/nvt.py

@@ -271,7 +271,7 @@ def body_fn(i, state):
                     )
                 ckpt = {"state": state, "step": step}
                 checkpoints.save_checkpoint(
-                    ckpt_dir=ckpt_dir,
+                    ckpt_dir=ckpt_dir.resolve(),
                     target=ckpt,
                     step=step,
                     overwrite=True,

apax/train/checkpoints.py

@@ -89,9 +89,9 @@ class CheckpointManager:
     def __init__(self) -> None:
         self.async_manager = checkpoints.AsyncManager()
 
-    def save_checkpoint(self, ckpt, epoch: int, path: str) -> None:
+    def save_checkpoint(self, ckpt, epoch: int, path: Path) -> None:
         checkpoints.save_checkpoint(
-            ckpt_dir=path,
+            ckpt_dir=path.resolve(),
             target=ckpt,
             step=epoch,
             overwrite=True,