python - Trouble implementing Hamitonian with QutiP - Stack Overflow

I am trying to compute the following function in QuTiP.

This function emulates an experiment that is used to determine the non-linearity of a cavity resonator. The pulse sequence involves applying a displacement on the cavity of amplitude alpha, waiting a variable time t, followed by another displacement pulse of the same amplitude but opposite phase. The amplitude of the displacement pulse is swept. The simulated data is fitted to the experimental data to extract the detuning between the drive frequency and the cavity frequency, as well as the non-linearity self-kerr term.

Here is my implementation:

def plot_kerr_effect_cavity(N, a_max, n_max, delta, kerr, tmax, dt, savefig=False):
    """
  Args: 
    N (int): Dimension of the Hilbert space.
    a_max (float): Maximum displacement amplitude.
    n_max (int): Maximum photon number to consider.
    delta (float): Detuning between the cavity and the drive frequency.
    kerr (float): Kerr nonlinearity parameter.
    tmax (float): Maximum time for the wait time.
    dt (float): Time step.
    """

    alphas = np.linspace(0,a_max,20)
    wait_times = np.arange(0, tmax+dt/2, dt)
    data = np.ones((len(alphas), len(wait_times)))
    

    for i, alpha in enumerate(alphas):
        psi_0 = displace(N,alpha) * basis(N,0)
        data[i,:] = non_linearity_exp(N, psi_0, n_max, alpha, kerr, wait_times, delta)

    fig, ax = plt.subplots()
    wait_times = wait_times * 1e6 # Convert to ns
    cax = ax.imshow(data, extent=[wait_times[0],wait_times[-1],alphas[0],alphas[-1]],            origin='lower', aspect='auto', cmap='seismic_r')
    ax.set_xlabel('Wait Time [$\mu$s]')
    ax.set_ylabel(r'Displacement Amplitude $\alpha$')
    ax.set_title('Cavity NonLinearity')
    fig.colorbar(cax, ax=ax, label='$P_e$')  
    plt.show()


def non_linearity_exp(N, psi_0, n_max, alpha,  kerr, wait_times, delta):

    basis_n = basis(N, n_max) # ket |n>
    ham = sum(2* np.pi * delta * basis(N, n)*basis(N, n).dag() for n in range(n_max))
    ham += 2*np.pi*kerr/2 * sum(n*(n-1)*basis(N, n)*basis(N, n).dag() for n in range(n_max))
    
    
    psi_t = sesolve(ham,psi_0, tlist=wait_times)
    result = []
    
    for i in range(len(wait_times)):
        D = displace(N, -alpha)  # Create the displacement operator
        result.append(np.abs(basis(N,0).dag() * D * psi_t.states[i]))
  
    return result

However, the resulting data does not look correct. The final plot should look something like this:

Instead of what I get with my version of the code:

Any thoughts?