Implement a GRAFS option for optimization

[gderossi]

The qutip-qtrl package currently supports the GRAPE and CRAB algorithms for pulse optimization. We would like to build upon this by adding the GRAFS algorithm. This will require understanding how the current algorithms are implemented.

See #1 for an exploration of required inputs, outputs, dependencies, and design of the GRAPE algorithm in QuTip, which will be used as a guide for implementing GRAFS.

[jmc595]

The GRAPE algorithm is a method used in quantum control to find optimal control pulses that guide a quantum system from an initial state to a desired final state.

[jmc595]

The code implements the GRAPE algorithm for calculating pulse sequences in quantum systems. The "progress_bar" variable is initialized to monitor the progress of the GRAPE algorithm during optimization iterations.

If "progress_bar" is already defined it continues using that progress bar, otherwise it creates a new one:

 progress_bar = progress_bar or BaseProgressBar(R)

Epsilon is calculated in order to control the step size of the optimization process. It influences how much the control pulses are adjusted in each iteration. If the parameter "eps" is not provided, the algorithm calculates a default value for it:

    if eps is None:
        eps = 0.1 * (2 * np.pi) / (times[-1])

It starts by initializing the control pulse matrix 'u' which represents the control fields that will be tuned over the course of the optimization process.

 u = np.zeros((R, J, M))

In both the "grape_unitary" and "cy_grape_unitary" functions, the NumPy 3D array "u" is initialized to be filled with zeroes, representing the control pulse matrix. "(R, J, M)" are the dimensions of the array, where "R" is the number of iterations or steps in the GRAPE algorithm, "J" represents the number of control fields or Hamiltonian operators that can be tuned via the control pulses, and "M" represents the number of time points for the control pulse evaluation.

[jmc595]

The initial conditions and parameters such as the static Hamiltonian H0, the control Hamiltonian operator H_ops, and the number of iterations R, are set in the function definitions of "grape_unitary" and "cy_grape_unitary". These parameters are passed as arguments to the functions and then used in the function bodies:

def grape_unitary(
    U,
    H0,
    H_ops,
    R,
    times,
    eps=None,
    u_start=None,
    u_limits=None,
    interp_kind="linear",
    use_interp=False,
    alpha=None,
    beta=None,
    phase_sensitive=True,
    progress_bar=None,
):

def cy_grape_unitary(
    U,
    H0,
    H_ops,
    R,
    times,
    eps=None,
    u_start=None,
    u_limits=None,
    interp_kind="linear",
    use_interp=False,
    alpha=None,
    beta=None,
    phase_sensitive=True,
    progress_bar=None,
):

[jmc595]

The initial control pulses, u_start, and their adjustments to respect specified limits, u_limits, along with the iterative updates to minimize the difference between the evolved unitary and the target unitary, are handled in the "grape_unitary" and "cy_grape_unitary" functions.

In the "grape_unitary" function, if initial control pulses ("u_start") and limits ("u_limits") are provided, the initial control pulses are adjusted to respect the limits. If any value in "u_start" is less than the lower limit or greater than the upper limit, it is set to the corresponding boundary value. Subsequently, the adjusted initial control pulses are assigned to the control pulse array 'u'. If only "u_start" is provided, without "u_limits", the initial control pulses are used without any adjustments. If neither "u_start" nor "u_limits" is provided, the control pulses remain unchanged from their initial values:

    if u_limits and u_start:
        # make sure that no values in u0 violates the u_limits conditions
        u_start = np.array(u_start)
        u_start[u_start < u_limits[0]] = u_limits[0]
        u_start[u_start > u_limits[1]] = u_limits[1]

    if u_start is not None:
        for idx, u0 in enumerate(u_start):
            u[0, idx, :] = u0

The main loop for iterations manages the calculation of unitary operations based on the control pulses and Hamiltonians either directly or using interpolation, depending on "use_interp" flag. The loop runs from 0 to R-1, where R is the number of GRAPE iterations:

for r in range(R - 1):

In the loop, the progress bar is updated to visualize the progress of the optimization process.

progress_bar.update()

It calculates the time interval, "dt", based on the time points provided in the "times" array.

dt = times[1] - times[0]

When the use_interp boolean variable is set to True, the algorithm uses interpolation to estimate control pulses between specified time points. This involves interpolation functions like ip_funcs, which is created for each control operator, H_ops, and each time point. The interpolation function is created using "interpld" method from the "scipy.interpolate" module. These interpolation methods help create a smooth approximation of the control pulses across time intervals:

        if use_interp:
            ip_funcs = [
                interp1d(
                    times,
                    u[r, j, :],
                    kind=interp_kind,
                    bounds_error=False,
                    fill_value=u[r, j, -1],
                )
                for j in range(J)
            ]

The "_H_t" function calculates a time-dependent Hamiltonian based on the interpolated control pulses and control operators. calculates the sum of contributions from each control operator, where each contribution is the product of the interpolated control pulse and its corresponding control operator. The result is converted to float if needed:

            def _H_t(t, args=None):
                return H0 + sum(
                    float(ip_funcs[j](t)) * H_ops[j] for j in range(J)
                )

