Fix stochastic checks

src/stochastic_master.jl

@@ -6,7 +6,7 @@ using ...bases, ...states, ...operators
 using ...operators_dense, ...operators_sparse
 using ...timeevolution
 using LinearAlgebra
-import ...timeevolution: integrate_stoch, recast!
+import ...timeevolution: integrate_stoch, recast!, QO_CHECKS
 import ...timeevolution.timeevolution_master: dmaster_h, dmaster_nh, dmaster_h_dynamic, check_master
 
 const DecayRates = Union{Vector{Float64}, Matrix{Float64}, Nothing}
@@ -55,7 +55,7 @@ function master(tspan, rho0::DenseOperator, H::Operator,
     dmaster_stoch(dx::DiffArray, t::Float64, rho::DenseOperator, drho::DenseOperator, n::Int) =
             dmaster_stochastic(dx, rho, C, Cdagger, drho, n)
 
-    isreducible = check_master(rho0, H, J, Jdagger, rates)
+    isreducible = check_master(rho0, H, J, Jdagger, rates) && check_master_stoch(rho0, C, Cdagger)
     if !isreducible
         dmaster_h_determ(t::Float64, rho::DenseOperator, drho::DenseOperator) =
             dmaster_h(rho, H, rates, J, Jdagger, drho, tmp)
@@ -158,8 +158,9 @@ end
 function dmaster_stoch_dynamic(dx::DiffArray, t::Float64, rho::DenseOperator,
             f::Function, drho::DenseOperator, n::Int)
     result = f(t, rho)
-    @assert 2 == length(result)
+    QO_CHECKS[] && @assert 2 == length(result)
     C, Cdagger = result
+    QO_CHECKS[] && check_master_stoch(rho, C, Cdagger)
     dmaster_stochastic(dx, rho, C, Cdagger, drho, n)
 end
 
@@ -174,6 +175,27 @@ function integrate_master_stoch(tspan, df::Function, dg::Function,
     integrate_stoch(tspan_, df, dg, x0, state, dstate, fout, n; kwargs...)
 end
 
+function check_master_stoch(rho0::DenseOperator, C::Vector, Cdagger::Vector)
+    @assert length(C) == length(Cdagger)
+    isreducible = true
+    for c=C
+        @assert isa(c, Operator)
+        if !(isa(c, DenseOperator) || isa(c, SparseOperator))
+            isreducible = false
+        end
+        check_samebases(rho0, c)
+    end
+    for c=Cdagger
+        @assert isa(c, Operator)
+        if !(isa(c, DenseOperator) || isa(c, SparseOperator))
+            isreducible = false
+        end
+        check_samebases(rho0, c)
+    end
+    isreducible
+end
+
+
 # TODO: Speed up by recasting to n-d arrays, remove vector methods
 function recast!(x::Union{Vector{ComplexF64}, SubArray{ComplexF64, 1}}, rho::DenseOperator)
     rho.data = reshape(x, size(rho.data))

src/stochastic_schroedinger.jl

@@ -45,6 +45,10 @@ function schroedinger(tspan, psi0::Ket, H::Operator, Hs::Vector;
     state = copy(psi0)
 
     check_schroedinger(psi0, H)
+    for h=Hs
+        check_schroedinger(psi0, h)
+    end
+
     dschroedinger_determ(t::Float64, psi::Ket, dpsi::Ket) = dschroedinger(psi, H, dpsi)
     dschroedinger_stoch(dx::DiffArray,
             t::Float64, psi::Ket, dpsi::Ket, n::Int) = dschroedinger_stochastic(dx, psi, Hs, dpsi, n)
@@ -130,14 +134,12 @@ end
 
 function dschroedinger_stochastic(dx::Vector{ComplexF64}, psi::Ket, Hs::Vector{T},
             dpsi::Ket, index::Int) where T <: Operator
-    QO_CHECKS[] && check_schroedinger(psi, Hs[index])
     recast!(dx, dpsi)
     dschroedinger(psi, Hs[index], dpsi)
 end
 function dschroedinger_stochastic(dx::Array{ComplexF64, 2}, psi::Ket, Hs::Vector{T},
             dpsi::Ket, n::Int) where T <: Operator
     for i=1:n
-        QO_CHECKS[] && check_schroedinger(psi, Hs[i])
         dx_i = @view dx[:, i]
         recast!(dx_i, dpsi)
         dschroedinger(psi, Hs[i], dpsi)
@@ -147,6 +149,11 @@ end
 function dschroedinger_stochastic(dx::DiffArray,
             t::Float64, psi::Ket, f::Function, dpsi::Ket, n::Int)
     ops = f(t, psi)
+    if QO_CHECKS[]
+        @inbounds for h=ops
+            check_schroedinger(psi, h)
+        end
+    end
     dschroedinger_stochastic(dx, psi, ops, dpsi, n)
 end
 

src/stochastic_semiclassical.jl

@@ -7,8 +7,9 @@ using ...operators_dense, ...operators_sparse
 using ...semiclassical
 import ...semiclassical: recast!, State, dmaster_h_dynamic
 using ...timeevolution
-import ...timeevolution: integrate_stoch
-import ...timeevolution.timeevolution_schroedinger: dschroedinger, dschroedinger_dynamic
+import ...timeevolution: integrate_stoch, QO_CHECKS
+import ...timeevolution.timeevolution_schroedinger: dschroedinger, dschroedinger_dynamic, check_schroedinger
+import ...stochastic.stochastic_master: check_master_stoch
 using ...stochastic
 using LinearAlgebra
 
@@ -207,6 +208,7 @@ function dschroedinger_stochastic(dx::Vector{ComplexF64}, t::Float64,
         dstate::State{Ket}, ::Int)
     H = fstoch_quantum(t, state.quantum, state.classical)
     recast!(dx, dstate)
+    QO_CHECKS[] && check_schroedinger(state.quantum, H[1])
     dschroedinger(state.quantum, H[1], dstate.quantum)
     recast!(dstate, dx)
 end
@@ -217,6 +219,7 @@ function dschroedinger_stochastic(dx::Array{ComplexF64, 2},
     for i=1:n
         dx_i = @view dx[:, i]
         recast!(dx_i, dstate)
+        QO_CHECKS[] && check_schroedinger(state.quantum, H[i])
         dschroedinger(state.quantum, H[i], dstate.quantum)
         recast!(dstate, dx_i)
     end
@@ -239,8 +242,9 @@ function dmaster_stoch_dynamic(dx::Vector{ComplexF64}, t::Float64,
             state::State{DenseOperator}, fstoch_quantum::Function,
             fstoch_classical::Nothing, dstate::State{DenseOperator}, ::Int)
     result = fstoch_quantum(t, state.quantum, state.classical)
-    @assert length(result) == 2
+    QO_CHECKS[] && @assert length(result) == 2
     C, Cdagger = result
+    QO_CHECKS[] && check_master_stoch(state.quantum, C, Cdagger)
     recast!(dx, dstate)
     operators.gemm!(1, C[1], state.quantum, 0, dstate.quantum)
     operators.gemm!(1, state.quantum, Cdagger[1], 1, dstate.quantum)
@@ -251,8 +255,9 @@ function dmaster_stoch_dynamic(dx::Array{ComplexF64, 2}, t::Float64,
             state::State{DenseOperator}, fstoch_quantum::Function,
             fstoch_classical::Nothing, dstate::State{DenseOperator}, n::Int)
     result = fstoch_quantum(t, state.quantum, state.classical)
-    @assert length(result) == 2
+    QO_CHECKS[] && @assert length(result) == 2
     C, Cdagger = result
+    QO_CHECKS[] && check_master_stoch(state.quantum, C, Cdagger)
     for i=1:n
         dx_i = @view dx[:, i]
         recast!(dx_i, dstate)