How to create analytical signals for simulation and optimization

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

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

plt.style.use(get_qctrl_style())

# Create the data flow graph describing the system.
graph = qctrl.create_graph()

# Define the duration of the pulse.
duration = 4e-6  # s

# Define the parameters of the pulse.
# Use a minimal segmentation for a more efficient sampling.
pulse = graph.signals.gaussian_pulse_pwc(
    duration=duration,
    segment_count=100,
    amplitude=2 * np.pi * 1.0e6,  # rad/s
    width=duration * 0.1,
    flat_duration=duration * 0.2,
    segmentation="MINIMAL",
    name="$\\Omega$",
)

# Add the pulse terms to the Hamiltonian.
hamiltonian = graph.hermitian_part(
    pulse * graph.pauli_kronecker_product([("M", 0)], subsystem_count=2)
)

# Add the qubit-qubit coupling terms to the Hamiltonian.
couplings = {"X": 8.0e6, "Y": 8.0e6, "Z": 9.0e6}  # rad/s

for i, coupling in couplings.items():
    hamiltonian += coupling * graph.pauli_kronecker_product(
        [(i, 0), (i, 1)], subsystem_count=2
    )

sample_times = np.linspace(0.0, duration, 100)

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

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

# Evolve the initial state with the unitary to get the final state.
evolved_states = unitaries @ initial_state
evolved_states.name = "states"

# Run simulation.
result = qctrl.functions.calculate_graph(
    graph=graph, output_node_names=["$\\Omega$", "states"]
)

# Plot the evolution of the populations of the qubits using `plot_population_dynamics`.
states = result.output.pop("states")
qubit_populations = np.abs(states["value"].squeeze()) ** 2

plot_population_dynamics(
    sample_times, {rf"$|{k:02b}\rangle$": qubit_populations[:, k] for k in range(4)}
)

# View the generated Gaussian pulse using the `plot_controls` function.
plot_controls(result.output)

graph = qctrl.create_graph()

# Define the duration of the signals.
duration = 3e-6  # s

# Define the number of segments for Pwc signals.
segment_count = 100

# Define α tanh ramp acting on sigma_z.
alpha_start_value = graph.optimization_variable(
    count=1, lower_bound=-3e7, upper_bound=3e7, name="alpha_start_value"
)
alpha_end_value = graph.optimization_variable(
    count=1, lower_bound=-3e7, upper_bound=3e7, name="alpha_end_value"
)
alpha_ramp_duration = graph.optimization_variable(
    count=1, lower_bound=0.1e-6, upper_bound=duration, name="alpha_ramp_duration"
)
alpha_center_time = graph.optimization_variable(
    count=1, lower_bound=0.0, upper_bound=duration, name="alpha_center_time"
)
alpha_ramp = graph.signals.tanh_ramp_pwc(
    duration=duration,
    segment_count=segment_count,
    start_value=alpha_start_value,
    end_value=alpha_end_value,
    ramp_duration=alpha_ramp_duration,
    center_time=alpha_center_time,
    name=r"$\alpha$",
)

# Define γ cosine pulse acting on sigma_-.
gamma_amplitude = graph.optimization_variable(
    count=1,
    lower_bound=-2.0 * np.pi * 1e6,  # rad/s
    upper_bound=2.0 * np.pi * 1e6,  # rad/s
    name="gamma_amplitude",
)
gamma_pulse = graph.signals.cosine_pulse_pwc(
    duration=duration,
    segment_count=segment_count,
    amplitude=gamma_amplitude,
    flat_duration=duration * 0.2,
    name=r"$\gamma$",
)

# Create pulse Hamiltonian.
hamiltonian = alpha_ramp * graph.pauli_matrix("Z")
hamiltonian += graph.hermitian_part(gamma_pulse * graph.pauli_matrix("M"))

# Add detuning term.
delta = 2 * np.pi * 0.25e6  # rad/s
hamiltonian += delta * graph.pauli_matrix("Z")

# Define target gate.
target = graph.target(operator=graph.pauli_matrix("X"))

# Add dephasing noise term.
zeta = 2 * np.pi * 20e3  # rad/s
dephasing = zeta * graph.pauli_matrix("Z")

# Robust infidelity.
robust_infidelity = graph.infidelity_pwc(
    hamiltonian=hamiltonian,
    noise_operators=[dephasing],
    target=target,
    name="robust_infidelity",
)

# Compute the unitary using `time_evolution_operators_pwc`.
unitary = graph.time_evolution_operators_pwc(
    hamiltonian=hamiltonian, sample_times=np.array([duration]), name="unitary"
)

result = qctrl.functions.calculate_optimization(
    graph=graph,
    cost_node_name="robust_infidelity",
    output_node_names=[
        "alpha_start_value",
        "alpha_end_value",
        "alpha_ramp_duration",
        "alpha_center_time",
        "gamma_amplitude",
        "unitary",
        r"$\alpha$",
        r"$\gamma$",
    ],
    optimization_count=8,
)

print(f"Optimized cost: {result.cost:.3e}")

# Check that the unitary operator associated with the evolution of the system Hamiltonian
# does indeed correspond to an X gate.
result.output.pop("unitary")["value"].round(5)

# Print out the converged signal parameters.
print("Alpha signal")
for param in ["start_value", "end_value"]:
    print(f"\t{param}: \t{result.output.pop(f'alpha_{param}')['value'][0]:.3e} rad/s")

for param in ["ramp_duration", "center_time"]:
    print(f"\t{param}: \t{result.output.pop(f'alpha_{param}')['value'][0]:.3e} s")

print("Gamma signal")
for param in ["amplitude"]:
    print(f"\t{param}: \t{result.output.pop(f'gamma_{param}')['value'][0]:.3e} rad/s")

# View the optimized signals using the `plot_controls` function.
plot_controls(result.output)