Make everything type stable

Author: albertomercurio

src/qobj/superoperators.jl

@@ -180,18 +180,5 @@ liouvillian(H::AbstractQuantumObject{Operator}, Id_cache::Diagonal = I(prod(H.di
 liouvillian(H::AbstractQuantumObject{SuperOperator}, Id_cache::Diagonal) = H
 
 function _sum_lindblad_dissipators(c_ops, Id_cache::Diagonal)
-    D = 0
-    # sum all the (time-independent) c_ops first
-    c_ops_ti = filter(op -> isa(op, QuantumObject), c_ops)
-    if !isempty(c_ops_ti)
-        D += mapreduce(op -> lindblad_dissipator(op, Id_cache), +, c_ops_ti)
-    end
-
-    # sum rest of the QobjEvo together
-    c_ops_td = filter(op -> isa(op, QuantumObjectEvolution), c_ops)
-    if !isempty(c_ops_td)
-        D += mapreduce(op -> lindblad_dissipator(op, Id_cache), +, c_ops_td)
-    end
-
-    return D
+    return sum(op -> lindblad_dissipator(op, Id_cache), c_ops; init = zero(spre(first(c_ops), Id_cache)))
 end

src/time_evolution/brmesolve.jl

@@ -3,7 +3,7 @@ export bloch_redfield_tensor, brterm, brmesolve
 @doc raw"""
     bloch_redfield_tensor(
         H::QuantumObject{Operator},
-        a_ops::Union{AbstractVector, Tuple},
+        a_ops::Union{AbstractVector, Tuple, Nothing},
         c_ops::Union{AbstractVector, Tuple, Nothing}=nothing;
         sec_cutoff::Real=0.1,
         fock_basis::Union{Val,Bool}=Val(false)
@@ -28,7 +28,7 @@ The return depends on `fock_basis`.
 """
 function bloch_redfield_tensor(
     H::QuantumObject{Operator},
-    a_ops::Union{AbstractVector,Tuple},
+    a_ops::Union{AbstractVector,Tuple,Nothing},
     c_ops::Union{AbstractVector,Tuple,Nothing} = nothing;
     sec_cutoff::Real = 0.1,
     fock_basis::Union{Val,Bool} = Val(false),
@@ -44,7 +44,9 @@ function bloch_redfield_tensor(
     # Check whether we can rotate the terms to the eigenbasis directly in the Hamiltonian space
     fock_basis_hamiltonian = getVal(fock_basis) && sec_cutoff == -1
 
-    R = isempty(a_ops) ? 0 : sum(x -> _brterm(rst, x[1], x[2], sec_cutoff, fock_basis_hamiltonian), a_ops)
+    R =
+        (isnothing(a_ops) || isempty(a_ops)) ? 0 :
+        sum(x -> _brterm(rst, x[1], x[2], sec_cutoff, fock_basis_hamiltonian), a_ops)
 
     # If in fock basis, we need to transform the terms back to the fock basis
     # Note: we can transform the terms in the Hamiltonian space only if sec_cutoff is -1

src/time_evolution/mesolve.jl

@@ -339,13 +339,7 @@ function mesolve_map(
     ens_prob = TimeEvolutionProblem(
         EnsembleProblem(
             prob.prob,
-            prob_func = (prob, i, repeat) -> remake(
-                prob,
-                f = deepcopy(prob.f.f),
-                u0 = iter[i][1],
-                p = iter[i][2:end],
-                callback = haskey(prob.kwargs, :callback) ? deepcopy(prob.kwargs[:callback]) : nothing,
-            ),
+            prob_func = (prob, i, repeat) -> _se_me_map_prob_func(prob, i, repeat, iter),
             output_func = _output_func[1],
             safetycopy = false,
         ),
@@ -357,7 +351,7 @@ function mesolve_map(
 
     # handle solution and make it become an Array of TimeEvolutionSol
     sol_vec =
-        [_gen_mesolve_solution(sol[:, i], prob.times, prob.dimensions, prob.kwargs.isoperket) for i in eachindex(sol)]
+        [_gen_mesolve_solution(sol[:, i], prob.times, prob.dimensions, prob.kwargs.isoperket) for i in eachindex(sol)] # map is type unstable
     return reshape(sol_vec, size(iter))
 end
 mesolve_map(

src/time_evolution/sesolve.jl

@@ -2,6 +2,23 @@ export sesolveProblem, sesolve, sesolve_map
 
 _sesolve_make_U_QobjEvo(H) = -1im * QuantumObjectEvolution(H, type = Operator())
 
+function _gen_sesolve_solution(sol, times, dimensions)
+    ψt = map(ϕ -> QuantumObject(ϕ, type = Ket(), dims = dimensions), sol.u)
+
+    kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
+
+    return TimeEvolutionSol(
+        times,
+        sol.t,
+        ψt,
+        _get_expvals(sol, SaveFuncSESolve),
+        sol.retcode,
+        sol.alg,
+        kwargs.abstol,
+        kwargs.reltol,
+    )
+end
+
 @doc raw"""
     sesolveProblem(
         H::Union{AbstractQuantumObject{Operator},Tuple},
@@ -160,23 +177,6 @@ function sesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit
     return _gen_sesolve_solution(sol, prob.times, prob.dimensions)
 end
 
-function _gen_sesolve_solution(sol, times, dimensions)
-    ψt = map(ϕ -> QuantumObject(ϕ, type = Ket(), dims = dimensions), sol.u)
-
-    kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
-
-    return TimeEvolutionSol(
-        times,
-        sol.t,
-        ψt,
-        _get_expvals(sol, SaveFuncSESolve),
-        sol.retcode,
-        sol.alg,
-        kwargs.abstol,
-        kwargs.reltol,
-    )
-end
-
 @doc raw"""
     sesolve_map(
         H::Union{AbstractQuantumObject{Operator},Tuple},
@@ -254,13 +254,7 @@ function sesolve_map(
     ens_prob = TimeEvolutionProblem(
         EnsembleProblem(
             prob.prob,
-            prob_func = (prob, i, repeat) -> remake(
-                prob,
-                f = deepcopy(prob.f.f),
-                u0 = iter[i][1],
-                p = iter[i][2:end],
-                callback = haskey(prob.kwargs, :callback) ? deepcopy(prob.kwargs[:callback]) : nothing,
-            ),
+            prob_func = (prob, i, repeat) -> _se_me_map_prob_func(prob, i, repeat, iter),
             output_func = _output_func[1],
             safetycopy = false,
         ),
@@ -271,7 +265,7 @@ function sesolve_map(
     sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj)
 
     # handle solution and make it become an Array of TimeEvolutionSol
-    sol_vec = [_gen_sesolve_solution(sol[:, i], prob.times, prob.dimensions) for i in eachindex(sol)]
+    sol_vec = [_gen_sesolve_solution(sol[:, i], prob.times, prob.dimensions) for i in eachindex(sol)] # map is type unstable
     return reshape(sol_vec, size(iter))
 end
 sesolve_map(H::Union{AbstractQuantumObject{Operator},Tuple}, ψ0::QuantumObject{Ket}, tlist::AbstractVector; kwargs...) =

src/time_evolution/time_evolution.jl

@@ -436,6 +436,18 @@ function _ensemble_dispatch_solve(
     return sol
 end
 
+# For mapped solvers
+function _se_me_map_prob_func(prob, i, repeat, iter)
+    f = deepcopy(prob.f.f)
+    u0 = iter[i][1]
+    p = iter[i][2:end]
+    if haskey(prob.kwargs, :callback)
+        return remake(prob, f = f, u0 = u0, p = p, callback = deepcopy(prob.kwargs[:callback]))
+    else
+        return remake(prob, f = f, u0 = u0, p = p)
+    end
+end
+
 #######################################
 #=
  Stochastic funcs

test/core-test/brmesolve.jl

@@ -98,7 +98,8 @@ end
     _z_c_op = √γ * sigmaz()
     _x_a_op = (sigmax(), spectra)
 
-    arg_sets = [[[_m_c_op], [], [_x_a_op]], [[_m_c_op], [_m_c_op], []], [[_m_c_op, _z_c_op], [_z_c_op], [_x_a_op]]]
+    arg_sets =
+        (([_m_c_op], nothing, [_x_a_op]), ([_m_c_op], [_m_c_op], nothing), ([_m_c_op, _z_c_op], [_z_c_op], [_x_a_op]))
 
     δ = 0
     ϵ = 0.5 * 2π

test/core-test/time_evolution.jl

@@ -154,7 +154,6 @@ end
     a = TESetup.a
     σz = TESetup.σz
     σm = TESetup.σm
-    ψ0 = TESetup.ψ0
     e_ops = TESetup.e_ops
 
     g = 0.01