Implement new standard outputs for atomistic models

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``.

metatensor-torch/src/atomistic/model.cpp

@@ -136,7 +136,9 @@ ModelOutput ModelOutputHolder::from_json(std::string_view json) {
 /******************************************************************************/
 
 std::unordered_set<std::string> KNOWN_OUTPUTS = {
- "energy"
+ "energy",
+ "dipole",
+ "energy_ensemble"
 };
 
 void ModelCapabilitiesHolder::set_outputs(torch::Dict<std::string, ModelOutput> outputs) {
@@ -1032,6 +1034,29 @@ static std::unordered_map<std::string, Quantity> 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) {

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");
  }
  };
 

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['<stored_property>']``.
  """
  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:
  """