Merge pull request #105 from jpmorganchase/fix_labs_circuit_gen

address #104

Author: bgupt

qokit/qaoa_circuit_labs.py

@@ -5,20 +5,17 @@
 # QAOA circuit for some Z objective
 from collections.abc import Sequence
 from qiskit import QuantumCircuit
+from qokit.labs import get_energy_term_indices
 from .qaoa_circuit import get_qaoa_circuit_from_terms, get_parameterized_qaoa_circuit_from_terms
 
 
-def get_qaoa_circuit(N: int, terms: Sequence, gammas: Sequence, betas: Sequence, save_statevector: bool = True) -> QuantumCircuit:
+def get_qaoa_circuit(N: int, gammas: Sequence, betas: Sequence, save_statevector: bool = True) -> QuantumCircuit:
     """Generates a circuit for Hamiltonian of the form \sum_{term \in terms} \prod_{j \in term} Z_j
 
     Parameters
     ----------
     N : int
         Number of qubits
-    terms : list of tuples
-        Each tuple corresponds to a term \prod_{j \in term} Z_j and contains indices
-        Example: for H = Z_0*Z_1 + Z_2*Z_3 + Z_0*Z_2*Z_4, terms = [(0,1), (2,3), (0,2,4)]
-        All indices must be less than N
     beta : list-like
         QAOA parameter beta
     gamma : list-like
@@ -30,10 +27,11 @@ def get_qaoa_circuit(N: int, terms: Sequence, gammas: Sequence, betas: Sequence,
     qc : qiskit.QuantumCircuit
         Quantum circuit implementing QAOA
     """
+    terms, _ = get_energy_term_indices(N)
     return get_qaoa_circuit_from_terms(N=N, terms=terms, gammas=gammas, betas=betas, save_statevector=save_statevector)
 
 
-def get_parameterized_qaoa_circuit(N: int, terms: Sequence, p: int, save_statevector: bool = True, return_parameter_vectors: bool = False) -> QuantumCircuit:
+def get_parameterized_qaoa_circuit(N: int, p: int, save_statevector: bool = True, return_parameter_vectors: bool = False) -> QuantumCircuit:
     """Generates a parameterized circuit for Hamiltonian of the form \sum_{term \in terms} \prod_{j \in term} Z_j
     This version is recommended for long circuits
 
@@ -69,6 +67,7 @@ def f(theta):
         (beta first, then gamma). To bind:
         qc.bind_parameters(np.hstack([angles['beta'], angles['gamma']]))
     """
+    terms, _ = get_energy_term_indices(N)
     return get_parameterized_qaoa_circuit_from_terms(
         N=N,
         terms=terms,

qokit/qaoa_objective_labs.py

@@ -190,8 +190,7 @@ def get_qaoa_labs_objective(
 
     if simulator == "qiskit":
         assert p is not None, "p must be passed if simulator == 'qiskit'"
-        terms, _ = get_energy_term_indices(N)
-        parameterized_circuit = get_parameterized_qaoa_circuit(N, terms, p)
+        parameterized_circuit = get_parameterized_qaoa_circuit(N, p)
         precomputed_diagonal_hamiltonian = None
     else:
         parameterized_circuit = None

tests/test_fast_simulators_labs.py

@@ -15,7 +15,7 @@
 from qokit.fur import QAOAFURXSimulatorC, QAOAFURXSimulator, QAOAFURXSimulatorGPU, QAOAFastSimulatorBase
 from qokit.labs import energy_vals_from_bitstring_general, get_energy_term_indices, get_terms_offset
 from qokit.utils import precompute_energies
-from qokit.qaoa_circuit_labs import get_qaoa_circuit
+from qokit.qaoa_circuit import get_qaoa_circuit_from_terms as get_qaoa_circuit
 
 SIMULATORS = get_available_simulators("x")
 print(SIMULATORS)

tests/test_qaoa_qiskit.py

@@ -11,10 +11,7 @@
     get_parameterized_qaoa_circuit,
     get_qaoa_circuit,
 )
-from qokit.labs import (
-    get_energy_term_indices,
-    negative_merit_factor_from_bitstring,
-)
+from qokit.labs import negative_merit_factor_from_bitstring
 from qokit.utils import (
     precompute_energies,
     obj_from_statevector,
@@ -36,9 +33,7 @@ def test_qaoa_parameterization():
     )
     row = df[(df["N"] == N) & (df["p"] == p)].squeeze()
 
-    terms_ix, offset = get_energy_term_indices(N)
-
-    qc = get_qaoa_circuit(N, terms_ix, row["gamma"], row["beta"])
+    qc = get_qaoa_circuit(N, row["gamma"], row["beta"])
     backend = AerSimulator(method="statevector")
     sv = np.asarray(backend.run(qc).result().get_statevector())
 
@@ -68,18 +63,17 @@ def test_parameterized_circuit():
     N = 10
     p = 50
     ramp = get_ramp(0.1663, p)
-    terms, offset = get_energy_term_indices(N)
 
     backend = AerSimulator(method="statevector")
 
-    qc_param = get_parameterized_qaoa_circuit(N, terms, p)
+    qc_param = get_parameterized_qaoa_circuit(N, p)
     qc1 = qc_param.assign_parameters(np.hstack([ramp["beta"], ramp["gamma"]]))
     f1 = obj_from_statevector(
         np.asarray(backend.run(qc1).result().get_statevector()),
         negative_merit_factor_from_bitstring,
     )
 
-    qc2 = get_qaoa_circuit(N, terms, ramp["gamma"], ramp["beta"])
+    qc2 = get_qaoa_circuit(N, ramp["gamma"], ramp["beta"])
     f2 = obj_from_statevector(
         np.asarray(backend.run(qc2).result().get_statevector()),
         negative_merit_factor_from_bitstring,