Mismatch between times and states in sesolve with callbacks

[LorenzoFioroni]

Bug Description

When using a callback function during time evolution with mesolve, the ODE integrator may evaluate the system at additional intermediate time points beyond those explicitly specified in the input tlist. Consequently, the solution object can contain more state vectors than entries in tlist. However, the times attribute of the solution object is not updated to reflect these additional evaluations. As a result, there may be a mismatch between the number of recorded states and the number of corresponding time points, making it unclear which time each state is associated with.

Code to Reproduce the Bug

using QuantumToolbox
using DifferentialEquations

H_tuple = (
    sum(multisite_operator(Val(3), i => sigmaz()) for i ∈ 1:3),
    ((multisite_operator(Val(3), i => sigmax()), (p, t) -> p.g[i]) for i ∈ 1:3)...
)
H_t = QobjEvo(H_tuple)
    
ψ0 = basis(8, 0, dims=(2,2,2))
t_list = collect(LinRange(0.0, 10.0, 100))

quench_disorder!(integrator) = begin
    integrator.p.g .= randn(3)
    return nothing
end
quench_callback = PeriodicCallback(quench_disorder!, 0.25)
initial_disorder = randn(3) 

problem = sesolveProblem(H_t, ψ0, t_list, params=(g=initial_disorder,), callback=quench_callback)
solution = sesolve(problem)

println("Number of states: ", length(solution.states))
println("Number of times: ", length(solution.times))

Code Output

Number of states: 178
Number of times: 100

Expected Behaviour

One would expect the solution object to be constructed using the states and corresponding time points as returned by the integrator. However, the current behavior is inconsistent: while the states are taken from the times at which the integrator evaluates the system, the times attribute of the solution object is retained as the original tlist provided to mesolve.

I think the problem is in line 208 from the snipped below. The times are initialised to prob.time instead of to those actually visited by the integrator.

QuantumToolbox.jl/src/time_evolution/mesolve.jl

Lines 207 to 215 in 0066cdb

	return TimeEvolutionSol(
	prob.times,
	ρt,
	_get_expvals(sol, SaveFuncMESolve),
	sol.retcode,
	sol.alg,
	NamedTuple(sol.prob.kwargs).abstol,
	NamedTuple(sol.prob.kwargs).reltol,
	)

Your Environment

Package information:
====================================
Julia              Ver. 1.11.5
QuantumToolbox     Ver. 0.32.1
SciMLOperators     Ver. 0.4.0
LinearSolve        Ver. 3.18.2
OrdinaryDiffEqCore Ver. 1.26.2

Additional Context

No response

[ytdHuang]

The sol.states should have the same length as sol.times.

I think it is better to change the way we store the states instead of changing the sol.times.

That is, the output of the example provided by @LorenzoFioroni should be:

Number of states: 100
Number of times: 100

If users have self-defined callbacks at different time point, they should store the state by using

Qobj(copy(integrator.u))

in their self-defined callback function.

@albertomercurio what do you think ?
I suggest we handle the states together with expect-callback, instead of convert all of them after solving ODE.

[ytdHuang]

But I'm quite surprised that we have this inconsistancy...

Our default ODE kwargs should be save_everystep = false and saveat = t_list in the example provided by @LorenzoFioroni

So the sol.u from ODE solver should only store the states corresponding to t_list.

Would this be a bug from SciML-side ?

[ytdHuang]

I did a small try with only OrdinaryDiffEq and DiffEqCallbacks:

using SciMLOperators
using OrdinaryDiffEq, DiffEqCallbacks

N = 10
A = MatrixOperator(rand(ComplexF64, N, N))
b = rand(ComplexF64, N)

tstart = 0
tend = 10
tlist = collect(tstart:1:tend)

kwargs = (
    saveat = tlist,
    save_everystep = false,
    save_end = true
)

# Define problem without extra callback
prob = ODEProblem(A, b, (tstart, tend); kwargs...)
sol = OrdinaryDiffEq.solve(prob, OrdinaryDiffEq.Tsit5())
println(length(sol.u))

# Define problem with extra callback
dummy!(integrator) = nothing
dummy_callback = PeriodicCallback(dummy!, 0.25)
prob = ODEProblem(A, b, (tstart, tend); callback=dummy_callback, kwargs...)
sol = OrdinaryDiffEq.solve(prob, OrdinaryDiffEq.Tsit5())
println(length(sol.u))

output:

11
80

[ytdHuang]

@LorenzoFioroni

I noticed that the callback saves are treated independently.

You can set save_positions in the callback to change that:

quench_callback = PeriodicCallback(quench_disorder!, 0.25; save_positions=(false, false))

[LorenzoFioroni]

I see, thank you.
This definitely provides a workaround in cases where the state right before or after the callback is not needed - as it is the case for me. Still, the inconsistency regarding the way QuantumToolbox treats times and states remains.
I still believe that either both times and states should be taken from the integration solution, or somehow the states should be filtered, such that at each specified time there is one associated state.

In other words: If save_positions=(true, true), the integrator's time and states fields are consistent with each other. The inconsistency arises from creating the TimeEvolutionSol object with the times given in input and the states taken from the integrator.

[ytdHuang]

The idea is to make tlist, expect, and states all consistent (align with qutip). Instead of making it consistent with "integration solution".

Since we store the expect using an internal callback corresponding to tlist, so I think we should also handle the storage of states within the same callback instead of at the end (from sol.u).

@albertomercurio What do you think ?

[albertomercurio]

I think it is useful to give the opportunity to define a callback that saves the states before and/or after the affect. If the user is using a custom callback, it's because the user is more or less advanced. So we can just return sol.t and Qobj(sol.u) if a custom callback is defined.

QuTiP doesn't have a callback (I think), so this behavior is already different. The standard user will never see this difference.

If it is ok for you, I could open a PR that does this.

[ytdHuang]

But that makes it inconsistent with sol.expect ...

[ytdHuang]

@LorenzoFioroni
We have release the new version (v0.33) with sol.times still corresponds to original user-specified tlist.

The SciML sol.t is now stored in a new field called times_states, that is, it can be accessed through our time evolution solution: sol.times_states.

[LorenzoFioroni]

This way we can have our cake and it to too, thanks!