From 61249ae0f47c5164154228faf2cfcb63665f2528 Mon Sep 17 00:00:00 2001 From: Mahfuza Humayra Mohona Date: Tue, 16 Jun 2026 19:33:51 +0600 Subject: [PATCH] Add embed_lazy functionality for CompositeBasis in QuantumOpticsBase --- docs/src/api.md | 4 ++ src/QuantumOpticsBase.jl | 2 +- src/operators_lazyproduct.jl | 5 ++ src/operators_lazytensor.jl | 126 +++++++++++++++++++++++++++++++++ src/time_dependent_operator.jl | 14 ++++ test/test_embed_lazy.jl | 117 ++++++++++++++++++++++++++++++ 6 files changed, 267 insertions(+), 1 deletion(-) create mode 100644 test/test_embed_lazy.jl diff --git a/docs/src/api.md b/docs/src/api.md index 05d062b6..1a568c78 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -220,6 +220,10 @@ variance embed ``` +```@docs +embed_lazy +``` + ```@docs permutesystems ``` diff --git a/src/QuantumOpticsBase.jl b/src/QuantumOpticsBase.jl index f289952e..8733705b 100644 --- a/src/QuantumOpticsBase.jl +++ b/src/QuantumOpticsBase.jl @@ -20,7 +20,7 @@ export Basis, GenericBasis, CompositeBasis, basis, dagger, normalize, normalize!, #operators AbstractOperator, DataOperator, expect, variance, - identityoperator, ptrace, reduced, embed, dense, tr, sparse, + identityoperator, ptrace, reduced, embed, embed_lazy, dense, tr, sparse, #operators_dense Operator, DenseOperator, DenseOpType, projector, dm, #operators_sparse diff --git a/src/operators_lazyproduct.jl b/src/operators_lazyproduct.jl index d5843aea..5356636c 100644 --- a/src/operators_lazyproduct.jl +++ b/src/operators_lazyproduct.jl @@ -148,3 +148,8 @@ end # GPU adaptation Adapt.adapt_structure(to, x::LazyProduct) = LazyProduct([Adapt.adapt(to, op) for op in x.operators], x.factor) + +embed_lazy(::CompositeBasis, ::CompositeBasis, ::Any, ::LazyProduct) = + throw(ArgumentError("embed_lazy is not supported for LazyProduct")) +embed_lazy(::CompositeBasis, ::Any, ::LazyProduct) = + throw(ArgumentError("embed_lazy is not supported for LazyProduct")) diff --git a/src/operators_lazytensor.jl b/src/operators_lazytensor.jl index 9ee029b4..7e9a50a2 100644 --- a/src/operators_lazytensor.jl +++ b/src/operators_lazytensor.jl @@ -219,6 +219,132 @@ end identityoperator(::Type{<:LazyTensor}, ::Type{T}, b1::Basis, b2::Basis) where {T<:Number} = LazyTensor(b1, b2, Int[], Tuple{}(), one(T)) +""" + embed_lazy(basis[, basis_r], indices, op) + embed_lazy(basis[, basis_r], indices, operators) + +Embed an operator into a larger composite basis without materializing the full matrix. + +Returns a [`LazyTensor`](@ref) for local [`DataOperator`](@ref) terms, preserves +[`LazySum`](@ref) structure, and re-embeds existing [`LazyTensor`](@ref) objects. + +Unlike [`embed`](@ref), dense and sparse operators are not converted to large embedded +sparse matrices. Unlike directly calling [`LazyTensor`](@ref), this function follows +[`embed`](@ref) index semantics (including basis compatibility checks) and also handles +[`LazySum`](@ref) and [`TimeDependentSum`](@ref). + +Note that `sx + 0.5*sz` on a shared basis returns an eager [`Operator`](@ref), not a +[`LazySum`](@ref); pass an explicit `LazySum` if you need a `LazySum` of embedded terms. + +[`LazyProduct`](@ref) is not supported. A single operator acting jointly on multiple +subsystems cannot be embedded lazily; use [`embed`](@ref) or pass local operators / +a [`LazyTensor`](@ref) with explicit structure instead. + +# Examples +```jldoctest +julia> using QuantumOpticsBase + +julia> b0 = SpinBasis(1//2); + +julia> b = tensor(b0, b0, b0); + +julia> sx = sigmax(b0); + +julia> op_lazy = embed_lazy(b, 2, sx); + +julia> op_lazy isa LazyTensor +true + +julia> dense(op_lazy) == dense(embed(b, 2, sx)) +true + +julia> embed_lazy(b, 2, LazySum(sx, 0.5 * sigmaz(b0))) isa LazySum +true +``` +""" +embed_lazy(basis::CompositeBasis, indices, op) = embed_lazy(basis, basis, indices, op) +embed_lazy(basis::CompositeBasis, indices, operators::AbstractVector{<:AbstractOperator}) = + embed_lazy(basis, basis, indices, operators) + +function _embed_lazy_check_composite(basis_l::CompositeBasis, basis_r::CompositeBasis) + N = length(basis_l.bases) + @assert N == length(basis_r.bases) + return N +end + +function _embed_lazy_check_single(basis_l::CompositeBasis, basis_r::CompositeBasis, index::Integer, op::AbstractOperator) + N = _embed_lazy_check_composite(basis_l, basis_r) + basis_l.bases[index] == op.basis_l || throw(IncompatibleBases()) + basis_r.bases[index] == op.basis_r || throw(IncompatibleBases()) + check_indices(N, index) +end + +function _embed_lazy_check_target_indices(N::Integer, target_indices) + check_sortedindices(N, target_indices) +end + +function _embed_lazy_ops(basis_l, basis_r, indices, ops::Tuple) + Tuple(embed_lazy(basis_l, basis_r, indices, o) for o in ops) +end +function _embed_lazy_ops(basis_l, basis_r, indices, ops) + [embed_lazy(basis_l, basis_r, indices, o) for o in ops] +end + +function embed_lazy(basis_l::CompositeBasis, basis_r::CompositeBasis, index::Integer, op::LazySum) + LazySum(basis_l, basis_r, op.factors, _embed_lazy_ops(basis_l, basis_r, index, op.operators)) +end + +function embed_lazy(basis_l::CompositeBasis, basis_r::CompositeBasis, + indices::AbstractVector, op::LazySum) + LazySum(basis_l, basis_r, op.factors, _embed_lazy_ops(basis_l, basis_r, indices, op.operators)) +end + +function embed_lazy(basis_l::CompositeBasis, basis_r::CompositeBasis, target_indices, op::LazyTensor) + target_indices = collect(target_indices) + N = _embed_lazy_check_composite(basis_l, basis_r) + length(target_indices) == length(op.basis_l.bases) || + throw(ArgumentError("embed_lazy target indices must match the number of subsystems of the LazyTensor")) + _embed_lazy_check_target_indices(N, target_indices) + sub_bl = reduce(tensor, basis_l.bases[target_indices]) + sub_br = reduce(tensor, basis_r.bases[target_indices]) + sub_bl == op.basis_l || throw(IncompatibleBases()) + sub_br == op.basis_r || throw(IncompatibleBases()) + global_indices = [target_indices[i] for i in op.indices] + LazyTensor(basis_l, basis_r, global_indices, op.operators, op.factor) +end + +function embed_lazy(basis_l::CompositeBasis, basis_r::CompositeBasis, index::Integer, op::LazyTensor) + length(op.basis_l.bases) == 1 || throw(ArgumentError("cannot embed_lazy a multi-subsystem LazyTensor at a single index")) + subop = only(op.operators) + _embed_lazy_check_single(basis_l, basis_r, index, subop) + LazyTensor(basis_l, basis_r, index, subop, op.factor) +end + +function embed_lazy(basis_l::CompositeBasis, basis_r::CompositeBasis, + indices::AbstractVector{<:Integer}, operators::AbstractVector{<:AbstractOperator}) + N = _embed_lazy_check_composite(basis_l, basis_r) + length(indices) == length(operators) || + throw(ArgumentError("Must have length(indices) == length(operators) in embed_lazy")) + _embed_lazy_check_target_indices(N, indices) + for (idx, o) in zip(indices, operators) + basis_l.bases[idx] == o.basis_l || throw(IncompatibleBases()) + basis_r.bases[idx] == o.basis_r || throw(IncompatibleBases()) + end + LazyTensor(basis_l, basis_r, collect(indices), Tuple(operators)) +end + +function embed_lazy(basis_l::CompositeBasis, basis_r::CompositeBasis, index::Integer, op::DataOperator) + _embed_lazy_check_single(basis_l, basis_r, index, op) + LazyTensor(basis_l, basis_r, index, op) +end + +function embed_lazy(basis_l::CompositeBasis, basis_r::CompositeBasis, + indices::AbstractVector{<:Integer}, op::DataOperator) + length(indices) == 1 || + throw(ArgumentError("embed_lazy cannot lazily embed a joint operator on multiple subsystems; use embed, pass a vector of local operators, or use a LazyTensor/LazySum with explicit structure")) + embed_lazy(basis_l, basis_r, only(indices), op) +end + ## LazyTensor global cache function lazytensor_default_cache_size() diff --git a/src/time_dependent_operator.jl b/src/time_dependent_operator.jl index 4f4beb5c..56bdf2a2 100644 --- a/src/time_dependent_operator.jl +++ b/src/time_dependent_operator.jl @@ -210,6 +210,20 @@ function embed(basis_l::CompositeBasis, basis_r::CompositeBasis, indices, o::Tim TimeDependentSum(coefficients(o), embed(basis_l, basis_r, indices, static_operator(o)), o.current_time) end +function embed_lazy(basis_l::CompositeBasis, basis_r::CompositeBasis, index::Integer, o::TimeDependentSum) + TimeDependentSum(coefficients(o), embed_lazy(basis_l, basis_r, index, static_operator(o)), current_time(o)) +end + +function embed_lazy(basis_l::CompositeBasis, basis_r::CompositeBasis, + indices::AbstractVector, o::TimeDependentSum) + TimeDependentSum(coefficients(o), embed_lazy(basis_l, basis_r, indices, static_operator(o)), current_time(o)) +end + +embed_lazy(basis::CompositeBasis, index::Integer, o::TimeDependentSum) = + embed_lazy(basis, basis, index, o) +embed_lazy(basis::CompositeBasis, indices::AbstractVector, o::TimeDependentSum) = + embed_lazy(basis, basis, indices, o) + function +(A::TimeDependentSum, B::TimeDependentSum) _check_same_time(A, B) TimeDependentSum( diff --git a/test/test_embed_lazy.jl b/test/test_embed_lazy.jl new file mode 100644 index 00000000..84d20789 --- /dev/null +++ b/test/test_embed_lazy.jl @@ -0,0 +1,117 @@ +@testitem "test_embed_lazy" begin +using Test +using QuantumOpticsBase +using LinearAlgebra, SparseArrays, Random + +@testset "embed_lazy" begin + +Random.seed!(0) + +b0 = SpinBasis(1//2) +b = tensor(b0, b0, b0, b0) + +sx = sigmax(b0) +sz = sigmaz(b0) + +psi = Ket(b, randn(ComplexF64, length(b))) + +@testset "single-site DataOperator" begin + op_eager = embed(b, 2, sx) + op_lazy = embed_lazy(b, 2, sx) + + @test op_lazy isa LazyTensor + @test dense(op_lazy) == dense(op_eager) + @test op_lazy * psi == op_eager * psi + + @test embed_lazy(b, 2, dense(sx)) isa LazyTensor + @test embed_lazy(b, 2, sparse(sx)) isa LazyTensor + @test dense(embed_lazy(b, 2, sparse(sx))) == dense(embed(b, 2, sparse(sx))) + @test dense(embed_lazy(b, [2], sx)) == dense(embed(b, 2, sx)) +end + +@testset "LazySum" begin + local_h = LazySum(sx, 0.5 * sz) + lazy_h = embed_lazy(b, 3, local_h) + eager_h = embed(b, 3, local_h) + + @test lazy_h isa LazySum + @test all(t -> t isa LazyTensor, lazy_h.operators) + @test lazy_h.factors == local_h.factors + @test dense(lazy_h) == dense(eager_h) + @test lazy_h * psi == eager_h * psi + + H = sum(embed_lazy(b, i, local_h) for i in 1:4) + H_eager = sum(embed(b, i, local_h) for i in 1:4) + @test H isa LazySum + @test dense(H) ≈ dense(H_eager) +end + +@testset "multi-site vector of operators" begin + op_lazy12 = embed_lazy(b, [1, 3], [sx, sz]) + op_eager12 = embed(b, [1, 3], [sx, sz]) + @test op_lazy12 isa LazyTensor + @test dense(op_lazy12) == dense(op_eager12) + @test op_lazy12 * psi == op_eager12 * psi +end + +@testset "asymmetric bases" begin + bf = FockBasis(2) + b_l = b0 ⊗ bf + b_r = bf ⊗ b0 + a = randoperator(bf, b0) + op_asym = embed_lazy(b_l, b_r, 2, a) + @test op_asym isa LazyTensor + @test dense(op_asym) == dense(embed(b_l, b_r, 2, a)) +end + +@testset "LazyTensor re-embedding" begin + b1a = GenericBasis(2) + b2a = GenericBasis(1) + b3a = GenericBasis(6) + b1b = GenericBasis(3) + b2b = GenericBasis(4) + b3b = GenericBasis(5) + b_l = b1a ⊗ b2a ⊗ b3a + b_r = b1b ⊗ b2b ⊗ b3b + op1 = randoperator(b1a, b1b) + op3 = randoperator(b3a, b3b) + x = LazyTensor(b_l, b_r, [1, 3], (op1, sparse(op3)), 0.3) + x_sub = LazyTensor(b1a ⊗ b3a, b1b ⊗ b3b, [1, 2], (op1, sparse(op3)), 0.3) + @test embed_lazy(b_l, b_r, [1, 3], x_sub) == x + + bsub = b0 ⊗ b0 + lt_partial = LazyTensor(bsub, [2], (sz,), 2.0) + lz = embed_lazy(b, [1, 3], lt_partial) + @test lz isa LazyTensor + @test lz.indices == [3] + @test lz.factor == 2.0 + @test dense(lz) ≈ dense(2.0 * embed(b, 3, sz)) +end + +@testset "TimeDependentSum" begin + subop = randoperator(b0) + td_op = TimeDependentSum(ComplexF64, (t -> cos(t)) => subop) + td_lazy = embed_lazy(b, 2, td_op) + @test td_lazy isa TimeDependentSum + @test dense(td_lazy) == dense(embed(b, 2, td_op)) + @test all(t -> t isa LazyTensor, static_operator(td_lazy).operators) +end + +@testset "errors" begin + @test_throws QuantumOpticsBase.IncompatibleBases embed_lazy(b, 1, destroy(FockBasis(2))) + @test_throws ArgumentError embed_lazy(b, 1, LazyProduct(sx)) + @test_throws ArgumentError embed_lazy(b, [1, 2], sx ⊗ sz) + @test_throws ArgumentError embed_lazy(b, [3, 1], [sz, sx]) + @test_throws ArgumentError embed_lazy(b, [3, 1], LazyTensor(b0 ⊗ b0, [1, 2], (sx, sz))) +end + +@testset "allocation" begin + b_chain = reduce(⊗, ntuple(_ -> SpinBasis(1//2), 12)) + alloc_lazy = @allocated embed_lazy(b_chain, 4, sx) + alloc_eager = @allocated embed(b_chain, 4, sx) + @test alloc_lazy < 10_000 + @test alloc_lazy < alloc_eager ÷ 10 +end + +end # testset +end