data_generation_fig5_6.py

# Copyright 2026 Henrik Gothen

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     https://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pickle as pkl

import numpy as np
from numpy.typing import NDArray
from qugradlab.pulses.invertible_functions.packaging import pack
from qugradlab.pulses.sampling import SampleTimes
from scipy.optimize import minimize
from thread_chunks import chunk

from pulsebasedcircuitelementsvqa.arg_parsing import parse_args
from pulsebasedcircuitelementsvqa.auxiliary_functions import (
    device_from_params,
    gate_gradients_wrapper,
)
from pulsebasedcircuitelementsvqa.custom_types import (
    DeviceParams,
    OptimizeResult,
    Result,
    TargetGate,
    drive_max,
    drive_random_scaling,
    flux_random_scaling,
    jmax,
    n,
    neighbour_zeeman_difference,
    reference_drive,
    simulate_in_rot_frame,
    tau,
)
from pulsebasedcircuitelementsvqa.gates_and_circuit_elements import (
    double_qubit_excitation as qubit_excitation,
)
from pulsebasedcircuitelementsvqa.pulse_form import get_pulse_form_no_g

num_cpus = parse_args()
print(f"num_cpus = {num_cpus}")

qubit_excitation_name = "double_qubit_excitation"
rwa_string = "_rwa"

savepath = "DataFig5/"

length = 100
nqudits = 4

factor = 1 if simulate_in_rot_frame else (nqudits - 1)

device_params = DeviceParams(
    qudits=nqudits,
    levels=2,
    exp_J=False,
    J_min=0,
    J_max=jmax,
    Z_max=drive_max,
    Z_detuning=neighbour_zeeman_difference / 2 * factor,
    use_graph=False,
)

initial_state_index = 1
device = device_from_params(device_params)
initial_state = device.hilbert_space.basis_vector(initial_state_index)

