From 9fbd1bf6b426b533bb4bbf2c6f1f198f69440739 Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Tue, 2 Jan 2024 08:25:46 +0000
Subject: [PATCH 1/9] Uncomment and fix bug in ucc_ansatz function

---
 src/qibochem/ansatz/ucc.py | 163 ++++++++++++++++++++-----------------
 1 file changed, 90 insertions(+), 73 deletions(-)

diff --git a/src/qibochem/ansatz/ucc.py b/src/qibochem/ansatz/ucc.py
index 8b3bffd..b1b08b0 100644
--- a/src/qibochem/ansatz/ucc.py
+++ b/src/qibochem/ansatz/ucc.py
@@ -6,6 +6,8 @@
 import openfermion
 from qibo import gates, models
 
+from qibochem.ansatz.hf_reference import hf_circuit
+
 
 def expi_pauli(n_qubits, theta, pauli_string):
     """
@@ -166,9 +168,9 @@ def generate_excitations(order, excite_from, excite_to, conserve_spin=True):
     if order > min(len(excite_from), len(excite_to)):
         return [[]]
 
+    # Generate all possible excitations first
     from itertools import combinations
 
-    # Generate all possible excitations first
     all_excitations = [
         [*_from, *_to] for _from in combinations(excite_from, order) for _to in combinations(excite_to, order)
     ]
@@ -228,7 +230,7 @@ def sort_excitations(excitations):
                 pair_excitations = sorted(pair_excitations)
                 ex_to_remove = pair_excitations.pop(0)
                 if ex_to_remove in copy_excitations:
-                    # 'Move' the first entry of same_mo_index from copy_excitations to result
+                    # 'Move' the first pair excitation from copy_excitations to result
                     index = copy_excitations.index(ex_to_remove)
                     result.append(copy_excitations.pop(index))
 
@@ -253,77 +255,92 @@ def sort_excitations(excitations):
                     result.append(copy_excitations.pop(index))
             prev = None
             continue
-        # Remove the first entry from the sorted list of remaining excitations and add it to result
-        prev = copy_excitations.pop(0)
-        result.append(prev)
+        if copy_excitations:
+            # Remove the first entry from the sorted list of remaining excitations and add it to result
+            prev = copy_excitations.pop(0)
+            result.append(prev)
     return result
 
 
-# def ucc_ansatz(molecule, excitation_level=None, excitations=None, thetas=None, trotter_steps=1, ferm_qubit_map=None):
-#     """
-#     Build a circuit corresponding to the UCC ansatz with multiple excitations for a given Molecule.
-#         If no excitations are given, it defaults to returning the full UCCSD circuit ansatz for the
-#         Molecule.
-#
-#     Args:
-#         molecule: Molecule of interest.
-#         excitation_level: Include excitations up to how many electrons, i.e. ("S", "D", "T", "Q")
-#             Ignored if `excitations` argument is given. Default is "D", i.e. double excitations
-#         excitations: List of excitations (e.g. `[[0, 1, 2, 3], [0, 1, 4, 5]]`) used to build the
-#             UCC circuit. Overrides the `excitation_level` argument
-#         thetas: Parameters for the excitations. Defaults to an array of zeros if not given
-#         trotter_steps: number of Trotter steps; i.e. number of times the UCC ansatz is applied with
-#             theta=theta/trotter_steps. Default is 1
-#         ferm_qubit_map: fermion-to-qubit transformation. Default is Jordan-Wigner ("jw")
-#
-#     Returns:
-#         circuit: Qibo Circuit corresponding to the UCC ansatz
-#     """
-#     # Get the number of electrons and virtual orbitals from the molecule argument
-#     n_elec = sum(molecule.nelec) if molecule.n_active_e is None else molecule.n_active_e
-#     n_orbs = molecule.nso if molecule.n_active_orbs is None else molecule.n_active_orbs
-#
-#     # Define the excitation level to be used if no excitations given
-#     if excitations is None:
-#         excitation_levels = ("S", "D", "T", "Q")
-#         if excitation_level is None:
-#             excitation_level = "D"
-#         else:
-#             # Check validity of input
-#             assert len(excitation_level) == 1 and excitation_level.upper() in excitation_levels
-#             # Note: Probably don't be too ambitious and try to do 'T'/'Q' at the moment...
-#             if excitation_level.upper() in ("T", "Q"):
-#                 raise NotImplementedError("Cannot handle triple and quadruple excitations!")
-#         # Get the (largest) order of excitation to use
-#         excitation_order = excitation_levels.index(excitation_level.upper()) + 1
-#
-#         # Generate and sort all the possible excitations
-#         excitations = []
-#         for order in range(excitation_order, 0, -1):  # Reversed to get higher excitations first
-#             excitations += sort_excitations(generate_excitations(order, range(0, n_elec), range(n_elec, n_orbs)))
-#     else:
-#         # Some checks to ensure the given excitations are valid
-#         assert all(len(_ex) % 2 == 0 for _ex in excitations), "Excitation with an odd number of elements found!"
-#
-#     # Check if thetas argument given, define to be all zeros if not
-#     # TODO: Unsure if want to use MP2 guess amplitudes for the doubles? Some say good, some say bad
-#     # Number of circuit parameters: S->2, D->8, (T/Q->32/128; Not sure?)
-#     n_parameters = 2 * len([_ex for _ex in excitations if len(_ex) == 2])  # Singles
-#     n_parameters += 8 * len([_ex for _ex in excitations if len(_ex) == 4])  # Doubles
-#     if thetas is None:
-#         thetas = np.zeros(n_parameters)
-#     else:
-#         # Check that number of circuit variables (i.e. thetas) matches the number of circuit parameters
-#         assert len(thetas) == n_parameters, "Number of input parameters doesn't match the number of circuit parameters!"
-#
-#     # Build the circuit
-#     circuit = models.Circuit(n_orbs)
-#     for excitation, theta in zip(excitations, thetas):
-#         # coeffs = []
-#         circuit += ucc_circuit(
-#             n_orbs, excitation, theta, trotter_steps=trotter_steps, ferm_qubit_map=ferm_qubit_map  # , coeffs=coeffs)
-#         )
-#         # if isinstance(all_coeffs, list):
-#         #     all_coeffs.append(np.array(coeffs))
-#
-#     return circuit
+def ucc_ansatz(
+    molecule,
+    excitation_level=None,
+    excitations=None,
+    thetas=None,
+    trotter_steps=1,
+    ferm_qubit_map=None,
+    include_hf=True,
+):
+    """
+    Build a circuit corresponding to the UCC ansatz with multiple excitations for a given Molecule.
+        If no excitations are given, it defaults to returning the full UCCSD circuit ansatz for the
+        Molecule.
+
+    Args:
+        molecule: Molecule of interest.
+        excitation_level: Include excitations up to how many electrons, i.e. ("S", "D", "T", "Q")
+            Ignored if `excitations` argument is given. Default is "D", i.e. double excitations
+        excitations: List of excitations (e.g. `[[0, 1, 2, 3], [0, 1, 4, 5]]`) used to build the
+            UCC circuit. Overrides the `excitation_level` argument
+        thetas: Parameters for the excitations. Defaults to an array of zeros if not given
+        trotter_steps: number of Trotter steps; i.e. number of times the UCC ansatz is applied with
+            theta=theta/trotter_steps. Default is 1
+        ferm_qubit_map: fermion-to-qubit transformation. Default: Jordan-Wigner ("jw")
+        include_hf: Whether or not to start the circuit with a Hartree-Fock circuit. Default: ``True``
+
+    Returns:
+        circuit: Qibo Circuit corresponding to the UCC ansatz
+    """
+    # Get the number of electrons and spin-orbitals from the molecule argument
+    n_elec = sum(molecule.nelec) if molecule.n_active_e is None else molecule.n_active_e
+    n_orbs = molecule.nso if molecule.n_active_orbs is None else molecule.n_active_orbs
+
+    # Define the excitation level to be used if no excitations given
+    if excitations is None:
+        excitation_levels = ("S", "D", "T", "Q")
+        if excitation_level is None:
+            excitation_level = "D"
+        else:
+            # Check validity of input
+            assert (
+                len(excitation_level) == 1 and excitation_level.upper() in excitation_levels
+            ), "Unknown input for excitation_level"
+            # Note: Probably don't be too ambitious and try to do 'T'/'Q' at the moment...
+            if excitation_level.upper() in ("T", "Q"):
+                raise NotImplementedError("Cannot handle triple and quadruple excitations!")
+        # Get the (largest) order of excitation to use
+        excitation_order = excitation_levels.index(excitation_level.upper()) + 1
+
+        # Generate and sort all the possible excitations
+        excitations = []
+        for order in range(excitation_order, 0, -1):  # Reversed to get higher excitations first
+            excitations += sort_excitations(generate_excitations(order, range(0, n_elec), range(n_elec, n_orbs)))
+    else:
+        # Some checks to ensure the given excitations are valid
+        assert all(len(_ex) % 2 == 0 for _ex in excitations), "Excitation with an odd number of elements found!"
+
+    # Check if thetas argument given, define to be all zeros if not
+    # TODO: Unsure if want to use MP2 guess amplitudes for the doubles? Some say good, some say bad
+    # Number of circuit parameters: S->2, D->8, (T/Q->32/128; Not sure?)
+    n_parameters = 2 * len([_ex for _ex in excitations if len(_ex) == 2])  # Singles
+    n_parameters += 8 * len([_ex for _ex in excitations if len(_ex) == 4])  # Doubles
+    if thetas is None:
+        thetas = np.zeros(n_parameters)
+    else:
+        # Check that number of circuit variables (i.e. thetas) matches the number of circuit parameters
+        assert len(thetas) == n_parameters, "Number of input parameters doesn't match the number of circuit parameters!"
+
+    # Build the circuit
+    if include_hf:
+        circuit = hf_circuit(n_orbs, n_elec, ferm_qubit_map=ferm_qubit_map)
+    else:
+        circuit = models.Circuit(n_orbs)
+    for excitation, theta in zip(excitations, thetas):
+        # coeffs = []
+        circuit += ucc_circuit(
+            n_orbs, excitation, theta, trotter_steps=trotter_steps, ferm_qubit_map=ferm_qubit_map  # , coeffs=coeffs)
+        )
+        # if isinstance(all_coeffs, list):
+        #     all_coeffs.append(np.array(coeffs))
+
+    return circuit

