Add embed_lazy functionality for CompositeBasis in QuantumOpticsBase

docs/src/api.md

@@ -220,6 +220,10 @@ variance
 embed
 ```
 
+```@docs
+embed_lazy
+```
+
 ```@docs
 permutesystems
 ```

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

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"))

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()

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(

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