From b059018265d25021513b42bebe832392a1644f36 Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Thu, 16 Dec 2021 13:06:41 -0800
Subject: [PATCH 1/8] Added extrapolation functions

---
 tangelo/toolboxes/post_processing/__init__.py |  1 +
 .../post_processing/extrapolation.py          | 65 +++++++++++++++++++
 .../tests/test_extrapolation.py               | 30 +++++++++
 3 files changed, 96 insertions(+)
 create mode 100644 tangelo/toolboxes/post_processing/extrapolation.py
 create mode 100644 tangelo/toolboxes/post_processing/tests/test_extrapolation.py

diff --git a/tangelo/toolboxes/post_processing/__init__.py b/tangelo/toolboxes/post_processing/__init__.py
index b73851e5e..dab23877f 100644
--- a/tangelo/toolboxes/post_processing/__init__.py
+++ b/tangelo/toolboxes/post_processing/__init__.py
@@ -13,3 +13,4 @@
 # limitations under the License.
 
 from .mc_weeny_rdm_purification import mcweeny_purify_2rdm
+from .extrapolation import diis, richardson
diff --git a/tangelo/toolboxes/post_processing/extrapolation.py b/tangelo/toolboxes/post_processing/extrapolation.py
new file mode 100644
index 000000000..c8275733d
--- /dev/null
+++ b/tangelo/toolboxes/post_processing/extrapolation.py
@@ -0,0 +1,65 @@
+# Copyright 2021 Good Chemistry Company.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+from scipy.optimize import minimize, root
+
+
+def diis(energies, coeffs):
+    """
+    DIIS extrapolation
+
+    Args:
+        energies (array-like): Energy expectation values for amplified noise rates
+        coeffs (array-like): Noise rate amplification factors
+
+    Returns:
+        float: Extrapolated energy
+    """
+    n = len(coeffs)
+    Eh = np.array(energies)
+    ck = np.array(coeffs)
+    B = np.ones((n+1, n+1))
+    B[n, n] = 0
+    B[:n, :n] = ck[:, None] @ ck[None, :]
+    b = np.zeros(n+1)
+    b[n] = 1
+    x = np.linalg.lstsq(B, b, rcond=None)[0]
+    x = x[:-1]
+    return np.dot(Eh, x)
+
+
+def richardson(energies, coeffs):
+    """
+    Richardson extrapolation as found in
+    Nature 567, 491-495 (2019) [arXiv:1805.04492]
+
+    Args:
+        energies (array-like): Energy expectation values for amplified noise rates
+        coeffs (array-like): Noise rate amplification factors
+
+    Returns:
+        float: Extrapolated energy
+    """
+    n = len(coeffs)
+    Eh = np.array(energies)
+    ck = np.array([coeffs**k for k in range(1, n+1)]).sum(axis=0)
+    B = np.ones((n+1, n+1))
+    B[n, n] = 0
+    B[:n, :n] = ck[:, None] @ ck[None, :]
+    b = np.zeros(n+1)
+    b[n] = 1
+    x = np.linalg.lstsq(B, b, rcond=None)[0]
+    x = x[:-1]
+    return np.dot(Eh, x)
diff --git a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
new file mode 100644
index 000000000..35d3beb82
--- /dev/null
+++ b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
@@ -0,0 +1,30 @@
+# Copyright 2021 Good Chemistry Company.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from tangelo.toolboxes.post_processing import diis, richardson
+
+
+class ExtrapolationTest(unittest.TestCase):
+
+    def test_diis(self):
+        self.assertAlmostEqual(reference, calculated, delta=1e-10)
+
+    def test_richardson(self):
+        self.assertAlmostEqual(reference, calculated, delta=1e-10)
+
+
+
+if __name__ == "__main__":
+    unittest.main()

From a645a3d476ecb15556ecc5454d5c4cc1e29a1da9 Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Thu, 16 Dec 2021 16:08:35 -0800
Subject: [PATCH 2/8] Added unit tests for extrapolation

