`mesolve` fails (or slow) while an equivalent QuTiP script works ok

[eunjongkim]

Describe the Issue!

Hi, I have been trying to translate my QuTiP script into QuantumToolbox.jl to see whether I could benefit from the expected speedup. However, I have been stuck at some weird issue where simple mesolve (time-independent Hamiltonian) is very slow (or fails). I was wondering if you have any insights on what is causing this?

Python Script (QuTiP):

from qutip import *

import numpy as np
from tqdm.notebook import tqdm

from matplotlib import pyplot as plt
import matplotlib.cm as cm
plt.style.use('default')

class TransmonROParams:
    """Abstract class to calculate theoretically predicted parameters for dispersive readout of transmon qubits """
    def __init__(self, omega_q, alpha, omega_r, g, kappa):
        self.omega_q = omega_q
        self.omega_r = omega_r
        self.g = g
        self.kappa = kappa
        self.alpha = alpha
    @property
    def delta_qr(self):
        """Qubit-Resonator Detuning"""
        return self.omega_q - self.omega_r
    
    @property
    def resonator_ringdown_time(self):
        """Resonator Ringdown Time"""
        return 1 / self.kappa

    @property
    def omega_r_dressed(self):
        """Dressed Resonator Frequency"""
        return self.omega_r - self.g ** 2 / (self.delta_qr + self.alpha)

    @property
    def omega_q_dressed(self):
        """Dressed Qubit Frequency"""
        return self.omega_q + self.g ** 2 / self.delta_qr
    
    @property
    def chi(self):
        """Dispersive shift"""
        return self.g ** 2 * self.alpha / (self.delta_qr * (self.delta_qr + self.alpha))
    
    @property
    def n_crit(self):
        """Critical Photon Number"""
        return self.delta_qr ** 2 / (4 * self.g ** 2)
    
    @property
    def gamma1_purcell(self):
        """Purcell Decay Rate"""
        return (self.g / self.delta_qr) ** 2 * self.kappa

    @property
    def t1_purcell(self):
        """Purcell-limited Lifetime"""
        return 1 / self.gamma1_purcell
    
    def alpha0_ss(self, omega_d, epsilon):
        """
        Steady-state resonator field amplitude corresponding to the qubit ground state,
        given an external drive with frequency `omega_d` and constant amplitude `epsilon`.
        """
        delta_r = self.omega_r_dressed - omega_d
        return - epsilon / (delta_r - self.chi - 1j * self.kappa / 2)

    def alpha1_ss(self, omega_d, epsilon):
        """
        Steady-state resonator field amplitude corresponding to the qubit excited state,
        given an external drive with frequency `omega_d` and constant amplitude `epsilon`.
        """
        delta_r = self.omega_r_dressed - omega_d
        return - epsilon / (delta_r + self.chi - 1j * self.kappa / 2)
    
# parameters
omega_q = 5 * (2 *np.pi)            # qubit transition frequency
omega_r = 7 * (2 * np.pi)           # readout resonator frequency
g = 0.15 * (2 * np.pi)              # qubit-resonator coupling
alpha = -0.2 * (2 * np.pi)          # anharmonicity
kappa = 2.046e-3 * (2 * np.pi)      # resonator decay rate

Delta = omega_q - omega_r           # qubit resonator detuning

# Create a transmon ro params instance
ro_params = TransmonROParams(omega_q, alpha, omega_r, g, kappa)

print(f'Resonator ringdown time: {ro_params.resonator_ringdown_time:.2f} ns')

# H space dimension
N_q, N_r = 4, 10                      # Hilbert space dimensions (qubit, resonator)

# operators
a = tensor(qeye(N_q), destroy(N_r))    # resonator annihilation operator
b = tensor(destroy(N_q), qeye(N_r))    # transmon annihilation operator

# Hamiltonians (time-independent)
H_q = omega_q * b.dag() * b + 0.5 * alpha * b.dag() * b.dag() * b * b    # bare transmon Hamiltonian
H_r = omega_r * a.dag() * a                                              # bare resonator Hamiltonian
#H_c = g * (a + a.dag()) * (b + b.dag())                                 # quantum Rabi model Hamiltonian (takes longer time to simulate)
H_c = g * (a * b.dag() + a.dag() * b)                                    # Jaynes-Cummings model Hamiltonian (faster to simulate)

# Total H
H = H_q + H_r + H_c

print(f"Dressed Qubit Freq. (Analytical): {ro_params.omega_q_dressed / (2 * np.pi):.4f} GHz")
print(f"Dressed Resonator Freq. (Analytical): {ro_params.omega_r_dressed / (2 * np.pi):.4f} GHz")
print(f'Dispersive shift (Analytical): {ro_params.chi / (2 * np.pi) / 0.001:.3f} MHz')

def dressed_state(s):
    """
    Return the eigenenergy and eigenstate of the Hamiltonian that has the largest overlap with the state `s`.
    """
    eig_energies, eig_states = H.eigenstates()

    overlaps = np.array([np.abs(s.overlap(ds)) for ds in eig_states])

    s_ind = np.argmax(overlaps)
    return eig_energies[s_ind], eig_states[s_ind]

E_g0, _ = dressed_state(tensor(basis(N_q, 0), tensor(basis(N_r, 0))))
E_g1, _ = dressed_state(tensor(basis(N_q, 0), tensor(basis(N_r, 1))))
E_e0, _ = dressed_state(tensor(basis(N_q, 1), tensor(basis(N_r, 0))))
E_e1, _ = dressed_state(tensor(basis(N_q, 1), tensor(basis(N_r, 1))))

omega_q_dressed_numerical = (E_e0 - E_g0)   # qubit frequency when the resonator is not occupied
omega_r_g_numerical = (E_g1 - E_g0)   # readout resonator frequency when qubit is in the ground state
omega_r_e_numerical = (E_e1 - E_e0)   # readout resonator frequency when qubit is in the excited state

print(f'Dressed Qubit Freq. (Numerical): {omega_q_dressed_numerical / (2 * np.pi):.4f} GHz"')
print(f'Resonator Freq. w/ Qubit in |g> (Numerical) = {omega_r_g_numerical/ (2 * np.pi):.4f} GHz')
print(f'Resonator Freq. w/ Qubit in |e> (Numerical) = {omega_r_e_numerical/ (2 * np.pi):.4f} GHz')
print(f'Dressed Resonator Freq. (Numerical): {0.5 * (omega_r_g_numerical + omega_r_e_numerical) / (2 * np.pi):.4f} GHz')
print(f'Dispersive Shift (Numerical) = {(omega_r_e_numerical - omega_r_g_numerical) / 2e-3 / (2 * np.pi):.4f} MHz')

print(f'Purcell-limited Lifetime (Analytical) = {ro_params.t1_purcell * 1e-3:.3f} us')
print(f'Purcell Decay Rate (Analytical) = {ro_params.gamma1_purcell / (2 * np.pi * 1e-6) :.3f} kHz' )

# Find dressed qubit excited state and the ground state
_, psi_e0_dressed = dressed_state(tensor(basis(N_q, 1), basis(N_r, 0)))
_, psi_g0_dressed = dressed_state(tensor(basis(N_q, 0), basis(N_r, 0)))

# Projectors onto the dressed qubit excited state and the ground state
proj_e0_dressed = psi_e0_dressed * psi_e0_dressed.dag()
proj_g0_dressed = psi_g0_dressed * psi_g0_dressed.dag()

tlist0 = np.arange(0, 50000, 50)   # list of times for time evolution, in units of ns

result0 = mesolve(
    [H], psi_e0_dressed, tlist0, c_ops=[np.sqrt(kappa) * a],  # no collapse operator for qubit
    e_ops=[proj_g0_dressed, proj_e0_dressed], options = {"progress_bar": 'tqdm', 'nsteps': 100000}
)

This QuTiP script takes about 1 min to execute on my desktop.

Julia Script (QuantumToolbox.jl):

using QuantumToolbox, CUDA
using CairoMakie, Printf
import DifferentialEquations: AutoTsit5, Rosenbrock23, AutoVern7, Rodas5, QNDF, CVODE_BDF, JVODE_Adams
import QuantumToolbox:mesolve
using LSODA

# struct to 
struct TransmonROParams{T}
    omega_q::T
    anharmonicity::T
    omega_r::T
    g::T
    kappa::T
end

delta_qr(p::TransmonROParams) = p.omega_q - p.omega_r
resonator_ringdown_time(p::TransmonROParams) = 1 / p.kappa
omega_r_dressed(p::TransmonROParams) = p.omega_r - p.g ^ 2 / (delta_qr(p) + p.anharmonicity)
omega_q_dressed(p::TransmonROParams) = p.omega_q + p.g ^ 2 / delta_qr(p)
chi(p::TransmonROParams) = p.g ^ 2 * p.anharmonicity / (delta_qr(p) * (delta_qr(p) + p.anharmonicity))
n_crit(p::TransmonROParams) = delta_qr(p) ^ 2 / (4 * p.g ^ 2)
gamma1_purcell(p::TransmonROParams) = (p.g / delta_qr(p)) ^ 2 * p.kappa
t1_purcell(p::TransmonROParams) = 1 / gamma1_purcell(p)
alpha0_ss(p::TransmonROParams, omega_d::Float64, epsilon::Float64) = begin
    delta_r = omega_q_dressed(p) - omega_d
    - epsilon / (delta_r - chi(p) - im * p.kappa / 2)
