Remove expvals from TimeEvolutionParameters

Author: albertomercurio

src/QuantumToolbox.jl

@@ -36,7 +36,8 @@ import SciMLBase:
     ContinuousCallback,
     DiscreteCallback,
     AbstractSciMLProblem,
-    AbstractODEIntegrator
+    AbstractODEIntegrator,
+    AbstractODESolution
 import StochasticDiffEq: StochasticDiffEqAlgorithm, SRA1
 import SciMLOperators:
     SciMLOperators,

src/time_evolution/callback_helpers.jl

@@ -4,13 +4,14 @@ This file contains helper functions for callbacks. The affect! function are defi
 
 ########## SESOLVE ##########
 
-struct SaveFuncSESolve{TE,PT<:Union{Nothing,ProgressBar},IT}
+struct SaveFuncSESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}}
     e_ops::TE
     progr::PT
     iter::IT
+    expvals::TEXPV
 end
 
-(f::SaveFuncSESolve)(integrator) = _save_func_sesolve(integrator, f.e_ops, f.progr, f.iter)
+(f::SaveFuncSESolve)(integrator) = _save_func_sesolve(integrator, f.e_ops, f.progr, f.iter, f.expvals)
 (f::SaveFuncSESolve{Nothing})(integrator) = _save_func_sesolve(integrator, f.progr)
 
 ##
@@ -29,8 +30,7 @@ function _save_func_sesolve(integrator, progr::Nothing)
 end
 
 # When e_ops is a list of operators
-function _save_func_sesolve(integrator, e_ops, progr, iter)
-    expvals = integrator.p.expvals
+function _save_func_sesolve(integrator, e_ops, progr, iter, expvals)
     ψ = integrator.u
     _expect = op -> dot(ψ, op, ψ)
     @. expvals[:, iter[]] = _expect(e_ops)
@@ -44,18 +44,40 @@ function _generate_sesolve_callback(e_ops, tlist, progress_bar)
 
     progr = getVal(progress_bar) ? ProgressBar(length(tlist), enable = getVal(progress_bar)) : nothing
 
-    _save_affect! = SaveFuncSESolve(e_ops_data, progr, Ref(1))
-    return PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
+    expvals = e_ops isa Nothing ? nothing : Array{ComplexF64}(undef, length(e_ops), length(tlist))
+
+    _save_affect! = SaveFuncSESolve(e_ops_data, progr, Ref(1), expvals)
+    return _PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
 end
 
+function _sesolve_get_expvals(sol::AbstractODESolution)
+    kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl
+    if hasproperty(kwargs, :callback)
+        return _sesolve_get_expvals(kwargs.callback)
+    else
+        return nothing
+    end
+end
+function _sesolve_get_expvals(cb::CallbackSet)
+    _cb = cb.discrete_callbacks[1]
+    return _sesolve_get_expvals(_cb)
+end
+_sesolve_get_expvals(cb::DiscreteCallback) = if cb.affect! isa SaveFuncSESolve
+        return cb.affect!.expvals
+    else
+        return nothing
+    end
+_sesolve_get_expvals(cb::ContinuousCallback) = nothing
+
 ########## MCSOLVE ##########
 
-struct SaveFuncMCSolve{TE,IT}
+struct SaveFuncMCSolve{TE,IT,TEXPV}
     e_ops::TE
     iter::IT
+    expvals::TEXPV
 end
 
-(f::SaveFuncMCSolve)(integrator) = _save_func_mcsolve(integrator, f.e_ops, f.iter)
+(f::SaveFuncMCSolve)(integrator) = _save_func_mcsolve(integrator, f.e_ops, f.iter, f.expvals)
 
 struct LindbladJump{T1,T2}
     c_ops::T1
@@ -66,8 +88,7 @@ end
 
 ##
 
-function _save_func_mcsolve(integrator, e_ops, iter)
-    expvals = integrator.p.expvals
+function _save_func_mcsolve(integrator, e_ops, iter, expvals)
     cache_mc = integrator.p.mcsolve_params.cache_mc
 
     copyto!(cache_mc, integrator.u)
@@ -100,12 +121,15 @@ function _generate_mcsolve_kwargs(e_ops, tlist, c_ops, jump_callback, kwargs)
     end
 
     if e_ops isa Nothing
+        # We are implicitly saying that we don't have a `ProgressBar`
         kwargs2 =
             haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(cb1, kwargs.callback),)) :
             merge(kwargs, (callback = cb1,))
         return kwargs2
     else
-        _save_affect! = SaveFuncMCSolve(get_data.(e_ops), Ref(1))
+        expvals = Array{ComplexF64}(undef, length(e_ops), length(tlist))
+
+        _save_affect! = SaveFuncMCSolve(get_data.(e_ops), Ref(1), expvals)
         cb2 = _PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
         kwargs2 =
             haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(cb1, cb2, kwargs.callback),)) :