---
 .../post_processing/tests/data/diis_test.dat  |  5 +++++
 .../tests/data/richardson_test.dat            |  5 +++++
 .../tests/test_extrapolation.py               | 20 ++++++++++++++++++-
 3 files changed, 29 insertions(+), 1 deletion(-)
 create mode 100644 tangelo/toolboxes/post_processing/tests/data/diis_test.dat
 create mode 100644 tangelo/toolboxes/post_processing/tests/data/richardson_test.dat

diff --git a/tangelo/toolboxes/post_processing/tests/data/diis_test.dat b/tangelo/toolboxes/post_processing/tests/data/diis_test.dat
new file mode 100644
index 000000000..0945ceb2e
--- /dev/null
+++ b/tangelo/toolboxes/post_processing/tests/data/diis_test.dat
@@ -0,0 +1,5 @@
+1.000000000000000000e+00 -1.047755735226631124e+00
+1.100000000000000089e+00 -1.043022885253066523e+00
+1.199999999999999956e+00 -1.033645681252726822e+00
+1.300000000000000044e+00 -1.030052450247785245e+00
+nan -1.065698319746184453e+00
diff --git a/tangelo/toolboxes/post_processing/tests/data/richardson_test.dat b/tangelo/toolboxes/post_processing/tests/data/richardson_test.dat
new file mode 100644
index 000000000..0945ceb2e
--- /dev/null
+++ b/tangelo/toolboxes/post_processing/tests/data/richardson_test.dat
@@ -0,0 +1,5 @@
+1.000000000000000000e+00 -1.047755735226631124e+00
+1.100000000000000089e+00 -1.043022885253066523e+00
+1.199999999999999956e+00 -1.033645681252726822e+00
+1.300000000000000044e+00 -1.030052450247785245e+00
+nan -1.065698319746184453e+00
diff --git a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
index 35d3beb82..2edc3a585 100644
--- a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
+++ b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
@@ -13,18 +13,36 @@
 # limitations under the License.
 
 import unittest
+import numpy as np
 from tangelo.toolboxes.post_processing import diis, richardson
 
 
 class ExtrapolationTest(unittest.TestCase):
 
     def test_diis(self):
+        """Test DIIS extrapolation on small sample data from Alejandro
+        """
+        with open(f"./data/diis_test.dat") as f:
+            data = np.loadtxt(f)
+            coeffs = data[:-1, 0]
+            energies = data[:-1, 1]
+            reference = data[-1, 1]
+
+        calculated = diis(energies, coeffs)
         self.assertAlmostEqual(reference, calculated, delta=1e-10)
 
     def test_richardson(self):
+        """Test Richardson extrapolation on small sample data from Alejandro
+        """
+        with open(f"./data/richardson_test.dat") as f:
+            data = np.loadtxt(f)
+            coeffs = data[:-1, 0]
+            energies = data[:-1, 1]
+            reference = data[-1, 1]
+
+        calculated = richardson(energies, coeffs)
         self.assertAlmostEqual(reference, calculated, delta=1e-10)
 
 
-
 if __name__ == "__main__":
     unittest.main()

From d50f361d3110bb370e67283a9b8ea2f61878f791 Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Thu, 16 Dec 2021 17:41:36 -0800
Subject: [PATCH 3/8] Corrected faulty tests

---
 .../toolboxes/post_processing/tests/data/diis_test.dat    | 2 +-
 .../toolboxes/post_processing/tests/test_extrapolation.py | 8 ++++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/tangelo/toolboxes/post_processing/tests/data/diis_test.dat b/tangelo/toolboxes/post_processing/tests/data/diis_test.dat
index 0945ceb2e..20d0c2191 100644
--- a/tangelo/toolboxes/post_processing/tests/data/diis_test.dat
+++ b/tangelo/toolboxes/post_processing/tests/data/diis_test.dat
@@ -2,4 +2,4 @@
 1.100000000000000089e+00 -1.043022885253066523e+00
 1.199999999999999956e+00 -1.033645681252726822e+00
 1.300000000000000044e+00 -1.030052450247785245e+00
