changeset 943:fb060e98ac0a feature/tensormapping_application_promotion

Remove more type assertions
author Jonatan Werpers <jonatan@werpers.com>
date Thu, 10 Mar 2022 16:57:01 +0100
parents 7829c09f8137
children 4a4ef4bf6cb9
files src/LazyTensors/lazy_tensor_operations.jl src/LazyTensors/tensor_mapping.jl test/LazyTensors/lazy_tensor_operations_test.jl
diffstat 3 files changed, 33 insertions(+), 14 deletions(-) [+]
line wrap: on
line diff
--- a/src/LazyTensors/lazy_tensor_operations.jl	Thu Mar 10 11:13:34 2022 +0100
+++ b/src/LazyTensors/lazy_tensor_operations.jl	Thu Mar 10 16:57:01 2022 +0100
@@ -52,8 +52,8 @@
 Base.adjoint(tm::TensorMapping) = LazyTensorMappingTranspose(tm)
 Base.adjoint(tmt::LazyTensorMappingTranspose) = tmt.tm
 
-apply(tmt::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{T,R}, I::Vararg{Any,D}) where {T,R,D} = apply_transpose(tmt.tm, v, I...)
-apply_transpose(tmt::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{T,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmt.tm, v, I...)
+apply(tmt::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D} = apply_transpose(tmt.tm, v, I...)
+apply_transpose(tmt::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmt.tm, v, I...)
 
 range_size(tmt::LazyTensorMappingTranspose) = domain_size(tmt.tm)
 domain_size(tmt::LazyTensorMappingTranspose) = range_size(tmt.tm)
@@ -69,8 +69,8 @@
 end
 # TODO: Boundschecking in constructor.
 
-apply(tmBinOp::LazyTensorMappingBinaryOperation{:+,T,R,D}, v::AbstractArray{T,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) + apply(tmBinOp.tm2, v, I...)
-apply(tmBinOp::LazyTensorMappingBinaryOperation{:-,T,R,D}, v::AbstractArray{T,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) - apply(tmBinOp.tm2, v, I...)
+apply(tmBinOp::LazyTensorMappingBinaryOperation{:+,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) + apply(tmBinOp.tm2, v, I...)
+apply(tmBinOp::LazyTensorMappingBinaryOperation{:-,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) - apply(tmBinOp.tm2, v, I...)
 
 range_size(tmBinOp::LazyTensorMappingBinaryOperation) = range_size(tmBinOp.tm1)
 domain_size(tmBinOp::LazyTensorMappingBinaryOperation) = domain_size(tmBinOp.tm1)
@@ -97,11 +97,11 @@
 range_size(tm::TensorMappingComposition) = range_size(tm.t1)
 domain_size(tm::TensorMappingComposition) = domain_size(tm.t2)
 
-function apply(c::TensorMappingComposition{T,R,K,D}, v::AbstractArray{T,D}, I::Vararg{Any,R}) where {T,R,K,D}
+function apply(c::TensorMappingComposition{T,R,K,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,K,D}
     apply(c.t1, c.t2*v, I...)
 end
 
-function apply_transpose(c::TensorMappingComposition{T,R,K,D}, v::AbstractArray{T,R}, I::Vararg{Any,D}) where {T,R,K,D}
+function apply_transpose(c::TensorMappingComposition{T,R,K,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,K,D}
     apply_transpose(c.t2, c.t1'*v, I...)
 end
 
@@ -134,7 +134,7 @@
 range_size(llm::LazyLinearMap) = size(llm.A)[[llm.range_indicies...]]
 domain_size(llm::LazyLinearMap) = size(llm.A)[[llm.domain_indicies...]]
 
-function apply(llm::LazyLinearMap{T,R,D}, v::AbstractArray{T,D}, I::Vararg{Any,R}) where {T,R,D}
+function apply(llm::LazyLinearMap{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D}
     view_index = ntuple(i->:,ndims(llm.A))
     for i ∈ 1:R
         view_index = Base.setindex(view_index, Int(I[i]), llm.range_indicies[i])
@@ -143,7 +143,7 @@
     return sum(A_view.*v)
 end
 
-function apply_transpose(llm::LazyLinearMap{T,R,D}, v::AbstractArray{T,R}, I::Vararg{Any,D}) where {T,R,D}
+function apply_transpose(llm::LazyLinearMap{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D}
     apply(LazyLinearMap(llm.A, llm.domain_indicies, llm.range_indicies), v, I...)
 end
 
@@ -166,8 +166,8 @@
 range_size(tmi::IdentityMapping) = tmi.size
 domain_size(tmi::IdentityMapping) = tmi.size
 
-apply(tmi::IdentityMapping{T,D}, v::AbstractArray{T,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
-apply_transpose(tmi::IdentityMapping{T,D}, v::AbstractArray{T,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
+apply(tmi::IdentityMapping{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
+apply_transpose(tmi::IdentityMapping{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
 
 """
     Base.:∘(tm, tmi)
@@ -260,7 +260,7 @@
     )
 end
 
-function apply(itm::InflatedTensorMapping{T,R,D}, v::AbstractArray{T,D}, I::Vararg{Any,R}) where {T,R,D}
+function apply(itm::InflatedTensorMapping{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D}
     dim_before = range_dim(itm.before)
     dim_domain = domain_dim(itm.tm)
     dim_range = range_dim(itm.tm)
@@ -272,7 +272,7 @@
     return apply(itm.tm, v_inner, inner_index...)
 end
 
-function apply_transpose(itm::InflatedTensorMapping{T,R,D}, v::AbstractArray{T,R}, I::Vararg{Any,D}) where {T,R,D}
+function apply_transpose(itm::InflatedTensorMapping{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D}
     dim_before = range_dim(itm.before)
     dim_domain = domain_dim(itm.tm)
     dim_range = range_dim(itm.tm)
--- a/src/LazyTensors/tensor_mapping.jl	Thu Mar 10 11:13:34 2022 +0100
+++ b/src/LazyTensors/tensor_mapping.jl	Thu Mar 10 16:57:01 2022 +0100
@@ -24,7 +24,7 @@
 export TensorMapping
 
 """
-    apply(t::TensorMapping{T,R,D}, v::AbstractArray{T,D}, I::Vararg) where {R,D,T}
+    apply(t::TensorMapping{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg) where {R,D,T}
 
 Return the result of the mapping for a given index.
 """
@@ -32,7 +32,7 @@
 export apply
 
 """
-    apply_transpose(t::TensorMapping{T,R,D}, v::AbstractArray{T,R}, I::Vararg) where {R,D,T}
+    apply_transpose(t::TensorMapping{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg) where {R,D,T}
 
 Return the result of the transposed mapping for a given index.
 """
--- a/test/LazyTensors/lazy_tensor_operations_test.jl	Thu Mar 10 11:13:34 2022 +0100
+++ b/test/LazyTensors/lazy_tensor_operations_test.jl	Thu Mar 10 16:57:01 2022 +0100
@@ -123,6 +123,8 @@
 
     @test range_size(A+B) == range_size(A) == range_size(B)
     @test domain_size(A+B) == domain_size(A) == domain_size(B)
+
+    @test ((A+B)*ComplexF64[1.1,1.2,1.3])[3] isa ComplexF64
 end
 
 
@@ -145,6 +147,9 @@
 
     v = rand(2)
     @test (Ã∘B̃)'*v ≈ B'*A'*v rtol=1e-14
+
+    @test (Ã∘B̃*ComplexF64[1.,2.,3.,4.])[1] isa ComplexF64
+    @test ((Ã∘B̃)'*ComplexF64[1.,2.])[1] isa ComplexF64
 end
 
 @testset "LazyLinearMap" begin
@@ -210,6 +215,10 @@
         @test I*v == v
         @test I'*v == v
 
+        v = rand(ComplexF64,sz...)
+        @test I*v == v
+        @test I'*v == v
+
         @test range_size(I) == sz
         @test domain_size(I) == sz
     end
@@ -338,6 +347,16 @@
             end
         end
 
+        @testset "application to other type" begin
+            tm = InflatedTensorMapping(I(3,2), A, I(4))
+
+            v = rand(ComplexF64, domain_size(tm)...)
+            @test (tm*v)[1,2,3,1] isa ComplexF64
+
+            v = rand(ComplexF64, domain_size(tm')...)
+            @test (tm'*v)[1,2,2,1] isa ComplexF64
+        end
+
         @testset "Inference of application" begin
             struct ScalingOperator{T,D} <: TensorMapping{T,D,D}
                 λ::T