Introduce EnsembleAlgorithm and remove extra ntraj definitions

Author: albertomercurio

src/time_evolution/smesolve.jl

@@ -158,7 +158,7 @@ end
         params = NullParameters(),
         rng::AbstractRNG = default_rng(),
         ntraj::Int = 500,
-        ensemblealg = EnsembleThreads(),
+        ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
         prob_func::Union{Function, Nothing} = nothing,
         output_func::Union{Tuple,Nothing} = nothing,
         progress_bar::Union{Val,Bool} = Val(true),
@@ -225,7 +225,7 @@ function smesolveEnsembleProblem(
     params = NullParameters(),
     rng::AbstractRNG = default_rng(),
     ntraj::Int = 500,
-    ensemblealg = EnsembleThreads(),
+    ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
     progress_bar::Union{Val,Bool} = Val(true),
@@ -282,7 +282,7 @@ end
         params = NullParameters(),
         rng::AbstractRNG = default_rng(),
         ntraj::Int = 500,
-        ensemblealg = EnsembleThreads(),
+        ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
         prob_func::Union{Function, Nothing} = nothing,
         output_func::Union{Tuple,Nothing} = nothing,
         progress_bar::Union{Val,Bool} = Val(true),
@@ -351,7 +351,7 @@ function smesolve(
     params = NullParameters(),
     rng::AbstractRNG = default_rng(),
     ntraj::Int = 500,
-    ensemblealg = EnsembleThreads(),
+    ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
     progress_bar::Union{Val,Bool} = Val(true),
@@ -383,7 +383,7 @@ function smesolve(
     ens_prob::TimeEvolutionProblem,
     alg::StochasticDiffEqAlgorithm = SRA1(),
     ntraj::Int = 500,
-    ensemblealg = EnsembleThreads(),
+    ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
 )
     sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj)
 

src/time_evolution/ssesolve.jl

@@ -157,7 +157,7 @@ end
         params = NullParameters(),
         rng::AbstractRNG = default_rng(),
         ntraj::Int = 500,
-        ensemblealg = EnsembleThreads(),
+        ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
         prob_func::Union{Function, Nothing} = nothing,
         output_func::Union{Tuple,Nothing} = nothing,
         progress_bar::Union{Val,Bool} = Val(true),
@@ -224,7 +224,7 @@ function ssesolveEnsembleProblem(
     params = NullParameters(),
     rng::AbstractRNG = default_rng(),
     ntraj::Int = 500,
-    ensemblealg = EnsembleThreads(),
+    ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
     progress_bar::Union{Val,Bool} = Val(true),
@@ -279,7 +279,7 @@ end
         params = NullParameters(),
         rng::AbstractRNG = default_rng(),
         ntraj::Int = 500,
-        ensemblealg = EnsembleThreads(),
+        ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
         prob_func::Union{Function, Nothing} = nothing,
         output_func::Union{Tuple,Nothing} = nothing,
         progress_bar::Union{Val,Bool} = Val(true),
@@ -351,7 +351,7 @@ function ssesolve(
     params = NullParameters(),
     rng::AbstractRNG = default_rng(),
     ntraj::Int = 500,
-    ensemblealg = EnsembleThreads(),
+    ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
     progress_bar::Union{Val,Bool} = Val(true),
@@ -382,7 +382,7 @@ function ssesolve(
     ens_prob::TimeEvolutionProblem,
     alg::StochasticDiffEqAlgorithm = SRA1(),
     ntraj::Int = 500,
-    ensemblealg = EnsembleThreads(),
+    ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
 )
     sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj)
 

test/core-test/time_evolution.jl

@@ -255,7 +255,7 @@
         sol_se = sesolve(H_dr_fr, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
         sol_me = mesolve(H_dr_fr, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
         sol_mc = mcsolve(H_dr_fr, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), rng = rng)
-        # sol_sse = ssesolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
+        # sol_sse = ssesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), rng = rng)
 
         # Time Evolution in the lab frame
 
@@ -268,7 +268,7 @@
         sol_se_td = sesolve(H_td, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false), params = p)
         sol_me_td = mesolve(H_td, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p)
         sol_mc_td = mcsolve(H_td, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng)
-        # sol_sse_td = ssesolve(H_td, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng)
+        # sol_sse_td = ssesolve(H_td, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng)
 
         @test sol_se.expect ≈ sol_se_td.expect atol = 1e-6 * length(tlist)
         @test sol_me.expect ≈ sol_me_td.expect atol = 1e-6 * length(tlist)
@@ -282,7 +282,7 @@
         sol_me_td2 = mesolve(L_td, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p)
         sol_mc_td2 = mcsolve(H_td2, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng)
         # sol_sse_td2 =
-        # ssesolve(H_td2, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng)
+        # ssesolve(H_td2, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng)
 
         @test sol_se.expect ≈ sol_se_td2.expect atol = 1e-6 * length(tlist)
         @test sol_me.expect ≈ sol_me_td2.expect atol = 1e-6 * length(tlist)