re-export some functions in SciMLOperators

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,28 @@ 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:
+    AbstractSciMLOperator,
+    MatrixOperator,
+    ScalarOperator,
+    ScaledOperator,
+    AddedOperator,
+    IdentityOperator,
+    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!,
@@ -33,17 +44,7 @@ import SciMLBase:
     ContinuousCallback,
     DiscreteCallback
 import StochasticDiffEq: StochasticDiffEqAlgorithm, SRA1
-import SciMLOperators:
-    AbstractSciMLOperator,
-    MatrixOperator,
-    ScalarOperator,
-    ScaledOperator,
-    AddedOperator,
-    IdentityOperator,
-    cache_operator,
-    update_coefficients!,
-    concretize,
-    isconstant
+import SciMLOperators: SciMLOperators, update_coefficients!, 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)

test/core-test/quantum_objects_evo.jl

@@ -174,6 +174,9 @@
         H_ti = coef1(p, t) * a + a' * a + coef2(p, t) * a'
         ψ = rand_ket(N)
         @test H_td(p, t) ≈ H_ti
+        @test iscached(H_td) == true
+        H_td = cache_operator(H_td, ψ)
+        @test iscached(H_td) == true
         @test H_td(ψ, p, t) ≈ H_ti * ψ
         @test isconstant(a) == true
         @test isconstant(H_td) == false
@@ -195,18 +198,18 @@
         L_td = @test_logs (:warn,) (:warn,) liouvillian(H_td, c_ops) # warnings from lazy tensor in `lindblad_dissipator(c_op2)`
         ρvec = mat2vec(rand_dm(N))
         @test L_td(p, t) ≈ L_ti
-        # TODO: L_td here is ComposedOperator and need to setup cache first for the following test
-        # TODO: (maybe can support `iscached` and `cache_operator` for QobjEvo in the future)
-        # @test iscached(L_td) == false
-        # @test L_td = cache_operator(L_td, ρvec)
-        # @test iscached(L_td) == true
-        # @test L_td(ρvec, p, t) ≈ L_ti * ρvec 
+        @test iscached(L_td) == false
+        L_td = cache_operator(L_td, ρvec)
+        @test iscached(L_td) == true
+        @test L_td(ρvec, p, t) ≈ L_ti * ρvec
         @test isconstant(L_td) == false
         @test issuper(L_td) == true
 
         @test_logs (:warn,) (:warn,) liouvillian(H_td * H_td) # warnings from lazy tensor
         @test_throws MethodError QobjEvo([[a, coef1], a' * a, [a', coef2]])
         @test_throws ArgumentError H_td(ρvec, p, t)
+        @test_throws ArgumentError cache_operator(H_td, ρvec)
         @test_throws ArgumentError L_td(ψ, p, t)
+        @test_throws ArgumentError cache_operator(L_td, ψ)
     end
 end