jpmorganchase · bgupt · Jul 1, 2025 · Jul 1, 2025 · Jul 1, 2025
diff --git a/qokit/labs.py b/qokit/labs.py
@@ -89,7 +89,7 @@
 
 def get_terms_offset(N: int):
     """
-    Retrun `terms` and `offset` for QAOA LABS problem
+    Retrun `offset` for QAOA LABS problem
 
     Parameters
     ----------
@@ -98,23 +98,17 @@ def get_terms_offset(N: int):
 
     Returns
     -------
-    terms : Sequence of `(float, term)`, where `term` is a tuple of ints.
-        Each term corresponds to a summand in the cost Hamiltonian
-        and th float value is the coefficient of this term.
-        e.g. if terms = [(0.5, (0,1)), (0.3, (0,1,2,3))]
-        the Hamiltonian is 0.5*Z0Z1 + 0.3*Z0Z1Z2Z3
 
     offset : int
         energy offset required due to constant factors (identity terms)
         not included in the Hamiltonian
 
     """
-    term_indices, offset = get_energy_term_indices(N)
-    terms = [(len(t), t) for t in term_indices]
-    return terms, offset
+    terms, offset = get_energy_term_indices_offset(N)
+    return offset
 
 
-def get_energy_term_indices(N: int):
+def get_term_indices(N: int) -> list:
     """Return indices of Pauli Zs in the LABS problem definition
 
     Parameters
@@ -125,10 +119,53 @@ def get_energy_term_indices(N: int):
     Returns
     -------
     terms : list of tuples
-        List of tuples, where each tuple defines a summand
+        List of ordered tuples, where each tuple defines a summand
         and contains indices of the Pauli Zs in the product
         e.g. if terms = [(0,1), (0,1,2,3), (1,2)]
         the Hamiltonian is Z0Z1 + Z0Z1Z2Z3 + Z1Z2
+    """
+    terms_with_weight, offset = get_energy_term_indices_offset(N)
+    terms = [term[1] for term in terms_with_weight]
+    return terms
+
+
+def get_terms(N: int) -> TermsType:
+    """Return terms definition of the LABS problem
+
+    Parameters
+    ----------
+    N : int
+        Problem size (number of spins)
+
+    Returns
+    -------
+    terms : TermsType
+        List of tuples (number, tuple) where the
+        tuple determines the location of Z operators
+        and the number is a scaling factor for the product.
+
+        e.g. if terms = [(2, (0,1)), (4, (0,1,2,3)), (2, (1,2))]
+        the Hamiltonian is 2*Z0Z1 + 4*Z0Z1Z2Z3 + 2*Z1Z2
+    """
+    terms, offset = get_energy_term_indices_offset(N)
+    return terms
+
+
+def get_energy_term_indices_offset(N: int):
+    """Return terms with indices and offset of Pauli Zs in the LABS problem definition
+
+    Parameters
+    ----------
+    N : int
+        Problem size (number of spins)
+
+    Returns
+    -------
+    terms : list of tuples
+        List of tuples, where each tuple defines a summand
+        and contains weight and indices of the Pauli Zs in the product
+        e.g. if terms = [(2, (0,1)), (4, (0,1,2,3)), (2, (1,2))]
+        the Hamiltonian is Z0Z1 + Z0Z1Z2Z3 + Z1Z2
 
     offset : int
         energy offset required due to constant factors (identity terms)
@@ -144,10 +181,10 @@ def get_energy_term_indices(N: int):
             # -1 because we index from 1
             # Drop duplicate terms, e.g. Z1Z2Z2Z3 should be just Z1Z3
             if i + k == j:
-                all_terms.append(tuple(sorted((i - 1, j + k - 1))))
+                all_terms.append((2, tuple(sorted((i - 1, j + k - 1)))))
             else:
-                all_terms.append(tuple(sorted((i - 1, i + k - 1, j - 1, j + k - 1))))
-    return set(all_terms), offset
+                all_terms.append((4, tuple(sorted((i - 1, i + k - 1, j - 1, j + k - 1)))))
+    return list(set(all_terms)), offset
 
 
 @njit()
@@ -208,8 +245,8 @@ def energy_vals_general(s: Sequence, terms: Iterable | None = None, offset: floa
         set(s) \in {+1, -1}
     terms : list-like, default None
     offset : float, default None
-        precomputed output of get_energy_term_indices
-        terms, offset = get_energy_term_indices(N)
+        precomputed output of get_energy_term_indices_offset
+        terms, offset = get_energy_term_indices_offset(N)
     check_parameters : bool, default True
         if set to False, no input validation is performed
     Returns
@@ -223,10 +260,11 @@ def energy_vals_general(s: Sequence, terms: Iterable | None = None, offset: floa
         assert set(s).issubset(set([-1, 1]))
     N = len(s)
     if terms is None or offset is None:
-        terms, offset = get_energy_term_indices(N)
+        terms, offset = get_energy_term_indices_offset(N)
     E_s = offset
     for term in terms:
-        E_s += (4 if len(term) == 4 else 2) * reduce(mul, [s[idx] for idx in term])
+        len_term, val_term = term
+        E_s += (4 if len(val_term) == 4 else 2) * reduce(mul, [s[idx] for idx in val_term])
     return E_s
 
 
@@ -240,8 +278,8 @@ def energy_vals_from_bitstring_general(x, terms: Sequence | None = None, offset:
         set(s) \in {0, 1}
     terms : list-like, default None
     offset : float, default None
-        precomputed output of get_energy_term_indices
-        terms, offset = get_energy_term_indices(N)
+        precomputed output of get_energy_term_indices_offset
+        terms, offset = get_energy_term_indices_offset(N)
     check_parameters : bool, default True
         if set to False, no input validation is performed
     Returns
@@ -262,8 +300,8 @@ def slow_merit_factor(s: Sequence, terms: Iterable | None = None, offset: float
         set(s) \in {-1, +1}
     terms : list-like, default None
     offset : float, default None
-        precomputed output of get_energy_term_indices
-        terms, offset = get_energy_term_indices(N)
+        precomputed output of get_energy_term_indices_offset
+        terms, offset = get_energy_term_indices_offset(N)
     check_parameters : bool, default True
         if set to False, no input validation is performed
 
@@ -278,10 +316,11 @@ def slow_merit_factor(s: Sequence, terms: Iterable | None = None, offset: float
         assert set(s).issubset(set([-1, 1]))
     N = len(s)
     if terms is None or offset is None:
-        terms, offset = get_energy_term_indices(N)
+        terms, offset = get_energy_term_indices_offset(N)
     E_s = offset
     for term in terms:
-        E_s += (4 if len(term) == 4 else 2) * reduce(mul, [s[idx] for idx in term])
+        len_term, val_term = term
+        E_s += len_term * reduce(mul, [s[idx] for idx in val_term])
     return N**2 / (2 * E_s)
 
 
@@ -328,46 +367,6 @@ def negative_merit_factor_from_bitstring(x, N: int | None = None) -> float:
     return -merit_factor(1 - 2 * x, N=N)
 
 
-def get_term_indices(N: int) -> list:
-    """Return indices of Pauli Zs in the LABS problem definition
-
-    Parameters
-    ----------
-    N : int
-        Problem size (number of spins)
-
-    Returns
-    -------
-    terms : list of tuples
-        List of ordered tuples, where each tuple defines a summand
-        and contains indices of the Pauli Zs in the product
-        e.g. if terms = [(0,1), (0,1,2,3), (1,2)]
-        the Hamiltonian is Z0Z1 + Z0Z1Z2Z3 + Z1Z2
-    """
-    return list(get_energy_term_indices(N)[0])
-
-
-    def get_terms(N: int) -> TermsType:
-    """Return terms definition of the LABS problem
-
-    Parameters
-    ----------
-    N : int
-        Problem size (number of spins)
-
-    Returns
-    -------
-    terms : TermsType
-        List of tuples (number, tuple) where the
-        tuple determines the location of Z operators
-        and the number is a scaling factor for the product.
-
-        e.g. if terms = [(2, (0,1)), (4, (0,1,2,3)), (2, (1,2))]
-        the Hamiltonian is 2*Z0Z1 + 4*Z0Z1Z2Z3 + 2*Z1Z2
-    """
-    return [(len(x), x) for x in get_term_indices(N)]
-
-
 def get_depth_optimized_terms(N: int) -> list:
     """Return indices of Pauli Zs in the LABS problem definition. The terms in the returned list are ordered to attempt to compress
     the circuit depth and increase parallelism.

