Add broadcasting and slicing using a lazy view structure

.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

.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 }}
+

.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'

.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`)))

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,
     )
@@ -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

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(

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",

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

src/SparseVariables.jl

@@ -6,15 +6,20 @@ 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 SparseArraySlice
+export slice
 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
@@ -29,19 +34,15 @@ export SafeInsert, UnsafeInsert
         # 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)
         end
         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

src/broadcast.jl

@@ -0,0 +1,129 @@
+# ------------------------------------------------------------------------------
+# 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. `K` is the key tuple type.
+All broadcast results are materialised as `SparseArray`.
+"""
+struct SparseBroadcastStyle{K} <: Broadcast.BroadcastStyle end
+
+function Base.BroadcastStyle(::Type{SA}) where {SA<:AbstractSparseArray}
+    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 incompatible key types",
+        ),
+    )
+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
+
+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(
+    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{K}},
+) where {K}
+    indices = _sparse_indices(bc)
+    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
+
+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