Working mcsolve

Author: albertomercurio

src/qobj/quantum_object_evo.jl

@@ -91,7 +91,7 @@ struct QuantumObjectEvolution{DT<:AbstractSciMLOperator,ObjType<:QuantumObjectTy
 end
 
 # Make the QuantumObjectEvolution, with the option to pre-multiply by a scalar
-function QuantumObjectEvolution(op_func_list::Tuple, α = true)
+function QuantumObjectEvolution(op_func_list::Tuple, α::Union{Nothing,Number} = nothing)
     op, data = _generate_data(op_func_list, α)
     dims = op.dims
     type = op.type
@@ -103,8 +103,15 @@ function QuantumObjectEvolution(op_func_list::Tuple, α = true)
     return QuantumObjectEvolution(data, type, dims)
 end
 
-QuantumObjectEvolution(op::QuantumObject, α = true) =
-    QuantumObjectEvolution(MatrixOperator(α * op.data), op.type, op.dims)
+QuantumObjectEvolution(op::QuantumObject, α::Union{Nothing,Number} = nothing) =
+    QuantumObjectEvolution(_make_SciMLOperator(op, α), op.type, op.dims)
+
+function QuantumObjectEvolution(op::QuantumObjectEvolution, α::Union{Nothing,Number} = nothing)
+    if α isa Nothing
+        return QuantumObjectEvolution(op.data, op.type, op.dims)
+    end
+    return QuantumObjectEvolution(α * op.data, op.type, op.dims)
+end
 
 @generated function _generate_data(op_func_list::Tuple, α)
     op_func_list_types = op_func_list.parameters
@@ -158,10 +165,18 @@ end
 function _make_SciMLOperator(op_func::Tuple, α)
     T = eltype(op_func[1])
     update_func = (a, u, p, t) -> op_func[2](p, t)
+    if α isa Nothing
+        return ScalarOperator(zero(T), update_func) * MatrixOperator(op_func[1].data)
+    end
     return ScalarOperator(zero(T), update_func) * MatrixOperator(α * op_func[1].data)
 end
 
-_make_SciMLOperator(op::QuantumObject, α) = MatrixOperator(α * op.data)
+function _make_SciMLOperator(op::QuantumObject, α)
+    if α isa Nothing
+        return MatrixOperator(op.data)
+    end
+    return MatrixOperator(α * op.data)
+end
 
 function (QO::QuantumObjectEvolution)(p, t)
     # We put 0 in the place of `u` because the time-dependence doesn't depend on the state

src/qobj/synonyms.jl

@@ -18,13 +18,14 @@ Note that this functions is same as `QuantumObject(A; type=type, dims=dims)`.
 Qobj(A; kwargs...) = QuantumObject(A; kwargs...)
 
 @doc raw"""
-    QobjEvo(op_func_list::Union{Tuple,QuantumObject}, α::Real=true)
+    QobjEvo(op_func_list::Union{Tuple,AbstractQuantumObject}, α::Union{Nothing,Number}=nothing)
 
 Generate [`QuantumObjectEvolution`](@ref)
 
 Note that this functions is same as `QuantumObjectEvolution(op_func_list)`. If `α` is provided, all the operators in `op_func_list` will be pre-multiplied by `α`.
 """
-QobjEvo(op_func_list::Union{Tuple,QuantumObject}, α = true) = QuantumObjectEvolution(op_func_list, α)
+QobjEvo(op_func_list::Union{Tuple,AbstractQuantumObject}, α::Union{Nothing,Number} = nothing) =
+    QuantumObjectEvolution(op_func_list, α)
 
 @doc raw"""
     shape(A::QuantumObject)

src/time_evolution/mcsolve.jl

@@ -191,41 +191,38 @@ If the environmental measurements register a quantum jump, the wave function und
 - `prob::ODEProblem`: The ODEProblem for the Monte Carlo wave function time evolution.
 """
 function mcsolveProblem(
-    H::QuantumObject{MT1,OperatorQuantumObject},
-    ψ0::QuantumObject{<:AbstractArray,KetQuantumObject},
+    H::Union{AbstractQuantumObject{DT1,OperatorQuantumObject},Tuple},
+    ψ0::QuantumObject{DT2,KetQuantumObject},
     tlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
     rng::AbstractRNG = default_rng(),
     jump_callback::TJC = ContinuousLindbladJumpCallback(),
     kwargs...,
-) where {MT1<:AbstractMatrix,TJC<:LindbladJumpCallbackType}
-    check_dims(H, ψ0)
-
+) where {DT1,DT2,TJC<:LindbladJumpCallbackType}
     haskey(kwargs, :save_idxs) &&
         throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
 
     c_ops isa Nothing &&
         throw(ArgumentError("The list of collapse operators must be provided. Use sesolveProblem instead."))
 
