LazyArraysFunctors.jl

using NumericFuns

# Array consisting of an element-wise transformation of another array
# Equivalent of vectorized functions
immutable LazyArray{T, N, A, F, V} <: AbstractArray{T, N}
    a::A
    fn::F
    args::V
end

LazyArray(a::AbstractArray, fn::Functor, args) = LazyArray{eltype(evaluate(fn, a[1], args...)), ndims(a), typeof(a), typeof(fn), typeof(args)}(a, fn, args)
LazyArray(a::AbstractArray, fn::Functor) = LazyArray(a, fn, ())

Base.size(a::LazyArray) = size(a.a)
# Recursive definition to follow the original array's structure
Base.similar(a::LazyArray, T=eltype(a)) = similar(a.a, T)
# A separate definition is needed since similar() for sparse matrices
# does not preserve the structure of nonzeros when passed 'dims'
Base.similar(a::LazyArray, T, dims::Dims) = similar(a.a, T, dims)
Base.getindex{T}(a::LazyArray{T}, i::Real) = evaluate(a.fn, a.a[i], a.args...)::T
Base.getindex(a::LazyArray, i::AbstractArray) = evaluate(a.fn, a.a[i...], a.args...)
Base.getindex(a::LazyArray, i...) = evaluate(a.fn, a.a[i...], a.args...)

# Array consisting of an element-wise combination of two arrays
# Equivalent of element-wise operators like .+ or .*
immutable LazyArray2{T, N, A1, A2, F, V} <: AbstractArray{T, N}
    a1::A1
    a2::A2
    fn::F
    args::V
end

function LazyArray2(a1::AbstractArray, a2::AbstractArray, fn::Functor, args)
    @assert size(a1) == size(a2)
    LazyArray2{eltype(evaluate(fn, a1[1], a2[1], args...)), ndims(a1), typeof(a1), typeof(a2), typeof(fn), typeof(args)}(a1, a2, fn, args)
end

LazyArray2(a1::AbstractArray, a2::AbstractArray, fn::Functor) = LazyArray2(a1, a2, fn, ())

Base.size(a::LazyArray2) = size(a.a1)

# Recursive definition to follow the original array's structure
function Base.similar(a::LazyArray2, T=eltype(a))
    @assert size(a.a1) == size(a.a2)

    # FIXME: handle arrays of same type but with different element types
    if typeof(a.a1) == typeof(a.a2)
        similar(a.a1, T)
    else
        # TODO: better promotion mechanism for cases where a non-standard
        # array type would be more appropriate
        Array(T, size(a.a1))
    end 
end

# A separate definition is needed since similar() for sparse matrices
# does not preserve the structure of nonzeros when passed 'dims'
function Base.similar(a::LazyArray2, T, dims::Dims)
    @assert size(a.a1) == size(a.a2)

    # FIXME: handle arrays of same type but with different element types
    if typeof(a.a1) == typeof(a.a2)
        similar(a.a1, T, dims)
    else
        # TODO: better promotion mechanism for cases where a non-standard
        # array type would be more appropriate
        Array(T, dims)
    end 
end

Base.similar(a::LazyArray2, dims::Dims) = similar(a, eltype(a), dims)

Base.getindex{T}(a::LazyArray2{T}, i::Real) = evaluate(a.fn, a.a1[i], a.a2[i], a.args...)::T
Base.getindex(a::LazyArray2, i::AbstractArray) = evaluate(a.fn, a.a1[i], a.a2[i], a.args...)
Base.getindex(a::LazyArray2, i...) = evaluate(a.fn, a.a1[i...], a.a2[i...], a.args...)


# Use linear indexing by default, should be overriden by specific array types
immutable EachIndex{T<:AbstractArray}
    a::T
end
eachindex(a::AbstractArray) = EachIndex(a)
# Recursive choice of iteration method, to follow the original array's structure
eachindex(a::LazyArray) = eachindex(a.a)

Base.length(e::EachIndex) = length(e.a)
Base.start(e::EachIndex) = 1
Base.next(e::EachIndex, i) = (i, i+1)
Base.done(e::EachIndex, i) = i > length(e.a)

# Iteration over nonzero entries of sparse matrices
immutable EachIndexSparseMatrixCSC
    S::SparseMatrixCSC
end
eachindex(a::SparseMatrixCSC) = EachIndexSparseMatrixCSC(a)

Base.length(e::EachIndexSparseMatrixCSC) = nnz(e.S)
Base.start(e::EachIndexSparseMatrixCSC) = 1
Base.next(e::EachIndexSparseMatrixCSC, i) = (NonzeroIndex(i), i+1)
Base.done(e::EachIndexSparseMatrixCSC, i) = i > nnz(e.S)

