Visualizing custom gates in Cirq

Question

I'm trying to visualize a custom gate in Cirq. I have through this page, which briefly mentions how to create custom gates in Cirq. However, I'm getting errors when trying to visualize it.

My current code defines a class CustomGate, which inherits from cirq.Gate class.

import numpy as np
import cirq 
class CustomGate(cirq.Gate):
    def __init__(self, unitary):
        self.unitary = unitary
        self.numQubits = int(np.log2(unitary.shape[0]))

    def _num_qubits_(self):
        return self.numQubits

    def _unitary_(self):
        return self.unitary

    def _circuit_diagram_info_(self, args='cirq.CircuitDiagramInfoArgs') -> 'cirq.CircuitDiagramInfo':
        return cirq.CircuitDiagramInfo(wire_symbols=("CG",), exponent=1.0, connected=True)

The usage is as follows:

# Define custom gate
customUnitary = np.eye(4)    # The custom unitary matrix would go here
CG = CustomGate(customUnitary)

# Setup circuit
q = cirq.LineQubit.range(2)
circuit=cirq.Circuit()
circuit.append(CG(q[0], q[1]))

# Visualize circuit
print(circuit)

I'm getting the following error:

Traceback (most recent call last):
  File "test.py", line 42, in <module>
    print(circuit)
  File "/Users/pgu/opt/anaconda3/lib/python3.7/site-packages/cirq/circuits/circuit.py", line 338, in __str__
    return self.to_text_diagram()
  File "/Users/pgu/opt/anaconda3/lib/python3.7/site-packages/cirq/circuits/circuit.py", line 1609, in to_text_diagram
    transpose=transpose)
  File "/Users/pgu/opt/anaconda3/lib/python3.7/site-packages/cirq/circuits/circuit.py", line 1668, in to_text_diagram_drawer
    get_circuit_diagram_info, include_tags)
  File "/Users/pgu/opt/anaconda3/lib/python3.7/site-packages/cirq/circuits/circuit.py", line 1924, in _draw_moment_in_diagram
    info = get_circuit_diagram_info(op, args)
  File "/Users/pgu/opt/anaconda3/lib/python3.7/site-packages/cirq/circuits/circuit.py", line 1818, in _get_operation_circuit_diagram_info_with_fallback
    info))
ValueError: Wanted diagram info from <__main__.CustomGate object at 0x10379bb90>.on(cirq.LineQubit(0), cirq.LineQubit(1)) for 2 qubits but got cirq.CircuitDiagramInfo(wire_symbols=('CG',), exponent=1.0, connected=True, exponent_qubit_index=None, auto_exponent_parens=True)

Answer 1

You're making a two qubit gate but the circuit diagram info only has one wire label. Change wire_symbols=("CG",) to wire_symbols=("CG", "#2").

Note that, since you're using default values for the diagram info you can just return the tuple return "CG", "#2" instead of bothering to wrap it in the more complicated data type.

Note that cirq does include cirq.MatrixGate, which takes an arbitrary unitary similar to what your custom gate is doing.