return problem feature added

Author: henhen724

Manifest.toml

@@ -13,11 +13,11 @@ Calculate MCWF trajectory where the Hamiltonian is given in hermitian form.
 For more information see: [`mcwf`](@ref)
 """
 function mcwf_h(tspan, psi0::Ket, H::AbstractOperator, J;
-        seed=rand(UInt), rates=nothing,
-        fout=nothing, Jdagger=dagger.(J),
-        tmp=copy(psi0),
-        display_beforeevent=false, display_afterevent=false,
-        kwargs...)
+    seed=rand(UInt), rates=nothing,
+    fout=nothing, Jdagger=dagger.(J),
+    tmp=copy(psi0),
+    display_beforeevent=false, display_afterevent=false,
+    kwargs...)
     _check_const(H)
     _check_const.(J)
     _check_const.(Jdagger)
@@ -47,9 +47,9 @@ H_{nh} = H - \\frac{i}{2} \\sum_k J^†_k J_k
 For more information see: [`mcwf`](@ref)
 """
 function mcwf_nh(tspan, psi0::Ket, Hnh::AbstractOperator, J;
-        seed=rand(UInt), fout=nothing,
-        display_beforeevent=false, display_afterevent=false,
-        kwargs...)
+    seed=rand(UInt), fout=nothing,
+    display_beforeevent=false, display_afterevent=false,
+    kwargs...)
     _check_const(Hnh)
     _check_const.(J)
     check_mcwf(psi0, Hnh, J, J, nothing)
@@ -107,10 +107,10 @@ of the jump operators with which the jump occured, respectively.
 * `kwargs...`: Further arguments are passed on to the ode solver.
 """
 function mcwf(tspan, psi0::Ket, H::AbstractOperator, J;
-        seed=rand(UInt), rates=nothing,
-        fout=nothing, Jdagger=dagger.(J),
-        display_beforeevent=false, display_afterevent=false,
-        kwargs...)
+    seed=rand(UInt), rates=nothing,
+    fout=nothing, Jdagger=dagger.(J),
+    display_beforeevent=false, display_afterevent=false,
+    kwargs...)
     _check_const(H)
     _check_const.(J)
     _check_const.(Jdagger)
@@ -132,12 +132,12 @@ function mcwf(tspan, psi0::Ket, H::AbstractOperator, J;
     else
         Hnh = copy(H)
         if isa(rates, Nothing)
-            for i=1:length(J)
-                Hnh -= complex(float(eltype(H)))(0.5im)*Jdagger[i]*J[i]
+            for i = 1:length(J)
+                Hnh -= complex(float(eltype(H)))(0.5im) * Jdagger[i] * J[i]
             end
         else
-            for i=1:length(J)
-                Hnh -= complex(float(eltype(H)))(0.5im*rates[i])*Jdagger[i]*J[i]
+            for i = 1:length(J)
+                Hnh -= complex(float(eltype(H)))(0.5im * rates[i]) * Jdagger[i] * J[i]
             end
         end
         dmcwf_nh_ = let Hnh = Hnh  # Hnh type often not inferrable
@@ -322,25 +322,26 @@ Integrate a single Monte Carlo wave function trajectory.
 * `kwargs`: Further arguments are passed on to the ode solver.
 """
 function integrate_mcwf(dmcwf::T, jumpfun::J, tspan,
-                        psi0, seed, fout;
-                        display_beforeevent=false, display_afterevent=false,
-                        display_jumps=false,
-                        rng_state=nothing,
-                        save_everystep=false, callback=nothing,
-                        saveat=tspan,
-                        alg=OrdinaryDiffEq.DP5(),
-                        kwargs...) where {T, J}
+    psi0, seed, fout;
+    display_beforeevent=false, display_afterevent=false,
+    display_jumps=false,
+    rng_state=nothing,
+    save_everystep=false, callback=nothing,
+    saveat=tspan,
+    alg=OrdinaryDiffEq.DP5(),
+    return_problem=false,
+    kwargs...) where {T,J}
 
     tspan_ = convert(Vector{float(eltype(tspan))}, tspan)
     # Display before or after events
-    function save_func!(affect!,integrator)
+    function save_func!(affect!, integrator)
         affect!.saveiter += 1
         copyat_or_push!(affect!.saved_values.t, affect!.saveiter, integrator.t)
         copyat_or_push!(affect!.saved_values.saveval, affect!.saveiter,
-            affect!.save_func(integrator.u, integrator.t, integrator),Val{false})
+            affect!.save_func(integrator.u, integrator.t, integrator), Val{false})
         return nothing
     end
