struct lazy_e{T} <: AbstractVector{T}
i::Int
n::Int
end
lazy_e(i::Int, n::Int, ::Type{T} = Int) where T = lazy_e{T}(i, n)
function Base.getindex(e::lazy_e{T}, i) where T
@boundscheck @assert i <= e.n
ifelse(i == e.i, one(T), zero(T))
end
Base.size(e::lazy_e) = (e.n,)
Then on Julia 0.7:
julia> fe(x) = (@inbounds d = x' * lazy_e(3,4); d)
fe (generic function with 1 method)
julia> using StaticArrays, BenchmarkTools
julia> x = @SVector randn(4)
4-element SArray{Tuple{4},Float64,1,4}:
-0.02571662495110571
-0.2620502046882458
-0.8456889511406026
0.017055806672449558
julia> y = randn(4)
4-element Array{Float64,1}:
-1.858893446204762
1.883085443252701
-0.7957485906904918
-0.5190800039346426
julia> @btime fe($x)
1.520 ns (0 allocations: 0 bytes)
-0.8456889511406026
julia> @btime fe($y)
8.301 ns (0 allocations: 0 bytes)
-0.7957485906904918
julia> lazy_e(2,5)
5-element lazy_e{Int64}:
0
1
0
0
0
julia> lazy_e(2,5)'
1×5 LinearAlgebra.Adjoint{Int64,lazy_e{Int64}}:
0 1 0 0 0
EDIT:
Ideally, you’d do this with a for loop and @eval
, and perhaps some other helper macros, but:
Base.:*(A::Matrix, e::lazy_e) = A[:, e.i]
Base.:*(A::Adjoint{T,Matrix{T}}, e::lazy_e) where T = A[:, e.i]
Base.:*(A::Transpose{T,Matrix{T}}, e::lazy_e) where T = A[:, e.i]
Base.:*(A::Adjoint{T,Vector{T}}, e::lazy_e) where T = A[e.i]
Base.:*(A::Transpose{T,Vector{T}}, e::lazy_e) where T = A[e.i]
Base.:*(A::Adjoint{T,lazy_e}, e::Matrix{T}) where T = A[e.i, :]
Base.:*(A::Transpose{T,lazy_e}, e::Matrix{T}) where T = A[e.i, :]
Base.:*(A::Adjoint{T,lazy_e}, e::Vector{T}) where T = A[e.i]
Base.:*(A::Transpose{T,lazy_e}, e::Vector{T}) where T = A[e.i]
I had to get really specific for dispatch to actually favor the newly defined multiplication methods, and I’m not sure why. This yielded:
julia> z = randn(3,4)
3×4 Array{Float64,2}:
0.478915 -0.542889 0.142261 -0.594415
0.509348 0.63491 1.77353 -0.545189
-0.659932 0.995295 1.706 -0.291947
julia> @btime $z * lazy_e(2,4)
26.288 ns (1 allocation: 112 bytes)
3-element Array{Float64,1}:
-0.5428894820340465
0.6349099814045577
0.9952951252882498
julia> @btime $z[:,2]
26.266 ns (1 allocation: 112 bytes)
3-element Array{Float64,1}:
-0.5428894820340465
0.6349099814045577
0.9952951252882498
The matrix multiplication functions could be edited to return views if you.