Printing the matrix of a circuit during optimization

Let’s say I have a circuit like this:

@qml.qnode(dev)
def circuit(weights):
    qml.RZ(weights, wires=0)
    return qml.expval(qml.PauliZ(0))

I can print the circuit diagram like this at each step of optimization:

print(circuit.draw())

But I would like the print the unitary matrix corresponding to this circuit. Is there any way I can achieve this?

Hi @andrew1, welcome to the forum!

Good question! This isn’t something that’s done by default in PennyLane due to the potential overhead of calculating and storing the unitary as the number of qubits increases. However, it’s possible to put together some code to recover the unitary, for example:

import pennylane as qml
from pennylane import numpy as np

qml.enable_tape()

wires = 2
dev = qml.device("default.qubit", wires=wires)


def template(x):
    qml.RX(x, wires=0)
    qml.RY(0.3, wires=1)
    qml.CNOT(wires=[0, 1])


def find_unitary(template, *args, **kwargs):

    @qml.qnode(dev)
    def find_output_state(input_state):
        qml.QubitStateVector(input_state, wires=range(wires))
        template(*args, **kwargs)
        return qml.state()
    
    input_states = np.eye(2 ** wires)
    
    return np.stack([find_output_state(state) for state in input_states]).T


U = find_unitary(template, 0.4)

print(U)

In the example above, you can update your circuit by changing the contents of template() (and the value of wires accordingly).

2 Likes

@Tom_Bromley Thank you! This works when I use the “default.qubit” device, but what if I wanted to use the “default.mixed” device? In this case, it returns an array with shape (32, 32, 32). (where I have wires=5).

Hi @andrew1 — If you are interested in getting the unitary of a circuit then you are automatically thinking about pure unitary evolution, thus you don’t need to use the mixed state simulator.

2 Likes

@Nicolas_Quesada right you are… thanks!

@andrew1, please let us know if there’s anything else we can help with!

1 Like