How to optimize controls on large sparse Hamiltonians

import matplotlib.pyplot as plt
import numpy as np
from qctrlvisualizer import get_qctrl_style, plot_controls

from qctrl import Qctrl

plt.style.use(get_qctrl_style())

# Start a Boulder Opal session.
qctrl = Qctrl()

# Define physical constraints.
omega_max = 2 * np.pi * 5e6  # rad/s
duration = 500e-9  # s
segment_count = 20
sample_times = np.linspace(0, duration, 100)
target_state = np.zeros(2**4, complex)
target_state[int("0000", 2)] = 1 / 2
target_state[int("0011", 2)] = 1j / 2
target_state[int("1100", 2)] = -1j / 2
target_state[int("1111", 2)] = 1 / 2


# Define coupling operators.
def create_operators(graph):
    sigma_x = graph.pauli_matrix("X")
    return [
        0.5 * graph.embed_operators([(sigma_x, n), (sigma_x, n + 1)], [2] * 4)
        for n in range(3)
    ]


# Build graph to estimate the Krylov subspace dimension.
graph = qctrl.create_graph()

# Find the spectral range for the Hamiltonian with highest omega value.
spectral_range = graph.spectral_range(operator=omega_max * sum(create_operators(graph)))

# Find the suggested Krylov subspace dimension for that spectral range
# with an error tolerance of 1e-5.
graph.estimated_krylov_subspace_dimension_lanczos(
    spectral_range=spectral_range,
    duration=duration,
    maximum_segment_duration=duration / segment_count,
    error_tolerance=1e-5,
    name="krylov_subspace_dimension",
)

# Estimate recommended Krylov subspace dimension.
krylov_subspace_dimension = qctrl.functions.calculate_graph(
    graph=graph, output_node_names=["krylov_subspace_dimension"]
).output["krylov_subspace_dimension"]["value"]

# Build optimization graph.
graph = qctrl.create_graph()

# Define initial state.
initial_state = graph.fock_state(16, 0)

# Define individual control signals (Omegas).
signals = [
    graph.utils.real_optimizable_pwc_signal(
        segment_count=segment_count,
        minimum=-omega_max,
        maximum=omega_max,
        duration=duration,
        name=rf"$\Omega_{n+1}$",
    )
    for n in range(3)
]

# Define individual terms of the Hamiltonian.
operators = create_operators(graph)
sparse_hamiltonian_terms = [
    graph.sparse_pwc_operator(signal=signal, operator=operator)
    for signal, operator in zip(signals, operators)
]

# Create total Hamiltonian by adding terms.
hamiltonian = graph.sparse_pwc_sum(sparse_hamiltonian_terms)

# Integrate the Schrödinger equation for the Hamiltonian,
# using the recommended choice of Krylov subspace dimension.
evolved_states = graph.state_evolution_pwc(
    initial_state=initial_state,
    hamiltonian=hamiltonian,
    krylov_subspace_dimension=krylov_subspace_dimension,
    sample_times=sample_times,
)

# Calculate the infidelities as a function of time.
sparse_infidelities = graph.state_infidelity(
    evolved_states, target_state, name="sparse_infidelities"
)

# Use the final state infidelity as the cost.
final_state_infidelity = sparse_infidelities[-1]
final_state_infidelity.name = "cost"

# Calculate the evolution using a dense Hamiltonian, for comparison.
dense_hamiltonian_terms = [
    signal * operator for signal, operator in zip(signals, operators)
]
dense_hamiltonian = graph.pwc_sum(dense_hamiltonian_terms)
time_evolution_operators = graph.time_evolution_operators_pwc(
    hamiltonian=dense_hamiltonian, sample_times=sample_times
)
dense_evolved_states = time_evolution_operators @ initial_state[:, None]
dense_infidelities = graph.state_infidelity(
    dense_evolved_states[..., 0], target_state, name="dense_infidelities"
)

# Run optimization.
result = qctrl.functions.calculate_optimization(
    graph=graph,
    cost_node_name="cost",
    output_node_names=["sparse_infidelities", "dense_infidelities"]
    + [rf"$\Omega_{n+1}$" for n in range(3)],
    optimization_count=5,
)

# Extract time-dependent infidelities.
sparse_infidelities = result.output.pop("sparse_infidelities")["value"].flatten()
dense_infidelities = result.output.pop("dense_infidelities")["value"].flatten()

# Plot and print results.
fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(sample_times / 1e-9, sparse_infidelities, label="Lanczos integration")
ax.plot(sample_times / 1e-9, dense_infidelities, ":", label="Exact integration")
ax.legend()
ax.set_xlabel("Time (ns)")
ax.set_ylabel("State infidelity")

print(f"\nRecommended Krylov subspace dimension:\t{krylov_subspace_dimension}")
print(f"Final state infidelity: {result.cost:.3e} (Lanczos integration)")
print(f"Final state infidelity: {dense_infidelities[-1]:.3e} (Exact integration)")

plot_controls(result.output)