docs(optimize): document post-hoc convergence check

examples/tutorials/high_level_tutorial.py

@@ -347,6 +347,10 @@
 that determines when the optimization is converged. By default, the `convergence_fn` will
 wait until the energy difference between steps is less than 1 meV.
 
+Note that `optimize` returns the final state even if `max_steps` is reached before
+convergence. To check per-system convergence afterwards, evaluate
+`convergence_fn(final_state)` on the returned state.
+
 Let's use the `optimize` function with the FIRE algorithm to relax our structures:
 """
 
@@ -359,6 +363,8 @@
 )
 
 final_atoms = final_state.to_atoms()
+# `optimize` may return `final_atoms` before convergence if `max_steps` is reached.
+# Evaluate `convergence_fn` on `final_state` to verify per-system convergence
 
 
 # %% [markdown]
@@ -409,6 +415,10 @@ def default_energy_convergence(state, last_energy):
 
 final_atoms = final_state.to_atoms()
 
+# Check per-system convergence on the returned state
+convergence_tensor = force_convergence_fn(final_state)
+print(f"Converged systems: {convergence_tensor.tolist()}")
+
 
 # %% [markdown]
 """

torch_sim/runners.py

@@ -593,6 +593,11 @@ def optimize[T: OptimState](  # noqa: C901, PLR0915
 
     Returns:
         T: Optimized system state
+
+    Notes:
+        The state is returned even if max_steps was reached before
+        `convergence_fn` was satisfied. To check per-system convergence,
+        evaluate `convergence_fn(final_state)` on the returned state
     """
     # create a default convergence function if one is not provided
     # TODO: document this behavior