How to tune the parameters of an optimization

import matplotlib.pyplot as plt
import numpy as np
import scipy
from time import time
from qctrlvisualizer import get_qctrl_style, plot_cost_histories

plt.style.use(get_qctrl_style())

from qctrl import Qctrl

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

# Define optimization parameters.
duration = 1.1e-6  # s
omega_max = 5.0e6 * (2 * np.pi)  # Hz
delta_range = 15.0e6 * (2 * np.pi)  # Hz

# Define physical parameters of the quantum system.
interaction_strength = 24.0e6 * (2 * np.pi)  # Hz
local_shifts = np.zeros((6,))
local_shifts[[0, -1]] = -4.5e6 * (2 * np.pi)  # Hz

# Build the terms of the Hamiltonian.
h_omega_values = [0.5] * 2**6 * 6
h_omega_indices = []
h_delta_values = []
h_delta_indices = []
h_fixed_values = []
h_fixed_indices = []

for j in range(2**6):
    state = np.array([int(m) for m in np.binary_repr(j, width=6)])

    # Global shift term (big delta).
    n_exc = np.sum(state)  # Number of excited atoms.
    h_delta_values.append(-n_exc)
    h_delta_indices.append([j, j])

    # Omega term.
    for i in range(6):
        # Couple "state" and "state with qubit i flipped".
        h_omega_indices.append([j, j ^ 2 ** (5 - i)])

    # Local shift term (small deltas).
    h_fixed = -np.sum(local_shifts * state)

    # Interaction term.
    for d in range(1, 6):
        # Number of excited atoms at distance d.
        n_d = np.dot(((state[:-d] - state[d:]) == 0), state[:-d])
        h_fixed += interaction_strength * n_d / d**6

    if h_fixed != 0.0:
        h_fixed_values.append(h_fixed)
        h_fixed_indices.append([j, j])

H_omega = scipy.sparse.coo_matrix(
    (h_omega_values, np.array(h_omega_indices).T), shape=[2**6, 2**6]
)
H_delta = scipy.sparse.coo_matrix(
    (h_delta_values, np.array(h_delta_indices).T), shape=[2**6, 2**6]
)
H_fixed = scipy.sparse.coo_matrix(
    (h_fixed_values, np.array(h_fixed_indices).T), shape=[2**6, 2**6]
)

def get_recommended_k(segment_count, H_omega, H_delta, H_fixed):
    """
    Return the appropriate Krylov subspace dimension (k) for a given set of parameters, as
    required for the integration. This is obtained for the Hamiltonian with the largest spectral range.
    """
    graph = qctrl.create_graph()

    maximal_range_hamiltonian = (
        omega_max * H_omega - delta_range * H_delta + H_fixed
    ).tocoo()

    # Calculate spectral range.
    spectral_range = graph.spectral_range(operator=maximal_range_hamiltonian)

    # Calculate the appropriate Krylov subspace dimension k.
    graph.estimated_krylov_subspace_dimension_lanczos(
        spectral_range=spectral_range,
        error_tolerance=1e-6,
        duration=duration,
        maximum_segment_duration=duration / segment_count,
        name="k",
    )

    result = qctrl.functions.calculate_graph(graph=graph, output_node_names=["k"])

    return result.output["k"]["value"]


def get_optimization_graph(segment_count, krylov_subspace_dimension):
    """
    Create an optimization graph for the requested number of segments.
    """
    graph = qctrl.create_graph()

    # Define pulses.
    omega_segments = graph.utils.real_optimizable_pwc_signal(
        segment_count=segment_count,
        minimum=0.0,
        maximum=omega_max,
        duration=duration,
        name="omega",
    )
    delta_segments = graph.utils.real_optimizable_pwc_signal(
        segment_count=segment_count,
        minimum=-delta_range,
        maximum=delta_range,
        duration=duration,
        name="delta",
    )
    fixed_segment = graph.pwc_signal(
        values=np.array([1]), duration=duration, name="fixed"
    )

    shift_delta = graph.sparse_pwc_operator(signal=delta_segments, operator=H_delta)
    shift_omega = graph.sparse_pwc_operator(signal=omega_segments, operator=H_omega)
    shift_fixed = graph.sparse_pwc_operator(signal=fixed_segment, operator=H_fixed)

    # Define initial state
    initial_state = graph.fock_state(2**6, 0)

    # Create target state.
    idx_even = sum(2**n for n in range(0, 6, 2))
    idx_odd = sum(2**n for n in range(1, 6, 2))
    target_state = np.zeros([2**6])
    target_state[[idx_even, idx_odd]] = 1.0 / np.sqrt(2.0)

    # Evolve the initial state.
    evolved_state = graph.state_evolution_pwc(
        initial_state=initial_state,
        hamiltonian=graph.sparse_pwc_sum([shift_fixed, shift_omega, shift_delta]),
        krylov_subspace_dimension=krylov_subspace_dimension,
    )

    # Calculate state infidelity.
    infidelity = graph.state_infidelity(target_state, evolved_state, name="infidelity")

    return graph

results = {}
segment_counts = [2**n for n in range(4, 10)]
calculation_times = []

# Run optimizations.
for segment_count in segment_counts:
    krylov_subspace_dimension = get_recommended_k(
        segment_count, H_omega, H_delta, H_fixed
    )
    graph = get_optimization_graph(segment_count, krylov_subspace_dimension)
    time_reference = time()
    results[segment_count] = qctrl.functions.calculate_optimization(
        graph=graph,
        cost_node_name="infidelity",
        output_node_names=["infidelity"],
        optimization_count=1,
        max_iteration_count=300,
        cost_history_scope="HISTORICAL_BEST",
    )
    calculation_times.append(time() - time_reference)

# Plot history of best costs.
plot_cost_histories(
    [
        result.cost_history.historical_best[0]
        for segment_count, result in results.items()
    ],
    labels=[str(n) for n in segment_counts],
    y_axis_log=True,
)
plt.suptitle("Best cost values histories for different numbers of segments")
plt.show()

# Plot final costs and calculation times.

costs = [result.cost for result in results.values()]

fig, axes = plt.subplots(1, 2, figsize=(15, 5))
labels = [str(s) for s in segment_counts]
axes[0].bar(labels, costs, log=True)
axes[1].bar(labels, calculation_times)

axes[0].set_title("Best cost achieved")
axes[0].set_xlabel("Segment count")
axes[0].set_ylabel("Final cost")

axes[1].set_title("Calculation time")
axes[1].set_xlabel("Segment count")
axes[1].set_ylabel("Runtime (s)")

plt.show()