First working case of sesolve

Author: albertomercurio

src/qobj/boolean_functions.jl

@@ -11,6 +11,7 @@ export isunitary
 Checks if the [`QuantumObject`](@ref) `A` is a [`BraQuantumObject`](@ref). Default case returns `false` for any other inputs.
 """
 isbra(A::QuantumObject) = A.type isa BraQuantumObject
+isbra(::Type{QuantumObject{DT,BraQuantumObject,N}}) where {DT,N} = true
 isbra(A) = false # default case
 
 @doc raw"""
@@ -19,6 +20,7 @@ isbra(A) = false # default case
 Checks if the [`QuantumObject`](@ref) `A` is a [`KetQuantumObject`](@ref). Default case returns `false` for any other inputs.
 """
 isket(A::QuantumObject) = A.type isa KetQuantumObject
+isket(::Type{QuantumObject{DT,KetQuantumObject,N}}) where {DT,N} = true
 isket(A) = false # default case
 
 @doc raw"""
@@ -27,6 +29,7 @@ isket(A) = false # default case
 Checks if the [`AbstractQuantumObject`](@ref) `A` is a [`OperatorQuantumObject`](@ref). Default case returns `false` for any other inputs.
 """
 isoper(A::AbstractQuantumObject) = A.type isa OperatorQuantumObject
+isoper(::Type{QuantumObject{DT,OperatorQuantumObject,N}}) where {DT,N} = true
 isoper(A) = false # default case
 
 @doc raw"""
@@ -35,6 +38,7 @@ isoper(A) = false # default case
 Checks if the [`QuantumObject`](@ref) `A` is a [`OperatorBraQuantumObject`](@ref). Default case returns `false` for any other inputs.
 """
 isoperbra(A::QuantumObject) = A.type isa OperatorBraQuantumObject
+isoperbra(::Type{QuantumObject{DT,OperatorBraQuantumObject,N}}) where {DT,N} = true
 isoperbra(A) = false # default case
 
 @doc raw"""
@@ -43,6 +47,7 @@ isoperbra(A) = false # default case
 Checks if the [`QuantumObject`](@ref) `A` is a [`OperatorKetQuantumObject`](@ref). Default case returns `false` for any other inputs.
 """
 isoperket(A::QuantumObject) = A.type isa OperatorKetQuantumObject
+isoperket(::Type{QuantumObject{DT,OperatorKetQuantumObject,N}}) where {DT,N} = true
 isoperket(A) = false # default case
 
 @doc raw"""
@@ -51,6 +56,7 @@ isoperket(A) = false # default case
 Checks if the [`AbstractQuantumObject`](@ref) `A` is a [`SuperOperatorQuantumObject`](@ref). Default case returns `false` for any other inputs.
 """
 issuper(A::AbstractQuantumObject) = A.type isa SuperOperatorQuantumObject
+issuper(::Type{QuantumObject{DT,SuperOperatorQuantumObject,N}}) where {DT,N} = true
 issuper(A) = false # default case
 
 @doc raw"""

src/qobj/eigsolve.jl

