Lazy identities via FillArrays (#94)

* Have identityoperator return an Eye * Handle Eyes in LazyTensor mul!() * Fix identityoperator semantics. Add a test. * Simpify _tpops_tuple * Add compat entry for FillArrays * Filter out all square eyes. Simplify code. * Try to make Aqua happier * Explicit eye tests * Only turn off ambiguity checking for reshape * Test left action. * Restore more efficient handling of isometries * Remember why we need to distinguish beta==0 * More specific Aqua exclusion * Fix test dependencies. * Handle alpha=0 and beta=0 better. * Managed to munge tests in merge. --------- Co-authored-by: Ashley Milsted <[email protected]>
qojulia · May 4, 2023 · c0336cf · c0336cf
1 parent 65c0d0a
commit c0336cf
Show file tree

Hide file tree

Showing 12 changed files with 233 additions and 123 deletions.
diff --git a/Project.toml b/Project.toml
@@ -5,6 +5,7 @@ version = "0.3.12"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 QuantumInterface = "5717a53b-5d69-4fa3-b976-0bf2f97ca1e5"
@@ -16,6 +17,7 @@ UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 [compat]
 Adapt = "1, 2, 3.3"
 FFTW = "1.2"
+FillArrays = "0.13, 1"
 LRUCache = "1"
 QuantumInterface = "0.1.0"
 Strided = "1, 2"
@@ -27,6 +29,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -37,4 +40,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 
 [targets]
-test = ["LinearAlgebra", "SparseArrays", "Random", "Test", "Aqua", "JET", "Adapt", "Dates", "FFTW", "LRUCache", "Strided", "UnsafeArrays"]
+test = ["LinearAlgebra", "SparseArrays", "Random", "Test", "Aqua", "JET", "Adapt", "Dates", "FFTW", "LRUCache", "Strided", "UnsafeArrays", "FillArrays"]
diff --git a/src/QuantumOpticsBase.jl b/src/QuantumOpticsBase.jl
@@ -1,6 +1,6 @@
 module QuantumOpticsBase
 
-using SparseArrays, LinearAlgebra, LRUCache, Strided, UnsafeArrays
+using SparseArrays, LinearAlgebra, LRUCache, Strided, UnsafeArrays, FillArrays
 import LinearAlgebra: mul!, rmul!
 
 import QuantumInterface: dagger, directsum, ⊕, dm, embed, expect, permutesystems,
@@ -17,7 +17,7 @@ export Basis, GenericBasis, CompositeBasis, basis,
         #operators_dense
                 Operator, DenseOperator, DenseOpType, projector, dm,
         #operators_sparse
-                SparseOperator, diagonaloperator, SparseOpType,
+                SparseOperator, diagonaloperator, SparseOpType, EyeOpType,
         #operators_lazysum
                 LazySum,
         #operators_lazyproduct

diff --git a/src/operators.jl b/src/operators.jl
@@ -350,7 +350,6 @@ to be used in the identity matrix.
 identityoperator(::Type{T}, ::Type{S}, b1::Basis, b2::Basis) where {T<:AbstractOperator,S} = throw(ArgumentError("Identity operator not defined for operator type $T."))
 identityoperator(::Type{T}, ::Type{S}, b::Basis) where {T<:AbstractOperator,S} = identityoperator(T,S,b,b)
 identityoperator(::Type{T}, bases::Basis...) where T<:AbstractOperator = identityoperator(T,eltype(T),bases...)
-identityoperator(b::Basis) = identityoperator(b,b)
 identityoperator(op::T) where {T<:AbstractOperator} = identityoperator(T, op.basis_l, op.basis_r)
 
 # Catch case where eltype cannot be inferred from type; this is a bit hacky

diff --git a/src/operators_lazyproduct.jl b/src/operators_lazyproduct.jl
@@ -85,6 +85,7 @@ function tensor(a::LazyProduct{B1, B2, F, T, KTL, BTR},b::Operator{B3,B4}) where
 end
 
 function mul!(result::Ket{B1},a::LazyProduct{B1,B2},b::Ket{B2},alpha,beta) where {B1,B2}
+    iszero(alpha) && (_zero_op_mul!(result.data, beta); return result)
     if length(a.operators)==1
         mul!(result,a.operators[1],b,a.factor*alpha,beta)
     else
@@ -98,6 +99,7 @@ function mul!(result::Ket{B1},a::LazyProduct{B1,B2},b::Ket{B2},alpha,beta) where
 end
 
 function mul!(result::Bra{B2},a::Bra{B1},b::LazyProduct{B1,B2},alpha,beta) where {B1,B2}
+    iszero(alpha) && (_zero_op_mul!(result.data, beta); return result)
     if length(b.operators)==1
         mul!(result, a, b.operators[1],b.factor*alpha,beta)
     else
@@ -111,6 +113,7 @@ function mul!(result::Bra{B2},a::Bra{B1},b::LazyProduct{B1,B2},alpha,beta) where
 end
 
 function mul!(result::Operator{B1,B3,T},a::LazyProduct{B1,B2},b::Operator{B2,B3},alpha,beta) where {B1,B2,B3,T}
+    iszero(alpha) && (_zero_op_mul!(result.data, beta); return result)
     if length(a.operators) == 1
         mul!(result,a.operators[1],b,a.factor*alpha,beta)
     else
@@ -127,6 +130,7 @@ function mul!(result::Operator{B1,B3,T},a::LazyProduct{B1,B2},b::Operator{B2,B3}
 end
 
 function mul!(result::Operator{B1,B3,T},a::Operator{B1,B2},b::LazyProduct{B2,B3},alpha,beta) where {B1,B2,B3,T}
+    iszero(alpha) && (_zero_op_mul!(result.data, beta); return result)
     if length(b.operators) == 1
         mul!(result, a, b.operators[1],b.factor*alpha,beta)
     else

diff --git a/src/operators_lazysum.jl b/src/operators_lazysum.jl
@@ -159,11 +159,20 @@ function embed(basis_l::CompositeBasis, basis_r::CompositeBasis, index::Integer,
     LazySum(basis_l, basis_r, op.factors, ((embed(basis_l, basis_r, index, o) for o in op.operators)...,)) # dispatch to fast-path single-index `embed`
 end
 
+function _zero_op_mul!(data, beta)
+    if iszero(beta)
+        fill!(data, zero(eltype(data)))
+    elseif !isone(beta)
+        data .*= beta
+    end
+    return data
+end
+
 # Fast in-place multiplication
 function mul!(result::Ket{B1},a::LazySum{B1,B2},b::Ket{B2},alpha,beta) where {B1,B2}
-    if length(a.operators) == 0
+    if length(a.operators) == 0 || iszero(alpha)
         _check_mul!_dim_compatibility(size(result), size(a), size(b))
-        result.data .*= beta
+        _zero_op_mul!(result.data, beta)
     else
         mul!(result,a.operators[1],b,alpha*a.factors[1],beta)
         for i=2:length(a.operators)
@@ -174,9 +183,9 @@ function mul!(result::Ket{B1},a::LazySum{B1,B2},b::Ket{B2},alpha,beta) where {B1
 end
 
 function mul!(result::Bra{B2},a::Bra{B1},b::LazySum{B1,B2},alpha,beta) where {B1,B2}
-    if length(b.operators) == 0
+    if length(b.operators) == 0 || iszero(alpha)
         _check_mul!_dim_compatibility(size(result), reverse(size(b)), size(a))
-        result.data .*= beta
+        _zero_op_mul!(result.data, beta)
     else
         mul!(result,a,b.operators[1],alpha*b.factors[1],beta)
         for i=2:length(b.operators)
@@ -187,9 +196,9 @@ function mul!(result::Bra{B2},a::Bra{B1},b::LazySum{B1,B2},alpha,beta) where {B1
 end
 
 function mul!(result::Operator{B1,B3},a::LazySum{B1,B2},b::Operator{B2,B3},alpha,beta) where {B1,B2,B3}
-    if length(a.operators) == 0
+    if length(a.operators) == 0 || iszero(alpha)
         _check_mul!_dim_compatibility(size(result), size(a), size(b))
-        result.data .*= beta
+        _zero_op_mul!(result.data, beta)
     else
         mul!(result,a.operators[1],b,alpha*a.factors[1],beta)
         for i=2:length(a.operators)
@@ -199,9 +208,9 @@ function mul!(result::Operator{B1,B3},a::LazySum{B1,B2},b::Operator{B2,B3},alpha
     return result
 end
 function mul!(result::Operator{B1,B3},a::Operator{B1,B2},b::LazySum{B2,B3},alpha,beta) where {B1,B2,B3}
-    if length(b.operators) == 0
+    if length(b.operators) == 0 || iszero(alpha)
         _check_mul!_dim_compatibility(size(result), size(a), size(b))
-        result.data .*= beta
+        _zero_op_mul!(result.data, beta)
     else
         mul!(result,a,b.operators[1],alpha*b.factors[1],beta)
         for i=2:length(b.operators)