@@ -178,35 +202,67 @@ function _mcsolve_get_e_ops(integrator::AbstractODEIntegrator)
     return cb.affect!.e_ops
 end
 
+function _mcsolve_get_expvals(sol::AbstractODESolution)
+    cb = NamedTuple(sol.prob.kwargs).callback
+    if _mcsolve_has_discrete_callbacks(cb)
+        return _mcsolve_get_expvals(cb)
+    else
+        return nothing
+    end
+end
+function _mcsolve_get_expvals(cb::CallbackSet)
+    idx = _mcsolve_has_continuous_jump(cb) ? 1 : 2
+    _cb = cb.discrete_callbacks[idx]
+    return _mcsolve_get_expvals(_cb)
+end
+_mcsolve_get_expvals(cb::DiscreteCallback) =
+    if cb.affect! isa SaveFuncMCSolve
+        return cb.affect!.expvals
+    else
+        nothing
+    end
+_mcsolve_get_expvals(cb::ContinuousCallback) = nothing
+
 #=
-    _mcsolve_callbacks_new_iter(prob, tlist)
+    _mcsolve_callbacks_new_iter_expvals(prob, tlist)
 
-Return the same callbacks of the `prob`, but with the `iter` variable reinitialized to 1.
+Return the same callbacks of the `prob`, but with the `iter` variable reinitialized to 1 and the `expvals` variable reinitialized to a new matrix.
 =#
-function _mcsolve_callbacks_new_iter(prob, tlist)
+function _mcsolve_callbacks_new_iter_expvals(prob, tlist)
     cb = prob.kwargs[:callback]
-    return _mcsolve_callbacks_new_iter(cb, tlist)
+    return _mcsolve_callbacks_new_iter_expvals(cb, tlist)
 end
-function _mcsolve_callbacks_new_iter(cb::CallbackSet, tlist)
+function _mcsolve_callbacks_new_iter_expvals(cb::CallbackSet, tlist)
     cb_continuous = cb.continuous_callbacks
     cb_discrete = cb.discrete_callbacks
 
-    if length(cb_continuous) > 0
+    if _mcsolve_has_continuous_jump(cb)
         idx = 1
         e_ops = cb_discrete[idx].affect!.e_ops
-        _save_affect! = SaveFuncMCSolve(e_ops, Ref(1))
+        expvals = similar(cb_discrete[idx].affect!.expvals)
+        _save_affect! = SaveFuncMCSolve(e_ops, Ref(1), expvals)
         cb_save = _PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
         return CallbackSet(cb_continuous..., cb_save, cb_discrete[2:end]...)
     else
         idx = 2
         e_ops = cb_discrete[idx].affect!.e_ops
-        _save_affect! = SaveFuncMCSolve(e_ops, Ref(1))
+        expvals = similar(cb_discrete[idx].affect!.expvals)
+        _save_affect! = SaveFuncMCSolve(e_ops, Ref(1), expvals)
         cb_save = _PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
         return CallbackSet(cb_continuous..., cb_discrete[1], cb_save, cb_discrete[3:end]...)
     end
 end
-_mcsolve_callbacks_new_iter(cb::ContinuousCallback, tlist) = cb
-_mcsolve_callbacks_new_iter(cb::DiscreteCallback, tlist) = cb
+_mcsolve_callbacks_new_iter_expvals(cb::ContinuousCallback, tlist) = cb # It is only the continuous LindbladJump callback  
+_mcsolve_callbacks_new_iter_expvals(cb::DiscreteCallback, tlist) = cb # It is only the discrete LindbladJump callback
+
+_mcsolve_has_discrete_callbacks(cb::CallbackSet) = length(cb.discrete_callbacks) > 0
+_mcsolve_has_discrete_callbacks(cb::ContinuousCallback) = false
+_mcsolve_has_discrete_callbacks(cb::DiscreteCallback) = true
+
+_mcsolve_has_continuous_jump(cb::CallbackSet) =
+    (length(cb.continuous_callbacks) > 0) && (cb.continuous_callbacks[1].affect! isa LindbladJump)
+_mcsolve_has_continuous_jump(cb::ContinuousCallback) = true
+_mcsolve_has_continuous_jump(cb::DiscreteCallback) = false
 
 ## Temporary function to avoid errors. Waiting for the PR In DiffEqCallbacks.jl to be merged.
 

src/time_evolution/mcsolve.jl

@@ -8,9 +8,7 @@ function _mcsolve_prob_func(prob, i, repeat, global_rng, seeds, tlist)
     traj_rng = typeof(global_rng)()
     seed!(traj_rng, seed)
 