-    t_l = convert(Vector{_FType(ψ0)}, tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl
+    tlist = convert(Vector{_FType(ψ0)}, tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl
 
-    H_eff = H - 1im * mapreduce(op -> op' * op, +, c_ops) / 2
+    H_eff_evo = QobjEvo(H) + QobjEvo(mapreduce(op -> op' * op, +, c_ops) / 2, -1im)
 
     if e_ops isa Nothing
-        expvals = Array{ComplexF64}(undef, 0, length(t_l))
+        expvals = Array{ComplexF64}(undef, 0, length(tlist))
         is_empty_e_ops_mc = true
-        e_ops2 = MT1[]
+        e_ops_data = ()
     else
-        expvals = Array{ComplexF64}(undef, length(e_ops), length(t_l))
-        e_ops2 = get_data.(e_ops)
+        expvals = Array{ComplexF64}(undef, length(e_ops), length(tlist))
+        e_ops_data = get_data.(e_ops)
         is_empty_e_ops_mc = isempty(e_ops)
     end
 
-    saveat = is_empty_e_ops_mc ? t_l : [t_l[end]]
+    saveat = is_empty_e_ops_mc ? tlist : [tlist[end]]
     # We disable the progress bar of the sesolveProblem because we use a global progress bar for all the trajectories
     default_values = (DEFAULT_ODE_SOLVER_OPTIONS..., saveat = saveat, progress_bar = Val(false))
     kwargs2 = merge(default_values, kwargs)
@@ -243,9 +240,9 @@ function mcsolveProblem(
 
     params2 = (
         expvals = expvals,
-        e_ops_mc = e_ops2,
+        e_ops_mc = e_ops_data,
         is_empty_e_ops_mc = is_empty_e_ops_mc,
-        progr_mc = ProgressBar(length(t_l), enable = false),
+        progr_mc = ProgressBar(length(tlist), enable = false),
         traj_rng = rng,
         c_ops = c_ops_data,
         c_ops_herm = c_ops_herm_data,
@@ -259,53 +256,51 @@ function mcsolveProblem(
         params...,
     )
 
-    return mcsolveProblem(H_eff, ψ0, t_l, alg, H_t, params2, jump_callback; kwargs2...)
+    return mcsolveProblem(H_eff_evo, ψ0, tlist, alg, params2, jump_callback; kwargs2...)
 end
 
 function mcsolveProblem(
-    H_eff::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
-    ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
-    t_l::AbstractVector,
+    H_eff_evo::QuantumObjectEvolution{DT1,OperatorQuantumObject},
+    ψ0::QuantumObject{DT2,KetQuantumObject},
+    tlist::AbstractVector,
     alg::OrdinaryDiffEqAlgorithm,
-    H_t::Union{Nothing,Function,TimeDependentOperatorSum},
     params::NamedTuple,
     jump_callback::DiscreteLindbladJumpCallback;
     kwargs...,
-) where {T1,T2}
+) where {DT1,DT2}
     cb1 = DiscreteCallback(LindbladJumpDiscreteCondition, LindbladJumpAffect!, save_positions = (false, false))
-    cb2 = PresetTimeCallback(t_l, _save_func_mcsolve, save_positions = (false, false))
+    cb2 = PresetTimeCallback(tlist, _save_func_mcsolve, save_positions = (false, false))
     kwargs2 = (; kwargs...)
     kwargs2 =
         haskey(kwargs2, :callback) ? merge(kwargs2, (callback = CallbackSet(cb1, cb2, kwargs2.callback),)) :
         merge(kwargs2, (callback = CallbackSet(cb1, cb2),))
 
-    return sesolveProblem(H_eff, ψ0, t_l; alg = alg, H_t = H_t, params = params, kwargs2...)
+    return sesolveProblem(H_eff_evo, ψ0, tlist; alg = alg, params = params, kwargs2...)
 end
 
 function mcsolveProblem(
-    H_eff::QuantumObject{<:AbstractArray,OperatorQuantumObject},
-    ψ0::QuantumObject{<:AbstractArray,KetQuantumObject},
-    t_l::AbstractVector,
+    H_eff_evo::QuantumObjectEvolution{DT1,OperatorQuantumObject},
+    ψ0::QuantumObject{DT2,KetQuantumObject},
+    tlist::AbstractVector,
     alg::OrdinaryDiffEqAlgorithm,
-    H_t::Union{Nothing,Function,TimeDependentOperatorSum},
     params::NamedTuple,
     jump_callback::ContinuousLindbladJumpCallback;
     kwargs...,
-)
+) where {DT1,DT2}
     cb1 = ContinuousCallback(
         LindbladJumpContinuousCondition,
         LindbladJumpAffect!,
         nothing,
         interp_points = jump_callback.interp_points,
         save_positions = (false, false),
     )
-    cb2 = PresetTimeCallback(t_l, _save_func_mcsolve, save_positions = (false, false))
+    cb2 = PresetTimeCallback(tlist, _save_func_mcsolve, save_positions = (false, false))
     kwargs2 = (; kwargs...)
     kwargs2 =
         haskey(kwargs2, :callback) ? merge(kwargs2, (callback = CallbackSet(cb1, cb2, kwargs2.callback),)) :
         merge(kwargs2, (callback = CallbackSet(cb1, cb2),))
 
-    return sesolveProblem(H_eff, ψ0, t_l; alg = alg, H_t = H_t, params = params, kwargs2...)
+    return sesolveProblem(H_eff_evo, ψ0, tlist; alg = alg, params = params, kwargs2...)
 end
 
 @doc raw"""
@@ -391,13 +386,12 @@ If the environmental measurements register a quantum jump, the wave function und
 - `prob::EnsembleProblem with ODEProblem`: The Ensemble ODEProblem for the Monte Carlo wave function time evolution.
 """
 function mcsolveEnsembleProblem(
-    H::QuantumObject{MT1,OperatorQuantumObject},
-    ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+    H::Union{AbstractQuantumObject{DT1,OperatorQuantumObject},Tuple},
+    ψ0::QuantumObject{DT2,KetQuantumObject},
     tlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
     rng::AbstractRNG = default_rng(),
     ntraj::Int = 1,
@@ -407,7 +401,7 @@ function mcsolveEnsembleProblem(
     output_func::Function = _mcsolve_dispatch_output_func(ensemble_method),
     progress_bar::Union{Val,Bool} = Val(true),
     kwargs...,
-) where {MT1<:AbstractMatrix,T2,TJC<:LindbladJumpCallbackType}
+) where {DT1,DT2,TJC<:LindbladJumpCallbackType}
     progr = ProgressBar(ntraj, enable = getVal(progress_bar))
     if ensemble_method isa EnsembleDistributed
         progr_channel::RemoteChannel{Channel{Bool}} = RemoteChannel(() -> Channel{Bool}(1))
@@ -429,7 +423,6 @@ function mcsolveEnsembleProblem(
             c_ops;
             alg = alg,
             e_ops = e_ops,
-            H_t = H_t,
             params = merge(params, (global_rng = rng, seeds = seeds)),
             rng = rng,
             jump_callback = jump_callback,
@@ -533,13 +526,12 @@ If the environmental measurements register a quantum jump, the wave function und
 - `sol::TimeEvolutionMCSol`: The solution of the time evolution. See also [`TimeEvolutionMCSol`](@ref)
 """
 function mcsolve(
-    H::QuantumObject{MT1,OperatorQuantumObject},
-    ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+    H::Union{AbstractQuantumObject{DT1,OperatorQuantumObject},Tuple},
+    ψ0::QuantumObject{DT2,KetQuantumObject},
     tlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
     rng::AbstractRNG = default_rng(),
     ntraj::Int = 1,
@@ -549,15 +541,14 @@ function mcsolve(
     output_func::Function = _mcsolve_dispatch_output_func(ensemble_method),
     progress_bar::Union{Val,Bool} = Val(true),
     kwargs...,
-) where {MT1<:AbstractMatrix,T2,TJC<:LindbladJumpCallbackType}
+) where {DT1,DT2,TJC<:LindbladJumpCallbackType}
     ens_prob_mc = mcsolveEnsembleProblem(
         H,
         ψ0,
         tlist,
         c_ops;
         alg = alg,
         e_ops = e_ops,
-        H_t = H_t,
         params = params,
         rng = rng,
         ntraj = ntraj,
@@ -569,11 +560,12 @@ function mcsolve(
         kwargs...,
     )
 
-    return mcsolve(ens_prob_mc; alg = alg, ntraj = ntraj, ensemble_method = ensemble_method)
+    return mcsolve(ens_prob_mc, tlist; alg = alg, ntraj = ntraj, ensemble_method = ensemble_method)
 end
 
 function mcsolve(
-    ens_prob_mc::EnsembleProblem;
+    ens_prob_mc::EnsembleProblem,
+    tlist::AbstractVector;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     ntraj::Int = 1,
     ensemble_method = EnsembleThreads(),
@@ -599,7 +591,7 @@ function mcsolve(
 
         return TimeEvolutionMCSol(
             ntraj,
-            _sol_1.prob.p.times,
+            tlist,
             states,
             expvals,
             expvals_all,

src/time_evolution/sesolve.jl

@@ -77,7 +77,7 @@ Generates the ODEProblem for the Schrödinger time evolution of a quantum system
 - `prob`: The `ODEProblem` for the Schrödinger time evolution of the system.
 """
 function sesolveProblem(
-    H::Union{QuantumObject{DT1,OperatorQuantumObject},Tuple},
+    H::Union{AbstractQuantumObject{DT1,OperatorQuantumObject},Tuple},
     ψ0::QuantumObject{DT2,KetQuantumObject},
     tlist::AbstractVector;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
@@ -169,7 +169,7 @@ Time evolution of a closed quantum system using the Schrödinger equation:
 - `sol::TimeEvolutionSol`: The solution of the time evolution. See also [`TimeEvolutionSol`](@ref)
 """
 function sesolve(
-    H::Union{QuantumObject{DT1,OperatorQuantumObject},Tuple},
+    H::Union{AbstractQuantumObject{DT1,OperatorQuantumObject},Tuple},
     ψ0::QuantumObject{DT2,KetQuantumObject},
     tlist::AbstractVector;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),

src/time_evolution/time_evolution.jl

@@ -71,12 +71,15 @@ A structure storing the results and some information from solving quantum trajec
 - `reltol::Real`: The relative tolerance which is used during the solving process.
 """
 struct TimeEvolutionMCSol{
-    TT<:Vector{<:Real},
+    TT<:AbstractVector{<:Real},
     TS<:AbstractVector,
     TE<:Matrix{ComplexF64},
     TEA<:Array{ComplexF64,3},
     TJT<:Vector{<:Vector{<:Real}},
     TJW<:Vector{<:Vector{<:Integer}},
+    AlgT<:OrdinaryDiffEqAlgorithm,
+    AT<:Real,
+    RT<:Real,
 }
     ntraj::Int
     times::TT
@@ -86,9 +89,9 @@ struct TimeEvolutionMCSol{
     jump_times::TJT
     jump_which::TJW
     converged::Bool
-    alg::OrdinaryDiffEqAlgorithm
-    abstol::Real
-    reltol::Real
+    alg::AlgT
+    abstol::AT
+    reltol::RT
 end
 
 function Base.show(io::IO, sol::TimeEvolutionMCSol)

src/time_evolution/time_evolution_dynamical.jl

@@ -618,7 +618,6 @@ function dsf_mcsolveEnsembleProblem(
     dsf_params::NamedTuple = NamedTuple();
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     e_ops::Function = (op_list, p) -> Vector{TOl}([]),
-    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
     ntraj::Int = 1,
     ensemble_method = EnsembleThreads(),
@@ -670,7 +669,6 @@ function dsf_mcsolveEnsembleProblem(
         c_ops₀;
         e_ops = e_ops₀,
         alg = alg,
-        H_t = H_t,
         params = params2,
         ntraj = ntraj,
         ensemble_method = ensemble_method,
@@ -690,7 +688,6 @@ end
         dsf_params::NamedTuple=NamedTuple();
         alg::OrdinaryDiffEqAlgorithm=Tsit5(),
         e_ops::Function=(op_list,p) -> Vector{TOl}([]),
-        H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
         params::NamedTuple=NamedTuple(),
         δα_list::Vector{<:Real}=fill(0.2, length(op_list)),
         ntraj::Int=1,
@@ -716,7 +713,6 @@ function dsf_mcsolve(
     dsf_params::NamedTuple = NamedTuple();
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     e_ops::Function = (op_list, p) -> Vector{TOl}([]),
-    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
     δα_list::Vector{<:Real} = fill(0.2, length(op_list)),
     ntraj::Int = 1,
@@ -736,7 +732,6 @@ function dsf_mcsolve(
         dsf_params;
         alg = alg,
         e_ops = e_ops,
-        H_t = H_t,
         params = params,
         ntraj = ntraj,
         ensemble_method = ensemble_method,
@@ -747,5 +742,5 @@ function dsf_mcsolve(
         kwargs...,
     )
 
-    return mcsolve(ens_prob_mc; alg = alg, ntraj = ntraj, ensemble_method = ensemble_method)
+    return mcsolve(ens_prob_mc, t_l; alg = alg, ntraj = ntraj, ensemble_method = ensemble_method)
 end