Merge branch 'main' into liouvillian

Author: ytdHuang

.github/workflows/Benchmarks.yml

@@ -37,7 +37,7 @@ jobs:
       - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
-          version: '1'
+          version: '1.11'
           arch: x64
       - uses: julia-actions/cache@v2
       - name: Run benchmark

.github/workflows/CleanPreviewDoc.yml

@@ -15,7 +15,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout gh-pages branch
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           ref: gh-pages
       - name: Delete preview(s) and reset history

.github/workflows/SpellCheck.yml

@@ -10,4 +10,4 @@ jobs:
       - name: Checkout Actions Repository
         uses: actions/checkout@v5
       - name: Check spelling
-        uses: crate-ci/typos@v1.37.2
+        uses: crate-ci/typos@v1.38.1

CHANGELOG.md

@@ -7,6 +7,7 @@ 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])
 - Generalize the definition of `liouvillian`. It no longer expects the Hamiltonian to be Hermitian. ([#541])
 
 ## [v0.37.0]
@@ -339,3 +340,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

benchmarks/Project.toml

@@ -1,4 +1,8 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 QuantumToolbox = "6c2fb7c5-b903-41d2-bc5e-5a7c320b9fab"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

benchmarks/autodiff.jl

@@ -0,0 +1,92 @@
+function benchmark_autodiff!(SUITE)
+    # Use harmonic oscillator system for both sesolve and mesolve
+    N = 20
+    a = destroy(N)
+    ψ0 = fock(N, 0)
+    tlist = range(0, 40, 100)
+
+    # ---- SESOLVE ----
+    # For direct Forward differentiation
+    function my_f_sesolve_direct(p)
+        H = p[1] * a' * a + p[2] * (a + a')
+        sol = sesolve(H, ψ0, tlist, progress_bar = Val(false))
+        return real(expect(a' * a, sol.states[end]))
+    end
+
+    # For SciMLSensitivity.jl (reverse mode with Zygote and Enzyme)
+    coef_Δ(p, t) = p[1]
+    coef_F(p, t) = p[2]
+    H_evo = QobjEvo(a' * a, coef_Δ) + QobjEvo(a + a', coef_F)
+
+    function my_f_sesolve(p)
+        sol = sesolve(
+            H_evo,
+            ψ0,
+            tlist,
+            progress_bar = Val(false),
+            params = p,
+            sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()),
+        )
+        return real(expect(a' * a, sol.states[end]))
+    end
+
+    # ---- MESOLVE ----
+    # For direct Forward differentiation
+    function my_f_mesolve_direct(p)
+        H = p[1] * a' * a + p[2] * (a + a')
+        c_ops = [sqrt(p[3]) * a]
+        sol = mesolve(H, ψ0, tlist, c_ops, progress_bar = Val(false))
+        return real(expect(a' * a, sol.states[end]))
+    end
+
+    # For SciMLSensitivity.jl (reverse mode with Zygote and Enzyme)
+    coef_γ(p, t) = sqrt(p[3])
+    c_ops = [QobjEvo(a, coef_γ)]
+    L = liouvillian(H_evo, c_ops)
+
+    function my_f_mesolve(p)
+        sol = mesolve(
+            L,
+            ψ0,
+            tlist,
+            progress_bar = Val(false),
+            params = p,
+            sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()),
+        )
+        return real(expect(a' * a, sol.states[end]))
+    end
+
+    # Parameters for benchmarks
+    params_sesolve = [1.0, 1.0]
+    params_mesolve = [1.0, 1.0, 1.0]
+
+    # Benchmark sesolve - Forward
+    SUITE["Autodiff"]["sesolve"]["Forward"] = @benchmarkable ForwardDiff.gradient($my_f_sesolve_direct, $params_sesolve)
+
+    # Benchmark sesolve - Reverse (Zygote)
+    SUITE["Autodiff"]["sesolve"]["Reverse (Zygote)"] = @benchmarkable Zygote.gradient($my_f_sesolve, $params_sesolve)
+
+    # Benchmark sesolve - Reverse (Enzyme)
+    SUITE["Autodiff"]["sesolve"]["Reverse (Enzyme)"] = @benchmarkable Enzyme.autodiff(
+        Enzyme.set_runtime_activity(Enzyme.Reverse),
+        Const($my_f_sesolve),
+        Active,
+        Duplicated($params_sesolve, dparams_sesolve),
+    ) setup=(dparams_sesolve = Enzyme.make_zero($params_sesolve))
+
+    # Benchmark mesolve - Forward
+    SUITE["Autodiff"]["mesolve"]["Forward"] = @benchmarkable ForwardDiff.gradient($my_f_mesolve_direct, $params_mesolve)
+
+    # Benchmark mesolve - Reverse (Zygote)
+    SUITE["Autodiff"]["mesolve"]["Reverse (Zygote)"] = @benchmarkable Zygote.gradient($my_f_mesolve, $params_mesolve)
+
+    # Benchmark mesolve - Reverse (Enzyme)
+    SUITE["Autodiff"]["mesolve"]["Reverse (Enzyme)"] = @benchmarkable Enzyme.autodiff(
+        Enzyme.set_runtime_activity(Enzyme.Reverse),
+        Const($my_f_mesolve),
+        Active,
+        Duplicated($params_mesolve, dparams_mesolve),
+    ) setup=(dparams_mesolve = Enzyme.make_zero($params_mesolve))
+
+    return nothing
+end

benchmarks/runbenchmarks.jl

@@ -3,6 +3,10 @@ using LinearAlgebra
 using SparseArrays
 using QuantumToolbox
 using SciMLBase: EnsembleSerial, EnsembleThreads
+using ForwardDiff
+using Zygote
+using Enzyme: Enzyme, Const, Active, Duplicated
+using SciMLSensitivity: BacksolveAdjoint, EnzymeVJP
 
 BLAS.set_num_threads(1)
 
@@ -14,13 +18,15 @@ include("dynamical_shifted_fock.jl")
 include("eigenvalues.jl")
 include("steadystate.jl")
 include("timeevolution.jl")
+include("autodiff.jl")
 
 benchmark_correlations_and_spectrum!(SUITE)
 benchmark_dfd!(SUITE)
 benchmark_dsf!(SUITE)
 benchmark_eigenvalues!(SUITE)
 benchmark_steadystate!(SUITE)
 benchmark_timeevolution!(SUITE)
+benchmark_autodiff!(SUITE)
 
 BenchmarkTools.tune!(SUITE)
 results = BenchmarkTools.run(SUITE, verbose = true)

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...)