From 2ca730b0e2bf73cd555a1453eb3efdc6360f1f57 Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Wed, 3 Jan 2024 03:04:33 +0000
Subject: [PATCH 2/9] Extend MP2 guess amplitude function for singles

Aka:just return 0.0 if input excitation is a single excitation
---
 src/qibochem/ansatz/ucc.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/qibochem/ansatz/ucc.py b/src/qibochem/ansatz/ucc.py
index b1b08b0..2e7711f 100644
--- a/src/qibochem/ansatz/ucc.py
+++ b/src/qibochem/ansatz/ucc.py
@@ -118,22 +118,28 @@ def ucc_circuit(n_qubits, excitation, theta=0.0, trotter_steps=1, ferm_qubit_map
 # Utility functions
 
 
-def mp2_amplitude(orbitals, orbital_energies, tei):
+def mp2_amplitude(excitation, orbital_energies, tei):
     """
-    Calculate the MP2 guess amplitudes to be used in the UCC doubles ansatz
-        In SO basis: t_{ij}^{ab} = (g_{ijab} - g_{ijba}) / (e_i + e_j - e_a - e_b)
+    Calculate the MP2 guess amplitude for a single UCC circuit: 0.0 for a single excitation.
+        for a double excitation (In SO basis): t_{ij}^{ab} = (g_{ijab} - g_{ijba}) / (e_i + e_j - e_a - e_b)
 
     Args:
-        orbitals: list of spin-orbitals representing a double excitation, must have exactly
-            4 elements
+        excitation: Iterable of spin-orbitals representing a excitation. Must have either 2 or 4 elements exactly,
+            representing a single or double excitation respectively.
         orbital_energies: eigenvalues of the Fock operator, i.e. orbital energies
         tei: Two-electron integrals in MO basis and second quantization notation
+
+    Returns:
+        MP2 guess amplitude (float)
     """
-    # Checks orbitals
-    assert len(orbitals) == 4, f"{orbitals} must have only 4 orbitals for a double excitation"
-    # Convert orbital indices to be in MO basis
-    mo_orbitals = [_orb // 2 for _orb in orbitals]
+    # Checks validity of excitation argument
+    assert len(orbitals) // 2 == 0 and len(orbitals) // 2 <= 2, f"{excitation} must have either 2 or 4 elements"
+    # If single excitation, can just return 0.0 directly
+    if len(orbitals) == 2:
+        return 0.0
 
+    # Convert orbital indices to be in MO basis
+    mo_orbitals = [orbitalin // 2 for orbital in excitation]
     # Numerator: g_ijab - g_ijba
     g_ijab = (
         tei[tuple(mo_orbitals)]  # Can index directly using the MO TEIs

From 6799e094b57ca0c58e14f5fae727b97608054d7a Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Wed, 3 Jan 2024 03:18:18 +0000
Subject: [PATCH 3/9] Fix minor bug

---
 src/qibochem/ansatz/ucc.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/qibochem/ansatz/ucc.py b/src/qibochem/ansatz/ucc.py
index 2e7711f..4dcd7ad 100644
--- a/src/qibochem/ansatz/ucc.py
+++ b/src/qibochem/ansatz/ucc.py
@@ -133,22 +133,22 @@ def mp2_amplitude(excitation, orbital_energies, tei):
         MP2 guess amplitude (float)
     """
     # Checks validity of excitation argument
-    assert len(orbitals) // 2 == 0 and len(orbitals) // 2 <= 2, f"{excitation} must have either 2 or 4 elements"
+    assert len(excitation) // 2 == 0 and len(excitation) // 2 <= 2, f"{excitation} must have either 2 or 4 elements"
     # If single excitation, can just return 0.0 directly
-    if len(orbitals) == 2:
+    if len(excitation) == 2:
         return 0.0
 
     # Convert orbital indices to be in MO basis
-    mo_orbitals = [orbitalin // 2 for orbital in excitation]
+    mo_orbitals = [orbital // 2 for orbital in excitation]
     # Numerator: g_ijab - g_ijba
     g_ijab = (
         tei[tuple(mo_orbitals)]  # Can index directly using the MO TEIs
-        if (orbitals[0] + orbitals[3]) % 2 == 0 and (orbitals[1] + orbitals[2]) % 2 == 0
+        if (excitation[0] + excitation[3]) % 2 == 0 and (excitation[1] + excitation[2]) % 2 == 0
         else 0.0
     )
     g_ijba = (
         tei[tuple(mo_orbitals[:2] + mo_orbitals[2:][::-1])]  # Reverse last two terms
-        if (orbitals[0] + orbitals[2]) % 2 == 0 and (orbitals[1] + orbitals[3]) % 2 == 0
+        if (excitation[0] + excitation[2]) % 2 == 0 and (excitation[1] + excitation[3]) % 2 == 0
         else 0.0
     )
     numerator = g_ijab - g_ijba

From 03043f663bdf667bb249ff7e82f2e2b3f5d73a9e Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Wed, 3 Jan 2024 03:29:03 +0000
Subject: [PATCH 4/9] Fix another stupid bug

Haiz, first time working in office in 2024, somehow the mind is sooooo sloppy and careless
---
 src/qibochem/ansatz/ucc.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qibochem/ansatz/ucc.py b/src/qibochem/ansatz/ucc.py
index 4dcd7ad..aca17d4 100644
--- a/src/qibochem/ansatz/ucc.py
+++ b/src/qibochem/ansatz/ucc.py
@@ -133,7 +133,7 @@ def mp2_amplitude(excitation, orbital_energies, tei):
         MP2 guess amplitude (float)
     """
     # Checks validity of excitation argument
-    assert len(excitation) // 2 == 0 and len(excitation) // 2 <= 2, f"{excitation} must have either 2 or 4 elements"
+    assert len(excitation) % 2 == 0 and len(excitation) // 2 <= 2, f"{excitation} must have either 2 or 4 elements"
     # If single excitation, can just return 0.0 directly
     if len(excitation) == 2:
         return 0.0

From 3721e10301486439492a962e8e1bd0c3493e27c8 Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Wed, 3 Jan 2024 06:31:21 +0000
Subject: [PATCH 5/9] Edit and expand UCC test cases

---
 tests/test_ucc.py | 180 +++++++++++++++++++++++++++++-----------------
 1 file changed, 115 insertions(+), 65 deletions(-)

diff --git a/tests/test_ucc.py b/tests/test_ucc.py
index 46b50c5..1c759e1 100644
--- a/tests/test_ucc.py
+++ b/tests/test_ucc.py
@@ -1,3 +1,5 @@
+from functools import partial
+
 import numpy as np
 import pytest
 from qibo import gates
@@ -8,6 +10,7 @@
     generate_excitations,
     mp2_amplitude,
     sort_excitations,
+    ucc_ansatz,
     ucc_circuit,
 )
 from qibochem.driver.molecule import Molecule
@@ -76,7 +79,11 @@ def test_sort_excitations_3():
         sort_excitations([[1, 2, 3, 4, 5, 6]])
 
 
-def test_mp2_amplitude():
+def test_mp2_amplitude_singles():
+    assert mp2_amplitude([0, 2], np.random.rand(4), np.random.rand(4, 4)) == 0.0
+
+
+def test_mp2_amplitude_doubles():
     h2 = Molecule([("H", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 0.7))])
     h2.run_pyscf()
     l = mp2_amplitude([0, 1, 2, 3], h2.eps, h2.tei)
@@ -85,7 +92,113 @@ def test_mp2_amplitude():
     assert np.isclose(l, ref_l)
 
 
-def test_uccsd():
+def test_ucc_jw_singles():
+    """Build a UCC singles (JW) circuit"""
+    rx_gate = partial(gates.RX, theta=-0.5 * np.pi, trainable=False)
+    cnot_cascade = [gates.CNOT(2, 1), gates.CNOT(1, 0), gates.RZ(0, 0.0), gates.CNOT(1, 0), gates.CNOT(2, 1)]
+    basis_rotation_gates = ([rx_gate(0), gates.H(2)], [gates.H(0), rx_gate(2)])
+
+    # Build control list of gates
+    nested_gate_list = [gate_list + cnot_cascade + gate_list for gate_list in basis_rotation_gates]
+    gate_list = [_gate for gate_list in nested_gate_list for _gate in gate_list]
+
+    # Test ucc_function
+    circuit = ucc_circuit(4, [0, 2], ferm_qubit_map="jw")
+    # Check gates are correct
+    assert all(
+        control.name == test.name and control.target_qubits == test.target_qubits
+        for control, test in zip(gate_list, list(circuit.queue))
+    )
+    # Check that only two parametrised gates
+    assert len(circuit.get_parameters()) == 2
+
+
+def test_ucc_jw_doubles():
+    """Build a UCC doubles (JW) circuit"""
+    rx_gate = partial(gates.RX, theta=-0.5 * np.pi, trainable=False)
+    cnot_cascade = [
+        gates.CNOT(3, 2),
+        gates.CNOT(2, 1),
+        gates.CNOT(1, 0),
+        gates.RZ(0, 0.0),
+        gates.CNOT(1, 0),
+        gates.CNOT(2, 1),
+        gates.CNOT(3, 2),
+    ]
+    basis_rotation_gates = (
+        [gates.H(0), gates.H(1), rx_gate(2), gates.H(3)],
+        [rx_gate(0), rx_gate(1), rx_gate(2), gates.H(3)],
+        [rx_gate(0), gates.H(1), gates.H(2), gates.H(3)],
+        [gates.H(0), rx_gate(1), gates.H(2), gates.H(3)],
+        [rx_gate(0), gates.H(1), rx_gate(2), rx_gate(3)],
+        [gates.H(0), rx_gate(1), rx_gate(2), rx_gate(3)],
+        [gates.H(0), gates.H(1), gates.H(2), rx_gate(3)],
+        [rx_gate(0), rx_gate(1), gates.H(2), rx_gate(3)],
+    )
+
+    # Build control list of gates
+    nested_gate_list = [gate_list + cnot_cascade + gate_list for gate_list in basis_rotation_gates]
+    gate_list = [_gate for gate_list in nested_gate_list for _gate in gate_list]
+
+    # Test ucc_function
+    circuit = ucc_circuit(4, [0, 1, 2, 3], ferm_qubit_map="jw")
+    # Check gates are correct
+    assert all(
+        control.name == test.name and control.target_qubits == test.target_qubits
+        for control, test in zip(gate_list, list(circuit.queue))
+    )
+    # Check that only two parametrised gates
+    assert len(circuit.get_parameters()) == 8
+
+
+def test_ucc_bk_singles():
+    """Build a UCC doubles (BK) circuit"""
+    rx_gate = partial(gates.RX, theta=-0.5 * np.pi, trainable=False)
+    cnot_cascade = [gates.CNOT(2, 1), gates.CNOT(1, 0), gates.RZ(0, 0.0), gates.CNOT(1, 0), gates.CNOT(2, 1)]
+    basis_rotation_gates = ([gates.H(0), rx_gate(1), gates.H(2)], [rx_gate(0), rx_gate(1), rx_gate(2)])
+
+    # Build control list of gates
+    nested_gate_list = [gate_list + cnot_cascade + gate_list for gate_list in basis_rotation_gates]
+    gate_list = [_gate for gate_list in nested_gate_list for _gate in gate_list]
+
+    # Test ucc_function
+    circuit = ucc_circuit(4, [0, 2], ferm_qubit_map="bk")
+    # Check gates are correct
+    assert all(
+        control.name == test.name and control.target_qubits == test.target_qubits
+        for control, test in zip(gate_list, list(circuit.queue))
+    )
+    # Check that only two parametrised gates
+    assert len(circuit.get_parameters()) == 2
+
+
+def test_ucc_ferm_qubit_map_error():
+    """If unknown fermion to qubit map used"""
+    with pytest.raises(KeyError):
+        ucc_circuit(2, [0, 1], ferm_qubit_map="Unknown")
+
+
+def test_ucc_ansatz_h2():
+    """Test the default arguments of ucc_ansatz using H2"""
+    mol = Molecule([("H", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 0.7))])
+    mol.run_pyscf()
+
+    # Build control circuit
+    control_circuit = hf_circuit(4, 2)
+    for excitation in ([0, 1, 2, 3], [0, 2], [1, 3]):
+        control_circuit += ucc_circuit(4, excitation)
+
+    test_circuit = ucc_ansatz(mol)
+
+    assert all(
+        control.name == test.name and control.target_qubits == test.target_qubits
+        for control, test in zip(list(control_circuit.queue), list(test_circuit.queue))
+    )
+    # Check that number of parametrised gates is the same
+    assert len(control_circuit.get_parameters()) == len(test_circuit.get_parameters())
+
+
+def test_ucc_ansatz_h2():
     # Define molecule and populate
     mol = Molecule([("Li", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 1.2))])
     mol.run_pyscf()
