From cc51830d37692cc5d859321bd94ea8b92b485325 Mon Sep 17 00:00:00 2001
From: Hatem Helal <hatemh@graphcore.ai>
Date: Tue, 23 Apr 2024 16:12:33 +0000
Subject: [PATCH] add minimise doc

---
 mess/hamiltonian.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/mess/hamiltonian.py b/mess/hamiltonian.py
index f7affc3..b684137 100644
--- a/mess/hamiltonian.py
+++ b/mess/hamiltonian.py
@@ -1,4 +1,5 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+from typing import Tuple
 import equinox as eqx
 import jax
 import jax.numpy as jnp
@@ -191,7 +192,19 @@ def orthonormalise(self, Z: FloatNxN) -> FloatNxN:
 
 
 @jax.jit
-def minimise(H: Hamiltonian):
+def minimise(H: Hamiltonian) -> Tuple[ScalarLike, FloatNxN, optx.Solution]:
+    """Solve for the electronic coefficients that minimise the total energy
+
+    Args:
+        H (Hamiltonian): the Hamiltonian for the given basis set and molecular structure
+
+    Returns:
+        Tuple[ScalarLike, FloatNxN, optimistix.Solution]: Tuple with elements:
+            - total energy in atomic units
+            - coefficient matrix C that minimises the Hamiltonian
+            - the optimistix.Solution object
+    """
+
     def f(Z, _):
         C = H.orthonormalise(Z)
         P = H.basis.density_matrix(C)