diff --git a/CHANGELOG.md b/CHANGELOG.md index c03b26d02..d31c9ea89 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/time_evolution/smesolve.jl b/src/time_evolution/smesolve.jl index 793b53b9e..5d50851d1 100644 --- a/src/time_evolution/smesolve.jl +++ b/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,8 +110,6 @@ 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) @@ -119,7 +117,16 @@ function smesolveProblem( 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) diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl index e6c5515d8..f3173dd04 100644 --- a/src/time_evolution/ssesolve.jl +++ b/src/time_evolution/ssesolve.jl @@ -1,64 +1,5 @@ export ssesolveProblem, ssesolveEnsembleProblem, ssesolve -# TODO: Merge this with _stochastic_prob_func -function _ssesolve_prob_func(prob, i, repeat) - internal_params = prob.p - - global_rng = internal_params.global_rng - seed = internal_params.seeds[i] - traj_rng = typeof(global_rng)() - seed!(traj_rng, seed) - - noise = RealWienerProcess!( - prob.tspan[1], - zeros(internal_params.n_sc_ops), - zeros(internal_params.n_sc_ops), - save_everystep = false, - rng = traj_rng, - ) - - return remake(prob, noise = noise, seed = seed) -end - -#= - struct DiffusionOperator - -A struct to represent the diffusion operator. This is used to perform the diffusion process on N different Wiener processes. -=# -struct DiffusionOperator{T,OpType<:Tuple{Vararg{AbstractSciMLOperator}}} <: AbstractSciMLOperator{T} - ops::OpType - function DiffusionOperator(ops::OpType) where {OpType} - T = mapreduce(eltype, promote_type, ops) - return new{T,OpType}(ops) - end -end - -@generated function update_coefficients!(L::DiffusionOperator, u, p, t) - ops_types = L.parameters[2].parameters - N = length(ops_types) - quote - Base.@nexprs $N i -> begin - update_coefficients!(L.ops[i], u, p, t) - end - - nothing - end -end - -@generated function LinearAlgebra.mul!(v::AbstractVecOrMat, L::DiffusionOperator, u::AbstractVecOrMat) - ops_types = L.parameters[2].parameters - N = length(ops_types) - quote - M = length(u) - S = size(v) - (S[1] == M && S[2] == $N) || throw(DimensionMismatch("The size of the output vector is incorrect.")) - Base.@nexprs $N i -> begin - mul!(@view(v[:, i]), L.ops[i], u) - end - return v - end -end - # TODO: Implement the three-argument dot function for SciMLOperators.jl # Currently, we are assuming a time-independent MatrixOperator function _ssesolve_update_coeff(u, p, t, op) @@ -66,21 +7,6 @@ function _ssesolve_update_coeff(u, p, t, op) return real(dot(u, op.A, u)) #this is en/2: /2 = Re end -# Output function with progress bar update -function _ssesolve_output_func_progress(sol, i) - next!(sol.prob.p.progr) - return _stochastic_output_func(sol, i) -end - -# Output function with distributed channel update for progress bar -function _ssesolve_output_func_distributed(sol, i) - put!(sol.prob.p.progr_channel, true) - return _stochastic_output_func(sol, i) -end - -_ssesolve_dispatch_output_func(::ET) where {ET<:Union{EnsembleSerial,EnsembleThreads}} = _ssesolve_output_func_progress -_ssesolve_dispatch_output_func(::EnsembleDistributed) = _ssesolve_output_func_distributed - _ScalarOperator_e(op, f = +) = ScalarOperator(one(eltype(op)), (a, u, p, t) -> f(_ssesolve_update_coeff(u, p, t, op))) _ScalarOperator_e2_2(op, f = +) = @@ -93,7 +19,7 @@ _ScalarOperator_e2_2(op, f = +) = tlist::AbstractVector, 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..., @@ -127,7 +53,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is th - `tlist`: List of times at which to save either the state or the expectation values of the system. - `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. @@ -149,7 +75,7 @@ function ssesolveProblem( tlist::AbstractVector, 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..., @@ -184,8 +110,6 @@ function ssesolveProblem( D_l = map(op -> op + _ScalarOperator_e(op, -) * IdentityOperator(prod(dims)), sc_ops_evo_data) 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, SaveFuncSSESolve) kwargs4 = _ssesolve_add_normalize_cb(kwargs3) @@ -194,7 +118,18 @@ function ssesolveProblem( 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)) - return SDEProblem{true}(K, D, ψ0, tspan, p; noise_rate_prototype = noise_rate_prototype, noise = noise, kwargs4...) + prob = SDEProblem{true}( + K, + D, + ψ0, + tspan, + params; + noise_rate_prototype = noise_rate_prototype, + noise = noise, + kwargs4..., + ) + + return TimeEvolutionProblem(prob, tlist, dims) end @doc raw""" @@ -204,12 +139,12 @@ end tlist::AbstractVector, 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(), - prob_func::Function = _ssesolve_prob_func, - output_func::Function = _ssesolve_dispatch_output_func(ensemble_method), + prob_func::Union{Function, Nothing} = nothing, + output_func::Union{Tuple,Nothing} = nothing, progress_bar::Union{Val,Bool} = Val(true), kwargs..., ) @@ -242,13 +177,13 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is t - `tlist`: List of times at which to save either the state or the expectation values of the system. - `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()`. - `jump_callback`: The Jump Callback type: Discrete or Continuous. The default is `ContinuousLindbladJumpCallback()`, which is more precise. -- `prob_func`: Function to use for generating the ODEProblem. -- `output_func`: Function to use for generating the output of a single trajectory. +- `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. @@ -269,52 +204,42 @@ function ssesolveEnsembleProblem( tlist::AbstractVector, 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(), - prob_func::Function = _ssesolve_prob_func, - output_func::Function = _ssesolve_dispatch_output_func(ensemble_method), + prob_func::Union{Function,Nothing} = nothing, + output_func::Union{Tuple,Nothing} = nothing, progress_bar::Union{Val,Bool} = Val(true), kwargs..., ) - progr = ProgressBar(ntraj, enable = getVal(progress_bar)) - if ensemble_method isa EnsembleDistributed - progr_channel::RemoteChannel{Channel{Bool}} = RemoteChannel(() -> Channel{Bool}(1)) - @async while take!(progr_channel) - next!(progr) - end - params = merge(params, (progr_channel = progr_channel,)) - else - params = merge(params, (progr_trajectories = progr,)) - end - - # Stop the async task if an error occurs - try - seeds = map(i -> rand(rng, UInt64), 1:ntraj) - prob_sse = ssesolveProblem( - H, - ψ0, - tlist, - sc_ops; - e_ops = e_ops, - params = merge(params, (global_rng = rng, seeds = seeds)), - rng = rng, - progress_bar = Val(false), - kwargs..., - ) - - # safetycopy is set to true because the K and D functions cannot be currently deepcopied. - # the memory overhead shouldn't be too large, compared to the safetycopy=false case. - ensemble_prob = EnsembleProblem(prob_sse, prob_func = prob_func, output_func = output_func, safetycopy = true) - - return ensemble_prob - catch e - if ensemble_method isa EnsembleDistributed - put!(progr_channel, false) - end - rethrow(e) - end + _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 + + prob_sme = ssesolveProblem( + H, + ψ0, + tlist, + sc_ops; + e_ops = e_ops, + params = params, + rng = rng, + progress_bar = Val(false), + kwargs..., + ) + + ensemble_prob = TimeEvolutionProblem( + EnsembleProblem(prob_sme, prob_func = _prob_func, output_func = _output_func[1], safetycopy = true), + prob_sme.times, + prob_sme.dimensions, + (progr = _output_func[2], channel = _output_func[3]), + ) + + return ensemble_prob end @doc raw""" @@ -325,12 +250,12 @@ 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(), - prob_func::Function = _ssesolve_prob_func, - output_func::Function = _ssesolve_dispatch_output_func(ensemble_method), + prob_func::Union{Function, Nothing} = nothing, + output_func::Union{Tuple,Nothing} = nothing, progress_bar::Union{Val,Bool} = Val(true), kwargs..., ) @@ -367,12 +292,12 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is th - `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. -- `output_func`: Function to use for generating the output of a single trajectory. +- `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. @@ -395,12 +320,12 @@ function ssesolve( 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(), - prob_func::Function = _ssesolve_prob_func, - output_func::Function = _ssesolve_dispatch_output_func(ensemble_method), + prob_func::Union{Function,Nothing} = nothing, + output_func::Union{Tuple,Nothing} = nothing, progress_bar::Union{Val,Bool} = Val(true), kwargs..., ) @@ -420,52 +345,39 @@ function ssesolve( kwargs..., ) - return ssesolve(ens_prob; alg = alg, ntraj = ntraj, ensemble_method = ensemble_method) + return ssesolve(ens_prob, alg, ntraj, ensemble_method) end function ssesolve( - ens_prob::EnsembleProblem; + ens_prob::TimeEvolutionProblem, alg::StochasticDiffEqAlgorithm = SRA1(), ntraj::Int = 1, ensemble_method = EnsembleThreads(), ) - # Stop the async task if an error occurs - try - sol = solve(ens_prob, alg, ensemble_method, trajectories = ntraj) - - if ensemble_method isa EnsembleDistributed - put!(sol[:, 1].prob.p.progr_channel, false) - end - - _sol_1 = sol[:, 1] - _expvals_sol_1 = _se_me_sse_get_expvals(_sol_1) - - normalize_states = Val(false) - dims = _sol_1.prob.p.Hdims - _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) - states = map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol)) - - expvals = - _se_me_sse_get_expvals(_sol_1) isa Nothing ? nothing : - dropdims(sum(expvals_all, dims = 3), dims = 3) ./ length(sol) - - return TimeEvolutionStochasticSol( - ntraj, - _sol_1.prob.p.times, - states, - expvals, - expvals_all, - sol.converged, - _sol_1.alg, - _sol_1.prob.kwargs[:abstol], - _sol_1.prob.kwargs[:reltol], - ) - catch e - if ensemble_method isa EnsembleDistributed - put!(ens_prob.prob.p.progr_channel, false) - end - rethrow(e) - end + sol = _ensemble_dispatch_solve(ens_prob, alg, ensemble_method, ntraj) + + _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) + states = map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol)) + + expvals = + _se_me_sse_get_expvals(_sol_1) isa Nothing ? nothing : + dropdims(sum(expvals_all, dims = 3), dims = 3) ./ length(sol) + + return TimeEvolutionStochasticSol( + ntraj, + ens_prob.times, + states, + expvals, + expvals_all, + sol.converged, + _sol_1.alg, + _sol_1.prob.kwargs[:abstol], + _sol_1.prob.kwargs[:reltol], + ) end diff --git a/src/time_evolution/time_evolution.jl b/src/time_evolution/time_evolution.jl index 20b7f9257..95f9b1089 100644 --- a/src/time_evolution/time_evolution.jl +++ b/src/time_evolution/time_evolution.jl @@ -296,9 +296,9 @@ function _ensemble_dispatch_output_func( end end -function _ensemble_dispatch_prob_func(rng, ntraj, tlist, prob_func) +function _ensemble_dispatch_prob_func(rng, ntraj, tlist, prob_func; kwargs...) seeds = map(i -> rand(rng, UInt64), 1:ntraj) - return (prob, i, repeat) -> prob_func(prob, i, repeat, rng, seeds, tlist) + return (prob, i, repeat) -> prob_func(prob, i, repeat, rng, seeds, tlist; kwargs...) end function _ensemble_dispatch_solve( @@ -336,17 +336,15 @@ end #= Stochastic funcs =# -function _stochastic_prob_func(prob, i, repeat, rng, seeds, tlist) - params = prob.prob.p - +function _stochastic_prob_func(prob, i, repeat, rng, seeds, tlist; kwargs...) seed = seeds[i] traj_rng = typeof(rng)() seed!(traj_rng, seed) noise = RealWienerProcess!( prob.prob.tspan[1], - zeros(params.n_sc_ops), - zeros(params.n_sc_ops), + zeros(kwargs[:n_sc_ops]), + zeros(kwargs[:n_sc_ops]), save_everystep = false, rng = traj_rng, ) @@ -357,6 +355,45 @@ end # Standard output function _stochastic_output_func(sol, i) = (sol, false) +#= + struct DiffusionOperator + +A struct to represent the diffusion operator. This is used to perform the diffusion process on N different Wiener processes. +=# +struct DiffusionOperator{T,OpType<:Tuple{Vararg{AbstractSciMLOperator}}} <: AbstractSciMLOperator{T} + ops::OpType + function DiffusionOperator(ops::OpType) where {OpType} + T = mapreduce(eltype, promote_type, ops) + return new{T,OpType}(ops) + end +end + +@generated function update_coefficients!(L::DiffusionOperator, u, p, t) + ops_types = L.parameters[2].parameters + N = length(ops_types) + quote + Base.@nexprs $N i -> begin + update_coefficients!(L.ops[i], u, p, t) + end + + nothing + end +end + +@generated function LinearAlgebra.mul!(v::AbstractVecOrMat, L::DiffusionOperator, u::AbstractVecOrMat) + ops_types = L.parameters[2].parameters + N = length(ops_types) + quote + M = length(u) + S = size(v) + (S[1] == M && S[2] == $N) || throw(DimensionMismatch("The size of the output vector is incorrect.")) + Base.@nexprs $N i -> begin + mul!(@view(v[:, i]), L.ops[i], u) + end + return v + end +end + ####################################### function liouvillian_floquet(