@@ -145,66 +258,3 @@ def electronic_energy(parameters):
     lih_uccsd_energy = -7.847535097575567
 
     assert vqe.fun == pytest.approx(lih_uccsd_energy)
-
-
-def test_ucc_bk():
-    """Build a UCC-BK circuit"""
-    # Control
-    gate_list = [
-        gates.H(0),
-        gates.RX(1, -0.5 * np.pi, trainable=False),
-        gates.H(2),
-        gates.CNOT(2, 1),
-        gates.CNOT(1, 0),
-        gates.RZ(0, 0.0),
-        gates.CNOT(1, 0),
-        gates.CNOT(2, 1),
-        gates.H(0),
-        gates.RX(1, -0.5 * np.pi, trainable=False),
-        gates.H(2),
-        gates.RX(0, -0.5 * np.pi, trainable=False),
-        gates.RX(1, -0.5 * np.pi, trainable=False),
-        gates.RX(2, -0.5 * np.pi, trainable=False),
-        gates.CNOT(2, 1),
-        gates.CNOT(1, 0),
-        gates.RZ(0, 0.0),
-        gates.CNOT(1, 0),
-        gates.CNOT(2, 1),
-        gates.RX(0, -0.5 * np.pi, trainable=False),
-        gates.RX(1, -0.5 * np.pi, trainable=False),
-        gates.RX(2, -0.5 * np.pi, trainable=False),
-    ]
-    # Test ucc_function
-    circuit = ucc_circuit(4, [0, 2], ferm_qubit_map="bk")
-    # Check gates are correct
-    assert all(
-        control.name == test.name and control.target_qubits == test.target_qubits
-        for control, test in zip(gate_list, list(circuit.queue))
-    )
-    # Check that only two parametrised gates
-    assert len(circuit.get_parameters()) == 2
-
-
-def test_ucc_ferm_qubit_map_error():
-    """If unknown fermion to qubit map used"""
-    with pytest.raises(KeyError):
-        ucc_circuit(2, [0, 1], ferm_qubit_map="Unknown")
-
-
-# def test_ucc_ansatz_default_args():
-#     mol = Molecule([("H", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 0.7))])
-#     mol.run_pyscf()
-#
-#     # Build control circuit
-#     control_circuit = hf_circuit(4, 2)
-#     for excitation in ([0, 1, 2, 3], [0, 2], [1, 3]):
-#         control_circuit += ucc_circuit(4, excitation)
-#
-#     test_circuit = ucc_ansatz(mol)
-#
-#     assert all(
-#         control.name == test.name and control.target_qubits == test.target_qubits
-#         for control, test in zip(list(control_circuit.queue), list(test_circuit.queue))
-#     )
-#     # Check that number of parametrised gates is the same
-#     assert len(control_circuit.get_parameters()) == len(test_circuit.get_parameters())