-    no_save_func!(affect!,integrator) = nothing
+    no_save_func!(affect!, integrator) = nothing
     save_before! = display_beforeevent ? save_func! : no_save_func!
     save_after! = display_afterevent ? save_func! : no_save_func!
 
@@ -349,8 +350,8 @@ function integrate_mcwf(dmcwf::T, jumpfun::J, tspan,
     jump_index = Int[]
 
     function jump_saver(t, i)
-        push!(jump_t,t)
-        push!(jump_index,i)
+        push!(jump_t, t)
+        push!(jump_index, i)
         return nothing
     end
     no_jump_saver(t, i) = nothing
@@ -362,51 +363,59 @@ function integrate_mcwf(dmcwf::T, jumpfun::J, tspan,
 
     fout_ = let state = state, fout = fout
         function fout_(x, t, integrator)
-            recast!(state,x)
+            recast!(state, x)
             fout(t, state)
         end
     end
 
     out_type = pure_inference(fout, Tuple{eltype(tspan_),typeof(state)})
-    out = DiffEqCallbacks.SavedValues(eltype(tspan_),out_type)
-    scb = DiffEqCallbacks.SavingCallback(fout_,out,saveat=tspan_,
-                                         save_everystep=save_everystep,
-                                         save_start = false)
+    out = DiffEqCallbacks.SavedValues(eltype(tspan_), out_type)
+    scb = DiffEqCallbacks.SavingCallback(fout_, out, saveat=tspan_,
+        save_everystep=save_everystep,
+        save_start=false)
 
     cb = jump_callback(jumpfun, seed, scb, save_before!, save_after!, save_t_index, psi0, rng_state)
-    full_cb = OrdinaryDiffEq.CallbackSet(callback,cb,scb)
+    full_cb = OrdinaryDiffEq.CallbackSet(callback, cb, scb)
 
     df_ = let state = state, dstate = dstate  # help inference along
         function df_(dx, x, p, t)
-            recast!(state,x)
-            recast!(dstate,dx)
+            recast!(state, x)
+            recast!(dstate, dx)
             dmcwf(t, state, dstate)
-            recast!(dx,dstate)
+            recast!(dx, dstate)
             return nothing
         end
     end
 
-    prob = OrdinaryDiffEq.ODEProblem{true}(df_, as_vector(psi0), (tspan_[1],tspan_[end]))
-
-    sol = OrdinaryDiffEq.solve(
-                prob,
-                alg;
-                reltol = 1.0e-6,
-                abstol = 1.0e-8,
-                save_everystep = false, save_start = false,
-                save_end = false,
-                callback=full_cb, kwargs...)
+    prob = OrdinaryDiffEq.ODEProblem{true}(df_, as_vector(psi0), (tspan_[1], tspan_[end]))
 
-    if display_jumps
-        return out.t, out.saveval, jump_t, jump_index
+    if return_problem
+        if display_jumps
+            return Dict("out" => out, "jump_t" => jump_t, "jump_index" => jump_index, "prob" => prob, "alg" => alg, "solve_kwargs" => (reltol=1.0e-6, abstol=1.0e-8, save_everystep=false, save_start=false, save_end=false, callback=full_cb, kwargs...))
+        else
+            return Dict("out" => out, "prob" => prob, "alg" => alg, "solve_kwargs" => (reltol=1.0e-6, abstol=1.0e-8, save_everystep=false, save_start=false, save_end=false, callback=full_cb, kwargs...))
+        end
     else
-        return out.t, out.saveval
+        sol = OrdinaryDiffEq.solve(
+            prob,
+            alg;
+            reltol=1.0e-6,
+            abstol=1.0e-8,
+            save_everystep=false, save_start=false,
+            save_end=false,
+            callback=full_cb, kwargs...)
+
+        if display_jumps
+            return out.t, out.saveval, jump_t, jump_index
+        else
+            return out.t, out.saveval
+        end
     end
 end
 
 function integrate_mcwf(dmcwf, jumpfun, tspan,
-                        psi0, seed, fout::Nothing;
-                        kwargs...)
+    psi0, seed, fout::Nothing;
+    kwargs...)
     function fout_(t, x)
         return normalize(x)
     end
@@ -427,43 +436,43 @@ mutable struct JumpRNGState{T<:Real,R<:AbstractRNG}
     rng::R
     threshold::T
 end
-function JumpRNGState(::Type{T}, seed) where T
+function JumpRNGState(::Type{T}, seed) where {T}
     rng = MersenneTwister(seed)
     threshold = rand(rng, T)
     JumpRNGState(rng, threshold)
 end
-roll!(s::JumpRNGState{T}) where T = (s.threshold = rand(s.rng, T))
+roll!(s::JumpRNGState{T}) where {T} = (s.threshold = rand(s.rng, T))
 threshold(s::JumpRNGState) = s.threshold
 
 function jump_callback(jumpfun::F, seed, scb, save_before!::G,
-                        save_after!::H, save_t_index::I, psi0, rng_state::JumpRNGState) where {F,G,H,I}
+    save_after!::H, save_t_index::I, psi0, rng_state::JumpRNGState) where {F,G,H,I}
 
     tmp = copy(psi0)
     psi_tmp = copy(psi0)
 