-nan -1.065698319746184453e+00
+nan -1.110479305772459568e+00
diff --git a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
index 2edc3a585..6835283c7 100644
--- a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
+++ b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
@@ -13,16 +13,20 @@
 # limitations under the License.
 
 import unittest
+import os
 import numpy as np
+
 from tangelo.toolboxes.post_processing import diis, richardson
 
+path_data = os.path.dirname(os.path.abspath(__file__)) + "/data"
+
 
 class ExtrapolationTest(unittest.TestCase):
 
     def test_diis(self):
         """Test DIIS extrapolation on small sample data from Alejandro
         """
-        with open(f"./data/diis_test.dat") as f:
+        with open(f"{path_data}/diis_test.dat") as f:
             data = np.loadtxt(f)
             coeffs = data[:-1, 0]
             energies = data[:-1, 1]
@@ -34,7 +38,7 @@ def test_diis(self):
     def test_richardson(self):
         """Test Richardson extrapolation on small sample data from Alejandro
         """
-        with open(f"./data/richardson_test.dat") as f:
+        with open(f"{path_data}/richardson_test.dat") as f:
             data = np.loadtxt(f)
             coeffs = data[:-1, 0]
             energies = data[:-1, 1]

From 6034260f048c1486d1c5324e4e7528d3a15d5ceb Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Mon, 17 Jan 2022 20:13:25 +0100
Subject: [PATCH 4/8] New revision of the extrapolation code. Moved common diis
 and Richardson code to separate function. Implemented explicit Richardson
 solution. Added the Richardson with exponent estimation.

---
 .../post_processing/extrapolation.py          | 109 +++++++++++++++---
 .../tests/test_extrapolation.py               |  27 ++---
 2 files changed, 102 insertions(+), 34 deletions(-)

diff --git a/tangelo/toolboxes/post_processing/extrapolation.py b/tangelo/toolboxes/post_processing/extrapolation.py
index c8275733d..9dcae3143 100644
--- a/tangelo/toolboxes/post_processing/extrapolation.py
+++ b/tangelo/toolboxes/post_processing/extrapolation.py
@@ -13,26 +13,64 @@
 # limitations under the License.
 
 import numpy as np
-from scipy.optimize import minimize, root
+import scipy.optimize as sp
 
 
 def diis(energies, coeffs):
     """
-    DIIS extrapolation
+    DIIS extrapolation, originally developped by Pulay in
+    Chemical Physics Letters 73, 393-398 (1980)
 
     Args:
         energies (array-like): Energy expectation values for amplified noise rates
         coeffs (array-like): Noise rate amplification factors
 
+    Returns:
+        float: Extrapolated energy
+    """
+    return extrapolation(energies, coeffs, 1)
+
+
+def richardson(energies, coeffs, estimate_exp=False):
+    """
+    General, DIIS-like extrapolation procedure as found in
+    Nature 567, 491-495 (2019) [arXiv:1805.04492]
+
+    Args:
+        energies (array-like): Energy expectation values for amplified noise rates
+        coeffs (array-like): Noise rate amplification factors
+
+    Returns:
+        float: Extrapolated energy
+    """
+    if estimate_exp==False:
+        return richardson_analytical(energies, coeffs)
+    else:
+        return richardson_with_exp_estimation(energies, coeffs)
+
+
+def extrapolation(energies, coeffs, N=None):
+    """
+    General, DIIS-like extrapolation procedure as found in
+    Nature 567, 491-495 (2019) [arXiv:1805.04492]
+
+    Args:
+        energies (array-like): Energy expectation values for amplified noise rates
+        coeffs (array-like): Noise rate amplification factors
+        N (int): Taylor expanion order; N=None for Richardson extrapolation (order determined from number of datapoints), N=1 for DIIS extrapolation
+
     Returns:
         float: Extrapolated energy
     """
     n = len(coeffs)
+    if N == None:
+        N = n-1
     Eh = np.array(energies)
     ck = np.array(coeffs)
+    ck = np.array([ck**k for k in range(1, N+1)])
     B = np.ones((n+1, n+1))
     B[n, n] = 0