@@ -378,7 +378,7 @@ function eigsolve_al(
     L = liouvillian(H, c_ops)
     prob = mesolveProblem(
         L,
-        QuantumObject(ρ0, type=Operator, dims = H.dims),
+        QuantumObject(ρ0, type = Operator, dims = H.dims),
         [zero(T), T];
         alg = alg,
         H_t = H_t,

src/qobj/quantum_object_evo.jl

@@ -90,18 +90,14 @@ struct QuantumObjectEvolution{DT<:AbstractSciMLOperator,ObjType<:QuantumObjectTy
     dims::SVector{N,Int}
 end
 
-function QuantumObjectEvolution(op_func_list::Tuple)
+# Make the QuantumObjectEvolution, with the option to pre-multiply by a scalar
+function QuantumObjectEvolution(op_func_list::Tuple, α = true)
     op = _get_op_func_first(op_func_list)
     dims = op.dims
     type = op.type
     T = eltype(op)
-    data = sum(op_func_list) do op_func
-        if op_func isa Tuple
-            return ScalarOperator(zero(T), update_func = (a, u, p, t) -> op_func[2](p, t)) * MatrixOperator(op_func[1].data)
-        else
-            return MatrixOperator(op_func.data)
-        end
-    end
+
+    data = _generate_data(T, op_func_list, α)
 
     # Preallocate the SciMLOperator cache using a dense vector as a reference
     v0 = sparse_to_dense(similar(op.data, size(op, 1)))
@@ -110,31 +106,75 @@ function QuantumObjectEvolution(op_func_list::Tuple)
     return QuantumObjectEvolution(data, type, dims)
 end
 
-QuantumObjectEvolution(op::QuantumObject) = QuantumObjectEvolution(MatrixOperator(op.data), op.type, op.dims)
-
-function _get_op_func_first(op_func_list)
-    dims_type_op = map(op_func_list) do op_func
-        if op_func isa Tuple
-            length(op_func) == 2 || throw(ArgumentError("The tuple must have two elements."))
-            ((isoper(op_func[1]) || issuper(op_func[1])) && op_func[2] isa Function) || throw(
+QuantumObjectEvolution(op::QuantumObject, α = true) =
+    QuantumObjectEvolution(MatrixOperator(α * op.data), op.type, op.dims)
+
+@generated function _get_op_func_first(op_func_list::Tuple)
+    op_func_list_types = op_func_list.parameters
+    N = length(op_func_list_types)
+    T = ()
+    dims_expr = ()
+    first_op = nothing
+    for i in 1:N
+        op_func_type = op_func_list_types[i]
+        if op_func_type <: Tuple
+            length(op_func_type.parameters) == 2 || throw(ArgumentError("The tuple must have two elements."))
+            op_type = op_func_type.parameters[1]
+            func_type = op_func_type.parameters[2]
+            ((isoper(op_type) || issuper(op_type)) && func_type <: Function) || throw(
                 ArgumentError(
                     "The first element must be a Operator or SuperOperator, and the second element must be a function.",
                 ),
             )
-            return (op_func[1].dims, op_func[1].type, op_func[1])
+            data_type = op_type.parameters[1]
+            T = (T..., eltype(data_type))
+            dims_expr = (dims_expr..., :(op_func_list[$i][1].dims))
+            if i == 1
+                first_op = :(op_func_list[$i][1])
+            end
         else
-            (isoper(op_func) || issuper(op_func)) || throw(ArgumentError("The element must be a QuantumObject."))
-            return (op_func.dims, op_func.type, op_func)
+            op_type = op_func_type
+            (isoper(op_type) || issuper(op_type)) ||
+                throw(ArgumentError("The element must be a Operator or SuperOperator."))
+            data_type = op_type.parameters[1]
+            T = (T..., eltype(data_type))
+            dims_expr = (dims_expr..., :(op_func_list[$i].dims))
+            if i == 1
+                first_op = :(op_func_list[$i])
+            end
         end
     end
 
-    length(unique(getindex.(dims_type_op, 1))) == 1 ||
-        throw(ArgumentError("The dimensions of the operators must be the same."))
+    length(unique(T)) == 1 || throw(ArgumentError("The types of the operators must be the same."))
+
+    quote
+        dims = tuple($(dims_expr...))
 
-    length(unique(getindex.(dims_type_op, 2))) == 1 ||
-        throw(ArgumentError("The types of the operators must be the same."))
+        length(unique(dims)) == 1 || throw(ArgumentError("The dimensions of the operators must be the same."))
 
-    return dims_type_op[1][3]
+        return $first_op
+    end
+end
+
+@generated function _generate_data(T, op_func_list::Tuple, α)
+    op_func_list_types = op_func_list.parameters
+    N = length(op_func_list_types)
+    data_expr = :(0)
+    for i in 1:N
+        op_func_type = op_func_list_types[i]
+        if op_func_type <: Tuple
+            data_expr = :(
+                $data_expr +
+                ScalarOperator(zero(T), op_func_list[$i][2]) * MatrixOperator(α * op_func_list[$i][1].data)
+            )
+        else
+            data_expr = :($data_expr + MatrixOperator(α * op_func_list[$i].data))
+        end
+    end
+
+    quote
+        return $data_expr
+    end
 end
 
 function (QO::QuantumObjectEvolution)(p, t)

src/qobj/superoperators.jl

@@ -89,16 +89,13 @@ the same function is applied multiple times with a known Hilbert space dimension
 
 See also [`spre`](@ref) and [`spost`](@ref).
 """
-function lindblad_dissipator(
-    O::QuantumObject{DT,OperatorQuantumObject},
-    Id_cache = I(size(O, 1)),
-) where {DT}
+function lindblad_dissipator(O::QuantumObject{DT,OperatorQuantumObject}, Id_cache = I(size(O, 1))) where {DT}
     Od_O = O' * O
     return sprepost(O, O') - spre(Od_O, Id_cache) / 2 - spost(Od_O, Id_cache) / 2
 end
 
 # It is already a SuperOperator
-lindblad_dissipator(O::QuantumObject{DT,SuperOperatorQuantumObject}, Id_cache=nothing) where {DT} = O
+lindblad_dissipator(O::QuantumObject{DT,SuperOperatorQuantumObject}, Id_cache = nothing) where {DT} = O
 
 @doc raw"""
     liouvillian(H::QuantumObject, c_ops::Union{Nothing,AbstractVector,Tuple}=nothing, Id_cache=I(prod(H.dims)))

src/qobj/synonyms.jl

@@ -18,13 +18,13 @@ Note that this functions is same as `QuantumObject(A; type=type, dims=dims)`.
 Qobj(A; kwargs...) = QuantumObject(A; kwargs...)
 
 @doc raw"""
-    QobjEvo(op_func_list::Union{Tuple,QuantumObject})
+    QobjEvo(op_func_list::Union{Tuple,QuantumObject}, α::Real=true)
 
 Generate [`QuantumObjectEvolution`](@ref)
 
-Note that this functions is same as `QuantumObjectEvolution(op_func_list)`.
+Note that this functions is same as `QuantumObjectEvolution(op_func_list)`. If `α` is provided, all the operators in `op_func_list` will be pre-multiplied by `α`.
 """
-QobjEvo(op_func_list::Union{Tuple,QuantumObject}) = QuantumObjectEvolution(op_func_list)
+QobjEvo(op_func_list::Union{Tuple,QuantumObject}, α = true) = QuantumObjectEvolution(op_func_list, α)
 
 @doc raw"""
     shape(A::QuantumObject)

src/time_evolution/sesolve.jl

@@ -4,7 +4,7 @@ function _save_func_sesolve(integrator)
     internal_params = integrator.p
     progr = internal_params.progr
 
-    if !internal_params.is_empty_e_ops
+    if internal_params.is_empty_e_ops
         e_ops = internal_params.e_ops
         expvals = internal_params.expvals
 
@@ -44,14 +44,13 @@ function _generate_sesolve_kwargs(e_ops, progress_bar::Val{false}, t_l, kwargs)
 end
 
 @doc raw"""
-    sesolveProblem(H::QuantumObject,
-        ψ0::QuantumObject,
-        tlist::AbstractVector;
-        alg::OrdinaryDiffEqAlgorithm=Tsit5()
-        e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-        H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
-        params::NamedTuple=NamedTuple(),
-        progress_bar::Union{Val,Bool}=Val(true),
+    sesolveProblem(H,
+        ψ0,
+        tlist;
+        alg=Tsit5()
+        e_ops = nothing,
+        params=NamedTuple(),
+        progress_bar=Val(true),
         kwargs...)
 
 Generates the ODEProblem for the Schrödinger time evolution of a quantum system:
@@ -62,7 +61,7 @@ Generates the ODEProblem for the Schrödinger time evolution of a quantum system
 
 # Arguments
 
-- `H::QuantumObject`: The Hamiltonian of the system ``\hat{H}``.
+- `H::Union{QuantumObject,Tuple}`: The Hamiltonian of the system ``\hat{H}``.
 - `ψ0::QuantumObject`: The initial state of the system ``|\psi(0)\rangle``.
 - `tlist::AbstractVector`: The time list of the evolution.
 - `alg::OrdinaryDiffEqAlgorithm`: The algorithm used for the time evolution.
@@ -85,47 +84,43 @@ Generates the ODEProblem for the Schrödinger time evolution of a quantum system
 - `prob`: The `ODEProblem` for the Schrödinger time evolution of the system.
 """
 function sesolveProblem(
-    H::Union{QuantumObject{MT1,OperatorQuantumObject},Tuple},
-    ψ0::QuantumObject{<:AbstractVector{T2},KetQuantumObject},
+    H::Union{QuantumObject{DT1,OperatorQuantumObject},Tuple},
+    ψ0::QuantumObject{DT2,KetQuantumObject},
     tlist::AbstractVector;
-    alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
     progress_bar::Union{Val,Bool} = Val(true),
     kwargs...,
-) where {MT1<:AbstractMatrix,T2}
-    check_dims(H, ψ0)
-
+) where {DT1,DT2}
     haskey(kwargs, :save_idxs) &&
         throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
 
-    is_time_dependent = !(H_t isa Nothing)
-
     ϕ0 = sparse_to_dense(_CType(ψ0), get_data(ψ0)) # Convert it to dense vector with complex element type
 
     t_l = convert(Vector{_FType(ψ0)}, tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl
 
-    U = -1im * get_data(H)
+    H_evo = QobjEvo(H, -1im) # pre-multiply by -i
+    isoper(H_evo) || throw(ArgumentError("The Hamiltonian must be an Operator."))
+    check_dims(H_evo, ψ0)
+    U = H_evo.data
+
     progr = ProgressBar(length(t_l), enable = getVal(progress_bar))
 
     if e_ops isa Nothing
         expvals = Array{ComplexF64}(undef, 0, length(t_l))
-        e_ops2 = MT1[]
+        e_ops_data = ()
         is_empty_e_ops = true
     else
         expvals = Array{ComplexF64}(undef, length(e_ops), length(t_l))
-        e_ops2 = get_data.(e_ops)
+        e_ops_data = get_data.(e_ops)
         is_empty_e_ops = isempty(e_ops)
     end
 
     p = (
-        U = U,
-        e_ops = e_ops2,
+        e_ops = e_ops_data,
         expvals = expvals,
         progr = progr,
-        Hdims = H.dims,
-        H_t = H_t,
+        Hdims = H_evo.dims,
         times = t_l,
         is_empty_e_ops = is_empty_e_ops,
         params...,
@@ -136,10 +131,8 @@ function sesolveProblem(
     kwargs2 = merge(default_values, kwargs)
     kwargs3 = _generate_sesolve_kwargs(e_ops, makeVal(progress_bar), t_l, kwargs2)
 
-    dudt! = is_time_dependent ? sesolve_td_dudt! : sesolve_ti_dudt!
-
     tspan = (t_l[1], t_l[end])
-    return ODEProblem{true,FullSpecialize}(dudt!, ϕ0, tspan, p; kwargs3...)
+    return ODEProblem{true}(U, ϕ0, tspan, p; kwargs3...)
 end
 
 @doc raw"""
@@ -184,27 +177,16 @@ Time evolution of a closed quantum system using the Schrödinger equation:
 - `sol::TimeEvolutionSol`: The solution of the time evolution. See also [`TimeEvolutionSol`](@ref)
 """
 function sesolve(
-    H::QuantumObject{MT1,OperatorQuantumObject},
-    ψ0::QuantumObject{<:AbstractVector{T2},KetQuantumObject},
+    H::Union{QuantumObject{DT1,OperatorQuantumObject},Tuple},
+    ψ0::QuantumObject{DT2,KetQuantumObject},
     tlist::AbstractVector;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
     progress_bar::Union{Val,Bool} = Val(true),
     kwargs...,
-) where {MT1<:AbstractMatrix,T2}
-    prob = sesolveProblem(
-        H,
-        ψ0,
-        tlist;
-        alg = alg,
-        e_ops = e_ops,
-        H_t = H_t,
-        params = params,
-        progress_bar = progress_bar,
-        kwargs...,
-    )
+) where {DT1,DT2}
+    prob = sesolveProblem(H, ψ0, tlist; e_ops = e_ops, params = params, progress_bar = progress_bar, kwargs...)
 
     return sesolve(prob, alg)
 end