Skip to content

Commit

Permalink
Preliminary support for matrix multiplication.
Browse files Browse the repository at this point in the history
- New methods for * for combinations of OffsetArray and AbstractArray
  Unfortunately this introduces method ambiguities with Base.
- As an alternative, tried defining a new promote_rule, but somehow it doesn't get invoked.

New convenience constructors
- Wrap any (non-offset) array as an OffsetArray with standard indices
- Remove layer of indirection when constructing with explicit offsets

Cosmetic edits to constructor section for improved readability
  • Loading branch information
Ryan Bennink committed Jan 3, 2020
1 parent e22f56e commit 945ef97
Showing 1 changed file with 36 additions and 2 deletions.
38 changes: 36 additions & 2 deletions src/OffsetArrays.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,32 +3,43 @@ VERSION < v"0.7.0-beta2.199" && __precompile__()
module OffsetArrays

using Base: Indices, tail, @propagate_inbounds
import Base: (*), convert, promote_rule

@static if !isdefined(Base, :IdentityUnitRange)
const IdentityUnitRange = Base.Slice
else
using Base: IdentityUnitRange
end

export OffsetArray, OffsetVector
export OffsetArray, OffsetVector, OffsetMatrix

struct OffsetArray{T,N,AA<:AbstractArray} <: AbstractArray{T,N}
parent::AA
offsets::NTuple{N,Int}
end
OffsetVector{T,AA<:AbstractArray} = OffsetArray{T,1,AA}
OffsetMatrix{T,AA<:AbstractArray} = OffsetArray{T,2,AA}

# Construct from existing array
OffsetArray(A::AbstractArray{T,N}, offsets::NTuple{N,Int}) where {T,N} =
OffsetArray{T,N,typeof(A)}(A, offsets)
OffsetArray(A::AbstractArray{T,N}, offsets::Vararg{Int,N}) where {T,N} =
OffsetArray(A, offsets)
OffsetArray(A::AbstractArray{T,0}) where {T} = OffsetArray{T,0,typeof(A)}(A, ())
OffsetArray(A::AbstractArray) = convert(OffsetArray, A)
OffsetArray(A::OffsetArray) = A

# Convert from existing array
convert(::Type{OffsetArray}, A::AbstractArray{T,N}) where {T,N} = OffsetArray(A, ntuple(x->0, N))

# Create an uninitialized OffsetArray with given element type
const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}
OffsetArray{T,N}(init::ArrayInitializer, inds::Indices{N}) where {T,N} =
OffsetArray{T,N,Array{T,N}}(Array{T,N}(init, map(indexlength, inds)), map(indexoffset, inds))
OffsetArray{T}(init::ArrayInitializer, inds::Indices{N}) where {T,N} = OffsetArray{T,N}(init, inds)
# Same thing, but taking multiple args for offsets/indices
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
OffsetArray{T}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
OffsetArray(A::AbstractArray{T,0}) where {T} = OffsetArray{T,0,typeof(A)}(A, ())

# OffsetVector constructors
OffsetVector(A::AbstractVector, offset) = OffsetArray(A, offset)
Expand Down Expand Up @@ -59,6 +70,7 @@ OffsetArray(A::AbstractArray{T,N}, inds::Vararg{AbstractUnitRange,N}) where {T,N
OffsetArray(A, inds)

# avoid a level of indirection when nesting OffsetArrays
OffsetArray(A::OffsetArray, offsets::NTuple{N,Int}) where {N} = OffsetArray(parent(A), offsets)
function OffsetArray(A::OffsetArray, inds::NTuple{N,AbstractUnitRange}) where {N}
OffsetArray(parent(A), inds)
end
Expand Down Expand Up @@ -283,6 +295,28 @@ end
no_offset_view(A::OffsetArray) = no_offset_view(parent(A))


# Quick hack for matrix multiplication.
# Ideally, one would instead improve LinearAlgebra's support of custom indexing.
function (*)(A::OffsetMatrix, B::OffsetMatrix)
matmult_check_axes(A, B)
C = OffsetArray(parent(A) * parent(B), (axes(A,1), axes(B,2)))
end

function (*)(A::OffsetMatrix, B::OffsetVector)
matmult_check_axes(A, B)
C = OffsetArray(parent(A) * parent(B), axes(A,1))
end
matmult_check_axes(A, B) = axes(A, 2) == axes(B, 1) || error("axes(A,2) must equal axes(B,1)")

(*)(A::OffsetMatrix, B::AbstractMatrix) = A * OffsetArray(B)
(*)(A::OffsetMatrix, B::AbstractVector) = A * OffsetArray(B)
(*)(A::AbstractMatrix, B::OffsetArray) = OffsetArray(A) * B
(*)(A::AbstractVector, B::OffsetArray) = OffsetArray(A) * B

# An alternative to the above four methods would be to use promote_rule, but it doesn't get invoked
# promote_rule(::Type{A1}, ::Type{A2}) where A1<:AbstractArray{<:Any,N} where A2<:OffsetArray{<:Any,N,A3} where {N,A3} = OffsetArray{eltype(promote_type(A1, A3)), N, promote_type(A1, A3)}


####
# work around for segfault in searchsorted*
# https://github.com/JuliaLang/julia/issues/33977
Expand Down

0 comments on commit 945ef97

Please sign in to comment.