-    B[:n, :n] = ck[:, None] @ ck[None, :]
+    B[:n, :n] = ck.T @ ck
     b = np.zeros(n+1)
     b[n] = 1
     x = np.linalg.lstsq(B, b, rcond=None)[0]
@@ -40,10 +78,10 @@ def diis(energies, coeffs):
     return np.dot(Eh, x)
 
 
-def richardson(energies, coeffs):
+def richardson_analytical(energies, coeffs):
     """
-    Richardson extrapolation as found in
-    Nature 567, 491-495 (2019) [arXiv:1805.04492]
+    Richardson extrapolation exlicit result as found in
+    Phys. Rev. Lett. 119, 180509 [arXiv:1612.02058]
 
     Args:
         energies (array-like): Energy expectation values for amplified noise rates
@@ -52,14 +90,55 @@ def richardson(energies, coeffs):
     Returns:
         float: Extrapolated energy
     """
-    n = len(coeffs)
     Eh = np.array(energies)
-    ck = np.array([coeffs**k for k in range(1, n+1)]).sum(axis=0)
-    B = np.ones((n+1, n+1))
-    B[n, n] = 0
-    B[:n, :n] = ck[:, None] @ ck[None, :]
-    b = np.zeros(n+1)
-    b[n] = 1
-    x = np.linalg.lstsq(B, b, rcond=None)[0]
-    x = x[:-1]
+    ck = np.array(coeffs)
+    x = np.array([np.prod((ai:=np.delete(ck, i))/(a - ai)) for i, a in enumerate(ck)])
     return np.dot(Eh, x)
+
+
+def richardson_with_exp_estimation(energies, coeffs):
+    """
+    Richardson extrapolation by recurrence, with exponent estimation
+
+    Args:
+        energies (array-like): Energy expectation values for amplified noise rates
+        coeffs (array-like): Noise rate amplification factors
+
+    Returns:
+        float: Extrapolated energy
+    """
+    n = len(coeffs)
+    p = 1
+    Eh = np.array(energies)
+    c = np.array(coeffs)
+    ck = np.array(coeffs)
+    p_old = 0
+    dp = 0
+    for i in range(n-1):
+        ti = ck[0]/ck[1]
+        si = ck[0]/ck[2]
+        if ((n > 2) & (i < (n-2))):
+            def f(k):
+                tk = np.sign(ti)*np.abs(ti)**k
+                sk = np.sign(si)*np.abs(si)**k
+                A1 = (tk*Eh[1] - Eh[0])/(tk - 1)
+                A2 = (sk*Eh[2] - Eh[0])/(sk - 1)
+                return (A1 - A2)**2
+            p = sp.minimize(f, p+1, method='BFGS', options={'disp':False}).x[0]
+            if (i == 0):
+                ck = c**p
+        else:
+            break
+        for j in range(n-i-1):
+            ti = (ck[j]/ck[j+1])
+            if (i > 0):
+                dp = p - p_old
+                if (dp == 0):
+                    break
+                ck[j] = ck[j]*(c[j]**dp - c[j+1]**dp)/(ti - 1)
+                ti = (ck[j]/ck[j+1])
+            else:
+                ck[j] = ck[j]*(c[j] - c[j+1])/(ti - 1)
+            Eh[j] = (ti*Eh[j+1] - Eh[j])/(ti - 1)
+        p_old = p
+    return(Eh[0])
diff --git a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
index 6835283c7..dd2c8a486 100644
--- a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
+++ b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
@@ -13,39 +13,28 @@
 # limitations under the License.
 
 import unittest
-import os
 import numpy as np
 
 from tangelo.toolboxes.post_processing import diis, richardson
 
-path_data = os.path.dirname(os.path.abspath(__file__)) + "/data"
-
+energies = [-1.0477557352266311, -1.0430228852530665, -1.0336456812527268, -1.0300524502477852]
+coeffs = [1., 1.1, 1.2, 1.3]
 
 class ExtrapolationTest(unittest.TestCase):
 
     def test_diis(self):