From 95cbc78c5ce439b46516c374cd748cf0d55fd8a9 Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Wed, 3 Jan 2024 06:32:29 +0000
Subject: [PATCH 6/9] Remove VQE-UCC test to shorten overall test times

---
 tests/test_ucc.py | 62 -----------------------------------------------
 1 file changed, 62 deletions(-)

diff --git a/tests/test_ucc.py b/tests/test_ucc.py
index 1c759e1..c767e7d 100644
--- a/tests/test_ucc.py
+++ b/tests/test_ucc.py
@@ -196,65 +196,3 @@ def test_ucc_ansatz_h2():
     )
     # Check that number of parametrised gates is the same
     assert len(control_circuit.get_parameters()) == len(test_circuit.get_parameters())
-
-
-def test_ucc_ansatz_h2():
-    # Define molecule and populate
-    mol = Molecule([("Li", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 1.2))])
-    mol.run_pyscf()
-
-    # Apply embedding
-    active_space = [1, 5]
-    mol.hf_embedding(active=[1, 5])
-    hamiltonian = mol.hamiltonian(oei=mol.embed_oei, tei=mol.embed_tei, constant=mol.inactive_energy)
-
-    # Set parameters for the rest of the experiment
-    n_qubits = mol.n_active_orbs
-    n_electrons = mol.n_active_e
-
-    excitations = generate_excitations(2, tuple(range(n_electrons)), tuple(range(n_electrons, n_qubits)))
-    excitations += generate_excitations(1, tuple(range(n_electrons)), tuple(range(n_electrons, n_qubits)))
-    n_excitations = len(excitations)
-
-    # Build circuit
-    circuit = hf_circuit(n_qubits, n_electrons)
-
-    # UCCSD: Circuit
-    all_coeffs = []
-    for _ex in excitations:
-        coeffs = []
-        circuit += ucc_circuit(n_qubits, _ex, coeffs=coeffs)
-        all_coeffs.append(coeffs)
-
-    def electronic_energy(parameters):
-        r"""
-        Loss function for the UCCSD ansatz
-        """
-        all_parameters = []
-
-        # UCC parameters
-        # Expand the parameters to match the total UCC ansatz manually
-        _ucc = parameters[:n_excitations]
-        # Manually group the related excitations together
-        ucc_parameters = [_ucc[0], _ucc[1], _ucc[2]]
-        # Need to iterate through each excitation this time
-        for _coeffs, _parameter in zip(all_coeffs, ucc_parameters):
-            # Convert a single value to a array with dimension=n_param_gates
-            ucc_parameter = np.repeat(_parameter, len(_coeffs))
-            # Multiply by coeffs
-            ucc_parameter *= _coeffs
-            all_parameters.append(ucc_parameter)
-
-        # Flatten all_parameters into a single list to set the circuit parameters
-        all_parameters = [_x for _param in all_parameters for _x in _param]
-        circuit.set_parameters(all_parameters)
-
-        return expectation(circuit, hamiltonian)
-
-    # Random initialization
-    params = np.zeros(n_excitations)
-    vqe = minimize(electronic_energy, params)
-
-    lih_uccsd_energy = -7.847535097575567
-
-    assert vqe.fun == pytest.approx(lih_uccsd_energy)