-    djumpnorm(x, t, integrator) = norm(x)^2 - (1-threshold(rng_state))
+    djumpnorm(x, t, integrator) = norm(x)^2 - (1 - threshold(rng_state))
 
     function dojump(integrator)
         x = integrator.u
         t = integrator.t
 
         affect! = scb.affect!
-        save_before!(affect!,integrator)
-        recast!(psi_tmp,x)
+        save_before!(affect!, integrator)
+        recast!(psi_tmp, x)
         i = jumpfun(rng_state.rng, t, psi_tmp, tmp)
         x .= tmp.data
-        save_after!(affect!,integrator)
-        save_t_index(t,i)
+        save_after!(affect!, integrator)
+        save_t_index(t, i)
 
         roll!(rng_state)
         return nothing
     end
 
-    return OrdinaryDiffEq.ContinuousCallback(djumpnorm,dojump,
-            save_positions = (false,false))
+    return OrdinaryDiffEq.ContinuousCallback(djumpnorm, dojump,
+        save_positions=(false, false))
 end
 jump_callback(jumpfun, seed, scb, save_before!,
-                        save_after!, save_t_index, psi0, ::Nothing) =
+    save_after!, save_t_index, psi0, ::Nothing) =
     jump_callback(jumpfun, seed, scb, save_before!,
         save_after!, save_t_index, psi0, JumpRNGState(real(eltype(psi0)), seed))
 
@@ -483,13 +492,13 @@ Default jump function.
 * `probs_tmp`: Temporary array for holding jump probailities.
 """
 function jump(rng, t, psi, J, psi_new, probs_tmp, rates::Nothing)
-    if length(J)==1
-        QuantumOpticsBase.mul!(psi_new,J[1],psi,true,false)
+    if length(J) == 1
+        QuantumOpticsBase.mul!(psi_new, J[1], psi, true, false)
         psi_new.data ./= norm(psi_new)
-        i=1
+        i = 1
     else
-        for i=1:length(J)
-            QuantumOpticsBase.mul!(psi_new,J[i],psi,true,false)
+        for i = 1:length(J)
+            QuantumOpticsBase.mul!(psi_new, J[i], psi, true, false)
             probs_tmp[i] = real(dot(psi_new.data, psi_new.data))
         end
         r = rand(rng)
@@ -501,19 +510,19 @@ function jump(rng, t, psi, J, psi_new, probs_tmp, rates::Nothing)
             cumulative_prob += p / total
             cumulative_prob > r && break
         end
-        QuantumOpticsBase.mul!(psi_new,J[i],psi,eltype(psi)(1/sqrt(probs_tmp[i])),zero(eltype(psi)))
+        QuantumOpticsBase.mul!(psi_new, J[i], psi, eltype(psi)(1 / sqrt(probs_tmp[i])), zero(eltype(psi)))
     end
     return i
 end
 
 function jump(rng, t, psi, J, psi_new, probs_tmp, rates::AbstractVector)
-    if length(J)==1
-        QuantumOpticsBase.mul!(psi_new,J[1],psi,eltype(psi)(sqrt(rates[1])),zero(eltype(psi)))
+    if length(J) == 1
+        QuantumOpticsBase.mul!(psi_new, J[1], psi, eltype(psi)(sqrt(rates[1])), zero(eltype(psi)))
         psi_new.data ./= norm(psi_new)
-        i=1
+        i = 1
     else
-        for i=1:length(J)
-            QuantumOpticsBase.mul!(psi_new,J[i],psi,eltype(psi)(sqrt(rates[i])),zero(eltype(psi)))
+        for i = 1:length(J)
+            QuantumOpticsBase.mul!(psi_new, J[i], psi, eltype(psi)(sqrt(rates[i])), zero(eltype(psi)))
             probs_tmp[i] = real(dot(psi_new.data, psi_new.data))
         end
         r = rand(rng)
@@ -525,7 +534,7 @@ function jump(rng, t, psi, J, psi_new, probs_tmp, rates::AbstractVector)
             cumulative_prob += p / total
             cumulative_prob > r && break
         end
-        QuantumOpticsBase.mul!(psi_new,J[i],psi,eltype(psi)(sqrt(rates[i]/probs_tmp[i])),zero(eltype(psi)))
+        QuantumOpticsBase.mul!(psi_new, J[i], psi, eltype(psi)(sqrt(rates[i] / probs_tmp[i])), zero(eltype(psi)))
     end
     return i
 end
@@ -540,19 +549,19 @@ the jump operators J.
 See also: [`mcwf`](@ref)
 """
 function dmcwf_h!(dpsi, H, J, Jdagger, rates::Nothing, psi, dpsi_cache)
