Clever design for basis arrays

General Usage
, ,
1

To compute Jacobians and Hessians in DifferentiationInterface.jl, I need to write a function basis(x::AbstractArray, inds) that computes parts of the standard basis in the appropriate vector space.
I would like it to have the following properties:

  1. basis(x, inds) returns an array with the same type as x, containing zeros everywhere and ones at the indices inds
  2. inds is obtained as a subset of eachindex(x), so its elements can be either integers or CartesianIndex, and its length is not statically known
  3. it natively works for weird array types like StaticArrays.jl or JLArrays.jl (GPU)
  4. it is efficient
  5. it is self-contained to avoid unnecessary load time

Sofar I’ve been relying on FillArrays.jl but it doesn’t satisfy criteria 2 and 5:

using FillArrays

function basis(
    x::AbstractArray{T,N}, inds::AbstractVector{<:CartesianIndex{N}}
) where {T,N}
    b = zero(x)
    for i in inds
        b += OneElement(one(T), Tuple(i), axes(x))
    end
    return b
end

I also have a self-contained version which doesn’t satisfy criteria 2 and 3 (because it calls similar, which changes the type of a StaticArrays.SArray to allow mutation):

using LinearAlgebra

function basis(
    x::AbstractArray{T}, inds::AbstractVector{<:Integer}
) where {T}
    b = zero(similar(x))
    I = Diagonal(vec(b) .+ one(T))
    vec(b) .+= sum(view(I, :, inds); dims=2)
    return b
end

Does anyone have a brilliant idea or is my wish list irreconcilable? I would already be happy with a basis(x, i) that works for a single index

2

What’s the issue with OneElement re. point 2? The fact that the constructor accepts indices that may potentially lie outside eachindex? Or that it doesn’t accept a CartesianIndex currently, which it probably should?

3

It doesn’t accept integer indices, only tuples (or equivalently CartesianIndex). Arguably not a big flaw, so perhaps I should just code my own version of OneElement and call it a day?

4

I think it’s reasonable for it to accept a CartesianIndex. A PR will be welcome!

2 Likes