-        """Test DIIS extrapolation on small sample data from Alejandro
+        """Test DIIS extrapolation on small sample data
         """
-        with open(f"{path_data}/diis_test.dat") as f:
-            data = np.loadtxt(f)
-            coeffs = data[:-1, 0]
-            energies = data[:-1, 1]
-            reference = data[-1, 1]
-
         calculated = diis(energies, coeffs)
-        self.assertAlmostEqual(reference, calculated, delta=1e-10)
+        self.assertAlmostEqual(-1.1104793057724596, calculated, delta=1e-10)
 
     def test_richardson(self):
-        """Test Richardson extrapolation on small sample data from Alejandro
+        """Test Richardson extrapolation on small sample data
         """
-        with open(f"{path_data}/richardson_test.dat") as f:
-            data = np.loadtxt(f)
-            coeffs = data[:-1, 0]
-            energies = data[:-1, 1]
-            reference = data[-1, 1]
-
         calculated = richardson(energies, coeffs)
-        self.assertAlmostEqual(reference, calculated, delta=1e-10)
+        self.assertAlmostEqual(-1.4545871813885753, calculated, delta=1e-10)
+        calculated = richardson(energies, coeffs, estimate_exp=True)
+        self.assertAlmostEqual(-1.056015997770032, calculated, delta=1e-10)
 
 
 if __name__ == "__main__":

From a42c22004230dbf8895ba243fe904f1a4a7c3aae Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Tue, 18 Jan 2022 00:06:19 +0100
Subject: [PATCH 5/8] Replaced the walrus assignment expression with a hack.
 Fixed codestyle.

---
 tangelo/toolboxes/post_processing/extrapolation.py       | 9 +++++----
 .../post_processing/tests/test_extrapolation.py          | 1 +
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tangelo/toolboxes/post_processing/extrapolation.py b/tangelo/toolboxes/post_processing/extrapolation.py
index 9dcae3143..390dcb803 100644
--- a/tangelo/toolboxes/post_processing/extrapolation.py
+++ b/tangelo/toolboxes/post_processing/extrapolation.py
@@ -43,7 +43,7 @@ def richardson(energies, coeffs, estimate_exp=False):
     Returns:
         float: Extrapolated energy
     """
-    if estimate_exp==False:
+    if estimate_exp is False:
         return richardson_analytical(energies, coeffs)
     else:
         return richardson_with_exp_estimation(energies, coeffs)
@@ -63,7 +63,7 @@ def extrapolation(energies, coeffs, N=None):
         float: Extrapolated energy
     """
     n = len(coeffs)
-    if N == None:
+    if N is None:
         N = n-1
     Eh = np.array(energies)
     ck = np.array(coeffs)
@@ -92,7 +92,8 @@ def richardson_analytical(energies, coeffs):
     """
     Eh = np.array(energies)
     ck = np.array(coeffs)
-    x = np.array([np.prod((ai:=np.delete(ck, i))/(a - ai)) for i, a in enumerate(ck)])
+    x = np.array([np.prod(ai/(a - ai)) for i, a in enumerate(ck)
+                      for ai in [np.delete(ck, i)]])
     return np.dot(Eh, x)
 
 
@@ -124,7 +125,7 @@ def f(k):
                 A1 = (tk*Eh[1] - Eh[0])/(tk - 1)
                 A2 = (sk*Eh[2] - Eh[0])/(sk - 1)
                 return (A1 - A2)**2
-            p = sp.minimize(f, p+1, method='BFGS', options={'disp':False}).x[0]
+            p = sp.minimize(f, p+1, method='BFGS', options={'disp': False}).x[0]
             if (i == 0):
                 ck = c**p
         else:
diff --git a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
index dd2c8a486..d1b172f54 100644
--- a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
+++ b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
@@ -20,6 +20,7 @@
 energies = [-1.0477557352266311, -1.0430228852530665, -1.0336456812527268, -1.0300524502477852]
 coeffs = [1., 1.1, 1.2, 1.3]
 
+
 class ExtrapolationTest(unittest.TestCase):
 
     def test_diis(self):

From 6a8c22fa71115b7e34bdd1be6126982dde992692 Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Wed, 19 Jan 2022 00:45:54 +0100
Subject: [PATCH 6/8] Implemented changes from the PR comments

---
 .../post_processing/extrapolation.py          | 53 +++++++++++--------
 1 file changed, 31 insertions(+), 22 deletions(-)

diff --git a/tangelo/toolboxes/post_processing/extrapolation.py b/tangelo/toolboxes/post_processing/extrapolation.py
index 390dcb803..2d800f887 100644
--- a/tangelo/toolboxes/post_processing/extrapolation.py
+++ b/tangelo/toolboxes/post_processing/extrapolation.py
@@ -44,12 +44,14 @@ def richardson(energies, coeffs, estimate_exp=False):
         float: Extrapolated energy
     """
     if estimate_exp is False:
