Prototype SciMLOperators-backed lazy operators

Author: cmdr-chara

Project.toml

@@ -19,10 +19,12 @@ UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 
 [weakdeps]
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
- 
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
+
 [extensions]
 QuantumOpticsBaseMakieExt = ["Makie"]
- 
+QuantumOpticsBaseSciMLOperatorsExt = "SciMLOperators"
+
 [compat]
 Adapt = "3.3, 4"
 FFTW = "1.2"
@@ -35,7 +37,8 @@ Makie = "0.24"
 QuantumInterface = "0.4"
 Random = "1"
 RecursiveArrayTools = "3.1, 4"
+SciMLOperators = "1"
 SparseArrays = "1"
 Strided = "1, 2"
 UnsafeArrays = "1"
-julia = "1.10"
+julia = "1.10"

benchmark/Project.toml

@@ -0,0 +1,9 @@
+name = "QuantumOpticsBaseBenchmarks"
+uuid = "729dd8e0-ae36-4b43-9a76-c0170d3614b7"
+
+[deps]
+BenchmarkTools = "6e4b80f9-dd1a-5c18-9360-3c73fbc0a7c0"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+QuantumOpticsBase = "4f57444f-1401-5e15-980d-4471b28d5678"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"

benchmark/sciml_lazy_operators.jl