immutable NonzeroIndex <: Number
    i::Int
end

Base.getindex(S::SparseMatrixCSC, i::NonzeroIndex) = nonzeros(S)[i.i]::T
Base.getindex{T}(a::LazyArray{T}, i::NonzeroIndex) = evaluate(a.fn, a.a[i], a.args...)::T
Base.setindex!(S::SparseMatrixCSC, x, i::NonzeroIndex) = nonzeros(S)[i.i] = x

function Base.collect(a::Union(LazyArray, LazyArray2))
    res = similar(a)
    for i in eachindex(a)
        @inbounds res[i] = a[i]
    end
    res
end

function collect!(newa, a::Union(LazyArray, LazyArray2))
    @assert size(a) == size(newa)
    for i in eachindex(a)
        @inbounds newa[i] = a[i]
    end
    newa
end


# Checks for LazyArray
a = rand(100, 100);
a1 = LazyArray(a, Divide(), (2,));
a2 = LazyArray(a1, SqrtFun());

@assert collect(a1) == a1 == a/2
@assert collect(a2) == a2 == sqrt(a/2)

@assert sum(a1, 2) == sum((a/2), 2)
@assert sum(a2, 2) == sum(sqrt(a/2), 2)


b = sprand(100, 100, .1);
b1 = LazyArray(b, Divide(), (2,));
b2 = LazyArray(b1, SqrtFun());

@assert collect(b1) == b1 == b/2
@assert collect(b2) == b2 == sqrt(b/2)

@assert sum(b1, 2) == sum((b/2), 2)
@assert sum(b2, 2) == sum(sqrt(b/2), 2)

# Performance tests for LazyArray
function test_lazy!(newa, a)
    a1 = LazyArray(a, Divide(), (2,))
    a2 = LazyArray(a1, SqrtFun())
    collect!(newa, a2)
    newa
end

function test_loop_dense!(newa, a)
    for i in 1:length(a)
        newa[i] = sqrt(a[i]/2)
    end
    newa
end

function test_loop_sparse!(newa, a)
    for i in 1:nnz(a)
        nonzeros(newa)[i] = sqrt(nonzeros(a)[i]/2)
    end
    newa
end

newa = similar(a);
newb = similar(b);

test_lazy!(newa, a);
@time for i in 1:1000 test_lazy!(newa, a) end

test_loop_dense!(newa, a);
@time for i in 1:1000 test_loop_dense!(newa, a) end

@time for i in 1:1000 sqrt(a/2) end

test_lazy!(newb, b);
@time for i in 1:1000 test_lazy!(newb, b) end

test_loop_sparse!(newb, b);
@time for i in 1:1000 test_loop_sparse!(newb, b) end

@time for i in 1:1000 sqrt(b/2) end


# Checks for LazyArray2
a3 = LazyArray2(a, a1, Add());
@assert collect(a3) == a3 == a .+ a1

b3 = LazyArray2(b, b1, Add());
@assert collect(b3) == b3 == b .+ b1

c = LazyArray2(a, b, Add());
@assert collect(c) == c == a .+ b

# Performance tests for LazyArray2
function test_lazy2!(newa, a1, a2)
    a = LazyArray2(a1, a2, Add())
    collect!(newa, a)
    newa
end

function test_loop_dense2!(newa, a1, a2)
    for i in 1:length(a1)
        @inbounds newa[i] = a1[i] + a2[i]
    end
    newa
end


a1 = rand(100, 100);
a2 = rand(100, 100);

test_lazy2!(newa, a1, a2);
@time for i in 1:1000 test_lazy2!(newa, a1, a2) end

test_loop_dense2!(newa, a1, a2);
@time for i in 1:1000 test_loop_dense2!(newa, a1, a2) end

@time for i in 1:1000 a1 .+ a2 end


# Performance test for row sums of squared differences
function test_lazy_sumsqdiff(x, y)
    a1 = LazyArray2(x, y, Subtract())
    a2 = LazyArray(a1, SqrFun(), ())
    sum(a2, 1)
end

function test_loop_dense_sumsqdiff(x, y)
    s = zeros(eltype(x), size(x, 1))
    for j in 1:size(x, 2), i in 1:size(x, 1)
        @inbounds s[j] += (x[i, j] - y[i, j])^2
    end
    s
end

@assert vec(test_lazy_sumsqdiff(a1, a2)) == test_loop_dense_sumsqdiff(a1, a2)

@time for i in 1:1000 test_lazy_sumsqdiff(a1, a2) end

@time for i in 1:1000 test_loop_dense_sumsqdiff(a1, a2) end

@time for i in 1:1000 sum((a1 .- a2).^2) end