Library

plot_state_populations(
    traj::NamedTrajectory;
    state_name::Symbol=:ψ̃,
    state_indices::Union{Nothing, AbstractVector{Int}}=nothing,
    control_name::Symbol=:u,
    subspace::Union{Nothing, AbstractVector{Int}}=nothing,
    kwargs...
)

Plot populations for multiple quantum states stored in a trajectory.

This function visualizes the time evolution of quantum state populations for trajectories containing multiple state trajectories (e.g., from sampling problems where multiple initial states are evolved). States are identified by a common prefix and numeric suffix pattern (e.g., :ψ̃1_system_1, :ψ̃2_system_1, etc.).

Mathematical Background

For a quantum state $|\psi(t)\rangle \in \mathcal{H}$ evolving under the Schrödinger equation, the population in computational basis state $|i\rangle$ is given by

\[P_i(t) = |\langle i|\psi(t)\rangle|^2 = |\psi_i(t)|^2\]

where $\psi_i(t)$ is the $i$-th component of the state vector in the computational basis.

For a normalized state, we have the conservation property:

\[\sum_{i=1}^{n} P_i(t) = \langle\psi(t)|\psi(t)\rangle = 1\]

When multiple states are being optimized simultaneously (as in sampling problems), this function plots the populations for each state, allowing comparison of how different initial conditions evolve under the same control protocol.

Key Properties:

• Populations are real and bounded: $P_i(t) \in [0,1]$
• Total probability is conserved: $\sum_i P_i(t) = 1$
• Can optionally restrict to a subspace (e.g., computational subspace excluding leakage states)

Arguments

• traj::NamedTrajectory: A trajectory containing one or more quantum states in isomorphism representation.

Keyword Arguments

• state_name::Symbol: The base name for state components. The function will find all trajectory components matching this prefix (e.g., :ψ̃ matches :ψ̃1_system_1, :ψ̃2_system_1, etc.). Default is :ψ̃.
• state_indices::Union{Nothing, AbstractVector{Int}}: If provided, only plot states with these indices (e.g., [1, 3] plots only the 1st and 3rd states). If nothing, plots all states matching the prefix. Default is nothing.
• control_name::Symbol: The name of the control signal component to include in the plot. Default is :u.
• subspace::Union{Nothing, AbstractVector{Int}}: If provided, only plot populations for these basis states (e.g., 1:2 for a qubit subspace). Useful for excluding leakage levels. Default is nothing (plot all levels).
• kwargs...: Additional keyword arguments passed to NamedTrajectories.plot.

Returns

• A Makie Figure object containing the population plots for all selected states.

Example

using NamedTrajectories
using PiccoloQuantumObjects
using PiccoloPlots

# Example: Two initial states evolving under the same controls
N = 100
Δt = 0.1

# Initial states
ψ1_init = ComplexF64[1, 0, 0]  # |0⟩
ψ2_init = ComplexF64[0, 1, 0]  # |1⟩

# Create trajectory with multiple states
traj = NamedTrajectory(
    (
        ψ̃1_system_1 = hcat([ket_to_iso(ψ1_init) for _ in 1:N]...),
        ψ̃2_system_1 = hcat([ket_to_iso(ψ2_init) for _ in 1:N]...),
        u = randn(2, N),
        Δt = fill(Δt, N),
    );
    controls = :u,
    timestep = :Δt,
)

# Plot populations for all states
plot_state_populations(traj)

# Plot only computational subspace (excluding 3rd level)
plot_state_populations(traj; subspace=1:2)

# Plot only first state
plot_state_populations(traj; state_indices=[1])

See also: plot_unitary_populations, NamedTrajectories.plot

plot_unitary_populations(
    traj::NamedTrajectory;
    unitary_columns::AbstractVector{Int}=1:2,
    unitary_name::Symbol=:Ũ⃗,
    control_name::Symbol=:u,
    kwargs...
)

Plot the state populations for specified columns of a unitary operator trajectory.

This function visualizes how the populations (squared magnitudes) of quantum states evolve over time for selected columns of a unitary matrix stored in a NamedTrajectory.

Mathematical Background

For a unitary operator $U(t) \in \mathcal{U}(n)$ evolving under a time-dependent Hamiltonian, this function plots the populations

\[P_{i,j}(t) = |U_{i,j}(t)|^2\]

where $U_{i,j}(t)$ is the $(i,j)$-th element of the unitary matrix at time $t$.

For a quantum system evolving according to the Schrödinger equation

\[\frac{d}{dt}U(t) = -iH(t)U(t), \quad U(0) = I\]

each column $j$ of $U(t)$ represents the time evolution of the initial basis state $|j\rangle$:

\[|\psi_j(t)\rangle = U(t)|j\rangle = \sum_{i=1}^{n} U_{i,j}(t)|i\rangle\]

The population $P_{i,j}(t) = |U_{i,j}(t)|^2$ gives the probability of finding the system in state $|i\rangle$ at time $t$ given that it started in state $|j\rangle$.

Key Properties:

• Unitarity ensures $\sum_{i=1}^{n} P_{i,j}(t) = 1$ for all $j$ and $t$ (probability conservation)
• $P_{i,j}(0) = \delta_{ij}$ (initially in definite state)
• Populations are real and bounded: $P_{i,j}(t) \in [0,1]$

The trajectory stores $U(t)$ in isomorphism representation $\tilde{U}(t)$, a vectorized form that preserves the operator structure while enabling efficient optimization algorithms.

Arguments

• traj::NamedTrajectory: A trajectory containing a unitary operator in isomorphism representation.

Keyword Arguments

• unitary_columns::AbstractVector{Int}: Indices of unitary matrix columns to plot. Each column $j$ corresponds to the evolution of basis state $|j\rangle$. Default is 1:2.
• unitary_name::Symbol: The name of the unitary operator component in the trajectory, stored as an isomorphism vector ($\tilde{U}$). Default is :Ũ⃗.
• control_name::Symbol: The name of the control signal component to include in the plot, typically the time-dependent control parameters $u(t)$ in $H(t) = H_0 + \sum_k u_k(t) H_k$. Default is :u.
• kwargs...: Additional keyword arguments passed to NamedTrajectories.plot, such as xlims, ylims, or Makie-specific plotting options.

Returns

• A Makie Figure object containing the population plots.

Example

using NamedTrajectories
using PiccoloQuantumObjects
using PiccoloPlots

# Define Hamiltonian: H = X + a₁(t)Z + a₂(t)Y
H_drift = PAULIS[:X]
H_drives = [PAULIS[:Z], PAULIS[:Y]]

# Generate control trajectory
N = 100
Δt = 0.1
ts = collect(0:Δt:Δt*(N-1))
u = 0.1 * randn(2, length(ts))

# Generate unitaries
Us = exp.(-im * [(H_drift + sum(u[:, k] .* H_drives)) * ts[k] for k = 1:N])

# Create trajectory
traj = NamedTrajectory(
    (
        Ũ⃗ = hcat(operator_to_iso_vec.(Us)...),
        u = u,
        Δt = ts,
    );
    controls = :u,
    timestep = :Δt,
)

# Plot populations for first two columns
plot_unitary_populations(traj)

# Plot only the first column
plot_unitary_populations(traj; unitary_columns=[1])

See also: NamedTrajectories.plot