+        # If no exponent estimation, run the direct Richardson solution
         return richardson_analytical(energies, coeffs)
     else:
+        # For exponent estimation run the Richardson recursive algorithm
         return richardson_with_exp_estimation(energies, coeffs)
 
 
-def extrapolation(energies, coeffs, N=None):
+def extrapolation(energies, coeffs, taylor_order=None):
     """
     General, DIIS-like extrapolation procedure as found in
     Nature 567, 491-495 (2019) [arXiv:1805.04492]
@@ -57,30 +59,33 @@ def extrapolation(energies, coeffs, N=None):
     Args:
         energies (array-like): Energy expectation values for amplified noise rates
         coeffs (array-like): Noise rate amplification factors
-        N (int): Taylor expanion order; N=None for Richardson extrapolation (order determined from number of datapoints), N=1 for DIIS extrapolation
+        taylor_order (int): Taylor expansion order; None for Richardson extrapolation (order determined from number of datapoints), 1 for DIIS extrapolation
 
     Returns:
         float: Extrapolated energy
     """
     n = len(coeffs)
-    if N is None:
-        N = n-1
+    if taylor_order is None:
+        # Determine the expansion order in case of Richardson extrapolation
+        taylor_order = n-1
     Eh = np.array(energies)
-    ck = np.array(coeffs)
-    ck = np.array([ck**k for k in range(1, N+1)])
+    coeffs = np.array(coeffs)
+    # Setup the linear system matrix
+    ck = np.array([coeffs**k for k in range(1, taylor_order+1)])
     B = np.ones((n+1, n+1))
     B[n, n] = 0
     B[:n, :n] = ck.T @ ck
+    # Setup the free coefficients
     b = np.zeros(n+1)
-    b[n] = 1
+    b[n] = 1 # For the Lagrange multiplier
+    # Solve  the DIIS equations by least squares
     x = np.linalg.lstsq(B, b, rcond=None)[0]
-    x = x[:-1]
-    return np.dot(Eh, x)
+    return np.dot(Eh, x[:-1])
 
 
 def richardson_analytical(energies, coeffs):
     """
-    Richardson extrapolation exlicit result as found in
+    Richardson extrapolation explicit result as found in
     Phys. Rev. Lett. 119, 180509 [arXiv:1612.02058]
 
     Args:
@@ -109,32 +114,36 @@ def richardson_with_exp_estimation(energies, coeffs):
         float: Extrapolated energy
     """
     n = len(coeffs)
-    p = 1
     Eh = np.array(energies)
     c = np.array(coeffs)
     ck = np.array(coeffs)
-    p_old = 0
-    dp = 0
+    p, p_old = 1, 0
+    
+    # Define a helper function for exponent optimization
+    def energy_diff(k, ti, si):
+        tk = np.sign(ti)*np.abs(ti)**k
+        sk = np.sign(si)*np.abs(si)**k
+        Et = (tk*Eh[1] - Eh[0])/(tk - 1)
+        Es = (sk*Eh[2] - Eh[0])/(sk - 1)
+        return (Et - Es)**2
+    
+    # Run the Richardson algorithm with exponent optimization
     for i in range(n-1):
         ti = ck[0]/ck[1]
         si = ck[0]/ck[2]