diff --git a/qokit/qaoa_circuit.py b/qokit/qaoa_circuit.py
@@ -27,10 +27,9 @@ def append_z_prod_term(qc: QuantumCircuit, term: Sequence, gamma: float) -> None
         # in labs, four-body terms appear two times more than two-body
         # there is also a global scaling factor of 2 for all terms (four and two), which is ignored here
         assert all(term[i] < term[i + 1] for i in range(len(term) - 1))
-        _gamma = 2 * gamma
         qc.cx(term[0], term[1])
         qc.cx(term[3], term[2])
-        qc.rzz(2 * _gamma, term[1], term[2])
+        qc.rzz(2 * gamma, term[1], term[2])
         qc.cx(term[3], term[2])
         qc.cx(term[0], term[1])
     elif term_weight == 2:
@@ -52,8 +51,8 @@ def append_x_term(qc: QuantumCircuit, q1: int, beta: float) -> None:
 def append_cost_operator_circuit(qc: QuantumCircuit, terms: Sequence, gamma: float) -> None:
     for term in terms:
         if len(term) == 2 and isinstance(term[1], tuple):
-            coeff, (i, j) = term
-            append_z_prod_term(qc, (i, j), gamma * coeff / 2)
+            coeff, term_tuple = term
+            append_z_prod_term(qc, term_tuple, gamma * coeff / 2)
         elif any([isinstance(i, tuple) for i in term]):
             raise ValueError(f"Invalid term received: {term}")
         else:

