Showing observable or basis with draw_mpl()

[YutoMorohoshi]

Feature details

(from https://discuss.pennylane.ai/t/can-draw-mpl-show-observable/5269)

I often need to change the observable (i.e. the basis for measurement) and investigate the effects of that. I use qml.draw_mpl() for debugging because it is easy to use, but it does not seem to display the observable or basis. Currently, if using qml.draw(), I can see which basis I am measuring on. If qml.draw_mpl() can do the same, it will make debugging even easier.

Implementation

No response

How important would you say this feature is?

1: Not important. Would be nice to have.

Additional information

import pennylane as qml
import matplotlib.pyplot as plt

observables = [qml.X(0)@qml.X(1), qml.Y(0)@qml.Y(1), qml.Z(0)@qml.Z(1)]

dev = qml.device("default.qubit", wires=2)
@qml.qnode(dev)
def circuit(x):
    qml.RX(x, wires=0)
    qml.Hadamard(wires=1)
    qml.CNOT(wires=[0, 1])

    # Even if these are set to observables[1] or [2], the output image will not change.
    return qml.expval(observables[0])

fig, ax = qml.draw_mpl(circuit, level='device')(0.5)
plt.show()

[CatalinaAlbornoz]

Thanks again for adding this here @YutoMorohoshi . We're looking into design options and possible next steps.