SK Model QAOA application benchmark (#290)

Author: mpharrigan

recirq/qaoa/__init__.py

@@ -11,3 +11,30 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from functools import lru_cache
+from typing import Optional
+
+from cirq.protocols.json_serialization import ObjectFactory, DEFAULT_RESOLVERS
+from .sk_model import (
+    SKModelQAOASpec,
+)
+
+
+@lru_cache()
+def _resolve_json(cirq_type: str) -> Optional[ObjectFactory]:
+    """Resolve the types of `recirq.qaoa.` json objects.
+
+    This is a Cirq JSON resolver suitable for appending to
+    `cirq.protocols.json_serialization.DEFAULT_RESOLVERS`.
+    """
+    if not cirq_type.startswith('recirq.qaoa.'):
+        return None
+
+    cirq_type = cirq_type[len('recirq.qaoa.'):]
+    return {k.__name__: k for k in [
+        SKModelQAOASpec,
+    ]}.get(cirq_type, None)
+
+
+DEFAULT_RESOLVERS.append(_resolve_json)

recirq/qaoa/classical_angle_optimization.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from timeit import default_timer as timer
+from typing import List
 
 import networkx as nx
 import numpy as np
@@ -53,7 +54,7 @@ def optimize_instance_interp_heuristic(graph: nx.Graph,
                                        param_guess_at_p1=None,
                                        node_to_index_map=None,
                                        dtype=np.complex128,
-                                       verbose=False):
+                                       verbose=False) -> List[OptimizationResult]:
     r"""
     Given a graph, find QAOA parameters that minimizes C=\sum_{<ij>} w_{ij} Z_i Z_j
 

recirq/qaoa/problem_circuits.py

@@ -14,7 +14,7 @@
     compile_to_non_negligible,
     validate_well_structured,
     compile_problem_unitary_to_swap_network, compile_swap_network_to_zzswap,
-    measure_with_final_permutation, compile_problem_unitary_to_arbitrary_zz)
+    measure_with_final_permutation, compile_problem_unitary_to_arbitrary_zz, ZZSwap)
 from recirq.qaoa.placement import place_on_device
 from recirq.qaoa.problems import HardwareGridProblem, SKProblem, ThreeRegularProblem
 
@@ -122,10 +122,17 @@ def get_routed_sk_model_circuit(
         qubits: List[cirq.Qid],
         gammas: Sequence[float],
         betas: Sequence[float],
+        *,
+        keep_zzswap_as_one_op=True,
 ) -> cirq.Circuit:
     """Get a QAOA circuit for a fully-connected problem using the linear swap
     network.
 
+    This leaves the circuit in an "uncompiled" form, using ZZ, Swap, and/or ZZSwap gates
+    for the two-qubit gate and will append a permutation gate for odd p depths. The former
+    should be compiled to hardware-native two-qubit gates; the latter can be absorbed into
+    analysis routines.
+
     See Also:
         :py:func:`get_compiled_sk_model_circuit`
 