end
alpha1_ss(p::TransmonROParams, omega_d::Float64, epsilon::Float64) = begin
    delta_r = omega_q_dressed(p) - omega_d
    - epsilon / (delta_r + chi(p) - im * p.kappa / 2)
end

# parameters
omega_q = 5 * 2π           # qubit transition frequency
omega_r = 7 * 2π           # readout resonator frequency
g = 0.15 * 2π              # qubit-resonator coupling
alpha = -0.2 * 2π          # anharmonicity
kappa = 2.046e-3 * 2π      # resonator decay rate

# Create a transmon ro params instance
ro_params = TransmonROParams(omega_q, alpha, omega_r, g, kappa)

@printf "Resonator ringdown time: %.2f ns" resonator_ringdown_time(ro_params)

# H space dimension
N_q, N_r = 4, 10                      # Hilbert space dimensions (qubit, resonator)

# operators
a = tensor(qeye(N_q), destroy(N_r))    # resonator annihilation operator
b = tensor(destroy(N_q), qeye(N_r))    # transmon annihilation operator

# Hamiltonians (time-independent)
H_q = omega_q * b' * b + 0.5 * alpha * b' * b' * b * b         # bare transmon Hamiltonian
H_r = omega_r * a' * a                                         # bare resonator Hamiltonian
#H_c = g * (a + a.dag()) * (b + b.dag())                       # quantum Rabi model Hamiltonian (takes longer time to simulate)
H_c = g * (a * b' + a' * b)                                    # Jaynes-Cummings model Hamiltonian (faster to simulate)

# Total H
H = H_q + H_r + H_c

@printf("Dressed Qubit Freq. (Analytical): %.4f GHz\n", omega_q_dressed(ro_params) / 2π)
@printf("Dressed Resonator Freq. (Analytical): %.4f GHz\n", omega_r_dressed(ro_params) / 2π)
@printf("Dispersive shift (Analytical): %.3f MHz\n", chi(ro_params) / 2π / 0.001)

r = eigenstates(H)

function dressed_state(s::QuantumObject)
    """
    Return the eigenenergy and eigenstate of the Hamiltonian that has the largest overlap with the state `s`.
    """
    eig_energies, eig_states = eigenstates(H)

    overlaps = [abs(fidelity(s, ds)) for ds in eig_states]

    s_ind = argmax(overlaps)
    eig_energies[s_ind], eig_states[s_ind]
end

E_g0, _ = dressed_state(tensor(basis(N_q, 0), tensor(basis(N_r, 0))))
E_g1, _ = dressed_state(tensor(basis(N_q, 0), tensor(basis(N_r, 1))))
E_e0, _ = dressed_state(tensor(basis(N_q, 1), tensor(basis(N_r, 0))))
E_e1, _ = dressed_state(tensor(basis(N_q, 1), tensor(basis(N_r, 1))))

omega_q_dressed_numerical = (E_e0 - E_g0)   # qubit frequency when the resonator is not occupied
omega_r_g_numerical = (E_g1 - E_g0)   # readout resonator frequency when qubit is in the ground state
omega_r_e_numerical = (E_e1 - E_e0)   # readout resonator frequency when qubit is in the excited state

@printf("Dressed Qubit Freq. (Numerical): %.4f GHz\n", omega_q_dressed_numerical / 2π)
@printf("Resonator Freq. w/ Qubit in |g> (Numerical): %.4f GHz\n", omega_r_g_numerical / 2π)
@printf("Resonator Freq. w/ Qubit in |e> (Numerical): %.4f GHz\n", omega_r_e_numerical / 2π)
@printf("Dressed Resonator Freq. (Numerical): %.4f GHz\n", 0.5 * (omega_r_g_numerical + omega_r_e_numerical) / 2π)
@printf("Dispersive Shift (Numerical): %.4f MHz\n", (omega_r_e_numerical - omega_r_g_numerical) / 2e-3 / 2π)

@printf("Purcell-limited Lifetime (Analytical): %.3f μs\n", t1_purcell(ro_params) * 1e-3)
@printf("Purcell Decay Rate (Analytical): %.3f kHz\n", gamma1_purcell(ro_params) / (2π * 1e-6))

# Find dressed qubit excited state and the ground state
_, psi_e0_dressed = dressed_state(tensor(basis(N_q, 1), basis(N_r, 0)))
_, psi_g0_dressed = dressed_state(tensor(basis(N_q, 0), basis(N_r, 0)))

# Projectors onto the dressed qubit excited state and the ground state
proj_e0_dressed = psi_e0_dressed * psi_e0_dressed'
proj_g0_dressed = psi_g0_dressed * psi_g0_dressed'

tlist0 = collect(StepRange(0, 50, 50000))   # list of times for time evolution, in units of ns

