changeset 403:618b7ee73b25 refactor/sbp_operators_tests/collect_and_compare

Merge in default and close branch before merge
author Jonatan Werpers <jonatan@werpers.com>
date Wed, 07 Oct 2020 12:31:54 +0200
parents 3cecbfb3d623 (current diff) 1936e38fe51e (diff)
children 48d57f185f86
files
diffstat 2 files changed, 88 insertions(+), 8 deletions(-) [+]
line wrap: on
line diff
diff -r 3cecbfb3d623 -r 618b7ee73b25 src/LazyTensors/lazy_tensor_operations.jl
--- a/src/LazyTensors/lazy_tensor_operations.jl	Sun Oct 04 19:41:14 2020 +0200
+++ b/src/LazyTensors/lazy_tensor_operations.jl	Wed Oct 07 12:31:54 2020 +0200
@@ -102,3 +102,38 @@
 # # Have i gone too crazy with the type parameters? Maybe they aren't all needed?
 
 # export →
+"""
+    LazyLinearMap{T,R,D,...}(A, range_indicies, domain_indicies)
+
+TensorMapping defined by the AbstractArray A. `range_indicies` and `domain_indicies` define which indicies of A should
+be considerd the range and domain of the TensorMapping. Each set of indices must be ordered in ascending order.
+
+For instance, if A is a m x n matrix, and range_size = (1,), domain_size = (2,), then the LazyLinearMap performs the
+standard matrix-vector product on vectors of size n.
+"""
+struct LazyLinearMap{T,R,D, RD, AA<:AbstractArray{T,RD}} <: TensorMapping{T,R,D}
+    A::AA
+    range_indicies::NTuple{R,Int}
+    domain_indicies::NTuple{D,Int}
+
+    function LazyLinearMap(A::AA, range_indicies::NTuple{R,Int}, domain_indicies::NTuple{D,Int}) where {T,R,D, RD, AA<:AbstractArray{T,RD}}
+        if !issorted(range_indicies) || !issorted(domain_indicies)
+            throw(DomainError("range_indicies and domain_indicies must be sorted in ascending order"))
+        end
+
+        return new{T,R,D,RD,AA}(A,range_indicies,domain_indicies)
+    end
+end
+export LazyLinearMap
+
+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{Index,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])
+    end
+    A_view = @view llm.A[view_index...]
+    return sum(A_view.*v)
+end
diff -r 3cecbfb3d623 -r 618b7ee73b25 test/testLazyTensors.jl
--- a/test/testLazyTensors.jl	Sun Oct 04 19:41:14 2020 +0200
+++ b/test/testLazyTensors.jl	Wed Oct 07 12:31:54 2020 +0200
@@ -118,15 +118,15 @@
 end
 
 @testset "LazyArray" begin
-	@testset "LazyConstantArray" begin
-	    @test LazyTensors.LazyConstantArray(3,(3,2)) isa LazyArray{Int,2}
+    @testset "LazyConstantArray" begin
+        @test LazyTensors.LazyConstantArray(3,(3,2)) isa LazyArray{Int,2}
 
-	    lca = LazyTensors.LazyConstantArray(3.0,(3,2))
-	    @test eltype(lca) == Float64
-	    @test ndims(lca) == 2
-	    @test size(lca) == (3,2)
-	    @test lca[2] == 3.0
-	end
+        lca = LazyTensors.LazyConstantArray(3.0,(3,2))
+        @test eltype(lca) == Float64
+        @test ndims(lca) == 2
+        @test size(lca) == (3,2)
+        @test lca[2] == 3.0
+    end
     struct DummyArray{T,D, T1<:AbstractArray{T,D}} <: LazyArray{T,D}
         data::T1
     end
@@ -193,6 +193,7 @@
     @test_throws DimensionMismatch v1 + v2
 end
 
+
 @testset "LazyFunctionArray" begin
     @test LazyFunctionArray(i->i^2, (3,)) == [1,4,9]
     @test LazyFunctionArray((i,j)->i*j, (3,2)) == [
@@ -212,4 +213,48 @@
 
 end
 
+@testset "LazyLinearMap" begin
+    # Test a standard matrix-vector product
+    # mapping vectors of size 4 to vectors of size 3.
+    A = rand(3,4)
+    Ã = LazyLinearMap(A, (1,), (2,))
+    v = rand(4)
+
+    @test à isa LazyLinearMap{T,1,1} where T
+    @test à isa TensorMapping{T,1,1} where T
+    @test range_size(Ã) == (3,)
+    @test domain_size(Ã) == (4,)
+
+    @test Ã*ones(4) ≈ A*ones(4) atol=5e-13
+    @test Ã*v ≈ A*v atol=5e-13
+
+    A = rand(2,3,4)
+    @test_throws DomainError LazyLinearMap(A, (3,1), (2,))
+
+    # Test more exotic mappings
+    B = rand(3,4,2)
+    # Map vectors of size 2 to matrices of size (3,4)
+    B̃ = LazyLinearMap(B, (1,2), (3,))
+    v = rand(2)
+
+    @test range_size(B̃) == (3,4)
+    @test domain_size(B̃) == (2,)
+    @test B̃ isa TensorMapping{T,2,1} where T
+    @test B̃*ones(2) ≈ B[:,:,1] + B[:,:,2] atol=5e-13
+    @test B̃*v ≈ B[:,:,1]*v[1] + B[:,:,2]*v[2] atol=5e-13
+
+    # Map matrices of size (3,2) to vectors of size 4
+    B̃ = LazyLinearMap(B, (2,), (1,3))
+    v = rand(3,2)
+
+    @test range_size(B̃) == (4,)
+    @test domain_size(B̃) == (3,2)
+    @test B̃ isa TensorMapping{T,1,2} where T
+    @test B̃*ones(3,2) ≈ B[1,:,1] + B[2,:,1] + B[3,:,1] +
+                        B[1,:,2] + B[2,:,2] + B[3,:,2] atol=5e-13
+    @test B̃*v ≈ B[1,:,1]*v[1,1] + B[2,:,1]*v[2,1] + B[3,:,1]*v[3,1] +
+                B[1,:,2]v[1,2] + B[2,:,2]*v[2,2] + B[3,:,2]*v[3,2] atol=5e-13
+
 end
+
+end