Better recasting for master

src/master.jl

@@ -181,15 +181,15 @@ master_nh_dynamic(tspan, psi0::Ket, f::Function; kwargs...) = master_nh_dynamic(
 
 
 # Recasting needed for the ODE solver is just providing the underlying data
-function recast!(x::Vector{Complex128}, rho::DenseOperator)
-    rho.data = reshape(x, size(rho.data))
+function recast!(x::Array{Complex128, 2}, rho::DenseOperator)
+    rho.data = x
 end
-recast!(rho::DenseOperator, x::Vector{Complex128}) = nothing
+recast!(rho::DenseOperator, x::Array{Complex128, 2}) = nothing
 
 function integrate_master(tspan, df::Function, rho0::DenseOperator,
                         fout::Union{Void, Function}; kwargs...)
     tspan_ = convert(Vector{Float64}, tspan)
-    x0 = reshape(rho0.data, length(rho0))
+    x0 = rho0.data
     state = DenseOperator(rho0.basis_l, rho0.basis_r, rho0.data)
     dstate = DenseOperator(rho0.basis_l, rho0.basis_r, rho0.data)
     integrate(tspan_, df, x0, state, dstate, fout; kwargs...)

src/stochastic_master.jl

@@ -500,9 +500,11 @@ function integrate_master_stoch(tspan, df::Function, dg::Function,
     integrate_stoch(tspan_, df, dg, x0, state, dstate, fout, n; kwargs...)
 end
 
-function recast!(x::SubArray{Complex128, 1}, rho::DenseOperator)
+# TODO: Speed up by recasting to n-d arrays, remove vector methods
+function recast!(x::Union{Vector{Complex128}, SubArray{Complex128, 1}}, rho::DenseOperator)
     rho.data = reshape(x, size(rho.data))
 end
 recast!(state::DenseOperator, x::SubArray{Complex128, 1}) = (x[:] = state.data)
+recast!(state::DenseOperator, x::Vector{Complex128}) = nothing
 
 end # module

src/timeevolution_base.jl

@@ -2,6 +2,8 @@ using ..metrics
 
 import OrdinaryDiffEq, DiffEqCallbacks, StochasticDiffEq
 
+const DiffArray = Union{Vector{Complex128}, Array{Complex128, 2}}
+
 function recast! end
 
 """
@@ -10,19 +12,19 @@ function recast! end
 
 Integrate using OrdinaryDiffEq
 """
-function integrate(tspan::Vector{Float64}, df::Function, x0::Vector{Complex128},
+function integrate(tspan::Vector{Float64}, df::Function, x0::DiffArray,
             state::T, dstate::T, fout::Function;
             alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm = OrdinaryDiffEq.DP5(),
             steady_state = false, tol = 1e-3, save_everystep = false,
             callback = nothing, kwargs...) where T
 
-    function df_(dx::Vector{Complex128}, x::Vector{Complex128}, p, t)
+    function df_(dx::DiffArray, x::DiffArray, p, t)
         recast!(x, state)
         recast!(dx, dstate)
         df(t, state, dstate)
         recast!(dstate, dx)
     end
-    function fout_(x::Vector{Complex128}, t::Float64, integrator)
+    function fout_(x::DiffArray, t::Float64, integrator)
         recast!(x, state)
         fout(t, state)
     end
@@ -64,7 +66,7 @@ function integrate(tspan::Vector{Float64}, df::Function, x0::Vector{Complex128},
     out.t,out.saveval
 end
 
-function integrate(tspan::Vector{Float64}, df::Function, x0::Vector{Complex128},
+function integrate(tspan::Vector{Float64}, df::Function, x0::DiffArray,
             state::T, dstate::T, ::Void; kwargs...) where T
     function fout(t::Float64, state::T)
         copy(state)