Make view(::AbstractWeights, ...) return an AbstractWeights

docs/src/weights.md

@@ -8,6 +8,7 @@ In statistical applications, it is not uncommon to assign weights to samples. To
 !!! note
     - The weight vector is a light-weight wrapper of the input vector. The input vector is NOT copied during construction.
     - The weight vector maintains the sum of weights, which is computed upon construction. If the value of the sum is pre-computed, one can supply it as the second argument to the constructor and save the time of computing the sum again.
+    - Views of weight vectors are also `AbstractWeights`, but their sum has to be recomputed by each call to `sum` as weights in the parent vector may have been mutated.
 
 
 ## Implementations

src/weights.jl

@@ -14,6 +14,7 @@ macro weights(name)
             sum::S
         end
         $(esc(name))(values::AbstractVector{<:Real}) = $(esc(name))(values, sum(values))
+        $(esc(:weightstype))(::Type{<:$(esc(name))}) = $(esc(name))
     end
 end
 
@@ -35,7 +36,33 @@ end
     W(v, sum(v))
 end
 
-Base.getindex(wv::W, ::Colon) where {W <: AbstractWeights} = W(copy(wv.values), sum(wv))
+Base.getindex(wv::AbstractWeights, ::Colon) = copy(wv)
+
+Base.copy(wv::W) where {W <: AbstractWeights} =
+    weightstype(W)(copy(wv.values), sum(wv))
+
+@propagate_inbounds function Base.view(wv::W, inds...) where {S <: Real, W <: AbstractWeights{S}}
+    @boundscheck checkbounds(wv, inds...)
+    @inbounds v = invoke(view, Tuple{AbstractArray, Vararg{Any}}, wv, inds...)
+    # Sum is not actually used but compute the right type for clarity
+    weightstype(W)(v, zero(S))
+end
+
+# This method is implemented for backward compatibility
+@propagate_inbounds function Base.view(wv::W, inds::Union{Integer, CartesianIndex}...) where
+    {S <: Real, W <: AbstractWeights{S}}
+    @boundscheck checkbounds(wv, inds...)
+    @inbounds invoke(view, Tuple{AbstractArray, Vararg{Any}}, wv, inds...)
+end
+
+# Always recompute the sum for views of AbstractWeights, as we cannot know whether
+# the parent array has been mutated
+Base.sum(wv::AbstractWeights{S, T, <:SubArray{T, <:Any, <:AbstractWeights}}) where
+    {S<:Real, T<:Real} = sum(wv.values)
+
+Base.copy(wv::W) where
+    {S<:Real, T<:Real, W<:AbstractWeights{S, T, <:SubArray{T, <:Any, <:AbstractWeights}}} =
+    weightstype(W)(copy(view(parent(wv.values).values, parentindices(wv.values)...)), wv.sum)
 
 @propagate_inbounds function Base.setindex!(wv::AbstractWeights, v::Real, i::Int)
     s = v - wv[i]
@@ -337,8 +364,12 @@ end
 
 for w in (AnalyticWeights, FrequencyWeights, ProbabilityWeights, Weights)
     @eval begin
-        Base.isequal(x::$w, y::$w) = isequal(x.sum, y.sum) && isequal(x.values, y.values)
-        Base.:(==)(x::$w, y::$w)   = (x.sum == y.sum) && (x.values == y.values)
+        Base.isequal(x::$w, y::$w) =
+            (x.values isa SubArray || y.values isa SubArray || isequal(x.values, y.values)) &&
+                isequal(x.values, y.values)
+        Base.:(==)(x::$w, y::$w)   =
+            (x.values isa SubArray || y.values isa SubArray || x.sum == y.sum) &&
+                (x.values == y.values)
     end
 end
 

test/weights.jl

@@ -37,7 +37,7 @@ weight_funcs = (weights, aweights, fweights, pweights)
     @test sum(sa, wv) === 7.0
 end
 
-@testset "$f, setindex!" for f in weight_funcs
+@testset "$f, getindex, view, setindex!" for f in weight_funcs
     w = [1., 2., 3.]
     wv = f(w)
 
@@ -46,6 +46,41 @@ end
     @test sum(wv) === 6.
     @test wv == w
 
+    @test wv[[1, 3]] == w[[1, 3]]
+    @test typeof(wv[[1, 3]]) === typeof(wv)
+    @test sum( wv[[1, 3]]) === sum(w[[1, 3]])
+
+    # Check copy
+    @test copy(wv) == wv
+    @test typeof(copy(wv)) === typeof(wv)
+
+    # Check view
+    @test view(wv, :) == wv
+    @test typeof(view(wv, :)) <: StatsBase.weightstype(typeof(wv))
+    @test view(wv, [1, 3]) == view(wv, [true, false, true]) == w[[1, 3]]
+    @test typeof(view(wv, [1, 3])) === typeof(view(wv, [true, false, true])) <:
+        StatsBase.weightstype(typeof(wv))
+    @test sum(view(wv, [1, 3])) === sum(view(wv, [true, false, true])) === sum(w[[1, 3]])
+    @test_throws BoundsError view(wv, [1, 5])
+    @test_throws BoundsError view(wv, [true, false, true, true])
+    v = view(wv, [1, 3])
+    wv[1] += 1
+    @test sum(v) == sum(view(wv, [1, 3]))
+    v[1] -= 1
+    @test wv[1] == 1
+    @test copy(v) == v[:] == v
+    @test typeof(copy(v)) === typeof(v[:]) == typeof(wv)
+    @test view(wv, [1, 2], 1) == wv[[1, 2]]
+    @test typeof(view(wv, [1, 2], 1)) <: StatsBase.weightstype(typeof(wv))
+    @test view(wv, 1) == fill(wv[1])
+    @test view(wv, 1) isa SubArray
+    @test view(wv, 1, 1) == fill(wv[1])
+    @test view(wv, 1, 1) isa SubArray
+    @test view(wv, CartesianIndex(1, 1)) == fill(wv[CartesianIndex(1, 1)])
+    @test view(wv, CartesianIndex(1, 1)) isa SubArray
+    @test view(wv, CartesianIndex(1, 1), 1) == fill(wv[CartesianIndex(1, 1), 1])
+    @test view(wv, CartesianIndex(1, 1), 1) isa SubArray
+
     # Test setindex! success
     @test (wv[1] = 4) === 4             # setindex! returns original val
     @test wv[1] === 4.                  # value correctly converted and set