Remove ProgressBar from ODE parameters

Author: albertomercurio

src/qobj/quantum_object_evo.jl

@@ -214,6 +214,8 @@ function QuantumObjectEvolution(
     return QuantumObjectEvolution(α * op.data, type, op.dims)
 end
 
+_all_equal(dims) = all(x -> x == first(dims), dims)
+
 #=
     _QobjEvo_generate_data(op_func_list::Tuple, α; f::Function=identity)
 
@@ -269,7 +271,7 @@ Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution`
     quote
         dims = tuple($(dims_expr...))
 
-        length(unique(dims)) == 1 || throw(ArgumentError("The dimensions of the operators must be the same."))
+        _all_equal(dims) || throw(ArgumentError("The dimensions of the operators must be the same."))
 
         data_expr_const = $qobj_expr_const isa Integer ? $qobj_expr_const : _make_SciMLOperator($qobj_expr_const, α)
 

src/time_evolution/callback_helpers.jl

@@ -4,49 +4,58 @@ This file contains helper functions for callbacks. The affect! function are defi
 
 ########## SESOLVE ##########
 
-struct SaveFuncSESolve{T1,T2}
-    e_ops::T1
-    is_empty_e_ops::T2
+struct SaveFuncSESolve{TE,PT<:Union{Nothing,ProgressBar},IT}
+    e_ops::TE
+    progr::PT
+    iter::IT
 end
 
-(f::SaveFuncSESolve)(integrator) = _save_func_sesolve(integrator, f.e_ops, f.is_empty_e_ops)
-(f::SaveFuncSESolve{Nothing})(integrator) = _save_func_sesolve(integrator)
+(f::SaveFuncSESolve)(integrator) = _save_func_sesolve(integrator, f.e_ops, f.progr, f.iter)
+(f::SaveFuncSESolve{Nothing})(integrator) = _save_func_sesolve(integrator, f.progr)
 
 ##
 
 # When e_ops is Nothing
-function _save_func_sesolve(integrator)
-    next!(integrator.p.progr)
+function _save_func_sesolve(integrator, progr)
+    next!(progr)
+    u_modified!(integrator, false)
+    return nothing
+end
+
+# When progr is Nothing
+function _save_func_sesolve(integrator, progr::Nothing)
     u_modified!(integrator, false)
     return nothing
 end
 
 # When e_ops is a list of operators
-function _save_func_sesolve(integrator, e_ops, is_empty_e_ops)
+function _save_func_sesolve(integrator, e_ops, progr, iter)
     expvals = integrator.p.expvals
-    progr = integrator.p.progr
-    if !is_empty_e_ops
-        ψ = integrator.u
-        _expect = op -> dot(ψ, op, ψ)
-        @. expvals[:, progr.counter[]+1] = _expect(e_ops)
-    end
-    return _save_func_sesolve(integrator)
+    ψ = integrator.u
+    _expect = op -> dot(ψ, op, ψ)
+    @. expvals[:, iter[]] = _expect(e_ops)
+    iter[] += 1
+
+    return _save_func_sesolve(integrator, progr)
 end
 
-function _generate_sesolve_callback(e_ops, tlist)
-    is_empty_e_ops = e_ops isa Nothing ? true : isempty(e_ops)
-    _save_affect! = SaveFuncSESolve(get_data.(e_ops), is_empty_e_ops)
+function _generate_sesolve_callback(e_ops, tlist, progress_bar)
+    e_ops_data = e_ops isa Nothing ? nothing : get_data.(e_ops)
+
+    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))
 end
 
 ########## MCSOLVE ##########
 
-struct SaveFuncMCSolve{T1,T2}
-    e_ops::T1
-    is_empty_e_ops::T2
+struct SaveFuncMCSolve{TE,IT}
+    e_ops::TE
+    iter::IT
 end
 
-(f::SaveFuncMCSolve)(integrator) = _save_func_mcsolve(integrator, f.e_ops, f.is_empty_e_ops)
+(f::SaveFuncMCSolve)(integrator) = _save_func_mcsolve(integrator, f.e_ops, f.iter)
 
 struct LindbladJump{T1,T2}
     c_ops::T1
@@ -57,18 +66,17 @@ end
 
 ##
 
-function _save_func_mcsolve(integrator, e_ops, is_empty_e_ops)
+function _save_func_mcsolve(integrator, e_ops, iter)
     expvals = integrator.p.expvals
-    progr = integrator.p.progr
     cache_mc = integrator.p.mcsolve_params.cache_mc
-    if !is_empty_e_ops
-        copyto!(cache_mc, integrator.u)
-        normalize!(cache_mc)
-        ψ = cache_mc
-        _expect = op -> dot(ψ, op, ψ)
-        @. expvals[:, progr.counter[]+1] = _expect(e_ops)
-    end
-    next!(progr)
+
+    copyto!(cache_mc, integrator.u)
+    normalize!(cache_mc)
+    ψ = cache_mc
+    _expect = op -> dot(ψ, op, ψ)
+    @. expvals[:, iter[]] = _expect(e_ops)
+    iter[] += 1
+
     u_modified!(integrator, false)
     return nothing
 end
@@ -97,8 +105,7 @@ function _generate_mcsolve_kwargs(e_ops, tlist, c_ops, jump_callback, kwargs)
             merge(kwargs, (callback = cb1,))
         return kwargs2
     else
-        is_empty_e_ops = isempty(e_ops)
-        _save_affect! = SaveFuncMCSolve(get_data.(e_ops), is_empty_e_ops)
+        _save_affect! = SaveFuncMCSolve(get_data.(e_ops), Ref(1))
         cb2 = _PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
         kwargs2 =
             haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(cb1, cb2, kwargs.callback),)) :
@@ -171,6 +178,36 @@ function _mcsolve_get_e_ops(integrator::AbstractODEIntegrator)
     return cb.affect!.e_ops
 end
 
+#=
+    _mcsolve_callbacks_new_iter(prob, tlist)
+
+Return the same callbacks of the `prob`, but with the `iter` variable reinitialized to 1.
+=#
+function _mcsolve_callbacks_new_iter(prob, tlist)
+    cb = prob.kwargs[:callback]
+    return _mcsolve_callbacks_new_iter(cb, tlist)
+end
+function _mcsolve_callbacks_new_iter(cb::CallbackSet, tlist)
+    cb_continuous = cb.continuous_callbacks
+    cb_discrete = cb.discrete_callbacks
+
+    if length(cb_continuous) > 0
+        idx = 1
+        e_ops = cb_discrete[idx].affect!.e_ops
+        _save_affect! = SaveFuncMCSolve(e_ops, Ref(1))
+        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))
+        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
+
 ## Temporary function to avoid errors. Waiting for the PR In DiffEqCallbacks.jl to be merged.
 
 import SciMLBase: INITIALIZE_DEFAULT, add_tstop!

src/time_evolution/mcsolve.jl

@@ -1,15 +1,14 @@
 export mcsolveProblem, mcsolveEnsembleProblem, mcsolve
 export ContinuousLindbladJumpCallback, DiscreteLindbladJumpCallback
 
-function _mcsolve_prob_func(prob, i, repeat, global_rng, seeds)
+function _mcsolve_prob_func(prob, i, repeat, global_rng, seeds, tlist)
     params = prob.p
 
     seed = seeds[i]
     traj_rng = typeof(global_rng)()
     seed!(traj_rng, seed)
 
     expvals = similar(params.expvals)
-    progr = ProgressBar(size(expvals, 2), enable = false)
 
     T = eltype(expvals)
 
@@ -24,16 +23,17 @@ function _mcsolve_prob_func(prob, i, repeat, global_rng, seeds)
         jump_times_which_idx = T[1],
     )
 
-    p = TimeEvolutionParameters(params.params, expvals, progr, mcsolve_params)
+    p = TimeEvolutionParameters(params.params, expvals, mcsolve_params)
 
     f = deepcopy(prob.f.f)
+    cb = _mcsolve_callbacks_new_iter(prob, tlist)
 
-    return remake(prob, f = f, p = p)
+    return remake(prob, f = f, p = p, callback = cb)
 end
 
-function _mcsolve_dispatch_prob_func(rng, ntraj)
+function _mcsolve_dispatch_prob_func(rng, ntraj, tlist)
     seeds = map(i -> rand(rng, UInt64), 1:ntraj)
-    return (prob, i, repeat) -> _mcsolve_prob_func(prob, i, repeat, rng, seeds)
+    return (prob, i, repeat) -> _mcsolve_prob_func(prob, i, repeat, rng, seeds, tlist)
 end
 
 # Standard output function
@@ -229,7 +229,7 @@ function mcsolveProblem(
         jump_which = jump_which,
         jump_times_which_idx = jump_times_which_idx,
     )
-    p = TimeEvolutionParameters(params, expvals, progr, mcsolve_params)
+    p = TimeEvolutionParameters(params, expvals, mcsolve_params)
 
     return sesolveProblem(H_eff_evo, ψ0, tlist; params = p, kwargs3...)
 end
@@ -330,7 +330,7 @@ function mcsolveEnsembleProblem(
     output_func::Union{Tuple,Nothing} = nothing,
     kwargs...,
 ) where {DT1,DT2,TJC<:LindbladJumpCallbackType}
-    _prob_func = prob_func isa Nothing ? _mcsolve_dispatch_prob_func(rng, ntraj) : prob_func
+    _prob_func = prob_func isa Nothing ? _mcsolve_dispatch_prob_func(rng, ntraj, tlist) : prob_func
     _output_func =
         output_func isa Nothing ? _mcsolve_dispatch_output_func(ensemble_method, progress_bar, ntraj) : output_func
 

src/time_evolution/sesolve.jl

@@ -9,18 +9,11 @@ function _merge_sesolve_kwargs_with_callback(kwargs, cb)
 end
 
 # Multiple dispatch depending on the progress_bar and e_ops types
-function _generate_sesolve_kwargs(e_ops, progress_bar::Val{true}, 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
-    cb = _generate_sesolve_callback(e_ops, tlist)
+function _generate_sesolve_kwargs(e_ops, progress_bar, tlist, kwargs)
+    cb = _generate_sesolve_callback(e_ops, tlist, progress_bar)
     return _merge_sesolve_kwargs_with_callback(kwargs, cb)
 end
+_generate_sesolve_kwargs(e_ops::Nothing, progress_bar::Val{false}, tlist, kwargs) = kwargs
 
 _sesolve_make_U_QobjEvo(H::QuantumObjectEvolution{<:MatrixOperator}) =
     QobjEvo(MatrixOperator(-1im * H.data.A), dims = H.dims, type = Operator)
@@ -88,8 +81,6 @@ function sesolveProblem(
     ψ0 = sparse_to_dense(_CType(ψ0), get_data(ψ0)) # Convert it to dense vector with complex element type
     U = H_evo.data
 
-    progr = ProgressBar(length(tlist), enable = getVal(progress_bar))
-
     if e_ops isa Nothing
         expvals = Array{ComplexF64}(undef, 0, length(tlist))
         is_empty_e_ops = true
@@ -99,14 +90,14 @@ function sesolveProblem(
     end
 
     if params isa TimeEvolutionParameters
-        (!getVal(progress_bar) && (e_ops isa Nothing)) || throw(
+        (e_ops isa Nothing) || throw(
             ArgumentError(
-                "The parameter `params` cannot be a TimeEvolutionParameters object when `e_ops` is not Nothing and `progress_bar` is true.",
+                "The parameter `params` cannot be a TimeEvolutionParameters object when `e_ops` is not Nothing",
             ),
         )
     end
 
-    p = params isa TimeEvolutionParameters ? params : TimeEvolutionParameters(params, expvals, progr)
+    p = params isa TimeEvolutionParameters ? params : TimeEvolutionParameters(params, expvals)
 
     saveat = is_empty_e_ops ? tlist : [tlist[end]]
     default_values = (DEFAULT_ODE_SOLVER_OPTIONS..., saveat = saveat)

src/time_evolution/time_evo_parameters.jl

@@ -1,22 +1,16 @@
 # 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
-
 @doc raw"""
     struct TimeEvolutionParameters
 
 A Julia constructor for handling the parameters of the time evolution of quantum systems.
 """
-struct TimeEvolutionParameters{ParT,TE<:AbstractMatrix,PT<:ProgressBar,MCST}
+struct TimeEvolutionParameters{ParT,TE<:AbstractMatrix,MCST}
     params::ParT
     expvals::TE
-    progr::PT
     mcsolve_params::MCST
 end
 
-TimeEvolutionParameters(params, expvals, progr) = TimeEvolutionParameters(params, expvals, progr, nothing)
+TimeEvolutionParameters(params, expvals) = TimeEvolutionParameters(params, expvals, nothing)
 
 #=
 By defining a custom `getproperty` method for the `TimeEvolutionParameters` struct, we can access the fields of `params` directly.
@@ -39,7 +33,7 @@ Base.getindex(obj::TimeEvolutionParameters, i::Int) = getindex(obj.params, i)
 Base.length(obj::TimeEvolutionParameters) = length(obj.params)
 
 # function Base.merge(a::TimeEvolutionParameters, b::NamedTuple)
-#     return TimeEvolutionParameters(merge(a.params, b), a.expvals, a.progr, a.mcsolve_params)
+#     return TimeEvolutionParameters(merge(a.params, b), a.expvals, a.mcsolve_params)
 # end
 
 ########## Mark the struct as a SciMLStructure ##########
@@ -75,12 +69,12 @@ end
 function replace(::Tunable, p::TimeEvolutionParameters{ParT}, newbuffer) where {ParT<:NamedTuple}
     @assert length(newbuffer) == length(p.params)
     new_params = NamedTuple{keys(p.params)}(Tuple(newbuffer))
-    return TimeEvolutionParameters(new_params, p.expvals, p.progr, p.mcsolve_params)
+    return TimeEvolutionParameters(new_params, p.expvals, p.mcsolve_params)
 end
 
 function replace(::Tunable, p::TimeEvolutionParameters{ParT}, newbuffer) where {ParT<:AbstractVector}
     @assert length(newbuffer) == length(p.params)
-    return TimeEvolutionParameters(newbuffer, p.expvals, p.progr, p.mcsolve_params)
+    return TimeEvolutionParameters(newbuffer, p.expvals, p.mcsolve_params)
 end
 
 function replace!(::Tunable, p::TimeEvolutionParameters{ParT}, newbuffer) where {ParT<:AbstractVector}

src/time_evolution/time_evolution_dynamical.jl

@@ -609,7 +609,7 @@ function _dsf_mcsolve_prob_func(prob, i, repeat)
         ),
     )
 
-    p = TimeEvolutionParameters(prm, expvals, progr, mcsolve_params)
+    p = TimeEvolutionParameters(prm, expvals, mcsolve_params)
 
     f = deepcopy(prob.f.f)