From 9089f85689217d7d71fe3b5f136de2791110ba63 Mon Sep 17 00:00:00 2001
From: Ivo Filot <ivo@ivofilot.nl>
Date: Tue, 5 Dec 2023 19:13:49 +0100
Subject: [PATCH] Fixing #27: Adding final update to energy vector

---
 pyqint/hf.py                       |  5 +++--
 tests/test_energy_decomposition.py | 24 ++++++++++++++++++++++++
 tests/test_hf.py                   |  3 ---
 3 files changed, 27 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_energy_decomposition.py

diff --git a/pyqint/hf.py b/pyqint/hf.py
index f9a3249..e766e14 100644
--- a/pyqint/hf.py
+++ b/pyqint/hf.py
@@ -134,7 +134,7 @@ def rhf(self, mol, basis, calc_forces=False, itermax=100,
 
             # calculate DIIS coefficients
             e = (F.dot(P.dot(S)) - S.dot(P.dot(F))).flatten()   # calculate error vector
-            #enorm = np.linalg.norm(e)                           # store error vector norm
+            #enorm = np.linalg.norm(e)                          # store error vector norm
             fmats_diis.append(F)                                # add Fock matrix to list
             pmat_diis.append(P)                                 # add density matrix to list
             evs_diis.append(e)
@@ -175,8 +175,9 @@ def rhf(self, mol, basis, calc_forces=False, itermax=100,
         end = time.time()
         time_stats['self_convergence'] = end - start
 
-        # update density matrix
+        # update density matrix and final energy
         P = np.einsum('ik,jk,k->ij', C, C, occ)
+        energies[-1] = 0.5 * np.einsum('ji,ij', P, T+V+F) + nuc_rep
 
         # build solution dictionary
         sol = {
diff --git a/tests/test_energy_decomposition.py b/tests/test_energy_decomposition.py
new file mode 100644
index 0000000..2fc4ec4
--- /dev/null
+++ b/tests/test_energy_decomposition.py
@@ -0,0 +1,24 @@
+import unittest
+from pyqint import MoleculeBuilder, HF
+import numpy as np
+
+class TestEnergyDecomposition(unittest.TestCase):
+
+    def test_hartree_fock_h2o(self):
+        """
+        Test Hartree-Fock calculation on water using STO-3G basis set
+        """
+        mol = MoleculeBuilder().from_name("h2o")
+
+        res = HF().rhf(mol, 'sto3g')
+        P = res['density']
+        T = res['kinetic']
+        V = res['nuclear']
+        H = res['fock']
+        enucrep = res['enucrep']
+        energy = res['energy']
+        
+        np.testing.assert_almost_equal(0.5 * np.einsum('ji,ij', P, T+V+H) + enucrep, energy, decimal=16)
+        
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file
diff --git a/tests/test_hf.py b/tests/test_hf.py
index 90ed714..0c39ec4 100644
--- a/tests/test_hf.py
+++ b/tests/test_hf.py
@@ -1,9 +1,6 @@
 import unittest
 from pyqint import PyQInt, cgf, gto, Molecule, HF
-from copy import deepcopy
 import numpy as np
-import multiprocessing
-import os
 
 class TestHF(unittest.TestCase):