Improve ensemble generation of ssesolve and change parameters on stochastic processes (#403)

Author: albertomercurio

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fix `smesolve` for specifying initial state as density matrix. ([#395])
 - Add more generic solver for `steadystate_floquet` to allow more linear solvers. ([#396])
 - Fix time evolution output when using `saveat` keyword argument. ([#398])
+- Improve ensemble generation of `ssesolve` and change parameters handling on stochastic processes. ([#403])
 
 ## [v0.26.0]
 Release date: 2025-02-09
@@ -128,3 +129,4 @@ Release date: 2024-11-13
 [#395]: https://github.com/qutip/QuantumToolbox.jl/issues/395
 [#396]: https://github.com/qutip/QuantumToolbox.jl/issues/396
 [#398]: https://github.com/qutip/QuantumToolbox.jl/issues/398
+[#403]: https://github.com/qutip/QuantumToolbox.jl/issues/403

src/time_evolution/smesolve.jl

@@ -17,7 +17,7 @@ _smesolve_ScalarOperator(op_vec) =
         c_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
         sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::NamedTuple = NamedTuple(),
+        params = NullParameters(),
         rng::AbstractRNG = default_rng(),
         progress_bar::Union{Val,Bool} = Val(true),
         kwargs...,
@@ -51,7 +51,7 @@ Above, ``\hat{C}_i`` represent the collapse operators related to pure dissipatio
 - `c_ops`: List of collapse operators ``\{\hat{C}_i\}_i``. It can be either a `Vector` or a `Tuple`.
 - `sc_ops`: List of stochastic collapse operators ``\{\hat{S}_n\}_n``. It can be either a `Vector` or a `Tuple`.
 - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`.
-- `params`: `NamedTuple` of parameters to pass to the solver.
+- `params`: `NullParameters` of parameters to pass to the solver.
 - `rng`: Random number generator for reproducibility.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
 - `kwargs`: The keyword arguments for the ODEProblem.
@@ -74,7 +74,7 @@ function smesolveProblem(
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
     sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::NamedTuple = NamedTuple(),
+    params = NullParameters(),
     rng::AbstractRNG = default_rng(),
     progress_bar::Union{Val,Bool} = Val(true),
     kwargs...,
@@ -110,16 +110,23 @@ function smesolveProblem(
     end
     D = DiffusionOperator(D_l)
 
-    p = (progr = progr, times = tlist, Hdims = dims, n_sc_ops = length(sc_ops), params...)
-
     kwargs2 = _merge_saveat(tlist, e_ops, DEFAULT_SDE_SOLVER_OPTIONS; kwargs...)
     kwargs3 = _generate_se_me_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2, SaveFuncMESolve)
 
     tspan = (tlist[1], tlist[end])
     noise =
         RealWienerProcess!(tlist[1], zeros(length(sc_ops)), zeros(length(sc_ops)), save_everystep = false, rng = rng)
     noise_rate_prototype = similar(ρ0, length(ρ0), length(sc_ops))
-    prob = SDEProblem{true}(K, D, ρ0, tspan, p; noise_rate_prototype = noise_rate_prototype, noise = noise, kwargs3...)
+    prob = SDEProblem{true}(
+        K,
+        D,
+        ρ0,
+        tspan,
+        params;
+        noise_rate_prototype = noise_rate_prototype,
+        noise = noise,
+        kwargs3...,
+    )
 
     return TimeEvolutionProblem(prob, tlist, dims)
 end
@@ -132,7 +139,7 @@ end
         c_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
         sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::NamedTuple = NamedTuple(),
+        params = NullParameters(),
         rng::AbstractRNG = default_rng(),
         ntraj::Int = 1,
         ensemble_method = EnsembleThreads(),
@@ -170,11 +177,11 @@ Above, ``\hat{C}_i`` represent the collapse operators related to pure dissipatio
 - `c_ops`: List of collapse operators ``\{\hat{C}_i\}_i``. It can be either a `Vector` or a `Tuple`.
 - `sc_ops`: List of stochastic collapse operators ``\{\hat{S}_n\}_n``. It can be either a `Vector` or a `Tuple`.
 - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`.
-- `params`: `NamedTuple` of parameters to pass to the solver.
+- `params`: `NullParameters` of parameters to pass to the solver.
 - `rng`: Random number generator for reproducibility.
 - `ntraj`: Number of trajectories to use.
 - `ensemble_method`: Ensemble method to use. Default to `EnsembleThreads()`.
-- `prob_func`: Function to use for generating the ODEProblem.
+- `prob_func`: Function to use for generating the SDEProblem.
 - `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.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
 - `kwargs`: The keyword arguments for the ODEProblem.
@@ -197,7 +204,7 @@ function smesolveEnsembleProblem(
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
     sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::NamedTuple = NamedTuple(),
+    params = NullParameters(),
     rng::AbstractRNG = default_rng(),
     ntraj::Int = 1,
     ensemble_method = EnsembleThreads(),
@@ -207,7 +214,8 @@ function smesolveEnsembleProblem(
     kwargs...,
 ) where {StateOpType<:Union{KetQuantumObject,OperatorQuantumObject}}
     _prob_func =
-        isnothing(prob_func) ? _ensemble_dispatch_prob_func(rng, ntraj, tlist, _stochastic_prob_func) : prob_func
+        isnothing(prob_func) ?
+        _ensemble_dispatch_prob_func(rng, ntraj, tlist, _stochastic_prob_func; n_sc_ops = length(sc_ops)) : prob_func
     _output_func =
         output_func isa Nothing ?
         _ensemble_dispatch_output_func(ensemble_method, progress_bar, ntraj, _stochastic_output_func) : output_func
@@ -244,7 +252,7 @@ end
         sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
         alg::StochasticDiffEqAlgorithm = SRA1(),
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        params::NamedTuple = NamedTuple(),
+        params = NullParameters(),
         rng::AbstractRNG = default_rng(),
         ntraj::Int = 1,
         ensemble_method = EnsembleThreads(),
@@ -283,11 +291,11 @@ Above, ``\hat{C}_i`` represent the collapse operators related to pure dissipatio
 - `sc_ops`: List of stochastic collapse operators ``\{\hat{S}_n\}_n``. It can be either a `Vector` or a `Tuple`.
 - `alg`: The algorithm to use for the stochastic differential equation. Default is `SRA1()`.
 - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`.
-- `params`: `NamedTuple` of parameters to pass to the solver.
+- `params`: `NullParameters` of parameters to pass to the solver.
 - `rng`: Random number generator for reproducibility.
 - `ntraj`: Number of trajectories to use.
 - `ensemble_method`: Ensemble method to use. Default to `EnsembleThreads()`.
-- `prob_func`: Function to use for generating the ODEProblem.
+- `prob_func`: Function to use for generating the SDEProblem.
 - `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.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
 - `kwargs`: The keyword arguments for the ODEProblem.
@@ -311,7 +319,7 @@ function smesolve(
     sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     alg::StochasticDiffEqAlgorithm = SRA1(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params::NamedTuple = NamedTuple(),
+    params = NullParameters(),
     rng::AbstractRNG = default_rng(),
     ntraj::Int = 1,
     ensemble_method = EnsembleThreads(),
@@ -351,7 +359,6 @@ function smesolve(
     _sol_1 = sol[:, 1]
     _expvals_sol_1 = _se_me_sse_get_expvals(_sol_1)
 
-    normalize_states = Val(false)
     dims = ens_prob.dimensions
     _expvals_all = _expvals_sol_1 isa Nothing ? nothing : map(i -> _se_me_sse_get_expvals(sol[:, i]), eachindex(sol))
     expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all)