Skip to content

Instantly share code, notes, and snippets.

@nalimilan
Last active August 29, 2015 14:09
Show Gist options
  • Save nalimilan/e737bc8b3b10288abdad to your computer and use it in GitHub Desktop.
Save nalimilan/e737bc8b3b10288abdad to your computer and use it in GitHub Desktop.
# 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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment