OrdinaryDiffEq fails to preserve norm in simple Schrödinger equation

[goerz]

I’m having an issue with OrdinaryDiffEq producing obviously wrong results (up to me doing something stupid, i.e., using it wrong). Consider the following MWE:

using OrdinaryDiffEq: OrdinaryDiffEq as ODE, ODEProblem, solve as ode_solve
using Plots

function f_ode(u, p, t)
	Δ = 1.0
	α = 4.6
	X = -1im * Float64[
         0.5 * α * t          Δ
               Δ       -0.5 * α * t
    ]
    return X * u
end

Ψ = [-0.999998+0.0im, 0.0021739+0.0im]  # eigenstate at t = -100

tlist = collect(range(-100.0, 100; length=1001))

ode_problem_direct = ODEProblem(ODE.ODEFunction(f_ode), Ψ, (tlist[begin], tlist[end]))

ode_sol_direct = ode_solve(ode_problem_direct, ODE.Tsit5(); abstol=1e-12, saveat=tlist)

popdyn_ode_direct = reduce(hcat, map(Ψ -> abs2.(Ψ), ode_sol_direct.u))

plot(tlist, popdyn_ode_direct'; label=["|⟨Ψ(t)|ϕ₀⟩|²" "|⟨Ψ(t)|ϕ₁⟩|²"], legend=:right, title="Projection onto canonical basis states", xlabel="time", ylabel="population")

That should be solving

\frac{\partial}{\partial t}|\Psi\rangle = -i H |\Psi\rangle

with the Hermitian

H = \begin{pmatrix} \frac{\alpha}{2} t & \Delta \\ \Delta & -\frac{\alpha}{2} t \end{pmatrix}\,.

Whatever the actual dynamics are, because H is Hermitian, the resulting dynamics should be unitary, i.e., this should be preserving the norm of |\Psi\rangle. The MWE does not:

This works fine using my own solver, which gives me this correct-looking result:

What am I doing wrong here with using OrdinaryDiffEq?

[JonasWickman]

The most immediate fix is to lower the relative tolerance:

ode_sol_direct = ode_solve(ode_problem_direct, ODE.Tsit5(); abstol=1e-12, reltol = 1e-10, saveat=tlist)

Typically, although I don’t think it affects the accuracy, there’s also no need to include the saveat = tlist. You can get get the values of the solution by interpolation by calling ode_sol_direct(t) after you’ve solved it.

This is not my area, but Tsit5() is probably not the best choice for solver for this problem, check the “Probelm Types” heading in the left bar of DifferentialEquations.jl: Efficient Differential Equation Solving in Julia · DifferentialEquations.jl . I don’t remember Quantum Mechanics stuff well enough to figure out if any of those types would match up with your system.

[ChrisRackauckas]

Standard RK methods are not norm preserving. As mentioned above, you can lower the tolerances and that should work for sufficiently short integrations. If you have long time integrations, you can use the symplectic Runge-Kutta methods or the methods designed for non-autonomous linear systems.

Dynamical, Hamiltonian, and 2nd Order ODE Solvers · DifferentialEquations.jl

Non-autonomous Linear ODE / Lie Group ODE Solvers · DifferentialEquations.jl

[Mason]

Indeed, I found while playing around with this that the tstops here actually hurt the accuracy, and asking for an abstol also hurt accuray (rather than a reltol).

In addition, Vern8() tends to be a more accurate solver for this sort of problem since it’s symplectic and thus will keep the Hamiltonian on the symplectic manifold, conserving various properties. (edit: you can of course do better with the sorts of solvers Chris mentioned above, but Vern8 does better than Tsit5 here and is easier to switch out).

Here’s what I see for the total norm of the wave-function over time after making some adjustments:

using OrdinaryDiffEq: OrdinaryDiffEq as ODE, ODEProblem, solve as ode_solve
using Plots

function f_ode(u, p, t)
	Δ = 1.0
	α = 4.6
	H = Hermitian([
        0.5 * α * t   Δ
        Δ            -0.5 * α * t
    ])
    -im * (H * u)
end

function tst(;tlist = range(-100, 100; length=1001), alg=ODE.Vern8())

    Ψ₀ = [-0.999998+0.0im, 0.0021739+0.0im]  # eigenstate at t = -100

    ode_problem_direct = ODEProblem(f_ode, Ψ₀, (tlist[begin], tlist[end]))

    Ψ = ode_solve(ode_problem_direct, alg;
                  reltol=1e-10)
    norms = map(t -> norm(Ψ(t)), tlist)
    plot(tlist, norms;
         label="⟨Ψ⟩²",
         legend=:right, title="Projection onto canonical basis states", xlabel="time", ylabel="population")

end 

which I think is a pretty impressive level of norm-conservation.

If it becomes a problem beyond that though, you can do tricks like using callbacks to re-normalize your solution.

[goerz]

Indeed, reltol fixes the problem. I don’t think I have a good grasp on setting abstol and reltol reliably. The way I’d usually think about errors in propagators would for the norm(Ψ - Ψ_exact) < err, where Ψ is the result of the propagator at final time T, Ψ_exact is the hypothetical analytic solution (in practice, obtained from a solver with all precisions cranked up to the max), and err is the error I’d like to prescribe (typically something like 1e-12). It’s a little non-ideal that there are two error parameters abstol, reltol that I have to tweak, and I have to get a better understanding of how to connect my err with abstol and reltol.

I suppose the key is the equation given in the “Basic step size control” of the CommonSolve.Solve docstring, but I have to have a closer look (Unless someone with more experience can give me some guidelines )

there’s also no need to include the saveat = tlist

Yeah, but there are reasons I need the states exactly on that time grid. But this is really just for the sake of the MWE / comparing with my normal solver. In my actual code, I use the integrator interface. But I do have to work on a specific time grid (because that’s how gradients in quantum control work).

This is not my area, but Tsit5() is probably not the best choice for solver for this problem

That’s interesting… I’m very much not an expert in general ODE solver methods. But to my knowledge when people have “traditionally” used ODE solvers for quantum dynamics (as an upgrade from something low-precision like Crank-Nicholson), they’ve typically used RK45. So I just sort of assumed Tsit5 would be the equivalent-but-better default.

@ChrisRackauckas

Standard RK methods are not norm preserving.

The title of the post may be a bit misleading, in that I care about “norm-preserving” only in so far as I care about “correct” in terms of the aforementioned norm(Ψ - Ψ_exact) < err, and that I know Ψ_exact to have norm 1. It’s not necessarily a problem that errors are reflected in a deviation of the norm. At least that’s pretty easy to detect, and thus possibly preferable to “grossly wrong, but with the correct norm”.

That’s not to say that I wouldn’t want to try solvers that inherently norm-preserving. My “workhorse propagator” of expanding exp(-i*H*dt) in Chebychev polynomials (the “correct” solution in my OP) is norm-preserving, after all. But that propagator is piecewise-constant, and the reason I ran into this whole issue was that I wanted to double-check with OrdinaryDiffEq that I wasn’t running into time-discretization errors. That is, compare with the solution for the time-continuous (adaptive time step) ODE solver. Which then, with the ill-advised approach of using OrdinaryDiffEq as a black box, gave me wildly inaccurate results.

@Mason

asking for an abstol also hurt accuracy

That’s weird, and a little concerning. But ultimately, I’m really just going to have to nail down how to reliably choose abstol and reltol based on my more general notion of “error”, and then see what is a good default ODE method for quantum dynamics.

This particular ODE might also be a bit pathological. I’m actually probing here a parameter regime where the adiabatic Landau-Zener transition (where everything would evolve very smoothly) breaks down. Still, it would be good to have a reliable ODE solver in my toolbox that doesn’t require much handholding.

I’ll definitely have a look at Vern8 and the other solvers Chris mentioned.

Thanks!

[Mason]

I was curious and tried this out using Magnus solvers and fixed time grid as you mentioned you needed that, and it seems to have performed very well:

using OrdinaryDiffEq, SciMLOperators, Plots

function tst2(;)
    tlist = range(-100, 100; length=1001)
    
    Ψ₀ = [-0.999998+0.0im, 0.0021739+0.0im]
    
    function update_func!(W, u, p, t)
        Δ = 1.0
	    α = 4.6
        W[1,1] = -im*(0.5 * α * t)
        W[2,1] = -im*(Δ)
        W[1,2] = -im*(Δ)
        W[2,2] = +im*(0.5 * α * t)
    end
    W = MatrixOperator(zeros(ComplexF64, 2, 2); update_func!) # -im * H
    
    prob = ODEProblem(W, Ψ₀, (first(tlist), last(tlist)))
    
    Ψ = solve(prob, MagnusMidpoint(), dt=step(tlist), saveat=tlist)

    popdyn_ode_direct = reduce(hcat, map(t -> abs2.(Ψ(t)), tlist))
    p = plot(tlist, popdyn_ode_direct'; label=["|⟨Ψ(t)|ϕ₀⟩|²" "|⟨Ψ(t)|ϕ₁⟩|²"], legend=:right, title="Projection onto canonical basis states", xlabel="time", ylabel="population")
    plot!(tlist, norm.(Ψ.(tlist)); label="|⟨Ψ(t)|Ψ(t)⟩|")
    p
end

Here’s a zoom-in on the drift of |⟨Ψ(t)|Ψ(t)⟩| and it is very precise:

This solver is also extremely fast so that’s nice (8x faster than Tsit5() for this example)

[Oscar_Smith]

So the key point about the 2 tolerances is that the step gets accepted if either passes. The easiest way to use this intuitiively is to set reltol=0 and then control your tolerance purely through abstol. The reason this isn’t the default is that abstol isn’t scale invariant so it can only be set well by the user. If you are looking for a more general tolerance setting, you should control your tolerance via reltol and set abstol to some value small enough that it’s able to be considered the same as 0.

[stevengj]

@Oscar_Smith

The easiest way to use this intuitiively is to set reltol=0 and then control your tolerance purely through abstol.

If the user specifies a positive abstol but not a reltol, it would be nice if the reltol were set to zero automatically. (This is how Base.isapprox works.)

[danielwe]

@goerz

It’s a little non-ideal that there are two error parameters abstol, reltol that I have to tweak, and I have to get a better understanding of how to connect my err with abstol and reltol

Firstly, both abstol and reltol are stepwise, not global. They bound the acceptable error for a single time step, but larger errors may accumulate over the course of the simulation.

For each step, abstol is a tolerance on the next value of norm(Ψ - Ψ_exact) conditional on the current Ψ being exact, while reltol is a tolerance on norm(Ψ - Ψ_exact) / norm(Ψ). Of course, you don’t actually know Ψ_exact, but if anything the error estimates lean toward the pessimistic side, so this is usually quite robust. (Specifically, for an order n algorithm the error is estimated by comparing the proposed step to one computed using an order n - 1 algorithm. This can really be thought of as a heuristic error estimate for the order n - 1 algorithm, and should thus be an overestimate for the order n algorithm unless something pathological is going on.)

abstol specifies what zero means to you in absolute terms, while reltol specifies the number of accurate digits you demand. If you know the typical magnitude of your state—say, the typical or maximum value of norm(Ψ) during the simulation—it’s natural to set abstol = scale * reltol. I think matlab ode solvers can even do this automatically: you specify a reltol, and the integrator keeps track of the largest state encountered and updates abstol accordingly. However, if the norm of your states is fairly constant, at least in order of magnitude terms, you might as well only specify reltol while keeping abstol = 0.

For quantum states where norm(Ψ) = 1, the distinction doesn’t really matter and abstol and reltol are interchangeable. In this case, I’d set abstol to zero and experiment with reltol.

[photor]

Your H is not Hermitian, unless alpha is purely imaginary.

[stevengj]

@photor

Your H is not Hermitian, unless alpha is purely imaginary.

I think X = -iH in the code.

[photor]

Still not OK, because Delta is real.

[stevengj]

H is real-symmetric. I think you are mis-reading something, @photor.

[goerz]

@stevengj

H is real-symmetric.

Yes, exactly. Which entails H= H^{\dagger}, that is, H_{ij} = H_{ji}^* = H_{ji}. And then U = \exp(-i H dt) is unitary, which is why the time evolution should be norm-preserving. That part, as the standard Schrödinger equation (and U being the standard time evolution operator), is definitely not what’s causing any problems

[Mason]

The Schrödinger equation is

{\partial \psi \over \partial t} = -i H \psi

His Hamiltonian is Hermitian, but one needs to absorb the factor of -i into it to get generator of time translations.

In ODEProblem, you need to provide an f such that

{\partial \psi \over \partial t} = f(\psi, p, t)