From ae95e7d711e167ff3eb1db45cf3a82e1c62cd931 Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Wed, 3 Jan 2024 07:18:12 +0000
Subject: [PATCH 7/9] Add tests for ucc_ansatz function

---
 tests/test_ucc.py | 92 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 91 insertions(+), 1 deletion(-)

diff --git a/tests/test_ucc.py b/tests/test_ucc.py
index c767e7d..acf8f06 100644
--- a/tests/test_ucc.py
+++ b/tests/test_ucc.py
@@ -6,7 +6,7 @@
 from scipy.optimize import minimize
 
 from qibochem.ansatz.hf_reference import hf_circuit
-from qibochem.ansatz.ucc import (  # ucc_ansatz
+from qibochem.ansatz.ucc import (
     generate_excitations,
     mp2_amplitude,
     sort_excitations,
@@ -178,6 +178,23 @@ def test_ucc_ferm_qubit_map_error():
         ucc_circuit(2, [0, 1], ferm_qubit_map="Unknown")
 
 
+def test_ucc_parameter_coefficients():
+    """Coefficients used to multiply the parameters in the UCC circuit. Note: may change in future!"""
+    # UCC-JW singles
+    control_values = (-1.0, 1.0)
+    coeffs = []
+    circuit = ucc_circuit(2, [0, 1], coeffs=coeffs)
+    # Check that the signs of the coefficients have been saved
+    assert all(control == test for control, test in zip(control_values, coeffs))
+
+    # UCC-JW doubles
+    control_values = (-0.25, 0.25, 0.25, 0.25, -0.25, -0.25, -0.25, 0.25)
+    coeffs = []
+    circuit = ucc_circuit(4, [0, 1, 2, 3], coeffs=coeffs)
+    # Check that the signs of the coefficients have been saved
+    assert all(control == test for control, test in zip(control_values, coeffs))
+
+
 def test_ucc_ansatz_h2():
     """Test the default arguments of ucc_ansatz using H2"""
     mol = Molecule([("H", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 0.7))])
@@ -196,3 +213,76 @@ def test_ucc_ansatz_h2():
     )
     # Check that number of parametrised gates is the same
     assert len(control_circuit.get_parameters()) == len(test_circuit.get_parameters())
