Plotting the GHZ state vectors

I am trying to reproduce the results of my own paper ( 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')())

def ghz_circuit():
    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

# 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()

    # 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[0].y = probs_real
    #[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


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?

Hey @Solomon!

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