# Copyright 2018 The Cirq Developers
#
# 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.

"""An optimization pass that combines adjacent single-qubit rotations."""

from typing import List, Tuple, Optional, cast

import numpy as np

from cirq import ops
from cirq.circuits import (
    Circuit,
    PointOptimizer,
    PointOptimizationSummary,
)
from cirq.extension import Extensions
from cirq.google.decompositions import two_qubit_matrix_to_native_gates


class MergeInteractions(PointOptimizer):
    """Combines adjacent constant single-qubit rotations."""

    def __init__(self,
                 tolerance: float = 1e-8,
                 allow_partial_czs: bool = True,
                 extensions: Extensions = None) -> None:
        self.tolerance = tolerance
        self.allow_partial_czs = allow_partial_czs
        self.extensions = extensions or Extensions()

    def optimization_at(self, circuit, index, op):
        if len(op.qubits) != 2:
            return None

        interaction_count, indices, matrix = (
            self._scan_two_qubit_ops_into_matrix(circuit, index, op.qubits))
        if interaction_count <= 1:
            return None

        # Find a max-3-cz construction.
        operations = two_qubit_matrix_to_native_gates(
            op.qubits[0],
            op.qubits[1],
            matrix,
            self.allow_partial_czs,
            self.tolerance)

        # TODO: don't replace if there's no benefit in CZ depth.

        return PointOptimizationSummary(
            clear_span=max(indices) + 1 - index,
            clear_qubits=op.qubits,
            new_operations=operations)

    def _op_to_matrix(self,
                      op: ops.Operation,
                      qubits: Tuple[ops.QubitId, ...]
                      ) -> Optional[Tuple[np.ndarray, bool]]:
        """Determines the effect of an operation on the given qubits.

        The operation must be a 1-qubit operation on one of the given qubits,
        or a 2-qubit operation on both of the given qubits. Also, the operation
        must have a known matrix. Otherwise None is returned.

        Args:
            op: The operation to understand.
            qubits: The qubits we care about. Order determines matrix tensor
                order.

        Returns:
            None, or else a tuple containing a matrix equivalent to the effect
            of the operation and a boolean indicating if the operation is a
            2-qubit interaction.
        """
        q1, q2 = qubits

        known = self.extensions.try_cast(op.gate, ops.KnownMatrixGate)
        if known is None:
            return None
        m = known.matrix()

        if op.qubits == qubits:
            return m, True
        if op.qubits == (q2, q1):
            return MergeInteractions._flip_kron_order(m), True
        if op.qubits == (q1,):
            return np.kron(np.eye(2), m), False
        if op.qubits == (q2,):
            return np.kron(m, np.eye(2)), False

        return None

    def _scan_two_qubit_ops_into_matrix(
            self,
            circuit: Circuit,
            index: Optional[int],
            qubits: Tuple[ops.QubitId, ...]
    ) -> Tuple[int, List[int], np.ndarray]:
        """Accumulates operations affecting the given pair of qubits.

        The scan terminates when it hits the end of the circuit, finds an
        operation without a known matrix, or finds an operation that interacts
        the given qubits with other qubits.

        Args:
            circuit: The circuit to scan for operations.
            index: The index to start scanning forward from.
            qubits: The pair of qubits we care about.

        Returns:
            A tuple containing:
                0. The number of 2-qubit operations that were scanned.
                1. The moment indices those operations were on.
                2. A matrix equivalent to the effect of the scanned operations.
        """

        product = np.eye(4, dtype=np.complex128)
        interaction_count = 0
        touched_indices = []

        while index is not None:
            operations = {circuit.operation_at(q, index) for q in qubits}
            op_data = [
                self._op_to_matrix(op, qubits)
                for op in operations
                if op
            ]

            # Stop at any non-constant or non-local interaction.
            if any(e is None for e in op_data):
                break
            present_op_data = cast(List[Tuple[np.ndarray, bool]], op_data)

            for op_mat, interacts in present_op_data:
                product = np.dot(op_mat, product)
                if interacts:
                    interaction_count += 1

            touched_indices.append(index)
            index = circuit.next_moment_operating_on(qubits, index + 1)

        return interaction_count, touched_indices, product

    @staticmethod
    def _flip_kron_order(mat4x4: np.ndarray) -> np.ndarray:
        """Given M = sum(kron(a_i, b_i)), returns M' = sum(kron(b_i, a_i))."""
        result = np.array([[0] * 4] * 4, dtype=np.complex128)
        order = [0, 2, 1, 3]
        for i in range(4):
            for j in range(4):
                result[order[i], order[j]] = mat4x4[i, j]
        return result

cirq/google/merge_interactions.py
# Copyright 2018 The Cirq Developers
#
# 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.

"""An optimization pass that combines adjacent single-qubit rotations."""

from typing import List, Tuple, Optional, cast

import numpy as np

