Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[SparseArrayDOKs] Add setindex_maybe_grow! and macro @maybe_grow #1434

Merged
merged 10 commits into from
May 14, 2024
Merged

[SparseArrayDOKs] Add setindex_maybe_grow! and macro @maybe_grow #1434

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
e74f478
Add maybe_grow features
emstoudenmire May 13, 2024
320ac28
Minor edits, add check in @maybe_grow macro
emstoudenmire May 13, 2024
4af4cc2
Improve syntax check for @maybe_grow
emstoudenmire May 13, 2024
41ef815
Improve readability
emstoudenmire May 13, 2024
625fa27
Use MacroTools in @maybe_grow
emstoudenmire May 13, 2024
0520e67
Move expression check to separate function
emstoudenmire May 14, 2024
7da8f09
Use dispatch to check setindex! expression
emstoudenmire May 14, 2024
93a65c6
Make dims field a Ref and remove mutable keyword
emstoudenmire May 14, 2024
c952653
Add and test constructor
emstoudenmire May 14, 2024
aba62fd
Update NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl
emstoudenmire May 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 27 additions & 1 deletion NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ using ..SparseArrayInterface:
SparseArrayInterface, AbstractSparseArray, getindex_zero_function

# TODO: Parametrize by `data`?
struct SparseArrayDOK{T,N,Zero} <: AbstractSparseArray{T,N}
mutable struct SparseArrayDOK{T,N,Zero} <: AbstractSparseArray{T,N}
emstoudenmire marked this conversation as resolved.
Show resolved Hide resolved
data::Dictionary{CartesianIndex{N},T}
dims::NTuple{N,Int}
zero::Zero
Expand Down Expand Up @@ -104,3 +104,29 @@ SparseArrayDOK{T}(a::AbstractArray) where {T} = SparseArrayDOK{T,ndims(a)}(a)
function SparseArrayDOK{T,N}(a::AbstractArray) where {T,N}
return SparseArrayInterface.sparse_convert(SparseArrayDOK{T,N}, a)
end

function Base.resize!(a::SparseArrayDOK{<:Any,N}, new_size::NTuple{N,Integer}) where {N}
a.dims = new_size
return a
end

function setindex_maybe_grow!(a::SparseArrayDOK{<:Any,N}, value, I::Vararg{Int,N}) where {N}
if any(I .> size(a))
resize!(a, max.(I, size(a)))
end
a[I...] = value
return a
end

macro maybe_grow(ex)
if ex.head != :(=)
emstoudenmire marked this conversation as resolved.
Show resolved Hide resolved
error("@maybe_grow must be used with setindex! syntax (@maybe_grow a[inds] = value)")
end
arr_name = esc(ex.args[1].args[1])
index = esc(ex.args[1].args[2:end])
value = esc(ex.args[2])
emstoudenmire marked this conversation as resolved.
Show resolved Hide resolved
e = quote
setindex_maybe_grow!($arr_name, $value, $index...)
end
return e
end
27 changes: 26 additions & 1 deletion NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,8 @@
# Slicing

using Test: @test, @testset, @test_broken
using NDTensors.SparseArrayDOKs: SparseArrayDOK, SparseMatrixDOK
using NDTensors.SparseArrayDOKs:
SparseArrayDOKs, SparseArrayDOK, SparseMatrixDOK, @maybe_grow
using NDTensors.SparseArrayInterface: storage_indices, nstored
using SparseArrays: SparseMatrixCSC, nnz
@testset "SparseArrayDOK (eltype=$elt)" for elt in
Expand Down Expand Up @@ -94,5 +95,29 @@ using SparseArrays: SparseMatrixCSC, nnz
end
end
end
@testset "Maybe Grow Feature" begin
a = SparseArrayDOK{elt,2}((0, 0))
SparseArrayDOKs.setindex_maybe_grow!(a, 230, 2, 3)
@test size(a) == (2, 3)
@test a[2, 3] == 230
# Test @maybe_grow macro
@maybe_grow a[5, 5] = 550
@test size(a) == (5, 5)
@test a[2, 3] == 230
@test a[5, 5] == 550
# Test that size remains same
# if we set at an index smaller than
# the maximum size:
@maybe_grow a[3, 4] = 340
@test size(a) == (5, 5)
@test a[2, 3] == 230
@test a[5, 5] == 550
@test a[3, 4] == 340
# Test vector case
v = SparseArrayDOK{elt,1}((0,))
@maybe_grow v[5] = 50
@test size(v) == (5,)
@test v[5] == 50
end
end
end
Loading