Introduce sesolve_map and mesolve_map

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main)
 
 - Fix `cite()` bibtex output. ([#552])
+- Implement `sesolve_map` and `mesolve_map` for solving multiple initial states and parameter sets in parallel. ([#554])
 - Add `qeye_like` and `qzero_like`, which are synonyms of `one` and `zero`. ([#555])
 - Add steadystate and DSF benchmarks. The `SteadyStateODESOlver` tolerances are lowered to `terminate_reltol=1e-4` and `terminate_abstol=1e-6` to improve speed at the cost of accuracy. ([#557])
 
@@ -328,5 +329,6 @@ Release date: 2024-11-13
 [#544]: https://github.com/qutip/QuantumToolbox.jl/issues/544
 [#546]: https://github.com/qutip/QuantumToolbox.jl/issues/546
 [#552]: https://github.com/qutip/QuantumToolbox.jl/issues/552
+[#554]: https://github.com/qutip/QuantumToolbox.jl/issues/554
 [#555]: https://github.com/qutip/QuantumToolbox.jl/issues/555
 [#557]: https://github.com/qutip/QuantumToolbox.jl/issues/557

docs/src/resources/api.md

@@ -215,6 +215,8 @@ mesolve
 mcsolve
 ssesolve
 smesolve
+sesolve_map
+mesolve_map
 dfd_mesolve
 liouvillian
 liouvillian_generalized

docs/src/users_guide/cluster.md

@@ -193,72 +193,44 @@ flush(stdout)
 end
 
 @everywhere begin
-    const Nc = 20
-    const ωc = 1.0
-    const g = 0.05
-    const γ = 0.01
-    const F = 0.01
+    Nc = 20
+    ωc = 1.0
+    g = 0.05
+    γ = 0.01
+    F = 0.01
 
-    const a = tensor(destroy(Nc), qeye(2))
+    a = tensor(destroy(Nc), qeye(2))
 
-    const σm = tensor(qeye(Nc), sigmam())
-    const σp = tensor(qeye(Nc), sigmap())
+    σm = tensor(qeye(Nc), sigmam())
+    σp = tensor(qeye(Nc), sigmap())
+    σz = tensor(qeye(Nc), sigmaz())
 
-    H(ωq) = ωc * a' * a + ωq * tensor(num(Nc), qeye(2)) + g * (a' * σm + a * σp)
+    ωq_fun(p, t) = p[1] # ωq
+    drive_fun(p, t) = p[3] * cos(p[2] * t) # F cos(ωd t)
 
-    coef(p, t) = p.F * cos(p.ωd * t) # coefficient for the time-dependent term
+    H = ωc * a' * a + QobjEvo(σz / 2, ωq_fun) + g * (a' * σm + a * σp) + QobjEvo(a + a', drive_fun)
 
-    const c_ops = [sqrt(γ) * a, sqrt(γ) * σm]
-    const e_ops = [a' * a]
+    c_ops = [sqrt(γ) * a, sqrt(γ) * σm]
+    e_ops = [a' * a]
 end
 
-# Define the ODE problem and the EnsembleProblem generation function
+ωq_list = range(ωc - 3*g, ωc + 3*g, 100)
+ωd_list = range(ωc - 3*g, ωc + 3*g, 100)
+F_list = [F]
 
-@everywhere begin
-    ωq_list = range(ωc - 3*g, ωc + 3*g, 100)
-    ωd_list = range(ωc - 3*g, ωc + 3*g, 100)
-
-    const iter = collect(Iterators.product(ωq_list, ωd_list))
-
-    function my_prob_func(prob, i, repeat, channel)
-        ωq, ωd = iter[i]
-        H_i = H(ωq)
-        H_d_i = H_i + QobjEvo(a + a', coef) # Hamiltonian with a driving term
-
-        L = liouvillian(H_d_i, c_ops).data # Make the Liouvillian
-
-        put!(channel, true) # Update the progress bar channel
-
-        remake(prob, f=L, p=(F = F, ωd = ωd))
-    end
-end
-
-ωq, ωd = iter[1]
-H0 = H(ωq) + QobjEvo(a + a', coef)
-ψ0 = tensor(fock(Nc, 0), basis(2, 1)) # Ground State
+ψ_list = [
+    tensor(fock(Nc, 0), basis(2, 0)),
+    tensor(fock(Nc, 0), basis(2, 1))
+]
 tlist = range(0, 20 / γ, 1000)
 
-prob = mesolveProblem(H0, ψ0, tlist, c_ops, e_ops=e_ops, progress_bar=Val(false), params=(F = F, ωd = ωd))
-
-### Just to print the progress bar
-progr = ProgressBar(length(iter))
-progr_channel::RemoteChannel{Channel{Bool}} = RemoteChannel(() -> Channel{Bool}(1))
-###
-ens_prob = EnsembleProblem(prob.prob, prob_func=(prob, i, repeat) -> my_prob_func(prob, i, repeat, progr_channel))
-
-
-@sync begin
-    @async while take!(progr_channel)
-        next!(progr)
-    end
-
-    @async begin
-        sol = solve(ens_prob, Tsit5(), EnsembleSplitThreads(), trajectories = length(iter))
-        put!(progr_channel, false)
-    end
-end
+sol = mesolve_map(H, ψ_list, tlist, c_ops; e_ops = e_ops, params = (ωq_list, ωd_list, F_list), ensemblealg = EnsembleSplitThreads())
 ```
 
-We are using the [`mesolveProblem`](@ref) function to define the master equation problem. We added some code to manage the progress bar, which is updated through a `RemoteChannel`. The `prob_func` argument of the `EnsembleProblem` function is used to define the function that generates the problem for each iteration. The `iter` variable contains the product of the `ωq_list` and `ωd_list` lists, which are used to sweep over the parameters. The `sol = solve(ens_prob, Tsit5(), EnsembleDistributed(), trajectories=length(iter))` command is used to solve the problem with the distributed ensemble method. The output of the script will be printed in the `output.out` file, which contains an output similar to the previous example.
+Notice that we are using the [`mesolve_map`](@ref) function, which internally uses the `SciMLBase.EnsembleProblem` function to parallelize the computation of [`mesolveProblem`](@ref). The result is an `Array` of [`TimeEvolutionSol`](@ref), where each element corresponds to a specific combination of initial states and parameters. One can access the solution for a specific combination of initial states and parameters using indexing. For example,
 
+```julia
+sol[i,j,k,l].expect
+```
 
+this will give the expectation values for the case of `ψ_list[i]`, `ωq_list[j]`, `ωd_list[k]`, and `F_list[l]`.

src/qobj/superoperators.jl

@@ -180,18 +180,5 @@ liouvillian(H::AbstractQuantumObject{Operator}, Id_cache::Diagonal = I(prod(H.di
 liouvillian(H::AbstractQuantumObject{SuperOperator}, Id_cache::Diagonal) = H
 
 function _sum_lindblad_dissipators(c_ops, Id_cache::Diagonal)
-    D = 0
-    # sum all the (time-independent) c_ops first
-    c_ops_ti = filter(op -> isa(op, QuantumObject), c_ops)
-    if !isempty(c_ops_ti)
-        D += mapreduce(op -> lindblad_dissipator(op, Id_cache), +, c_ops_ti)
-    end
-
-    # sum rest of the QobjEvo together
-    c_ops_td = filter(op -> isa(op, QuantumObjectEvolution), c_ops)
-    if !isempty(c_ops_td)
-        D += mapreduce(op -> lindblad_dissipator(op, Id_cache), +, c_ops_td)
-    end
-
-    return D
+    return sum(op -> lindblad_dissipator(op, Id_cache), c_ops; init = zero(spre(first(c_ops), Id_cache)))
 end

src/time_evolution/brmesolve.jl

@@ -3,7 +3,7 @@ export bloch_redfield_tensor, brterm, brmesolve
 @doc raw"""
     bloch_redfield_tensor(
         H::QuantumObject{Operator},
-        a_ops::Union{AbstractVector, Tuple},
+        a_ops::Union{AbstractVector, Tuple, Nothing},
         c_ops::Union{AbstractVector, Tuple, Nothing}=nothing;
         sec_cutoff::Real=0.1,
         fock_basis::Union{Val,Bool}=Val(false)
@@ -28,7 +28,7 @@ The return depends on `fock_basis`.
 """
 function bloch_redfield_tensor(
     H::QuantumObject{Operator},
-    a_ops::Union{AbstractVector,Tuple},
+    a_ops::Union{AbstractVector,Tuple,Nothing},
     c_ops::Union{AbstractVector,Tuple,Nothing} = nothing;
     sec_cutoff::Real = 0.1,
     fock_basis::Union{Val,Bool} = Val(false),
@@ -44,7 +44,9 @@ function bloch_redfield_tensor(
     # Check whether we can rotate the terms to the eigenbasis directly in the Hamiltonian space
     fock_basis_hamiltonian = getVal(fock_basis) && sec_cutoff == -1
 
-    R = isempty(a_ops) ? 0 : sum(x -> _brterm(rst, x[1], x[2], sec_cutoff, fock_basis_hamiltonian), a_ops)
+    R =
+        (isnothing(a_ops) || isempty(a_ops)) ? 0 :
+        sum(x -> _brterm(rst, x[1], x[2], sec_cutoff, fock_basis_hamiltonian), a_ops)
 
     # If in fock basis, we need to transform the terms back to the fock basis
     # Note: we can transform the terms in the Hamiltonian space only if sec_cutoff is -1

src/time_evolution/callback_helpers/callback_helpers.jl

@@ -139,7 +139,7 @@ Return the Callback that is responsible for saving the expectation values of the
 =#
 function _get_save_callback(sol::AbstractODESolution, method::Type{SF}) where {SF<:AbstractSaveFunc}
     kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl
-    if hasproperty(kwargs, :callback)
+    if hasproperty(kwargs, :callback) && !isnothing(kwargs.callback)
         return _get_save_callback(kwargs.callback, method)
     else
         return nothing

src/time_evolution/mesolve.jl

@@ -1,10 +1,31 @@
-export mesolveProblem, mesolve
+export mesolveProblem, mesolve, mesolve_map
 
 _mesolve_make_L_QobjEvo(H::Union{QuantumObject,Nothing}, c_ops) = QobjEvo(liouvillian(H, c_ops); type = SuperOperator())
 _mesolve_make_L_QobjEvo(H::Union{QuantumObjectEvolution,Tuple}, c_ops) = liouvillian(QobjEvo(H), c_ops)
 _mesolve_make_L_QobjEvo(H::Nothing, c_ops::Nothing) = throw(ArgumentError("Both H and
 c_ops are Nothing. You are probably running the wrong function."))
 
+function _gen_mesolve_solution(sol, times, dimensions, isoperket::Val)
+    if getVal(isoperket)
+        ρt = map(ϕ -> QuantumObject(ϕ, type = OperatorKet(), dims = dimensions), sol.u)
+    else
+        ρt = map(ϕ -> QuantumObject(vec2mat(ϕ), type = Operator(), dims = dimensions), sol.u)
+    end
+
+    kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
+
+    return TimeEvolutionSol(
+        times,
+        sol.t,
+        ρt,
+        _get_expvals(sol, SaveFuncMESolve),
+        sol.retcode,
+        sol.alg,
+        kwargs.abstol,
+        kwargs.reltol,
+    )
+end
+
 @doc raw"""
     mesolveProblem(
         H::Union{AbstractQuantumObject,Tuple},
@@ -207,23 +228,143 @@ end
 function mesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit5(); kwargs...)
     sol = solve(prob.prob, alg; kwargs...)
 
-    # No type instabilities since `isoperket` is a Val, and so it is known at compile time
-    if getVal(prob.kwargs.isoperket)
-        ρt = map(ϕ -> QuantumObject(ϕ, type = OperatorKet(), dims = prob.dimensions), sol.u)
-    else
-        ρt = map(ϕ -> QuantumObject(vec2mat(ϕ), type = Operator(), dims = prob.dimensions), sol.u)
+    return _gen_mesolve_solution(sol, prob.times, prob.dimensions, prob.kwargs.isoperket)
+end
+
+@doc raw"""
+    mesolve_map(
+        H::Union{AbstractQuantumObject,Tuple},
+        ψ0::Union{QuantumObject,AbstractVector{<:QuantumObject}},
+        tlist::AbstractVector,
+        c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
+        alg::OrdinaryDiffEqAlgorithm = Tsit5(),
+        ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
+        e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
+        params::Union{NullParameters,Tuple} = NullParameters(),
+        progress_bar::Union{Val,Bool} = Val(true),
+        kwargs...,
+    )
+
+Solve the master equation for multiple initial states and parameter sets using ensemble simulation.
+
+This function computes the time evolution for all combinations (Cartesian product) of initial states and parameter sets, solving the Lindblad master equation (see [`mesolve`](@ref)):
+
+```math
+\frac{\partial \hat{\rho}(t)}{\partial t} = -i[\hat{H}, \hat{\rho}(t)] + \sum_n \mathcal{D}(\hat{C}_n) [\hat{\rho}(t)]
+```
+
+where
+
+```math
+\mathcal{D}(\hat{C}_n) [\hat{\rho}(t)] = \hat{C}_n \hat{\rho}(t) \hat{C}_n^\dagger - \frac{1}{2} \hat{C}_n^\dagger \hat{C}_n \hat{\rho}(t) - \frac{1}{2} \hat{\rho}(t) \hat{C}_n^\dagger \hat{C}_n
+```
+
+for each combination in the ensemble.
+
+# Arguments
+
+- `H`: Hamiltonian of the system ``\hat{H}``. It can be either a [`QuantumObject`](@ref), a [`QuantumObjectEvolution`](@ref), or a `Tuple` of operator-function pairs.
+- `ψ0`: Initial state(s) of the system. Can be a single [`QuantumObject`](@ref) or a `Vector` of initial states. It can be either a [`Ket`](@ref), [`Operator`](@ref) or [`OperatorKet`](@ref).
+- `tlist`: List of time points 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`.
+- `alg`: The algorithm for the ODE solver. The default is `Tsit5()`.
+- `ensemblealg`: Ensemble algorithm to use for parallel computation. Default is `EnsembleThreads()`.
+- `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`.
+- `params`: A `Tuple` of parameter sets. Each element should be an `AbstractVector` representing the sweep range for that parameter. The function will solve for all combinations of initial states and parameter sets.
+- `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
+- `kwargs`: The keyword arguments for the ODEProblem.
+
+# Notes
+
+- The function returns an array of solutions with dimensions matching the Cartesian product of initial states and parameter sets.
+- If `ψ0` is a vector of `m` states and `params = (p1, p2, ...)` where `p1` has length `n1`, `p2` has length `n2`, etc., the output will be of size `(m, n1, n2, ...)`.
+- If `H` is an [`Operator`](@ref), `ψ0` is a [`Ket`](@ref) and `c_ops` is `Nothing`, the function will call [`sesolve_map`](@ref) instead.
+- See [`mesolve`](@ref) for more details.
+
+# Returns
+
+- An array of [`TimeEvolutionSol`](@ref) objects with dimensions `(length(ψ0), length(params[1]), length(params[2]), ...)`.
+"""
+function mesolve_map(
+    H::Union{AbstractQuantumObject{HOpType},Tuple},
+    ψ0::AbstractVector{<:QuantumObject{StateOpType}},
+    tlist::AbstractVector,
+    c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
+    alg::OrdinaryDiffEqAlgorithm = Tsit5(),
+    ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
+    e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
+    params::Union{NullParameters,Tuple} = NullParameters(),
+    progress_bar::Union{Val,Bool} = Val(true),
+    kwargs...,
+) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet}}
+    (isoper(H) && all(isket, ψ0) && isnothing(c_ops)) && return sesolve_map(
+        H,
+        ψ0,
+        tlist;
+        alg = alg,
+        ensemblealg = ensemblealg,
+        e_ops = e_ops,
+        params = params,
+        progress_bar = progress_bar,
+        kwargs...,
+    )
+
+    # mapping initial states and parameters
+    # Convert to appropriate format based on state type
+    ψ0_iter = map(ψ0) do state
+        T = _complex_float_type(eltype(state))
+        if isoperket(state)
+            to_dense(T, copy(state.data))
+        else
+            to_dense(T, mat2vec(ket2dm(state).data))
+        end
     end
+    iter =
+        params isa NullParameters ? collect(Iterators.product(ψ0_iter, [params])) :
+        collect(Iterators.product(ψ0_iter, params...))
+    ntraj = length(iter)
 
-    kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
+    # we disable the progress bar of the mesolveProblem because we use a global progress bar for all the trajectories
+    prob = mesolveProblem(
+        H,
+        first(ψ0),
+        tlist,
+        c_ops;
+        e_ops = e_ops,
+        params = first(iter)[2:end],
+        progress_bar = Val(false),
+        kwargs...,
+    )
 
-    return TimeEvolutionSol(
+    # generate and solve ensemble problem
+    _output_func = _ensemble_dispatch_output_func(ensemblealg, progress_bar, ntraj, _standard_output_func) # setup global progress bar
+    ens_prob = TimeEvolutionProblem(
+        EnsembleProblem(
+            prob.prob,
+            prob_func = (prob, i, repeat) -> _se_me_map_prob_func(prob, i, repeat, iter),
+            output_func = _output_func[1],
+            safetycopy = false,
+        ),
         prob.times,
-        sol.t,
-        ρt,
-        _get_expvals(sol, SaveFuncMESolve),
-        sol.retcode,
-        sol.alg,
-        kwargs.abstol,
-        kwargs.reltol,
+        prob.dimensions,
+        (progr = _output_func[2], channel = _output_func[3], isoperket = prob.kwargs.isoperket),
     )
+    sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj)
+
+    # handle solution and make it become an Array of TimeEvolutionSol
+    sol_vec =
+        [_gen_mesolve_solution(sol[:, i], prob.times, prob.dimensions, prob.kwargs.isoperket) for i in eachindex(sol)] # map is type unstable
+    if params isa NullParameters # if no parameters specified, just return a Vector
+        return sol_vec
+    else
+        return reshape(sol_vec, size(iter))
+    end
 end
+mesolve_map(
+    H::Union{AbstractQuantumObject{HOpType},Tuple},
+    ψ0::QuantumObject{StateOpType},
+    tlist::AbstractVector,
+    c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
+    kwargs...,
+) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet}} =
+    mesolve_map(H, [ψ0], tlist, c_ops; kwargs...)