diff --git a/Project.toml b/Project.toml index 3bb147f0..3e662d46 100644 --- a/Project.toml +++ b/Project.toml @@ -26,6 +26,7 @@ SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" [extensions] QuantumOpticsBaseMakieExt = ["Makie"] +QuantumOpticsBaseSciMLOperatorsExt = ["SciMLOperators"] [compat] Adapt = "4.1" diff --git a/benchmark/Project.toml b/benchmark/Project.toml new file mode 100644 index 00000000..49efcfbc --- /dev/null +++ b/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" diff --git a/benchmark/sciml_lazyoperators.jl b/benchmark/sciml_lazyoperators.jl new file mode 100644 index 00000000..b584f348 --- /dev/null +++ b/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 diff --git a/ext/QuantumOpticsBaseSciMLOperatorsExt.jl b/ext/QuantumOpticsBaseSciMLOperatorsExt.jl new file mode 100644 index 00000000..a50767f8 --- /dev/null +++ b/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 diff --git a/src/QuantumOpticsBase.jl b/src/QuantumOpticsBase.jl index f289952e..bccded3f 100644 --- a/src/QuantumOpticsBase.jl +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index e30c1813..5e387554 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/runtests.jl b/test/runtests.jl index ad6776ba..e94f825c 100644 --- a/test/runtests.jl +++ b/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 diff --git a/test/test_sciml_lazyoperators.jl b/test/test_sciml_lazyoperators.jl new file mode 100644 index 00000000..6e42934a --- /dev/null +++ b/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