Find and replace gate

Question

I have a collection of circuits with cirq that use CNOT(q1, q2). I would like to be able to find and replace all the instance of this gate in the circuit collection and replace them with a composite gate of single qubit on each qubit and another entanglement gate. So I want to do substitution like this one:

-o-        -X-x-Z- 
 |     =>     |
-o-        -Z-x-Y-

where I have tried to represent two different entanglement gates.

Is there an efficient way to do this using some cirq function (other than a basic loop over each moment of the circuit) ? I was able to use the findall_operations_with_gate_type() function, but was not able to replace these gate.

Answer 1

If you're able to install the pre-release version of Cirq via the pip install --pre cirq command, you can do this with the recently added Transformers. I've posted a snippet of how a Transformer would solve your problem. Please feel fee to replace cirq.XX in my example with your entangling gate.

import cirq

def map_func(op: cirq.Operation, _: int) -> cirq.OP_TREE:
    if op.gate == cirq.CNOT:
        yield cirq.X(op.qubits[0])
        yield cirq.Z(op.qubits[1])
        yield cirq.XX(op.qubits[0], op.qubits[1])
        yield cirq.Z(op.qubits[0])
        yield cirq.Y(op.qubits[1])
    else:
        yield op

qubits = cirq.LineQubit.range(2)
circuit = cirq.Circuit(cirq.H(qubits[0]), cirq.CNOT(qubits[0], qubits[1]))
print(circuit)
# prints
# 0: ───H───@───
#           │
# 1: ───────X───

circuit_mapped = cirq.map_operations_and_unroll(circuit, map_func)
print(circuit_mapped)
# prints
# 0: ───H───X───XX───Z───
#               │
# 1: ───────Z───XX───Y───