-        if ((n > 2) & (i < (n-2))):
-            def f(k):
-                tk = np.sign(ti)*np.abs(ti)**k
-                sk = np.sign(si)*np.abs(si)**k
-                A1 = (tk*Eh[1] - Eh[0])/(tk - 1)
-                A2 = (sk*Eh[2] - Eh[0])/(sk - 1)
-                return (A1 - A2)**2
-            p = sp.minimize(f, p+1, method='BFGS', options={'disp': False}).x[0]
+        if ((n > 2) and (i < (n-2))):
+            # Minimize the energy difference to determine the optimal exponent
+            p = sp.minimize(energy_diff, p+1, args=(ti, si), method='BFGS', options={'disp': False}).x[0]
             if (i == 0):
                 ck = c**p
         else:
             break
+        
         for j in range(n-i-1):
             ti = (ck[j]/ck[j+1])
             if (i > 0):
                 dp = p - p_old
-                if (dp == 0):
+                if (np.isclose(dp, 0)):
                     break
                 ck[j] = ck[j]*(c[j]**dp - c[j+1]**dp)/(ti - 1)
                 ti = (ck[j]/ck[j+1])

From 9c971569615b2d737f674d4a6ebc10d3a2e29e09 Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Wed, 19 Jan 2022 00:52:00 +0100
Subject: [PATCH 7/8] Codestyle fixes

---
 tangelo/toolboxes/post_processing/extrapolation.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tangelo/toolboxes/post_processing/extrapolation.py b/tangelo/toolboxes/post_processing/extrapolation.py
index 2d800f887..6d85e4dd5 100644
--- a/tangelo/toolboxes/post_processing/extrapolation.py
+++ b/tangelo/toolboxes/post_processing/extrapolation.py
@@ -77,7 +77,7 @@ def extrapolation(energies, coeffs, taylor_order=None):
     B[:n, :n] = ck.T @ ck
     # Setup the free coefficients
     b = np.zeros(n+1)
-    b[n] = 1 # For the Lagrange multiplier
+    b[n] = 1  # For the Lagrange multiplier
     # Solve  the DIIS equations by least squares
     x = np.linalg.lstsq(B, b, rcond=None)[0]
     return np.dot(Eh, x[:-1])
@@ -118,7 +118,7 @@ def richardson_with_exp_estimation(energies, coeffs):
     c = np.array(coeffs)
     ck = np.array(coeffs)
     p, p_old = 1, 0
-    
+
     # Define a helper function for exponent optimization
     def energy_diff(k, ti, si):
         tk = np.sign(ti)*np.abs(ti)**k
@@ -126,7 +126,7 @@ def energy_diff(k, ti, si):
         Et = (tk*Eh[1] - Eh[0])/(tk - 1)
         Es = (sk*Eh[2] - Eh[0])/(sk - 1)
         return (Et - Es)**2
-    
+
     # Run the Richardson algorithm with exponent optimization
     for i in range(n-1):
         ti = ck[0]/ck[1]
@@ -138,7 +138,7 @@ def energy_diff(k, ti, si):
                 ck = c**p
         else:
             break
-        
+
         for j in range(n-i-1):
             ti = (ck[j]/ck[j+1])
             if (i > 0):

From 302c559d79765013825a9e2228d6a1cc57d58e9f Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Wed, 19 Jan 2022 22:48:45 +0100
Subject: [PATCH 8/8] Reduced precision in tests and removed old test datafiles

---
 tangelo/toolboxes/post_processing/extrapolation.py        | 4 ++++
 .../toolboxes/post_processing/tests/data/diis_test.dat    | 5 -----
 .../post_processing/tests/data/richardson_test.dat        | 5 -----
 .../toolboxes/post_processing/tests/test_extrapolation.py | 8 ++++----
 4 files changed, 8 insertions(+), 14 deletions(-)
 delete mode 100644 tangelo/toolboxes/post_processing/tests/data/diis_test.dat
 delete mode 100644 tangelo/toolboxes/post_processing/tests/data/richardson_test.dat