@@ -0,0 +1,119 @@
+using BenchmarkTools
+using LinearAlgebra
+using QuantumOpticsBase
+using Random
+using SciMLOperators
+
+const SUITE = BenchmarkGroup()
+
+for family in ("lazy_sum", "lazy_product", "lazy_tensor", "mixed_hamiltonian")
+    SUITE[family] = BenchmarkGroup([family])
+end
+
+function chain_basis(n)
+    site = SpinBasis(1//2)
+    return tensor(ntuple(_ -> site, n)...)
+end
+
+function random_ket(basis, seed)
+    rng = MersenneTwister(seed)
+    return Ket(basis, randn(rng, ComplexF64, length(basis)))
+end
+
+function chain_sum_case(n)
+    site = SpinBasis(1//2)
+    basis = tensor(ntuple(_ -> site, n)...)
+    sx = sigmax(site)
+    sz = sigmaz(site)
+    terms = AbstractOperator[LazyTensor(basis, i, isodd(i) ? sx : sz) for i in 1:n]
+    return LazySum(terms...)
+end
+
+function product_case(n)
+    site = SpinBasis(1//2)
+    basis = tensor(ntuple(_ -> site, n)...)
+    sx = sigmax(site)
+    sz = sigmaz(site)
+    terms = AbstractOperator[LazyTensor(basis, i, isodd(i) ? sx : sz) for i in 1:n]
+    return LazyProduct(terms...)
+end
+
+function two_site_product_case(n)
+    site = SpinBasis(1//2)
+    basis = tensor(ntuple(_ -> site, n)...)
+    mid = max(1, n ÷ 2)
+    return LazyProduct(LazyTensor(basis, mid, sigmax(site)), LazyTensor(basis, mid + 1, sigmaz(site)))
+end
+
+function tensor_case(n)
+    site = SpinBasis(1//2)
+    basis = tensor(ntuple(_ -> site, n)...)
+    return LazyTensor(basis, [1, n], (sigmax(site), sigmaz(site)))
+end
+
+function mixed_hamiltonian_case(n)
+    site = SpinBasis(1//2)
+    basis = tensor(ntuple(_ -> site, n)...)
+    sx = sigmax(site)
+    sz = sigmaz(site)
+    terms = AbstractOperator[]
+
+    for i in 1:(n - 1)
+        push!(terms, LazyProduct(LazyTensor(basis, i, sx), LazyTensor(basis, i + 1, sx)))
+        push!(terms, LazyProduct(LazyTensor(basis, i, sz), LazyTensor(basis, i + 1, sz)))
+    end
+
+    for i in 1:n
+        push!(terms, 0.3 * LazyTensor(basis, i, sz))
+    end
+
+    return LazySum(terms...)
+end
+
+function add_apply_benchmarks!(suite, name, current, psi)
+    group = BenchmarkGroup([name])
+    sciml_uncached = sciml_lazy_operator(current)
+    sciml_cached = cache_sciml_lazy_operator(sciml_uncached, psi)
+
+    out_current = Ket(current.basis_l)
+    out_uncached = Ket(current.basis_l)
+    out_cached = Ket(current.basis_l)
+
+    group["current"] = @benchmarkable mul!($out_current, $current, $psi, 1, 0)
+    group["sciml_uncached"] = @benchmarkable mul!($out_uncached, $sciml_uncached, $psi, 1, 0)
+    group["sciml_cached"] = @benchmarkable mul!($out_cached, $sciml_cached, $psi, 1, 0)
+
+    suite[name] = group
+    return group
+end
+
+for n in (4, 6, 8)
+    op = chain_sum_case(n)
+    add_apply_benchmarks!(SUITE["lazy_sum"], "n=$n", op, random_ket(op.basis_r, n))
+end
+
+for n in (4, 6)
+    op = product_case(n)
+    add_apply_benchmarks!(SUITE["lazy_product"], "n=$n", op, random_ket(op.basis_r, 20 + n))
+end
+
+for n in (4, 6, 8)
+    op = two_site_product_case(n)
+    add_apply_benchmarks!(SUITE["lazy_product"], "two_site n=$n", op, random_ket(op.basis_r, 30 + n))
+end
+
+for n in (4, 6, 8)
+    op = tensor_case(n)
+    add_apply_benchmarks!(SUITE["lazy_tensor"], "n=$n", op, random_ket(op.basis_r, 40 + n))
+end
+
+for n in (4, 6)
+    op = mixed_hamiltonian_case(n)
+    add_apply_benchmarks!(SUITE["mixed_hamiltonian"], "n=$n", op, random_ket(op.basis_r, 60 + n))
+end
+
+if abspath(PROGRAM_FILE) == @__FILE__
+    results = run(SUITE; verbose=true)
+    show(stdout, MIME("text/plain"), results)
+    println()
+end

ext/QuantumOpticsBaseSciMLOperatorsExt.jl

@@ -0,0 +1,230 @@
+module QuantumOpticsBaseSciMLOperatorsExt
+
+import Base: +, -, *, /, copy, eltype, size, transpose
+import LinearAlgebra: adjoint, mul!
+import SparseArrays
+import QuantumOpticsBase
+import QuantumOpticsBase: Bra, Ket
+import SciMLOperators
+
+const QOB = QuantumOpticsBase
+
+struct SciMLLazyOperator{BL,BR,O} <: QOB.LazyOperator{BL,BR}
+    basis_l::BL
+    basis_r::BR
+    operator::O
+
+    function SciMLLazyOperator(basis_l::BL, basis_r::BR, operator::O) where {BL,BR,O}
+        expected = (length(basis_l), length(basis_r))
+        size(operator) == expected || throw(DimensionMismatch("SciML operator has size $(size(operator)), expected $expected from bases."))
+        return new{BL,BR,O}(basis_l, basis_r, operator)
+    end
+end
+
+copy(op::SciMLLazyOperator) = SciMLLazyOperator(op.basis_l, op.basis_r, copy(op.operator))
+eltype(op::SciMLLazyOperator) = eltype(op.operator)
+size(op::SciMLLazyOperator) = (length(op.basis_l), length(op.basis_r))
+size(op::SciMLLazyOperator, dim::Integer) = size(op)[dim]
+
+function _check_same_bases(a, b)
+    a.basis_l == b.basis_l && a.basis_r == b.basis_r || throw(QOB.IncompatibleBases())
+    return nothing
+end
+
+function _check_multiplicable(a, b)
+    a.basis_r == b.basis_l || throw(QOB.IncompatibleBases())
+    return nothing
+end
+
+_wrap(basis_l, basis_r, operator) = SciMLLazyOperator(basis_l, basis_r, operator)
+
+function _operator_number_type(op)
+    T = eltype(op)
+    return T <: Number ? T : ComplexF64
+end
+
+function _matrix_identity(::Type{T}, basis_l, basis_r) where {T}
+    if length(basis_l) == length(basis_r)
+        return SciMLOperators.IdentityOperator(length(basis_l))
+    end
+    return SciMLOperators.MatrixOperator(QOB.identityoperator(T, basis_l, basis_r).data)
+end
+
+_sciml_operator(op::SciMLLazyOperator) = op.operator
+_sciml_operator(op::QOB.DataOperator) = SciMLOperators.MatrixOperator(op.data)
+
+function _sciml_operator(op::QOB.LazySum)
+    if isempty(op.operators)
+        return SciMLOperators.NullOperator(length(op.basis_l), length(op.basis_r))
+    end
+
+    result = op.factors[1] * _sciml_operator(op.operators[1])
+    for i in 2:length(op.operators)
+        result = result + op.factors[i] * _sciml_operator(op.operators[i])
+    end
+    return result
+end
+
+function _combined_lazytensor_product(op::QOB.LazyProduct)
+    length(op.operators) == 2 || return nothing
+    all(operator -> operator isa QOB.LazyTensor, op.operators) || return nothing
+
+    result = op.operators[1]
+    for i in 2:length(op.operators)
+        result = result * op.operators[i]
+        result isa QOB.LazyTensor || return nothing
+    end
+    return result
+end
+
+function _sciml_operator(op::QOB.LazyProduct)
+    combined = _combined_lazytensor_product(op)
+    combined === nothing || return op.factor * _sciml_operator(combined)
+
+    result = _sciml_operator(op.operators[1])
+    for i in 2:length(op.operators)
+        result = result * _sciml_operator(op.operators[i])
+    end
+    return op.factor * result
+end
+
+function _sciml_operator(op::QOB.LazyTensor)
+    T = _operator_number_type(op)
+    factors = map(1:length(op.basis_l.bases)) do index
+        op_index = findfirst(isequal(index), op.indices)
+        if op_index === nothing
+            return _matrix_identity(T, op.basis_l.bases[index], op.basis_r.bases[index])
+        end
+        return _sciml_operator(op.operators[op_index])
+    end
+
+    # QuantumOpticsBase stores tensor data as kron(right, left), which matches
+    # SciMLOperators when the factor order is reversed.
+    tensor_operator = SciMLOperators.TensorProductOperator(reverse(factors)...)
+    return op.factor * tensor_operator
+end
+
+function _sciml_operator(op::QOB.AbstractOperator)
+    return SciMLOperators.MatrixOperator(QOB.dense(op).data)
+end
+
+function QOB.sciml_lazy_operator(op::QOB.AbstractOperator; cache=nothing)
+    wrapped = _wrap(op.basis_l, op.basis_r, _sciml_operator(op))
+    cache === nothing && return wrapped
+    return QOB.cache_sciml_lazy_operator(wrapped, cache)
+end
+
+function QOB.sciml_lazy_operator(op::SciMLLazyOperator; cache=nothing)
+    cache === nothing && return op
+    return QOB.cache_sciml_lazy_operator(op, cache)
+end
+
+function QOB.cache_sciml_lazy_operator(op::SciMLLazyOperator, state::Ket)
+    return QOB.cache_sciml_lazy_operator(op, state.data)
+end
+
+function QOB.cache_sciml_lazy_operator(op::SciMLLazyOperator, data::AbstractVecOrMat)
+    return _wrap(op.basis_l, op.basis_r, SciMLOperators.cache_operator(op.operator, data))
+end
+
+function QOB.cache_sciml_lazy_operator(op::QOB.AbstractOperator, state)
+    return QOB.cache_sciml_lazy_operator(QOB.sciml_lazy_operator(op), state)
+end
+
+function QOB.dense(op::SciMLLazyOperator)
+    return QOB.DenseOperator(op.basis_l, op.basis_r, Matrix(op.operator))
+end
+
+function SparseArrays.sparse(op::SciMLLazyOperator)
+    return QOB.SparseOperator(op.basis_l, op.basis_r, SparseArrays.sparse(Matrix(op.operator)))
+end
+
+function -(op::SciMLLazyOperator)
+    return _wrap(op.basis_l, op.basis_r, -op.operator)
+end
+
+function +(a::SciMLLazyOperator, b::SciMLLazyOperator)
+    _check_same_bases(a, b)
+    return _wrap(a.basis_l, a.basis_r, a.operator + b.operator)
+end
+
+function -(a::SciMLLazyOperator, b::SciMLLazyOperator)
+    _check_same_bases(a, b)
+    return _wrap(a.basis_l, a.basis_r, a.operator - b.operator)
+end
+
+function +(a::SciMLLazyOperator, b::QOB.AbstractOperator)
+    return a + QOB.sciml_lazy_operator(b)
+end
+
+function +(a::QOB.AbstractOperator, b::SciMLLazyOperator)
+    return QOB.sciml_lazy_operator(a) + b
+end
+
+function -(a::SciMLLazyOperator, b::QOB.AbstractOperator)
+    return a - QOB.sciml_lazy_operator(b)
+end
+
+function -(a::QOB.AbstractOperator, b::SciMLLazyOperator)
+    return QOB.sciml_lazy_operator(a) - b
+end
+
+function *(a::SciMLLazyOperator, b::SciMLLazyOperator)
+    _check_multiplicable(a, b)
+    return _wrap(a.basis_l, b.basis_r, a.operator * b.operator)
+end
+
+function *(a::SciMLLazyOperator{B1,B2}, b::QOB.Operator{B2,B3,T}) where {B1,B2,B3,T}
+    return a * QOB.sciml_lazy_operator(b)
+end
+
+function *(a::QOB.Operator{B1,B2,T}, b::SciMLLazyOperator{B2,B3}) where {B1,B2,B3,T}
+    return QOB.sciml_lazy_operator(a) * b
+end
+
+function *(a::SciMLLazyOperator, b::QOB.AbstractOperator)
+    return a * QOB.sciml_lazy_operator(b)
+end
+
+function *(a::QOB.AbstractOperator, b::SciMLLazyOperator)
+    return QOB.sciml_lazy_operator(a) * b
+end
+
+*(a::Number, b::SciMLLazyOperator) = _wrap(b.basis_l, b.basis_r, a * b.operator)
+*(a::SciMLLazyOperator, b::Number) = _wrap(a.basis_l, a.basis_r, a.operator * b)
+/(a::SciMLLazyOperator, b::Number) = _wrap(a.basis_l, a.basis_r, a.operator / b)
+
+function QOB.dagger(op::SciMLLazyOperator)
+    return _wrap(op.basis_r, op.basis_l, adjoint(op.operator))
+end
+
+function transpose(op::SciMLLazyOperator)
+    return _wrap(op.basis_r, op.basis_l, transpose(op.operator))
+end
+
+function _operator_for_mul(op, data)
+    SciMLOperators.iscached(op) && return op
+    return SciMLOperators.cache_operator(op, data)
+end
+
+function mul!(result::Ket{B1}, op::SciMLLazyOperator{B1,B2}, state::Ket{B2}, alpha, beta) where {B1,B2}
+    sciml_op = _operator_for_mul(op.operator, state.data)
+    if isone(alpha) && iszero(beta)
+        mul!(result.data, sciml_op, state.data)
+        return result
+    end
+    mul!(result.data, sciml_op, state.data, alpha, beta)
+    return result
+end
+
+function mul!(result::Bra{B2}, state::Bra{B1}, op::SciMLLazyOperator{B1,B2}, alpha, beta) where {B1,B2}
+    sciml_op = _operator_for_mul(transpose(op.operator), state.data)
+    if isone(alpha) && iszero(beta)
+        mul!(result.data, sciml_op, state.data)
+        return result
+    end
+    mul!(result.data, sciml_op, state.data, alpha, beta)
+    return result
+end
+
+end # module