Support time-dependent lindblad_dissipator (#280)

Author: ytdHuang

docs/src/api.md

@@ -28,9 +28,10 @@ SuperOperatorQuantumObject
 SuperOperator
 QuantumObject
 QuantumObjectEvolution
-size
-eltype
-length
+Base.size
+Base.eltype
+Base.length
+SciMLOperators.cache_operator
 ```
 
 ## [Qobj boolean functions](@id doc-API:Qobj-boolean-functions)
@@ -46,7 +47,8 @@ LinearAlgebra.ishermitian
 LinearAlgebra.issymmetric
 LinearAlgebra.isposdef
 isunitary
-isconstant
+SciMLOperators.iscached
+SciMLOperators.isconstant
 ```
 
 ## [Qobj arithmetic and attributes](@id doc-API:Qobj-arithmetic-and-attributes)

src/QuantumToolbox.jl

@@ -3,17 +3,19 @@ module QuantumToolbox
 # Re-export:
 #   1. StaticArraysCore.SVector for the type of dims
 #   2. basic functions in LinearAlgebra and SparseArrays
+#   3. some functions in SciMLOperators
 import Reexport: @reexport
 @reexport import StaticArraysCore: SVector
 @reexport using LinearAlgebra
 @reexport using SparseArrays
+@reexport import SciMLOperators: cache_operator, iscached, isconstant
 
 # other functions in LinearAlgebra
 import LinearAlgebra: BlasReal, BlasInt, BlasFloat, BlasComplex, checksquare
 import LinearAlgebra.BLAS: @blasfunc
 import LinearAlgebra.LAPACK: hseqr!
 
-# SciML packages (for OrdinaryDiffEq and LinearSolve)
+# SciML packages (for QobjEvo, OrdinaryDiffEq, and LinearSolve)
 import SciMLBase:
     solve,
     solve!,
@@ -34,16 +36,15 @@ import SciMLBase:
     DiscreteCallback
 import StochasticDiffEq: StochasticDiffEqAlgorithm, SRA1
 import SciMLOperators:
+    SciMLOperators,
     AbstractSciMLOperator,
     MatrixOperator,
     ScalarOperator,
     ScaledOperator,
     AddedOperator,
     IdentityOperator,
-    cache_operator,
     update_coefficients!,
-    concretize,
-    isconstant
+    concretize
 import LinearSolve: LinearProblem, SciMLLinearSolveAlgorithm, KrylovJL_MINRES, KrylovJL_GMRES
 import DiffEqBase: get_tstops
 import DiffEqCallbacks: PeriodicCallback, PresetTimeCallback, TerminateSteadyState

src/qobj/boolean_functions.jl

@@ -3,7 +3,7 @@ All boolean functions for checking the data or type in `QuantumObject`
 =#
 
 export isket, isbra, isoper, isoperbra, isoperket, issuper
-export isunitary, isconstant
+export isunitary
 
 @doc raw"""
     isbra(A)
@@ -91,8 +91,15 @@ isunitary(U::QuantumObject{<:AbstractArray{T}}; kwargs...) where {T} =
     isoper(U) ? isapprox(U.data * U.data', I(size(U, 1)); kwargs...) : false
 
 @doc raw"""
-    isconstant(A::AbstractQuantumObject)
+    SciMLOperators.iscached(A::AbstractQuantumObject)
+
+Test whether the [`AbstractQuantumObject`](@ref) `A` has preallocated caches for inplace evaluations.
+"""
+SciMLOperators.iscached(A::AbstractQuantumObject) = iscached(A.data)
+
+@doc raw"""
+    SciMLOperators.isconstant(A::AbstractQuantumObject)
 
 Test whether the [`AbstractQuantumObject`](@ref) `A` is constant in time. For a [`QuantumObject`](@ref), this function returns `true`, while for a [`QuantumObjectEvolution`](@ref), this function returns `true` if the operator is contant in time.
 """
-isconstant(A::AbstractQuantumObject) = isconstant(A.data)
+SciMLOperators.isconstant(A::AbstractQuantumObject) = isconstant(A.data)

src/qobj/quantum_object.jl

@@ -3,6 +3,7 @@ This file defines the QuantumObject (Qobj) structure.
 It also implements the fundamental functions in Julia standard library:
     - Base: show, real, imag, Vector, Matrix
     - SparseArrays: sparse, nnz, nonzeros, rowvals, droptol!, dropzeros, dropzeros!, SparseVector, SparseMatrixCSC
+    - SciMLOperators: cache_operator
 =#
 
 export QuantumObject
@@ -167,6 +168,41 @@ SparseArrays.droptol!(A::QuantumObject{<:AbstractSparseArray}, tol::Real) = (dro
 SparseArrays.dropzeros(A::QuantumObject{<:AbstractSparseArray}) = QuantumObject(dropzeros(A.data), A.type, A.dims)
 SparseArrays.dropzeros!(A::QuantumObject{<:AbstractSparseArray}) = (dropzeros!(A.data); return A)
 
+@doc raw"""
+    SciMLOperators.cached_operator(L::AbstractQuantumObject, u)
+
+Allocate caches for [`AbstractQuantumObject`](@ref) `L` for in-place evaluation with `u`-like input vectors.
+
+Here, `u` can be in either the following types:
+- `AbstractVector`
+- [`Ket`](@ref)-type [`QuantumObject`](@ref) (if `L` is an [`Operator`](@ref))
+- [`OperatorKet`](@ref)-type [`QuantumObject`](@ref) (if `L` is a [`SuperOperator`](@ref))
+"""
+SciMLOperators.cache_operator(
+    L::AbstractQuantumObject{DT,OpType},
+    u::AbstractVector,
+) where {DT,OpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject}} =
+    get_typename_wrapper(L)(cache_operator(L.data, sparse_to_dense(similar(u))), L.type, L.dims)
+
+function SciMLOperators.cache_operator(
+    L::AbstractQuantumObject{DT1,OpType},
+    u::QuantumObject{DT2,SType},
+) where {
+    DT1,
+    DT2,
+    OpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject},
+    SType<:Union{KetQuantumObject,OperatorKetQuantumObject},
+}
+    check_dims(L, u)
+
+    if isoper(L) && isoperket(u)
+        throw(ArgumentError("The input state `u` must be a Ket if `L` is an Operator."))
+    elseif issuper(L) && isket(u)
+        throw(ArgumentError("The input state `u` must be an OperatorKet if `L` is a SuperOperator."))
+    end
+    return cache_operator(L, u.data)
+end
+
 # data type conversions
 Base.Vector(A::QuantumObject{<:AbstractVector}) = QuantumObject(Vector(A.data), A.type, A.dims)
 Base.Vector{T}(A::QuantumObject{<:AbstractVector}) where {T<:Number} = QuantumObject(Vector{T}(A.data), A.type, A.dims)

src/qobj/quantum_object_base.jl

@@ -209,8 +209,6 @@ function _check_QuantumObject(type::OperatorBraQuantumObject, dims, m::Int, n::I
     return nothing
 end
 
-get_typename_wrapper(A::AbstractQuantumObject) = Base.typename(typeof(A)).wrapper
-
 # functions for getting Float or Complex element type
 _FType(A::AbstractQuantumObject) = _FType(eltype(A))
 _CType(A::AbstractQuantumObject) = _CType(eltype(A))

src/qobj/quantum_object_evo.jl

@@ -158,9 +158,8 @@ function QuantumObjectEvolution(
     op_func_list::Tuple,
     α::Union{Nothing,Number} = nothing;
     type::Union{Nothing,QuantumObjectType} = nothing,
-    f::Function = identity,
 )
-    op, data = _QobjEvo_generate_data(op_func_list, α; f = f)
+    op, data = _QobjEvo_generate_data(op_func_list, α)
     dims = op.dims
     if type isa Nothing
         type = op.type
@@ -177,24 +176,21 @@ function QuantumObjectEvolution(
     op::QuantumObject,
     α::Union{Nothing,Number} = nothing;
     type::Union{Nothing,QuantumObjectType} = nothing,
-    f::Function = identity,
 )
     if type isa Nothing
         type = op.type
     end
-    return QuantumObjectEvolution(_make_SciMLOperator(op, α, f = f), type, op.dims)
+    return QuantumObjectEvolution(_make_SciMLOperator(op, α), type, op.dims)
 end
 
 function QuantumObjectEvolution(
     op::QuantumObjectEvolution,
     α::Union{Nothing,Number} = nothing;
     type::Union{Nothing,QuantumObjectType} = nothing,
-    f::Function = identity,
 )
-    f !== identity && throw(ArgumentError("The function `f` is not supported for QuantumObjectEvolution inputs."))
     if type isa Nothing
         type = op.type
-    else
+    elseif type != op.type
         throw(
             ArgumentError(
                 "The type of the QuantumObjectEvolution object cannot be changed when using another QuantumObjectEvolution object as input.",
@@ -217,7 +213,7 @@ Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution`
 - `α`: A scalar to pre-multiply the operators.
 - `f::Function=identity`: A function to pre-apply to the operators.
 =#
-@generated function _QobjEvo_generate_data(op_func_list::Tuple, α; f::Function = identity)
+@generated function _QobjEvo_generate_data(op_func_list::Tuple, α)
     op_func_list_types = op_func_list.parameters
     N = length(op_func_list_types)
 
@@ -242,9 +238,9 @@ Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution`
             data_type = op_type.parameters[1]
             dims_expr = (dims_expr..., :($op.dims))
             if i == 1
-                first_op = :(f($op))
+                first_op = :($op)
             end
-            data_expr = :($data_expr + _make_SciMLOperator(op_func_list[$i], α, f = f))
+            data_expr = :($data_expr + _make_SciMLOperator(op_func_list[$i], α))
         else
             op_type = op_func_type
             (isoper(op_type) || issuper(op_type)) ||
@@ -255,7 +251,7 @@ Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution`
             if i == 1
                 first_op = :(op_func_list[$i])
             end
-            qobj_expr_const = :($qobj_expr_const + f(op_func_list[$i]))
+            qobj_expr_const = :($qobj_expr_const + op_func_list[$i])
         end
     end
 
@@ -272,20 +268,20 @@ Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution`
     end
 end
 
-function _make_SciMLOperator(op_func::Tuple, α; f::Function = identity)
+function _make_SciMLOperator(op_func::Tuple, α)
     T = eltype(op_func[1])
     update_func = (a, u, p, t) -> op_func[2](p, t)
     if α isa Nothing
-        return ScalarOperator(zero(T), update_func) * MatrixOperator(f(op_func[1]).data)
+        return ScalarOperator(zero(T), update_func) * MatrixOperator(op_func[1].data)
     end
-    return ScalarOperator(zero(T), update_func) * MatrixOperator(α * f(op_func[1]).data)
+    return ScalarOperator(zero(T), update_func) * MatrixOperator(α * op_func[1].data)
 end
 
-function _make_SciMLOperator(op::QuantumObject, α; f::Function = identity)
+function _make_SciMLOperator(op::QuantumObject, α)
     if α isa Nothing
-        return MatrixOperator(f(op).data)
+        return MatrixOperator(op.data)
     end
-    return MatrixOperator(α * f(op.data))
+    return MatrixOperator(α * op.data)
 end
 
 @doc raw"""
@@ -375,15 +371,11 @@ function (A::QuantumObjectEvolution)(
     check_dims(ψout, A)
 
     if isoper(A) && isoperket(ψin)
-        throw(
-            ArgumentError(
-                "The input state must be a KetQuantumObject if the QuantumObjectEvolution object is an Operator.",
-            ),
-        )
+        throw(ArgumentError("The input state must be a Ket if the QuantumObjectEvolution object is an Operator."))
     elseif issuper(A) && isket(ψin)
         throw(
             ArgumentError(
-                "The input state must be an OperatorKetQuantumObject if the QuantumObjectEvolution object is a SuperOperator.",
+                "The input state must be an OperatorKet if the QuantumObjectEvolution object is a SuperOperator.",
             ),
         )
     end