-    QuantumOpticsBase.mul!(dpsi,H,psi,eltype(psi)(-im),zero(eltype(psi)))
-    for i=1:length(J)
-        QuantumOpticsBase.mul!(dpsi_cache,J[i],psi,true,false)
-        QuantumOpticsBase.mul!(dpsi,Jdagger[i],dpsi_cache,eltype(psi)(-0.5),one(eltype(psi)))
+    QuantumOpticsBase.mul!(dpsi, H, psi, eltype(psi)(-im), zero(eltype(psi)))
+    for i = 1:length(J)
+        QuantumOpticsBase.mul!(dpsi_cache, J[i], psi, true, false)
+        QuantumOpticsBase.mul!(dpsi, Jdagger[i], dpsi_cache, eltype(psi)(-0.5), one(eltype(psi)))
     end
     return dpsi
 end
 
 function dmcwf_h!(dpsi, H, J, Jdagger, rates::AbstractVector, psi, dpsi_cache)
-    QuantumOpticsBase.mul!(dpsi,H,psi,eltype(psi)(-im),zero(eltype(psi)))
-    for i=1:length(J)
-        QuantumOpticsBase.mul!(dpsi_cache,J[i],psi,eltype(psi)(rates[i]),zero(eltype(psi)))
-        QuantumOpticsBase.mul!(dpsi,Jdagger[i],dpsi_cache,eltype(psi)(-0.5),one(eltype(psi)))
+    QuantumOpticsBase.mul!(dpsi, H, psi, eltype(psi)(-im), zero(eltype(psi)))
+    for i = 1:length(J)
+        QuantumOpticsBase.mul!(dpsi_cache, J[i], psi, eltype(psi)(rates[i]), zero(eltype(psi)))
+        QuantumOpticsBase.mul!(dpsi, Jdagger[i], dpsi_cache, eltype(psi)(-0.5), one(eltype(psi)))
     end
     return dpsi
 end
@@ -568,13 +577,13 @@ function check_mcwf(psi0, H, J, Jdagger, rates)
     if !(isa(H, DenseOpType) || isa(H, SparseOpType))
         isreducible = false
     end
-    for j=J
+    for j = J
         @assert isa(j, AbstractOperator)
         if !(isa(j, DenseOpType) || isa(j, SparseOpType))
             isreducible = false
         end
     end
-    for j=Jdagger
+    for j = Jdagger
         @assert isa(j, AbstractOperator)
         if !(isa(j, DenseOpType) || isa(j, SparseOpType))
             isreducible = false
@@ -605,11 +614,11 @@ corresponding set of jump operators is calculated.
 function diagonaljumps(rates::AbstractMatrix, J)
     @assert length(J) == size(rates)[1] == size(rates)[2]
     d, v = eigen(rates)
-    d, [sum([v[j, i]*J[j] for j=1:length(d)]) for i=1:length(d)]
+    d, [sum([v[j, i] * J[j] for j = 1:length(d)]) for i = 1:length(d)]
 end
 
-function diagonaljumps(rates::AbstractMatrix, J::Vector{T}) where T<:Union{LazySum,LazyTensor,LazyProduct}
+function diagonaljumps(rates::AbstractMatrix, J::Vector{T}) where {T<:Union{LazySum,LazyTensor,LazyProduct}}
     @assert length(J) == size(rates)[1] == size(rates)[2]
     d, v = eigen(rates)
-    d, [LazySum([v[j, i]*J[j] for j=1:length(d)]...) for i=1:length(d)]
+    d, [LazySum([v[j, i] * J[j] for j = 1:length(d)]...) for i = 1:length(d)]
 end