from cirq import ops
from cirq.circuits import (
    Circuit,
    PointOptimizer,
    PointOptimizationSummary,
)
from cirq.extension import Extensions
from cirq.google.decompositions import two_qubit_matrix_to_native_gates


class MergeInteractions(PointOptimizer):
    """Combines adjacent constant single-qubit rotations."""

    def __init__(self,
                 tolerance: float = 1e-8,
                 allow_partial_czs: bool = True,
                 extensions: Extensions = None) -> None:
        self.tolerance = tolerance
        self.allow_partial_czs = allow_partial_czs
        self.extensions = extensions or Extensions()

    def optimization_at(self, circuit, index, op):
        if len(op.qubits) != 2:
            return None

        interaction_count, indices, matrix = (
            self._scan_two_qubit_ops_into_matrix(circuit, index, op.qubits))
        if interaction_count <= 1:
            return None

        # Find a max-3-cz construction.
        operations = two_qubit_matrix_to_native_gates(
            op.qubits[0],
            op.qubits[1],
            matrix,
            self.allow_partial_czs,
            self.tolerance)

        # TODO: don't replace if there's no benefit in CZ depth.

        return PointOptimizationSummary(
            clear_span=max(indices) + 1 - index,
            clear_qubits=op.qubits,
            new_operations=operations)

    def _op_to_matrix(self,
                      op: ops.Operation,
                      qubits: Tuple[ops.QubitId, ...]
                      ) -> Optional[Tuple[np.ndarray, bool]]:
        """Determines the effect of an operation on the given qubits.

        The operation must be a 1-qubit operation on one of the given qubits,
        or a 2-qubit operation on both of the given qubits. Also, the operation
        must have a known matrix. Otherwise None is returned.

        Args:
            op: The operation to understand.
            qubits: The qubits we care about. Order determines matrix tensor
                order.

        Returns:
            None, or else a tuple containing a matrix equivalent to the effect
            of the operation and a boolean indicating if the operation is a
            2-qubit interaction.
        """
        q1, q2 = qubits

        known = self.extensions.try_cast(op.gate, ops.KnownMatrixGate)
        if known is None:
            return None
        m = known.matrix()

        if op.qubits == qubits:
            return m, True
        if op.qubits == (q2, q1):
            return MergeInteractions._flip_kron_order(m), True
        if op.qubits == (q1,):
            return np.kron(m, np.eye(2)), False
        if op.qubits == (q2,):
            return np.kron(np.eye(2), m), False

        return None

    def _scan_two_qubit_ops_into_matrix(
            self,
            circuit: Circuit,
            index: Optional[int],
            qubits: Tuple[ops.QubitId, ...]
    ) -> Tuple[int, List[int], np.ndarray]:
        """Accumulates operations affecting the given pair of qubits.

        The scan terminates when it hits the end of the circuit, finds an
        operation without a known matrix, or finds an operation that interacts
        the given qubits with other qubits.

        Args:
            circuit: The circuit to scan for operations.
            index: The index to start scanning forward from.
            qubits: The pair of qubits we care about.

        Returns:
            A tuple containing:
                0. The number of 2-qubit operations that were scanned.
                1. The moment indices those operations were on.
                2. A matrix equivalent to the effect of the scanned operations.
        """

        product = np.eye(4, dtype=np.complex128)
        interaction_count = 0
        touched_indices = []

        while index is not None:
            operations = {circuit.operation_at(q, index) for q in qubits}
            op_data = [
                self._op_to_matrix(op, qubits)
                for op in operations
                if op
            ]

            # Stop at any non-constant or non-local interaction.
            if any(e is None for e in op_data):
                break
            present_op_data = cast(List[Tuple[np.ndarray, bool]], op_data)

            for op_mat, interacts in present_op_data:
                product = np.dot(op_mat, product)
                if interacts:
                    interaction_count += 1

            touched_indices.append(index)
            index = circuit.next_moment_operating_on(qubits, index + 1)

        return interaction_count, touched_indices, product

    @staticmethod
    def _flip_kron_order(mat4x4: np.ndarray) -> np.ndarray:
        """Given M = sum(kron(a_i, b_i)), returns M' = sum(kron(b_i, a_i))."""
        result = np.array([[0] * 4] * 4, dtype=np.complex128)
        order = [0, 2, 1, 3]
        for i in range(4):
            for j in range(4):
                result[order[i], order[j]] = mat4x4[i, j]
        return result

