changeset 374:99296cbb7bcd

Merge feature/lazy_function
author Jonatan Werpers <jonatan@werpers.com>
date Wed, 30 Sep 2020 21:08:03 +0200
parents 8e55dee6a1a1 (current diff) 125ca96b781d (diff)
children f65809a26a17
files
diffstat 2 files changed, 43 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
--- a/src/LazyTensors/lazy_array.jl	Mon Sep 28 22:56:54 2020 +0200
+++ b/src/LazyTensors/lazy_array.jl	Wed Sep 30 21:08:03 2020 +0200
@@ -17,6 +17,30 @@
 Base.getindex(lca::LazyConstantArray{T,D}, I::Vararg{Int,D}) where {T,D} = lca.val
 
 """
+    LazyFunctionArray{F<:Function,T, D} <: LazyArray{T,D}
+
+A lazy array where each element is defined by a function f(i,j,...)
+"""
+struct LazyFunctionArray{F<:Function,T, D} <: LazyArray{T,D}
+    f::F
+    size::NTuple{D,Int}
+end
+export LazyFunctionArray
+
+function LazyFunctionArray(f::F, size::NTuple{D,Int}) where {F<:Function,D}
+    T = typeof(f(ones(D)...))
+    return LazyFunctionArray{F,T,D}(f,size)
+end
+
+Base.size(lfa::LazyFunctionArray) = lfa.size
+
+function Base.getindex(lfa::LazyFunctionArray{F,T,D}, I::Vararg{Int,D}) where {F,T,D}
+    @boundscheck checkbounds(lfa, I...)
+    return lfa.f(I...)
+end
+
+
+"""
     LazyElementwiseOperation{T,D,Op} <: LazyArray{T,D}
 Struct allowing for lazy evaluation of elementwise operations on AbstractArrays.
 
--- a/test/testLazyTensors.jl	Mon Sep 28 22:56:54 2020 +0200
+++ b/test/testLazyTensors.jl	Wed Sep 30 21:08:03 2020 +0200
@@ -193,4 +193,23 @@
     @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)) == [
+        1 2;
+        2 4;
+        3 6;
+    ]
+
+    @test size(LazyFunctionArray(i->i^2, (3,))) == (3,)
+    @test size(LazyFunctionArray((i,j)->i*j, (3,2))) == (3,2)
+
+    @inferred LazyFunctionArray(i->i^2, (3,))[2]
+
+    @test_throws BoundsError LazyFunctionArray(i->i^2, (3,))[4]
+    @test_throws BoundsError LazyFunctionArray((i,j)->i*j, (3,2))[4,2]
+    @test_throws BoundsError LazyFunctionArray((i,j)->i*j, (3,2))[2,3]
+
 end
+
+end