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.

[CatalinaAlbornoz]

Hi @YutoMorohoshi !

We looked into design options and we think the one below could be a good one to start with. We might look at better designs in the future but for now we should just take into account the following:

• We will need to make sure the text is actually fitting inside the figure, so we may need to elongate the figure in the x direction if the text starts getting longer.
• We will want it to align with the center of the line, even with larger text sizes.
• Directly to the right of the measure symbol we should include the measurement and type similarly to qml.draw(). For example: $\langle X \rangle$, Probs, State, etc.

We don't have capacity to work on this right now but please let us know if you'd like to work on this yourself.
Thanks again for making this enhancement suggestion!

[YutoMorohoshi]

Thanks @CatalinaAlbornoz for providing of concrete actions!
I will try to work on it according to your ideas. This is my first time fixing OSS, so it will probably take a little time to PR, but I would be grateful if you could give me the chance.

[CatalinaAlbornoz]

Yes, go ahead and give it a try @YutoMorohoshi ! This blog post can give you a good guide to start. And feel free to ask here if you have any questions about the process!