Make time evolution solvers compatible with automatic differentiation

Project.toml

@@ -21,6 +21,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
+SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
@@ -53,6 +54,7 @@ Random = "1"
 Reexport = "1"
 SciMLBase = "2"
 SciMLOperators = "0.3"
+SciMLStructures = "1.5.0"
 SparseArrays = "1"
 SpecialFunctions = "2"
 StaticArraysCore = "1"

src/QuantumToolbox.jl

@@ -45,6 +45,7 @@ import SciMLOperators:
     IdentityOperator,
     update_coefficients!,
     concretize
+import SciMLStructures: isscimlstructure, ismutablescimlstructure, hasportion, canonicalize, replace, replace!, Tunable
 import LinearSolve: LinearProblem, SciMLLinearSolveAlgorithm, KrylovJL_MINRES, KrylovJL_GMRES
 import DiffEqBase: get_tstops
 import DiffEqCallbacks: PeriodicCallback, PresetTimeCallback, TerminateSteadyState
@@ -87,6 +88,7 @@ include("qobj/superoperators.jl")
 include("qobj/synonyms.jl")
 
 # time evolution
+include("time_evolution/time_evo_parameters.jl")
 include("time_evolution/time_evolution.jl")
 include("time_evolution/mesolve.jl")
 include("time_evolution/lr_mesolve.jl")

src/time_evolution/mcsolve.jl

@@ -270,6 +270,8 @@ function mcsolveProblem(
         jump_which = jump_which,
         jump_times_which_init_size = jump_times_which_init_size,
         jump_times_which_idx = Ref(1),
+        times = tlist, # Temporary fix
+        Hdims = H_eff_evo.dims, # Temporary fix
         params...,
     )
 
