Add option to renormalize state in schroedinger

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
 

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