`smesolve` result inconsistent with one from qutip

[eunjongkim]

Describe the Issue!

Hi, I was trying to transition from QuTiP to QuantumToolbox.jl. So far, the speedup in mesolve was quite amazing, but I got stuck at an issue where the smesolve outcome is not consistent with the result from qutip. I wonder if any of you could help me figure this out.

QuTiP case

script (takes about 1hr on my desktop computer):

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')

from scipy import special
from scipy.optimize import curve_fit


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

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

# 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

# time-dependent coefficient for drive
def Hrd_env(t: np.ndarray, args: dict) -> np.ndarray:
    """
    Envelope function for the readout drive
    """
    epsilon = args['epsilon']   # readout pulse amplitude
    T = args['T']   # readout pulse duration
    return epsilon * (0 <= t) * (t <= T)

def Hrd_coeff(t: np.ndarray, args: dict) -> np.ndarray:
    omega_d = args['omega_d']   # readout drive frequency
    return 2 * Hrd_env(t, args) * np.cos(omega_d * t)

def Hrd_coeff_m(t: np.ndarray, args: dict) -> np.ndarray: 
    omega_d = args['omega_d']   # readout drive frequency
    return Hrd_env(t, args) * np.exp(-1j * omega_d * t)

def Hrd_coeff_p(t: np.ndarray, args: dict) -> np.ndarray:
    omega_d = args['omega_d']   # readout drive frequency
    return Hrd_env(t, args) * np.exp(+1j * omega_d * t)

ro_len = 400                          # readout pulse length, in units of ns
integration_time = 500                     # readout integration length, in units of ns
dt = 1                                      # time step for simulation, in units of ns

# readout signal integration settings
tlist = np.arange(0, integration_time, dt)     # times for simulation
drive_phase = np.exp(1.0j * omega_d * tlist)   # phase of the readout resonator drive

# t_analysis_step = len(tlist) // 4               # dividing len(t_ro) and discard the remainder. NOTE: integration is performed at a temporal resolution of a clock cycle = 4ns
# tlist_used = tlist[0::t_analysis_step]          # time samples used in the analysis: t[0], t[t_analysis_step-1], t[2*t_analysis_step-1], ...

# resonator drive amplitude: how is this chosen?
ro_amp = 0.0015 * (2 * np.pi)

# arguments to be passed onto the solver
solver_args = {'epsilon': ro_amp, 'T': ro_len, 'omega_d': omega_d}

### smesolve part (DEBUG)
## choose ntraj appropriately to avoid running out of memory 
ntraj = 12 ## choose ntraj accordingly to get the # of points you want per each simulation

## choose batches
batch_num = 1 ## choose batch accordingly to not run out of memory each time

c_ops = []
sc_ops = [np.sqrt(kappa / 2) * a, -1j * np.sqrt(kappa / 2) * a]       # stochastic evolution operator (readout signal)

def run_batch(batch_i: int):
    """Function to implement readout simulation for each batch """
    np.random.seed(batch_i)  ## choose a different but fixed seed each batch so it's reproducible

    ro_trace = np.zeros((2, ntraj, len(tlist) - 1), dtype=complex)

    for prep_i, init_q_state in enumerate([0, 1]):
        psi = tensor(basis(N_q, init_q_state), basis(N_r, 0))

        result = smesolve(
            [H, [a + a.dag(), Hrd_coeff]], psi, tlist,
            c_ops=c_ops, sc_ops=sc_ops,  ## measurement op
            e_ops=[], args = solver_args, ntraj=ntraj,
            options = {
                "store_measurement": "end",
                "store_states": False,
                "progress_bar": "tqdm",
                "map": "loky",
                "num_cpus" : 12,
                "dt" : 2e-4 # make this small enough to avoid producing nans
            },  ## false if many ntraj or memory runs out
        )
        print(result.measurement[0].shape)
        for traj_i in range(ntraj):
            ro_trace[prep_i, traj_i, :] = (
                result.measurement[traj_i][0, :].real
                + 1j * result.measurement[traj_i][1, :].real
            )

    return ro_trace

def remove_rotating_phase(ro_trace):
    ro_trace_env = np.zeros((2, ntraj, len(tlist) - 1), dtype=complex)
    for prep_i, init_q_state in enumerate([0, 1]):
        for traj_i in range(ntraj):
            ro_trace_env[prep_i, traj_i, :] = ro_trace[prep_i, traj_i, :] * drive_phase[1:]
    return ro_trace_env


def demodulate(ro_trace_env):
    weight = np.ones(len(tlist) - 1)
    integrated_signal = np.trapezoid(ro_trace_env * weight, tlist[1:], axis=-1)
    return integrated_signal

ro_trace = run_batch(0)
ro_trace_env = remove_rotating_phase(ro_trace)
integrated_signal = demodulate(ro_trace_env)


