From e83bdf30125fbf659d31eb8c3aa18ce8cab10f8f Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Tue, 18 Jun 2024 05:29:14 +0200 Subject: [PATCH 1/5] Store additional properties in ASE calculator --- .../torch/atomistic/ase_calculator.py | 28 +++++++++++++------ 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py index d7fd18677..fc6b7bb99 100644 --- a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py +++ b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py @@ -62,6 +62,7 @@ def __init__( extensions_directory=None, check_consistency=False, device=None, + properties_to_store: Optional[List[str]] = None, ): """ :param model: model to use for the calculation. This can be a file path, a @@ -73,6 +74,12 @@ def __init__( running, defaults to False. :param device: torch device to use for the calculation. If ``None``, we will try the options in the model's ``supported_device`` in order. + :param properties_to_store: list of model outputs to store as results of the ASE + calculator at every step. This is useful when you want to store properties + that are not used in the propagation of the dynamics and/or are not standard + ASE properties ('energy', 'forces', 'stress', 'stresses', 'dipole', + 'charges', 'magmom', 'magmoms', 'free_energy', 'energies'). These properties + will be available as ``atoms.calc.results['']``. """ super().__init__() @@ -147,6 +154,9 @@ def __init__( # We do our own check to verify if a property is implemented in `calculate()`, # so we pretend to be able to compute all properties ASE knows about. self.implemented_properties = ALL_ASE_PROPERTIES + self.properties_to_store = ( + properties_to_store if properties_to_store is not None else [] + ) def todict(self): if "model_path" not in self.parameters: @@ -242,8 +252,10 @@ def calculate( system_changes=system_changes, ) + properties_to_calculate = properties + self.properties_to_store + with record_function("ASECalculator::prepare_inputs"): - outputs = _ase_properties_to_metatensor_outputs(properties) + outputs = _ase_properties_to_metatensor_outputs(properties_to_calculate) capabilities = self._model.capabilities() for name in outputs.keys(): if name not in capabilities.outputs: @@ -257,11 +269,11 @@ def calculate( ) do_backward = False - if "forces" in properties: + if "forces" in properties_to_calculate: do_backward = True positions.requires_grad_(True) - if "stress" in properties: + if "stress" in properties_to_calculate: do_backward = True scaling = torch.eye(3, requires_grad=True, dtype=self._dtype) @@ -271,7 +283,7 @@ def calculate( cell = cell @ scaling - if "stresses" in properties: + if "stresses" in properties_to_calculate: raise NotImplementedError("'stresses' are not implemented yet") run_options = ModelEvaluationOptions( @@ -322,14 +334,14 @@ def calculate( self.results = {} - if "energies" in properties: + if "energies" in properties_to_calculate: energies_values = energies.detach().reshape(-1) energies_values = energies_values.to(device="cpu").to( dtype=torch.float64 ) self.results["energies"] = energies_values.numpy() - if "energy" in properties: + if "energy" in properties_to_calculate: energy_values = energy.detach() energy_values = energy_values.to(device="cpu").to(dtype=torch.float64) self.results["energy"] = energy_values.numpy()[0, 0] @@ -339,12 +351,12 @@ def calculate( energy.backward(-torch.ones_like(energy)) with record_function("ASECalculator::convert_outputs"): - if "forces" in properties: + if "forces" in properties_to_calculate: forces_values = system.positions.grad.reshape(-1, 3) forces_values = forces_values.to(device="cpu").to(dtype=torch.float64) self.results["forces"] = forces_values.numpy() - if "stress" in properties: + if "stress" in properties_to_calculate: stress_values = -scaling.grad.reshape(3, 3) / atoms.cell.volume stress_values = stress_values.to(device="cpu").to(dtype=torch.float64) self.results["stress"] = _full_3x3_to_voigt_6_stress( From 4e9c5b07b1e6628934f6fcff38b98f01f476f08c Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Tue, 18 Jun 2024 05:44:17 +0200 Subject: [PATCH 2/5] Allow to use vesin in ASE calculator --- .../torch/atomistic/ase_calculator.py | 25 +++++++++++++++---- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py index fc6b7bb99..390b573d9 100644 --- a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py +++ b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py @@ -30,6 +30,13 @@ ) +try: + import vesin + HAS_VESIN = True +except ImportError: + HAS_VESIN = False + + if os.environ.get("METATENSOR_IMPORT_FOR_SPHINX", "0") == "0": # this can not be imported when building the documentation from .. import sum_over_samples # isort: skip @@ -435,11 +442,19 @@ def _ase_properties_to_metatensor_outputs(properties): def _compute_ase_neighbors(atoms, options, dtype, device): - nl_i, nl_j, nl_S, nl_D = ase.neighborlist.neighbor_list( - "ijSD", - atoms, - cutoff=options.engine_cutoff(engine_length_unit="angstrom"), - ) + + if (np.all(atoms.pbc) or np.all(~atoms.pbc)) and HAS_VESIN: + nl_i, nl_j, nl_S, nl_D = vesin.ase_neighbor_list( + "ijSD", + atoms, + cutoff=options.engine_cutoff(engine_length_unit="angstrom"), + ) + else: + nl_i, nl_j, nl_S, nl_D = ase.neighborlist.neighbor_list( + "ijSD", + atoms, + cutoff=options.engine_cutoff(engine_length_unit="angstrom"), + ) selected = [] for pair_i, (i, j, S) in enumerate(zip(nl_i, nl_j, nl_S)): From abec5d1d4f40faea66b1e3a32b6a7491471aa623 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Fri, 28 Jun 2024 16:32:57 +0200 Subject: [PATCH 3/5] Add storing mechanism for requests every n steps --- .../torch/atomistic/ase_calculator.py | 48 +++++++++++++++---- 1 file changed, 40 insertions(+), 8 deletions(-) diff --git a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py index 77422200a..a95a904be 100644 --- a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py +++ b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py @@ -37,13 +37,6 @@ HAS_VESIN = False -try: - import vesin - HAS_VESIN = True -except ImportError: - HAS_VESIN = False - - if os.environ.get("METATENSOR_IMPORT_FOR_SPHINX", "0") == "0": # this can not be imported when building the documentation from .. import sum_over_samples # isort: skip @@ -175,6 +168,7 @@ def __init__( self.properties_to_store = ( properties_to_store if properties_to_store is not None else [] ) + self.additional_properties_to_store = [] def todict(self): if "model_path" not in self.parameters: @@ -270,7 +264,9 @@ def calculate( system_changes=system_changes, ) - properties_to_calculate = properties + self.properties_to_store + properties_to_calculate = ( + properties + self.properties_to_store + self.additional_properties_to_store + ) with record_function("ASECalculator::prepare_inputs"): outputs = _ase_properties_to_metatensor_outputs(properties_to_calculate) @@ -383,6 +379,42 @@ def calculate( stress_values.numpy() ) + def request_properties_every_n_steps( + self, + dyn: ase.md.MolecularDynamics, + properties: List[str], + n: int, + ): + """ + Makes a property available every n steps of the dynamics. + + :param dyn: ASE molecular dynamics object + :param properties: list of properties to be made available + at regular intervals + :param n: number of steps between each property calculation + """ + + # prepare for step 0, where the properties must be available + self.additional_properties_to_store.extend(properties) + + def _request_properties(): + self.additional_properties_to_store.extend(properties) + + def _unrequest_properties(): + for prop in properties: + self.additional_properties_to_store.remove(prop) + + def _manage_additional_properties(): + step_count = dyn.nsteps + if step_count % n == n - 1: + _request_properties() + elif step_count % n == 0: + _unrequest_properties() + else: + pass + + dyn.attach(_manage_additional_properties, interval=1, mode="step") + def _find_best_device(devices: List[str]) -> torch.device: """ From ccde6eb2639c5d52c431e577cfe54cbcb5087154 Mon Sep 17 00:00:00 2001 From: Filippo Bigi <98903385+frostedoyster@users.noreply.github.com> Date: Sat, 29 Jun 2024 07:27:20 +0200 Subject: [PATCH 4/5] Update ase_calculator.py --- .../metatensor/torch/atomistic/ase_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py index a95a904be..58fe23c18 100644 --- a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py +++ b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py @@ -381,7 +381,7 @@ def calculate( def request_properties_every_n_steps( self, - dyn: ase.md.MolecularDynamics, + dyn: ase.md.md.MolecularDynamics, properties: List[str], n: int, ): From de2b751304549fef19f7a7ac695b61dd4aa45727 Mon Sep 17 00:00:00 2001 From: Filippo Bigi <98903385+frostedoyster@users.noreply.github.com> Date: Sat, 29 Jun 2024 07:39:08 +0200 Subject: [PATCH 5/5] Update ase_calculator.py --- .../metatensor/torch/atomistic/ase_calculator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py index 58fe23c18..c5f178832 100644 --- a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py +++ b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py @@ -21,6 +21,7 @@ import ase # isort: skip +import ase.md # isort: skip import ase.neighborlist # isort: skip import ase.calculators.calculator # isort: skip from ase.calculators.calculator import ( # isort: skip