Plotting the GHZ state vectors

Hello,
I am trying to reproduce the results of my own paper (https://arxiv.org/pdf/2302.12889.pdf) in which I used Yao.jl, in PennyLane. Here is the code snippet:

import pennylane as qml
import plotly.graph_objects as go
from IPython.display import display
from tqdm import tqdm
import time

# Quantum circuit and device setup
wires = 2
dev4 = qml.device("default.qubit", wires=wires, shots=1000)

def Q_Plot(cirq_0):
    fig, ax = qml.draw_mpl(cirq_0, expansion_strategy='device')()
    print(qml.draw(cirq_0, expansion_strategy='device')())
    # plt.show()
    fig.show()

@qml.qnode(dev4)
def ghz_circuit():
    qml.Hadamard(wires=0)
    qml.CNOT(wires=[0, 1])
    # qml.CNOT(wires=[0, 2])

    return qml.state()

# Initialize list to store results
all_results = []

# Create figure and histogram plots for the real and imaginary parts
fig_real = go.FigureWidget()
fig_real.add_bar(x=[f'Basis {i}' for i in range(2 ** wires)], y=[0] * (2 ** wires))
fig_real.update_layout(xaxis_title='Basis State', yaxis_title='Probability', title='Distribution of Real Part')

fig_imag = go.FigureWidget()


# Display the figures
display(fig_real)

# Generate N=1000 calls to the quantum circuit and plot their distribution in real time
N = 100
for _ in tqdm(range(N)):
    result = ghz_circuit()
    all_results.append(result)

    # Separate the real and imaginary parts of the state vector
    probs_real = abs(result.real) ** 2
    # probs_imag = abs(result.imag) ** 2

    # Update the histogram plots with new data
    fig_real.data[0].y = probs_real
    # fig_imag.data[0].y = probs_imag

    # Update the plot layout
    fig_real.update_layout(title=f'Distribution of Real Part ({len(all_results)}/{N})')
    # fig_imag.update_layout(title=f'Distribution of Imaginary Part ({len(all_results)}/{N})')

    # Pause to allow for interactive plotting
    time.sleep(0.01)

    
Q_Plot(ghz_circuit)

This is the resulting plot:

And this is what I am trying to plot (from my paper):

Is there a better / built-in method to do this?
Thanks.

Hey @Solomon!

You can use qml.counts! There are good examples in the link I attached :slight_smile:. Let me know if this helps!