some fixes and minor refactoring of arithmetic to allow generalization

Author: ablaom

src/methods.jl

@@ -388,45 +388,48 @@ const ERR_DIFFERENT_SAMPLE_SPACES = ArgumentError(
     "Adding two `UnivariateFinite` objects whose "*
     "sample spaces have different labellings is not allowed. ")
 
-import Base: +, *, /
+import Base: +, *, /, -
 
-function +(d1::U, d2::U) where U <: UnivariateFinite
+function _plus(d1, d2, T, N...) # N... is single integer or absent
     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 UnivariateFinite(S, decoder, prob_given_ref)
+    return T(S, decoder, prob_given_ref, N...)
 end
++(d1::U, d2::U) where U <: UnivariateFinite = _plus(d1, d2, UnivariateFinite)
 
-function -(d::UnivariateFinite)
+function _minus(d, T, N...)
     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 UnivariateFinite(S, decoder, prob_given_ref)
+    return T(S, decoder, prob_given_ref, N...)
 end
+-(d::UnivariateFinite) = _minus(d, UnivariateFinite)
 
-function -(d1::U, d2::U) where U <: UnivariateFinite
+function _minus(d1, d2, T, N...)
     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 UnivariateFinite(S, decoder, prob_given_ref)
+    return T(S, decoder, prob_given_ref, N...)
 end
+-(d1::U, d2::U) where U <: UnivariateFinite = _minus(d1, d2, UnivariateFinite)
 
 # 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 *(d::UnivariateFinite, x::Real)
+function _times(d, x, T, N...)
     S = d.scitype
     decoder = d.decoder
     prob_given_ref = copy(d.prob_given_ref)
@@ -435,6 +438,7 @@ function *(d::UnivariateFinite, x::Real)
     end
     return UnivariateFinite(d.scitype, decoder, prob_given_ref)
 end
-*(x::Real, d::UnivariateFinite) = d*x
+*(d::UnivariateFinite, x::Real) = _times(d, x, UnivariateFinite)
 
-/(d::UnivariateFinite, x::Real) = d*inv(x)
+*(x::Real, d::SingletonOrArray) = d*x
+/(d::SingletonOrArray, x::Real) = d*inv(x)

src/types.jl

@@ -203,6 +203,10 @@ end
 
 const UnivariateFiniteVector{S,V,R,P} = UnivariateFiniteArray{S,V,R,P,1}
 
+# private:
+const SingletonOrArray{S,V,R,P} = Union{UnivariateFinite{S,V,R,P},
+                                        UnivariateFiniteArray{S,V,R,P}}
+
 
 # # CHECKS AND ERROR MESSAGES
 
@@ -273,6 +277,7 @@ 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

@@ -277,10 +277,11 @@ end
     # @test v ≈ v_close
 end
 
-@tesset "arithmetic" begin
-    L = ["yes", "no"]
-    d1 = UnivariateFinite(L, rand(rng, 2), pool=missing)
-    d2 = UnivariateFinite(L, rand(rng, 2), pool=missing)
+L = ["yes", "no"]
+d1 = UnivariateFinite(L, rand(rng, 2), pool=missing)
+d2 = UnivariateFinite(L, rand(rng, 2), pool=missing)
+
+@testset "arithmetic" begin
 
     # addition and subtraction:
     for op in [:+, :-]
@@ -295,8 +296,10 @@ end
     @test pdf(d_neg, L) == -pdf(d1, L)
 
     # multiplication by scalar:
-    d3 = d1/42
-    @test pdf(d3, L) ≈ pdf(d1, L)/42
+    d3 = d1*42
+    @test pdf(d3, L) ≈ pdf(d1, L)*42
+    d3 = 42*d1
+    @test pdf(d3, L) ≈ pdf(d1, L)*42
 
     # division by scalar:
     d3 = d1/42