Default to computing energy gradients everytime the energy is requested

Author: Luthaf

python/examples/4-profiling.py

@@ -134,33 +134,37 @@ def forward(
 #     </details>
 #
 
-
 # %%
 #
-# If you are trying to profile your own model, you can start here and create a
+# If you are trying to profile your own model, you can start here and create
 # ``MetatomicCalculator`` with your own model.
 
+# %%
+#
+# Profiling energy calculation
+# ----------------------------
+#
+# We will start with an energy-only calculator, which can be enabled with
+# ``do_gradients_with_energy=False``.
 
-atoms.calc = MetatomicCalculator("exported-model.pt")
+atoms.calc = MetatomicCalculator("exported-model.pt", do_gradients_with_energy=False)
 
 # %%
 #
 # Before trying to profile the code, it is a good idea to run it a couple of times to
 # allow torch to warmup internally.
 
-atoms.get_forces()
-atoms.get_potential_energy()
+for _ in range(10):
+    # force the model to re-run everytime, otherwise ASE caches calculation results
+    atoms.rattle(1e-6)
+    atoms.get_potential_energy()
 
 # %%
 #
-# Profiling energy calculation
-# ----------------------------
-#
 # Now we can run code using :py:func:`torch.profiler.profile` to collect statistic on
-# how long each function takes to run. We randomize the positions to force ASE to
-# recompute the energy of the system
+# how long each function takes to run.
 
-atoms.positions += np.random.rand(*atoms.positions.shape)
+atoms.rattle(1e-6)
 with torch.profiler.profile() as energy_profiler:
     atoms.get_potential_energy()
 
@@ -202,7 +206,14 @@ def forward(
 # Let's now do the same, but computing the forces for this system. This mean we should
 # now see some time spent in the ``backward()`` function, on top of everything else.
 
-atoms.positions += np.random.rand(*atoms.positions.shape)
+atoms.calc = MetatomicCalculator("exported-model.pt")
+
+# warmup
+for _ in range(10):
+    atoms.rattle(1e-6)
+    atoms.get_forces()
+
+atoms.rattle(1e-6)
 with torch.profiler.profile() as forces_profiler:
     atoms.get_forces()
 

python/metatomic_torch/metatomic/torch/ase_calculator.py

@@ -68,6 +68,7 @@ def __init__(
         check_consistency=False,
         device=None,
         non_conservative=False,
+        do_gradients_with_energy=True,
     ):
         """
         :param model: model to use for the calculation. This can be a file path, a
@@ -87,10 +88,18 @@ 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 non_conservative: if ``True``, the model will be asked to
-            compute non-conservative forces and stresses. This can afford a speed-up,
+        :param non_conservative: if ``True``, the model will be asked to compute
+            non-conservative forces and stresses. This can afford a speed-up,
             potentially at the expense of physical correctness (especially in molecular
             dynamics simulations).
+        :param do_gradients_with_energy: if ``True``, this calculator will always
+            compute the energy gradients (forces and stress) when the energy is
+            requested (e.g. through ``atoms.get_potential_energy()``). Because the
+            results of a calculation are cached by ASE, this means future calls to
+            ``atom.get_forces()`` will return immediately, without needing to execute
+            the model again. If you are mainly interested in the energy, you can set
+            this to ``False`` and enjoy a faster model. Forces will still be calculated
+            if requested with ``atoms.get_forces()``.
         """
         super().__init__()
 
@@ -175,6 +184,7 @@ def __init__(
         self._device = device
         self._model = model.to(device=self._device)
         self._non_conservative = non_conservative
+        self._do_gradients_with_energy = do_gradients_with_energy
 
         # 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.
@@ -327,6 +337,11 @@ def calculate(
         if "stresses" in properties:
             raise NotImplementedError("'stresses' are not implemented yet")
 
+        if self._do_gradients_with_energy:
+            if calculate_energies or calculate_energy:
+                calculate_forces = True
+                calculate_stress = True
+
         with record_function("MetatomicCalculator::prepare_inputs"):
             outputs = self._ase_properties_to_metatensor_outputs(
                 properties,
@@ -409,6 +424,19 @@ def calculate(
             assert len(energy.block().gradients_list()) == 0
             assert energy.block().values.shape == (1, 1)
 
+        if do_backward:
+            if energy.block().values.grad_fn is None:
+                # did the user actually request a gradient, or are we trying to
+                # compute one just for efficiency?
+                if "forces" in properties or "stress" in properties:
+                    # the user asked for it, let it fail below
+                    pass
+                else:
+                    # we added the calculation, let's remove it
+                    do_backward = False
+                    calculate_forces = False
+                    calculate_stress = False
+
         with record_function("MetatomicCalculator::run_backward"):
             if do_backward:
                 energy.block().values.backward()