enter image description hereI want to use gradient ascent pulse engineering (GRAPE) to optimize the CZ, however, my program cannot work properly, it has no ideal result. In fact, it is the reproduce of Fig. 1(d) in article “Time-Optimal Two- and Three-Qubit Gates for Rydberg Atoms_-Jandura_Pupillo_-2022”.enter image description here. Here is my matlab code, I cannot find why it is wrong:

clear all; clc; close all;

dim = 4;
Natom = 1;

%% Define basis
s01 = [1;0;0;0];   s0r = [0;1;0;0];   s11 = [0;0;1;0];   sW = [0;0;0;1];
s01_0r = s01*s0r';    s11_W = s11*sW';

Omega = 1;   theta = pi;

% Hamiltonian definition simplified
H = Omega/2*s01*s0r' + Omega/sqrt(2)*s11*sW' + (Omega/2*s01*s0r' + Omega/sqrt(2)*s11*sW')';

%% Start Gate and Target Gate
U0 = eye(dim);
U_F = exp(1i*theta)*s01*s01'+exp(1i*(2*theta+pi))*s11*s11';

%% Time list
t_list = linspace(0, 7, 101);
dt = t_list(2) - t_list(1);

% Initial Pulses with random values
phi = randn(1, length(t_list));
phi = max(min(phi, pi), -pi);

%% Fidelity initialization
fidelity = 0;
infid = zeros(1, 1501);

for iter = 1:1501
    % Bidirectional Evolution
    Omega = 1;
    A = U0;
    A_list = cell(1, length(t_list));
    for i = 1:length(t_list)
        varphi = phi(i);
        HH = Omega/2 * exp(1i*varphi) * s01 * s0r' + Omega/sqrt(2) * exp(1i*varphi) * s11 * sW' + (Omega/2 * exp(1i*varphi) * s01 * s0r' + Omega/sqrt(2) * exp(1i*varphi) * s11 * sW')';
        U_indt = expm(-1i * HH * dt);
        A = U_indt * A;
        A_list{i} = A;
    end

    B = U_F;
    B_list = cell(1, length(t_list));
    for i = length(t_list):-1:1
        varphi = phi(i);
        HH = Omega/2 * exp(1i*varphi) * s01 * s0r' + Omega/sqrt(2) * exp(1i*varphi) * s11 * sW' + (Omega/2 * exp(1i*varphi) * s01 * s0r' + Omega/sqrt(2) * exp(1i*varphi) * s11 * sW')';
        U_indt = expm(-1i * HH * dt);
        B = U_indt * B;
        B_list{i} = B;
    end

    % Gradient Calculation simplified
    partial_phi = zeros(1, length(t_list));
    for i = 1:length(t_list)
        varphi = phi(i);
        H = Omega/2 * exp(1i*varphi) * s01 * s0r' + Omega/sqrt(2) * exp(1i*varphi) * s11 * sW' + (Omega/2 * exp(1i*varphi) * s01 * s0r' + Omega/sqrt(2) * exp(1i*varphi) * s11 * sW')';
        partial_phi(i) = -2*real(trace(B_list{i}' * 1i * dt * H * A_list{i}) * trace(A_list{i}' * B_list{i}));
    end

    % Update Pulses with decreasing learning rate
    lr = 0.2 / (1 - 0.01 * iter);
    phi = phi + lr * partial_phi; % Note the minus sign for gradient descent
    phi = max(min(phi, pi), -pi);

    % Fidelity calculation
    fidelity = abs(trace(A_list{end}' * U_F)) / dim; % Adjusted for general dimension
    infid(iter) = 1 - fidelity;

    % Plot Pulses (simplified for demonstration)
    if mod(iter, 300) == 0 || iter == 15501
        fprintf('Iteration %d, Fidelity: %.4fn', iter, fidelity);
    end
end

% Final fidelity and infidelity plot
fprintf('Final Fidelity: %.4fn', fidelity);
figure; plot(infid); xlabel('Iteration'); ylabel('Infidelity');

I want to plot the Fig.1 (d), and can understand GRAPE code