@@ -135,11 +142,18 @@ def get_routed_sk_model_circuit(
         qubits: The qubits to use in construction of the circuit.
         gammas: Gamma angles to use as parameters for problem unitaries
         betas: Beta angles to use as parameters for driver unitaries
+        keep_zzswap_as_one_op: If True, use `recirq.qaoa.gates_and_compilation.ZZSwap`
+            custom, composite gate. This is required for using `get_compiled_sk_model_circuit`
+            and `compile_to_syc`. Otherwise, decompose each ZZSwap operation into a ZZPowGate
+            and SWAP gate. This is useful if you plan to use vanilla Cirq transformers.
     """
     circuit = get_generic_qaoa_circuit(problem_graph, qubits, gammas, betas)
     circuit = compile_problem_unitary_to_swap_network(circuit)
     circuit = compile_swap_network_to_zzswap(circuit)
     circuit = compile_driver_unitary_to_rx(circuit)
+    if not keep_zzswap_as_one_op:
+        circuit = cirq.expand_composite(
+            circuit, no_decomp=lambda op: not isinstance(op.gate, ZZSwap))
     return circuit
 
 

recirq/qaoa/sk_model/__init__.py

@@ -0,0 +1 @@
+from .sk_model import *

recirq/qaoa/sk_model/sk_model.py

@@ -0,0 +1,163 @@
+# Copyright 2022 Google
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+from dataclasses import dataclass
+from typing import List, Tuple, Iterable, Sequence
+
+import networkx as nx
+import numpy as np
+
+import cirq
+from cirq.protocols import dataclass_json_dict
+from cirq_google.workflow import QuantumExecutable, BitstringsMeasurement, QuantumExecutableGroup, \
+    ExecutableSpec
+from recirq.qaoa.classical_angle_optimization import optimize_instance_interp_heuristic
+from recirq.qaoa.problem_circuits import get_routed_sk_model_circuit
+
+
+def _graph_from_row_major_upper_triangular(
+        all_to_all_couplings: Sequence[float], *, n: int
+) -> nx.Graph:
+    """Get `all_to_all_couplings` in the form of a NetworkX graph."""
+    if not len(all_to_all_couplings) == n * (n - 1) / 2:
+        raise ValueError("Number of couplings does not match the number of nodes.")
+
+    g = nx.Graph()
+    for (u, v), coupling in zip(itertools.combinations(range(n), r=2), all_to_all_couplings):
+        g.add_edge(u, v, weight=coupling)
+    return g
+
+
+def _all_to_all_couplings_from_graph(graph: nx.Graph) -> Tuple[int, ...]:
+    """Given a networkx graph, turn it into a tuple of all-to-all couplings."""
+    n = graph.number_of_nodes()
+    if not sorted(graph.nodes) == sorted(range(n)):
+        raise ValueError("Nodes must be contiguous and zero-indexed.")
+
+    edges = graph.edges
+    return tuple(edges[u, v]['weight'] for u, v in itertools.combinations(range(n), r=2))
+
+
+@dataclass(frozen=True)
+class SKModelQAOASpec(ExecutableSpec):
+    """ExecutableSpec for running SK-model QAOA.
+
+    QAOA uses alternating applications of a problem-specific entangling unitary and a
+    problem-agnostic driver unitary. It is a variational algorithm, but for this spec
+    we rely on optimizing the angles via classical simulation.
+
+    The SK model is an all-to-all 2-body spin problem that we can route using the
+    "swap network" to require only linear connectivity (but circuit depth scales with problem
+    size)
+
+    Args:
+        n_nodes: The number of nodes in the SK problem. This is equal to the number of qubits.
+        all_to_all_couplings: The n(n-1)/2 pairwise coupling constants that defines the problem
+            as a serializable tuple of the row-major upper triangular coupling matrix.
+        p_depth: The depth hyperparemeter that presecribes the number of U_problem * U_driver
+            repetitions.
+        n_repetitions: The number of shots to take when running the circuits.
+        executable_family: `recirq.qaoa.sk_model`.
+
+    """
+
+    n_nodes: int
+    all_to_all_couplings: Tuple[int, ...]
+    p_depth: int
+    n_repetitions: int
+    executable_family: str = 'recirq.qaoa.sk_model'
+
+    def __post_init__(self):
+        object.__setattr__(self, 'all_to_all_couplings', tuple(self.all_to_all_couplings))
+
+    def get_graph(self) -> nx.Graph:
+        """Get `all_to_all_couplings` in the form of a NetworkX graph."""
+        return _graph_from_row_major_upper_triangular(self.all_to_all_couplings, n=self.n_nodes)
+
+    @staticmethod
+    def get_all_to_all_couplings_from_graph(graph: nx.Graph) -> Tuple[int, ...]:
+        """Given a networkx graph, turn it into a tuple of all-to-all couplings."""
+        return _all_to_all_couplings_from_graph(graph)
+
+    @classmethod
+    def _json_namespace_(cls):
+        return 'recirq.qaoa'
+
+    def _json_dict_(self):
+        return dataclass_json_dict(self, namespace=self._json_namespace_())
+
+
+def _classically_optimize_qaoa_parameters(graph: nx.Graph, *, n: int, p_depth: int):
+    param_guess = [
+        np.arccos(np.sqrt((1 + np.sqrt((n - 2) / (n - 1))) / 2)),
+        -np.pi / 8
+    ]
+
+    optima = optimize_instance_interp_heuristic(
+        graph=graph,
+        # Potential performance improvement: To optimize for a given p_depth,
+        # we also find the optima for lower p values.
+        # You could cache these instead of re-finding for each executable.
+        p_max=p_depth,
+        param_guess_at_p1=param_guess,
+        verbose=True,
+    )
+    # The above returns a list, but since we asked for p_max = spec.p_depth,
+    # we always want the last one.
+    optimum = optima[-1]
+    assert optimum.p == p_depth
+    return optimum
+
+
+def sk_model_qaoa_spec_to_exe(
+        spec: SKModelQAOASpec,
+) -> QuantumExecutable:
+    """Create a full `QuantumExecutable` from a given `SKModelQAOASpec`
+
+    Args:
+        spec: The spec
+
+    Returns:
+        a QuantumExecutable corresponding to the input specification.
+    """
+    n = spec.n_nodes
+    graph = spec.get_graph()
+
+    # Get params
+    optimum = _classically_optimize_qaoa_parameters(graph, n=n, p_depth=spec.p_depth)
+
+    # Make the circuit
+    qubits = cirq.LineQubit.range(n)
+    circuit = get_routed_sk_model_circuit(
+        graph, qubits, optimum.gammas, optimum.betas, keep_zzswap_as_one_op=False)
+
+    # QAOA code optionally finishes with a QubitPermutationGate, which we want to
+    # absorb into measurement. Maybe at some point this can be part of
+    # `cg.BitstringsMeasurement`, but for now we'll do it implicitly in the analysis code.
+    if spec.p_depth % 2 == 1:
+        assert len(circuit[-1]) == 1
+        permute_op, = circuit[-1]
+        assert isinstance(permute_op.gate, cirq.QubitPermutationGate)
+        circuit = circuit[:-1]
+
+    # Measure
+    circuit += cirq.measure(*qubits, key='z')
+
+    return QuantumExecutable(
+        spec=spec,
+        problem_topology=cirq.LineTopology(n),
+        circuit=circuit,
+        measurement=BitstringsMeasurement(spec.n_repetitions),
+    )

recirq/qaoa/sk_model/sk_model_test.py

@@ -0,0 +1,95 @@
+# Copyright 2022 Google
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import networkx as nx
+import numpy as np
+import pytest
+
+import cirq
+from recirq.qaoa.sk_model.sk_model import _graph_from_row_major_upper_triangular, \
+    _all_to_all_couplings_from_graph, SKModelQAOASpec, sk_model_qaoa_spec_to_exe
+
+
+def test_graph_from_row_major_upper_triangular():
+    couplings = np.array([
+        [0, 1, 2, 3],
+        [0, 0, 4, 5],
+        [0, 0, 0, 6],
+        [0, 0, 0, 0],
+    ])
+    flat_couplings = np.arange(1, 6 + 1)
+
+    graph1 = nx.from_numpy_array(couplings)
+    graph2 = _graph_from_row_major_upper_triangular(flat_couplings, n=4)
+    assert sorted(graph1.nodes) == sorted(graph2.nodes)
+    assert sorted(graph1.edges) == sorted(graph2.edges)
+
+    for u, v, w in graph1.edges.data('weight'):
+        assert w == graph2.edges[u, v]['weight']
+
+
+def test_graph_from_row_major_upper_triangular_bad():
+    with pytest.raises(ValueError):
+        _graph_from_row_major_upper_triangular([1, 2, 3, 4], n=2)
+
+
+def test_all_to_all_couplings_from_graph():
+    g = nx.Graph()
+    g.add_edge(0, 1, weight=1)
+    g.add_edge(0, 2, weight=2)
+    g.add_edge(1, 2, weight=3)
+    couplings = _all_to_all_couplings_from_graph(g)
+    assert couplings == (1, 2, 3)
+
+
+def test_all_to_all_couplings_from_graph_missing():
+    g = nx.Graph()
+    g.add_edge(0, 1, weight=1)
+    # g.add_edge(0, 2, weight=2)
+    g.add_edge(1, 2, weight=3)
+    with pytest.raises(KeyError):
+        _ = _all_to_all_couplings_from_graph(g)
+
+
+def test_all_to_all_couplings_from_graph_bad_nodes():
+    g = nx.Graph()
+    g.add_edge(10, 11, weight=1)
+    g.add_edge(10, 12, weight=2)
+    g.add_edge(11, 12, weight=3)
+    with pytest.raises(ValueError):
+        _ = _all_to_all_couplings_from_graph(g)
+
+
+@pytest.mark.parametrize('n', [5, 6])
+def test_graph_round_trip(n):
+    couplings = tuple(np.random.choice([0, 1], size=n * (n - 1) // 2))
+    assert couplings == _all_to_all_couplings_from_graph(
+        _graph_from_row_major_upper_triangular(couplings, n=n))
+
+
+def test_spec_to_exe():
+    spec = SKModelQAOASpec(
+        n_nodes=3, all_to_all_couplings=[1, -1, 1], p_depth=1, n_repetitions=1_000
+    )
+    assert isinstance(spec.all_to_all_couplings, tuple)
+    assert hash(spec) is not None
+    exe = sk_model_qaoa_spec_to_exe(spec)
+    init_hadamard_depth = 1
+    zz_swap_depth = 2  # zz + swap
+    driver_depth = 1
+    measure_depth = 1
+    assert len(
+        exe.circuit) == init_hadamard_depth + zz_swap_depth * 3 + driver_depth + measure_depth
+    assert exe.spec == spec
+    assert exe.problem_topology == cirq.LineTopology(3)