Add SciMLOperators-backed lazy operators and benchmarks

Project.toml

@@ -26,6 +26,7 @@ SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 
 [extensions]
 QuantumOpticsBaseMakieExt = ["Makie"]
+QuantumOpticsBaseSciMLOperatorsExt = ["SciMLOperators"]
 
 [compat]
 Adapt = "4.1"

benchmark/Project.toml

@@ -0,0 +1,4 @@
+[deps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+QuantumOpticsBase = "4f57444f-1401-5e15-980d-4471b28d5678"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"

benchmark/sciml_lazyoperators.jl

@@ -0,0 +1,88 @@
+import Pkg
+Pkg.activate(@__DIR__)
+Pkg.develop(path=joinpath(@__DIR__, ".."))
+
+using BenchmarkTools
+using SciMLOperators
+using QuantumOpticsBase
+using Random
+
+const RNG = MersenneTwister(20260609)
+
+function spin_chain(n)
+    b = SpinBasis(1//2)
+    foldl(⊗, fill(b, n))
+end
+
+function local_spin_term(chain, idx, op)
+    LazyTensor(chain, chain, [idx], (op,))
+end
+
+function structured_hamiltonian(n)
+    chain = spin_chain(n)
+    sx = sigmax(SpinBasis(1//2))
+    sz = sigmaz(SpinBasis(1//2))
+    terms = Any[]
+    for i in 1:n
+        push!(terms, local_spin_term(chain, i, isodd(i) ? sx : sz))
+    end
+    if n >= 2
+        push!(terms, LazyProduct(local_spin_term(chain, 1, sx), local_spin_term(chain, 2, sz)))
+    end
+    LazySum(terms...)
+end
+
+function benchmark_group()
+    group = BenchmarkGroup()
+    group["lazy_sum"] = BenchmarkGroup()
+    group["lazy_product"] = BenchmarkGroup()
+    group["lazy_tensor"] = BenchmarkGroup()
+    group["mixed_hamiltonian"] = BenchmarkGroup()
+
+    for n in (4, 6, 8)
+        chain = spin_chain(n)
+        ψ = Ket(chain, randn(RNG, ComplexF64, length(chain)))
+        H = structured_hamiltonian(n)
+        Hs = QuantumOpticsBase.sciml_lazy_operator(H)
+        Hsc = cache_operator(Hs, ψ.data)
+
+        group["lazy_sum"][n] = BenchmarkGroup()
+        group["lazy_sum"][n]["current"] = @benchmarkable $H * $ψ
+        group["lazy_sum"][n]["sciml"] = @benchmarkable $Hs * $ψ
+        group["lazy_sum"][n]["sciml_cached"] = @benchmarkable $Hsc * $ψ
+
+        sx = sigmax(SpinBasis(1//2))
+        sz = sigmaz(SpinBasis(1//2))
+
+        P = LazyProduct(local_spin_term(chain, 1, sx),
+                        local_spin_term(chain, 2, sz))
+        Ps = QuantumOpticsBase.sciml_lazy_operator(P)
+        Psc = cache_operator(Ps, ψ.data)
+
+        group["lazy_product"][n] = BenchmarkGroup()
+        group["lazy_product"][n]["current"] = @benchmarkable $P * $ψ
+        group["lazy_product"][n]["sciml"] = @benchmarkable $Ps * $ψ
+        group["lazy_product"][n]["sciml_cached"] = @benchmarkable $Psc * $ψ
+
+        T = LazyTensor(chain, chain, [1, n], (sx, sz))
+        Ts = QuantumOpticsBase.sciml_lazy_operator(T)
+        Tsc = cache_operator(Ts, ψ.data)
+
+        group["lazy_tensor"][n] = BenchmarkGroup()
+        group["lazy_tensor"][n]["current"] = @benchmarkable $T * $ψ
+        group["lazy_tensor"][n]["sciml"] = @benchmarkable $Ts * $ψ
+        group["lazy_tensor"][n]["sciml_cached"] = @benchmarkable $Tsc * $ψ
+
+        group["mixed_hamiltonian"][n] = BenchmarkGroup()
+        group["mixed_hamiltonian"][n]["current"] = @benchmarkable $H * $ψ
+        group["mixed_hamiltonian"][n]["sciml"] = @benchmarkable $Hs * $ψ
+        group["mixed_hamiltonian"][n]["sciml_cached"] = @benchmarkable $Hsc * $ψ
+    end
+
+    return group
+end
+
+if abspath(PROGRAM_FILE) == @__FILE__
+    results = run(benchmark_group(); verbose = true)
+    display(results)
+end

ext/QuantumOpticsBaseSciMLOperatorsExt.jl

@@ -0,0 +1,174 @@
+module QuantumOpticsBaseSciMLOperatorsExt
+
+import Base: copy, size, eltype, +, -, *, /
+import LinearAlgebra: mul!, adjoint, I
+
+import QuantumOpticsBase
+import QuantumOpticsBase: AbstractOperator, Basis, CompositeBasis, DataOperator, DenseOperator,
+    SparseOperator, LazySum, LazyProduct, LazyTensor, Ket, Bra, Operator, basis, dense, sparse,
+    dagger, identityoperator, suboperator, check_samebases, sciml_lazy_operator
+
+import SciMLOperators
+using SciMLOperators: MatrixOperator, TensorProductOperator, cache_operator
+
+struct SciMLLazyOperator{BL,BR,O} <: AbstractOperator{BL,BR}
+    basis_l::BL
+    basis_r::BR
+    op::O
+end
+
+QuantumOpticsBase.basis(op::SciMLLazyOperator) = op.basis_l
+
+Base.copy(op::SciMLLazyOperator) = SciMLLazyOperator(op.basis_l, op.basis_r, copy(op.op))
+Base.eltype(op::SciMLLazyOperator) = eltype(op.op)
+Base.size(op::SciMLLazyOperator) = size(op.op)
+Base.size(op::SciMLLazyOperator, d::Int) = size(op.op, d)
+
+dense(op::SciMLLazyOperator) = DenseOperator(op.basis_l, op.basis_r, Matrix(op.op))
+sparse(op::SciMLLazyOperator) = SparseOperator(op.basis_l, op.basis_r, sparse(op.op))
+
+function _sciml_identity(b1::Basis, b2::Basis, T)
+    MatrixOperator(Matrix{T}(I, length(b1), length(b2)))
+end
+
+_sciml_raw(op::SciMLLazyOperator) = op.op
+_sciml_raw(op::DataOperator) = MatrixOperator(op.data)
+_sciml_raw(op::AbstractOperator) = MatrixOperator(dense(op).data)
+
+function _sciml_raw(op::LazySum)
+    if isempty(op.operators)
+        T = eltype(op.factors)
+        return MatrixOperator(zeros(T, length(op.basis_l), length(op.basis_r)))
+    end
+    terms = map(i -> op.factors[i] * _sciml_raw(op.operators[i]), eachindex(op.operators))
+    return reduce(+, terms)
+end
+
+function _sciml_raw(op::LazyProduct)
+    if isempty(op.operators)
+        T = eltype(op)
+        return MatrixOperator(Matrix{T}(I, length(op.basis_l), length(op.basis_r)))
+    end
+    return reduce(*, (_sciml_raw(o) for o in op.operators))
+end
+
+function _sciml_raw(op::LazyTensor)
+    raw_ops = Any[]
+    sizehint!(raw_ops, length(op.basis_l.bases))
+    for idx in 1:length(op.basis_l.bases)
+        term = findfirst(isequal(idx), op.indices)
+        if term === nothing
+            push!(raw_ops, _sciml_identity(op.basis_l.bases[idx], op.basis_r.bases[idx], eltype(op)))
+        else
+            push!(raw_ops, _sciml_raw(suboperator(op, idx)))
+        end
+    end
+    raw = TensorProductOperator(reverse(raw_ops)...)
+    return isone(op.factor) ? raw : op.factor * raw
+end
+
+sciml_lazy_operator(op::SciMLLazyOperator) = op
+sciml_lazy_operator(op::LazySum) = SciMLLazyOperator(op.basis_l, op.basis_r, _sciml_raw(op))
+sciml_lazy_operator(op::LazyProduct) = SciMLLazyOperator(op.basis_l, op.basis_r, _sciml_raw(op))
+sciml_lazy_operator(op::LazyTensor) = SciMLLazyOperator(op.basis_l, op.basis_r, _sciml_raw(op))
+sciml_lazy_operator(op::DataOperator) = SciMLLazyOperator(op.basis_l, op.basis_r, _sciml_raw(op))
+sciml_lazy_operator(op::AbstractOperator) = SciMLLazyOperator(op.basis_l, op.basis_r, _sciml_raw(op))
+
+function SciMLOperators.cache_operator(op::SciMLLazyOperator, sample)
+    SciMLLazyOperator(op.basis_l, op.basis_r, cache_operator(op.op, sample))
+end
+
+function mul!(result::Ket{B1}, op::SciMLLazyOperator{B1,B2}, state::Ket{B2}, alpha, beta) where {B1,B2}
+    cached = cache_operator(op.op, state.data)
+    mul!(result.data, cached, state.data, alpha, beta)
+    return result
+end
+
+function mul!(result::Bra{B2}, state::Bra{B1}, op::SciMLLazyOperator{B1,B2}, alpha, beta) where {B1,B2}
+    cached = cache_operator(adjoint(op.op), state.data)
+    mul!(result.data, cached, state.data, alpha, beta)
+    return result
+end
+
+function mul!(result::Operator{B1,B3}, op::SciMLLazyOperator{B1,B2}, rhs::Operator{B2,B3}, alpha, beta) where {B1,B2,B3}
+    mul!(result.data, dense(op).data, rhs.data, alpha, beta)
+    return result
+end
+
+function mul!(result::Operator{B1,B3}, lhs::Operator{B1,B2}, op::SciMLLazyOperator{B2,B3}, alpha, beta) where {B1,B2,B3}
+    mul!(result.data, lhs.data, dense(op).data, alpha, beta)
+    return result
+end
+
+function +(a::SciMLLazyOperator{B1,B2}, b::SciMLLazyOperator{B1,B2}) where {B1,B2}
+    check_samebases(a, b)
+    SciMLLazyOperator(a.basis_l, a.basis_r, a.op + b.op)
+end
+
+function +(a::SciMLLazyOperator, b::AbstractOperator)
+    check_samebases(a, b)
+    SciMLLazyOperator(a.basis_l, a.basis_r, a.op + _sciml_raw(b))
+end
+
+function +(a::AbstractOperator, b::SciMLLazyOperator)
+    check_samebases(a, b)
+    SciMLLazyOperator(a.basis_l, a.basis_r, _sciml_raw(a) + b.op)
+end
+
+function -(a::SciMLLazyOperator{B1,B2}, b::SciMLLazyOperator{B1,B2}) where {B1,B2}
+    check_samebases(a, b)
+    SciMLLazyOperator(a.basis_l, a.basis_r, a.op - b.op)
+end
+
+function -(a::SciMLLazyOperator, b::AbstractOperator)
+    check_samebases(a, b)
+    SciMLLazyOperator(a.basis_l, a.basis_r, a.op - _sciml_raw(b))
+end
+
+function -(a::AbstractOperator, b::SciMLLazyOperator)
+    check_samebases(a, b)
+    SciMLLazyOperator(a.basis_l, a.basis_r, _sciml_raw(a) - b.op)
+end
+
+-(a::SciMLLazyOperator) = SciMLLazyOperator(a.basis_l, a.basis_r, -a.op)
+
+function *(a::SciMLLazyOperator{B1,B2}, b::SciMLLazyOperator{B2,B3}) where {B1,B2,B3}
+    check_samebases(a.basis_r, b.basis_l)
+    SciMLLazyOperator(a.basis_l, b.basis_r, a.op * b.op)
+end
+
+function *(a::SciMLLazyOperator, b::AbstractOperator)
+    check_samebases(a.basis_r, b.basis_l)
+    SciMLLazyOperator(a.basis_l, b.basis_r, a.op * _sciml_raw(b))
+end
+
+function *(a::AbstractOperator, b::SciMLLazyOperator)
+    check_samebases(a.basis_r, b.basis_l)
+    SciMLLazyOperator(a.basis_l, b.basis_r, _sciml_raw(a) * b.op)
+end
+
+*(a::SciMLLazyOperator, b::Number) = SciMLLazyOperator(a.basis_l, a.basis_r, a.op * b)
+*(a::Number, b::SciMLLazyOperator) = SciMLLazyOperator(b.basis_l, b.basis_r, a * b.op)
+/(a::SciMLLazyOperator, b::Number) = SciMLLazyOperator(a.basis_l, a.basis_r, a.op / b)
+
+dagger(op::SciMLLazyOperator) = SciMLLazyOperator(op.basis_r, op.basis_l, adjoint(op.op))
+
+function tensor(a::SciMLLazyOperator, b::SciMLLazyOperator)
+    SciMLLazyOperator(tensor(a.basis_l, b.basis_l), tensor(a.basis_r, b.basis_r),
+                      TensorProductOperator(b.op, a.op))
+end
+
+function tensor(a::SciMLLazyOperator, b::AbstractOperator)
+    SciMLLazyOperator(tensor(a.basis_l, b.basis_l), tensor(a.basis_r, b.basis_r),
+                      TensorProductOperator(_sciml_raw(b), a.op))
+end
+
+function tensor(a::AbstractOperator, b::SciMLLazyOperator)
+    SciMLLazyOperator(tensor(a.basis_l, b.basis_l), tensor(a.basis_r, b.basis_r),
+                      TensorProductOperator(b.op, _sciml_raw(a)))
+end
+
+identityoperator(::Type{<:SciMLLazyOperator}, ::Type{S}, b1::Basis, b2::Basis) where {S<:Number} =
+    sciml_lazy_operator(identityoperator(S, b1, b2))
+
+end # module

src/QuantumOpticsBase.jl

@@ -27,12 +27,14 @@ export Basis, GenericBasis, CompositeBasis, basis,
                 SparseOperator, diagonaloperator, SparseOpType, EyeOpType,
         #operators_lazysum
                 LazySum,
-        #operators_lazyproduct
+                #operators_lazyproduct
                 LazyProduct,
-        #operators_lazytensor
+                #operators_lazytensor
                 LazyTensor, lazytensor_use_cache, lazytensor_clear_cache,
                 lazytensor_cachesize, lazytensor_disable_cache, lazytensor_enable_cache,
-        #states_lazyket
+                # SciML prototype
+                sciml_lazy_operator,
+                #states_lazyket
                 LazyKet,
         #time_dependent_operators
                 AbstractTimeDependentOperator, TimeDependentSum, set_time!,
@@ -107,4 +109,6 @@ include("printing.jl")
 include("apply.jl")
 include("visualization.jl")
 
+function sciml_lazy_operator end
+
 end # module

test/Project.toml

@@ -11,8 +11,10 @@ OrdinaryDiffEqLowOrderRK = "1344f307-1e59-4825-a18e-ace9aa3fa4c6"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 QuantumInterface = "5717a53b-5d69-4fa3-b976-0bf2f97ca1e5"
 QuantumOptics = "6e0679c1-51ea-5a7c-ac74-d61b76210b0c"
+QuantumOpticsBase = "4f57444f-1401-5e15-980d-4471b28d5678"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/runtests.jl

@@ -71,4 +71,6 @@ testfilter = ti -> begin
 end
 println("Starting tests with $(Threads.nthreads()) threads out of `Sys.CPU_THREADS = $(Sys.CPU_THREADS)`...")
 
+include("test_sciml_lazyoperators.jl")
+
 @run_package_tests filter=testfilter

test/test_sciml_lazyoperators.jl

@@ -0,0 +1,49 @@
+using Random
+using Test
+using SciMLOperators
+using QuantumOpticsBase
+
+Test.@testset "test_sciml_lazyoperators" begin
+    rng = MersenneTwister(42)
+    b = SpinBasis(1//2)
+    chain = b ⊗ b ⊗ b ⊗ b
+
+    sx = sigmax(b)
+    sz = sigmaz(b)
+
+    H_current = LazySum(
+        LazyTensor(chain, chain, [1], (sx,)),
+        LazyTensor(chain, chain, [2], (sz,)),
+        LazyProduct(
+            LazyTensor(chain, chain, [3], (sx,)),
+            LazyTensor(chain, chain, [4], (sz,)),
+        ),
+    )
+
+    H_sciml = QuantumOpticsBase.sciml_lazy_operator(H_current)
+    H_dense = dense(H_current)
+    ψ = Ket(chain, randn(rng, ComplexF64, length(chain)))
+
+    Test.@test H_sciml.basis_l == H_current.basis_l
+    Test.@test H_sciml.basis_r == H_current.basis_r
+    Test.@test dense(H_sciml).data ≈ H_dense.data
+    Test.@test (H_sciml * ψ).data ≈ (H_current * ψ).data
+    Test.@test (2.0 * H_sciml * ψ).data ≈ ((2.0 * H_current) * ψ).data
+    Test.@test (dagger(H_sciml) * ψ).data ≈ (dagger(H_current) * ψ).data
+    Test.@test dense(H_sciml + H_sciml).data ≈ dense(2 * H_current).data
+
+    A = LazyTensor(chain, chain, [1], (sx,))
+    B = LazyTensor(chain, chain, [2], (sz,))
+    A_sciml = QuantumOpticsBase.sciml_lazy_operator(A)
+    B_sciml = QuantumOpticsBase.sciml_lazy_operator(B)
+    Test.@test dense(A_sciml * B_sciml).data ≈ dense(A * B).data
+
+    local_dense = DenseOperator(chain, chain, randn(rng, ComplexF64, length(chain), length(chain)))
+    local_sparse = SparseOperator(chain, sparse(local_dense.data))
+
+    Test.@test dense(QuantumOpticsBase.sciml_lazy_operator(local_dense)).data ≈ local_dense.data
+    Test.@test dense(QuantumOpticsBase.sciml_lazy_operator(local_sparse)).data ≈ local_sparse.data
+
+    cached = cache_operator(H_sciml, ψ.data)
+    Test.@test (cached * ψ).data ≈ (H_sciml * ψ).data
+end