# cannot run
result0 = mesolve(
    H, psi_e0_dressed, tlist0, [sqrt(kappa) * a],  # no collapse operator for qubit
    e_ops=[proj_g0_dressed, proj_e0_dressed]
)

• My observation is that increasing maxiters in the solver option does help in reducing the failure, but still the execution takes infinite amount of time. Also, reducing the range of sweep to sth like StepRange(0, 50, 5000) runs ok. This seems to be related to ODE integration error building up.
• Changing the alg option didn't seem to help much. Also, I noticed that options like lsoda imported from LSODA.jl cannot be handled by the mesolver due to type mismatch. In the case of QuTiP where things worked fine, I was using the default adams integrator with order 12.

[eunjongkim]

FYI, here are my version info

QuTiP:

QuTiP: Quantum Toolbox in Python
================================
Copyright (c) QuTiP team 2011 and later.
Current admin team: Alexander Pitchford, Nathan Shammah, Shahnawaz Ahmed, Neill Lambert, Eric Giguère, Boxi Li, Jake Lishman, Simon Cross, Asier Galicia, Paul Menczel, and Patrick Hopf.
Board members: Daniel Burgarth, Robert Johansson, Anton F. Kockum, Franco Nori and Will Zeng.
Original developers: R. J. Johansson & P. D. Nation.
Previous lead developers: Chris Granade & A. Grimsmo.
Currently developed through wide collaboration. See https://github.com/qutip for details.

QuTiP Version:      5.1.0
Numpy Version:      2.2.3
Scipy Version:      1.15.2
Cython Version:     3.0.12
Matplotlib Version: 3.10.1
Python Version:     3.12.9
Number of CPUs:     24
BLAS Info:          Generic
INTEL MKL Ext:      None
Platform Info:      Windows (AMD64)
Installation path:  C:\Users\ekim7\anaconda3\envs\qutip-env\Lib\site-packages\qutip
================================================================================
Please cite QuTiP in your publication.
================================================================================
For your convenience a bibtex reference can be easily generated using `qutip.cite()`

QuantumToolbox:

 QuantumToolbox.jl: Quantum Toolbox in Julia
≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡≡
Copyright © QuTiP team 2022 and later.
Current admin team:
    Alberto Mercurio and Yi-Te Huang

Package information:
====================================
Julia              Ver. 1.11.4
QuantumToolbox     Ver. 0.29.1
SciMLOperators     Ver. 0.3.13
LinearSolve        Ver. 3.7.2
OrdinaryDiffEqCore Ver. 1.22.0

System information:
====================================
OS       : Windows (x86_64-w64-mingw32)
CPU      : 24 × 13th Gen Intel(R) Core(TM) i7-13700
Memory   : 31.739 GB
WORD_SIZE: 64
LIBM     : libopenlibm
LLVM     : libLLVM-16.0.6 (ORCJIT, alderlake)
BLAS     : libopenblas64_.dll (ilp64)
Threads  : 1 (on 24 virtual cores)

[albertomercurio]

Hello @eunjongkim,

Thank you for reporting this issue.

The problem was related to the fact that QuTiP automatically removes small elements from the quantum objects. Thus, the initial state in qutip is more clean than the one in QuantumToolbox. This makes the ODE in QuantumToolbox more unstable.

To fix it, you can just use tidyup! to the objects you want to clean, for example

_, psi_e0_dressed = dressed_state(tensor(basis(N_q, 1), basis(N_r, 0)))
_, psi_g0_dressed = dressed_state(tensor(basis(N_q, 0), basis(N_r, 0)))

tidyup!(psi_e0_dressed)
tidyup!(psi_g0_dressed)

# Projectors onto the dressed qubit excited state and the ground state
proj_e0_dressed = psi_e0_dressed * psi_e0_dressed'
proj_g0_dressed = psi_g0_dressed * psi_g0_dressed'

tlist0 = collect(StepRange(0, 50, 50000))   # list of times for time evolution, in units of ns

# cannot run
result0 = mesolve(
    H, psi_e0_dressed, tlist0, [sqrt(kappa) * a],  # no collapse operator for qubit
    e_ops=[proj_g0_dressed, proj_e0_dressed],
    maxiters=Int(1e6),
)

I hope that this fixes your issue.

[albertomercurio]

Moreover, you don't need to set such a fine time spacing, as the internal ODE solver automatically chooses them. You can even set a two-points time list in principle. The number of points of tlist only determines the number of points for the expect field of the output.

If you want to specify additional points at which the solver should stop, you can add the tstops keyword argument.

[eunjongkim]

Hello @albertomercurio your suggestion fixes the issue. Thanks for helping out and saving a lot of time for me :-)

[albertomercurio]

Glad to have fixed it.

I’ll close this issue if there are no other related problems.