diff --git a/qokit/qaoa_circuit_labs.py b/qokit/qaoa_circuit_labs.py
@@ -5,7 +5,7 @@
 # QAOA circuit for some Z objective
 from collections.abc import Sequence
 from qiskit import QuantumCircuit
-from qokit.labs import get_energy_term_indices
+from qokit.labs import get_energy_term_indices_offset
 from .qaoa_circuit import get_qaoa_circuit_from_terms, get_parameterized_qaoa_circuit_from_terms
 
 
@@ -27,7 +27,7 @@ def get_qaoa_circuit(N: int, gammas: Sequence, betas: Sequence, save_statevector
     qc : qiskit.QuantumCircuit
         Quantum circuit implementing QAOA
     """
-    terms, _ = get_energy_term_indices(N)
+    terms, _ = get_energy_term_indices_offset(N)
     return get_qaoa_circuit_from_terms(N=N, terms=terms, gammas=gammas, betas=betas, save_statevector=save_statevector)
 
 
@@ -67,7 +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)
+    terms, _ = get_energy_term_indices_offset(N)
     return get_parameterized_qaoa_circuit_from_terms(
         N=N,
         terms=terms,

diff --git a/qokit/qaoa_objective_labs.py b/qokit/qaoa_objective_labs.py
@@ -9,7 +9,7 @@
 from pathlib import Path
 
 from .labs import (
-    get_energy_term_indices,
+    get_energy_term_indices_offset,
     negative_merit_factor_from_bitstring,
     true_optimal_energy,
     energy_vals_from_bitstring,
@@ -183,7 +183,7 @@ def get_qaoa_labs_objective(
     # TODO: needs to generate parameterized circuit and check that the precomputed stuff is loaded correctly
     # Otherwise pass directly to get_qaoa_objective
 
-    terms_ix, offset = get_energy_term_indices(N)
+    terms_ix, offset = get_energy_term_indices_offset(N)
 
     if precomputed_negative_merit_factors is None:
         precomputed_negative_merit_factors = get_precomputed_labs_merit_factors(N)

diff --git a/tests/test_fast_simulators_labs.py b/tests/test_fast_simulators_labs.py
@@ -13,7 +13,7 @@
 from qokit.fur.lazy_import import MPI
 from qokit.fur import QAOAFURXSimulatorGPUMPI, get_available_simulators
 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.labs import energy_vals_from_bitstring_general, get_energy_term_indices_offset
 from qokit.utils import precompute_energies
 from qokit.qaoa_circuit import get_qaoa_circuit_from_terms as get_qaoa_circuit
 
@@ -86,7 +86,7 @@ def test_phase_operator(N, p, simclass):
     terms = [(3, [0, 1, 2]), (1, [0])]
     terms_without_weights = [t[1] for t in terms]
     offset = 2
-    f = partial(energy_vals_from_bitstring_general, terms=terms_without_weights, offset=offset)
+    f = partial(energy_vals_from_bitstring_general, terms=terms, offset=offset)
     precomputed_energies = precompute_energies(f, N) - offset
 
     sim = simclass(N, costs=precomputed_energies)
@@ -101,7 +101,7 @@ def test_phase_operator_labs(N, p, simclass):
     """
     gammas = np.random.uniform(0, np.pi, p)
     betas = [0] * p
-    terms_ix, offset = get_energy_term_indices(N)
+    terms_ix, offset = get_energy_term_indices_offset(N)
     precomputed_energies = -(N**2) / (2 * get_precomputed_labs_merit_factors(N)) - offset
 
     sim = simclass(N, costs=precomputed_energies)
@@ -114,7 +114,7 @@ def test_against_qiskit_p_1(simclass):
     gamma = np.random.uniform(0, np.pi)
 
     N = 5
-    terms_ix, offset = get_energy_term_indices(N)
+    terms_ix, offset = get_energy_term_indices_offset(N)
     precomputed_energies = -(N**2) / (2 * get_precomputed_labs_merit_factors(N)) - offset
 
     sim = simclass(N, costs=precomputed_energies)
@@ -125,7 +125,7 @@ def test_against_qiskit_p_1(simclass):
 def test_against_qiskit_p_3(simclass):
     p = 3
     N = 5
-    terms, offset = get_energy_term_indices(N)
+    terms, offset = get_energy_term_indices_offset(N)
     precomputed_energies = -(N**2) / (2 * get_precomputed_labs_merit_factors(N)) - offset
 
     sim = simclass(N, costs=precomputed_energies)
@@ -141,8 +141,7 @@ def test_against_qiskit_p_3(simclass):
 def test_terms_api(simclass: typing.Type[QAOAFastSimulatorBase], N):
     """Test that simulator can precompute energies when given wegihted Terms"""
     p = 1
-    terms, offset = get_terms_offset(N)
-    terms_ix, _ = get_energy_term_indices(N)
+    terms, offset = get_energy_term_indices_offset(N)
 
     sim = simclass(N, terms=terms)
     precomputed_energies = -(N**2) / (2 * get_precomputed_labs_merit_factors(N)) - offset
@@ -152,4 +151,4 @@ def test_terms_api(simclass: typing.Type[QAOAFastSimulatorBase], N):
         betas = np.random.uniform(0, np.pi, p)
         gammas = np.random.uniform(0, np.pi, p)
 
-        _check_simulator_against_qiskit(sim, N, terms_ix, gammas, betas)
+        _check_simulator_against_qiskit(sim, N, terms, gammas, betas)
diff --git a/tests/test_labs.py b/tests/test_labs.py
@@ -15,14 +15,14 @@
     get_gate_optimized_terms_greedy,
     true_optimal_energy,
     energy_vals,
-    get_energy_term_indices,
+    get_energy_term_indices_offset,
     energy_vals_general,
 )
 
 
 def test_slow_merit_factor():
     for N in range(3, 12):