+
+
+def test_ucc_ansatz_embedding():
+    """Test the default arguments of ucc_ansatz using LiH with HF embedding applied, but without the HF circuit"""
+    mol = Molecule([("Li", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 1.4))])
+    mol.run_pyscf()
+    mol.hf_embedding(active=[1, 2, 5])
+
+    # Generate all possible excitations
+    excitations = []
+    for order in range(2, 0, -1):
+        # 2 electrons, 6 spin-orbitals
+        excitations += sort_excitations(generate_excitations(order, range(0, 2), range(2, 6)))
+    # Build control circuit
+    control_circuit = hf_circuit(6, 0)
+    for excitation in excitations:
+        control_circuit += ucc_circuit(6, excitation)
+
+    test_circuit = ucc_ansatz(mol, include_hf=False)
+
+    assert all(
+        control.name == test.name and control.target_qubits == test.target_qubits
+        for control, test in zip(list(control_circuit.queue), list(test_circuit.queue))
+    )
+    # Check that number of parametrised gates is the same
+    assert len(control_circuit.get_parameters()) == len(test_circuit.get_parameters())
+
+
+def test_ucc_ansatz_excitations():
+    """Test the `excitations` argument of ucc_ansatz"""
+    mol = Molecule([("Li", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 1.4))])
+    mol.run_pyscf()
+    mol.hf_embedding(active=[1, 2, 5])
+
+    # Generate all possible excitations
+    excitations = [[0, 1, 2, 3], [0, 1, 4, 5]]
+    # Build control circuit
+    control_circuit = hf_circuit(6, 2)
+    for excitation in excitations:
+        control_circuit += ucc_circuit(6, excitation)
+
+    test_circuit = ucc_ansatz(mol, excitations=excitations)
+
+    assert all(
+        control.name == test.name and control.target_qubits == test.target_qubits
+        for control, test in zip(list(control_circuit.queue), list(test_circuit.queue))
+    )
+    # Check that number of parametrised gates is the same
+    assert len(control_circuit.get_parameters()) == len(test_circuit.get_parameters())
+
+
+def test_ucc_ansatz_error_checks():
+    """Test the checks for input validity"""
+    mol = Molecule([("Li", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 1.4))])
+    # Define number of electrons and spin-obritals by hand
+    mol.nelec = (2, 2)
+    mol.nso = 12
+
+    # Excitation level check: Input is in ("S", "D", "T", "Q")
+    with pytest.raises(AssertionError):
+        ucc_ansatz(mol, "Z")
+
+    # Excitation level check: Excitation > doubles
+    with pytest.raises(NotImplementedError):
+        ucc_ansatz(mol, "T")
+
+    # Excitations list check: excitation must have an even number of elements
+    with pytest.raises(AssertionError):
+        ucc_ansatz(mol, excitations=[[0]])
+
+    # Input parameter check: Must have correct number of input parameters
+    with pytest.raises(AssertionError):
+        ucc_ansatz(mol, excitations=[[0, 1, 2, 3]], thetas=np.zeros(2))

From 0407d4e58e7bde300a730516a5c3f2cf9ef11538 Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Wed, 3 Jan 2024 08:06:31 +0000
Subject: [PATCH 8/9] Use MP2 amplitudes to initialse ucc_ansatz

---
 src/qibochem/ansatz/ucc.py |  8 ++++++--
 tests/test_ucc.py          | 17 ++++++++++++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/qibochem/ansatz/ucc.py b/src/qibochem/ansatz/ucc.py
index aca17d4..63f8d15 100644
--- a/src/qibochem/ansatz/ucc.py
+++ b/src/qibochem/ansatz/ucc.py
@@ -276,6 +276,7 @@ def ucc_ansatz(
     trotter_steps=1,
     ferm_qubit_map=None,
     include_hf=True,
+    use_mp2_guess=True,
 ):
     """
     Build a circuit corresponding to the UCC ansatz with multiple excitations for a given Molecule.
@@ -293,6 +294,7 @@ def ucc_ansatz(
             theta=theta/trotter_steps. Default is 1
         ferm_qubit_map: fermion-to-qubit transformation. Default: Jordan-Wigner ("jw")
         include_hf: Whether or not to start the circuit with a Hartree-Fock circuit. Default: ``True``
+        use_mp2_guesses: Whether or not to use MP2 amplitudes as the initial guess parameter. Default: ``True``
 
     Returns:
         circuit: Qibo Circuit corresponding to the UCC ansatz
@@ -326,12 +328,14 @@ def ucc_ansatz(
         assert all(len(_ex) % 2 == 0 for _ex in excitations), "Excitation with an odd number of elements found!"
 
     # Check if thetas argument given, define to be all zeros if not
-    # TODO: Unsure if want to use MP2 guess amplitudes for the doubles? Some say good, some say bad
     # Number of circuit parameters: S->2, D->8, (T/Q->32/128; Not sure?)
     n_parameters = 2 * len([_ex for _ex in excitations if len(_ex) == 2])  # Singles
     n_parameters += 8 * len([_ex for _ex in excitations if len(_ex) == 4])  # Doubles
     if thetas is None:
-        thetas = np.zeros(n_parameters)
+        if use_mp2_guess:
+            thetas = np.array([mp2_amplitude(excitation, molecule.eps, molecule.tei) for excitation in excitations])
+        else:
+            thetas = np.zeros(n_parameters)
     else:
         # Check that number of circuit variables (i.e. thetas) matches the number of circuit parameters
         assert len(thetas) == n_parameters, "Number of input parameters doesn't match the number of circuit parameters!"
diff --git a/tests/test_ucc.py b/tests/test_ucc.py
index acf8f06..9765ef7 100644
--- a/tests/test_ucc.py
+++ b/tests/test_ucc.py
@@ -202,7 +202,8 @@ def test_ucc_ansatz_h2():
 
     # Build control circuit
     control_circuit = hf_circuit(4, 2)
-    for excitation in ([0, 1, 2, 3], [0, 2], [1, 3]):
+    excitations = ([0, 1, 2, 3], [0, 2], [1, 3])
+    for excitation in excitations:
         control_circuit += ucc_circuit(4, excitation)
 
     test_circuit = ucc_ansatz(mol)
@@ -214,6 +215,20 @@ def test_ucc_ansatz_h2():
     # Check that number of parametrised gates is the same
     assert len(control_circuit.get_parameters()) == len(test_circuit.get_parameters())
 
+    # Then check that the circuit parameters are the MP2 guess parameters
+    # Get the MP2 amplitudes first, then expand the list based on the excitation type
+    mp2_guess_amplitudes = []
+    for excitation in excitations:
+        mp2_guess_amplitudes += [
+            mp2_amplitude(excitation, mol.eps, mol.tei) for _ in range(2 ** (2 * (len(excitation) // 2) - 1))
+        ]
+    mp2_guess_amplitudes = np.array(mp2_guess_amplitudes)
+    coeffs = np.array([-0.25, 0.25, 0.25, 0.25, -0.25, -0.25, -0.25, 0.25, 0.0, 0.0, 0.0, 0.0])
+    mp2_guess_amplitudes *= coeffs
+    # Need to flatten the output of circuit.get_parameters() to compare it to mp2_guess_amplitudes
+    test_parameters = np.array([_x for _tuple in test_circuit.get_parameters() for _x in _tuple])
+    assert np.allclose(mp2_guess_amplitudes, test_parameters)
+
 
 def test_ucc_ansatz_embedding():
     """Test the default arguments of ucc_ansatz using LiH with HF embedding applied, but without the HF circuit"""

From 9c5003e39d34136b14ac81a62ae1a1dd972c620d Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Wed, 3 Jan 2024 08:43:06 +0000
Subject: [PATCH 9/9] Documentation for ucc_ansatz

Also added remaining test coverage
---
 doc/source/api-reference/ansatz.rst       |  1 +
 src/qibochem/ansatz/basis_rotation.py     |  2 +-
 src/qibochem/ansatz/hardware_efficient.py |  2 +-
 src/qibochem/ansatz/hf_reference.py       |  2 +-
 src/qibochem/ansatz/ucc.py                | 30 +++++++++++------------
 tests/test_ucc.py                         |  6 ++++-
 6 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/doc/source/api-reference/ansatz.rst b/doc/source/api-reference/ansatz.rst
index 48ecacd..fb931f8 100644
--- a/doc/source/api-reference/ansatz.rst
+++ b/doc/source/api-reference/ansatz.rst
@@ -21,6 +21,7 @@ Unitary Coupled Cluster
 
 .. autofunction:: qibochem.ansatz.ucc.ucc_circuit
 
+.. autofunction:: qibochem.ansatz.ucc.ucc_ansatz
 
 Basis rotation
 --------------
diff --git a/src/qibochem/ansatz/basis_rotation.py b/src/qibochem/ansatz/basis_rotation.py
index 887480f..c5a540a 100644
--- a/src/qibochem/ansatz/basis_rotation.py
+++ b/src/qibochem/ansatz/basis_rotation.py
@@ -129,7 +129,7 @@ def br_circuit(n_qubits, parameters, n_occ):
         n_occ: Number of occupied orbitals
 
     Returns:
-        Qibo ``Circuit`` corresponding to the basis rotation ansatz between the occupied and virtual orbitals
+        Qibo ``Circuit``: Circuit corresponding to the basis rotation ansatz between the occupied and virtual orbitals
     """
     assert len(parameters) == (n_occ * (n_qubits - n_occ)), "Need len(parameters) == (n_occ * n_virt)"
     # Unitary rotation matrix \exp(\kappa)
diff --git a/src/qibochem/ansatz/hardware_efficient.py b/src/qibochem/ansatz/hardware_efficient.py
index 5b13e17..89a5fef 100644
--- a/src/qibochem/ansatz/hardware_efficient.py
+++ b/src/qibochem/ansatz/hardware_efficient.py
@@ -15,7 +15,7 @@ def hea(n_layers, n_qubits, parameter_gates=["RY", "RZ"], coupling_gates="CZ"):
             valid two-qubit ``Qibo`` gate. Default: ``"CZ"``
 
     Returns:
-        List of gates corresponding to the hardware-efficient ansatz
+        ``list``: Gates corresponding to the hardware-efficient ansatz
     """
     gate_list = []
 
diff --git a/src/qibochem/ansatz/hf_reference.py b/src/qibochem/ansatz/hf_reference.py
index 5683a50..db64f6a 100644
--- a/src/qibochem/ansatz/hf_reference.py
+++ b/src/qibochem/ansatz/hf_reference.py
@@ -59,7 +59,7 @@ def hf_circuit(n_qubits, n_electrons, ferm_qubit_map=None):
         ferm_qubit_map: Fermion to qubit map. Must be either Jordan-Wigner (``jw``) or Brayvi-Kitaev (``bk``). Default value is ``jw``.
 
     Returns:
-        Qibo ``Circuit`` initialized in a HF reference state
+        Qibo ``Circuit``: Circuit initialized in a HF reference state
     """
     # Which fermion-to-qubit map to use
     if ferm_qubit_map is None:
diff --git a/src/qibochem/ansatz/ucc.py b/src/qibochem/ansatz/ucc.py
index 63f8d15..a234c1b 100644
--- a/src/qibochem/ansatz/ucc.py
+++ b/src/qibochem/ansatz/ucc.py
@@ -68,13 +68,13 @@ def ucc_circuit(n_qubits, excitation, theta=0.0, trotter_steps=1, ferm_qubit_map
         excitation: Iterable of orbitals involved in the excitation; must have an even number of elements
             E.g. ``[0, 1, 2, 3]`` represents the excitation of electrons in orbitals ``(0, 1)`` to ``(2, 3)``
         theta: UCC parameter. Defaults to 0.0
-        trotter_steps: Number of Trotter steps; i.e. number of times the UCC ansatz is applied with \theta = \theta / trotter_steps
+        trotter_steps: Number of Trotter steps; i.e. number of times the UCC ansatz is applied with ``theta`` = ``theta / trotter_steps``. Default: 1
         ferm_qubit_map: Fermion-to-qubit transformation. Default is Jordan-Wigner (``jw``).
         coeffs: List to hold the coefficients for the rotation parameter in each Pauli string.
             May be useful in running the VQE. WARNING: Will be modified in this function
 
     Returns:
-        Qibo ``Circuit`` corresponding to a single UCC excitation
+        Qibo ``Circuit``: Circuit corresponding to a single UCC excitation
     """
     # Check size of orbitals input
     n_orbitals = len(excitation)
@@ -279,25 +279,25 @@ def ucc_ansatz(
     use_mp2_guess=True,
 ):
     """
-    Build a circuit corresponding to the UCC ansatz with multiple excitations for a given Molecule.
-        If no excitations are given, it defaults to returning the full UCCSD circuit ansatz for the
-        Molecule.
+    Convenience function for buildng a circuit corresponding to the UCC ansatz with multiple excitations for a given ``Molecule``.
+    If no excitations are given, it defaults to returning the full UCCSD circuit ansatz.
 
     Args:
-        molecule: Molecule of interest.
-        excitation_level: Include excitations up to how many electrons, i.e. ("S", "D", "T", "Q")
-            Ignored if `excitations` argument is given. Default is "D", i.e. double excitations
-        excitations: List of excitations (e.g. `[[0, 1, 2, 3], [0, 1, 4, 5]]`) used to build the
-            UCC circuit. Overrides the `excitation_level` argument
-        thetas: Parameters for the excitations. Defaults to an array of zeros if not given
+        molecule: The ``Molecule`` of interest.
+        excitation_level: Include excitations up to how many electrons, i.e. ``"S"`` or ``"D"``.
+            Ignored if ``excitations`` argument is given. Default: ``"D"``, i.e. double excitations
+        excitations: List of excitations (e.g. ``[[0, 1, 2, 3], [0, 1, 4, 5]]``) used to build the
+            UCC circuit. Overrides the ``excitation_level`` argument
+        thetas: Parameters for the excitations. Default value depends on the ``use_mp2_guess`` argument.
         trotter_steps: number of Trotter steps; i.e. number of times the UCC ansatz is applied with
-            theta=theta/trotter_steps. Default is 1
-        ferm_qubit_map: fermion-to-qubit transformation. Default: Jordan-Wigner ("jw")
+            ``theta`` = ``theta / trotter_steps``. Default: 1
+        ferm_qubit_map: fermion-to-qubit transformation. Default: Jordan-Wigner (``"jw"``)
         include_hf: Whether or not to start the circuit with a Hartree-Fock circuit. Default: ``True``
-        use_mp2_guesses: Whether or not to use MP2 amplitudes as the initial guess parameter. Default: ``True``
+        use_mp2_guess: Whether to use MP2 amplitudes or a numpy zero array as the initial guess parameter. Default: ``True``;
+            use the MP2 amplitudes as the default guess parameters
 
     Returns:
-        circuit: Qibo Circuit corresponding to the UCC ansatz
+        Qibo ``Circuit``: Circuit corresponding to an UCC ansatz
     """
     # Get the number of electrons and spin-orbitals from the molecule argument
     n_elec = sum(molecule.nelec) if molecule.n_active_e is None else molecule.n_active_e
diff --git a/tests/test_ucc.py b/tests/test_ucc.py
index 9765ef7..9a5d688 100644
--- a/tests/test_ucc.py
+++ b/tests/test_ucc.py
@@ -246,7 +246,7 @@ def test_ucc_ansatz_embedding():
     for excitation in excitations:
         control_circuit += ucc_circuit(6, excitation)
 
-    test_circuit = ucc_ansatz(mol, include_hf=False)
+    test_circuit = ucc_ansatz(mol, include_hf=False, use_mp2_guess=False)
 
     assert all(
         control.name == test.name and control.target_qubits == test.target_qubits
@@ -255,6 +255,10 @@ def test_ucc_ansatz_embedding():
     # Check that number of parametrised gates is the same
     assert len(control_circuit.get_parameters()) == len(test_circuit.get_parameters())
 
+    # Check that the circuit parameters are all zeros
+    test_parameters = np.array([_x for _tuple in test_circuit.get_parameters() for _x in _tuple])
+    assert np.allclose(test_parameters, np.zeros(len(test_parameters)))
+
 
 def test_ucc_ansatz_excitations():
     """Test the `excitations` argument of ucc_ansatz"""