The "U_list" is designed to store a series of matrices representing the unitary evolution of the quantum system over discrete time steps. For each "idx" in the range, the code performs calculations for each discrete time interval. "_H_t(times[idx])" calculates the time-dependent Hamiltonian at that specific point. This Hamiltonian is then multiplied by "-1j" (the imaginary unit in Python) and "dt", the step size, creating a complex number in order to accurately describe how the system changes over time. The matrix exponential of this complex number is then computed using the ".expm" function. The resulting matrix from this matrix exponential is stored in "U_list". This process is repeated for each time point, generating a list of matrices that describe the unitary evolution of the quantum system over the entire time interval:

              U_list = [
                      (-1j * _H_t(times[idx]) * dt).expm() for idx in range(M - 1)
                  ]

If use_interp is not set to True, the Hamiltonian, "H_t", is constructed, without interpolation, at each time step, by combining the effects of various pulses. "H0" is added as the static part of the Hamiltonian and "u[r, j, idx] * H_ops[j]" calculates the influence of the j-th control pulse at time step "idx", multiplied by "u[r, j, idx]", where "u[r, j, idx]" represents a set of control coefficients.

def _H_idx(idx):
                return H0 + sum([u[r, j, idx] * H_ops[j] for j in range(J)])

The algorithm then computes the forward evolutions and backward evolutions. The forward evolution, "U_f_list", represents the cumulative evolution from the initial state up to each subsequent step, while the backwards evolution, "U_b_list", represents the cumulative evolution from each subsequent step back to the initial state.

"U_f_list" and "U_b_list" are initialized as empty lists that store the forward and backward evolution operators at each step. As the algorithm progresses, each element in U_f_list will represent the accumulated forward evolution operator up to a certain point in time. Each element in U_b_list will represent the accumulated backward evolution operator from a certain point in time back to the initial state:

        U_f_list = []
        U_b_list = []

"U_f" and "U_b", are initialized as identity operators of the same dimension as the unitary operators in "U_list", thus ensuring they start as neutral operators that don't affect the quantum state but can accumulate transformations as dictated by the sequence of unitary operations in "U_list":

        U_f = qeye(U.dims[0])
        U_b = qeye(U.dims[0])

The forward evolution is computed through a loop that iterates over the length of "U_list". Within each iteration, "U_f" is updated by left multiplying it with the "n"-th unitary operator from "U_list", which accumulates the forward evolution up to the "n"-th operator. "@" denotes matrix multiplication in the Qutip library, where the unitary operator "U_list[n]" is applied to "U_f" from the left. The updated "U_f" is then appended to "U_f_list":

      for n in range(M - 1):
            U_f = U_list[n] @ U_f
            U_f_list.append(U_f)

Simultaneously, the current "U_b", which begins as the identity operator, is inserted at the beginning of "U_b_list", in order to apply inverse transformations in reverse order compared to the forward evolution. "U_b" is updated by by right-multiplying it with the Hermitian adjoint, ".dag()" (i.e. dagger), of the corresponding unitary operator from "U_list", ultimately reversing the action of the operator. This part of the loop accumulates the backward evolution, starting from the end of the list and moving towards the beginning:

            U_b_list.insert(0, U_b)
            U_b = _data.adjoint(U_list[M - 2 - n]) @ U_b

The algorithm then loops through each control field, "J", and each time step, "M". For each control field and each time step, "P" represents the quantum state after applying the backward evolution described by "U_b_list[m]" to the current quantum state "U". "Q" evaluates how applying the Hamiltonian operator, "H_ops[j]", at a given time step, "m", affects the original state "U_f_list[m]", at a small time increment, "dt":

        for j in range(J):
            for m in range(M - 1):
                P = U_b_list[m] @ U
                Q = 1j * dt * H_ops[j] @ U_f_list[m]

There are two different methods for computing the gradient, "du", depending on whether "phase_sensitive" is "True" or "False". When "phase_sensitive" is "True", the method simply calculates "du" as "-_overlap(P, Q)", where "_overlap(P, Q)" measures the overlap between "P" and "Q" and the "-" helps indicate that the system should adjust the pulses in the opposite direction of how "P" and "Q" are overlapping. When "phase_sensitive" is "False", the method for computing "du" is more complex: "_overlap(P, Q)" still computes the overlap between "P" and "Q", but "_overlap(U_f_list[m], P)" introduces another overlap calculation between "U_f_list[m]" and "P". These overlaps are multiply together with a coefficient "-2" which serves to adjust the magnitude of the gradient and help scale the impact of the overlaps:

                if phase_sensitive:
                    du = -_overlap(P, Q)
                else:
                    du = -2 * _overlap(P, Q) * _overlap(U_f_list[m], P)

[jmc595]

@gderossi Can you explain why there is a negative sign when calculating the gradient here:

                if phase_sensitive:
                    du = -_overlap(P, Q)

[jmc595]

Flag because I don't fully understand: why algorithm finds overlap between P & Q / P & U_f_list specifically:

if phase_sensitive:
                    du = -_overlap(P, Q)
                else:
                    du = -2 * _overlap(P, Q) * _overlap(U_f_list[m], P)

[gderossi]

@gderossi Can you explain why there is a negative sign when calculating the gradient here:

                if phase_sensitive:
                    du = -_overlap(P, Q)

According to slide 8 on these slides, the negative sign is just part of the gradient. There might be a derivation in the original GRAPE paper, but I wouldn't worry about going into that level of detail unless you think it will be important for implementing GRAFS.