diff --git a/src/stochastic_schroedinger.jl b/src/stochastic_schroedinger.jl index a577884e..8eefd172 100644 --- a/src/stochastic_schroedinger.jl +++ b/src/stochastic_schroedinger.jl @@ -8,6 +8,8 @@ using ...timeevolution import ...timeevolution: integrate_stoch, recast! import ...timeevolution.timeevolution_schroedinger: dschroedinger, dschroedinger_dynamic, check_schroedinger +import DiffEqCallbacks + const DiffArray = Union{Vector{Complex128}, Array{Complex128, 2}} """ @@ -25,10 +27,14 @@ Integrate stochastic Schrödinger equation. * `fout=nothing`: If given, this function `fout(t, state)` is called every time an output should be displayed. ATTENTION: The given state is neither normalized nor permanent! +* `normalize=false`: Specify whether or not to normalize the state after each + time step taken by the solver. * `kwargs...`: Further arguments are passed on to the ode solver. """ function schroedinger(tspan, psi0::Ket, H::Operator, Hs::Vector; fout::Union{Function,Void}=nothing, + normalize_state::Bool=false, + calback=nothing, kwargs...) tspan_ = convert(Vector{Float64}, tspan) @@ -42,7 +48,17 @@ function schroedinger(tspan, psi0::Ket, H::Operator, Hs::Vector; dschroedinger_stoch(dx::DiffArray, t::Float64, psi::Ket, dpsi::Ket, n::Int) = dschroedinger_stochastic(dx, psi, Hs, dpsi, n) + if normalize_state + norm_func(u::Vector{Complex128}, t::Float64, integrator) = normalize!(u) + ncb = DiffEqCallbacks.FunctionCallingCallback(norm_func; + func_everystep=true, + func_start=false) + else + ncb = nothing + end + integrate_stoch(tspan_, dschroedinger_determ, dschroedinger_stoch, x0, state, dstate, fout, n; + ncb=ncb, kwargs...) end schroedinger(tspan, psi0::Ket, H::Operator, Hs::Operator; kwargs...) = schroedinger(tspan, psi0, H, [Hs]; kwargs...) @@ -69,10 +85,13 @@ Integrate stochastic Schrödinger equation with dynamic Hamiltonian. from the function output. NOTE: Set this number if you want to avoid an initial calculation of the function output! +* `normalize=false`: Specify whether or not to normalize the state after each + time step taken by the solver. * `kwargs...`: Further arguments are passed on to the ode solver. """ function schroedinger_dynamic(tspan, psi0::Ket, fdeterm::Function, fstoch::Function; fout::Union{Function,Void}=nothing, noise_processes::Int=0, + normalize_state::Bool=false, kwargs...) tspan_ = convert(Vector{Float64}, tspan) @@ -92,7 +111,18 @@ function schroedinger_dynamic(tspan, psi0::Ket, fdeterm::Function, fstoch::Funct t::Float64, psi::Ket, dpsi::Ket, n::Int) = dschroedinger_stochastic(dx, t, psi, fstoch, dpsi, n) - integrate_stoch(tspan, dschroedinger_determ, dschroedinger_stoch, x0, state, dstate, fout, n; + if normalize_state + norm_func(u::Vector{Complex128}, t::Float64, integrator) = normalize!(u) + ncb = DiffEqCallbacks.FunctionCallingCallback(norm_func; + func_everystep=true, + func_start=false) + else + ncb = nothing + end + + integrate_stoch(tspan, dschroedinger_determ, dschroedinger_stoch, x0, state, + dstate, fout, n; + ncb=ncb, kwargs...) end diff --git a/src/stochastic_semiclassical.jl b/src/stochastic_semiclassical.jl index e6b85a8b..c8630259 100644 --- a/src/stochastic_semiclassical.jl +++ b/src/stochastic_semiclassical.jl @@ -12,6 +12,8 @@ import ...timeevolution.timeevolution_schroedinger: dschroedinger, dschroedinger using ...stochastic import ...stochastic.stochastic_master: dneumann, dwiseman, dwiseman_nl, dlindblad +import DiffEqCallbacks + const DecayRates = Union{Vector{Float64}, Matrix{Float64}, Void} const DiffArray = Union{Vector{Complex128}, Array{Complex128, 2}} @@ -49,6 +51,8 @@ Integrate time-dependent Schrödinger equation coupled to a classical system. stochastic function `fstoch_classical` only. Must be set for non-diagonal classical noise or combinations of quantum and classical noise. See the documentation for details. +* `normalize=false`: Specify whether or not to normalize (quantum part of) the + state after each time step taken by the solver. * `kwargs...`: Further arguments are passed on to the ode solver. """ function schroedinger_semiclassical(tspan, state0::State{Ket}, fquantum::Function, @@ -57,6 +61,7 @@ function schroedinger_semiclassical(tspan, state0::State{Ket}, fquantum::Functio fout::Union{Function,Void}=nothing, noise_processes::Int=0, noise_prototype_classical=nothing, + normalize_state::Bool=false, kwargs...) tspan_ = convert(Vector{Float64}, tspan) dschroedinger_det(t::Float64, state::State{Ket}, dstate::State{Ket}) = semiclassical.dschroedinger_dynamic(t, state, fquantum, fclassical, dstate) @@ -86,11 +91,24 @@ function schroedinger_semiclassical(tspan, state0::State{Ket}, fquantum::Functio end end + if n > 0 && normalize_state + function norm_func(u::Vector{Complex128}, t::Float64, integrator) + u_quantum = @view u[1:n] + normalize!(u_quantum) + end + ncb = DiffEqCallbacks.FunctionCallingCallback(norm_func; + func_everystep=true, + func_start=false) + else + ncb = nothing + end + dschroedinger_stoch(dx::DiffArray, t::Float64, state::State{Ket}, dstate::State{Ket}, n::Int) = dschroedinger_stochastic(dx, t, state, fstoch_quantum, fstoch_classical, dstate, n) integrate_stoch(tspan_, dschroedinger_det, dschroedinger_stoch, x0, state, dstate, fout, n; noise_prototype_classical = noise_prototype_classical, + ncb=ncb, kwargs...) end