Merge branch 'dev' into bal

Author: M-R-Schaefer

README.md

@@ -39,12 +39,12 @@ pip install --upgrade pip
 
 CUDA 12 installation. Wheels only available on linux.
 ```bash
-pip install --upgrade "jax[cuda12_pip]==0.4.13" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install --upgrade "jax[cuda12_pip]==0.4.14" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 CUDA 11 installation. Wheels only available on linux.
 ```bash
-pip install --upgrade "jax[cuda11_pip]==0.4.13" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install --upgrade "jax[cuda11_pip]==0.4.14" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 See the [Jax installation instructions](https://github.com/google/jax#installation) for more details.

apax/md/ase_calc.py

@@ -8,6 +8,7 @@
 from ase.calculators.calculator import Calculator, all_changes
 from flax.traverse_util import flatten_dict
 from jax_md import partition, quantity, space
+from matscipy.neighbours import neighbour_list
 
 from apax.model import ModelBuilder
 from apax.train.checkpoints import restore_parameters
@@ -28,32 +29,41 @@ def maybe_vmap(apply, params, Z):
     return energy_fn
 
 
-def build_energy_neighbor_fns(atoms, config, params, dr_threshold):
+def build_energy_neighbor_fns(atoms, config, params, dr_threshold, neigbor_from_jax):
+    r_max = config.model.r_max
     atomic_numbers = jnp.asarray(atoms.numbers)
-    box = jnp.asarray(atoms.get_cell().array, dtype=jnp.float32)
+    box = jnp.asarray(atoms.cell.array, dtype=jnp.float64)
+    neigbor_from_jax = neighbor_calculable_with_jax(box, r_max)
     box = box.T
-
-    if np.all(box < 1e-6):
-        displacement_fn, _ = space.free()
-    else:
-        displacement_fn, _ = space.periodic_general(box, fractional_coordinates=True)
+    displacement_fn = None
+    neighbor_fn = None
+
+    if neigbor_from_jax:
+        if np.all(box < 1e-6):
+            displacement_fn, _ = space.free()
+        else:
+            displacement_fn, _ = space.periodic_general(box, fractional_coordinates=True)
+
+        neighbor_fn = jax_md_reduced.partition.neighbor_list(
+            displacement_fn,
+            box,
+            config.model.r_max,
+            dr_threshold,
+            fractional_coordinates=True,
+            disable_cell_list=True,
+            format=partition.Sparse,
+        )
 
     Z = jnp.asarray(atomic_numbers)
     n_species = 119  # int(np.max(Z) + 1)
     builder = ModelBuilder(config.model.get_dict(), n_species=n_species)
+
     model = builder.build_energy_derivative_model(
         apply_mask=True, init_box=np.array(box), inference_disp_fn=displacement_fn
     )
+
     energy_fn = maybe_vmap(model.apply, params, Z)
-    neighbor_fn = jax_md_reduced.partition.neighbor_list(
-        displacement_fn,
-        box,
-        config.model.r_max,
-        dr_threshold,
-        fractional_coordinates=True,
-        disable_cell_list=True,
-        format=partition.Sparse,
-    )
+
     return energy_fn, neighbor_fn
 
 
@@ -62,7 +72,7 @@ def process_stress(results, box):
     results = {
         # We should properly check whether CP2K uses the ASE cell convention
         # for tetragonal strain, it doesn't matter whether we transpose or not
-        k: (val.T / V if k.startswith("stress") else val)
+        k: val.T / V if k.startswith("stress") else val
         for k, val in results.items()
     }
     return results
@@ -83,7 +93,6 @@ def ensemble(positions, Z, idx, box, offsets):
 class ASECalculator(Calculator):
     """
     ASE Calculator for apax models.
-    DOES NOT SUPPORT CUTOFFS LARGER THAN MIN(BOX SIZE / 2)!
     """
 
     implemented_properties = [
@@ -96,6 +105,7 @@ def __init__(
         model_dir: Union[Path, list[Path]],
         dr_threshold: float = 0.5,
         transformations: Callable = [],
+        padding_factor: float = 1.5,
         **kwargs
     ):
         Calculator.__init__(self, **kwargs)
@@ -105,6 +115,7 @@ def __init__(
         self.n_models = 1 if isinstance(model_dir, (Path, str)) else len(model_dir)
 
         self.model_config, self.params = restore_parameters(model_dir)
+        self.padding_factor = padding_factor
 
         if self.model_config.model.calc_stress:
             self.implemented_properties.append("stress")
@@ -119,13 +130,18 @@ def __init__(
         self.step = None
         self.neighbor_fn = None
         self.neighbors = None
+        self.offsets = None
 
     def initialize(self, atoms):
+        box = jnp.asarray(atoms.cell.array, dtype=jnp.float64)
+        self.r_max = self.model_config.model.r_max
+        self.neigbor_from_jax = neighbor_calculable_with_jax(box, self.r_max)
         model, neighbor_fn = build_energy_neighbor_fns(
             atoms,
             self.model_config,
             self.params,
             self.dr_threshold,
+            self.neigbor_from_jax,
         )
 
         if self.is_ensemble:
@@ -134,7 +150,99 @@ def initialize(self, atoms):
         for transformation in self.transformations:
             model = transformation.apply(model, self.n_models)
 
-        Z = jnp.asarray(atoms.numbers)
+        self.step = get_step_fn(model, atoms, self.neigbor_from_jax)
+        self.neighbor_fn = neighbor_fn
+
+    def set_neighbours_and_offsets(self, atoms, box):
+        idxs_i, idxs_j, offsets = neighbour_list("ijS", atoms, self.r_max)
+
+        if len(idxs_i) > self.padded_length:
+            print("neighbor list overflowed, reallocating.")
+            self.padded_length = int(len(idxs_i) * self.padding_factor)
+            self.initialize(atoms)
+
+        zeros_to_add = self.padded_length - len(idxs_i)
+
+        self.neighbors = np.array([idxs_i, idxs_j], dtype=np.int32)
+        self.neighbors = np.pad(self.neighbors, ((0, 0), (0, zeros_to_add)), "constant")
+
+        offsets = np.matmul(offsets, box)
+        self.offsets = np.pad(offsets, ((0, zeros_to_add), (0, 0)), "constant")
+
+    def calculate(self, atoms, properties=["energy"], system_changes=all_changes):
+        Calculator.calculate(self, atoms, properties, system_changes)
+        positions = jnp.asarray(atoms.positions, dtype=jnp.float64)
+        box = jnp.asarray(atoms.cell.array, dtype=jnp.float64)
+
+        # setup model and neighbours
+        if self.step is None:
+            self.initialize(atoms)
+
+            if self.neigbor_from_jax:
+                self.neighbors = self.neighbor_fn.allocate(positions)
+            else:
+                idxs_i = neighbour_list("i", atoms, self.r_max)
+                self.padded_length = int(len(idxs_i) * self.padding_factor)
+
+        elif "numbers" in system_changes:
+            self.initialize(atoms)
+
+            if self.neigbor_from_jax:
+                self.neighbors = self.neighbor_fn.allocate(positions)
+
+        elif "cell" in system_changes:
+            neigbor_from_jax = neighbor_calculable_with_jax(box, self.r_max)
+            if self.neigbor_from_jax != neigbor_from_jax:
+                self.initialize(atoms)
+
+        # predict
+        if self.neigbor_from_jax:
+            results, self.neighbors = self.step(positions, self.neighbors, box)
+
+            if self.neighbors.did_buffer_overflow:
+                print("neighbor list overflowed, reallocating.")
+                self.initialize(atoms)
+                self.neighbors = self.neighbor_fn.allocate(positions)
+
+                results, self.neighbors = self.step(positions, self.neighbors, box)
+
+        else:
+            self.set_neighbours_and_offsets(atoms, box)
+            positions = np.array(space.transform(np.linalg.inv(box), atoms.positions))
+
+            results = self.step(positions, self.neighbors, box, self.offsets)
+
+        self.results = {k: np.array(v, dtype=np.float64) for k, v in results.items()}
+        self.results["energy"] = self.results["energy"].item()
+
+
+def neighbor_calculable_with_jax(box, r_max):
+    if np.all(box < 1e-6):
+        return True
+    else:
+        # all lettice vector combinations to calculate all three plane distances
+        a_vec_list = [box[0], box[0], box[1]]
+        b_vec_list = [box[1], box[2], box[2]]
+        c_vec_list = [box[2], box[1], box[0]]
+
+        height = []
+        for i in range(3):
+            normvec = np.cross(a_vec_list[i], b_vec_list[i])
+            projection = (
+                c_vec_list[i]
+                - np.sum(normvec * c_vec_list[i]) / np.sum(normvec**2) * normvec
+            )
+            height.append(np.linalg.norm(c_vec_list[i] - projection))
+
+        if np.min(height) / 2 > r_max:
+            return True
+        else:
+            return False
+
+
+def get_step_fn(model, atoms, neigbor_from_jax):
+    Z = jnp.asarray(atoms.numbers)
+    if neigbor_from_jax:
 
         @jax.jit
         def step_fn(positions, neighbor, box):
@@ -145,33 +253,24 @@ def step_fn(positions, neighbor, box):
                 neighbor = neighbor.update(positions, box=box)
             else:
                 neighbor = neighbor.update(positions)
-            offsets = jnp.full([neighbor.idx.shape[1], 3], 0)
 
+            offsets = jnp.full([neighbor.idx.shape[1], 3], 0)
             results = model(positions, Z, neighbor.idx, box, offsets)
 
             if "stress" in results.keys():
                 results = process_stress(results, box)
 
             return results, neighbor
 
-        self.step = step_fn
-        self.neighbor_fn = neighbor_fn
+    else:
 
-    def calculate(self, atoms, properties=["energy"], system_changes=all_changes):
-        Calculator.calculate(self, atoms, properties, system_changes)
-        positions = jnp.asarray(atoms.positions, dtype=jnp.float64)
-        box = jnp.asarray(atoms.cell.array, dtype=jnp.float64)
-        if self.step is None or "numbers" in system_changes:
-            self.initialize(atoms)
-            self.neighbors = self.neighbor_fn.allocate(positions)
+        @jax.jit
+        def step_fn(positions, neighbor, box, offsets):
+            results = model(positions, Z, neighbor, box, offsets)
 
-        results, self.neighbors = self.step(positions, self.neighbors, box)
+            if "stress" in results.keys():
+                results = process_stress(results, box)
 
-        if self.neighbors.did_buffer_overflow:
-            print("neighbor list overflowed, reallocating.")
-            self.initialize(atoms)
-            self.neighbors = self.neighbor_fn.allocate(positions)
-            results, self.neighbors = self.step(positions, self.neighbors, box)
+            return results
 
-        self.results = {k: np.array(v, dtype=np.float64) for k, v in results.items()}
-        self.results["energy"] = self.results["energy"].item()
+    return step_fn

apax/md/function_transformations.py

@@ -67,7 +67,7 @@ class GaussianAcceleratedMolecularDynamics(FunctionTransformation):
     https://pubs.acs.org/doi/10.1021/acs.jctc.5b00436
 
     Parameters:
-    energy_target: Target potential energy below wich to apply the boost potential.
+    energy_target: Target potential energy below which to apply the boost potential.
     spring_constant: Spring constant of the boost potential.
     """
 

apax/optimizer/get_optimizer.py

@@ -17,7 +17,7 @@ def map_nested_fn(fn: Callable[[str, Any], dict]) -> Callable[[dict], dict]:
 
     def map_fn(nested_dict):
         return {
-            k: (map_fn(v) if isinstance(v, dict) else fn(k, v))
+            k: map_fn(v) if isinstance(v, dict) else fn(k, v)
             for k, v in nested_dict.items()
         }