diff --git a/tangelo/toolboxes/post_processing/extrapolation.py b/tangelo/toolboxes/post_processing/extrapolation.py
index 6d85e4dd5..f9bf8f048 100644
--- a/tangelo/toolboxes/post_processing/extrapolation.py
+++ b/tangelo/toolboxes/post_processing/extrapolation.py
@@ -70,14 +70,17 @@ def extrapolation(energies, coeffs, taylor_order=None):
         taylor_order = n-1
     Eh = np.array(energies)
     coeffs = np.array(coeffs)
+
     # Setup the linear system matrix
     ck = np.array([coeffs**k for k in range(1, taylor_order+1)])
     B = np.ones((n+1, n+1))
     B[n, n] = 0
     B[:n, :n] = ck.T @ ck
+
     # Setup the free coefficients
     b = np.zeros(n+1)
     b[n] = 1  # For the Lagrange multiplier
+
     # Solve  the DIIS equations by least squares
     x = np.linalg.lstsq(B, b, rcond=None)[0]
     return np.dot(Eh, x[:-1])
@@ -139,6 +142,7 @@ def energy_diff(k, ti, si):
         else:
             break
 
+        # Run the main Richardson loop
         for j in range(n-i-1):
             ti = (ck[j]/ck[j+1])
             if (i > 0):
diff --git a/tangelo/toolboxes/post_processing/tests/data/diis_test.dat b/tangelo/toolboxes/post_processing/tests/data/diis_test.dat
deleted file mode 100644
index 20d0c2191..000000000
--- a/tangelo/toolboxes/post_processing/tests/data/diis_test.dat
+++ /dev/null
@@ -1,5 +0,0 @@
-1.000000000000000000e+00 -1.047755735226631124e+00
-1.100000000000000089e+00 -1.043022885253066523e+00
-1.199999999999999956e+00 -1.033645681252726822e+00
-1.300000000000000044e+00 -1.030052450247785245e+00
-nan -1.110479305772459568e+00
diff --git a/tangelo/toolboxes/post_processing/tests/data/richardson_test.dat b/tangelo/toolboxes/post_processing/tests/data/richardson_test.dat
deleted file mode 100644
index 0945ceb2e..000000000
--- a/tangelo/toolboxes/post_processing/tests/data/richardson_test.dat
+++ /dev/null
@@ -1,5 +0,0 @@
-1.000000000000000000e+00 -1.047755735226631124e+00
-1.100000000000000089e+00 -1.043022885253066523e+00
-1.199999999999999956e+00 -1.033645681252726822e+00
-1.300000000000000044e+00 -1.030052450247785245e+00
-nan -1.065698319746184453e+00
diff --git a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
index d1b172f54..6f8ffd88f 100644
--- a/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
+++ b/tangelo/toolboxes/post_processing/tests/test_extrapolation.py
@@ -17,7 +17,7 @@
 
 from tangelo.toolboxes.post_processing import diis, richardson
 
-energies = [-1.0477557352266311, -1.0430228852530665, -1.0336456812527268, -1.0300524502477852]
+energies = [-1.04775574, -1.04302289, -1.03364568, -1.03005245]
 coeffs = [1., 1.1, 1.2, 1.3]
 
 
@@ -27,15 +27,15 @@ def test_diis(self):
         """Test DIIS extrapolation on small sample data
         """
         calculated = diis(energies, coeffs)
-        self.assertAlmostEqual(-1.1104793057724596, calculated, delta=1e-10)
+        self.assertAlmostEqual(-1.11047933, calculated, delta=1e-6)
 
     def test_richardson(self):
         """Test Richardson extrapolation on small sample data
         """
         calculated = richardson(energies, coeffs)
-        self.assertAlmostEqual(-1.4545871813885753, calculated, delta=1e-10)
+        self.assertAlmostEqual(-1.45459036, calculated, delta=1e-6)
         calculated = richardson(energies, coeffs, estimate_exp=True)
-        self.assertAlmostEqual(-1.056015997770032, calculated, delta=1e-10)
+        self.assertAlmostEqual(-1.05601603, calculated, delta=1e-6)
 
 
 if __name__ == "__main__":