Merge pull request #11 from jpmorganchase/fix/9_objective_gpu_precompute

fix precomputation in get_objective_maxcut.py

qokit/fur/__init__.py

@@ -29,7 +29,18 @@
 }
 
 
-def get_available_simulator_names(type="x"):
+def get_available_simulator_names(type: str = "x") -> list:
+    """
+    Return names of available simulators
+
+    Parameters
+    ----------
+        type: type of QAOA mixer to simulate
+
+    Returns
+    -------
+        List of available simulators
+    """
     family = SIMULATORS.get(type, None)
     if family is None:
         raise ValueError(f"Unknown simulator type: {type}")
@@ -45,7 +56,18 @@ def get_available_simulator_names(type="x"):
     return available
 
 
-def get_available_simulators(type="x"):
+def get_available_simulators(type: str = "x") -> list:
+    """
+    Return (uninitialized) classes of available simulators
+
+    Parameters
+    ----------
+        type: type of QAOA mixer to simulate
+
+    Returns
+    -------
+        List of available simulators
+    """
     available_names = get_available_simulator_names(type=type)
     return [SIMULATORS[type][s] for s in available_names]
 

qokit/maxcut.py

@@ -5,6 +5,7 @@
 """
 Helper functions for the Maximum Cut (MaxCut) problem
 """
+from qokit.fur.qaoa_simulator_base import TermsType
 import numpy as np
 import networkx as nx
 
@@ -21,6 +22,30 @@ def maxcut_obj(x: np.ndarray, w: np.ndarray) -> float:
     return np.sum(w * X)  # type: ignore
 
 
+def get_maxcut_terms(G: nx.Graph) -> TermsType:
+    """Get terms corresponding to cost function value
+
+    .. math::
+
+        S = \\sum_{(i,j,w)\\in G} w*(1-s_i*s_j)/2
+
+    Args:
+        G: MaxCut problem graph
+    Returns:
+        terms to be used in the simulation
+    """
+    if nx.is_weighted(G):
+        terms = [(-float(G[u][v]["weight"]) / 2, (int(u), int(v))) for u, v, *_ in G.edges()]
+        total_w = sum([float(G[u][v]["weight"]) for u, v, *_ in G.edges()])
+
+    else:
+        terms = [(-1 / 2, (int(e[0]), int(e[1]))) for e in G.edges()]
+        total_w = int(G.number_of_edges())
+    N = G.number_of_nodes()
+    terms.append((+total_w / 2, tuple()))
+    return terms
+
+
 def get_adjacency_matrix(G: nx.Graph) -> np.ndarray:
     """Get adjacency matrix to be used in maxcut_obj
     Args:

qokit/qaoa_objective_maxcut.py

@@ -8,7 +8,7 @@
 import warnings
 
 from .utils import precompute_energies
-from .maxcut import maxcut_obj, get_adjacency_matrix
+from .maxcut import maxcut_obj, get_adjacency_matrix, get_maxcut_terms
 
 from .qaoa_circuit_maxcut import get_parameterized_qaoa_circuit
 from .qaoa_objective import get_qaoa_objective
@@ -53,13 +53,14 @@ def get_qaoa_maxcut_objective(
     f : callable
         Function returning the negative of expected value of QAOA with parameters theta
     """
+    terms = None
 
     if precomputed_cuts is not None and G is not None:
         warnings.warn("If precomputed_cuts is passed, G is ignored")
 
     if precomputed_cuts is None:
         assert G is not None, "G must be passed if precomputed_cuts is None"
-        precomputed_cuts = precompute_energies(maxcut_obj, N, w=get_adjacency_matrix(G))
+        terms = get_maxcut_terms(G)
 
     if simulator == "qiskit":
         assert G is not None, "G must be passed if simulator == 'qiskit'"
@@ -73,6 +74,7 @@ def get_qaoa_maxcut_objective(
         p=p,
         precomputed_diagonal_hamiltonian=precomputed_cuts,
         precomputed_objectives=precomputed_cuts,
+        terms=terms,
         precomputed_optimal_bitstrings=precomputed_optimal_bitstrings,
         parameterized_circuit=parameterized_circuit,
         parameterization=parameterization,

tests/test_maxcut.py

@@ -8,19 +8,23 @@
 from pathlib import Path
 from functools import partial
 from qiskit.providers.aer import AerSimulator
+import pytest
 
-from qokit.maxcut import maxcut_obj, get_adjacency_matrix
+from qokit.maxcut import maxcut_obj, get_adjacency_matrix, get_maxcut_terms
 
 from qokit.qaoa_objective_maxcut import get_qaoa_maxcut_objective
 
 from qokit.qaoa_circuit_maxcut import get_qaoa_circuit, get_parameterized_qaoa_circuit
 from qokit.utils import brute_force, precompute_energies
 from qokit.parameter_utils import get_sk_gamma_beta, get_fixed_gamma_beta
+import qokit
+from qokit.fur import get_available_simulators
 
 test_maxcut_folder = Path(__file__).parent
 
 
 qiskit_backend = AerSimulator(method="statevector")
+SIMULATORS = get_available_simulators("x") + get_available_simulators("xyring") + get_available_simulators("xycomplete")
 
 
 def test_maxcut_obj():
@@ -80,6 +84,20 @@ def test_maxcut_weighted_qaoa_obj():
         assert np.isclose(precomputed_cuts.dot(np.abs(sv) ** 2), row["Expected cut of QAOA"])
 
 
+@pytest.mark.parametrize("simclass", SIMULATORS)
+def test_maxcut_precompute(simclass):
+    N = 4
+    G = nx.random_regular_graph(3, N)
+    print(G.edges())
+    for u, v, w in G.edges(data=True):
+        w["weight"] = np.random.rand()
+    precomputed_cuts = precompute_energies(maxcut_obj, N, w=get_adjacency_matrix(G))
+    terms = get_maxcut_terms(G)
+    sim = simclass(N, terms=terms)
+    cuts = sim.get_cost_diagonal()
+    assert np.allclose(precomputed_cuts, cuts, atol=1e-6)
+
+
 def test_sk_ini_maxcut():
     N = 10
     for d, max_p in [(3, 5), (5, 5)]: