How can I specialize similar and copy on a offset custom array?

General Usage Performance
, ,
1

Hi All! :smiley:

I have a Spinor struct, that Inherits from AbstractVector, but it’s index go from 0 to 2^{\lfloor dim/2\rfloor}-1, with dim a part of the Spinor struct, for this I’m using CustomUnitRanges.jl. The Spinor struct contains a SparseArray to save the values.

Reading in the manual I found that I needed t specialize similar in two different ways, but I was not getting arround about it, so I looked on the OffsetArray implementation and tried something. That worked, but when I was doing copy, It would allocate all the zeros of the SparseArray.

This is extremely problematic, because the plan is to have a dim as large as possible (maybe 1024 o even more), and the benefit of the SparseArray is to never have to allocate all the values. For the kind of problem I’m solving it’s extremely rare to have too many of the values as non-zeros.

This is a somewhat minimal example that is reproducible with my actual implementation. I added the show routines so I was easier to see in the REPL.

module Spinors

# needs to install CustomUnitRanges
# using Pkg
# Pkg.add("CustomUnitRanges")

using CustomUnitRanges: filename_for_zerorange
include(filename_for_zerorange)

using SparseArrays

export Spinor, spinor

struct Spinor{Tv<:Number} <: AbstractVector{Tv}
    ti::NamedTuple{(:dim,),Tuple{Int}}
    v::SparseVector{Tv,Int}  # Values

    function Spinor{Tv}(
        v::SparseVector{Tv,Ti},
        ti::NamedTuple{(:dim,),Tuple{Int}},
    ) where {Tv<:Number,Ti<:Integer}
        ti.dim >= 0 || throw(ArgumentError("The dimension of the spinor must be positive"))
        d = ti.dim >> 1
        r = d == 0 ? 0 : 1 << d
        spv = sparsevec(SparseArrays.nonzeroinds(v), SparseArrays.nonzeros(v), r)
        new{Tv}(ti, spv)
    end
end

spinor(
    coeff::AbstractVector{Tv},
    idx::AbstractVector{Ti};
    dim,
) where {Tv<:Number,Ti<:Integer} = Spinor{Tv}(sparsevec(idx .+ 1, coeff), (; dim))

Base.IndexStyle(::Spinor) = IndexLinear()

Base.size(spn::Spinor) = (length(spn),)

Base.axes(spn::Spinor) = (Base.Slice(map(ZeroRange, length(spn))),)

Base.length(spn::Spinor) = 1 << (spn.ti.dim >> 1)

const SpnAxis = Base.Slice{ZeroRange{Int}}

function Base.similar(
    spn::Spinor,
    ::Type{T},
    inds::Tuple{SpnAxis,Vararg{SpnAxis}},
) where {T}
    B = similar(spn.v, T, map(length, inds))
    return spinor(B, 0:(length(spn) - 1), dim = spn.ti.dim)
end

@inline function Base.getindex(spn::Spinor, i::Int)
    @boundscheck checkbounds(spn, i)
    spn.v[i + 1]
end

@inline function Base.getindex(spn::Spinor, I::AbstractUnitRange)
    @boundscheck checkbounds(spn, I)
    Spinor{eltype(spn)}(spn.v[I .+ 1], (; dim))
end

Base.getindex(spn::Spinor, ::Colon) = copy(spn)

@inline function Base.setindex!(spn::Spinor{Tv}, v::Tv, i::Int) where {Tv}
    @boundscheck checkbounds(spn, i)
    spn.v[i + 1] = v
end

SparseArrays.nonzeroinds(spn::Spinor) = SparseArrays.nonzeroinds(spn.v) .- 1

SparseArrays.nonzeros(spn::Spinor) = SparseArrays.nonzeros(spn.v)

SparseArrays.dropzeros!(spn::Spinor) = SparseArrays.dropzeros!(spn.v)

function Base.show(io::IO, ::MIME"text/plain", x::Spinor)
    n = length(SparseArrays.nonzeros(x.v))
    print(
        io,
        typeof(x),
        "($(x.ti.dim))",
        " with ",
        n,
        n == 1 ? " entry" : " entries",
        " stored",
    )
    if count(!iszero, x.v) > 0
        println(io, ":")
        print(IOContext(io, :typeinfo => eltype(x)), x)
    end
end

Base.show(io::IO, x::Spinor) = show(convert(IOContext, io), x)
function Base.show(io::IOContext, x::Spinor)
    n = length(x)
    nzind = SparseArrays.nonzeroinds(x)
    nzval = nonzeros(x)
    if isempty(nzind)
        return show(io, MIME("text/plain"), x)
    end
    limit = get(io, :limit, false)::Bool
    half_screen_rows = limit ? div(displaysize(io)[1] - 8, 2) : typemax(Int)
    pad = ndigits(n)
    if !haskey(io, :compact)
        io = IOContext(io, :compact => true)
    end
    for k in eachindex(nzind)
        if k < half_screen_rows || k > length(nzind) - half_screen_rows
            print(io, "  ", '[', rpad(nzind[k], pad), "]  =  ")
            if isassigned(nzval, Int(k))
                show(io, nzval[k])
            else
                print(io, Base.undef_ref_str)
            end
            k != length(nzind) && println(io)
        elseif k == half_screen_rows
            println(io, "   ", " "^pad, "   \u22ee")
        end
    end
end

end # module

What I would want to happen is

using Spinors

spn1 = spinor([1], [0], dim=16)
#Spinor{Int64}(16) with 1 entry stored:

spn1[:] # or copy(spn1)
#Spinor{Int64}(16) with 1 entry stored:

What hapen is

using Spinors

spn1 = spinor([1], [0], dim=16)
#Spinor{Int64}(16) with 1 entry stored:

spn1[:] # or copy(spn1)
#Spinor{Int64}(16) with 256 entries stored:

Of course I can do dropzeros!(spn1[:]), but what I want is to never need to allocate all that memory.

I’m using the latest CustomUnitRanges, and my Julia version is

julia> versioninfo()
Julia Version 1.5.1
Commit 697e782ab8 (2020-08-25 20:08 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Core(TM) i7-6500U CPU @ 2.50GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-9.0.1 (ORCJIT, skylake)
Environment:
  JULIA_NUM_THREADS = 4

P.D. I could try to use OffsetArrays, but the index of the SparseArray will not be an Int, but a large Int (maybe BigInt, or something else with Integer interface), where I will be continuously packing and unpacking the bits on other Int’s for indexing. And I have no idea on how or if I can do that with OffsetArray.

Thanks in advance. ñ.ñ :slight_smile: