Introduce new methods of sesolve_map and mesolve_map for advanced usage (#565)

Author: ytdHuang

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main)
 
+- Introduce new methods of `sesolve_map` and `mesolve_map` for advanced usage. Users can now customize their own `iter`ator structure, `prob_func` and `output_func`. ([#565])
+
 ## [v0.37.0]
 Release date: 2025-10-12
 
@@ -336,3 +338,4 @@ Release date: 2024-11-13
 [#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
+[#565]: https://github.com/qutip/QuantumToolbox.jl/issues/565

src/time_evolution/mesolve.jl

@@ -319,10 +319,11 @@ function mesolve_map(
             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)
+    if params isa NullParameters
+        iter = collect(Iterators.product(ψ0_iter, [params])) |> vec # convert nx1 Matrix into Vector
+    else
+        iter = collect(Iterators.product(ψ0_iter, params...))
+    end
 
     # we disable the progress bar of the mesolveProblem because we use a global progress bar for all the trajectories
     prob = mesolveProblem(
@@ -336,35 +337,49 @@ function mesolve_map(
         kwargs...,
     )
 
-    # generate and solve ensemble problem
-    _output_func = _ensemble_dispatch_output_func(ensemblealg, progress_bar, ntraj, _standard_output_func) # setup global progress bar
+    return mesolve_map(prob, iter, alg, ensemblealg; progress_bar = progress_bar)
+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...)
+
+# this method is for advanced usage
+# User can define their own iterator structure, prob_func and output_func
+#   - `prob_func`: Function to use for generating the ODEProblem.
+#   - `output_func`: a `Tuple` containing the `Function` to use for generating the output of a single trajectory, the (optional) `ProgressBar` object, and the (optional) `RemoteChannel` object.
+#
+# Return: An array of TimeEvolutionSol objects with the size same as the given iter.
+function mesolve_map(
+    prob::TimeEvolutionProblem{<:ODEProblem},
+    iter::AbstractArray,
+    alg::OrdinaryDiffEqAlgorithm = Tsit5(),
+    ensemblealg::EnsembleAlgorithm = EnsembleThreads();
+    prob_func::Union{Function,Nothing} = nothing,
+    output_func::Union{Tuple,Nothing} = nothing,
+    progress_bar::Union{Val,Bool} = Val(true),
+)
+    # generate ensemble problem
+    ntraj = length(iter)
+    _prob_func = isnothing(prob_func) ? (prob, i, repeat) -> _se_me_map_prob_func(prob, i, repeat, iter) : prob_func
+    _output_func =
+        isnothing(output_func) ?
+        _ensemble_dispatch_output_func(ensemblealg, progress_bar, ntraj, _standard_output_func) : output_func
     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,
-        ),
+        EnsembleProblem(prob.prob, prob_func = _prob_func, output_func = _output_func[1], safetycopy = false),
         prob.times,
         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
+    return reshape(sol_vec, size(iter))
 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...)

src/time_evolution/sesolve.jl

@@ -235,10 +235,11 @@ function sesolve_map(
 )
     # mapping initial states and parameters
     ψ0_iter = map(get_data, ψ0)
-    iter =
-        params isa NullParameters ? collect(Iterators.product(ψ0_iter, [params])) :
-        collect(Iterators.product(ψ0_iter, params...))
-    ntraj = length(iter)
+    if params isa NullParameters
+        iter = collect(Iterators.product(ψ0_iter, [params])) |> vec # convert nx1 Matrix into Vector
+    else
+        iter = collect(Iterators.product(ψ0_iter, params...))
+    end
 
     # we disable the progress bar of the sesolveProblem because we use a global progress bar for all the trajectories
     prob = sesolveProblem(
@@ -251,28 +252,42 @@ function sesolve_map(
         kwargs...,
     )
 
-    # generate and solve ensemble problem
-    _output_func = _ensemble_dispatch_output_func(ensemblealg, progress_bar, ntraj, _standard_output_func) # setup global progress bar
+    return sesolve_map(prob, iter, alg, ensemblealg; progress_bar = progress_bar)
+end
+sesolve_map(H::Union{AbstractQuantumObject{Operator},Tuple}, ψ0::QuantumObject{Ket}, tlist::AbstractVector; kwargs...) =
+    sesolve_map(H, [ψ0], tlist; kwargs...)
+
+# this method is for advanced usage
+# User can define their own iterator structure, prob_func and output_func
+#   - `prob_func`: Function to use for generating the ODEProblem.
+#   - `output_func`: a `Tuple` containing the `Function` to use for generating the output of a single trajectory, the (optional) `ProgressBar` object, and the (optional) `RemoteChannel` object.
+#
+# Return: An array of TimeEvolutionSol objects with the size same as the given iter.
+function sesolve_map(
+    prob::TimeEvolutionProblem{<:ODEProblem},
+    iter::AbstractArray,
+    alg::OrdinaryDiffEqAlgorithm = Tsit5(),
+    ensemblealg::EnsembleAlgorithm = EnsembleThreads();
+    prob_func::Union{Function,Nothing} = nothing,
+    output_func::Union{Tuple,Nothing} = nothing,
+    progress_bar::Union{Val,Bool} = Val(true),
+)
+    # generate ensemble problem
+    ntraj = length(iter)
+    _prob_func = isnothing(prob_func) ? (prob, i, repeat) -> _se_me_map_prob_func(prob, i, repeat, iter) : prob_func
+    _output_func =
+        isnothing(output_func) ?
+        _ensemble_dispatch_output_func(ensemblealg, progress_bar, ntraj, _standard_output_func) : output_func
     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,
-        ),
+        EnsembleProblem(prob.prob, prob_func = _prob_func, output_func = _output_func[1], safetycopy = false),
         prob.times,
         prob.dimensions,
         (progr = _output_func[2], channel = _output_func[3]),
     )
+
     sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj)
 
     # handle solution and make it become an Array of TimeEvolutionSol
     sol_vec = [_gen_sesolve_solution(sol[:, i], prob.times, prob.dimensions) 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
+    return reshape(sol_vec, size(iter))
 end
-sesolve_map(H::Union{AbstractQuantumObject{Operator},Tuple}, ψ0::QuantumObject{Ket}, tlist::AbstractVector; kwargs...) =
-    sesolve_map(H, [ψ0], tlist; kwargs...)