-        terms, offset = get_energy_term_indices(N)
+        terms, offset = get_energy_term_indices_offset(N)
         s = np.random.choice([-1, 1], size=N)
         assert np.isclose(
             slow_merit_factor(s, terms=terms, offset=offset, check_parameters=True),
@@ -61,10 +61,12 @@ def test_energy_vals():
     ):
         mf_from_en = N**2 / (2 * true_optimal_energy[N])
         assert np.around(mf_from_en, decimals=3) == true_optimal_mf[N]
-    for N in range(3, 12):
-        terms, offset = get_energy_term_indices(N)
+    for N in range(5, 12):
+        terms, offset = get_energy_term_indices_offset(N)
         s = np.random.choice([-1, 1], size=N)
+        _energy_vals_general = energy_vals_general(s, terms=terms, offset=offset, check_parameters=True)
+        _energy_vals = energy_vals(s, N=N)
         assert np.isclose(
-            energy_vals_general(s, terms=terms, offset=offset, check_parameters=True),
-            energy_vals(s, N=N),
+            _energy_vals_general,
+            _energy_vals,
         )
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -3,14 +3,14 @@
 # // Copyright : JP Morgan Chase & Co
 ###############################################################################
 import numpy as np
-from qokit.labs import get_energy_term_indices, negative_merit_factor_from_bitstring
+from qokit.labs import get_energy_term_indices_offset, negative_merit_factor_from_bitstring
 from qokit.utils import precompute_energies_parallel, precompute_energies
 from qokit.utils import reverse_array_index_bit_order
 
 
 def test_precompute_energies():
     N = 5
-    terms, offset = get_energy_term_indices(N)
+    terms, offset = get_energy_term_indices_offset(N)
     ens1 = precompute_energies(negative_merit_factor_from_bitstring, N, N)
     ens2 = precompute_energies_parallel(negative_merit_factor_from_bitstring, N, 2)