sample_count = 800_000 // 100 if simulate_in_rot_frame else 800_000
sample_count_adjusted = 1 + (1 + (sample_count) // length) * (length)

flux_signals = np.ones((length, (device_params["qudits"] - 1))) * (-1.0 + 1e-4)
drive_signals = np.zeros((length, device_params["qudits"], 2))
drive_signals[:, 1, 0] = 1e-4

n_random_init = 5  # defines how many random initialisations per pulse duration T are tried.
random_seed = 232
rng = np.random.default_rng(random_seed)
flux_signals_random = (
    rng.uniform(0, 1, (*flux_signals.shape, n_random_init)) * flux_random_scaling - 1
)
drive_signals_random = (
    0.5 - rng.uniform(0, 1, (*drive_signals.shape, n_random_init))
) * drive_random_scaling
flux_inits = np.concatenate([flux_signals[:, :, None], flux_signals_random], axis=-1)
drive_inits = np.concatenate([drive_signals[:, :, :, None], drive_signals_random], axis=-1)

thetas = np.arange(0.0, 0.2501, 0.0025) * 2 * np.pi
# We are starting somewhere in the mid theta range and then propagate the pulses towards
# smaller and larger thetas.
right_thetas = thetas[30:]
left_thetas = thetas[:31][::-1]


def find_pulses_for_fixed_T(T: float, xinit: NDArray, thetas: NDArray):
    # for the result to be propagated into the next higher theta I want two things to be true:
    # 1. The optimizer deems the result a success
    # 2. The infidelity is comparable to the one of earlier theta
    propagated_initialization = xinit  # in the first theta, xinit is the best we've got
    loginfidelity_from_previous = (
        -5
    )  # this also implies: Don't even get started with this random init if it can't beat -4
    result_propagation_spoiled = False
    inithash = hash(
        np.sum(xinit)
    )  # this has to be xinit so I can later figure out which results form a chain.
    secondary_inithash = "nosecondary"
    for target_theta in thetas:
        if result_propagation_spoiled:
            continue
        samples_per_point = 1 + sample_count // length
        sample_times = SampleTimes(T=T, number_sample_points=samples_per_point * (length) + 1)
        generate_pulse_form = get_pulse_form_no_g(
            length,
            T,
            sample_times.dt,
            samples_per_point,
            n=n,
            tau=tau,
        )
        driven_device = device.pulse_form(generate_pulse_form)

        if simulate_in_rot_frame:
            # gates are defined in rotating frame, hence just use identity here
            U_rot_frame = np.diag(np.ones(device.hilbert_space.dim))
        else:
            U_rot_frame = np.diag(np.exp(-1j * sample_times.T * np.diag(device.H0)))
        target_gate: TargetGate = TargetGate(
            U=U_rot_frame @ qubit_excitation(target_theta),
            name=qubit_excitation_name,
            param=target_theta,
        )
        figname_appendix = (
            f"{target_gate.name}_param{np.round(target_gate.param, 2)}_len{length}_"  # type: ignore
            f"gfactor{np.round(device.max_drive_strength/reference_drive)}_{sample_count}_{sample_count_adjusted}"
        )  # type: ignore

        res: OptimizeResult = minimize(
            gate_gradients_wrapper,
            x0=propagated_initialization,
            args=(
                nqudits,
                length,
                initial_state,
                target_gate,
                driven_device,
                propagated_initialization,
                sample_times,
                device_params,
                sample_count,
                sample_count_adjusted,
                rwa_string,
                savepath,
            ),
            method="BFGS",
            jac=True,
            options={"gtol": 1e-6, "xrtol": 4e-5},
        )
        res.hess_inv = np.zeros((2, 2))
        myres = Result(
            res=res,
            xinit=propagated_initialization,
            T=T,
            target_gate=target_gate,
            device_params=device_params,
            length=length,
            sample_count=sample_count,
            sample_count_adjusted=sample_count_adjusted,
            optimization_strategy="GRAPE_no_g" + rwa_string,
        )

        with open(f"{savepath}/log_file_finished", mode="a") as file:
            file.write(
                f"T{T}..inithash{inithash}..secondaryinithash{secondary_inithash}..target_theta{target_theta}\n"
            )

        with open(
            f"{savepath}/result_T{T:.2f}_{inithash}_{secondary_inithash}_{figname_appendix}",
            mode="wb+",
        ) as file:
            pkl.dump(myres, file)

        threshold = -4

        if res.success:
            if np.log10(res.fun) < loginfidelity_from_previous + 3:
                if np.log10(res.fun) < threshold:
                    propagated_initialization = res.x
                    loginfidelity_from_previous = np.log10(res.fun)
                else:
                    print(
                        f"Breaking because result didnt come below log10(infidelity) "
                        f"(={np.log10(res.fun)}:.2f) < {threshold}, T={T}"
                    )
                    break
            else:
                print(
                    f"Breaking because result was too much worse than previous, "
                    f"theta={target_theta}"
                )
                break
        else:
            print(
                f"Breaking because no optim. success at theta = {target_gate.param}, "
                f"res.success = {res.success}, T={T}"
            )
            break


xinits = [
    np.concatenate([pack(list(flux_inits[:, :, k])), np.zeros(nqudits)])
    for k in range(flux_inits.shape[-1])
]

Ts = np.arange(1000, 1021, 20)[::-1]

xinits_rollout = xinits * len(Ts)
Ts_rollout = np.array([[T] * (n_random_init + 1) for T in Ts]).flatten()
Tx_pairs = [[tr, xr] for tr, xr in zip(Ts_rollout, xinits_rollout, strict=False)]


if __name__ == "__main__":
    with open(f"{savepath}/log_file_finished", mode="w+") as file:
        file.write("")

    def find_pulses_right(a, b):
        find_pulses_for_fixed_T(a, b, thetas=right_thetas)

    chunk(
        find_pulses_right,
        Tx_pairs,
        chunk_size=num_cpus,
        progress_bar=True,
    )

    def find_pulses_left(a, b):
        find_pulses_for_fixed_T(a, b, thetas=left_thetas)

    chunk(
        find_pulses_left,
        Tx_pairs,
        chunk_size=num_cpus,
        progress_bar=True,
    )