cirq/google/merge_interactions.py
diff --git a/cirq/google/merge_interactions.py b/cirq/google/merge_interactions.py
index 0af943cba29..b293038adba 100644
--- a/cirq/google/merge_interactions.py
+++ b/cirq/google/merge_interactions.py
@@ -95,9 +95,9 @@ def _op_to_matrix(self,
         if op.qubits == (q2, q1):
             return MergeInteractions._flip_kron_order(m), True
         if op.qubits == (q1,):
-            return np.kron(np.eye(2), m), False
-        if op.qubits == (q2,):
             return np.kron(m, np.eye(2)), False
+        if op.qubits == (q2,):
+            return np.kron(np.eye(2), m), False
 
         return None
 
diff --git a/cirq/google/merge_interactions_test.py b/cirq/google/merge_interactions_test.py
index 66e1c44886a..4a8651ae979 100644
--- a/cirq/google/merge_interactions_test.py
+++ b/cirq/google/merge_interactions_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from cirq import testing
 from cirq import circuits
 from cirq import ops
 from cirq.google import ExpZGate, MergeInteractions, MergeRotations
@@ -39,6 +40,16 @@ def assert_optimizes(before, after):
         print(after)
     assert before == after
 
+def assert_optimization_not_broken(circuit):
+    """Check that the unitary matrix for the input circuit is the same (up to
+    global phase and rounding error) as the unitary matrix of the optimized
+    circuit."""
+    u_before = circuit.to_unitary_matrix()
+    MergeInteractions().optimize_circuit(circuit)
+    u_after = circuit.to_unitary_matrix()
+
+    testing.assert_allclose_up_to_global_phase(u_before, u_after, atol=1e-8)
+
 
 def test_clears_paired_cnot():
     q0 = ops.QubitId()
@@ -84,3 +95,49 @@ def test_ignores_czs_separated_by_outer_cz():
             circuits.Moment([ops.CZ(q00, q10)]),
             circuits.Moment([ops.CZ(q00, q01)]),
         ]))
+
+
+def test_cnots_separated_by_single_gates_correct():
+    q0 = ops.QubitId()
+    q1 = ops.QubitId()
+    assert_optimization_not_broken(
+        circuits.Circuit.from_ops(
+            ops.CNOT(q0, q1),
+            ops.H(q1),
+            ops.CNOT(q0, q1),
+        ))
+
+
+def test_czs_separated_by_single_gates_correct():
+    q0 = ops.QubitId()
+    q1 = ops.QubitId()
+    assert_optimization_not_broken(
+        circuits.Circuit.from_ops(
+            ops.CZ(q0, q1),
+            ops.X(q1),
+            ops.X(q1),
+            ops.X(q1),
+            ops.CZ(q0, q1),
+        ))
+
+
+def test_inefficient_circuit_correct():
+    t = 0.1
+    v = 0.11
+    q0 = ops.QubitId()
+    q1 = ops.QubitId()
+    assert_optimization_not_broken(
+        circuits.Circuit.from_ops(
+            ops.H(q1),
+            ops.CNOT(q0, q1),
+            ops.H(q1),
+            ops.CNOT(q0, q1),
+            ops.CNOT(q1, q0),
+            ops.H(q0),
+            ops.CNOT(q0, q1),
+            ops.Z(q0)**t, ops.Z(q1)**-t,
+            ops.CNOT(q0, q1),
+            ops.H(q0), ops.Z(q1)**v,
+            ops.CNOT(q0, q1),
+            ops.Z(q0)**-v, ops.Z(q1)**-v,
+        ))

MergeInteractions optimization is creating bad circuits
Not sure how this made it past the tests yet.

```python
def main():
    circuit = make_inefficient_circuit()
    print("BEFORE:")
    print(circuit.to_unitary_matrix().round(3))
    print()
    cirq.google.MergeInteractions().optimize_circuit(circuit)
    print("AFTER:")
    print(circuit.to_unitary_matrix().round(3))

def make_inefficient_circuit(t=0.1, v=0.11):
    from cirq import H, CNOT, Z
    a = cirq.NamedQubit('a')
    b = cirq.NamedQubit('b')

    return cirq.Circuit.from_ops(
        H(b),
        CNOT(a, b),
        H(b),
        CNOT(a, b),
        CNOT(b, a),
        H(a),
        CNOT(a, b),
        Z(a)**t, Z(b)**-t,
        CNOT(a, b),
        H(a), Z(b)**v,
        CNOT(a, b),
        Z(a)**-v, Z(b)**-v)
```

```
BEFORE:
[[ 1.   -0.j     0.   -0.j     0.   +0.j     0.   -0.j   ]
 [-0.   -0.j    -0.   -0.309j  0.951-0.j    -0.   +0.j   ]
 [ 0.   -0.j     0.951-0.j    -0.   -0.309j  0.   +0.j   ]
 [-0.   +0.j     0.   +0.j     0.   -0.j    -0.771+0.637j]]

AFTER:
[[-0.354-0.354j -0.354-0.354j -0.227-0.446j  0.227+0.446j]
 [-0.452-0.213j -0.452-0.213j  0.364+0.342j -0.364-0.342j]
 [-0.354-0.354j  0.354+0.354j -0.446-0.227j -0.446-0.227j]
 [-0.452-0.213j  0.452+0.213j  0.496+0.063j  0.496+0.063j]]
```