# normalization of integrated signals
integrated_signal_norm = (
    integrated_signal / np.sqrt(ro_params.kappa) / np.trapezoid(np.ones(len(tlist) - 1), tlist[1:])
)

# mean and standard deviation
means = np.mean(integrated_signal_norm, axis=1)
stdevs = np.std(integrated_signal_norm, axis=1)

# draw figure
# plot IQ blobs in the complex plane
plt.figure()
ax = plt.gca()
for init_q_state in [0, 1]:
    # single-shot readout outcomes
    plt.scatter(
        integrated_signal_norm[init_q_state].real,
        integrated_signal_norm[init_q_state].imag,
        label=f'Prep. $|{init_q_state}\\rangle$'
    )

    # 1-sigma circle centered at each mean
    c = plt.Circle(
        (means[init_q_state].real, means[init_q_state].imag),
        stdevs[init_q_state],
        edgecolor=f'C{init_q_state}', alpha=0.5, lw=2, facecolor='None'
    )
    ax.add_patch(c)

plt.axis('equal')
plt.xlabel(r'$I$')
plt.ylabel(r'$Q$')

plt.legend()
plt.show()

Result

Two blobs separated well, with I and Q values consistent with theory expectations.

QuTiP version info

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()`

equivalent Julia script (takes about 11 min on my desktop):

using QuantumToolbox, Trapz, Statistics, Random, DifferentialEquations
using CairoMakie, Printf

## transmon readout parameters
mutable struct TransmonROParams
    omega_q::Float64
    anharmonicity::Float64
    omega_r::Float64
    g::Float64
    kappa::Float64
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_r_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_r_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)
omega_d = omega_r_dressed(ro_params)       # readout drive frequency

# 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

# time-dependent coefficient for drive
Hrd_env(p::NamedTuple, t) = p.ϵ * (0 .<= t) .* (t .<= p.T)
Hrd_coeff(p::NamedTuple, t) = 2 * Hrd_env(p, t) .* cos.(p.ωd * t)
Hrd_coeff_m(p::NamedTuple, t) = Hrd_env(p, t) .* exp.(- im * p.ωd * t)
Hrd_coeff_m(p::NamedTuple, t) = Hrd_env(p, t) .* exp.(+ im * p.ωd * t);

# Readout Parameters
ro_len = 400                          # readout pulse length, in units of ns
integration_time = 500                     # readout integration length, in units of ns
dt = 1                                      # time step for simulation, in units of ns

# readout signal integration settings
tlist = collect(0:dt:integration_time)     # times for simulation
drive_phase = exp.(im * omega_d * tlist)   # phase of the readout resonator drive

# resonator drive amplitude
ro_amp = 0.0015 * 2π

# arguments to be passed onto the solver
solver_params = (ϵ = ro_amp, T = ro_len, ωd = omega_d)

### smesolve part (DEBUG)
## choose ntraj appropriately to avoid running out of memory 
ntraj = 12 ## choose ntraj accordingly to get the # of points you want per each simulation

## choose batches
batch_num = 1 ## choose batch accordingly to not run out of memory each time

c_ops = []
sc_ops = [sqrt(ro_params.kappa / 2) * a, -im * sqrt(ro_params.kappa / 2) * a]

function run_batch(batch_i:: Int; kwargs...)
    Random.seed!(batch_i)

    ro_trace = zeros(ComplexF64, 2, ntraj, length(tlist) - 1)
    for (prep_i, init_q_state) in enumerate([0, 1])
        psi = tensor(
            basis(N_q, init_q_state), basis(N_r, 0)
        )

        result = smesolve(
            H + QobjEvo(a + a', Hrd_coeff), psi, tlist, c_ops, sc_ops;
            params = solver_params, ntraj = ntraj, store_measurement = Val(true),
            kwargs...
        )
        for traj_i = 1:ntraj
            ro_trace[prep_i, traj_i, :] = real(result.measurement[1, traj_i, :]) + im * real(result.measurement[2, traj_i, :])
        end
    end
    ro_trace
end

function remove_rotating_phase(ro_trace::Array)
    ro_trace_env = zeros(ComplexF64, 2, ntraj, length(tlist) - 1)
    for (prep_i, state) in enumerate([0, 1])
        for traj_i = 1:ntraj
            ro_trace_env[prep_i, traj_i, :] = ro_trace[prep_i, traj_i, :] .* drive_phase[2:end]
        end
    end
    ro_trace_env
end

function demodulate(ro_trace_env::Array)
    weight = ones(ComplexF64, length(tlist) - 1)
    trapz(tlist[2:end], ro_trace_env, Val(3))
end

ro_trace = run_batch(0, maxiters = Int(1e7)) #; alg=DifferentialEquations.RKMilGeneral(), dt=2e-4, reltol=1e-5, abstol=1e-5)#; alg=DifferentialEquations.EM(), dt=5e-5, maxiters = Int(1e7))
ro_trace_env = remove_rotating_phase(ro_trace)
integrated_signal = demodulate(ro_trace_env)

# normalization of integrated signals
integrated_signal_norm = (
    integrated_signal / sqrt(ro_params.kappa) / trapz(tlist[2:end], ones(length(tlist) - 1))
)

# mean and standard deviation
means = Statistics.mean(integrated_signal_norm, dims=2)
stdevs = Statistics.std(integrated_signal_norm, dims=2)

# draw figure
fig = Figure()
ax = Axis(fig[1, 1], xlabel=L"$I$", ylabel=L"$Q$", aspect=DataAspect())
scatter!(real(integrated_signal_norm[1, :]), imag(integrated_signal_norm[1, :]), label = L"Prep. $|0\rangle$")
scatter!(real(integrated_signal_norm[2, :]), imag(integrated_signal_norm[2, :]), label = L"Prep. $|1\rangle$")
arc!(Point2f(real(means[1]), imag(means[1])), stdevs[1], -π, π)
arc!(Point2f(real(means[2]), imag(means[2])), stdevs[2], -π, π)
axislegend(ax)
display(fig)

Result

Note that not only the distribution, but the range of I and Q values are very different from QuTiP results.

Version info:

 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  : 12 (on 24 virtual cores)

It is true that I am using a time-dependent Hamiltonian that rotates very fast---but I am not sure if that is a problem. I have tried a few different stochastic solver options and parameters but wasn't successful.

[ytdHuang]

I'm not so familiar with the model you are considering @@

Could you start with an easier model ?

Like the minimum code that produces the inconsistancy.

[eunjongkim]

Thanks for your reply, @ytdHuang . Yes, I am working on creating a simpler model to describe the issue. This seems to be a mixture of a few different issues.

[albertomercurio]

Hi @eunjongkim,

Thank you for showing this issue.

I probably found the main origin of this. I have created a Pull Request here to try to fix it. Could you try to test your code with that version? You can do it by installing the specific version of that branch with

pkg> add https://github.com/albertomercurio/QuantumToolbox.jl/tree/dev/patch-2

If I run your code with that branch, I get this result now

If everything is ok, we can merge that PR and release a new version.

[ytdHuang]

@eunjongkim

If PR #453 fixes this issue, we would probably announce a patch release.

Remember that you need to re-install the standard release version of QuantumToolbox after you finish the testing with that branch.

[eunjongkim]

Hi @albertomercurio @ytdHuang ! I appreciate your prompt feedback.
I was able to obtain a very similar result with your , with maxiters=Int(1e7), adaptive=false and dt=2e-4

One remaining issue (or concern) is that the result seems to be somewhat inaccurate as I remove the keyword arguments adaptive=False and dt=1e-4 (adaptive time steps, I guess?).

I think this is ok, but I wonder if additional care could be taken when using adaptive time stepping?

[ytdHuang]

@eunjongkim

Do you mean that you need to specify adaptive=false and dt=2e-4 even using the branch made by @albertomercurio ?

Or what you are saying is that you can obtain the correct results using our standard release by specifying adaptive=false and dt=2e-4.

[eunjongkim]

@ytdHuang
I mean I still need to specify adaptive=false and dt=2e-4 in the branch made by @albertomercurio in order to get the correct result.

Note:

• I cannot obtain any physically meaningful result in the standard release, even if I add those keyword arguments
• I can obtain some physical result in the branch by @albertomercurio with the above algorithm options, but the result is slightly different when I remove those keyword arguments. That leads me to question whether adaptive time stepping may cause additional errors.

[albertomercurio]

We still need to find the best default options for the SDE solver that are a compromise between speed and accuracy.

In your case, I guess that you can instead just starting tuning abstol and reltol first. In my branch, the default values are 0.001 and 0.01, respectively.

This is the result using abstol=1e-3 and reltol=2e-3

Is it better?

Moreover, I have alse a few suggestions. I see that you don't have c_ops. In this case you could use ssesolve. We are planning to do it automatically in the future, if no c_ops are set. I don't know what QuTiP does in this case.

Also, if sc_ops is a single QuantumObject rather than a Vector. In that case, a specific solver for Diagonal noise is chosen. In your case you have two sc_ops, so the non-diagonal noise solver is used. Would the implementation of the heterodyne allow to set only one sc_ops?

[eunjongkim]

Hi @albertomercurio , I was able to reproduce your result with abstol=1e-3 and reltol=2e-3 and without specifying dt value. Yes, this made it better.

Also, I appreciate your suggestions on the functionalities. At this time I am not using c_ops but I plan to add this going forward. For this case though, I think using ssesolve, as you suggested, is a great idea.

Also, for me the implementation of heterodyne would for sure allow me to use sc_ops as a single QuantumObject rather than a Vector.

[albertomercurio]

Great. I will set the default tolerances to abstol=1e-3 and reltol=2e-3. I will then announce a new patch release.