diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 2649dbbf..45ce9021 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -67,10 +67,13 @@ steps: - echo "Julia depot path $${JULIA_DEPOT_PATH}" - julia --project=$(mktemp -d) -e ' using Pkg; - pkg"dev ."; + Pkg.develop(PackageSpec(path=pwd())); Pkg.add("{{matrix.PACKAGE}}"); Pkg.build("{{matrix.PACKAGE}}"); - Pkg.test("{{matrix.PACKAGE}}");' + Pkg.test("{{matrix.PACKAGE}}"; test_fn=() -> begin + Pkg.add(PackageSpec(name="LoweredCodeUtils", version="3.5.3")); + Pkg.pin("LoweredCodeUtils"); + end);' matrix: setup: PACKAGE: ["QuantumOptics", "WaveguideQED", "QuantumSymbolics"] diff --git a/Project.toml b/Project.toml index 3bc5d145..fc911a82 100644 --- a/Project.toml +++ b/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.20" SparseArrays = "1" Strided = "1, 2" UnsafeArrays = "1" -julia = "1.10" \ No newline at end of file +julia = "1.10" diff --git a/benchmark/Project.toml b/benchmark/Project.toml new file mode 100644 index 00000000..acb21adf --- /dev/null +++ b/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" diff --git a/benchmark/sciml_lazy_operators.jl b/benchmark/sciml_lazy_operators.jl new file mode 100644 index 00000000..7424e52c --- /dev/null +++ b/benchmark/sciml_lazy_operators.jl @@ -0,0 +1,150 @@ +using BenchmarkTools +using LinearAlgebra +using QuantumOpticsBase +using Random +using SciMLOperators +using SparseArrays + +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, div(n, 2)) + return LazyProduct(LazyTensor(basis, mid, sigmax(site)), LazyTensor(basis, mid + 1, sigmaz(site))) +end + +function tensor_factors(site, storage) + factors = (sigmax(site), sigmaz(site)) + return storage == :dense ? dense.(factors) : sparse.(factors) +end + +function middle_sites(n) + left = max(2, div(n, 2)) + return (left, left + 1) +end + +function tensor_case(n; spin=1//2, storage=:sparse, sites=(1, n)) + site = SpinBasis(spin) + basis = tensor(ntuple(_ -> site, n)...) + return LazyTensor(basis, collect(sites), tensor_factors(site, storage)) +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 + +function add_tensor_benchmarks!(n; spin=1//2, storage=:sparse, sites=(1, n), seed=40 + n) + op = tensor_case(n; spin=spin, storage=storage, sites=sites) + site_label = sites == (1, n) ? "edge" : "middle" + spin_label = spin == 1//2 ? "spin=1/2" : "spin=$spin" + name = "$(String(storage)) $site_label $spin_label n=$n" + return add_apply_benchmarks!(SUITE["lazy_tensor"], name, op, random_ket(op.basis_r, seed)) +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) + for storage in (:sparse, :dense) + add_tensor_benchmarks!(n; storage=storage, sites=(1, n), seed=40 + n + (storage == :dense ? 100 : 0)) + add_tensor_benchmarks!(n; storage=storage, sites=middle_sites(n), seed=80 + n + (storage == :dense ? 100 : 0)) + end +end + +for n in (3, 4) + for storage in (:sparse, :dense) + add_tensor_benchmarks!(n; spin=1, storage=storage, sites=(1, n), seed=120 + n + (storage == :dense ? 100 : 0)) + end +end + +for storage in (:sparse, :dense) + add_tensor_benchmarks!(4; spin=1, storage=storage, sites=middle_sites(4), seed=160 + (storage == :dense ? 100 : 0)) +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 diff --git a/ext/QuantumOpticsBaseSciMLOperatorsExt.jl b/ext/QuantumOpticsBaseSciMLOperatorsExt.jl new file mode 100644 index 00000000..d994ee97 --- /dev/null +++ b/ext/QuantumOpticsBaseSciMLOperatorsExt.jl @@ -0,0 +1,255 @@ +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::Int) = size(op)[dim] +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.LazyOperator) + return a + QOB.sciml_lazy_operator(b) +end + +function +(a::QOB.LazyOperator, 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::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.LazyOperator) + return a - QOB.sciml_lazy_operator(b) +end + +function -(a::QOB.LazyOperator, 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, b::QOB.LazyOperator) + return a * QOB.sciml_lazy_operator(b) +end + +function *(a::QOB.LazyOperator, b::SciMLLazyOperator) + return QOB.sciml_lazy_operator(a) * b +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 diff --git a/src/QuantumOpticsBase.jl b/src/QuantumOpticsBase.jl index f289952e..87c50ea8 100644 --- a/src/QuantumOpticsBase.jl +++ b/src/QuantumOpticsBase.jl @@ -32,6 +32,7 @@ export Basis, GenericBasis, CompositeBasis, basis, #operators_lazytensor LazyTensor, lazytensor_use_cache, lazytensor_clear_cache, lazytensor_cachesize, lazytensor_disable_cache, lazytensor_enable_cache, + sciml_lazy_operator, cache_sciml_lazy_operator, #states_lazyket LazyKet, #time_dependent_operators @@ -107,4 +108,12 @@ include("printing.jl") include("apply.jl") include("visualization.jl") +function sciml_lazy_operator(args...; kwargs...) + throw(ArgumentError("SciMLOperators-backed lazy operators require SciMLOperators.jl. Load it with `using SciMLOperators` before calling `sciml_lazy_operator`.")) +end + +function cache_sciml_lazy_operator(args...; kwargs...) + throw(ArgumentError("SciMLOperators-backed lazy operators require SciMLOperators.jl. Load it with `using SciMLOperators` before calling `cache_sciml_lazy_operator`.")) +end + end # module diff --git a/test/Project.toml b/test/Project.toml index 7b9df2fa..3ce2073f 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -7,13 +7,20 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LoweredCodeUtils = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b" +Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +OrdinaryDiffEqLowOrderRK = "1344f307-1e59-4825-a18e-ace9aa3fa4c6" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" QuantumInterface = "5717a53b-5d69-4fa3-b976-0bf2f97ca1e5" QuantumOptics = "6e0679c1-51ea-5a7c-ac74-d61b76210b0c" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a" -UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6" \ No newline at end of file +UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6" + +[compat] +LoweredCodeUtils = "=3.5.3" diff --git a/test/test_sciml_broadcast_interfaces.jl b/test/test_sciml_broadcast_interfaces.jl index 411b05e5..e06101b4 100644 --- a/test/test_sciml_broadcast_interfaces.jl +++ b/test/test_sciml_broadcast_interfaces.jl @@ -2,6 +2,7 @@ using Test using QuantumOptics using OrdinaryDiffEq +using OrdinaryDiffEqLowOrderRK: DP5 @testset "sciml interface" begin @@ -22,7 +23,7 @@ prob_data! = ODEProblem(schrod_data!, u0, (t₀, t₁)) sol = solve(prob!, DP5(); reltol = 1.0e-8, abstol = 1.0e-10, save_everystep=false) sol_data = solve(prob_data!, DP5(); reltol = 1.0e-8, abstol = 1.0e-10, save_everystep=false) -@test sol[end].data ≈ sol_data[end] +@test sol.u[end].data ≈ sol_data.u[end] # dense operator ODE problem σ₋ = sigmam(ℋ) @@ -59,7 +60,7 @@ prob_data! = ODEProblem(lind_data!, m0, (t₀, t₁)) sol = solve(prob!, DP5(); reltol = 1.0e-8, abstol = 1.0e-10, save_everystep=false) sol_data = solve(prob_data!, DP5(); reltol = 1.0e-8, abstol = 1.0e-10, save_everystep=false) -@test sol[end].data ≈ sol_data[end] +@test sol.u[end].data ≈ sol_data.u[end] end -end \ No newline at end of file +end diff --git a/test/test_sciml_lazy_operators.jl b/test/test_sciml_lazy_operators.jl new file mode 100644 index 00000000..b1d69ee2 --- /dev/null +++ b/test/test_sciml_lazy_operators.jl @@ -0,0 +1,134 @@ +@testitem "test_sciml_lazy_operators" begin +using Test +using LinearAlgebra, Random, SparseArrays +using QuantumOpticsBase +using SciMLOperators + +@testset "SciMLOperators-backed lazy operators" begin + Random.seed!(522) + + state_distance(a::StateVector, b::StateVector) = norm(a - b) + op_distance(a::AbstractOperator, b::AbstractOperator) = abs(tracedistance_nh(dense(a), dense(b))) + + b0 = SpinBasis(1//2) + b = tensor(b0, b0, b0, b0) + sx = sigmax(b0) + sz = sigmaz(b0) + + current = LazySum( + LazyTensor(b, 1, sx), + LazyTensor(b, 2, sz), + LazyProduct(LazyTensor(b, 3, sx), LazyTensor(b, 4, sz)), + ) + sciml = sciml_lazy_operator(current) + psi = Ket(b, randn(ComplexF64, length(b))) + + @test op_distance(sciml, current) < 1e-12 + @test state_distance(sciml * psi, current * psi) < 1e-12 + + out_current = Ket(b) + out_sciml = Ket(b) + mul!(out_current, current, psi, 0.25 - 0.5im, 0) + mul!(out_sciml, sciml, psi, 0.25 - 0.5im, 0) + @test state_distance(out_sciml, out_current) < 1e-12 + + cached = cache_sciml_lazy_operator(sciml, psi) + out_cached = Ket(b) + mul!(out_cached, cached, psi, 1, 0) + @test state_distance(out_cached, current * psi) < 1e-12 + + a = LazyTensor(b, 1, sx) + c = LazyTensor(b, 2, sz) + a_sciml = sciml_lazy_operator(a) + c_sciml = sciml_lazy_operator(c) + + @test size(a_sciml) == size(a) + @test size(a_sciml, 1) == size(a, 1) + @test size(a_sciml, 2) == size(a, 2) + @test eltype(a_sciml) <: Number + @test op_distance(copy(a_sciml), a) < 1e-12 + @test op_distance(sparse(a_sciml), sparse(a)) < 1e-12 + + @test sciml_lazy_operator(a_sciml) === a_sciml + @test op_distance(sciml_lazy_operator(a_sciml; cache=psi), a) < 1e-12 + @test op_distance(sciml_lazy_operator(a; cache=psi), a) < 1e-12 + @test op_distance(cache_sciml_lazy_operator(a, psi), a) < 1e-12 + + @test op_distance(2a_sciml + c_sciml / 3, 2a + c / 3) < 1e-12 + @test op_distance(a_sciml + c, a + c) < 1e-12 + @test op_distance(a + c_sciml, a + c) < 1e-12 + @test op_distance(a_sciml - c, a - c) < 1e-12 + @test op_distance(a - c_sciml, a - c) < 1e-12 + @test op_distance(a_sciml + dense(c), a + dense(c)) < 1e-12 + @test op_distance(dense(a) + c_sciml, dense(a) + c) < 1e-12 + @test op_distance(a_sciml - dense(c), a - dense(c)) < 1e-12 + @test op_distance(dense(a) - c_sciml, dense(a) - c) < 1e-12 + @test op_distance(-a_sciml, -a) < 1e-12 + @test op_distance(a_sciml * 2, a * 2) < 1e-12 + @test op_distance(2 * a_sciml, 2 * a) < 1e-12 + @test op_distance(a_sciml * c_sciml, a * c) < 1e-12 + @test op_distance(a_sciml * c, a * c) < 1e-12 + @test op_distance(a * c_sciml, a * c) < 1e-12 + @test op_distance(a_sciml * dense(c), a * dense(c)) < 1e-12 + @test op_distance(dense(a) * c_sciml, dense(a) * c) < 1e-12 + @test op_distance(dagger(a_sciml), dagger(a)) < 1e-12 + @test isapprox(dense(transpose(a_sciml)).data, transpose(dense(a).data)) + @test op_distance(dense(cached), current) < 1e-12 + + empty_sum = LazySum(ComplexF64, b, b) + @test op_distance(sciml_lazy_operator(empty_sum), empty_sum) < 1e-12 + + three_product = LazyProduct( + LazyTensor(b, 1, sx), + LazyTensor(b, 2, sz), + LazyTensor(b, 3, sx), + ) + @test op_distance(sciml_lazy_operator(three_product), three_product) < 1e-12 + @test op_distance(sciml_lazy_operator(LazyProduct(dense(a), dense(c))), dense(a) * dense(c)) < 1e-12 + + direct_sum = LazyDirectSum(sx, sz) + direct_sum_sciml = sciml_lazy_operator(direct_sum) + @test op_distance(direct_sum_sciml, direct_sum) < 1e-12 + @test op_distance(direct_sum_sciml + direct_sum, direct_sum + direct_sum) < 1e-12 + @test op_distance(direct_sum + direct_sum_sciml, direct_sum + direct_sum) < 1e-12 + @test op_distance(direct_sum_sciml - direct_sum, dense(direct_sum) - dense(direct_sum)) < 1e-12 + @test op_distance(direct_sum - direct_sum_sciml, dense(direct_sum) - dense(direct_sum)) < 1e-12 + @test op_distance(direct_sum_sciml * direct_sum, direct_sum * direct_sum) < 1e-12 + @test op_distance(direct_sum * direct_sum_sciml, direct_sum * direct_sum) < 1e-12 + + bra = dagger(psi) + bra_current = Bra(b) + bra_sciml = Bra(b) + mul!(bra_current, bra, current, 0.5 + 0.25im, 0) + mul!(bra_sciml, bra, sciml, 0.5 + 0.25im, 0) + @test norm(bra_sciml.data - bra_current.data) < 1e-12 + + bra_fast_current = Bra(b) + bra_fast_sciml = Bra(b) + mul!(bra_fast_current, bra, current, 1, 0) + mul!(bra_fast_sciml, bra, sciml, 1, 0) + @test norm(bra_fast_sciml.data - bra_fast_current.data) < 1e-12 + + b_other = SpinBasis(1) + other_sciml = sciml_lazy_operator(sigmax(b_other)) + @test_throws QuantumOpticsBase.IncompatibleBases a_sciml + other_sciml + @test_throws QuantumOpticsBase.IncompatibleBases a_sciml * other_sciml + + ext = Base.get_extension(QuantumOpticsBase, :QuantumOpticsBaseSciMLOperatorsExt) + small_operator = SciMLOperators.MatrixOperator(zeros(ComplexF64, 2, 2)) + @test_throws DimensionMismatch ext.SciMLLazyOperator(b, b, small_operator) + + b1l = GenericBasis(2) + b1r = GenericBasis(4) + b2 = GenericBasis(3) + bl = tensor(b1l, b2) + br = tensor(b1r, b2) + local_op = randoperator(b2) + rectangular = LazyTensor(bl, br, [2], (local_op,), 0.7) + rectangular_sciml = sciml_lazy_operator(rectangular) + psi_rectangular = Ket(br, randn(ComplexF64, length(br))) + + @test op_distance(rectangular_sciml, rectangular) < 1e-12 + @test state_distance(rectangular_sciml * psi_rectangular, rectangular * psi_rectangular) < 1e-12 +end +end