How to simulate quantum dynamics subject to noise with graphs

import matplotlib.pyplot as plt
import numpy as np
from qctrlopencontrols import new_corpse_control
from qctrlvisualizer import display_bloch_sphere_from_density_matrices, get_qctrl_style

from qctrl import Qctrl

# Start a Boulder Opal session
qctrl = Qctrl()

plt.style.use(get_qctrl_style())

# Define the non-Markovian noise power spectrum
def pink_spectrum(frequencies, frequency_cutoff):
    return 1 / (frequencies + frequency_cutoff)


frequency_step = 2 * 1e3
frequencies = np.arange(0, 2 * 1e6, frequency_step)

power_densities = 4e9 * pink_spectrum(
    frequencies=frequencies, frequency_cutoff=0.05 * 1e6
)

fig = plt.figure(figsize=(10, 5))
fig.suptitle("Dephasing noise spectral density")
plt.plot(frequencies / 1e6, power_densities * 1e6, color="#EB64B6")
plt.fill_between(frequencies / 1e6, 0, power_densities * 1e6, color="#F790DB")
plt.xlabel("Frequency (MHz)")
plt.ylabel("Power density (1/MHz)")
plt.show()

# Define control parameters
omega_max = 2 * np.pi * 1e6  # Hz
total_rotation = np.pi / 2

# Import pulse from Q-CTRL Open Controls
pulse = new_corpse_control(
    rabi_rotation=total_rotation,
    azimuthal_angle=0.0,
    maximum_rabi_rate=omega_max,
    name="CORPSE",
)

# Run this number of simulations with different noise trajectories
simulation_count = 80


def get_computation_graph():
    graph = qctrl.create_graph()

    # Define initial state
    initial_state = graph.fock_state(2, 0)[:, None]

    # Signal for Omega extracted from the imported pulse
    omega_signal = graph.pwc(
        values=pulse.rabi_rates * np.exp(1j * pulse.azimuthal_angles),
        durations=pulse.durations,
        name="omega",
    )

    # System Hamiltonian
    system_hamiltonian = graph.hermitian_part(omega_signal * graph.pauli_matrix("M"))

    # Noise Hamiltonian. Generate random noise trajectories.
    noise_stf = graph.random_colored_noise_stf_signal(
        power_spectral_density=power_densities,
        frequency_step=frequency_step,
        batch_shape=(simulation_count,),
    )
    noise_pwc = graph.discretize_stf(
        stf=noise_stf, duration=pulse.duration, segment_count=50
    )
    noise_hamiltonian = noise_pwc * graph.pauli_matrix("Z")

    # Combine the system and the noise Hamiltonian to get the total Hamiltonian.
    # We can automatically combine a batched and a non-batched term.
    hamiltonian = system_hamiltonian + noise_hamiltonian

    # Compute the unitary using `time_evolution_operators_pwc`
    unitaries = graph.time_evolution_operators_pwc(
        hamiltonian=hamiltonian, sample_times=sample_times
    )

    # Use the unitary to evolve the initial state into the final state.
    final_states = unitaries @ initial_state

    # Calculate the density matrix from simulated states, ρ = |ψ><ψ|
    density_matrices = graph.outer_product(final_states[..., 0], final_states[..., 0])

    # Calculate the average density matrix
    average_density_matrices = graph.sum(density_matrices, axis=0) / simulation_count
    average_density_matrices.name = "average_density_matrices"

    # Calculate the sigma_x expectation
    sigma_x_expectations = graph.real(
        graph.expectation_value(final_states[..., 0], graph.pauli_matrix("X")),
        name="sigma_x_expectations",
    )

    return graph

# Define sample times for the output
sample_times = np.linspace(0, np.sum(pulse.durations), 100)

# Run simulation
simulation_result = qctrl.functions.calculate_graph(
    graph=get_computation_graph(),
    output_node_names=["average_density_matrices", "sigma_x_expectations"],
)

average_density_matrices = simulation_result.output["average_density_matrices"]["value"]
sigma_x_expectations = simulation_result.output["sigma_x_expectations"]["value"]

# Plot average sigma_x expectation
fig = plt.figure(figsize=(10, 5))
plt.suptitle(r"Average $\sigma_x$ expectation")
plt.plot(sample_times * 1e6, np.average(sigma_x_expectations, axis=0))
plt.axhline(y=0.0, c="k", lw=0.5)
plt.xlabel("Time (µs)")
plt.ylabel(r"$\overline{\langle\sigma_x\rangle}$")
plt.show()

display_bloch_sphere_from_density_matrices(density_matrices=average_density_matrices)

# Plot individual sigma_x expectations

fig = plt.figure(figsize=(10, 5))
fig.suptitle(r"$\sigma_x$ expectations")
plt.plot(sample_times * 1e6, sigma_x_expectations[0:10].T)
plt.axhline(y=0, c="k", lw=0.5)
plt.xlabel("Time (µs)")
plt.ylabel(r"$\langle\sigma_x\rangle$")
plt.show()