make arithmetic for UnivariateFiniteArray performant

Author: ablaom

src/methods.jl

@@ -390,55 +390,61 @@ const ERR_DIFFERENT_SAMPLE_SPACES = ArgumentError(
 
 import Base: +, *, /, -
 
-function _plus(d1, d2, T, N...) # N... is single integer or absent
+function _plus(d1, d2, T)
     classes(d1) == classes(d2) || throw(ERR_DIFFERENT_SAMPLE_SPACES)
     S = d1.scitype
     decoder = d1.decoder
     prob_given_ref = copy(d1.prob_given_ref)
     for ref in keys(prob_given_ref)
         prob_given_ref[ref] += d2.prob_given_ref[ref]
     end
-    return T(S, decoder, prob_given_ref, N...)
+    return T(S, decoder, prob_given_ref)
 end
 +(d1::U, d2::U) where U <: UnivariateFinite = _plus(d1, d2, UnivariateFinite)
++(d1::U, d2::U) where U <: UnivariateFiniteArray =
+    _plus(d1, d2, UnivariateFiniteArray)
 
-function _minus(d, T, N...)
+function _minus(d, T)
     S = d.scitype
     decoder = d.decoder
     prob_given_ref = copy(d.prob_given_ref)
     for ref in keys(prob_given_ref)
         prob_given_ref[ref] = -prob_given_ref[ref]
     end
-    return T(S, decoder, prob_given_ref, N...)
+    return T(S, decoder, prob_given_ref)
 end
 -(d::UnivariateFinite) = _minus(d, UnivariateFinite)
+-(d::UnivariateFiniteArray) = _minus(d, UnivariateFiniteArray)
 
-function _minus(d1, d2, T, N...)
+function _minus(d1, d2, T)
     classes(d1) == classes(d2) || throw(ERR_DIFFERENT_SAMPLE_SPACES)
     S = d1.scitype
     decoder = d1.decoder
     prob_given_ref = copy(d1.prob_given_ref)
     for ref in keys(prob_given_ref)
         prob_given_ref[ref] -= d2.prob_given_ref[ref]
     end
-    return T(S, decoder, prob_given_ref, N...)
+    return T(S, decoder, prob_given_ref)
 end
 -(d1::U, d2::U) where U <: UnivariateFinite = _minus(d1, d2, UnivariateFinite)
+-(d1::U, d2::U) where U <: UnivariateFiniteArray =
+    _minus(d1, d2, UnivariateFiniteArray)
 
 # TODO: remove type restrction on `x` in the following methods if
 # https://github.com/JuliaStats/Distributions.jl/issues/1438 is
 # resolved. Currently we'd have a method ambiguity
 
-function _times(d, x, T, N...)
+function _times(d, x, T)
     S = d.scitype
     decoder = d.decoder
     prob_given_ref = copy(d.prob_given_ref)
     for ref in keys(prob_given_ref)
         prob_given_ref[ref] *= x
     end
-    return UnivariateFinite(d.scitype, decoder, prob_given_ref)
+    return T(d.scitype, decoder, prob_given_ref)
 end
 *(d::UnivariateFinite, x::Real) = _times(d, x, UnivariateFinite)
+*(d::UnivariateFiniteArray, x::Real) = _times(d, x, UnivariateFiniteArray)
 
 *(x::Real, d::SingletonOrArray) = d*x
 /(d::SingletonOrArray, x::Real) = d*inv(x)

src/types.jl

@@ -277,7 +277,6 @@ function _augment_probs(::Val{true},
                         probs::AbstractArray{P,N}) where {P,N}
     _check_probs_01(probs)
     aug_size = [size(probs)..., 2]
-    @show probs P
     augmentation = one(P) .- probs
     all(0 .<= augmentation .<= 1) || throw(ERR_AUG)
     aug_probs = Array{P}(undef, aug_size...)

test/methods.jl

@@ -287,23 +287,55 @@ d2 = UnivariateFinite(L, rand(rng, 2), pool=missing)
     for op in [:+, :-]
         quote
             s = $op(d1, d2 )
-            @test $op(pdf(d1, L), pdf(d2, L)) ≈ pdf(s, L)
+            @test $op(pdf.(d1, L), pdf.(d2, L)) ≈ pdf.(s, L)
         end |> eval
     end
 
     # negative:
     d_neg = -d1
-    @test pdf(d_neg, L) == -pdf(d1, L)
+    @test pdf.(d_neg, L) == -pdf.(d1, L)
 
     # multiplication by scalar:
     d3 = d1*42
-    @test pdf(d3, L) ≈ pdf(d1, L)*42
+    @test pdf.(d3, L) ≈ pdf.(d1, L)*42
     d3 = 42*d1
-    @test pdf(d3, L) ≈ pdf(d1, L)*42
+    @test pdf.(d3, L) ≈ pdf.(d1, L)*42
 
     # division by scalar:
     d3 = d1/42
-    @test pdf(d3, L) ≈ pdf(d1, L)/42
+    @test pdf.(d3, L) ≈ pdf.(d1, L)/42
+end
+
+p = [0.1, 0.9]
+P = vcat(fill(p', 10^5)...)
+slow = fill(UnivariateFinite(L, p, pool=missing), 10^5)
+fast = UnivariateFinite(L, P, pool=missing)
+# @assert pdf(slow, L) == pdf(fast, L)
+
+@testset "performant arithmetic for UnivariateFiniteArray" begin
+    @test pdf(slow + slow, L) == pdf(fast + fast, L)
+    t_slow = @elapsed @eval slow + slow
+    t_fast = @elapsed @eval fast + fast
+    @test t_slow/t_fast > 10
+
+    @test pdf(slow - slow, L) == pdf(fast - fast, L)
+    t_slow = @elapsed @eval slow - slow
+    t_fast = @elapsed @eval fast - fast
+    @test t_slow/t_fast > 10
+
+    @test pdf(42*slow, L) == pdf(42*fast, L)
+    @test pdf(slow*42, L) == pdf(fast*42, L)
+    t_slow = @elapsed @eval 42*slow
+    t_fast = @elapsed @eval 42*fast
+    @test t_slow/t_fast > 10
+    t_slow = @elapsed @eval slow*42
+    t_fast = @elapsed @eval fast*42
+    @test t_slow/t_fast > 10
+
+    @test pdf(slow/42, L) == pdf(fast/42, L)
+    t_slow = @elapsed @eval slow/42
+    t_fast = @elapsed @eval fast/42
+    @test t_slow/t_fast > 10
 end
 
 end # module