-    expvals = similar(params.expvals)
-
-    T = eltype(expvals)
+    T = eltype(prob.u0)
 
     mcsolve_params = (
         traj_rng = traj_rng,
@@ -23,10 +21,10 @@ function _mcsolve_prob_func(prob, i, repeat, global_rng, seeds, tlist)
         jump_times_which_idx = T[1],
     )
 
-    p = TimeEvolutionParameters(params.params, expvals, mcsolve_params)
+    p = TimeEvolutionParameters(params.params, mcsolve_params)
 
     f = deepcopy(prob.f.f)
-    cb = _mcsolve_callbacks_new_iter(prob, tlist)
+    cb = _mcsolve_callbacks_new_iter_expvals(prob, tlist)
 
     return remake(prob, f = f, p = p, callback = cb)
 end
@@ -106,7 +104,7 @@ end
         tlist::AbstractVector,
         c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::Union{NamedTuple,AbstractVector} = eltype(ψ0)[],
+        params = eltype(ψ0)[],
         rng::AbstractRNG = default_rng(),
         jump_callback::TJC = ContinuousLindbladJumpCallback(),
         kwargs...,
@@ -153,7 +151,7 @@ If the environmental measurements register a quantum jump, the wave function und
 - `tlist`: List of times at which to save either the state or the expectation values of the system.
 - `c_ops`: List of collapse operators ``\{\hat{C}_n\}_n``. It can be either a `Vector` or a `Tuple`.
 - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`.
-- `params`: `NamedTuple` or `AbstractVector` of parameters to pass to the solver.
+- `params`: `NamedTuple` or `AbstractVector` of parameters to pass to the solver. For more advanced usage, any custom struct can be used.
 - `rng`: Random number generator for reproducibility.
 - `jump_callback`: The Jump Callback type: Discrete or Continuous. The default is `ContinuousLindbladJumpCallback()`, which is more precise.
 - `kwargs`: The keyword arguments for the ODEProblem.
@@ -175,7 +173,7 @@ function mcsolveProblem(
     tlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::Union{NamedTuple,AbstractVector} = eltype(ψ0)[],
+    params = eltype(ψ0)[],
     rng::AbstractRNG = default_rng(),
     jump_callback::TJC = ContinuousLindbladJumpCallback(),
     kwargs...,
@@ -192,13 +190,7 @@ function mcsolveProblem(
 
     T = Base.promote_eltype(H_eff_evo, ψ0)
 
-    if e_ops isa Nothing
-        expvals = Array{T}(undef, 0, length(tlist))
-        is_empty_e_ops = true
-    else
-        expvals = Array{T}(undef, length(e_ops), length(tlist))
-        is_empty_e_ops = isempty(e_ops)
-    end
+    is_empty_e_ops = e_ops isa Nothing ? true : isempty(e_ops)
 
     saveat = is_empty_e_ops ? tlist : [tlist[end]]
     # We disable the progress bar of the sesolveProblem because we use a global progress bar for all the trajectories
@@ -217,8 +209,6 @@ function mcsolveProblem(
     random_n = similar(ψ0.data, T, 1) # We could use a Ref, but we have to keep the same type for all the parameters due to SciMLStructures.jl.
     random_n[1] = rand(rng)
 
-    progr = ProgressBar(length(tlist), enable = false)
-
     mcsolve_params = (
         traj_rng = rng,
         random_n = random_n,
@@ -229,7 +219,7 @@ function mcsolveProblem(
         jump_which = jump_which,
         jump_times_which_idx = jump_times_which_idx,
     )
-    p = TimeEvolutionParameters(params, expvals, mcsolve_params)
+    p = TimeEvolutionParameters(params, mcsolve_params)
 
     return sesolveProblem(H_eff_evo, ψ0, tlist; params = p, kwargs3...)
 end
@@ -524,13 +514,15 @@ function mcsolve(
 
     dims = ens_prob_mc.dims
     _sol_1 = sol[:, 1]
+    _expvals_sol_1 = _mcsolve_get_expvals(_sol_1)
 
-    expvals_all = mapreduce(i -> sol[:, i].prob.p.expvals, (x, y) -> cat(x, y, dims = 3), eachindex(sol))
+    _expvals_all = _expvals_sol_1 isa Nothing ? nothing : map(i -> _mcsolve_get_expvals(sol[:, i]), eachindex(sol))
+    expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all)
     states = map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol))
     jump_times = map(i -> real.(sol[:, i].prob.p.mcsolve_params.jump_times), eachindex(sol))
     jump_which = map(i -> round.(Int, sol[:, i].prob.p.mcsolve_params.jump_which), eachindex(sol))
 
-    expvals = dropdims(sum(expvals_all, dims = 3), dims = 3) ./ length(sol)
+    expvals = _expvals_sol_1 isa Nothing ? nothing : dropdims(sum(expvals_all, dims = 3), dims = 3) ./ length(sol)
 
     return TimeEvolutionMCSol(
         ntraj,

src/time_evolution/sesolve.jl

@@ -2,7 +2,7 @@ export sesolveProblem, sesolve
 
 function _merge_sesolve_kwargs_with_callback(kwargs, cb)
     kwargs2 =
-        haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(kwargs.callback, cb),)) :
+        haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(cb, kwargs.callback),)) :
         merge(kwargs, (callback = cb,))
 
     return kwargs2
@@ -25,7 +25,7 @@ _sesolve_make_U_QobjEvo(H) = QobjEvo(H, -1im)
         ψ0::QuantumObject{DT2,KetQuantumObject},
         tlist::AbstractVector;
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::Union{NamedTuple, AbstractVector, TimeEvolutionParameters} = eltype(ψ0)[],
+        params = eltype(ψ0)[],
         progress_bar::Union{Val,Bool} = Val(true),
         inplace::Union{Val,Bool} = Val(true),
         kwargs...,
@@ -43,7 +43,7 @@ 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` or `AbstractVector` of parameters to pass to the solver. For more advanced usage, you can use the [`TimeEvolutionParameters`](@ref) struct.
+- `params`: `NamedTuple` or `AbstractVector` of parameters to pass to the solver. For more advanced usage, any custom struct can be used.
 - `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.
@@ -64,7 +64,7 @@ function sesolveProblem(
     ψ0::QuantumObject{DT2,KetQuantumObject},
     tlist::AbstractVector;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::Union{NamedTuple,AbstractVector,TimeEvolutionParameters} = eltype(ψ0)[],
+    params = eltype(ψ0)[],
     progress_bar::Union{Val,Bool} = Val(true),
     inplace::Union{Val,Bool} = Val(true),
     kwargs...,
@@ -78,34 +78,19 @@ function sesolveProblem(
     isoper(H_evo) || throw(ArgumentError("The Hamiltonian must be an Operator."))
     check_dims(H_evo, ψ0)
 
-    ψ0 = sparse_to_dense(_CType(ψ0), get_data(ψ0)) # Convert it to dense vector with complex element type
+    T = Base.promote_eltype(H_evo, ψ0)
+    ψ0 = sparse_to_dense(_CType(T), get_data(ψ0)) # Convert it to dense vector with complex element type
     U = H_evo.data
 
-    if e_ops isa Nothing
-        expvals = Array{ComplexF64}(undef, 0, length(tlist))
-        is_empty_e_ops = true
-    else
-        expvals = Array{ComplexF64}(undef, length(e_ops), length(tlist))
-        is_empty_e_ops = isempty(e_ops)
-    end
-
-    if params isa TimeEvolutionParameters
-        (e_ops isa Nothing) || throw(
-            ArgumentError(
-                "The parameter `params` cannot be a TimeEvolutionParameters object when `e_ops` is not Nothing",
-            ),
-        )
-    end
-
-    p = params isa TimeEvolutionParameters ? params : TimeEvolutionParameters(params, expvals)
+    is_empty_e_ops = (e_ops isa Nothing) ? true : isempty(e_ops)
 
     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])
-    prob = ODEProblem{getVal(inplace),FullSpecialize}(U, ψ0, tspan, p; kwargs3...)
+    prob = ODEProblem{getVal(inplace),FullSpecialize}(U, ψ0, tspan, params; kwargs3...)
 
     return TimeEvolutionProblem(prob, tlist, H_evo.dims)
 end
@@ -117,7 +102,7 @@ end
         tlist::AbstractVector;
         alg::OrdinaryDiffEqAlgorithm = Tsit5(),
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::Union{NamedTuple, AbstractVector} = eltype(ψ0)[],
+        params = eltype(ψ0)[],
         progress_bar::Union{Val,Bool} = Val(true),
         inplace::Union{Val,Bool} = Val(true),
         kwargs...,
@@ -136,7 +121,7 @@ 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` or `AbstractVector` of parameters to pass to the solver.
+- `params`: `NamedTuple` or `AbstractVector` of parameters to pass to the solver. For more advanced usage, any custom struct can be used.
 - `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.
@@ -159,7 +144,7 @@ function sesolve(
     tlist::AbstractVector;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::Union{NamedTuple,AbstractVector} = eltype(ψ0)[],
+    params = eltype(ψ0)[],
     progress_bar::Union{Val,Bool} = Val(true),
     inplace::Union{Val,Bool} = Val(true),
     kwargs...,
@@ -186,7 +171,7 @@ function sesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit
     return TimeEvolutionSol(
         prob.times,
         ψt,
-        sol.prob.p.expvals,
+        _sesolve_get_expvals(sol),
         sol.retcode,
         sol.alg,
         NamedTuple(sol.prob.kwargs).abstol,