From 78033636a2f2e3f53213b652cb55cf6c8aa9c4ef Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Sun, 16 Jun 2024 09:09:22 +0200 Subject: [PATCH] Add new properties --- docs/src/atomistic/outputs.rst | 84 +++++++++++++++++++ metatensor-torch/src/atomistic/model.cpp | 27 +++++- metatensor-torch/tests/atomistic.cpp | 2 +- .../torch/atomistic/ase_calculator.py | 42 ++++++++-- 4 files changed, 145 insertions(+), 10 deletions(-) diff --git a/docs/src/atomistic/outputs.rst b/docs/src/atomistic/outputs.rst index c4f8998ef..fd18f38b5 100644 --- a/docs/src/atomistic/outputs.rst +++ b/docs/src/atomistic/outputs.rst @@ -119,3 +119,87 @@ The following gradients can be defined and requested with - ``["xyz_1", "xyz_2"]`` - Both ``"xyz_1"`` and ``"xyz_2"`` have values ``[0, 1, 2]``, and correspond to the two axes of the 3x3 strain matrix :math:`\epsilon`. + + +Energy ensemble +^^^^^^^^^^^^^^^ + +An ensemble of energies is associated with the ``"energy_ensemble"`` key in the +model outputs, and must have the following metadata: + +.. list-table:: Metadata for energy ensemble output + :widths: 2 3 7 + :header-rows: 1 + + * - Metadata + - Names + - Description + + * - keys + - ``"_"`` + - the energy ensemble keys must have a single dimension named ``"_"``, with a + single entry set to ``0``. The energy ensemble is always a + :py:class:`metatensor.torch.TensorMap` with a single block. + + * - samples + - ``["system", "atom"]`` or ``["system"]`` + - if doing ``per_atom`` output, the sample names must be ``["system", + "atom"]``, otherwise the sample names must be ``["system"]``. + + ``"system"`` must range from 0 to the number of systems given as input to + the model. ``"atom"`` must range between 0 and the number of + atoms/particles in the corresponding system. If ``selected_atoms`` is + provided, then only the selected atoms for each system should be part of + the samples. + + * - components + - + - the ensemble of energies must not have any components + + * - properties + - ``"ensemble_member"`` + - the energy ensemble must have a single property dimension named + ``"ensemble_member"``, with entries ranging from 0 to the number of + members of the ensemble. + + +Dipole +^^^^^^ + +Electric dipole moments are represented by the ``"dipole"`` key in the +model outputs, and must have the following metadata: + +.. list-table:: Metadata for dipole output + :widths: 2 3 7 + :header-rows: 1 + + * - Metadata + - Names + - Description + + * - keys + - ``"_"`` + - the dipole keys must have a single dimension named ``"_"``, with a + single entry set to ``0``. The dipole is always a + :py:class:`metatensor.torch.TensorMap` with a single block. + + * - samples + - ``["system", "atom"]`` or ``["system"]`` + - if doing ``per_atom`` output, the sample names must be ``["system", + "atom"]``, otherwise the sample names must be ``["system"]``. + + ``"system"`` must range from 0 to the number of systems given as input to + the model. ``"atom"`` must range between 0 and the number of + atoms/particles in the corresponding system. If ``selected_atoms`` is + provided, then only the selected atoms for each system should be part of + the samples. + + * - components + - ``["xyz"]`` + - the dipole must have a single component named ``"xyz"`` with values 0, 1, + 2; indicating the components of the dipole moment along x, y, and z. + + * - properties + - ``"dipole"`` + - the energy must have a single property dimension named ``"dipole"``, with + a single entry set to ``0``. diff --git a/metatensor-torch/src/atomistic/model.cpp b/metatensor-torch/src/atomistic/model.cpp index dea720f7f..66392edce 100644 --- a/metatensor-torch/src/atomistic/model.cpp +++ b/metatensor-torch/src/atomistic/model.cpp @@ -136,7 +136,9 @@ ModelOutput ModelOutputHolder::from_json(std::string_view json) { /******************************************************************************/ std::unordered_set KNOWN_OUTPUTS = { - "energy" + "energy", + "dipole", + "energy_ensemble" }; void ModelCapabilitiesHolder::set_outputs(torch::Dict outputs) { @@ -1032,6 +1034,29 @@ static std::unordered_map KNOWN_QUANTITIES = { {"J", "Joule"}, {"Ry", "Rydberg"}, }}}, + {"energy_ensemble", Quantity{/* name */ "energy", /* baseline */ "eV", { + {"eV", 1.0}, + {"meV", 1000.0}, + {"Hartree", 0.03674932247495664}, + {"kcal/mol", 23.060548012069496}, + {"kJ/mol", 96.48533288249877}, + {"Joule", 1.60218e-19}, + {"Rydberg", 0.07349864435130857}, + }, { + // alternative names + {"J", "Joule"}, + {"Ry", "Rydberg"}, + }}}, + {"dipole", Quantity{/* name */ "dipole", /* baseline */ "D", { + {"Debye", 1.0}, + {"Coulomb-meter", 1000.0}, + {"atomic units", 0.03674932247495664}, + }, { + // alternative names + {"D", "Debye"}, + {"C-m", "Coulomb-meter"}, + {"a.u.", "atomic units"}, + }}}, }; bool metatensor_torch::valid_quantity(const std::string& quantity) { diff --git a/metatensor-torch/tests/atomistic.cpp b/metatensor-torch/tests/atomistic.cpp index 534091f7e..be44f1d5d 100644 --- a/metatensor-torch/tests/atomistic.cpp +++ b/metatensor-torch/tests/atomistic.cpp @@ -93,7 +93,7 @@ TEST_CASE("Models metadata") { struct WarningHandler: public torch::WarningHandler { virtual ~WarningHandler() override = default; void process(const torch::Warning& warning) override { - CHECK(warning.msg() == "unknown quantity 'unknown', only [energy length] are supported"); + CHECK(warning.msg() == "unknown quantity 'unknown', only [dipole energy_ensemble energy length] are supported"); } }; diff --git a/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py b/python/metatensor-torch/metatensor/torch/atomistic/ase_calculator.py index d7fd18677..337e8ea7c 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: @@ -243,7 +253,9 @@ def calculate( ) with record_function("ASECalculator::prepare_inputs"): - outputs = _ase_properties_to_metatensor_outputs(properties) + outputs = _ase_properties_to_metatensor_outputs( + properties + self.properties_to_store + ) 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 + self.properties_to_store: do_backward = True positions.requires_grad_(True) - if "stress" in properties: + if "stress" in properties + self.properties_to_store: 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 + self.properties_to_store: 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 + self.properties_to_store: 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 + self.properties_to_store: energy_values = energy.detach() energy_values = energy_values.to(device="cpu").to(dtype=torch.float64) self.results["energy"] = energy_values.numpy()[0, 0] @@ -339,18 +351,32 @@ def calculate( energy.backward(-torch.ones_like(energy)) with record_function("ASECalculator::convert_outputs"): - if "forces" in properties: + if "forces" in properties + self.properties_to_store: 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 + self.properties_to_store: 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( stress_values.numpy() ) + if "dipole" in properties + self.properties_to_store: + dipole_values = outputs["dipole"].block().values.detach().reshape(3) + dipole_values = dipole_values.to(device="cpu").to(dtype=torch.float64) + self.results["dipole"] = dipole_values.numpy() + + if "energy_ensemble" in properties + self.properties_to_store: + energy_ensemble_values = ( + outputs["energy_ensemble"].block().values.detach().flatten() + ) + energy_ensemble_values = energy_ensemble_values.to(device="cpu").to( + dtype=torch.float64 + ) + self.results["energy_ensemble"] = energy_ensemble_values.numpy() + def _find_best_device(devices: List[str]) -> torch.device: """