fix mapped solvers with NullParameters

Author: ytdHuang

src/time_evolution/mesolve.jl

@@ -240,7 +240,7 @@ end
         alg::OrdinaryDiffEqAlgorithm = Tsit5(),
         ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::Tuple = (NullParameters(),),
+        params::Union{NullParameters,Tuple} = NullParameters(),
         progress_bar::Union{Val,Bool} = Val(true),
         kwargs...,
     )
@@ -293,7 +293,7 @@ function mesolve_map(
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::Tuple = (NullParameters(),),
+    params::Union{NullParameters,Tuple} = NullParameters(),
     progress_bar::Union{Val,Bool} = Val(true),
     kwargs...,
 ) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet}}
@@ -319,7 +319,9 @@ function mesolve_map(
             to_dense(T, mat2vec(ket2dm(state).data))
         end
     end
-    iter = collect(Iterators.product(ψ0_iter, params...))
+    iter =
+        params isa NullParameters ? collect(Iterators.product(ψ0_iter, [params])) :
+        collect(Iterators.product(ψ0_iter, params...))
     ntraj = length(iter)
 
     # we disable the progress bar of the mesolveProblem because we use a global progress bar for all the trajectories
@@ -352,7 +354,11 @@ function mesolve_map(
     # 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
-    return reshape(sol_vec, size(iter))
+    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},

src/time_evolution/sesolve.jl

@@ -185,7 +185,7 @@ end
         alg::OrdinaryDiffEqAlgorithm = Tsit5(),
         ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::Tuple = (NullParameters(),),
+        params::Union{NullParameters,Tuple} = NullParameters(),
         progress_bar::Union{Val,Bool} = Val(true),
         kwargs...,
     )
@@ -229,13 +229,15 @@ function sesolve_map(
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::Tuple = (NullParameters(),),
+    params::Union{NullParameters,Tuple} = NullParameters(),
     progress_bar::Union{Val,Bool} = Val(true),
     kwargs...,
 )
     # mapping initial states and parameters
     ψ0_iter = map(get_data, ψ0)
-    iter = collect(Iterators.product(ψ0_iter, params...))
+    iter =
+        params isa NullParameters ? collect(Iterators.product(ψ0_iter, [params])) :
+        collect(Iterators.product(ψ0_iter, params...))
     ntraj = length(iter)
 
     # we disable the progress bar of the sesolveProblem because we use a global progress bar for all the trajectories
@@ -266,7 +268,11 @@ function sesolve_map(
 
     # 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
-    return reshape(sol_vec, size(iter))
+    if params isa NullParameters # if no parameters specified, just return a Vector
+        return sol_vec
+    else
+        return reshape(sol_vec, size(iter))
+    end
 end
 sesolve_map(H::Union{AbstractQuantumObject{Operator},Tuple}, ψ0::QuantumObject{Ket}, tlist::AbstractVector; kwargs...) =
     sesolve_map(H, [ψ0], tlist; kwargs...)

test/core-test/time_evolution.jl

@@ -171,9 +171,11 @@ end
     ωq_fun(p, t) = p[2]
     H = QobjEvo(a' * a, ωc_fun) + QobjEvo(σz / 2, ωq_fun) + g * (a' * σm + a * σm')
 
-    sols1 = sesolve_map(H, ψ_0_e, tlist; e_ops = e_ops, params = (ωc_list, ωq_list))
+    sols0 = sesolve_map(TESetup.H, ψ0_list, tlist; e_ops = e_ops) # no params, but test progress_bar
+    sols1 = sesolve_map(H, ψ_0_e, tlist; e_ops = e_ops, params = (ωc_list, ωq_list), progress_bar = Val(false))
     sols2 = sesolve_map(H, ψ0_list, tlist; e_ops = e_ops, params = (ωc_list, ωq_list), progress_bar = Val(false))
-
+    @test size(sols0) == (2,)
+    @test sols0 isa Vector{<:TimeEvolutionSol}
     @test size(sols1) == (1, 3, 4)
     @test sols1 isa Array{<:TimeEvolutionSol}
     @test size(sols2) == (2, 3, 4)
@@ -193,6 +195,7 @@ end
     end
 
     @testset "Type Inference sesolve_map" begin
+        @inferred sesolve_map(TESetup.H, ψ0_list, tlist; e_ops = e_ops, progress_bar = Val(false)) # no params
         @inferred sesolve_map(H, ψ0_list, tlist; e_ops = e_ops, params = (ωc_list, ωq_list), progress_bar = Val(false))
     end
 end
@@ -324,14 +327,18 @@ end
     ωq_fun(p, t) = p[2]
     H = QobjEvo(a' * a, ωc_fun) + QobjEvo(σz / 2, ωq_fun) + g * (a' * σm + a * σm')
 
+    # Test with multiple initial states but no params (this also tests progress_bar)
+    sols0 = mesolve_map(TESetup.H, ψ0_list, tlist, c_ops; e_ops = e_ops)
     # Test with single initial state
-    sols1 = mesolve_map(H, ψ_0_e, tlist, c_ops; e_ops = e_ops, params = (ωc_list, ωq_list))
+    sols1 = mesolve_map(H, ψ_0_e, tlist, c_ops; e_ops = e_ops, params = (ωc_list, ωq_list), progress_bar = Val(false))
     # Test with multiple initial states
     sols2 = mesolve_map(H, ψ0_list, tlist, c_ops; e_ops = e_ops, params = (ωc_list, ωq_list), progress_bar = Val(false))
 
     # Test redirect to sesolve_map when c_ops is nothing
     sols3 = mesolve_map(H, ψ0_list, tlist; e_ops = e_ops, params = (ωc_list, ωq_list), progress_bar = Val(false))
 
+    @test size(sols0) == (2,)
+    @test sols0 isa Vector{<:TimeEvolutionSol}
     @test size(sols1) == (1, 3, 4)
     @test sols1 isa Array{<:TimeEvolutionSol}
     @test size(sols2) == (2, 3, 4)
@@ -368,6 +375,7 @@ end
     @test sols5 isa Array{<:TimeEvolutionSol}
 
     @testset "Type Inference mesolve_map" begin
+        @inferred mesolve_map(TESetup.H, ψ0_list, tlist, c_ops; e_ops = e_ops, progress_bar = Val(false)) # no params
         @inferred mesolve_map(
             H,
             ψ0_list,

test/runtests.jl

@@ -5,7 +5,7 @@ using Pkg
 const GROUP_LIST = String["All", "Core", "Code-Quality", "AutoDiff_Ext", "Makie_Ext", "CUDA_Ext"]
 
 const GROUP = get(ENV, "GROUP", "All")
-(GROUP in GROUP_LIST) || throw(ArgumentError("Unknown GROUP = $GROUP"))
+(GROUP in GROUP_LIST) || throw(ArgumentError("Unknown GROUP = $GROUP\nThe allowed groups are: $GROUP_LIST\n"))
 
 # Core tests
 if (GROUP == "All") || (GROUP == "Core")