LazyArrays.jl

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

LazyArray(a::AbstractArray, fn::Function, args) = LazyArray{eltype(fn(a[1], args...)), ndims(a), typeof(a)}(a, fn, args)
LazyArray(a::AbstractArray, fn::Function) = 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(a::LazyArray, i::Real) = a.fn(a.a[i], a.args...)
Base.getindex(a::LazyArray, i::AbstractArray) = a.fn(a.a[i...], a.args...)
Base.getindex(a::LazyArray, i...) = 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} <: AbstractArray{T, N}
    a1::A1
    a2::A2
    fn::Function
    args
end

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

LazyArray2(a1::AbstractArray, a2::AbstractArray, fn::Function) = 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(a::LazyArray2, i::Real) = a.fn(a.a1[i], a.a2[i], a.args...)
Base.getindex(a::LazyArray2, i::AbstractArray) = a.fn(a.a1[i...], a.a2[i...], a.args...)
Base.getindex(a::LazyArray2, i...) = a.fn(a.a[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]
Base.getindex(a::LazyArray, i::NonzeroIndex) = a.fn(a.a[i], a.args...)
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


# Checks for LazyArray
a = rand(100, 100);
a1 = LazyArray(a, /, (2,));
a2 = LazyArray(a1, ^, (2,));

@assert collect(a1) == a1 == a/2
@assert collect(a2) == a2 == (a/2).^2

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


b = sprand(100, 100, .2);
b1 = LazyArray(b, /, (2,));
b2 = LazyArray(b1, ^, (2,));

@assert collect(b1) == b1 == b/2
@assert collect(b2) == b2 == (b/2).^2

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


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

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

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


# Performance tests
function test_lazy(a)
    a1 = LazyArray(a, /, (2,))
    a2 = LazyArray(a1, ^, (2,))
    collect(a2)
end

function test_loop_dense(a)
    b = similar(a)
    for i in 1:length(a)
        b[i] = (a[i]/2)^2
    end
    b
end

function test_loop_sparse(a)
    b = similar(a)
    for i in 1:nnz(a)
        nonzeros(b)[i] = (nonzeros(a)[i]/2)^2
    end
    b
end

@time for i in 1:1000 test_lazy(a) end
# elapsed time: 3.235718017 seconds (1192272000 bytes allocated, 32.20% gc time)
@time for i in 1:1000 test_loop_dense(a) end
# elapsed time: 0.1004941 seconds (80048000 bytes allocated, 67.97% gc time)

@time for i in 1:1000 test_lazy(b) end
# elapsed time: 1.327816608 seconds (368032000 bytes allocated, 22.45% gc time)
@time for i in 1:1000 test_loop_sparse(b) end
# elapsed time: 0.048461734 seconds (33264000 bytes allocated, 70.05% gc time)