diff --git a/mlcalcdriver/base/posinp.py b/mlcalcdriver/base/posinp.py index e47c18f..8e23798 100644 --- a/mlcalcdriver/base/posinp.py +++ b/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( diff --git a/mlcalcdriver/calculators/schnetpack_patch.py b/mlcalcdriver/calculators/schnetpack_patch.py index 319202b..567db8a 100644 --- a/mlcalcdriver/calculators/schnetpack_patch.py +++ b/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 diff --git a/mlcalcdriver/interfaces/schnetpack_interface.py b/mlcalcdriver/interfaces/schnetpack_interface.py index df36449..4c31aa7 100644 --- a/mlcalcdriver/interfaces/schnetpack_interface.py +++ b/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) diff --git a/mlcalcdriver/workflows/phonon.py b/mlcalcdriver/workflows/phonon.py index 490d461..8b4714b 100644 --- a/mlcalcdriver/workflows/phonon.py +++ b/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