Create ConvertToXmonGates optimizer with support for 2-qubit operations (#256)

Author: Strilanc

cirq/google/__init__.py

@@ -14,7 +14,6 @@
 
 from cirq.google.convert_to_xmon_gates import (
     ConvertToXmonGates,
-    xmon_gate_ext,
 )
 from cirq.google.decompositions import (
     controlled_op_to_native_gates,
@@ -39,6 +38,9 @@
 from cirq.google.xmon_device import (
     XmonDevice,
 )
+from cirq.google.xmon_gate_extensions import (
+    xmon_gate_ext,
+)
 from cirq.google.xmon_gates import (
     Exp11Gate,
     ExpWGate,

cirq/google/convert_to_xmon_gates.py

@@ -13,30 +13,99 @@
 # limitations under the License.
 
 from cirq import ops
-from cirq.circuits.optimization_pass import PointOptimizer, \
-    PointOptimizationSummary
+from cirq.circuits.optimization_pass import (
+    PointOptimizationSummary,
+    PointOptimizer,
+)
+from cirq.extension import Extensions
+from cirq.google.decompositions import (
+    single_qubit_matrix_to_native_gates,
+    two_qubit_matrix_to_native_gates,
+)
 from cirq.google.xmon_gate_extensions import xmon_gate_ext
 from cirq.google.xmon_gates import XmonGate
 
 
 class ConvertToXmonGates(PointOptimizer):
-    """Pointwise converts a circuit to XmonGates if possible."""
+    """Attempts to convert strange gates into XmonGates.
 
-    def __init__(self, ignore_cast_failures=True):
-        self.ignore_cast_failures = ignore_cast_failures
+    First, checks if the given extensions are able to cast the gate into an
+        XmonGate instance.
+
+    Second, checks if the given extensions are able to cast the gate into a
+        CompositeGate instance. If so, recurses on the decomposition.
+
+    Third, checks if the given extensions are able to cast the gate into a
+        KnownMatrixGate instance. If so, and the gate is a 1-qubit or 2-qubit
+        gate, then performs circuit synthesis of the operation.
+
+    Fourth, if ignore_failures is set, gives up and returns the gate unchanged.
+        Otherwise raises a TypeError.
+    """
+
+    def __init__(self,
+                 extensions: Extensions=None,
+                 ignore_failures=False) -> None:
+        """
+        Args:
+            extensions: The extensions instance to use when trying to
+                cast gates to known types. Defaults to the standard xmon
+                gate extension.
+            ignore_failures: If set, gates that fail to convert are forwarded
+                unchanged. If not set, conversion failures raise a TypeError.
+        """
+        self.extensions = extensions or xmon_gate_ext
+        self.ignore_failures = ignore_failures
+
+    def _convert_one(self, op: ops.Operation) -> ops.OP_TREE:
+        # Already supported?
+        if isinstance(op.gate, XmonGate):
+            return op
+
+        # Maybe we know how to wrap it?
+        xmon = self.extensions.try_cast(op.gate, XmonGate)
+        if xmon is not None:
+            return xmon.on(*op.qubits)
+
+        # Provides a decomposition?
+        composite = self.extensions.try_cast(op.gate, ops.CompositeGate)
+        if composite is not None:
+            return composite.default_decompose(op.qubits)
+
+        # Known matrix?
+        mat = self.extensions.try_cast(op.gate, ops.KnownMatrixGate)
+        if mat is not None and len(op.qubits) == 1:
+            gates = single_qubit_matrix_to_native_gates(mat.matrix())
+            return [g.on(op.qubits[0]) for g in gates]
+        if mat is not None and len(op.qubits) == 2:
+            return two_qubit_matrix_to_native_gates(
+                op.qubits[0],
+                op.qubits[1],
+                mat.matrix(),
+                allow_partial_czs=True)
+
+        # Just let it be?
+        if self.ignore_failures:
+            return op
+
+        raise TypeError("Don't know how to work with {!r}. "
+                        "It isn't an XmonGate, "
+                        "1-qubit KnownMatrixGate, "
+                        "2-qubit KnownMatrixGate, "
+                        "or CompositeGate.".format(op))
+
+    def convert(self, op: ops.Operation) -> ops.OP_TREE:
+        converted = self._convert_one(op)
+        if converted is op:
+            return converted
+        return [self.convert(e) for e in ops.flatten_op_tree(converted)]
 
     def optimization_at(self, circuit, index, op):
-        if self.ignore_cast_failures:
-            c = xmon_gate_ext.try_cast(op.gate, XmonGate)
-            if c is None:
-                return None
-        else:
-            c = xmon_gate_ext.cast(op.gate, XmonGate)
-
-        if c is op.gate:
+        converted = self.convert(op)
+        if converted is op:
             return None
 
         return PointOptimizationSummary(
             clear_span=1,
-            new_operations=ops.Operation(c, op.qubits),
+            new_operations=converted,
             clear_qubits=op.qubits)

cirq/google/eject_z_test.py

@@ -20,10 +20,10 @@
 
 def assert_optimizes(before, after):
     pre_optimizations = [
-        ConvertToXmonGates(),
+        ConvertToXmonGates(ignore_failures=True)
     ]
     followup_optimizations = [
-        ConvertToXmonGates(),
+        ConvertToXmonGates(ignore_failures=True),
         circuits.DropEmptyMoments()
     ]
 

cirq/google/engine/engine.py

@@ -29,7 +29,7 @@
 from google.protobuf.json_format import MessageToDict
 
 from cirq.api.google.v1 import program_pb2
-from cirq.circuits import Circuit, ExpandComposite
+from cirq.circuits import Circuit
 from cirq.circuits.drop_empty_moments import DropEmptyMoments
 from cirq.devices import Device
 from cirq.google.convert_to_xmon_gates import ConvertToXmonGates
@@ -188,14 +188,9 @@ def run_sweep(self,
             if not device:
                 raise TypeError('device is required when running a circuit')
             # Convert to a schedule.
-            expand = ExpandComposite()
-            convert = ConvertToXmonGates(ignore_cast_failures=False)
-            drop = DropEmptyMoments()
-
             circuit_copy = Circuit(program.moments)
-            expand.optimize_circuit(circuit_copy)
-            convert.optimize_circuit(circuit_copy)
-            drop.optimize_circuit(circuit_copy)
+            ConvertToXmonGates().optimize_circuit(circuit_copy)
+            DropEmptyMoments().optimize_circuit(circuit_copy)
 
             schedule = moment_by_moment_schedule(device, circuit_copy)
         elif isinstance(program, Schedule):

cirq/google/sim/xmon_simulator.py

@@ -25,22 +25,21 @@
     results = sim.run(circuit)
 """
 
-import functools
-import math
-
 from collections import defaultdict, Iterable
-from typing import Tuple  # pylint: disable=unused-import
+import math
 from typing import Dict, Iterator, List, Sequence, Union, cast
+from typing import Tuple  # pylint: disable=unused-import
 
+import functools
 import numpy as np
 
-from cirq.circuits import Circuit, ExpandComposite
+from cirq.circuits import Circuit
 from cirq.circuits.drop_empty_moments import DropEmptyMoments
-from cirq.ops import raw_types
-from cirq.schedules import Schedule
 from cirq.google import xmon_gates
 from cirq.google.convert_to_xmon_gates import ConvertToXmonGates
 from cirq.google.sim.xmon_stepper import Stepper
+from cirq.ops import raw_types
+from cirq.schedules import Schedule
 from cirq.study import ParamResolver, Sweep, Sweepable, TrialResult
 
 
@@ -85,7 +84,7 @@ class SimulatorTrialResult(TrialResult):
     Attributes:
         measurements: A dictionary from measurement gate key to measurement
             results ordered by the qubits acted upon by the measurement gate.
-        final_state: The final state (wave function) of the system after
+        final_states: The final states (wave function) of the system after
             the trial finishes.
     """
 
@@ -133,10 +132,10 @@ def run(
                 a canonical ordering of the qubits. This canonical ordering
                 is used to define the wave function.
             initial_state: If an int, the state is set to the computational
-                basis state corresponding corresponding to this state. Otherwise
-                if this is a np.ndarray it is the full initial state. In this
-                case it must be the correct size, be normalized (an L2 norm of
-                1), and have a dtype of np.complex64.
+                basis state corresponding corresponding to this state.
+                Otherwise  if this is a np.ndarray it is the full initial
+                state. In this case it must be the correct size, be normalized
+                (an L2 norm of 1), and have a dtype of np.complex64.
 
         Returns:
             Results for this run.
@@ -164,10 +163,10 @@ def run_sweep(
                 a canonical ordering of the qubits. This canonical ordering
                 is used to define the wave function.
             initial_state: If an int, the state is set to the computational
-                basis state corresponding corresponding to this state. Otherwise
-                if this is a np.ndarray it is the full initial state. In this
-                case it must be the correct size, be normalized (an L2 norm of
-                1), and have a dtype of np.complex64.
+                basis state corresponding corresponding to this state.
+                Otherwise if this is a np.ndarray it is the full initial state.
+                In this case it must be the correct size, be normalized (an L2
+                norm of 1), and have a dtype of np.complex64.
 
         Returns:
             List of trial results for this run, one for each possible parameter
@@ -230,10 +229,10 @@ def moment_steps(
                 a canonical ordering of the qubits. This canonical ordering
                 is used to define the wave function.
             initial_state: If an int, the state is set to the computational
-                basis state corresponding corresponding to this state. Otherwise
-                if this is a np.ndarray it is the full initial state. In this
-                case it must be the correct size, be normalized (an L2 norm of
-                1), and have a dtype of np.complex64.
+                basis state corresponding corresponding to this state.
+                Otherwise if this is a np.ndarray it is the full initial state.
+                In this case it must be the correct size, be normalized (an L2
+                norm of 1), and have a dtype of np.complex64.
             param_resolver: A ParamResolver for determining values of
                 Symbols.
 
@@ -285,21 +284,16 @@ def simulator_iterator(
     qubit_map = {q: i for i, q in enumerate(qubits)}
 
     # TODO: Use one optimization pass.
-    expand = ExpandComposite()
-    convert = ConvertToXmonGates(ignore_cast_failures=False)
-    drop = DropEmptyMoments()
-
     circuit_copy = Circuit(circuit.moments)
-    expand.optimize_circuit(circuit_copy)
-    convert.optimize_circuit(circuit_copy)
-    drop.optimize_circuit(circuit_copy)
+    ConvertToXmonGates().optimize_circuit(circuit_copy)
+    DropEmptyMoments().optimize_circuit(circuit_copy)
     validate_unique_measurement_keys(circuit_copy)
 
-    with Stepper(
-        num_qubits=len(qubits),
-        num_prefix_qubits=options.num_prefix_qubits,
-        initial_state=initial_state,
-        min_qubits_before_shard=options.min_qubits_before_shard) as stepper:
+    with Stepper(num_qubits=len(qubits),
+                 num_prefix_qubits=options.num_prefix_qubits,
+                 initial_state=initial_state,
+                 min_qubits_before_shard=options.min_qubits_before_shard
+                 ) as stepper:
         for moment in circuit_copy.moments:
             measurements = defaultdict(list)  # type: Dict[str, List[bool]]
             phase_map = {}  # type: Dict[Tuple[int, ...], float]
@@ -330,9 +324,8 @@ def simulator_iterator(
                             result = not result
                         measurements[gate.key].append(result)
                 else:
-                    raise TypeError(
-                        'Gate %s is not a gate supported by the xmon simulator.'
-                        % gate)
+                    raise TypeError('{!r} is not supported by the '
+                                    'xmon simulator.'.format(gate))
             stepper.simulate_phases(phase_map)
             yield StepResult(stepper, qubit_map, measurements)
 
@@ -414,9 +407,9 @@ def set_state(self, state: Union[int, np.ndarray]):
     def merge(a: 'StepResult', b: 'StepResult') -> 'StepResult':
         """Merges measurement results of last_result into a new Result.
 
-        The measurement results are merges such that measurements with duplicate
-        keys have the results of last_result before those of this objects
-        results.
+        The measurement results are merges such that measurements with
+        duplicate keys have the results of last_result before those of this
+        objects' results.
 
         Args:
             a: First result to merge.
@@ -425,7 +418,7 @@ def merge(a: 'StepResult', b: 'StepResult') -> 'StepResult':
         Returns:
             A new StepResult with merged measurements.
         """
-        new_measurements = {}  # type: Dict[str, np.ndarray]
+        new_measurements = {}  # type: Dict[str, list]
         for d in [a.measurements, b.measurements]:
             for key, results in d.items():
                 if key not in new_measurements:

cirq/google/sim/xmon_simulator_test.py

@@ -22,14 +22,13 @@
 import numpy as np
 import pytest
 
-
 from cirq.circuits import Circuit
-from cirq.linalg import allclose_up_to_global_phase
 from cirq.devices import UnconstrainedDevice
 from cirq.google import (
     ExpWGate, ExpZGate, Exp11Gate, XmonMeasurementGate, XmonQubit,
 )
 from cirq.google.sim import xmon_simulator
+from cirq.linalg import allclose_up_to_global_phase
 from cirq.ops import op_tree
 from cirq.ops import raw_types
 from cirq.ops.common_gates import CNOT, H, X, Y, Z, CZ

cirq/google/sim/xmon_stepper.py

@@ -256,7 +256,7 @@ def reset_state(self, reset_state):
                 args.append(kwargs)
             self._pool.map(_reset_state, args)
 
-    def simulate_phases(self, phase_map: Dict[Tuple[int], float]):
+    def simulate_phases(self, phase_map: Dict[Tuple[int, ...], float]):
         """Simulate a set of phase gates on the xmon architecture.
 
         Args:

cirq/google/xmon_gates.py

@@ -99,6 +99,9 @@ def to_proto(self, *qubits):
         op.measurement.key = self.key
         return op
 
+    def __repr__(self):
+        return 'XmonMeasurementGate({})'.format(repr(self.key))
+
 
 class Exp11Gate(XmonGate,
                 ops.AsciiDiagrammableGate,
@@ -302,8 +305,10 @@ def ascii_wire_symbols(self):
         return 'Z',
 
     def ascii_exponent(self):
-        if self.half_turns in [-0.5, -0.25, 0.25, 0.5]:
+        if self.half_turns in [0.25, 0.5]:
             return 1
+        if self.half_turns in [-0.5, -0.25]:
+            return -1
         return self.half_turns
 
     def try_cast_to(self, desired_type):
@@ -329,6 +334,11 @@ def matrix(self):
             raise ValueError("Don't have a known matrix.")
         return ops.RotZGate(half_turns=self.half_turns).matrix()
 
+    def trace_distance_bound(self):
+        if isinstance(self.half_turns, Symbol):
+            return 1
+        return abs(self.half_turns) * 3.5
+
     def to_proto(self, *qubits):
         if len(qubits) != 1:
             raise ValueError('Wrong number of qubits.')