@@ -291,7 +293,7 @@ function mcsolveProblem(
         haskey(kwargs2, :callback) ? merge(kwargs2, (callback = CallbackSet(cb1, cb2, kwargs2.callback),)) :
         merge(kwargs2, (callback = CallbackSet(cb1, cb2),))
 
-    return sesolveProblem(H_eff_evo, ψ0, tlist; params = params, kwargs2...)
+    return sesolveProblem(H_eff_evo, ψ0, tlist; params = params, kwargs2...).prob # Temporary fix
 end
 
 function mcsolveProblem(
@@ -315,7 +317,7 @@ function mcsolveProblem(
         haskey(kwargs2, :callback) ? merge(kwargs2, (callback = CallbackSet(cb1, cb2, kwargs2.callback),)) :
         merge(kwargs2, (callback = CallbackSet(cb1, cb2),))
 
-    return sesolveProblem(H_eff_evo, ψ0, tlist; params = params, kwargs2...)
+    return sesolveProblem(H_eff_evo, ψ0, tlist; params = params, kwargs2...).prob # Temporary fix
 end
 
 @doc raw"""

src/time_evolution/sesolve.jl

@@ -1,39 +1,55 @@
 export sesolveProblem, sesolve
 
+# When e_ops is Nothing
 function _save_func_sesolve(integrator)
-    internal_params = integrator.p
-    progr = internal_params.progr
-
-    if !internal_params.is_empty_e_ops
-        e_ops = internal_params.e_ops
-        expvals = internal_params.expvals
+    next!(integrator.p.progr)
+    return u_modified!(integrator, false)
+end
 
+# When e_ops is a list of operators
+function _save_func_sesolve(integrator, e_ops, is_empty_e_ops)
+    expvals = integrator.p.expvals
+    progr = integrator.p.progr
+    if !is_empty_e_ops
         ψ = integrator.u
-        _expect = op -> dot(ψ, op, ψ)
+        _expect = op -> dot(ψ, get_data(op), ψ)
         @. expvals[:, progr.counter[]+1] = _expect(e_ops)
     end
-    next!(progr)
-    return u_modified!(integrator, false)
+    return _save_func_sesolve(integrator)
+end
+
+# Generate the callback depending on the e_ops type
+function _generate_sesolve_callback(e_ops::Nothing, tlist)
+    f = integrator -> _save_func_sesolve(integrator)
+    return PresetTimeCallback(tlist, f, save_positions = (false, false))
 end
 
-function _generate_sesolve_kwargs_with_callback(tlist, kwargs)
-    cb1 = PresetTimeCallback(tlist, _save_func_sesolve, save_positions = (false, false))
+function _generate_sesolve_callback(e_ops, tlist)
+    is_empty_e_ops = isempty(e_ops)
+    f = integrator -> _save_func_sesolve(integrator, e_ops, is_empty_e_ops)
+    return PresetTimeCallback(tlist, f, save_positions = (false, false))
+end
+
+function _merge_sesolve_kwargs_with_callback(kwargs, cb)
     kwargs2 =
-        haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(kwargs.callback, cb1),)) :
-        merge(kwargs, (callback = cb1,))
+        haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(kwargs.callback, cb),)) :
+        merge(kwargs, (callback = cb,))
 
     return kwargs2
 end
 
+# Multiple dispatch depending on the progress_bar and e_ops types
 function _generate_sesolve_kwargs(e_ops, progress_bar::Val{true}, tlist, kwargs)
-    return _generate_sesolve_kwargs_with_callback(tlist, kwargs)
+    cb = _generate_sesolve_callback(e_ops, tlist)
+    return _merge_sesolve_kwargs_with_callback(kwargs, cb)
 end
 
 function _generate_sesolve_kwargs(e_ops, progress_bar::Val{false}, tlist, kwargs)
     if e_ops isa Nothing
         return kwargs
     end
-    return _generate_sesolve_kwargs_with_callback(tlist, kwargs)
+    cb = _generate_sesolve_callback(e_ops, tlist)
+    return _merge_sesolve_kwargs_with_callback(kwargs, cb)
 end
 
 _sesolve_make_U_QobjEvo(H::QuantumObjectEvolution{<:MatrixOperator}) =
@@ -46,8 +62,9 @@ _sesolve_make_U_QobjEvo(H) = QobjEvo(H, -1im)
         ψ0::QuantumObject{DT2,KetQuantumObject},
         tlist::AbstractVector;
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::NamedTuple = NamedTuple(),
+        params::Union{NamedTuple, AbstractVector} = NamedTuple(),
         progress_bar::Union{Val,Bool} = Val(true),
+        inplace::Union{Val,Bool} = Val(true),
         kwargs...,
     )
 
@@ -63,8 +80,9 @@ Generate the ODEProblem for the Schrödinger time evolution of a quantum system:
 - `ψ0`: Initial state of the system ``|\psi(0)\rangle``.
 - `tlist`: List of times at which to save either the state or the expectation values of the system.
 - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`.
-- `params`: `NamedTuple` of parameters to pass to the solver.
+- `params`: `NamedTuple` or `AbstractVector` of parameters to pass to the solver.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
+- `inplace`: Whether to use the inplace version of the ODEProblem. The default is `Val(true)`.
 - `kwargs`: The keyword arguments for the ODEProblem.
 
 # Notes
@@ -83,8 +101,9 @@ function sesolveProblem(
     ψ0::QuantumObject{DT2,KetQuantumObject},
     tlist::AbstractVector;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::NamedTuple = NamedTuple(),
+    params::Union{NamedTuple,AbstractVector} = NamedTuple(),
     progress_bar::Union{Val,Bool} = Val(true),
+    inplace::Union{Val,Bool} = Val(true),
     kwargs...,
 ) where {DT1,DT2}
     haskey(kwargs, :save_idxs) &&
@@ -103,31 +122,23 @@ function sesolveProblem(
 
     if e_ops isa Nothing
         expvals = Array{ComplexF64}(undef, 0, length(tlist))
-        e_ops_data = ()
         is_empty_e_ops = true
     else
         expvals = Array{ComplexF64}(undef, length(e_ops), length(tlist))
-        e_ops_data = get_data.(e_ops)
         is_empty_e_ops = isempty(e_ops)
     end
 
-    p = (
-        e_ops = e_ops_data,
-        expvals = expvals,
-        progr = progr,
-        times = tlist,
-        Hdims = H_evo.dims,
-        is_empty_e_ops = is_empty_e_ops,
-        params...,
-    )
+    p = QuantumTimeEvoParameters(expvals, progr, params)
 
     saveat = is_empty_e_ops ? tlist : [tlist[end]]
     default_values = (DEFAULT_ODE_SOLVER_OPTIONS..., saveat = saveat)
     kwargs2 = merge(default_values, kwargs)
     kwargs3 = _generate_sesolve_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2)
 
     tspan = (tlist[1], tlist[end])
-    return ODEProblem{true,FullSpecialize}(U, ψ0, tspan, p; kwargs3...)
+    prob = ODEProblem{getVal(inplace),FullSpecialize}(U, ψ0, tspan, p; kwargs3...)
+
+    return QuantumTimeEvoProblem(prob, tlist, H_evo.dims)
 end
 
 @doc raw"""
@@ -137,8 +148,9 @@ end
         tlist::AbstractVector;
         alg::OrdinaryDiffEqAlgorithm = Tsit5(),
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::NamedTuple = NamedTuple(),
+        params::Union{NamedTuple, AbstractVector} = NamedTuple(),
         progress_bar::Union{Val,Bool} = Val(true),
+        inplace::Union{Val,Bool} = Val(true),
         kwargs...,
     )
 
@@ -155,8 +167,9 @@ Time evolution of a closed quantum system using the Schrödinger equation:
 - `tlist`: List of times at which to save either the state or the expectation values of the system.
 - `alg`: The algorithm for the ODE solver. The default is `Tsit5()`.
 - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`.
-- `params`: `NamedTuple` of parameters to pass to the solver.
+- `params`: `NamedTuple` or `AbstractVector` of parameters to pass to the solver.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
+- `inplace`: Whether to use the inplace version of the ODEProblem. The default is `Val(true)`.
 - `kwargs`: The keyword arguments for the ODEProblem.
 
 # Notes
@@ -177,22 +190,32 @@ function sesolve(
     tlist::AbstractVector;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::NamedTuple = NamedTuple(),
+    params::Union{NamedTuple,AbstractVector} = NamedTuple(),
     progress_bar::Union{Val,Bool} = Val(true),
+    inplace::Union{Val,Bool} = Val(true),
     kwargs...,
 ) where {DT1,DT2}
-    prob = sesolveProblem(H, ψ0, tlist; e_ops = e_ops, params = params, progress_bar = progress_bar, kwargs...)
+    prob = sesolveProblem(
+        H,
+        ψ0,
+        tlist;
+        e_ops = e_ops,
+        params = params,
+        progress_bar = progress_bar,
+        inplace = inplace,
+        kwargs...,
+    )
 
     return sesolve(prob, alg)
 end
 
-function sesolve(prob::ODEProblem, alg::OrdinaryDiffEqAlgorithm = Tsit5())
-    sol = solve(prob, alg)
+function sesolve(prob::QuantumTimeEvoProblem, alg::OrdinaryDiffEqAlgorithm = Tsit5())
+    sol = solve(prob.prob, alg)
 
-    ψt = map(ϕ -> QuantumObject(ϕ, type = Ket, dims = sol.prob.p.Hdims), sol.u)
+    ψt = map(ϕ -> QuantumObject(ϕ, type = Ket, dims = prob.dims), sol.u)
 
     return TimeEvolutionSol(
-        sol.prob.p.times,
+        prob.times,
         ψt,
         sol.prob.p.expvals,
         sol.retcode,

src/time_evolution/ssesolve.jl

@@ -184,11 +184,9 @@ function ssesolveProblem(
 
     if e_ops isa Nothing
         expvals = Array{ComplexF64}(undef, 0, length(tlist))
-        e_ops_data = ()
         is_empty_e_ops = true
     else
         expvals = Array{ComplexF64}(undef, length(e_ops), length(tlist))
-        e_ops_data = get_data.(e_ops)
         is_empty_e_ops = isempty(e_ops)
     end
 
@@ -205,16 +203,7 @@ function ssesolveProblem(
     D_l = map(op -> op + _ScalarOperator_e(op, -) * IdentityOperator(prod(dims)), sc_ops_evo_data)
     D = DiffusionOperator(D_l)
 
-    p = (
-        e_ops = e_ops_data,
-        expvals = expvals,
-        progr = progr,
-        times = tlist,
-        Hdims = dims,
-        is_empty_e_ops = is_empty_e_ops,
-        n_sc_ops = length(sc_ops),
-        params...,
-    )
+    p = (expvals = expvals, progr = progr, times = tlist, Hdims = dims, n_sc_ops = length(sc_ops), params...)
 
     saveat = is_empty_e_ops ? tlist : [tlist[end]]
     default_values = (DEFAULT_SDE_SOLVER_OPTIONS..., saveat = saveat)

src/time_evolution/time_evo_parameters.jl

@@ -0,0 +1,54 @@
+# This function should be implemented after Julia v1.12
+Base.@constprop :aggressive function _delete_field(a::NamedTuple{an}, field::Symbol) where {an}
+    names = Base.diff_names(an, (field,))
+    return NamedTuple{names}(a)
+end
+
+struct QuantumTimeEvoParameters{TE<:AbstractMatrix,PT<:ProgressBar,ParT}
+    expvals::TE
+    progr::PT
+    params::ParT
+
+    function QuantumTimeEvoParameters(expvals, progr, params)
+        _expvals = expvals
+        _progr = progr
+        _params = params
+
+        # We replace the fields if they are aleady in the `params` struct
+        # Then, we remove them from the `params` struct
+        if :expvals ∈ fieldnames(typeof(_params))
+            _expvals = _params.expvals
+            _params = _delete_field(_params, :expvals)
+        end
+        if :progr ∈ fieldnames(typeof(_params))
+            _progr = _params.progr
+            _params = _delete_field(_params, :progr)
+        end
+
+        return new{typeof(_expvals),typeof(_progr),typeof(_params)}(_expvals, _progr, _params)
+    end
+end
+
+#=
+By defining a custom `getproperty` method for the `QuantumTimeEvoParameters` struct, we can access the fields of `params` directly.
+=#
+function Base.getproperty(obj::QuantumTimeEvoParameters, field::Symbol)
+    if field ∈ fieldnames(typeof(obj))
+        getfield(obj, field)
+    elseif field ∈ fieldnames(typeof(obj.params))
+        getfield(obj.params, field)
+    else
+        throw(KeyError("Field $field not found in QuantumTimeEvoParameters or params."))
+    end
+end
+
+#=
+It also supports `params` as a `Vector`, so we implement the `getindex` method for the `QuantumTimeEvoParameters` struct.
+=#
+Base.getindex(obj::QuantumTimeEvoParameters, i::Int) = getindex(obj.params, i)
+
+Base.length(obj::QuantumTimeEvoParameters) = length(obj.params)
+
+function Base.merge(a::QuantumTimeEvoParameters, b::NamedTuple)
+    return QuantumTimeEvoParameters(a.expvals, a.progr, merge(a.params, b))
+end

src/time_evolution/time_evolution.jl

@@ -5,6 +5,12 @@ export liouvillian_floquet, liouvillian_generalized
 const DEFAULT_ODE_SOLVER_OPTIONS = (abstol = 1e-8, reltol = 1e-6, save_everystep = false, save_end = true)
 const DEFAULT_SDE_SOLVER_OPTIONS = (abstol = 1e-2, reltol = 1e-2, save_everystep = false, save_end = true)
 
+struct QuantumTimeEvoProblem{PT,TT<:AbstractVector,DT<:AbstractVector}
+    prob::PT
+    times::TT
+    dims::DT
+end
+
 @doc raw"""
     struct TimeEvolutionSol
 

test/core-test/time_evolution.jl

@@ -16,7 +16,7 @@
         sol2 = sesolve(H, psi0, t_l, progress_bar = Val(false))
         sol3 = sesolve(H, psi0, t_l, e_ops = e_ops, saveat = t_l, progress_bar = Val(false))
         sol_string = sprint((t, s) -> show(t, "text/plain", s), sol)
-        @test prob.f.f isa MatrixOperator
+        @test prob.prob.f.f isa MatrixOperator
         @test sum(abs.(sol.expect[1, :] .- sin.(η * t_l) .^ 2)) / length(t_l) < 0.1
         @test length(sol.times) == length(t_l)
         @test length(sol.states) == 1