Working sesolve

Author: albertomercurio

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_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 _generate_sesolve_kwargs_with_callback(tlist, kwargs)
-    cb1 = PresetTimeCallback(tlist, _save_func_sesolve, save_positions = (false, false))
+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}) =
@@ -103,31 +119,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{true,FullSpecialize}(U, ψ0, tspan, p; kwargs3...)
+
+    return QuantumTimeEvoProblem(prob, tlist, H_evo.dims)
 end
 
 @doc raw"""
@@ -186,13 +194,13 @@ function sesolve(
     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_evolution.jl

@@ -5,6 +5,60 @@ 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)
 
+# 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
+
+function Base.merge(a::QuantumTimeEvoParameters, b::NamedTuple)
+    return QuantumTimeEvoParameters(a.expvals, a.progr, merge(a.params, b))
+end
+
+struct QuantumTimeEvoProblem{PT,TT<:AbstractVector,DT<:AbstractVector}
+    prob::PT
+    times::TT
+    dims::DT
+end
+
 @doc raw"""
     struct TimeEvolutionSol