From a76f19eb4a9b41e54b974ca0bbafc750ebad30c4 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Fri, 22 May 2026 11:24:13 +0200 Subject: [PATCH 01/22] Add two types of filtered view on IndexedVarArray --- src/SparseVariables.jl | 4 + src/indexedarray.jl | 299 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 299 insertions(+), 4 deletions(-) diff --git a/src/SparseVariables.jl b/src/SparseVariables.jl index 96b2e63..c1cfbab 100644 --- a/src/SparseVariables.jl +++ b/src/SparseVariables.jl @@ -12,6 +12,10 @@ include("tables.jl") export SparseArray export IndexedVarArray +export IndexedVarArrayView +export filter_view +export IndexedVarArrayViewAlt +export filter_view_alt export insertvar! export unsafe_insertvar! export SafeInsert, UnsafeInsert diff --git a/src/indexedarray.jl b/src/indexedarray.jl index 9f52440..f3a0959 100644 --- a/src/indexedarray.jl +++ b/src/indexedarray.jl @@ -65,7 +65,7 @@ end """ unsafe_insertvar!(var::indexedVarArray{V,N,T}, index...) -Insert a new variable with the given index withouth checking if the index is valid or +Insert a new variable with the given index withouth checking if the index is valid or already assigned. """ function unsafe_insertvar!(var::IndexedVarArray{V,N,T}, index...) where {V,N,T} @@ -103,10 +103,9 @@ function build_cache!(cache, pat, sa::IndexedVarArray{V,N,T}) where {V,N,T} return cache end -function _select_cached(sa::IndexedVarArray{V,N,T}, pat) where {V,N,T} +function _select_cached(sa::IndexedVarArray{V,N,T}, pat)::Vector{T} where {V,N,T} # TODO: Benchmark to find good cutoff-value for caching - # TODO: Return same type for type stability - length(_data(sa)) < 100 && return _select_gen(keys(_data(sa)), pat) + length(_data(sa)) < 100 && return collect(T, _select_gen(keys(_data(sa)), pat)) cache = _getcache(sa, pat)::Dictionary{_decode_nonslices(sa, pat),Vector{T}} build_cache!(cache, pat, sa) vals = _dropslices_gen(pat) @@ -210,3 +209,295 @@ end function Base.lastindex(sa::IndexedVarArray, d) return last(sort(sa.index_names[d])) end + +# ------------------------------------------------------------------------------ +# IndexedVarArrayView +# ------------------------------------------------------------------------------ + +""" + IndexedVarArrayView{V,N,T,MT,FT} + +A lazy, filtered view into an `IndexedVarArray` where some dimensions are fixed +to specific values and the rest are free (marked with `Colon`). Iterates as an +`AbstractDict` mapping projected keys (`FT`, covering only free dimensions) to +`V` (variable refs). Backed by the parent's `index_cache`, so iteration is O(1) +once the cache is warm. + +Create via `filter_view(iva, mask...)`. +""" +struct IndexedVarArrayView{V<:AbstractVariableRef,N,T,MT<:Tuple,FT<:Tuple} <: + AbstractDict{FT,V} + parent::IndexedVarArray{V,N,T} + mask::MT +end + +""" + filter_view(iva::IndexedVarArray, mask...) + +Return a lazy `IndexedVarArrayView` over entries of `iva` matching `mask`. +Use `:` for free (wildcard) dimensions and exact values for fixed dimensions. +The keys of the view are projected tuples covering only the free dimensions. + +Unlike `Base.view`, the result is an `AbstractDict{FT, V}` (not an +`AbstractArray`) where `FT` is a tuple of the free-dimension value types. + +# Example +```julia +v = filter_view(flow, :, c, p, t) # one free dim: factory +for ((f,), var) in v; ...; end # projected key is a 1-tuple +sum(values(v)) # sum of matching VariableRefs +``` +""" +function filter_view(iva::IndexedVarArray{V,N,T}, mask...) where {V,N,T} + return _make_view(iva, tuple(mask...)) +end + +@generated function _make_view( + iva::IndexedVarArray{V,N,T}, + mask::MT, +) where {V,N,T,MT<:Tuple} + fieldcount(MT) != N && return :(throw(BoundsError(iva, mask))) + free = [fieldtypes(T)[i] for i in 1:N if fieldtypes(MT)[i] === Colon] + FT = Tuple{free...} + return :(IndexedVarArrayView{$V,$N,$T,$MT,$FT}(iva, mask)) +end + +# Project a full key T down to the free dimensions FT. +@generated function _project_key(key::T, ::Type{MT}) where {T,MT} + free_idx = [i for i in 1:fieldcount(T) if fieldtypes(MT)[i] === Colon] + FT = Tuple{[fieldtypes(T)[i] for i in free_idx]...} + return :($(Expr(:tuple, [:(key[$i]) for i in free_idx]...))::$FT) +end + +# Reconstruct a full key T from the fixed values in mask and the free key FT. +@generated function _reconstruct_key(mask::MT, free_key::FT, ::Type{T}) where {MT,FT,T} + parts = Vector{Expr}(undef, fieldcount(T)) + fi = 1 + for i in 1:fieldcount(T) + if fieldtypes(MT)[i] === Colon + parts[i] = :(free_key[$fi]) + fi += 1 + else + parts[i] = :(mask[$i]) + end + end + return :($(Expr(:tuple, parts...))::$T) +end + +# Returns the matching full keys from the parent's cache. +function _view_matching_keys(v::IndexedVarArrayView{V,N,T})::Vector{T} where {V,N,T} + return _select_cached(v.parent, v.mask) +end + +function Base.iterate(v::IndexedVarArrayView{V,N,T,MT,FT}) where {V,N,T,MT,FT} + matching = _view_matching_keys(v) + isempty(matching) && return nothing + key = matching[1] + return (_project_key(key, MT) => v.parent[key], (matching, 2)) +end + +function Base.iterate( + v::IndexedVarArrayView{V,N,T,MT,FT}, + state::Tuple{Vector{T},Int}, +) where {V,N,T,MT,FT} + matching, pos = state + pos > length(matching) && return nothing + key = matching[pos] + return (_project_key(key, MT) => v.parent[key], (matching, pos + 1)) +end + +function Base.getindex( + v::IndexedVarArrayView{V,N,T,MT,FT}, + free_key::FT, +) where {V,N,T,MT,FT} + return v.parent[_reconstruct_key(v.mask, free_key, T)] +end + +function Base.haskey( + v::IndexedVarArrayView{V,N,T,MT,FT}, + free_key::FT, +) where {V,N,T,MT,FT} + return haskey(_data(v.parent), _reconstruct_key(v.mask, free_key, T)) +end + +Base.length(v::IndexedVarArrayView) = length(_view_matching_keys(v)) + +Base.keys(v::IndexedVarArrayView{V,N,T,MT,FT}) where {V,N,T,MT,FT} = + [_project_key(k, MT) for k in _view_matching_keys(v)] + +Base.values(v::IndexedVarArrayView) = [v.parent[k] for k in _view_matching_keys(v)] + +""" + sum(v::IndexedVarArrayView) + +Sum the variable refs in the view. Returns `zero(V)` for an empty view, +preserving the same behaviour as `sum(iva[mask...])` before views were introduced. +""" +function Base.sum(v::IndexedVarArrayView{V}) where {V} + result = zero(AffExpr) + for k in _view_matching_keys(v) + JuMP.add_to_expression!(result, v.parent[k]) + end + return result +end + +""" + IndexedVarArrayViewAlt{V,N,T,NF,MT,FT} + +A lazy, filtered view into an `IndexedVarArray` implementing the +`AbstractSparseArray{V,NF}` interface, where `NF` is the number of free +(Colon) dimensions. Keys are projected tuples covering only the free dimensions. + +Unlike `IndexedVarArrayView`, iteration yields **values only** (not key-value +pairs), matching standard `AbstractArray` semantics. Use `pairs(v)` or +`eachindex(v)` to access projected keys alongside values. + +Create via `filter_view_alt(iva, mask...)`. +""" +struct IndexedVarArrayViewAlt{V<:AbstractVariableRef,N,T,NF,MT<:Tuple,FT<:Tuple} <: + AbstractSparseArray{V,NF} + parent::IndexedVarArray{V,N,T} + mask::MT +end + +""" + filter_view_alt(iva::IndexedVarArray, mask...) + +Return a lazy `IndexedVarArrayViewAlt` over entries of `iva` matching `mask`. +Use `:` for free (wildcard) dimensions and exact values for fixed dimensions. + +Unlike `filter_view`, the result is an `AbstractSparseArray{V,NF}` where `NF` +is the number of free dimensions. Iterates values only; use `pairs(v)` for +projected-key/value pairs, or `eachindex(v)` for projected keys. + +# Example +```julia +v = filter_view_alt(flow, :, c, p, t) # NF=1, one free dimension +sum(v) # sum of matching VariableRefs +for (k, var) in pairs(v); ...; end # k is a 1-tuple projected key +var = v[f, c] # splatted index lookup +``` +""" +function filter_view_alt(iva::IndexedVarArray{V,N,T}, mask...) where {V,N,T} + return _make_view_alt(iva, tuple(mask...)) +end + +@generated function _make_view_alt( + iva::IndexedVarArray{V,N,T}, + mask::MT, +) where {V,N,T,MT<:Tuple} + fieldcount(MT) != N && return :(throw(BoundsError(iva, mask))) + free = [fieldtypes(T)[i] for i in 1:N if fieldtypes(MT)[i] === Colon] + NF = length(free) + FT = Tuple{free...} + return :(IndexedVarArrayViewAlt{$V,$N,$T,$NF,$MT,$FT}(iva, mask)) +end + +function _view_matching_keys(v::IndexedVarArrayViewAlt{V,N,T})::Vector{T} where {V,N,T} + return _select_cached(v.parent, v.mask) +end + +# Iterator traits: length is known but size() is not meaningful +Base.IteratorSize(::Type{<:IndexedVarArrayViewAlt}) = Base.HasLength() +Base.IteratorEltype(::Type{<:IndexedVarArrayViewAlt}) = Base.HasEltype() +Base.eltype(::Type{<:IndexedVarArrayViewAlt{V}}) where {V} = V + +# Iteration: values only (AbstractArray semantics) +function Base.iterate(v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} + matching = _view_matching_keys(v) + isempty(matching) && return nothing + return (v.parent[matching[1]], (matching, 2)) +end + +function Base.iterate( + v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}, + state::Tuple{Vector{T},Int}, +) where {V,N,T,NF,MT,FT} + matching, pos = state + pos > length(matching) && return nothing + return (v.parent[matching[pos]], (matching, pos + 1)) +end + +# getindex by FT tuple: v[(f, c)] +function Base.getindex( + v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}, + free_key::FT, +) where {V,N,T,NF,MT,FT} + return v.parent[_reconstruct_key(v.mask, free_key, T)] +end + +# getindex by splatted args: v[f, c] or v[f] (NF==1) +# Generated at compile time — reconstructs the full key by interleaving +# the fixed mask values and the free positional arguments. +@generated function Base.getindex( + v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}, + idx..., +) where {V,N,T,NF,MT,FT} + if length(idx) != NF + return :(throw(BoundsError(v, idx))) + end + parts = Expr[] + fi = 1 + for i in 1:fieldcount(T) + if fieldtypes(MT)[i] === Colon + push!(parts, :(idx[$fi])) + fi += 1 + else + push!(parts, :(v.mask[$i])) + end + end + return :(v.parent[$(Expr(:tuple, parts...))]) +end + +# Block mutation — views are read-only +Base.setindex!(::IndexedVarArrayViewAlt, _, _...) = + error("IndexedVarArrayViewAlt is read-only") + +# size is not meaningful for sparse tuple-keyed arrays +function Base.size(::IndexedVarArrayViewAlt) + return error( + "`Base.size` is not implemented for `IndexedVarArrayViewAlt` because it " * + "is conceptually a sparse dictionary with NF-dimensional keys. " * + "Use `length` for the number of entries.", + ) +end + +function Base.haskey( + v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}, + free_key::FT, +) where {V,N,T,NF,MT,FT} + return haskey(_data(v.parent), _reconstruct_key(v.mask, free_key, T)) +end + +Base.length(v::IndexedVarArrayViewAlt) = length(_view_matching_keys(v)) + +Base.keys(v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} = + [_project_key(k, MT) for k in _view_matching_keys(v)] + +Base.values(v::IndexedVarArrayViewAlt) = [v.parent[k] for k in _view_matching_keys(v)] + +# eachindex returns projected FT tuples (same as keys) +Base.eachindex(v::IndexedVarArrayViewAlt) = keys(v) + +Base.pairs(v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} = + [_project_key(k, MT) => v.parent[k] for k in _view_matching_keys(v)] + +function Base.firstindex(v::IndexedVarArrayViewAlt, d) + return minimum(k[d] for k in _view_matching_keys(v)) +end +function Base.lastindex(v::IndexedVarArrayViewAlt, d) + return maximum(k[d] for k in _view_matching_keys(v)) +end + +# sum: build AffExpr directly — avoids _data() from AbstractSparseArray default +function Base.sum(v::IndexedVarArrayViewAlt{V}) where {V} + result = zero(AffExpr) + for k in _view_matching_keys(v) + JuMP.add_to_expression!(result, v.parent[k]) + end + return result +end + +# show: override AbstractSparseArray defaults which call _data() +Base.show(io::IO, ::MIME"text/plain", v::IndexedVarArrayViewAlt) = summary(io, v) +Base.show(io::IO, v::IndexedVarArrayViewAlt) = summary(io, v) From 33c59d169ca2dec7dfed9b316b11492b374b8429 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Fri, 22 May 2026 11:36:30 +0200 Subject: [PATCH 02/22] Add parameter for cache cutoff --- src/SparseVariables.jl | 1 + src/indexedarray.jl | 18 ++++++++++++++++-- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/src/SparseVariables.jl b/src/SparseVariables.jl index c1cfbab..6b7e76b 100644 --- a/src/SparseVariables.jl +++ b/src/SparseVariables.jl @@ -19,6 +19,7 @@ export filter_view_alt export insertvar! export unsafe_insertvar! export SafeInsert, UnsafeInsert +export set_cache_cutoff! @setup_workload begin # Putting some things in `setup` can reduce the size of the diff --git a/src/indexedarray.jl b/src/indexedarray.jl index f3a0959..876f978 100644 --- a/src/indexedarray.jl +++ b/src/indexedarray.jl @@ -103,9 +103,23 @@ function build_cache!(cache, pat, sa::IndexedVarArray{V,N,T}) where {V,N,T} return cache end +# Minimum number of entries before the index cache is used; below this a +# linear scan is cheaper. Tune with set_cache_cutoff! or calibrate with +# benchmark/cutoff_benchmark.jl. +const _CACHE_CUTOFF = Ref{Int}(100) + +""" + set_cache_cutoff!(n::Int) + +Set the minimum number of entries in an `IndexedVarArray` at which +`filter_view` / `filter_view_alt` switch from a linear scan to the pre-built +index cache. Smaller values favour caching; larger values favour the linear +scan for small arrays. Default: `100`. +""" +set_cache_cutoff!(n::Int) = (_CACHE_CUTOFF[] = n; nothing) + function _select_cached(sa::IndexedVarArray{V,N,T}, pat)::Vector{T} where {V,N,T} - # TODO: Benchmark to find good cutoff-value for caching - length(_data(sa)) < 100 && return collect(T, _select_gen(keys(_data(sa)), pat)) + length(_data(sa)) < _CACHE_CUTOFF[] && return collect(T, _select_gen(keys(_data(sa)), pat)) cache = _getcache(sa, pat)::Dictionary{_decode_nonslices(sa, pat),Vector{T}} build_cache!(cache, pat, sa) vals = _dropslices_gen(pat) From e31044a2a3e2093e240caf5981ae6a9f0314c430 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Fri, 22 May 2026 13:00:46 +0200 Subject: [PATCH 03/22] Renaming and removing filtering with dict behaviour --- src/SparseVariables.jl | 6 +- src/indexedarray.jl | 353 ++++++++++++++++++++++++----------------- 2 files changed, 205 insertions(+), 154 deletions(-) diff --git a/src/SparseVariables.jl b/src/SparseVariables.jl index 6b7e76b..e3b81ac 100644 --- a/src/SparseVariables.jl +++ b/src/SparseVariables.jl @@ -12,10 +12,8 @@ include("tables.jl") export SparseArray export IndexedVarArray -export IndexedVarArrayView -export filter_view -export IndexedVarArrayViewAlt -export filter_view_alt +export IndexedVarArraySlice +export slice export insertvar! export unsafe_insertvar! export SafeInsert, UnsafeInsert diff --git a/src/indexedarray.jl b/src/indexedarray.jl index 876f978..fa95433 100644 --- a/src/indexedarray.jl +++ b/src/indexedarray.jl @@ -112,7 +112,7 @@ const _CACHE_CUTOFF = Ref{Int}(100) set_cache_cutoff!(n::Int) Set the minimum number of entries in an `IndexedVarArray` at which -`filter_view` / `filter_view_alt` switch from a linear scan to the pre-built +selection switches from a linear scan to the pre-built index cache. Smaller values favour caching; larger values favour the linear scan for small arrays. Default: `100`. """ @@ -224,58 +224,6 @@ function Base.lastindex(sa::IndexedVarArray, d) return last(sort(sa.index_names[d])) end -# ------------------------------------------------------------------------------ -# IndexedVarArrayView -# ------------------------------------------------------------------------------ - -""" - IndexedVarArrayView{V,N,T,MT,FT} - -A lazy, filtered view into an `IndexedVarArray` where some dimensions are fixed -to specific values and the rest are free (marked with `Colon`). Iterates as an -`AbstractDict` mapping projected keys (`FT`, covering only free dimensions) to -`V` (variable refs). Backed by the parent's `index_cache`, so iteration is O(1) -once the cache is warm. - -Create via `filter_view(iva, mask...)`. -""" -struct IndexedVarArrayView{V<:AbstractVariableRef,N,T,MT<:Tuple,FT<:Tuple} <: - AbstractDict{FT,V} - parent::IndexedVarArray{V,N,T} - mask::MT -end - -""" - filter_view(iva::IndexedVarArray, mask...) - -Return a lazy `IndexedVarArrayView` over entries of `iva` matching `mask`. -Use `:` for free (wildcard) dimensions and exact values for fixed dimensions. -The keys of the view are projected tuples covering only the free dimensions. - -Unlike `Base.view`, the result is an `AbstractDict{FT, V}` (not an -`AbstractArray`) where `FT` is a tuple of the free-dimension value types. - -# Example -```julia -v = filter_view(flow, :, c, p, t) # one free dim: factory -for ((f,), var) in v; ...; end # projected key is a 1-tuple -sum(values(v)) # sum of matching VariableRefs -``` -""" -function filter_view(iva::IndexedVarArray{V,N,T}, mask...) where {V,N,T} - return _make_view(iva, tuple(mask...)) -end - -@generated function _make_view( - iva::IndexedVarArray{V,N,T}, - mask::MT, -) where {V,N,T,MT<:Tuple} - fieldcount(MT) != N && return :(throw(BoundsError(iva, mask))) - free = [fieldtypes(T)[i] for i in 1:N if fieldtypes(MT)[i] === Colon] - FT = Tuple{free...} - return :(IndexedVarArrayView{$V,$N,$T,$MT,$FT}(iva, mask)) -end - # Project a full key T down to the free dimensions FT. @generated function _project_key(key::T, ::Type{MT}) where {T,MT} free_idx = [i for i in 1:fieldcount(T) if fieldtypes(MT)[i] === Colon] @@ -298,105 +246,45 @@ end return :($(Expr(:tuple, parts...))::$T) end -# Returns the matching full keys from the parent's cache. -function _view_matching_keys(v::IndexedVarArrayView{V,N,T})::Vector{T} where {V,N,T} - return _select_cached(v.parent, v.mask) -end - -function Base.iterate(v::IndexedVarArrayView{V,N,T,MT,FT}) where {V,N,T,MT,FT} - matching = _view_matching_keys(v) - isempty(matching) && return nothing - key = matching[1] - return (_project_key(key, MT) => v.parent[key], (matching, 2)) -end - -function Base.iterate( - v::IndexedVarArrayView{V,N,T,MT,FT}, - state::Tuple{Vector{T},Int}, -) where {V,N,T,MT,FT} - matching, pos = state - pos > length(matching) && return nothing - key = matching[pos] - return (_project_key(key, MT) => v.parent[key], (matching, pos + 1)) -end - -function Base.getindex( - v::IndexedVarArrayView{V,N,T,MT,FT}, - free_key::FT, -) where {V,N,T,MT,FT} - return v.parent[_reconstruct_key(v.mask, free_key, T)] -end - -function Base.haskey( - v::IndexedVarArrayView{V,N,T,MT,FT}, - free_key::FT, -) where {V,N,T,MT,FT} - return haskey(_data(v.parent), _reconstruct_key(v.mask, free_key, T)) -end - -Base.length(v::IndexedVarArrayView) = length(_view_matching_keys(v)) - -Base.keys(v::IndexedVarArrayView{V,N,T,MT,FT}) where {V,N,T,MT,FT} = - [_project_key(k, MT) for k in _view_matching_keys(v)] - -Base.values(v::IndexedVarArrayView) = [v.parent[k] for k in _view_matching_keys(v)] - -""" - sum(v::IndexedVarArrayView) - -Sum the variable refs in the view. Returns `zero(V)` for an empty view, -preserving the same behaviour as `sum(iva[mask...])` before views were introduced. -""" -function Base.sum(v::IndexedVarArrayView{V}) where {V} - result = zero(AffExpr) - for k in _view_matching_keys(v) - JuMP.add_to_expression!(result, v.parent[k]) - end - return result -end - """ - IndexedVarArrayViewAlt{V,N,T,NF,MT,FT} + IndexedVarArraySlice{V,N,T,NF,MT,FT} A lazy, filtered view into an `IndexedVarArray` implementing the `AbstractSparseArray{V,NF}` interface, where `NF` is the number of free (Colon) dimensions. Keys are projected tuples covering only the free dimensions. +Iterates values only; use `pairs(v)` or `eachindex(v)` for projected keys alongside values. -Unlike `IndexedVarArrayView`, iteration yields **values only** (not key-value -pairs), matching standard `AbstractArray` semantics. Use `pairs(v)` or -`eachindex(v)` to access projected keys alongside values. - -Create via `filter_view_alt(iva, mask...)`. +Create via `slice(iva, mask...)`. """ -struct IndexedVarArrayViewAlt{V<:AbstractVariableRef,N,T,NF,MT<:Tuple,FT<:Tuple} <: +struct IndexedVarArraySlice{V<:AbstractVariableRef,N,T,NF,MT<:Tuple,FT<:Tuple} <: AbstractSparseArray{V,NF} parent::IndexedVarArray{V,N,T} mask::MT end """ - filter_view_alt(iva::IndexedVarArray, mask...) + slice(iva::IndexedVarArray, mask...) -Return a lazy `IndexedVarArrayViewAlt` over entries of `iva` matching `mask`. +Return a lazy `IndexedVarArraySlice` over entries of `iva` matching `mask`. Use `:` for free (wildcard) dimensions and exact values for fixed dimensions. -Unlike `filter_view`, the result is an `AbstractSparseArray{V,NF}` where `NF` -is the number of free dimensions. Iterates values only; use `pairs(v)` for -projected-key/value pairs, or `eachindex(v)` for projected keys. +The result is an `AbstractSparseArray{V,NF}` where `NF` is the number of free +dimensions. Iterates values only; use `pairs(v)` for projected-key/value pairs, +or `eachindex(v)` for projected keys. # Example ```julia -v = filter_view_alt(flow, :, c, p, t) # NF=1, one free dimension +v = slice(flow, :, c, p, t) # NF=1, one free dimension sum(v) # sum of matching VariableRefs for (k, var) in pairs(v); ...; end # k is a 1-tuple projected key var = v[f, c] # splatted index lookup ``` """ -function filter_view_alt(iva::IndexedVarArray{V,N,T}, mask...) where {V,N,T} - return _make_view_alt(iva, tuple(mask...)) +function slice(iva::IndexedVarArray{V,N,T}, mask...) where {V,N,T} + return _make_slice(iva, tuple(mask...)) end -@generated function _make_view_alt( +@generated function _make_slice( iva::IndexedVarArray{V,N,T}, mask::MT, ) where {V,N,T,MT<:Tuple} @@ -404,27 +292,27 @@ end free = [fieldtypes(T)[i] for i in 1:N if fieldtypes(MT)[i] === Colon] NF = length(free) FT = Tuple{free...} - return :(IndexedVarArrayViewAlt{$V,$N,$T,$NF,$MT,$FT}(iva, mask)) + return :(IndexedVarArraySlice{$V,$N,$T,$NF,$MT,$FT}(iva, mask)) end -function _view_matching_keys(v::IndexedVarArrayViewAlt{V,N,T})::Vector{T} where {V,N,T} +function _view_matching_keys(v::IndexedVarArraySlice{V,N,T})::Vector{T} where {V,N,T} return _select_cached(v.parent, v.mask) end # Iterator traits: length is known but size() is not meaningful -Base.IteratorSize(::Type{<:IndexedVarArrayViewAlt}) = Base.HasLength() -Base.IteratorEltype(::Type{<:IndexedVarArrayViewAlt}) = Base.HasEltype() -Base.eltype(::Type{<:IndexedVarArrayViewAlt{V}}) where {V} = V +Base.IteratorSize(::Type{<:IndexedVarArraySlice}) = Base.HasLength() +Base.IteratorEltype(::Type{<:IndexedVarArraySlice}) = Base.HasEltype() +Base.eltype(::Type{<:IndexedVarArraySlice{V}}) where {V} = V # Iteration: values only (AbstractArray semantics) -function Base.iterate(v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} +function Base.iterate(v::IndexedVarArraySlice{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} matching = _view_matching_keys(v) isempty(matching) && return nothing return (v.parent[matching[1]], (matching, 2)) end function Base.iterate( - v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}, + v::IndexedVarArraySlice{V,N,T,NF,MT,FT}, state::Tuple{Vector{T},Int}, ) where {V,N,T,NF,MT,FT} matching, pos = state @@ -434,17 +322,28 @@ end # getindex by FT tuple: v[(f, c)] function Base.getindex( - v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}, + v::IndexedVarArraySlice{V,N,T,NF,MT,FT}, free_key::FT, ) where {V,N,T,NF,MT,FT} return v.parent[_reconstruct_key(v.mask, free_key, T)] end +# Disambiguate: AbstractSparseArray defines getindex(sa, ::NTuple{N,Any}) where N=NF, +# which overlaps with the FT method above when FT <: NTuple{NF,Any}. +# This more-specific overload resolves the ambiguity for broadcasting and other +# callers that hold the key as an unparameterised NTuple. +function Base.getindex( + v::IndexedVarArraySlice{V,N,T,NF,MT,FT}, + idx::NTuple{NF,Any}, +) where {V,N,T,NF,MT,FT} + return v.parent[_reconstruct_key(v.mask, idx, T)] +end + # getindex by splatted args: v[f, c] or v[f] (NF==1) # Generated at compile time — reconstructs the full key by interleaving # the fixed mask values and the free positional arguments. @generated function Base.getindex( - v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}, + v::IndexedVarArraySlice{V,N,T,NF,MT,FT}, idx..., ) where {V,N,T,NF,MT,FT} if length(idx) != NF @@ -464,47 +363,47 @@ end end # Block mutation — views are read-only -Base.setindex!(::IndexedVarArrayViewAlt, _, _...) = - error("IndexedVarArrayViewAlt is read-only") +Base.setindex!(::IndexedVarArraySlice, _, _...) = + error("IndexedVarArraySlice is read-only") # size is not meaningful for sparse tuple-keyed arrays -function Base.size(::IndexedVarArrayViewAlt) +function Base.size(::IndexedVarArraySlice) return error( - "`Base.size` is not implemented for `IndexedVarArrayViewAlt` because it " * + "`Base.size` is not implemented for `IndexedVarArraySlice` because it " * "is conceptually a sparse dictionary with NF-dimensional keys. " * "Use `length` for the number of entries.", ) end function Base.haskey( - v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}, + v::IndexedVarArraySlice{V,N,T,NF,MT,FT}, free_key::FT, ) where {V,N,T,NF,MT,FT} return haskey(_data(v.parent), _reconstruct_key(v.mask, free_key, T)) end -Base.length(v::IndexedVarArrayViewAlt) = length(_view_matching_keys(v)) +Base.length(v::IndexedVarArraySlice) = length(_view_matching_keys(v)) -Base.keys(v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} = +Base.keys(v::IndexedVarArraySlice{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} = [_project_key(k, MT) for k in _view_matching_keys(v)] -Base.values(v::IndexedVarArrayViewAlt) = [v.parent[k] for k in _view_matching_keys(v)] +Base.values(v::IndexedVarArraySlice) = [v.parent[k] for k in _view_matching_keys(v)] # eachindex returns projected FT tuples (same as keys) -Base.eachindex(v::IndexedVarArrayViewAlt) = keys(v) +Base.eachindex(v::IndexedVarArraySlice) = keys(v) -Base.pairs(v::IndexedVarArrayViewAlt{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} = +Base.pairs(v::IndexedVarArraySlice{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} = [_project_key(k, MT) => v.parent[k] for k in _view_matching_keys(v)] -function Base.firstindex(v::IndexedVarArrayViewAlt, d) +function Base.firstindex(v::IndexedVarArraySlice, d) return minimum(k[d] for k in _view_matching_keys(v)) end -function Base.lastindex(v::IndexedVarArrayViewAlt, d) +function Base.lastindex(v::IndexedVarArraySlice, d) return maximum(k[d] for k in _view_matching_keys(v)) end # sum: build AffExpr directly — avoids _data() from AbstractSparseArray default -function Base.sum(v::IndexedVarArrayViewAlt{V}) where {V} +function Base.sum(v::IndexedVarArraySlice{V}) where {V} result = zero(AffExpr) for k in _view_matching_keys(v) JuMP.add_to_expression!(result, v.parent[k]) @@ -513,5 +412,159 @@ function Base.sum(v::IndexedVarArrayViewAlt{V}) where {V} end # show: override AbstractSparseArray defaults which call _data() -Base.show(io::IO, ::MIME"text/plain", v::IndexedVarArrayViewAlt) = summary(io, v) -Base.show(io::IO, v::IndexedVarArrayViewAlt) = summary(io, v) +Base.show(io::IO, ::MIME"text/plain", v::IndexedVarArraySlice) = summary(io, v) +Base.show(io::IO, v::IndexedVarArraySlice) = summary(io, v) + +# ------------------------------------------------------------------------------ +# Broadcasting +# Follows the pattern of JuMP.Containers.SparseAxisArray. +# The result of any broadcast over these types is always a plain SparseArray. +# ------------------------------------------------------------------------------ + +""" + IVABroadcastStyle{N,K} <: Broadcast.BroadcastStyle + +Shared broadcasting style for `IndexedVarArray` and `IndexedVarArraySlice`. +`N` is the effective key dimensionality and `K` is the concrete key tuple type. +All broadcast results are materialised as `SparseArray`. +""" +struct IVABroadcastStyle{N,K} <: Broadcast.BroadcastStyle end + +Base.BroadcastStyle(::Type{<:IndexedVarArray{V,N,T}}) where {V,N,T} = + IVABroadcastStyle{N,T}() + +Base.BroadcastStyle( + ::Type{<:IndexedVarArraySlice{V,N,T,NF,MT,FT}}, +) where {V,N,T,NF,MT,FT} = IVABroadcastStyle{NF,FT}() + +# Disallow mixing with other array types. +function Base.BroadcastStyle(::IVABroadcastStyle, ::Base.BroadcastStyle) + return throw( + ArgumentError( + "Cannot broadcast IndexedVarArray or a view with another array of a different type", + ), + ) +end + +# Scalar (0-d) broadcasting is allowed. +function Base.BroadcastStyle( + style::IVABroadcastStyle, + ::Base.Broadcast.DefaultArrayStyle{0}, +) + return style +end + +# Fix ambiguity with Unknown. +function Base.BroadcastStyle(::IVABroadcastStyle, ::Base.Broadcast.Unknown) + return throw( + ArgumentError( + "Cannot broadcast IndexedVarArray or a view with an unknown broadcast style", + ), + ) +end + +# Bypass the default instantiate which calls axes(). +function Base.Broadcast.instantiate( + bc::Base.Broadcast.Broadcasted{<:IVABroadcastStyle}, +) + return bc +end + +# ── Internal helpers ────────────────────────────────────────────────────────── + +# Apply a broadcast tree to a single key. +_iva_getindex(x::IndexedVarArray, key) = x[key] +_iva_getindex(x::IndexedVarArraySlice, key) = x[key] +_iva_getindex(x::Any, ::Any) = x +_iva_getindex(x::Ref, ::Any) = x[] + +function _iva_getindex( + bc::Base.Broadcast.Broadcasted{<:IVABroadcastStyle}, + key, +) + return bc.f(_iva_get_args(bc.args, key)...) +end + +function _iva_get_args(args::Tuple, key) + return (_iva_getindex(first(args), key), _iva_get_args(Base.tail(args), key)...) +end +_iva_get_args(::Tuple{}, ::Any) = () + +# Verify x has the same key set as ref_keys. +function _iva_check_same_keys(ref_keys, x::IndexedVarArray, args...) + if length(ref_keys) != length(_data(x)) || + any(k -> !haskey(_data(x), k), ref_keys) + throw(ArgumentError("Cannot broadcast IndexedVarArrays with different indices")) + end + return _iva_check_same_keys(ref_keys, args...) +end + +function _iva_check_same_keys( + ref_keys, + x::IndexedVarArraySlice, + args..., +) + if length(ref_keys) != length(x) || any(k -> !haskey(x, k), ref_keys) + throw( + ArgumentError("Cannot broadcast IndexedVarArray views with different indices"), + ) + end + return _iva_check_same_keys(ref_keys, args...) +end + +_iva_check_same_keys(ref_keys, ::Any, args...) = _iva_check_same_keys(ref_keys, args...) +_iva_check_same_keys(::Any) = nothing + +# Recursively extract the key set from the first IVA-family object found. +function _iva_indices(bc::Base.Broadcast.Broadcasted{<:IVABroadcastStyle}, rest...) + return _iva_indices(bc.args..., rest...) +end + +function _iva_indices(x::IndexedVarArray, rest...) + ks = collect(keys(_data(x))) + _iva_check_same_keys(ks, rest...) + return ks +end + +function _iva_indices(x::IndexedVarArraySlice, rest...) + ks = keys(x) # Vector{FT} + _iva_check_same_keys(ks, rest...) + return ks +end + +_iva_indices(::Any, rest...) = _iva_indices(rest...) # skip scalars + +# ── Materialise ─────────────────────────────────────────────────────────────── + +function Base.copy( + bc::Base.Broadcast.Broadcasted{IVABroadcastStyle{N,K}}, +) where {N,K} + indices = _iva_indices(bc) + isempty(indices) && return SparseArray(Dictionary{K,Any}()) + vals = [_iva_getindex(bc, k) for k in indices] + return SparseArray(Dictionary(indices, vals)) +end + +# Prevent scalar broadcast from reducing to a 0-d result. +for _IVAType in ( + :IndexedVarArray, + :IndexedVarArraySlice, +) + @eval begin + function Base.Broadcast.broadcast_preserving_zero_d( + f, + A::$_IVAType, + As..., + ) + return broadcast(f, A, As...) + end + function Base.Broadcast.broadcast_preserving_zero_d( + f, + x, + A::$_IVAType, + As..., + ) + return broadcast(f, x, A, As...) + end + end +end From 26888ce358ef95839e9c284270410ea0f4ee7bbb Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Fri, 22 May 2026 14:06:27 +0200 Subject: [PATCH 04/22] Move broadcasting and slicing into separate files --- src/SparseVariables.jl | 4 +- src/broadcast.jl | 112 ++++++++++++++ src/indexedarray.jl | 342 +---------------------------------------- src/slice.jl | 198 ++++++++++++++++++++++++ src/sparsearray.jl | 18 +++ 5 files changed, 338 insertions(+), 336 deletions(-) create mode 100644 src/broadcast.jl create mode 100644 src/slice.jl diff --git a/src/SparseVariables.jl b/src/SparseVariables.jl index e3b81ac..c139bf7 100644 --- a/src/SparseVariables.jl +++ b/src/SparseVariables.jl @@ -6,13 +6,15 @@ using LinearAlgebra using PrecompileTools include("sparsearray.jl") +include("slice.jl") +include("broadcast.jl") include("dictionaries.jl") include("indexedarray.jl") include("tables.jl") export SparseArray export IndexedVarArray -export IndexedVarArraySlice +export SparseArraySlice export slice export insertvar! export unsafe_insertvar! diff --git a/src/broadcast.jl b/src/broadcast.jl new file mode 100644 index 0000000..5bc1a9b --- /dev/null +++ b/src/broadcast.jl @@ -0,0 +1,112 @@ +# ------------------------------------------------------------------------------ +# Broadcasting over AbstractSparseArray +# Follows the pattern of JuMP.Containers.SparseAxisArray. +# The result of any broadcast is always a plain SparseArray. +# ------------------------------------------------------------------------------ + +""" + SparseBroadcastStyle{N,K} <: Broadcast.BroadcastStyle + +Broadcasting style for all `AbstractSparseArray` subtypes. `N` is the key +dimensionality and `K` is the key tuple type. All broadcast results are +materialised as `SparseArray`. +""" +struct SparseBroadcastStyle{N,K} <: Broadcast.BroadcastStyle end + +Base.BroadcastStyle(::Type{SA}) where {SA<:AbstractSparseArray} = + SparseBroadcastStyle{ndims(SA),_keytype(SA)}() + +# Disallow mixing with other array types. +function Base.BroadcastStyle(::SparseBroadcastStyle, ::Base.BroadcastStyle) + return throw( + ArgumentError( + "Cannot broadcast a SparseArray with another array of a different type", + ), + ) +end + +# Scalar (0-d) broadcasting is allowed. +function Base.BroadcastStyle( + style::SparseBroadcastStyle, + ::Base.Broadcast.DefaultArrayStyle{0}, +) + return style +end + +# Fix ambiguity with Unknown. +function Base.BroadcastStyle(::SparseBroadcastStyle, ::Base.Broadcast.Unknown) + return throw( + ArgumentError( + "Cannot broadcast a SparseArray with an unknown broadcast style", + ), + ) +end + +# Bypass the default instantiate which calls axes(). +function Base.Broadcast.instantiate( + bc::Base.Broadcast.Broadcasted{<:SparseBroadcastStyle}, +) + return bc +end + +# ── Internal helpers ────────────────────────────────────────────────────────── + +_sparse_getindex(x::AbstractSparseArray, key) = x[key] +_sparse_getindex(x::Any, ::Any) = x +_sparse_getindex(x::Ref, ::Any) = x[] + +function _sparse_getindex( + bc::Base.Broadcast.Broadcasted{<:SparseBroadcastStyle}, + key, +) + return bc.f(_sparse_get_args(bc.args, key)...) +end + +function _sparse_get_args(args::Tuple, key) + return (_sparse_getindex(first(args), key), _sparse_get_args(Base.tail(args), key)...) +end +_sparse_get_args(::Tuple{}, ::Any) = () + +function _sparse_check_same_keys(ref_keys, x::AbstractSparseArray, args...) + if length(ref_keys) != length(x) || any(k -> !haskey(x, k), ref_keys) + throw(ArgumentError("Cannot broadcast SparseArrays with different indices")) + end + return _sparse_check_same_keys(ref_keys, args...) +end + +_sparse_check_same_keys(ref_keys, ::Any, args...) = + _sparse_check_same_keys(ref_keys, args...) +_sparse_check_same_keys(::Any) = nothing + +function _sparse_indices( + bc::Base.Broadcast.Broadcasted{<:SparseBroadcastStyle}, + rest..., +) + return _sparse_indices(bc.args..., rest...) +end + +function _sparse_indices(x::AbstractSparseArray, rest...) + ks = collect(keys(x)) + _sparse_check_same_keys(ks, rest...) + return ks +end + +_sparse_indices(::Any, rest...) = _sparse_indices(rest...) + +# ── Materialise ─────────────────────────────────────────────────────────────── + +function Base.copy( + bc::Base.Broadcast.Broadcasted{SparseBroadcastStyle{N,K}}, +) where {N,K} + indices = _sparse_indices(bc) + isempty(indices) && return SparseArray(Dictionary{K,Any}()) + vals = [_sparse_getindex(bc, k) for k in indices] + return SparseArray(Dictionary(indices, vals)) +end + +function Base.Broadcast.broadcast_preserving_zero_d(f, A::AbstractSparseArray, As...) + return broadcast(f, A, As...) +end +function Base.Broadcast.broadcast_preserving_zero_d(f, x, A::AbstractSparseArray, As...) + return broadcast(f, x, A, As...) +end diff --git a/src/indexedarray.jl b/src/indexedarray.jl index fa95433..20772c3 100644 --- a/src/indexedarray.jl +++ b/src/indexedarray.jl @@ -14,6 +14,7 @@ struct SafeInsert end struct UnsafeInsert end _data(sa::IndexedVarArray) = sa.data +_keytype(::Type{<:IndexedVarArray{V,N,T}}) where {V,N,T} = T already_defined(var, index) = haskey(_data(var), index) @@ -224,347 +225,18 @@ function Base.lastindex(sa::IndexedVarArray, d) return last(sort(sa.index_names[d])) end -# Project a full key T down to the free dimensions FT. -@generated function _project_key(key::T, ::Type{MT}) where {T,MT} - free_idx = [i for i in 1:fieldcount(T) if fieldtypes(MT)[i] === Colon] - FT = Tuple{[fieldtypes(T)[i] for i in free_idx]...} - return :($(Expr(:tuple, [:(key[$i]) for i in free_idx]...))::$FT) -end - -# Reconstruct a full key T from the fixed values in mask and the free key FT. -@generated function _reconstruct_key(mask::MT, free_key::FT, ::Type{T}) where {MT,FT,T} - parts = Vector{Expr}(undef, fieldcount(T)) - fi = 1 - for i in 1:fieldcount(T) - if fieldtypes(MT)[i] === Colon - parts[i] = :(free_key[$fi]) - fi += 1 - else - parts[i] = :(mask[$i]) - end - end - return :($(Expr(:tuple, parts...))::$T) -end - -""" - IndexedVarArraySlice{V,N,T,NF,MT,FT} - -A lazy, filtered view into an `IndexedVarArray` implementing the -`AbstractSparseArray{V,NF}` interface, where `NF` is the number of free -(Colon) dimensions. Keys are projected tuples covering only the free dimensions. -Iterates values only; use `pairs(v)` or `eachindex(v)` for projected keys alongside values. - -Create via `slice(iva, mask...)`. -""" -struct IndexedVarArraySlice{V<:AbstractVariableRef,N,T,NF,MT<:Tuple,FT<:Tuple} <: - AbstractSparseArray{V,NF} - parent::IndexedVarArray{V,N,T} - mask::MT -end - -""" - slice(iva::IndexedVarArray, mask...) - -Return a lazy `IndexedVarArraySlice` over entries of `iva` matching `mask`. -Use `:` for free (wildcard) dimensions and exact values for fixed dimensions. - -The result is an `AbstractSparseArray{V,NF}` where `NF` is the number of free -dimensions. Iterates values only; use `pairs(v)` for projected-key/value pairs, -or `eachindex(v)` for projected keys. - -# Example -```julia -v = slice(flow, :, c, p, t) # NF=1, one free dimension -sum(v) # sum of matching VariableRefs -for (k, var) in pairs(v); ...; end # k is a 1-tuple projected key -var = v[f, c] # splatted index lookup -``` -""" -function slice(iva::IndexedVarArray{V,N,T}, mask...) where {V,N,T} - return _make_slice(iva, tuple(mask...)) -end - -@generated function _make_slice( - iva::IndexedVarArray{V,N,T}, - mask::MT, -) where {V,N,T,MT<:Tuple} - fieldcount(MT) != N && return :(throw(BoundsError(iva, mask))) - free = [fieldtypes(T)[i] for i in 1:N if fieldtypes(MT)[i] === Colon] - NF = length(free) - FT = Tuple{free...} - return :(IndexedVarArraySlice{$V,$N,$T,$NF,$MT,$FT}(iva, mask)) -end - -function _view_matching_keys(v::IndexedVarArraySlice{V,N,T})::Vector{T} where {V,N,T} +# Override _view_matching_keys for IndexedVarArray parent: use index cache. +function _view_matching_keys( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, +)::Vector{T} where {P<:IndexedVarArray,V,N,T,NF,MT,FT} return _select_cached(v.parent, v.mask) end -# Iterator traits: length is known but size() is not meaningful -Base.IteratorSize(::Type{<:IndexedVarArraySlice}) = Base.HasLength() -Base.IteratorEltype(::Type{<:IndexedVarArraySlice}) = Base.HasEltype() -Base.eltype(::Type{<:IndexedVarArraySlice{V}}) where {V} = V - -# Iteration: values only (AbstractArray semantics) -function Base.iterate(v::IndexedVarArraySlice{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} - matching = _view_matching_keys(v) - isempty(matching) && return nothing - return (v.parent[matching[1]], (matching, 2)) -end - -function Base.iterate( - v::IndexedVarArraySlice{V,N,T,NF,MT,FT}, - state::Tuple{Vector{T},Int}, -) where {V,N,T,NF,MT,FT} - matching, pos = state - pos > length(matching) && return nothing - return (v.parent[matching[pos]], (matching, pos + 1)) -end - -# getindex by FT tuple: v[(f, c)] -function Base.getindex( - v::IndexedVarArraySlice{V,N,T,NF,MT,FT}, - free_key::FT, -) where {V,N,T,NF,MT,FT} - return v.parent[_reconstruct_key(v.mask, free_key, T)] -end - -# Disambiguate: AbstractSparseArray defines getindex(sa, ::NTuple{N,Any}) where N=NF, -# which overlaps with the FT method above when FT <: NTuple{NF,Any}. -# This more-specific overload resolves the ambiguity for broadcasting and other -# callers that hold the key as an unparameterised NTuple. -function Base.getindex( - v::IndexedVarArraySlice{V,N,T,NF,MT,FT}, - idx::NTuple{NF,Any}, -) where {V,N,T,NF,MT,FT} - return v.parent[_reconstruct_key(v.mask, idx, T)] -end - -# getindex by splatted args: v[f, c] or v[f] (NF==1) -# Generated at compile time — reconstructs the full key by interleaving -# the fixed mask values and the free positional arguments. -@generated function Base.getindex( - v::IndexedVarArraySlice{V,N,T,NF,MT,FT}, - idx..., -) where {V,N,T,NF,MT,FT} - if length(idx) != NF - return :(throw(BoundsError(v, idx))) - end - parts = Expr[] - fi = 1 - for i in 1:fieldcount(T) - if fieldtypes(MT)[i] === Colon - push!(parts, :(idx[$fi])) - fi += 1 - else - push!(parts, :(v.mask[$i])) - end - end - return :(v.parent[$(Expr(:tuple, parts...))]) -end - -# Block mutation — views are read-only -Base.setindex!(::IndexedVarArraySlice, _, _...) = - error("IndexedVarArraySlice is read-only") - -# size is not meaningful for sparse tuple-keyed arrays -function Base.size(::IndexedVarArraySlice) - return error( - "`Base.size` is not implemented for `IndexedVarArraySlice` because it " * - "is conceptually a sparse dictionary with NF-dimensional keys. " * - "Use `length` for the number of entries.", - ) -end - -function Base.haskey( - v::IndexedVarArraySlice{V,N,T,NF,MT,FT}, - free_key::FT, -) where {V,N,T,NF,MT,FT} - return haskey(_data(v.parent), _reconstruct_key(v.mask, free_key, T)) -end - -Base.length(v::IndexedVarArraySlice) = length(_view_matching_keys(v)) - -Base.keys(v::IndexedVarArraySlice{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} = - [_project_key(k, MT) for k in _view_matching_keys(v)] - -Base.values(v::IndexedVarArraySlice) = [v.parent[k] for k in _view_matching_keys(v)] - -# eachindex returns projected FT tuples (same as keys) -Base.eachindex(v::IndexedVarArraySlice) = keys(v) - -Base.pairs(v::IndexedVarArraySlice{V,N,T,NF,MT,FT}) where {V,N,T,NF,MT,FT} = - [_project_key(k, MT) => v.parent[k] for k in _view_matching_keys(v)] - -function Base.firstindex(v::IndexedVarArraySlice, d) - return minimum(k[d] for k in _view_matching_keys(v)) -end -function Base.lastindex(v::IndexedVarArraySlice, d) - return maximum(k[d] for k in _view_matching_keys(v)) -end - -# sum: build AffExpr directly — avoids _data() from AbstractSparseArray default -function Base.sum(v::IndexedVarArraySlice{V}) where {V} +# JuMP-efficient sum: build AffExpr directly via add_to_expression!. +function Base.sum(v::SparseArraySlice{<:IndexedVarArray,V}) where {V<:AbstractVariableRef} result = zero(AffExpr) for k in _view_matching_keys(v) JuMP.add_to_expression!(result, v.parent[k]) end return result end - -# show: override AbstractSparseArray defaults which call _data() -Base.show(io::IO, ::MIME"text/plain", v::IndexedVarArraySlice) = summary(io, v) -Base.show(io::IO, v::IndexedVarArraySlice) = summary(io, v) - -# ------------------------------------------------------------------------------ -# Broadcasting -# Follows the pattern of JuMP.Containers.SparseAxisArray. -# The result of any broadcast over these types is always a plain SparseArray. -# ------------------------------------------------------------------------------ - -""" - IVABroadcastStyle{N,K} <: Broadcast.BroadcastStyle - -Shared broadcasting style for `IndexedVarArray` and `IndexedVarArraySlice`. -`N` is the effective key dimensionality and `K` is the concrete key tuple type. -All broadcast results are materialised as `SparseArray`. -""" -struct IVABroadcastStyle{N,K} <: Broadcast.BroadcastStyle end - -Base.BroadcastStyle(::Type{<:IndexedVarArray{V,N,T}}) where {V,N,T} = - IVABroadcastStyle{N,T}() - -Base.BroadcastStyle( - ::Type{<:IndexedVarArraySlice{V,N,T,NF,MT,FT}}, -) where {V,N,T,NF,MT,FT} = IVABroadcastStyle{NF,FT}() - -# Disallow mixing with other array types. -function Base.BroadcastStyle(::IVABroadcastStyle, ::Base.BroadcastStyle) - return throw( - ArgumentError( - "Cannot broadcast IndexedVarArray or a view with another array of a different type", - ), - ) -end - -# Scalar (0-d) broadcasting is allowed. -function Base.BroadcastStyle( - style::IVABroadcastStyle, - ::Base.Broadcast.DefaultArrayStyle{0}, -) - return style -end - -# Fix ambiguity with Unknown. -function Base.BroadcastStyle(::IVABroadcastStyle, ::Base.Broadcast.Unknown) - return throw( - ArgumentError( - "Cannot broadcast IndexedVarArray or a view with an unknown broadcast style", - ), - ) -end - -# Bypass the default instantiate which calls axes(). -function Base.Broadcast.instantiate( - bc::Base.Broadcast.Broadcasted{<:IVABroadcastStyle}, -) - return bc -end - -# ── Internal helpers ────────────────────────────────────────────────────────── - -# Apply a broadcast tree to a single key. -_iva_getindex(x::IndexedVarArray, key) = x[key] -_iva_getindex(x::IndexedVarArraySlice, key) = x[key] -_iva_getindex(x::Any, ::Any) = x -_iva_getindex(x::Ref, ::Any) = x[] - -function _iva_getindex( - bc::Base.Broadcast.Broadcasted{<:IVABroadcastStyle}, - key, -) - return bc.f(_iva_get_args(bc.args, key)...) -end - -function _iva_get_args(args::Tuple, key) - return (_iva_getindex(first(args), key), _iva_get_args(Base.tail(args), key)...) -end -_iva_get_args(::Tuple{}, ::Any) = () - -# Verify x has the same key set as ref_keys. -function _iva_check_same_keys(ref_keys, x::IndexedVarArray, args...) - if length(ref_keys) != length(_data(x)) || - any(k -> !haskey(_data(x), k), ref_keys) - throw(ArgumentError("Cannot broadcast IndexedVarArrays with different indices")) - end - return _iva_check_same_keys(ref_keys, args...) -end - -function _iva_check_same_keys( - ref_keys, - x::IndexedVarArraySlice, - args..., -) - if length(ref_keys) != length(x) || any(k -> !haskey(x, k), ref_keys) - throw( - ArgumentError("Cannot broadcast IndexedVarArray views with different indices"), - ) - end - return _iva_check_same_keys(ref_keys, args...) -end - -_iva_check_same_keys(ref_keys, ::Any, args...) = _iva_check_same_keys(ref_keys, args...) -_iva_check_same_keys(::Any) = nothing - -# Recursively extract the key set from the first IVA-family object found. -function _iva_indices(bc::Base.Broadcast.Broadcasted{<:IVABroadcastStyle}, rest...) - return _iva_indices(bc.args..., rest...) -end - -function _iva_indices(x::IndexedVarArray, rest...) - ks = collect(keys(_data(x))) - _iva_check_same_keys(ks, rest...) - return ks -end - -function _iva_indices(x::IndexedVarArraySlice, rest...) - ks = keys(x) # Vector{FT} - _iva_check_same_keys(ks, rest...) - return ks -end - -_iva_indices(::Any, rest...) = _iva_indices(rest...) # skip scalars - -# ── Materialise ─────────────────────────────────────────────────────────────── - -function Base.copy( - bc::Base.Broadcast.Broadcasted{IVABroadcastStyle{N,K}}, -) where {N,K} - indices = _iva_indices(bc) - isempty(indices) && return SparseArray(Dictionary{K,Any}()) - vals = [_iva_getindex(bc, k) for k in indices] - return SparseArray(Dictionary(indices, vals)) -end - -# Prevent scalar broadcast from reducing to a 0-d result. -for _IVAType in ( - :IndexedVarArray, - :IndexedVarArraySlice, -) - @eval begin - function Base.Broadcast.broadcast_preserving_zero_d( - f, - A::$_IVAType, - As..., - ) - return broadcast(f, A, As...) - end - function Base.Broadcast.broadcast_preserving_zero_d( - f, - x, - A::$_IVAType, - As..., - ) - return broadcast(f, x, A, As...) - end - end -end diff --git a/src/slice.jl b/src/slice.jl new file mode 100644 index 0000000..ae6d542 --- /dev/null +++ b/src/slice.jl @@ -0,0 +1,198 @@ +# Match key against mask at compile time; Colon positions are free (skipped). +@generated function _matches_mask(key::T, mask::MT) where {T,MT} + checks = Expr[] + for i in 1:fieldcount(T) + if fieldtypes(MT)[i] !== Colon + push!(checks, :(key[$i] != mask[$i] && return false)) + end + end + return quote + $(checks...) + return true + end +end + +# Project a full key T down to the free dimensions FT. +@generated function _project_key(key::T, ::Type{MT}) where {T,MT} + free_idx = [i for i in 1:fieldcount(T) if fieldtypes(MT)[i] === Colon] + FT = Tuple{[fieldtypes(T)[i] for i in free_idx]...} + return :($(Expr(:tuple, [:(key[$i]) for i in free_idx]...))::$FT) +end + +# Reconstruct a full key T from the fixed values in mask and the free key FT. +@generated function _reconstruct_key(mask::MT, free_key::FT, ::Type{T}) where {MT,FT,T} + parts = Vector{Expr}(undef, fieldcount(T)) + fi = 1 + for i in 1:fieldcount(T) + if fieldtypes(MT)[i] === Colon + parts[i] = :(free_key[$fi]) + fi += 1 + else + parts[i] = :(mask[$i]) + end + end + return :($(Expr(:tuple, parts...))::$T) +end + +""" + SparseArraySlice{P,V,N,T,NF,MT,FT} + +A lazy, mask-filtered view of any `AbstractSparseArray`. `P` is the concrete +parent type, `NF` is the number of free (Colon) dimensions, and `FT` is the +projected key tuple type covering only the free dimensions. Implements +`AbstractSparseArray{V,NF}`. + +Create via `slice(sa, mask...)`. +""" +struct SparseArraySlice{ + P<:AbstractSparseArray, + V, + N, + T, + NF, + MT<:Tuple, + FT<:Tuple, +} <: AbstractSparseArray{V,NF} + parent::P + mask::MT +end + +_keytype(::Type{<:SparseArraySlice{P,V,N,T,NF,MT,FT}}) where {P,V,N,T,NF,MT,FT} = FT + +""" + slice(sa::AbstractSparseArray, mask...) + +Return a lazy `SparseArraySlice` over entries of `sa` matching `mask`. Use `:` +for free (wildcard) dimensions and exact values for fixed dimensions. The result +is an `AbstractSparseArray{V,NF}` where `NF` is the number of free dimensions. + +# Example +```julia +v = slice(sa, :, "foo", :) # NF=2, two free dimensions +sum(v) +for (k, val) in pairs(v); ...; end +``` +""" +function slice(sa::AbstractSparseArray, mask...) + return _make_slice(sa, tuple(mask...)) +end + +@generated function _make_slice(sa::P, mask::MT) where {P<:AbstractSparseArray,MT<:Tuple} + K = _keytype(P) + N = ndims(P) + V = eltype(P) + fieldcount(MT) != N && return :(throw(BoundsError(sa, mask))) + free = [fieldtypes(K)[i] for i in 1:N if fieldtypes(MT)[i] === Colon] + NF = length(free) + FT = Tuple{free...} + return :(SparseArraySlice{$P,$V,$N,$K,$NF,$MT,$FT}(sa, mask)) +end + +# Default: linear scan. Subtypes may override for cached lookup. +function _view_matching_keys( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, +)::Vector{T} where {P,V,N,T,NF,MT,FT} + return [k for k in keys(_data(v.parent)) if _matches_mask(k, v.mask)] +end + +# Iterator traits +Base.IteratorSize(::Type{<:SparseArraySlice}) = Base.HasLength() +Base.IteratorEltype(::Type{<:SparseArraySlice}) = Base.HasEltype() +Base.eltype(::Type{<:SparseArraySlice{P,V}}) where {P,V} = V + +# Iteration: values only (AbstractArray semantics) +function Base.iterate(v::SparseArraySlice{P,V,N,T,NF,MT,FT}) where {P,V,N,T,NF,MT,FT} + matching = _view_matching_keys(v) + isempty(matching) && return nothing + return (v.parent[matching[1]], (matching, 2)) +end + +function Base.iterate( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, + state::Tuple{Vector{T},Int}, +) where {P,V,N,T,NF,MT,FT} + matching, pos = state + pos > length(matching) && return nothing + return (v.parent[matching[pos]], (matching, pos + 1)) +end + +# getindex by FT tuple: v[(f, c)] +function Base.getindex( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, + free_key::FT, +) where {P,V,N,T,NF,MT,FT} + return v.parent[_reconstruct_key(v.mask, free_key, T)] +end + +# Disambiguate vs AbstractSparseArray's NTuple method. +function Base.getindex( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, + idx::NTuple{NF,Any}, +) where {P,V,N,T,NF,MT,FT} + return v.parent[_reconstruct_key(v.mask, idx, T)] +end + +# Splatted: v[f, c] or v[f] (NF==1) +@generated function Base.getindex( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, + idx..., +) where {P,V,N,T,NF,MT,FT} + if length(idx) != NF + return :(throw(BoundsError(v, idx))) + end + parts = Expr[] + fi = 1 + for i in 1:fieldcount(T) + if fieldtypes(MT)[i] === Colon + push!(parts, :(idx[$fi])) + fi += 1 + else + push!(parts, :(v.mask[$i])) + end + end + return :(v.parent[$(Expr(:tuple, parts...))]) +end + +Base.setindex!(::SparseArraySlice, _, _...) = error("SparseArraySlice is read-only") + +function Base.size(::SparseArraySlice) + return error( + "`Base.size` is not implemented for `SparseArraySlice`. " * + "Use `length` for the number of entries.", + ) +end + +function Base.haskey( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, + free_key::FT, +) where {P,V,N,T,NF,MT,FT} + return haskey(_data(v.parent), _reconstruct_key(v.mask, free_key, T)) +end + +Base.length(v::SparseArraySlice) = length(_view_matching_keys(v)) + +Base.keys(v::SparseArraySlice{P,V,N,T,NF,MT,FT}) where {P,V,N,T,NF,MT,FT} = + [_project_key(k, MT) for k in _view_matching_keys(v)] + +Base.values(v::SparseArraySlice) = [v.parent[k] for k in _view_matching_keys(v)] + +Base.eachindex(v::SparseArraySlice) = keys(v) + +Base.pairs(v::SparseArraySlice{P,V,N,T,NF,MT,FT}) where {P,V,N,T,NF,MT,FT} = + [_project_key(k, MT) => v.parent[k] for k in _view_matching_keys(v)] + +function Base.firstindex(v::SparseArraySlice, d) + return minimum(k[d] for k in _view_matching_keys(v)) +end +function Base.lastindex(v::SparseArraySlice, d) + return maximum(k[d] for k in _view_matching_keys(v)) +end + +function Base.sum(v::SparseArraySlice{P,V}) where {P,V} + ks = _view_matching_keys(v) + isempty(ks) && return zero(V) + return sum(v.parent[k] for k in ks) +end + +Base.show(io::IO, ::MIME"text/plain", v::SparseArraySlice) = summary(io, v) +Base.show(io::IO, v::SparseArraySlice) = summary(io, v) diff --git a/src/sparsearray.jl b/src/sparsearray.jl index 5f6d9a7..b0c7f88 100644 --- a/src/sparsearray.jl +++ b/src/sparsearray.jl @@ -115,3 +115,21 @@ function SparseArray{T,N,K}() where {T,N,K<:NTuple{N,Any}} end _data(sa::SparseArray) = sa.data + +# ------------------------------------------------------------------------------ +# _keytype interface +# Each concrete AbstractSparseArray subtype must implement _keytype(::Type{<:SA}). +# ------------------------------------------------------------------------------ + +""" + _keytype(sa) / _keytype(::Type{<:AbstractSparseArray}) + +Return the key tuple type used by the sparse array. Required by `slice` and +broadcasting. Implement `_keytype(::Type{MySA})` for every concrete subtype. +""" +_keytype(::Type{<:AbstractSparseArray}) = + error("_keytype not implemented for this AbstractSparseArray subtype") +_keytype(sa::AbstractSparseArray) = _keytype(typeof(sa)) +_keytype(::Type{<:SparseArray{T,N,K}}) where {T,N,K} = K + +Base.haskey(sa::AbstractSparseArray, k) = haskey(_data(sa), k) From 00c4c62a0a3657e1e37a98333f26f89745a472ad Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Fri, 22 May 2026 14:06:49 +0200 Subject: [PATCH 05/22] Add tests for slicing and broadcast --- test/runtests.jl | 211 ++++++++++++++++++++++++++++++++++++++++++++++- test/testdata.jl | 12 +++ 2 files changed, 222 insertions(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index 742a538..5e7bcfb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,4 @@ -using Base: product +using Base: product using Dictionaries using HiGHS using JuMP @@ -329,3 +329,212 @@ end @test sum(x) == sum(x[:, :]) @test typeof(sum(x)) <: GenericAffExpr{Float64,MockVariableRef} end + +const _test_sa = testdata_sa() + +@testset "SparseArraySlice on SparseArray" begin + sa = _test_sa + + # _keytype + @test SV._keytype(sa) == Tuple{String,Int} + + # Types + v = slice(sa, "ford", :) + @test v isa SparseArraySlice + @test v isa SV.AbstractSparseArray + @test SV._keytype(v) == Tuple{Int} # projected key type + @test eltype(v) == Int + + # length + @test length(slice(sa, "ford", :)) == 2 + @test length(slice(sa, :, 2001)) == 2 + @test length(slice(sa, :, :)) == 5 + @test length(slice(sa, "xxx", :)) == 0 + + # keys / values / eachindex / pairs + ks = sort(keys(v)) + @test ks == [(2000,), (2001,)] + @test eachindex(v) == keys(v) + @test sort(values(v)) == [100, 150] + ps = Dict(pairs(v)) + @test ps[(2000,)] == 100 + @test ps[(2001,)] == 150 + + # getindex + @test v[(2000,)] == 100 # FT-tuple + @test v[NTuple{1,Any}((2001,))] == 150 # NTuple{NF,Any} + @test v[2000] == 100 # splatted (NF==1) + @test v[2001] == 150 + + v2 = slice(sa, :, :) # NF==2 + @test v2["ford", 2000] == 100 + @test v2["bmw", 2002] == 300 + + # haskey + @test haskey(v, (2000,)) + @test !haskey(v, (1999,)) + + # sum + @test sum(v) == 250 + @test sum(slice(sa, :, 2001)) == 350 + @test sum(slice(sa, "xxx", :)) == 0 + + # firstindex / lastindex (d = parent-dimension index) + @test SV.firstindex(v, 2) == 2000 + @test SV.lastindex(v, 2) == 2001 + + # iteration (values only) + @test sum(val for val in v) == 250 + @test Base.IteratorSize(typeof(v)) == Base.HasLength() + @test Base.IteratorEltype(typeof(v)) == Base.HasEltype() + + # show / summary + @test occursin("SparseArraySlice", sprint(summary, v)) + + # read-only + @test_throws MethodError (v[(2000,)] = 999) + @test_throws ErrorException size(v) + + # wrong mask length + @test_throws BoundsError slice(sa, "ford", :, :) + + # empty slice + ve = slice(sa, "xxx", :) + @test length(ve) == 0 + @test isempty(keys(ve)) + @test isempty(values(ve)) + @test sum(ve) == 0 +end + +@testset "SparseArraySlice on IndexedVarArray" begin + (; cars, year, car_cost) = testdata1(false) + m = Model() + @variable(m, x[c = cars, y = year]; container = IndexedVarArray) + for k in keys(car_cost) + insertvar!(x, k...) + end + + v = slice(x, "ford", :) + + # type and _keytype + @test v isa SparseArraySlice + @test SV._keytype(x) == Tuple{String,Int} + @test SV._keytype(v) == Tuple{Int} + + # length / keys + @test length(v) == 2 + @test sort(keys(v)) == [(2000,), (2001,)] + + # JuMP sum returns AffExpr + @test sum(v) isa AffExpr + @test length(sum(v).terms) == 2 + + @test sum(slice(x, :, 2001)) isa AffExpr + @test length(sum(slice(x, :, 2001)).terms) == 2 + + # empty JuMP slice sum returns zero(AffExpr) + @test sum(slice(x, "xxx", :)) == zero(AffExpr) +end + +@testset "Broadcasting SparseArray" begin + sa = _test_sa + + # scalar broadcast + r = sa .* 2 + @test r isa SparseArray + @test length(r) == 5 + @test r["ford", 2000] == 200 + @test r["lotus", 1957] == 1000 + + r2 = 2 .* sa + @test r2["bmw", 2001] == 400 + + r3 = sa .+ 10 + @test r3["ford", 2001] == 160 + + # element-wise binary + r4 = sa .+ sa + @test r4["ford", 2000] == 200 + @test r4["bmw", 2002] == 600 + + # function broadcast + r5 = sqrt.(sa .* 1.0) + @test r5 isa SparseArray + @test r5["ford", 2000] ≈ sqrt(100.0) + + # result type + @test Base.BroadcastStyle(typeof(sa)) isa SV.SparseBroadcastStyle + + # key mismatch error + sa2 = SparseArray(Dict(("a", 1) => 1)) + @test_throws ArgumentError sa .+ sa2 + + # empty array broadcast + empty_sa = SparseArray(Dictionary{Tuple{String,Int},Int}()) + r_empty = empty_sa .* 2 + @test r_empty isa SparseArray + @test length(r_empty) == 0 +end + +@testset "Broadcasting SparseArraySlice" begin + sa = _test_sa + + v = slice(sa, "ford", :) + r = v .* 2 + @test r isa SparseArray + @test length(r) == 2 + @test r[(2000,)] == 200 + @test r[(2001,)] == 300 + + # slice .+ slice (same keys) + r2 = v .+ v + @test r2[(2000,)] == 200 + + # NF=2 slice broadcast + v2 = slice(sa, :, :) + r3 = v2 .* 3 + @test r3["ford", 2000] == 300 + @test r3["lotus", 1957] == 1500 + + # key mismatch between two slices + vbmw = slice(sa, "bmw", :) + @test_throws ArgumentError v .+ vbmw +end + +@testset "Broadcasting IndexedVarArray" begin + (; cars, year, car_cost) = testdata1(false) + m = Model() + @variable(m, x[c = cars, y = year] >= 0; container = IndexedVarArray) + for k in keys(car_cost) + insertvar!(x, k...) + end + @objective(m, Min, sum(x[c, y] for (c, y) in keys(car_cost))) + @constraint(m, sum(x[:, :]) == 1) + set_optimizer(m, HiGHS.Optimizer) + set_optimizer_attribute(m, MOI.Silent(), true) + optimize!(m) + + # value.(iva) → SparseArray + vals = value.(x) + @test vals isa SparseArray + @test length(vals) == length(x) + @test isapprox(sum(values(vals)), 1.0; atol = 1e-6) + + # value.(slice) → SparseArray with projected keys + vslice = value.(slice(x, "ford", :)) + @test vslice isa SparseArray + @test length(vslice) == 2 + @test eltype(vslice) == Float64 + + # iva .+ iva → SparseArray{AffExpr} + aff = x .+ x + @test aff isa SparseArray + @test length(aff) == length(x) + @test first(values(aff)) isa AffExpr + + # key mismatch error + m2 = Model() + @variable(m2, y[c = ["lotus"], yr = [1957]]; container = IndexedVarArray) + insertvar!(y, "lotus", 1957) + @test_throws ArgumentError x .+ y +end diff --git a/test/testdata.jl b/test/testdata.jl index 5da7620..bb59522 100644 --- a/test/testdata.jl +++ b/test/testdata.jl @@ -26,6 +26,18 @@ function testdata(N = 998) ) end +function testdata_sa() + return SparseVariables.SparseArray( + Dict( + ("ford", 2000) => 100, + ("ford", 2001) => 150, + ("bmw", 2001) => 200, + ("bmw", 2002) => 300, + ("lotus", 1957) => 500, + ), + ) +end + function testdata1(addlotus = true) cars = ["ford", "bmw", "opel"] year = [2000, 2001, 2002, 2003] From c41d29b4e6dc22566c9b815febd4f14ad62084e4 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Fri, 22 May 2026 19:55:33 +0200 Subject: [PATCH 06/22] Simplify duplicated code --- src/indexedarray.jl | 54 +++++---------------------------------------- src/slice.jl | 32 ++++++++++----------------- 2 files changed, 18 insertions(+), 68 deletions(-) diff --git a/src/indexedarray.jl b/src/indexedarray.jl index 20772c3..32fb1f3 100644 --- a/src/indexedarray.jl +++ b/src/indexedarray.jl @@ -73,30 +73,10 @@ function unsafe_insertvar!(var::IndexedVarArray{V,N,T}, index...) where {V,N,T} return insertvar!(var, UnsafeInsert(), index...) end -joinex(ex1, ex2) = :($ex1..., $ex2...) -@generated function _active(idx::I, pat::P) where {I,P} - ids = fieldtypes(I) - ps = fieldtypes(P) - exs = [] - for i in 1:length(ids) - if ps[i] != Colon - if i > 2 - push!(exs, :(a1 = idx[$i],)) - else - push!(exs, :(idx[$i],)) - end - end - end - for i in 1:length(exs)-1 - exs[i+1] = joinex(exs[i], exs[i+1]) - end - return :(tuple($(exs[end])...)) -end - function build_cache!(cache, pat, sa::IndexedVarArray{V,N,T}) where {V,N,T} if isempty(cache) for v in keys(sa) - vred = _active(v, pat) + vred = _project_fixed(v, typeof(pat)) nv = get!(cache, vred, T[]) push!(nv, v) end @@ -105,9 +85,9 @@ function build_cache!(cache, pat, sa::IndexedVarArray{V,N,T}) where {V,N,T} end # Minimum number of entries before the index cache is used; below this a -# linear scan is cheaper. Tune with set_cache_cutoff! or calibrate with +# linear scan is assumed cheaper. Tune with set_cache_cutoff! or calibrate with # benchmark/cutoff_benchmark.jl. -const _CACHE_CUTOFF = Ref{Int}(100) +_CACHE_CUTOFF::Int = 100 """ set_cache_cutoff!(n::Int) @@ -117,13 +97,13 @@ selection switches from a linear scan to the pre-built index cache. Smaller values favour caching; larger values favour the linear scan for small arrays. Default: `100`. """ -set_cache_cutoff!(n::Int) = (_CACHE_CUTOFF[] = n; nothing) +set_cache_cutoff!(n::Int) = (global _CACHE_CUTOFF = n; nothing) function _select_cached(sa::IndexedVarArray{V,N,T}, pat)::Vector{T} where {V,N,T} - length(_data(sa)) < _CACHE_CUTOFF[] && return collect(T, _select_gen(keys(_data(sa)), pat)) + length(_data(sa)) < _CACHE_CUTOFF && return collect(T, _select_gen(keys(_data(sa)), pat)) cache = _getcache(sa, pat)::Dictionary{_decode_nonslices(sa, pat),Vector{T}} build_cache!(cache, pat, sa) - vals = _dropslices_gen(pat) + vals = _project_fixed(pat, typeof(pat)) return get!(cache, vals, T[]) end @@ -133,28 +113,6 @@ bin2int(v) = bin2int(v, Dim{length(v)}()) w = reverse([2^(i - 1) for i in 1:N]) return :(dot($w, v)) end - -function _dropslices(t::P) where {P} - return Tuple(ti for ti in t if ti != Colon()) -end - -@generated function _dropslices_gen(pat::P) where {P} - ps = fieldtypes(P) - exs = [] - for i in 1:length(ps) - if ps[i] != Colon - if i > 2 # Workaround for slurping of iterables (like strings) when passing to joinex - push!(exs, :(a2 = pat[$i],)) - else - push!(exs, :(pat[$i],)) - end - end - end - for i in 1:length(exs)-1 - exs[i+1] = joinex(exs[i], exs[i+1]) - end - return exs[end] -end """ _get_cache_index(::P) diff --git a/src/slice.jl b/src/slice.jl index ae6d542..4c8a472 100644 --- a/src/slice.jl +++ b/src/slice.jl @@ -1,22 +1,14 @@ -# Match key against mask at compile time; Colon positions are free (skipped). -@generated function _matches_mask(key::T, mask::MT) where {T,MT} - checks = Expr[] - for i in 1:fieldcount(T) - if fieldtypes(MT)[i] !== Colon - push!(checks, :(key[$i] != mask[$i] && return false)) - end - end - return quote - $(checks...) - return true - end +# Project tpl to the non-Colon (fixed) positions of MT. +@generated function _project_fixed(tpl::T, ::Type{MT}) where {T,MT} + inds = [i for i in 1:fieldcount(T) if fieldtypes(MT)[i] !== Colon] + return :($(Expr(:tuple, [:(tpl[$i]) for i in inds]...))) end -# Project a full key T down to the free dimensions FT. -@generated function _project_key(key::T, ::Type{MT}) where {T,MT} - free_idx = [i for i in 1:fieldcount(T) if fieldtypes(MT)[i] === Colon] - FT = Tuple{[fieldtypes(T)[i] for i in free_idx]...} - return :($(Expr(:tuple, [:(key[$i]) for i in free_idx]...))::$FT) +# Project tpl to the Colon (free) positions of MT. +@generated function _project_free(tpl::T, ::Type{MT}) where {T,MT} + inds = [i for i in 1:fieldcount(T) if fieldtypes(MT)[i] === Colon] + FT = Tuple{[fieldtypes(T)[i] for i in inds]...} + return :($(Expr(:tuple, [:(tpl[$i]) for i in inds]...))::$FT) end # Reconstruct a full key T from the fixed values in mask and the free key FT. @@ -92,7 +84,7 @@ end function _view_matching_keys( v::SparseArraySlice{P,V,N,T,NF,MT,FT}, )::Vector{T} where {P,V,N,T,NF,MT,FT} - return [k for k in keys(_data(v.parent)) if _matches_mask(k, v.mask)] + return collect(T, _select_gen(keys(_data(v.parent)), v.mask)) end # Iterator traits @@ -172,14 +164,14 @@ end Base.length(v::SparseArraySlice) = length(_view_matching_keys(v)) Base.keys(v::SparseArraySlice{P,V,N,T,NF,MT,FT}) where {P,V,N,T,NF,MT,FT} = - [_project_key(k, MT) for k in _view_matching_keys(v)] + [_project_free(k, MT) for k in _view_matching_keys(v)] Base.values(v::SparseArraySlice) = [v.parent[k] for k in _view_matching_keys(v)] Base.eachindex(v::SparseArraySlice) = keys(v) Base.pairs(v::SparseArraySlice{P,V,N,T,NF,MT,FT}) where {P,V,N,T,NF,MT,FT} = - [_project_key(k, MT) => v.parent[k] for k in _view_matching_keys(v)] + [_project_free(k, MT) => v.parent[k] for k in _view_matching_keys(v)] function Base.firstindex(v::SparseArraySlice, d) return minimum(k[d] for k in _view_matching_keys(v)) From 1f90bf38e2ad003fe37059d7f12743b4ea3363db Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Sat, 23 May 2026 12:32:01 +0200 Subject: [PATCH 07/22] Slicing of an AbstractSparseArray returns a SparseArraysSlice also for indexing notation --- src/dictionaries.jl | 26 ++++--------- src/indexedarray.jl | 11 +++++- src/slice.jl | 92 ++++++++++++++++++++++++++++++++------------- test/runtests.jl | 11 ++++++ 4 files changed, 94 insertions(+), 46 deletions(-) diff --git a/src/dictionaries.jl b/src/dictionaries.jl index d065710..8c3d2db 100644 --- a/src/dictionaries.jl +++ b/src/dictionaries.jl @@ -1,7 +1,7 @@ """ make_filter_fun(c, pos) -Return function to use for filtering depending on the type and value of `c` +Return function to use for filtering depending on the type and value of `c` to apply at position `pos` """ make_filter_fun(c, pos) = x -> x[pos] == c @@ -38,7 +38,7 @@ function indices_fun(some_tuple) end """ - _select_rowwise(a, pattern) + _select_rowwise(a, pattern) Filter iterable data a by tuple `pattern` by row (slow) """ @@ -47,7 +47,7 @@ function _select_rowwise(a, pattern) end """ - _select_colwise(a, pattern) + _select_colwise(a, pattern) Filter iterable data a by tuple `pattern` by column (recursively) """ @@ -59,7 +59,7 @@ end _select_gen(a, pattern) Filter iterable data `a` by tuple `pattern` by row, using generated function for speed. -See more straight-forward implementations `_select_rowwise` and `_select_colwise` for reference. +See more straight-forward implementations `_select_rowwise` and `_select_colwise` for reference. """ function _select_gen(a, pattern) return filter(x -> _select_generated(pattern, x), a) @@ -68,7 +68,7 @@ end """ _select_gen_perm(a, pattern, perm) Filter iterable `data` byt tuple `pattern` by row using generated function that permutes the sequence of -evaluation by the permutation tuple `perm` for improved control as this can give performance advantages, +evaluation by the permutation tuple `perm` for improved control as this can give performance advantages, depending on the uniqueness of the search pattern and cost of function evaluation. ## Example @@ -226,33 +226,23 @@ Works on types because it is used in generated function """ isfixed(t) = true isfixed(::Type{T} where {T<:Function}) = false -isfixed(::Type{T} where {T<:UnitRange}) = false -iscolon(t) = false -iscolon(::Type{T} where {T<:Colon}) = true +isfixed(::Type{T} where {T<:AbstractRange}) = false @generated function _getindex( sa::AbstractSparseArray{T,N}, tpl::Tuple, ) where {T,N} lookup = true - slice = true for t in fieldtypes(tpl) if !isfixed(t) lookup = false - if !iscolon(t) - slice = false - end end end if lookup return :(get(_data(sa), tpl, zero(T))) - elseif !slice - return :(retval = select(_data(sa), tpl); - length(retval) > 0 ? retval : zero(T)) - else # Return selection or zero if empty to avoid reduction of empty iterate - return :(retval = _select_var(sa, tpl); - length(retval) > 0 ? retval : zero(T)) + else + return :(_make_slice(sa, tpl)) end end diff --git a/src/indexedarray.jl b/src/indexedarray.jl index 32fb1f3..c20e656 100644 --- a/src/indexedarray.jl +++ b/src/indexedarray.jl @@ -99,8 +99,13 @@ scan for small arrays. Default: `100`. """ set_cache_cutoff!(n::Int) = (global _CACHE_CUTOFF = n; nothing) +@generated function _is_cacheable_pattern(::Type{P}) where {P<:Tuple} + return :($(all(t == Colon || isfixed(t) for t in fieldtypes(P)))) +end + function _select_cached(sa::IndexedVarArray{V,N,T}, pat)::Vector{T} where {V,N,T} length(_data(sa)) < _CACHE_CUTOFF && return collect(T, _select_gen(keys(_data(sa)), pat)) + _is_cacheable_pattern(typeof(pat)) || return collect(T, _select_gen(keys(_data(sa)), pat)) cache = _getcache(sa, pat)::Dictionary{_decode_nonslices(sa, pat),Vector{T}} build_cache!(cache, pat, sa) vals = _project_fixed(pat, typeof(pat)) @@ -190,8 +195,10 @@ function _view_matching_keys( return _select_cached(v.parent, v.mask) end -# JuMP-efficient sum: build AffExpr directly via add_to_expression!. -function Base.sum(v::SparseArraySlice{<:IndexedVarArray,V}) where {V<:AbstractVariableRef} +# JuMP-efficient sum: build AffExpr directly via add_to_expression! for the +# standard VariableRef type. Custom AbstractVariableRef subtypes fall back to +# the generic slice sum implementation. +function Base.sum(v::SparseArraySlice{<:IndexedVarArray,VariableRef}) result = zero(AffExpr) for k in _view_matching_keys(v) JuMP.add_to_expression!(result, v.parent[k]) diff --git a/src/slice.jl b/src/slice.jl index 4c8a472..5300b31 100644 --- a/src/slice.jl +++ b/src/slice.jl @@ -1,22 +1,29 @@ -# Project tpl to the non-Colon (fixed) positions of MT. +# Generated functions in this file cannot rely on helpers defined later in the +# include order, so selector classification stays local. +_is_exact_selector_type(::Type) = true +_is_exact_selector_type(::Type{<:Function}) = false +_is_exact_selector_type(::Type{<:AbstractRange}) = false +_is_exact_selector_type(::Type{<:Colon}) = false + +# Project tpl to the exact-match positions of MT. @generated function _project_fixed(tpl::T, ::Type{MT}) where {T,MT} - inds = [i for i in 1:fieldcount(T) if fieldtypes(MT)[i] !== Colon] + inds = [i for i in 1:fieldcount(T) if _is_exact_selector_type(fieldtypes(MT)[i])] return :($(Expr(:tuple, [:(tpl[$i]) for i in inds]...))) end -# Project tpl to the Colon (free) positions of MT. +# Project tpl to the non-exact (free) positions of MT. @generated function _project_free(tpl::T, ::Type{MT}) where {T,MT} - inds = [i for i in 1:fieldcount(T) if fieldtypes(MT)[i] === Colon] + inds = [i for i in 1:fieldcount(T) if !_is_exact_selector_type(fieldtypes(MT)[i])] FT = Tuple{[fieldtypes(T)[i] for i in inds]...} return :($(Expr(:tuple, [:(tpl[$i]) for i in inds]...))::$FT) end -# Reconstruct a full key T from the fixed values in mask and the free key FT. +# Reconstruct a full key T from the exact values in mask and the free key FT. @generated function _reconstruct_key(mask::MT, free_key::FT, ::Type{T}) where {MT,FT,T} parts = Vector{Expr}(undef, fieldcount(T)) fi = 1 for i in 1:fieldcount(T) - if fieldtypes(MT)[i] === Colon + if !_is_exact_selector_type(fieldtypes(MT)[i]) parts[i] = :(free_key[$fi]) fi += 1 else @@ -50,13 +57,29 @@ struct SparseArraySlice{ end _keytype(::Type{<:SparseArraySlice{P,V,N,T,NF,MT,FT}}) where {P,V,N,T,NF,MT,FT} = FT +_parent_keytype(::Type{<:SparseArraySlice{P,V,N,T,NF,MT,FT}}) where {P,V,N,T,NF,MT,FT} = FT + +@generated function _parent_keytype(::Type{P}) where {P<:AbstractSparseArray} + :data in fieldnames(P) || + error("cannot infer key type for this AbstractSparseArray subtype") + D = fieldtype(P, :data) + return :($(D.parameters[1])) +end + +function _matches_free_key( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, + free_key, +) where {P,V,N,T,NF,MT,FT} + return _select_generated(_project_free(v.mask, MT), free_key) +end """ slice(sa::AbstractSparseArray, mask...) Return a lazy `SparseArraySlice` over entries of `sa` matching `mask`. Use `:` -for free (wildcard) dimensions and exact values for fixed dimensions. The result -is an `AbstractSparseArray{V,NF}` where `NF` is the number of free dimensions. +for wildcard dimensions, exact values for fixed dimensions, and predicates or +ranges for filtered dimensions. The result is an `AbstractSparseArray{V,NF}` +where `NF` is the number of non-exact dimensions. # Example ```julia @@ -70,11 +93,11 @@ function slice(sa::AbstractSparseArray, mask...) end @generated function _make_slice(sa::P, mask::MT) where {P<:AbstractSparseArray,MT<:Tuple} - K = _keytype(P) + K = _parent_keytype(P) N = ndims(P) V = eltype(P) fieldcount(MT) != N && return :(throw(BoundsError(sa, mask))) - free = [fieldtypes(K)[i] for i in 1:N if fieldtypes(MT)[i] === Colon] + free = [fieldtypes(K)[i] for i in 1:N if !_is_exact_selector_type(fieldtypes(MT)[i])] NF = length(free) FT = Tuple{free...} return :(SparseArraySlice{$P,$V,$N,$K,$NF,$MT,$FT}(sa, mask)) @@ -113,6 +136,7 @@ function Base.getindex( v::SparseArraySlice{P,V,N,T,NF,MT,FT}, free_key::FT, ) where {P,V,N,T,NF,MT,FT} + _matches_free_key(v, free_key) || return zero(V) return v.parent[_reconstruct_key(v.mask, free_key, T)] end @@ -121,28 +145,17 @@ function Base.getindex( v::SparseArraySlice{P,V,N,T,NF,MT,FT}, idx::NTuple{NF,Any}, ) where {P,V,N,T,NF,MT,FT} + _matches_free_key(v, idx) || return zero(V) return v.parent[_reconstruct_key(v.mask, idx, T)] end # Splatted: v[f, c] or v[f] (NF==1) -@generated function Base.getindex( +function Base.getindex( v::SparseArraySlice{P,V,N,T,NF,MT,FT}, idx..., ) where {P,V,N,T,NF,MT,FT} - if length(idx) != NF - return :(throw(BoundsError(v, idx))) - end - parts = Expr[] - fi = 1 - for i in 1:fieldcount(T) - if fieldtypes(MT)[i] === Colon - push!(parts, :(idx[$fi])) - fi += 1 - else - push!(parts, :(v.mask[$i])) - end - end - return :(v.parent[$(Expr(:tuple, parts...))]) + length(idx) == NF || throw(BoundsError(v, idx)) + return v[idx] end Base.setindex!(::SparseArraySlice, _, _...) = error("SparseArraySlice is read-only") @@ -158,6 +171,7 @@ function Base.haskey( v::SparseArraySlice{P,V,N,T,NF,MT,FT}, free_key::FT, ) where {P,V,N,T,NF,MT,FT} + _matches_free_key(v, free_key) || return false return haskey(_data(v.parent), _reconstruct_key(v.mask, free_key, T)) end @@ -186,5 +200,31 @@ function Base.sum(v::SparseArraySlice{P,V}) where {P,V} return sum(v.parent[k] for k in ks) end -Base.show(io::IO, ::MIME"text/plain", v::SparseArraySlice) = summary(io, v) +function Base.summary(io::IO, v::SparseArraySlice) + num_entries = length(v) + return print( + io, + "SparseArraySlice with ", + num_entries, + isone(num_entries) ? " entry" : " entries", + " matching ", + v.mask, + ) +end + +function Base.show(io::IO, ::MIME"text/plain", v::SparseArraySlice) + summary(io, v) + if !iszero(length(v)) + println(io, ":") + entries = pairs(v) + if length(entries) > 20 + show(io, first(entries, 10)) + print(io, "\n ⋮\n") + show(io, last(entries, 10)) + else + show(io, entries) + end + end +end + Base.show(io::IO, v::SparseArraySlice) = summary(io, v) diff --git a/test/runtests.jl b/test/runtests.jl index 5e7bcfb..0955672 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -77,6 +77,10 @@ end @test car_cost["bmw", 2001] == 200 @test car_cost["bmw", 2003] == 0 + @test car_cost[endswith("s"), <(2000)] isa SparseArraySlice + @test length(car_cost[endswith("s"), <(2000)]) == 1 + @test car_cost[endswith("s"), <(2000)]["lotus", 1957] == 500 + @test car_cost[endswith("s"), <(2000)]["bmw", 2001] == 0 @test length(car_cost) == 5 @test car_cost["lotus", 1957] == 500 @@ -181,6 +185,10 @@ end # Slicing and lookup @test length(z["bmw", :]) == 2 @test length(z[:, 2001]) == 2 + @test z[endswith("w"), isodd] isa SparseArraySlice + @test length(z[endswith("w"), isodd]) == 1 + @test haskey(z[endswith("w"), isodd], ("bmw", 2001)) + @test !haskey(z[endswith("w"), isodd], ("bmw", 2002)) @test typeof(z["bmw", 2001]) == VariableRef @test z["bmw", 20] == 0 @@ -390,6 +398,9 @@ const _test_sa = testdata_sa() # show / summary @test occursin("SparseArraySlice", sprint(summary, v)) + @test occursin("matching (\"ford\", Colon())", sprint(summary, v)) + @test occursin("(2000,) => 100", sprint(show, MIME("text/plain"), v)) + @test occursin("(2001,) => 150", sprint(show, MIME("text/plain"), v)) # read-only @test_throws MethodError (v[(2000,)] = 999) From 802db0ee7428493037df16f9d135800f095af7f4 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Sat, 23 May 2026 12:33:15 +0200 Subject: [PATCH 08/22] Invalidate cache also for unsafe insert --- src/indexedarray.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/indexedarray.jl b/src/indexedarray.jl index c20e656..1475287 100644 --- a/src/indexedarray.jl +++ b/src/indexedarray.jl @@ -60,6 +60,7 @@ function insertvar!( ::UnsafeInsert, index..., ) where {V,N,T} + clear_cache!(var) return var[index] = var.f(index...) end From da0895a1c58e7946aaac0c1c418f0c3c16068b7c Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Sat, 23 May 2026 12:36:58 +0200 Subject: [PATCH 09/22] Format fixes --- src/SparseVariables.jl | 8 ++--- src/broadcast.jl | 34 +++++++++++++++----- src/indexedarray.jl | 11 +++++-- src/slice.jl | 70 +++++++++++++++++++++++++++++------------- src/sparsearray.jl | 7 +++-- test/runtests.jl | 50 +++++++++++++++--------------- 6 files changed, 114 insertions(+), 66 deletions(-) diff --git a/src/SparseVariables.jl b/src/SparseVariables.jl index c139bf7..5414888 100644 --- a/src/SparseVariables.jl +++ b/src/SparseVariables.jl @@ -34,11 +34,7 @@ export set_cache_cutoff! # all calls in this block will be precompiled, regardless of whether # they belong to your package or not (on Julia 1.8 and higher) - @variable( - m, - x[r = rs, i = is, st = sts, sy = sys]; - container = IndexedVarArray - ) + @variable(m, x[r=rs, i=is, st=sts, sy=sys]; container = IndexedVarArray) for r in rs, i in is, st in sts, sy in sys insertvar!(x, r, i, st, sy) unsafe_insertvar!(x, r, i, st, sy) @@ -46,7 +42,7 @@ export set_cache_cutoff! x[:, 1, :, :] x[10, :, :, :] x[1, :, :, :a] - @variable(m, y[i = rs, j = rs, k = rs]; container = IndexedVarArray) + @variable(m, y[i=rs, j=rs, k=rs]; container = IndexedVarArray) for i in rs, j in rs, k in rs insertvar!(y, i, j, k) end diff --git a/src/broadcast.jl b/src/broadcast.jl index 5bc1a9b..a7e4dbe 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -13,8 +13,9 @@ materialised as `SparseArray`. """ struct SparseBroadcastStyle{N,K} <: Broadcast.BroadcastStyle end -Base.BroadcastStyle(::Type{SA}) where {SA<:AbstractSparseArray} = - SparseBroadcastStyle{ndims(SA),_keytype(SA)}() +function Base.BroadcastStyle(::Type{SA}) where {SA<:AbstractSparseArray} + return SparseBroadcastStyle{ndims(SA),_keytype(SA)}() +end # Disallow mixing with other array types. function Base.BroadcastStyle(::SparseBroadcastStyle, ::Base.BroadcastStyle) @@ -63,19 +64,27 @@ function _sparse_getindex( end function _sparse_get_args(args::Tuple, key) - return (_sparse_getindex(first(args), key), _sparse_get_args(Base.tail(args), key)...) + return ( + _sparse_getindex(first(args), key), + _sparse_get_args(Base.tail(args), key)..., + ) end _sparse_get_args(::Tuple{}, ::Any) = () function _sparse_check_same_keys(ref_keys, x::AbstractSparseArray, args...) if length(ref_keys) != length(x) || any(k -> !haskey(x, k), ref_keys) - throw(ArgumentError("Cannot broadcast SparseArrays with different indices")) + throw( + ArgumentError( + "Cannot broadcast SparseArrays with different indices", + ), + ) end return _sparse_check_same_keys(ref_keys, args...) end -_sparse_check_same_keys(ref_keys, ::Any, args...) = - _sparse_check_same_keys(ref_keys, args...) +function _sparse_check_same_keys(ref_keys, ::Any, args...) + return _sparse_check_same_keys(ref_keys, args...) +end _sparse_check_same_keys(::Any) = nothing function _sparse_indices( @@ -104,9 +113,18 @@ function Base.copy( return SparseArray(Dictionary(indices, vals)) end -function Base.Broadcast.broadcast_preserving_zero_d(f, A::AbstractSparseArray, As...) +function Base.Broadcast.broadcast_preserving_zero_d( + f, + A::AbstractSparseArray, + As..., +) return broadcast(f, A, As...) end -function Base.Broadcast.broadcast_preserving_zero_d(f, x, A::AbstractSparseArray, As...) +function Base.Broadcast.broadcast_preserving_zero_d( + f, + x, + A::AbstractSparseArray, + As..., +) return broadcast(f, x, A, As...) end diff --git a/src/indexedarray.jl b/src/indexedarray.jl index 1475287..164561b 100644 --- a/src/indexedarray.jl +++ b/src/indexedarray.jl @@ -104,9 +104,14 @@ set_cache_cutoff!(n::Int) = (global _CACHE_CUTOFF = n; nothing) return :($(all(t == Colon || isfixed(t) for t in fieldtypes(P)))) end -function _select_cached(sa::IndexedVarArray{V,N,T}, pat)::Vector{T} where {V,N,T} - length(_data(sa)) < _CACHE_CUTOFF && return collect(T, _select_gen(keys(_data(sa)), pat)) - _is_cacheable_pattern(typeof(pat)) || return collect(T, _select_gen(keys(_data(sa)), pat)) +function _select_cached( + sa::IndexedVarArray{V,N,T}, + pat, +)::Vector{T} where {V,N,T} + length(_data(sa)) < _CACHE_CUTOFF && + return collect(T, _select_gen(keys(_data(sa)), pat)) + _is_cacheable_pattern(typeof(pat)) || + return collect(T, _select_gen(keys(_data(sa)), pat)) cache = _getcache(sa, pat)::Dictionary{_decode_nonslices(sa, pat),Vector{T}} build_cache!(cache, pat, sa) vals = _project_fixed(pat, typeof(pat)) diff --git a/src/slice.jl b/src/slice.jl index 5300b31..15e7d1f 100644 --- a/src/slice.jl +++ b/src/slice.jl @@ -7,19 +7,28 @@ _is_exact_selector_type(::Type{<:Colon}) = false # Project tpl to the exact-match positions of MT. @generated function _project_fixed(tpl::T, ::Type{MT}) where {T,MT} - inds = [i for i in 1:fieldcount(T) if _is_exact_selector_type(fieldtypes(MT)[i])] + inds = [ + i for i in 1:fieldcount(T) if _is_exact_selector_type(fieldtypes(MT)[i]) + ] return :($(Expr(:tuple, [:(tpl[$i]) for i in inds]...))) end # Project tpl to the non-exact (free) positions of MT. @generated function _project_free(tpl::T, ::Type{MT}) where {T,MT} - inds = [i for i in 1:fieldcount(T) if !_is_exact_selector_type(fieldtypes(MT)[i])] + inds = [ + i for + i in 1:fieldcount(T) if !_is_exact_selector_type(fieldtypes(MT)[i]) + ] FT = Tuple{[fieldtypes(T)[i] for i in inds]...} return :($(Expr(:tuple, [:(tpl[$i]) for i in inds]...))::$FT) end # Reconstruct a full key T from the exact values in mask and the free key FT. -@generated function _reconstruct_key(mask::MT, free_key::FT, ::Type{T}) where {MT,FT,T} +@generated function _reconstruct_key( + mask::MT, + free_key::FT, + ::Type{T}, +) where {MT,FT,T} parts = Vector{Expr}(undef, fieldcount(T)) fi = 1 for i in 1:fieldcount(T) @@ -43,21 +52,22 @@ projected key tuple type covering only the free dimensions. Implements Create via `slice(sa, mask...)`. """ -struct SparseArraySlice{ - P<:AbstractSparseArray, - V, - N, - T, - NF, - MT<:Tuple, - FT<:Tuple, -} <: AbstractSparseArray{V,NF} +struct SparseArraySlice{P<:AbstractSparseArray,V,N,T,NF,MT<:Tuple,FT<:Tuple} <: + AbstractSparseArray{V,NF} parent::P mask::MT end -_keytype(::Type{<:SparseArraySlice{P,V,N,T,NF,MT,FT}}) where {P,V,N,T,NF,MT,FT} = FT -_parent_keytype(::Type{<:SparseArraySlice{P,V,N,T,NF,MT,FT}}) where {P,V,N,T,NF,MT,FT} = FT +function _keytype( + ::Type{<:SparseArraySlice{P,V,N,T,NF,MT,FT}}, +) where {P,V,N,T,NF,MT,FT} + return FT +end +function _parent_keytype( + ::Type{<:SparseArraySlice{P,V,N,T,NF,MT,FT}}, +) where {P,V,N,T,NF,MT,FT} + return FT +end @generated function _parent_keytype(::Type{P}) where {P<:AbstractSparseArray} :data in fieldnames(P) || @@ -92,12 +102,18 @@ function slice(sa::AbstractSparseArray, mask...) return _make_slice(sa, tuple(mask...)) end -@generated function _make_slice(sa::P, mask::MT) where {P<:AbstractSparseArray,MT<:Tuple} +@generated function _make_slice( + sa::P, + mask::MT, +) where {P<:AbstractSparseArray,MT<:Tuple} K = _parent_keytype(P) N = ndims(P) V = eltype(P) fieldcount(MT) != N && return :(throw(BoundsError(sa, mask))) - free = [fieldtypes(K)[i] for i in 1:N if !_is_exact_selector_type(fieldtypes(MT)[i])] + free = [ + fieldtypes(K)[i] for + i in 1:N if !_is_exact_selector_type(fieldtypes(MT)[i]) + ] NF = length(free) FT = Tuple{free...} return :(SparseArraySlice{$P,$V,$N,$K,$NF,$MT,$FT}(sa, mask)) @@ -116,7 +132,9 @@ Base.IteratorEltype(::Type{<:SparseArraySlice}) = Base.HasEltype() Base.eltype(::Type{<:SparseArraySlice{P,V}}) where {P,V} = V # Iteration: values only (AbstractArray semantics) -function Base.iterate(v::SparseArraySlice{P,V,N,T,NF,MT,FT}) where {P,V,N,T,NF,MT,FT} +function Base.iterate( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, +) where {P,V,N,T,NF,MT,FT} matching = _view_matching_keys(v) isempty(matching) && return nothing return (v.parent[matching[1]], (matching, 2)) @@ -158,7 +176,9 @@ function Base.getindex( return v[idx] end -Base.setindex!(::SparseArraySlice, _, _...) = error("SparseArraySlice is read-only") +function Base.setindex!(::SparseArraySlice, _, _...) + return error("SparseArraySlice is read-only") +end function Base.size(::SparseArraySlice) return error( @@ -177,15 +197,21 @@ end Base.length(v::SparseArraySlice) = length(_view_matching_keys(v)) -Base.keys(v::SparseArraySlice{P,V,N,T,NF,MT,FT}) where {P,V,N,T,NF,MT,FT} = - [_project_free(k, MT) for k in _view_matching_keys(v)] +function Base.keys( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, +) where {P,V,N,T,NF,MT,FT} + return [_project_free(k, MT) for k in _view_matching_keys(v)] +end Base.values(v::SparseArraySlice) = [v.parent[k] for k in _view_matching_keys(v)] Base.eachindex(v::SparseArraySlice) = keys(v) -Base.pairs(v::SparseArraySlice{P,V,N,T,NF,MT,FT}) where {P,V,N,T,NF,MT,FT} = - [_project_free(k, MT) => v.parent[k] for k in _view_matching_keys(v)] +function Base.pairs( + v::SparseArraySlice{P,V,N,T,NF,MT,FT}, +) where {P,V,N,T,NF,MT,FT} + return [_project_free(k, MT) => v.parent[k] for k in _view_matching_keys(v)] +end function Base.firstindex(v::SparseArraySlice, d) return minimum(k[d] for k in _view_matching_keys(v)) diff --git a/src/sparsearray.jl b/src/sparsearray.jl index b0c7f88..b9c4d0f 100644 --- a/src/sparsearray.jl +++ b/src/sparsearray.jl @@ -127,8 +127,11 @@ _data(sa::SparseArray) = sa.data Return the key tuple type used by the sparse array. Required by `slice` and broadcasting. Implement `_keytype(::Type{MySA})` for every concrete subtype. """ -_keytype(::Type{<:AbstractSparseArray}) = - error("_keytype not implemented for this AbstractSparseArray subtype") +function _keytype(::Type{<:AbstractSparseArray}) + return error( + "_keytype not implemented for this AbstractSparseArray subtype", + ) +end _keytype(sa::AbstractSparseArray) = _keytype(typeof(sa)) _keytype(::Type{<:SparseArray{T,N,K}}) where {T,N,K} = K diff --git a/test/runtests.jl b/test/runtests.jl index 0955672..961c2f9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,4 @@ -using Base: product +using Base: product using Dictionaries using HiGHS using JuMP @@ -49,7 +49,7 @@ end @variable( m, - car_vars[maker = cars, year = years, color = colors, kms = kms]; + car_vars[maker=cars, year=years, color=colors, kms=kms]; container = IndexedVarArray ) @test typeof(car_vars) == @@ -118,14 +118,14 @@ end (; car_cost) = testdata1() m = Model() - @variable(m, y[c = cars, i = years]; container = IndexedVarArray) + @variable(m, y[c=cars, i=years]; container = IndexedVarArray) for (c, i) in collect(keys(car_cost)) insertvar!(y, c, i) end @test typeof(y) == IndexedVarArray{VariableRef,2,Tuple{String,Int}} - @variable(m, w[c = cars, i = years], Bin; container = IndexedVarArray) + @variable(m, w[c=cars, i=years], Bin; container = IndexedVarArray) for (c, i) in collect(keys(car_cost)) insertvar!(w, c, i) end @@ -168,7 +168,7 @@ end m = Model() (; cars, year, car_cost) = testdata1(false) - @variable(m, z[cars = cars, year = year]; container = IndexedVarArray) + @variable(m, z[cars=cars, year=year]; container = IndexedVarArray) for (cr, yr) in keys(car_cost) insertvar!(z, cr, yr) @@ -198,7 +198,7 @@ end @test length(z) == 5 # Alternative constructor - @variable(m, z2[cars = cars, year = year], container = IndexedVarArray) + @variable(m, z2[cars=cars, year=year], container = IndexedVarArray) for k in keys(car_cost) insertvar!(z2, k...) end @@ -209,7 +209,7 @@ end @variable( m, - z3[cars = cars, year = years, color = colors, km = kms]; + z3[cars=cars, year=years, color=colors, km=kms]; container = IndexedVarArray ) for k in indices @@ -239,7 +239,7 @@ end (; cars, year, car_cost) = testdata1(false) m = Model() - @variable(m, y[car = cars, year = year] >= 0; container = IndexedVarArray) + @variable(m, y[car=cars, year=year] >= 0; container = IndexedVarArray) for c in cars insertvar!(y, c, 2002) end @@ -279,7 +279,7 @@ end # Test JuMP Extension m = Model() - @variable(m, x[i = 1:3, j = 100:102] >= 0, container = IndexedVarArray) + @variable(m, x[i=1:3, j=100:102] >= 0, container = IndexedVarArray) @test length(x) == 0 insertvar!(x, 1, 100) @test length(x) == 1 @@ -324,7 +324,7 @@ end m = Model() @variable( m, - x[i = 1:3, j = 100:102] >= 0, + x[i=1:3, j=100:102] >= 0, Mocking(), container = IndexedVarArray ) @@ -355,9 +355,9 @@ const _test_sa = testdata_sa() # length @test length(slice(sa, "ford", :)) == 2 - @test length(slice(sa, :, 2001)) == 2 - @test length(slice(sa, :, :)) == 5 - @test length(slice(sa, "xxx", :)) == 0 + @test length(slice(sa, :, 2001)) == 2 + @test length(slice(sa, :, :)) == 5 + @test length(slice(sa, "xxx", :)) == 0 # keys / values / eachindex / pairs ks = sort(keys(v)) @@ -369,17 +369,17 @@ const _test_sa = testdata_sa() @test ps[(2001,)] == 150 # getindex - @test v[(2000,)] == 100 # FT-tuple + @test v[(2000,)] == 100 # FT-tuple @test v[NTuple{1,Any}((2001,))] == 150 # NTuple{NF,Any} - @test v[2000] == 100 # splatted (NF==1) - @test v[2001] == 150 + @test v[2000] == 100 # splatted (NF==1) + @test v[2001] == 150 v2 = slice(sa, :, :) # NF==2 - @test v2["ford", 2000] == 100 - @test v2["bmw", 2002] == 300 + @test v2["ford", 2000] == 100 + @test v2["bmw", 2002] == 300 # haskey - @test haskey(v, (2000,)) + @test haskey(v, (2000,)) @test !haskey(v, (1999,)) # sum @@ -389,7 +389,7 @@ const _test_sa = testdata_sa() # firstindex / lastindex (d = parent-dimension index) @test SV.firstindex(v, 2) == 2000 - @test SV.lastindex(v, 2) == 2001 + @test SV.lastindex(v, 2) == 2001 # iteration (values only) @test sum(val for val in v) == 250 @@ -407,7 +407,7 @@ const _test_sa = testdata_sa() @test_throws ErrorException size(v) # wrong mask length - @test_throws BoundsError slice(sa, "ford", :, :) + @test_throws BoundsError slice(sa,"ford",:,:) # empty slice ve = slice(sa, "xxx", :) @@ -420,7 +420,7 @@ end @testset "SparseArraySlice on IndexedVarArray" begin (; cars, year, car_cost) = testdata1(false) m = Model() - @variable(m, x[c = cars, y = year]; container = IndexedVarArray) + @variable(m, x[c=cars, y=year]; container = IndexedVarArray) for k in keys(car_cost) insertvar!(x, k...) end @@ -466,7 +466,7 @@ end # element-wise binary r4 = sa .+ sa @test r4["ford", 2000] == 200 - @test r4["bmw", 2002] == 600 + @test r4["bmw", 2002] == 600 # function broadcast r5 = sqrt.(sa .* 1.0) @@ -515,7 +515,7 @@ end @testset "Broadcasting IndexedVarArray" begin (; cars, year, car_cost) = testdata1(false) m = Model() - @variable(m, x[c = cars, y = year] >= 0; container = IndexedVarArray) + @variable(m, x[c=cars, y=year] >= 0; container = IndexedVarArray) for k in keys(car_cost) insertvar!(x, k...) end @@ -545,7 +545,7 @@ end # key mismatch error m2 = Model() - @variable(m2, y[c = ["lotus"], yr = [1957]]; container = IndexedVarArray) + @variable(m2, y[c=["lotus"], yr=[1957]]; container = IndexedVarArray) insertvar!(y, "lotus", 1957) @test_throws ArgumentError x .+ y end From d3e5f36426ca41e71438f13ec6242d8db3065c36 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Sat, 23 May 2026 12:52:49 +0200 Subject: [PATCH 10/22] Formatting --- tutorial/sparse_tutorial.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tutorial/sparse_tutorial.jl b/tutorial/sparse_tutorial.jl index bc2c05e..48c444c 100644 --- a/tutorial/sparse_tutorial.jl +++ b/tutorial/sparse_tutorial.jl @@ -6,7 +6,7 @@ using InteractiveUtils # This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error). macro bind(def, element) - quote + return quote local iv = try Base.loaded_modules[Base.PkgId( Base.UUID("6e696c72-6542-2067-7265-42206c756150"), From f7baf89d6cbe6d3f86668095c5039ae344edd86d Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Wed, 24 Jun 2026 13:27:10 +0200 Subject: [PATCH 11/22] Simplify slicing structure --- src/indexedarray.jl | 4 +- src/slice.jl | 96 +++++++++++++++++++-------------------------- test/Project.toml | 1 + 3 files changed, 44 insertions(+), 57 deletions(-) diff --git a/src/indexedarray.jl b/src/indexedarray.jl index 164561b..aff3a55 100644 --- a/src/indexedarray.jl +++ b/src/indexedarray.jl @@ -196,8 +196,8 @@ end # Override _view_matching_keys for IndexedVarArray parent: use index cache. function _view_matching_keys( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, -)::Vector{T} where {P<:IndexedVarArray,V,N,T,NF,MT,FT} + v::SparseArraySlice{P,V,NF,MT}, +) where {P<:IndexedVarArray,V,NF,MT} return _select_cached(v.parent, v.mask) end diff --git a/src/slice.jl b/src/slice.jl index 15e7d1f..618342a 100644 --- a/src/slice.jl +++ b/src/slice.jl @@ -43,43 +43,40 @@ end end """ - SparseArraySlice{P,V,N,T,NF,MT,FT} + SparseArraySlice{P,V,NF,MT} A lazy, mask-filtered view of any `AbstractSparseArray`. `P` is the concrete -parent type, `NF` is the number of free (Colon) dimensions, and `FT` is the -projected key tuple type covering only the free dimensions. Implements +parent type, `NF` is the number of non-exact dimensions, and `MT` is the +mask type. Implements `AbstractSparseArray{V,NF}`. Create via `slice(sa, mask...)`. """ -struct SparseArraySlice{P<:AbstractSparseArray,V,N,T,NF,MT<:Tuple,FT<:Tuple} <: +struct SparseArraySlice{P<:AbstractSparseArray,V,NF,MT<:Tuple} <: AbstractSparseArray{V,NF} parent::P mask::MT end function _keytype( - ::Type{<:SparseArraySlice{P,V,N,T,NF,MT,FT}}, -) where {P,V,N,T,NF,MT,FT} - return FT -end -function _parent_keytype( - ::Type{<:SparseArraySlice{P,V,N,T,NF,MT,FT}}, -) where {P,V,N,T,NF,MT,FT} - return FT + ::Type{<:SparseArraySlice{P,V,NF,MT}}, +) where {P,V,NF,MT} + return _free_keytype(MT, _keytype(P)) end -@generated function _parent_keytype(::Type{P}) where {P<:AbstractSparseArray} - :data in fieldnames(P) || - error("cannot infer key type for this AbstractSparseArray subtype") - D = fieldtype(P, :data) - return :($(D.parameters[1])) +@generated function _free_keytype(::Type{MT}, ::Type{T}) where {MT,T} + inds = [ + i for + i in 1:fieldcount(T) if !_is_exact_selector_type(fieldtypes(MT)[i]) + ] + FT = Tuple{[fieldtypes(T)[i] for i in inds]...} + return :($FT) end function _matches_free_key( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, + v::SparseArraySlice{P,V,NF,MT}, free_key, -) where {P,V,N,T,NF,MT,FT} +) where {P,V,NF,MT} return _select_generated(_project_free(v.mask, MT), free_key) end @@ -106,23 +103,18 @@ end sa::P, mask::MT, ) where {P<:AbstractSparseArray,MT<:Tuple} - K = _parent_keytype(P) N = ndims(P) V = eltype(P) fieldcount(MT) != N && return :(throw(BoundsError(sa, mask))) - free = [ - fieldtypes(K)[i] for - i in 1:N if !_is_exact_selector_type(fieldtypes(MT)[i]) - ] - NF = length(free) - FT = Tuple{free...} - return :(SparseArraySlice{$P,$V,$N,$K,$NF,$MT,$FT}(sa, mask)) + NF = count(i -> !_is_exact_selector_type(fieldtypes(MT)[i]), 1:N) + return :(SparseArraySlice{$P,$V,$NF,$MT}(sa, mask)) end # Default: linear scan. Subtypes may override for cached lookup. function _view_matching_keys( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, -)::Vector{T} where {P,V,N,T,NF,MT,FT} + v::SparseArraySlice{P,V,NF,MT}, +) where {P,V,NF,MT} + T = _keytype(P) return collect(T, _select_gen(keys(_data(v.parent)), v.mask)) end @@ -133,17 +125,17 @@ Base.eltype(::Type{<:SparseArraySlice{P,V}}) where {P,V} = V # Iteration: values only (AbstractArray semantics) function Base.iterate( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, -) where {P,V,N,T,NF,MT,FT} + v::SparseArraySlice, +) matching = _view_matching_keys(v) isempty(matching) && return nothing return (v.parent[matching[1]], (matching, 2)) end function Base.iterate( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, - state::Tuple{Vector{T},Int}, -) where {P,V,N,T,NF,MT,FT} + v::SparseArraySlice, + state::Tuple{Vector,Int}, +) matching, pos = state pos > length(matching) && return nothing return (v.parent[matching[pos]], (matching, pos + 1)) @@ -151,27 +143,20 @@ end # getindex by FT tuple: v[(f, c)] function Base.getindex( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, - free_key::FT, -) where {P,V,N,T,NF,MT,FT} + v::SparseArraySlice{P,V,NF,MT}, + free_key::Tuple, +) where {P,V,NF,MT} + length(free_key) == NF || throw(BoundsError(v, free_key)) + T = _keytype(P) _matches_free_key(v, free_key) || return zero(V) return v.parent[_reconstruct_key(v.mask, free_key, T)] end -# Disambiguate vs AbstractSparseArray's NTuple method. -function Base.getindex( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, - idx::NTuple{NF,Any}, -) where {P,V,N,T,NF,MT,FT} - _matches_free_key(v, idx) || return zero(V) - return v.parent[_reconstruct_key(v.mask, idx, T)] -end - # Splatted: v[f, c] or v[f] (NF==1) function Base.getindex( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, + v::SparseArraySlice{P,V,NF,MT}, idx..., -) where {P,V,N,T,NF,MT,FT} +) where {P,V,NF,MT} length(idx) == NF || throw(BoundsError(v, idx)) return v[idx] end @@ -188,9 +173,10 @@ function Base.size(::SparseArraySlice) end function Base.haskey( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, - free_key::FT, -) where {P,V,N,T,NF,MT,FT} + v::SparseArraySlice{P,V,NF,MT}, + free_key::Tuple, +) where {P,V,NF,MT} + T = _keytype(P) _matches_free_key(v, free_key) || return false return haskey(_data(v.parent), _reconstruct_key(v.mask, free_key, T)) end @@ -198,8 +184,8 @@ end Base.length(v::SparseArraySlice) = length(_view_matching_keys(v)) function Base.keys( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, -) where {P,V,N,T,NF,MT,FT} + v::SparseArraySlice{P,V,NF,MT}, +) where {P,V,NF,MT} return [_project_free(k, MT) for k in _view_matching_keys(v)] end @@ -208,8 +194,8 @@ Base.values(v::SparseArraySlice) = [v.parent[k] for k in _view_matching_keys(v)] Base.eachindex(v::SparseArraySlice) = keys(v) function Base.pairs( - v::SparseArraySlice{P,V,N,T,NF,MT,FT}, -) where {P,V,N,T,NF,MT,FT} + v::SparseArraySlice{P,V,NF,MT}, +) where {P,V,NF,MT} return [_project_free(k, MT) => v.parent[k] for k in _view_matching_keys(v)] end diff --git a/test/Project.toml b/test/Project.toml index 7f33635..7988d16 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -2,4 +2,5 @@ Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4" HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" +SparseVariables = "2749762c-80ed-4b14-8f33-f0736679b02b" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" From 82ae7fc414ace923e0106ec555dafd5f118443e2 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Wed, 24 Jun 2026 13:28:09 +0200 Subject: [PATCH 12/22] Format fix --- src/slice.jl | 30 +++++++----------------------- test/runtests.jl | 2 +- 2 files changed, 8 insertions(+), 24 deletions(-) diff --git a/src/slice.jl b/src/slice.jl index 618342a..fed6e98 100644 --- a/src/slice.jl +++ b/src/slice.jl @@ -58,9 +58,7 @@ struct SparseArraySlice{P<:AbstractSparseArray,V,NF,MT<:Tuple} <: mask::MT end -function _keytype( - ::Type{<:SparseArraySlice{P,V,NF,MT}}, -) where {P,V,NF,MT} +function _keytype(::Type{<:SparseArraySlice{P,V,NF,MT}}) where {P,V,NF,MT} return _free_keytype(MT, _keytype(P)) end @@ -111,9 +109,7 @@ end end # Default: linear scan. Subtypes may override for cached lookup. -function _view_matching_keys( - v::SparseArraySlice{P,V,NF,MT}, -) where {P,V,NF,MT} +function _view_matching_keys(v::SparseArraySlice{P,V,NF,MT}) where {P,V,NF,MT} T = _keytype(P) return collect(T, _select_gen(keys(_data(v.parent)), v.mask)) end @@ -124,18 +120,13 @@ Base.IteratorEltype(::Type{<:SparseArraySlice}) = Base.HasEltype() Base.eltype(::Type{<:SparseArraySlice{P,V}}) where {P,V} = V # Iteration: values only (AbstractArray semantics) -function Base.iterate( - v::SparseArraySlice, -) +function Base.iterate(v::SparseArraySlice) matching = _view_matching_keys(v) isempty(matching) && return nothing return (v.parent[matching[1]], (matching, 2)) end -function Base.iterate( - v::SparseArraySlice, - state::Tuple{Vector,Int}, -) +function Base.iterate(v::SparseArraySlice, state::Tuple{Vector,Int}) matching, pos = state pos > length(matching) && return nothing return (v.parent[matching[pos]], (matching, pos + 1)) @@ -153,10 +144,7 @@ function Base.getindex( end # Splatted: v[f, c] or v[f] (NF==1) -function Base.getindex( - v::SparseArraySlice{P,V,NF,MT}, - idx..., -) where {P,V,NF,MT} +function Base.getindex(v::SparseArraySlice{P,V,NF,MT}, idx...) where {P,V,NF,MT} length(idx) == NF || throw(BoundsError(v, idx)) return v[idx] end @@ -183,9 +171,7 @@ end Base.length(v::SparseArraySlice) = length(_view_matching_keys(v)) -function Base.keys( - v::SparseArraySlice{P,V,NF,MT}, -) where {P,V,NF,MT} +function Base.keys(v::SparseArraySlice{P,V,NF,MT}) where {P,V,NF,MT} return [_project_free(k, MT) for k in _view_matching_keys(v)] end @@ -193,9 +179,7 @@ Base.values(v::SparseArraySlice) = [v.parent[k] for k in _view_matching_keys(v)] Base.eachindex(v::SparseArraySlice) = keys(v) -function Base.pairs( - v::SparseArraySlice{P,V,NF,MT}, -) where {P,V,NF,MT} +function Base.pairs(v::SparseArraySlice{P,V,NF,MT}) where {P,V,NF,MT} return [_project_free(k, MT) => v.parent[k] for k in _view_matching_keys(v)] end diff --git a/test/runtests.jl b/test/runtests.jl index 961c2f9..e77e2e9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -407,7 +407,7 @@ const _test_sa = testdata_sa() @test_throws ErrorException size(v) # wrong mask length - @test_throws BoundsError slice(sa,"ford",:,:) + @test_throws BoundsError slice(sa, "ford", :, :) # empty slice ve = slice(sa, "xxx", :) From 37fee7606b908ee215fff07a5e02c46719b58c87 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Wed, 24 Jun 2026 13:56:43 +0200 Subject: [PATCH 13/22] Update ci workflows --- .github/workflows/ci.yml | 6 ++++-- .github/workflows/codecov.yml | 2 -- .github/workflows/documentation.yml | 2 +- .github/workflows/format_check.yml | 4 ++-- 4 files changed, 7 insertions(+), 7 deletions(-) delete mode 100644 .github/workflows/codecov.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f9c8abb..8b93ca3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -49,6 +49,8 @@ jobs: with: depwarn: error - uses: julia-actions/julia-processcoverage@v1 - - uses: codecov/codecov-action@v3 + - uses: codecov/codecov-action@v5 with: - file: lcov.info + files: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + diff --git a/.github/workflows/codecov.yml b/.github/workflows/codecov.yml deleted file mode 100644 index 3217852..0000000 --- a/.github/workflows/codecov.yml +++ /dev/null @@ -1,2 +0,0 @@ -codecov: - token: 3c305759-ecc3-419d-aee0-20ab9a53d1e9 \ No newline at end of file diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml index 76aff3f..31c3146 100644 --- a/.github/workflows/documentation.yml +++ b/.github/workflows/documentation.yml @@ -9,7 +9,7 @@ jobs: build: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@latest with: version: '1' diff --git a/.github/workflows/format_check.yml b/.github/workflows/format_check.yml index 88045c8..4f91f39 100644 --- a/.github/workflows/format_check.yml +++ b/.github/workflows/format_check.yml @@ -13,13 +13,13 @@ jobs: - uses: julia-actions/setup-julia@latest with: version: '1' - - uses: actions/checkout@v1 + - uses: actions/checkout@v4 - name: Format check shell: julia --color=yes {0} run: | using Pkg # If you update the version, also update the style guide docs. - Pkg.add(PackageSpec(name="JuliaFormatter", version="1.0.13")) + Pkg.add(PackageSpec(name="JuliaFormatter", version="2")) using JuliaFormatter format("."; verbose = true) out = String(read(Cmd(`git diff`))) From c2728e55a5ea802f37d87084d39224ce89f3f9bd Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Wed, 24 Jun 2026 14:09:19 +0200 Subject: [PATCH 14/22] Fix format --- benchmark/benchmarks.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 1263277..97933e5 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -296,7 +296,7 @@ function model_sparse(F, C, P, T, D, U, V, W) # Variable creation @variable( m, - x[factory = F, customer = C, product = P, period = T], + x[factory=F, customer=C, product=P, period=T], container = IndexedVarArray ) @@ -338,10 +338,10 @@ function model_sparse_aa(F, C, P, T, D, U, V, W) @variable( m, x[ - factory = F, - customer = C, - product = P, - period = T; + factory=F, + customer=C, + product=P, + period=T; W[factory, product] == 1 && (factory, product, period) in keys(D), ] >= 0, ) From a7ed014b7973edd793d7470e5c7a0ebac469122b Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Wed, 24 Jun 2026 14:48:21 +0200 Subject: [PATCH 15/22] Make slices writable --- src/indexedarray.jl | 2 +- src/slice.jl | 22 ++++++++++++++++++++-- test/runtests.jl | 12 ++++++++---- 3 files changed, 29 insertions(+), 7 deletions(-) diff --git a/src/indexedarray.jl b/src/indexedarray.jl index aff3a55..edc708d 100644 --- a/src/indexedarray.jl +++ b/src/indexedarray.jl @@ -67,7 +67,7 @@ end """ unsafe_insertvar!(var::indexedVarArray{V,N,T}, index...) -Insert a new variable with the given index withouth checking if the index is valid or +Insert a new variable with the given index without checking if the index is valid or already assigned. """ function unsafe_insertvar!(var::IndexedVarArray{V,N,T}, index...) where {V,N,T} diff --git a/src/slice.jl b/src/slice.jl index fed6e98..eb3310f 100644 --- a/src/slice.jl +++ b/src/slice.jl @@ -149,8 +149,26 @@ function Base.getindex(v::SparseArraySlice{P,V,NF,MT}, idx...) where {P,V,NF,MT} return v[idx] end -function Base.setindex!(::SparseArraySlice, _, _...) - return error("SparseArraySlice is read-only") +# Forward mutation to parent array +function Base.setindex!( + v::SparseArraySlice{P,V,NF,MT}, + val, + free_key::Tuple, +) where {P,V,NF,MT} + length(free_key) == NF || throw(BoundsError(v, free_key)) + T = _keytype(P) + v.parent[_reconstruct_key(v.mask, free_key, T)] = val + return val +end + +# Splatted version +function Base.setindex!( + v::SparseArraySlice{P,V,NF,MT}, + val, + idx..., +) where {P,V,NF,MT} + length(idx) == NF || throw(BoundsError(v, idx)) + return setindex!(v, val, idx) end function Base.size(::SparseArraySlice) diff --git a/test/runtests.jl b/test/runtests.jl index e77e2e9..4e3563b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -396,16 +396,20 @@ const _test_sa = testdata_sa() @test Base.IteratorSize(typeof(v)) == Base.HasLength() @test Base.IteratorEltype(typeof(v)) == Base.HasEltype() + # setindex + v2["bmw", 2002] = 200 + @test v2["bmw", 2002] == 200 + @test sa["bmw", 2002] == 200 + v2[("bmw", 2002)] = 300 + @test v2[("bmw", 2002)] == 300 + @test sa[("bmw", 2002)] == 300 + # show / summary @test occursin("SparseArraySlice", sprint(summary, v)) @test occursin("matching (\"ford\", Colon())", sprint(summary, v)) @test occursin("(2000,) => 100", sprint(show, MIME("text/plain"), v)) @test occursin("(2001,) => 150", sprint(show, MIME("text/plain"), v)) - # read-only - @test_throws MethodError (v[(2000,)] = 999) - @test_throws ErrorException size(v) - # wrong mask length @test_throws BoundsError slice(sa, "ford", :, :) From 9665183479a99e778dda8cba877735eae2027063 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Jul 2026 14:51:48 +0000 Subject: [PATCH 16/22] Add memoization for _view_matching_keys in SparseArraySlice --- src/slice.jl | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/slice.jl b/src/slice.jl index eb3310f..c2af607 100644 --- a/src/slice.jl +++ b/src/slice.jl @@ -56,6 +56,10 @@ struct SparseArraySlice{P<:AbstractSparseArray,V,NF,MT<:Tuple} <: AbstractSparseArray{V,NF} parent::P mask::MT + _cache::Ref{Any} + function SparseArraySlice{P,V,NF,MT}(parent::P, mask::MT) where {P,V,NF,MT} + return new{P,V,NF,MT}(parent, mask, Ref{Any}(nothing)) + end end function _keytype(::Type{<:SparseArraySlice{P,V,NF,MT}}) where {P,V,NF,MT} @@ -108,10 +112,14 @@ end return :(SparseArraySlice{$P,$V,$NF,$MT}(sa, mask)) end -# Default: linear scan. Subtypes may override for cached lookup. +# Default: linear scan with memoization. Subtypes may override for cached lookup. function _view_matching_keys(v::SparseArraySlice{P,V,NF,MT}) where {P,V,NF,MT} + cached = v._cache[] + cached !== nothing && return cached::Vector{_keytype(P)} T = _keytype(P) - return collect(T, _select_gen(keys(_data(v.parent)), v.mask)) + keys_vec = collect(T, _select_gen(keys(_data(v.parent)), v.mask)) + v._cache[] = keys_vec + return keys_vec end # Iterator traits From 517eb2ad984e995fcc1961a63085168099943e82 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Jul 2026 15:02:43 +0000 Subject: [PATCH 17/22] Generalize make_filter_fun to handle AbstractRange --- src/dictionaries.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/dictionaries.jl b/src/dictionaries.jl index 8c3d2db..1556dad 100644 --- a/src/dictionaries.jl +++ b/src/dictionaries.jl @@ -12,7 +12,8 @@ make_filter_fun(c) = x -> x == c make_filter_fun(c::Base.Fix2) = x -> c(x) make_filter_fun(c::Function) = x -> c(x) make_filter_fun(c::Colon) = x -> true -make_filter_fun(c::UnitRange) = x -> (x ≥ c.start && x ≤ c.stop) +make_filter_fun(c::AbstractRange, pos) = x -> x[pos] in c +make_filter_fun(c::AbstractRange) = x -> x in c """ recursive_filter(fs, data) From 67da7269f4fc33e4633f2266b33aeccc59458dfb Mon Sep 17 00:00:00 2001 From: Lars Hellemo Date: Thu, 9 Jul 2026 13:37:52 +0200 Subject: [PATCH 18/22] Formatting tutorial --- tutorial/sparse_tutorial.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tutorial/sparse_tutorial.jl b/tutorial/sparse_tutorial.jl index 48c444c..5c0794b 100644 --- a/tutorial/sparse_tutorial.jl +++ b/tutorial/sparse_tutorial.jl @@ -186,8 +186,8 @@ begin @constraint( m, sum( - x[f, c, p, t] for (f, c, p, t) in - filter(i -> i[1] == f̄ && i[4] == t̄, indices) + x[f, c, p, t] for + (f, c, p, t) in filter(i -> i[1] == f̄ && i[4] == t̄, indices) ) ≤ U[f̄, t̄] ) end @@ -244,8 +244,8 @@ begin @constraint( m, sum( - x[(f, c, p, t)] for (f, c, p, t) in - filter(i -> i[1] == f̄ && i[4] == t̄, indices) + x[(f, c, p, t)] for + (f, c, p, t) in filter(i -> i[1] == f̄ && i[4] == t̄, indices) ) ≤ U[f̄, t̄] ) end From 8f78d1848beebc202f0ac3b9c9a812a074f545e6 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Mon, 27 Jul 2026 13:52:02 +0200 Subject: [PATCH 19/22] Refactor SparseBroadcastStyle and enhance broadcasting tests with JuMP scalar --- src/broadcast.jl | 23 +++++++++++------------ test/runtests.jl | 22 ++++++++++++++++++++++ 2 files changed, 33 insertions(+), 12 deletions(-) diff --git a/src/broadcast.jl b/src/broadcast.jl index a7e4dbe..d4e4e12 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -7,21 +7,20 @@ """ SparseBroadcastStyle{N,K} <: Broadcast.BroadcastStyle -Broadcasting style for all `AbstractSparseArray` subtypes. `N` is the key -dimensionality and `K` is the key tuple type. All broadcast results are -materialised as `SparseArray`. +Broadcasting style for all `AbstractSparseArray` subtypes. `K` is the key tuple type. +All broadcast results are materialised as `SparseArray`. """ -struct SparseBroadcastStyle{N,K} <: Broadcast.BroadcastStyle end +struct SparseBroadcastStyle{K} <: Broadcast.BroadcastStyle end function Base.BroadcastStyle(::Type{SA}) where {SA<:AbstractSparseArray} - return SparseBroadcastStyle{ndims(SA),_keytype(SA)}() + return SparseBroadcastStyle{_keytype(SA)}() end # Disallow mixing with other array types. function Base.BroadcastStyle(::SparseBroadcastStyle, ::Base.BroadcastStyle) return throw( ArgumentError( - "Cannot broadcast a SparseArray with another array of a different type", + "Cannot broadcast a SparseArray with incompatible key types", ), ) end @@ -50,8 +49,7 @@ function Base.Broadcast.instantiate( return bc end -# ── Internal helpers ────────────────────────────────────────────────────────── - +# Internal helpers _sparse_getindex(x::AbstractSparseArray, key) = x[key] _sparse_getindex(x::Any, ::Any) = x _sparse_getindex(x::Ref, ::Any) = x[] @@ -102,13 +100,14 @@ end _sparse_indices(::Any, rest...) = _sparse_indices(rest...) -# ── Materialise ─────────────────────────────────────────────────────────────── +# Materialise function Base.copy( - bc::Base.Broadcast.Broadcasted{SparseBroadcastStyle{N,K}}, -) where {N,K} + bc::Base.Broadcast.Broadcasted{SparseBroadcastStyle{K}}, +) where {K} indices = _sparse_indices(bc) - isempty(indices) && return SparseArray(Dictionary{K,Any}()) + T = Base.Broadcast.combine_eltypes(bc.f, bc.args) + isempty(indices) && return SparseArray(Dictionary{K,T}()) vals = [_sparse_getindex(bc, k) for k in indices] return SparseArray(Dictionary(indices, vals)) end diff --git a/test/runtests.jl b/test/runtests.jl index 4e3563b..74ecdf7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -491,6 +491,28 @@ end @test length(r_empty) == 0 end +@testset "Broadcasting SparseArray with JuMP scalar" begin + sa = _test_sa # SparseArray{Int} keyed by (String, Int) + m = Model() + @variable(m, t) + + # sa .* scalar-variable + ra = sa .* t + @test ra isa SparseArray + @test eltype(ra) <: JuMP.AbstractJuMPScalar + @test JuMP.isequal_canonical(ra["ford", 2000], 100 * t) + @test length(ra) == length(sa) + + # scalar-variable on the left + ra2 = t .* sa + @test JuMP.isequal_canonical(ra2["bmw", 2001], 200 * t) + + # sa .+ scalar-expression + rb = sa .+ (2t + 1) + @test rb isa SparseArray + @test JuMP.isequal_canonical(rb["ford", 2000], 100 + 2t + 1) +end + @testset "Broadcasting SparseArraySlice" begin sa = _test_sa From d5acc8c69984163b94fa739ea0dfdebfe1e41ec0 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Mon, 27 Jul 2026 13:55:30 +0200 Subject: [PATCH 20/22] Align formatting with JuMP standard --- .JuliaFormatter.toml | 1 + JuliaFormatter.toml | 8 -------- benchmark/benchmarks.jl | 8 ++++---- benchmark/transport.jl | 8 ++++---- docs/make.jl | 2 +- docs/notebook_juliacon2022.jl | 8 ++++---- tutorial/sparse_tutorial.jl | 24 ++++++++++++++---------- 7 files changed, 28 insertions(+), 31 deletions(-) delete mode 100644 JuliaFormatter.toml diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml index a700a07..f196f10 100644 --- a/.JuliaFormatter.toml +++ b/.JuliaFormatter.toml @@ -5,4 +5,5 @@ always_for_in = true always_use_return = true margin = 80 remove_extra_newlines = true +separate_kwargs_with_semicolon = true short_to_long_function_def = true diff --git a/JuliaFormatter.toml b/JuliaFormatter.toml deleted file mode 100644 index 74a0a71..0000000 --- a/JuliaFormatter.toml +++ /dev/null @@ -1,8 +0,0 @@ -# Configuration file for JuliaFormatter.jl -# For more information, see: https://domluna.github.io/JuliaFormatter.jl/stable/config/ - -always_for_in = true -always_use_return = true -margin = 80 -remove_extra_newlines = true -short_to_long_function_def = true \ No newline at end of file diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 97933e5..afccf0e 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -455,7 +455,7 @@ REPS = 5 # ╔═╡ 04570ea7-885c-4d0e-be88-eb2a5f77da90 begin - res = DataFrame(Method = Symbol[], NC = Int[], Time = Float64[]) + res = DataFrame(; Method = Symbol[], NC = Int[], Time = Float64[]) @progress for nc in 5:10:100 for method in [ model_standard, @@ -476,7 +476,7 @@ end # ╔═╡ cc097148-23b1-4584-a150-c7f22376b65c begin - sparsity = DataFrame(Method = Symbol[], DP = Float64[], Time = Float64[]) + sparsity = DataFrame(; Method = Symbol[], DP = Float64[], Time = Float64[]) @progress for dp in 0.05:0.05:1.0 for method in [ model_standard, @@ -511,10 +511,10 @@ end # ╔═╡ b0aa0499-e920-4014-b0b2-ce8ea3da7c95 function plot(df, x = :NC, y = :Time) - CairoMakie.activate!(type = "svg") + CairoMakie.activate!(; type = "svg") return draw( data(df) * - mapping(x, y => "Time (s)", color = :Method, marker = :Method) * + mapping(x, y => "Time (s)"; color = :Method, marker = :Method) * (visual(Lines) + visual(Scatter)), ) end diff --git a/benchmark/transport.jl b/benchmark/transport.jl index 4f19faf..47b326f 100644 --- a/benchmark/transport.jl +++ b/benchmark/transport.jl @@ -360,7 +360,7 @@ function create_vars_indexedtable(m, pp) C, P, T, - V, + V; names = [:factory, :customer, :product, :period, :var], pkey = [:factory, :customer, :product, :period], ) @@ -371,7 +371,7 @@ function create_constraints_indexedtable(m, pp) flow = m[:flow] # Production capacity - pc_table = groupby(collect, flow, (:factory, :product), select = :var) + pc_table = groupby(collect, flow, (:factory, :product); select = :var) for r in rows(pc_table) if (r.factory, r.product) in keys(pp.prodcap) @constraint( @@ -384,7 +384,7 @@ function create_constraints_indexedtable(m, pp) # Customer demand cpp_table = - groupby(collect, flow, (:customer, :product, :period), select = :var) + groupby(collect, flow, (:customer, :product, :period); select = :var) for r in rows(cpp_table) if (r.customer, r.product, r.period) in keys(pp.demand) @constraint( @@ -397,7 +397,7 @@ function create_constraints_indexedtable(m, pp) # Transport capacity fc_table = - groupby(collect, flow, (:factory, :customer, :period), select = :var) + groupby(collect, flow, (:factory, :customer, :period); select = :var) for r in rows(fc_table) if (r.factory, r.customer) in keys(pp.flowcap) @constraint( diff --git a/docs/make.jl b/docs/make.jl index 8314890..ec7879a 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -9,7 +9,7 @@ pages = [ "API reference" => "reference/api.md", ] -Documenter.makedocs( +Documenter.makedocs(; sitename = "SparseVariables", format = Documenter.HTML(; prettyurls = get(ENV, "CI", "false") == "true", diff --git a/docs/notebook_juliacon2022.jl b/docs/notebook_juliacon2022.jl index 8ea55b6..3cf6ac0 100644 --- a/docs/notebook_juliacon2022.jl +++ b/docs/notebook_juliacon2022.jl @@ -408,7 +408,7 @@ REPS = 5 # ╔═╡ 04570ea7-885c-4d0e-be88-eb2a5f77da90 begin - res = DataFrame(Method = Symbol[], NC = Int[], Time = Float64[]) + res = DataFrame(; Method = Symbol[], NC = Int[], Time = Float64[]) @progress for nc in 5:10:100 for method in [ model_standard, @@ -428,7 +428,7 @@ end # ╔═╡ cc097148-23b1-4584-a150-c7f22376b65c begin - sparsity = DataFrame(Method = Symbol[], DP = Float64[], Time = Float64[]) + sparsity = DataFrame(; Method = Symbol[], DP = Float64[], Time = Float64[]) @progress for dp in 0.05:0.05:1.0 for method in [ model_standard, @@ -462,10 +462,10 @@ end # ╔═╡ b0aa0499-e920-4014-b0b2-ce8ea3da7c95 function plot(df, x = :NC, y = :Time) - CairoMakie.activate!(type = "svg") + CairoMakie.activate!(; type = "svg") return draw( data(df) * - mapping(x, y => "Time (s)", color = :Method, marker = :Method) * + mapping(x, y => "Time (s)"; color = :Method, marker = :Method) * (visual(Lines) + visual(Scatter)), ) end diff --git a/tutorial/sparse_tutorial.jl b/tutorial/sparse_tutorial.jl index 5c0794b..1cc7aec 100644 --- a/tutorial/sparse_tutorial.jl +++ b/tutorial/sparse_tutorial.jl @@ -186,8 +186,8 @@ begin @constraint( m, sum( - x[f, c, p, t] for - (f, c, p, t) in filter(i -> i[1] == f̄ && i[4] == t̄, indices) + x[f, c, p, t] for (f, c, p, t) in + filter(i -> i[1] == f̄ && i[4] == t̄, indices) ) ≤ U[f̄, t̄] ) end @@ -244,8 +244,8 @@ begin @constraint( m, sum( - x[(f, c, p, t)] for - (f, c, p, t) in filter(i -> i[1] == f̄ && i[4] == t̄, indices) + x[(f, c, p, t)] for (f, c, p, t) in + filter(i -> i[1] == f̄ && i[4] == t̄, indices) ) ≤ U[f̄, t̄] ) end @@ -377,7 +377,7 @@ md" # ╔═╡ c49a3599-65fd-442b-b5c9-625b87e05efa begin - res = DataFrame(Method = Symbol[], NC = Int[], Time = Float64[]) + res = DataFrame(; Method = Symbol[], NC = Int[], Time = Float64[]) @progress for nc in 5:5:50 for method in [ model_standard, @@ -398,10 +398,10 @@ res # ╔═╡ 5ca68304-0c21-4344-a92b-0594c04674a4 function plot(df, x = :NC, y = :Time) - CairoMakie.activate!(type = "svg") + CairoMakie.activate!(; type = "svg") return draw( data(df) * - mapping(x, y, color = :Method, marker = :Method) * + mapping(x, y; color = :Method, marker = :Method) * (visual(Lines) + visual(Scatter)), ) end @@ -416,7 +416,7 @@ md" # ╔═╡ 1d68c8c0-1dbc-4d8b-97ae-3db1d2b06a4f begin - sparsity = DataFrame(Method = Symbol[], DP = Float64[], Time = Float64[]) + sparsity = DataFrame(; Method = Symbol[], DP = Float64[], Time = Float64[]) @progress for dp in 0.05:0.05:1.0 for method in [ model_standard, @@ -442,8 +442,12 @@ plot(sparsity, :DP, :Time) # ╔═╡ deefc30f-4846-49db-8ce8-69b4aec14924 begin - large = - DataFrame(Method = Symbol[], nc = Int[], vars = Int[], Time = Float64[]) + large = DataFrame(; + Method = Symbol[], + nc = Int[], + vars = Int[], + Time = Float64[], + ) @progress for nc in 500:500:5000 for method in [model_incremental, model_sparse] GC.gc() From bd0f4d8b2be8dc508c328279d5f7e7dd83c3c803 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Mon, 27 Jul 2026 14:34:00 +0200 Subject: [PATCH 21/22] Enhance SparseArraySlice with cache invalidation tests and update first/last index handling --- src/slice.jl | 16 ++++++++++++---- test/runtests.jl | 43 ++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 52 insertions(+), 7 deletions(-) diff --git a/src/slice.jl b/src/slice.jl index c2af607..9492ef0 100644 --- a/src/slice.jl +++ b/src/slice.jl @@ -90,6 +90,13 @@ for wildcard dimensions, exact values for fixed dimensions, and predicates or ranges for filtered dimensions. The result is an `AbstractSparseArray{V,NF}` where `NF` is the number of non-exact dimensions. +!!! note + + Slices materialise their matching keys lazily on first access and cache + them. Mutating through the slice invalidates the cache, but edits made + directly to the parent array afterwards are not tracked — re-create the + slice if the parent changes. + # Example ```julia v = slice(sa, :, "foo", :) # NF=2, two free dimensions @@ -166,6 +173,7 @@ function Base.setindex!( length(free_key) == NF || throw(BoundsError(v, free_key)) T = _keytype(P) v.parent[_reconstruct_key(v.mask, free_key, T)] = val + v._cache[] = nothing # invalidate: a new matching key may have been added return val end @@ -209,11 +217,11 @@ function Base.pairs(v::SparseArraySlice{P,V,NF,MT}) where {P,V,NF,MT} return [_project_free(k, MT) => v.parent[k] for k in _view_matching_keys(v)] end -function Base.firstindex(v::SparseArraySlice, d) - return minimum(k[d] for k in _view_matching_keys(v)) +function Base.firstindex(v::SparseArraySlice{P,V,NF,MT}, d) where {P,V,NF,MT} + return minimum(_project_free(k, MT)[d] for k in _view_matching_keys(v)) end -function Base.lastindex(v::SparseArraySlice, d) - return maximum(k[d] for k in _view_matching_keys(v)) +function Base.lastindex(v::SparseArraySlice{P,V,NF,MT}, d) where {P,V,NF,MT} + return maximum(_project_free(k, MT)[d] for k in _view_matching_keys(v)) end function Base.sum(v::SparseArraySlice{P,V}) where {P,V} diff --git a/test/runtests.jl b/test/runtests.jl index 74ecdf7..2873510 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -340,6 +340,43 @@ end const _test_sa = testdata_sa() +@testset "SparseArraySlice cache invalidation" begin + sa = SparseArray( + Dict( + ("ford", 2000) => 100, + ("ford", 2001) => 150, + ("bmw", 2001) => 200, + ), + ) + + v = slice(sa, "ford", :) + @test length(v) == 2 + + # Add a matching entry through the slice + v[(2002,)] = 175 + @test length(v) == 3 # cache invalidated + recomputed + @test v[(2002,)] == 175 + @test Set(keys(v)) == Set([(2000,), (2001,), (2002,)]) + @test sort(values(v)) == [100, 150, 175] + @test sum(v) == 100 + 150 + 175 + + # Add another matching entry using splatted setindex! + v[2003] = 50 + @test length(v) == 4 + @test v[2003] == 50 + @test sa["ford", 2003] == 50 # parent updated + @test length(sa) == 5 + + # Update an existing free key: count unchanged, value refreshed + v[(2000,)] = 999 + @test length(v) == 4 + @test v[(2000,)] == 999 + + # Mutating through the slice must not affect a non-matching dim + @test v[(2001,)] == 150 + @test sa["bmw", 2001] == 200 +end + @testset "SparseArraySlice on SparseArray" begin sa = _test_sa @@ -387,9 +424,9 @@ const _test_sa = testdata_sa() @test sum(slice(sa, :, 2001)) == 350 @test sum(slice(sa, "xxx", :)) == 0 - # firstindex / lastindex (d = parent-dimension index) - @test SV.firstindex(v, 2) == 2000 - @test SV.lastindex(v, 2) == 2001 + # firstindex / lastindex (d = free-dimension index, 1:NF) + @test SV.firstindex(v, 1) == 2000 + @test SV.lastindex(v, 1) == 2001 # iteration (values only) @test sum(val for val in v) == 250 From fd39a91c99647d3361494f2544c2cac07fa5fec0 Mon Sep 17 00:00:00 2001 From: Truls Flatberg Date: Mon, 27 Jul 2026 14:42:41 +0200 Subject: [PATCH 22/22] Format fix --- tutorial/sparse_tutorial.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tutorial/sparse_tutorial.jl b/tutorial/sparse_tutorial.jl index 1cc7aec..d20c8c3 100644 --- a/tutorial/sparse_tutorial.jl +++ b/tutorial/sparse_tutorial.jl @@ -186,8 +186,8 @@ begin @constraint( m, sum( - x[f, c, p, t] for (f, c, p, t) in - filter(i -> i[1] == f̄ && i[4] == t̄, indices) + x[f, c, p, t] for + (f, c, p, t) in filter(i -> i[1] == f̄ && i[4] == t̄, indices) ) ≤ U[f̄, t̄] ) end @@ -244,8 +244,8 @@ begin @constraint( m, sum( - x[(f, c, p, t)] for (f, c, p, t) in - filter(i -> i[1] == f̄ && i[4] == t̄, indices) + x[(f, c, p, t)] for + (f, c, p, t) in filter(i -> i[1] == f̄ && i[4] == t̄, indices) ) ≤ U[f̄, t̄] ) end