From 5a9acd4e2e836cca6e1b4e3183b521eb3396beb7 Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Tue, 20 Aug 2024 07:21:00 +0000
Subject: [PATCH] Update code for single excitation

---
 src/qibochem/ansatz/givens_excitation.py | 29 +++++++++++++++++++++---
 tests/test_givens_excitation.py          | 28 ++++++++++++++++++++---
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/src/qibochem/ansatz/givens_excitation.py b/src/qibochem/ansatz/givens_excitation.py
index 3d8c3efc..e891d15b 100644
--- a/src/qibochem/ansatz/givens_excitation.py
+++ b/src/qibochem/ansatz/givens_excitation.py
@@ -1,5 +1,6 @@
 """
-Module documentation
+Circuit ansatz for representing a fermionic excitation as a Givens rotation by Arrazola et al.
+Reference: https://doi.org/10.22331/q-2022-06-20-742
 """
 
 from qibo import Circuit, gates
@@ -10,6 +11,28 @@
 # Helper functions
 
 
+def single_excitation_gate(sorted_orbitals, theta):
+    """
+    Decomposition of a Givens single excitation gate into single qubit rotations and CNOTs. In principle, should be
+    identical to gates.GIVENS(qubit0, qubit1, theta)
+
+    Args:
+        sorted_orbitals (list): Sorted list of orbitals involved in the excitation
+        theta (float): Rotation angle
+
+    Returns:
+        (list): List of gates representing the decomposition of the Givens' single excitation gate
+    """
+    result = []
+    result.append(gates.CNOT(sorted_orbitals[0], sorted_orbitals[1]))
+    result.append(gates.RY(sorted_orbitals[0], 0.5 * theta))
+    result.append(gates.CNOT(sorted_orbitals[1], sorted_orbitals[0]))
+    result.append(gates.RY(sorted_orbitals[0], -0.5 * theta))
+    result.append(gates.CNOT(sorted_orbitals[1], sorted_orbitals[0]))
+    result.append(gates.CNOT(sorted_orbitals[0], sorted_orbitals[1]))
+    return result
+
+
 def double_excitation_gate(sorted_orbitals, theta):
     """
     Decomposition of a Givens double excitation gate into single qubit rotations and CNOTs
@@ -56,7 +79,7 @@ def double_excitation_gate(sorted_orbitals, theta):
 # Main function
 def givens_excitation_circuit(n_qubits, excitation, theta=None):
     """
-    Circuit ansatz corresponding to the Givens rotation excitation from Arrazola et al. (https://doi.org/10.22331/q-2022-06-20-742) for a single excitation.
+    Circuit ansatz corresponding to the Givens rotation excitation from Arrazola et al.
 
     Args:
         n_qubits: Number of qubits in the circuit
@@ -76,7 +99,7 @@ def givens_excitation_circuit(n_qubits, excitation, theta=None):
 
     circuit = Circuit(n_qubits)
     if len(excitation) == 2:
-        circuit.add(gates.GIVENS(excitation[0], excitation[1], theta))
+        circuit.add(single_excitation_gate(sorted_orbitals, theta))
     elif len(excitation) == 4:
         circuit.add(double_excitation_gate(sorted_orbitals, theta))
     else:
diff --git a/tests/test_givens_excitation.py b/tests/test_givens_excitation.py
index a16ea21a..2344a202 100644
--- a/tests/test_givens_excitation.py
+++ b/tests/test_givens_excitation.py
@@ -11,11 +11,33 @@
     double_excitation_gate,
     givens_excitation_ansatz,
     givens_excitation_circuit,
+    single_excitation_gate,
 )
 from qibochem.ansatz.util import generate_excitations, mp2_amplitude, sort_excitations
 from qibochem.driver import Molecule
 
 
+def test_single_excitation_gate():
+    # Hardcoded test
+    theta = 0.1
+
+    control_gates = [
+        gates.CNOT(0, 1),
+        gates.RY(0, 0.5 * theta),
+        gates.CNOT(1, 0),
+        gates.RY(0, -0.5 * theta),
+        gates.CNOT(1, 0),
+        gates.CNOT(0, 1),
+    ]
+    test_list = single_excitation_gate([0, 1, 2, 3], 0.1)
+
+    # Check gates are correct
+    assert all(
+        control.name == test.name and control.target_qubits == test.target_qubits
+        for control, test in zip(control_gates, test_list)
+    )
+
+
 def test_double_excitation_gate():
     # Hardcoded test
     theta = 0.0
@@ -62,7 +84,7 @@ def test_double_excitation_gate():
 @pytest.mark.parametrize(
     "excitation,expected",
     [
-        ([0, 2], [gates.GIVENS(0, 2, 0.0)]),
+        ([0, 2], single_excitation_gate([0, 2], 0.0)),
         ([0, 1, 2, 3], double_excitation_gate([0, 1, 2, 3], 0.0)),
     ],
 )
@@ -106,8 +128,8 @@ def test_givens_excitation_ansatz_h2():
     # 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 = [mp2_amplitude([0, 1, 2, 3], mol.eps, mol.tei) for _ in range(8)]  # Doubles
-    mp2_guess_amplitudes += [0.0, 0.0]  # Singles
-    coeffs = np.array([-0.125, 0.125, -0.125, 0.125, 0.125, -0.125, 0.125, -0.125, 1.0, 1.0])
+    mp2_guess_amplitudes += [0.0, 0.0, 0.0, 0.0]  # Singles
+    coeffs = np.array([-0.125, 0.125, -0.125, 0.125, 0.125, -0.125, 0.125, -0.125, 1.0, 1.0, 1.0, 1.0])
     mp2_guess_amplitudes = coeffs * np.array(mp2_guess_amplitudes)
     # 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])