improvements for lanczos

mlcalcdriver/base/posinp.py

@@ -634,7 +634,7 @@ def translate(self, vector):
         """
         new_positions = self.positions + np.array(vector)
         atoms = [
-            Atom(atom.type, pos.tolist())
+            Atom(atom.type, pos, isotope=atom.isotope)
             for atom, pos in zip(self.atoms, new_positions)
         ]
         return Posinp(

mlcalcdriver/calculators/schnetpack_patch.py

@@ -48,6 +48,7 @@ def __init__(
         subgrid=None,
         sparse=False,
         atomic_environments=None,
+        patches=None,
     ):
         super().__init__(
             model_dir=model_dir,
@@ -59,6 +60,7 @@ def __init__(
         self.n_interaction = len(self.model.representation.interactions)
         self.subgrid = subgrid
         self.atomic_environments = atomic_environments
+        self.patches = patches
         self.sparse = sparse
         self._convert_model()
 
@@ -94,6 +96,19 @@ def atomic_environments(self, atomic_environments):
         else:
             self._atomic_environments = [None] * np.prod(self.subgrid)
 
+    @property
+    def patches(self):
+        return self._patches
+
+    @patches.setter
+    def patches(self, patches):
+        if patches is not None:
+            for el in patches:
+                assert len(el) == np.prod(self.subgrid)
+            self._patches = patches
+        else:
+            self._patches = None
+
     def run(
         self,
         property,
@@ -146,15 +161,23 @@ def run(
         )
 
         # Split the configuration according to the subgrid
-        at_to_patches = AtomsToPatches(
-            cutoff=self.cutoff, n_interaction=self.n_interaction, grid=self.subgrid
-        )
-        (
-            subcells,
-            subcells_main_idx,
-            original_cell_idx,
-            complete_subcell_copy_idx,
-        ) = at_to_patches.split_atoms(atoms)
+        if self.patches is None:
+            at_to_patches = AtomsToPatches(
+                cutoff=self.cutoff, n_interaction=self.n_interaction, grid=self.subgrid
+            )
+            (
+                subcells,
+                subcells_main_idx,
+                original_cell_idx,
+                complete_subcell_copy_idx,
+            ) = at_to_patches.split_atoms(atoms)
+        else:
+            (
+                subcells,
+                subcells_main_idx,
+                original_cell_idx,
+                complete_subcell_copy_idx,
+            ) = self.patches
 
         # Pass each subcell independantly
         results = []
@@ -220,7 +243,7 @@ def run(
         elif property == "forces":
             forces = np.zeros((len(atoms), 3), dtype=np.float32)
             for i in range(len(results)):
-                forces[original_cell_idx[i]] = results[i]["forces"][0][
+                forces[original_cell_idx[i]] = results[i]["forces"].squeeze()[
                     subcells_main_idx[i]
                 ]
             predictions["forces"] = forces

mlcalcdriver/interfaces/schnetpack_interface.py

@@ -63,11 +63,11 @@ def get_properties(self, idx):
 
 def fast_convert(atoms, atomic_environment):
     inputs = {}
-    inputs[Properties.Z] = atoms.numbers.astype(np.int)
+    inputs[Properties.Z] = atoms.numbers.astype(int)
     inputs[Properties.R] = atoms.positions.astype(np.float32)
 
     nbh_idx, offsets = atomic_environment
-    inputs[Properties.neighbors] = nbh_idx.astype(np.int)
+    inputs[Properties.neighbors] = nbh_idx.astype(int)
 
     inputs[Properties.cell] = np.array(atoms.cell.array, dtype=np.float32)
     inputs[Properties.cell_offset] = offsets.astype(np.float32)

mlcalcdriver/workflows/phonon.py

@@ -416,7 +416,7 @@ def run(self, use_jax=False, sparse_kwargs={}):
             self._solve_sparse_hessian(kwargs=sparse_kwargs)
 
     def _solve_sparse_hessian(self, kwargs):
-        self.dyn_mat = self._compute_sparse_dyn_mat()
+        self._compute_sparse_dyn_mat()
         self.energies, self.normal_modes = self._solve_sparse_dyn_mat(kwargs=kwargs)
         self.energies *= HA_TO_CMM1