More efficient PoissonBinomial (#1285)

Author: devmotion

src/univariate/discrete/poissonbinomial.jl

@@ -1,4 +1,3 @@
-# TODO: this distribution may need clean-up
 """
     PoissonBinomial(p)
 
@@ -24,35 +23,45 @@ External links:
 * [Poisson-binomial distribution on Wikipedia](http://en.wikipedia.org/wiki/Poisson_binomial_distribution)
 
 """
-struct PoissonBinomial{T<:Real} <: DiscreteUnivariateDistribution
-    p::Vector{T}
-    pmf::Vector{T}
-
-    function PoissonBinomial{T}(p::AbstractArray) where {T <: Real}
-        pb = poissonbinomial_pdf(p)
-        @assert isprobvec(pb)
-        new{T}(p, pb)
+mutable struct PoissonBinomial{T<:Real,P<:AbstractVector{T}} <: DiscreteUnivariateDistribution
+    p::P
+    pmf::Union{Nothing,Vector{T}} # lazy computation of the probability mass function
+
+    function PoissonBinomial{T}(p::AbstractVector{T}; check_args=true) where {T <: Real}
+        check_args && @check_args(PoissonBinomial, all(x -> zero(x) <= x <= one(x), p))
+        return new{T,typeof(p)}(p, nothing)
     end
 end
 
-function PoissonBinomial(p::AbstractArray{T}; check_args=true) where {T <: Real}
-    if check_args
-        for i in eachindex(p)
-            @check_args(PoissonBinomial, 0 <= p[i] <= 1)
+function PoissonBinomial(p::AbstractVector{T}; check_args=true) where {T<:Real}
+    return PoissonBinomial{T}(p; check_args=check_args)
+end
+
+function Base.getproperty(d::PoissonBinomial, x::Symbol)
+    if x === :pmf
+        z = getfield(d, :pmf)
+        if z === nothing
+            y = poissonbinomial_pdf(d.p)
+            isprobvec(y) || error("probability mass function is not normalized")
+            setfield!(d, :pmf, y)
+            return y
+        else
+            return z
         end
+    else
+        return getfield(d, x)
     end
-    return PoissonBinomial{T}(p)
 end
 
 @distr_support PoissonBinomial 0 length(d.p)
 
 #### Conversions
 
-function PoissonBinomial(::Type{PoissonBinomial{T}}, p::Vector{S}) where {T, S}
-    return PoissonBinomial(Vector{T}(p))
+function PoissonBinomial(::Type{PoissonBinomial{T}}, p::AbstractVector{S}) where {T, S}
+    return PoissonBinomial(AbstractVector{T}(p))
 end
 function PoissonBinomial(::Type{PoissonBinomial{T}}, d::PoissonBinomial{S}) where {T, S}
-    return PoissonBinomial(Vector{T}(d.p), check_args=false)
+    return PoissonBinomial(AbstractVector{T}(d.p), check_args=false)
 end
 
 #### Parameters
@@ -67,7 +76,7 @@ partype(::PoissonBinomial{T}) where {T} = T
 #### Properties
 
 mean(d::PoissonBinomial) = sum(succprob(d))
-var(d::PoissonBinomial) = sum(succprob(d) .* failprob(d))
+var(d::PoissonBinomial) = sum(p * (1 - p) for p in succprob(d))
 
 function skewness(d::PoissonBinomial{T}) where {T}
     v = zero(T)
@@ -91,23 +100,27 @@ function kurtosis(d::PoissonBinomial{T}) where {T}
     s / v / v
 end
 
-entropy(d::PoissonBinomial) = entropy(Categorical(d.pmf))
+entropy(d::PoissonBinomial) = entropy(d.pmf)
 median(d::PoissonBinomial) = median(Categorical(d.pmf)) - 1
 mode(d::PoissonBinomial) = argmax(d.pmf) - 1
-modes(d::PoissonBinomial) = [x  - 1 for x in modes(Categorical(d.pmf))]
+modes(d::PoissonBinomial) = modes(DiscreteNonParametric(support(d), d.pmf))
 
 #### Evaluation
 
 quantile(d::PoissonBinomial, x::Float64) = quantile(Categorical(d.pmf), x) - 1
 
-function mgf(d::PoissonBinomial{T}, t::Real) where {T}
-    p,  = params(d)
-    prod(one(T) .- p .+ p .* exp(t))
+function mgf(d::PoissonBinomial, t::Real)
+    expm1_t = expm1(t)
+    mapreduce(*, succprob(d)) do p
+        1 + p * expm1_t
+    end
 end
 
-function cf(d::PoissonBinomial{T}, t::Real) where {T}
-    p,  = params(d)
-    prod(one(T) .- p .+ p .* cis(t))
+function cf(d::PoissonBinomial, t::Real)
+    cis_t = cis(t)
+    mapreduce(*, succprob(d)) do p
+        1 - p + p * cis_t
+    end
 end
 
 pdf(d::PoissonBinomial, k::Real) = insupport(d, k) ? d.pmf[k+1] : zero(eltype(d.pmf))
@@ -120,19 +133,17 @@ logpdf(d::PoissonBinomial, k::Real) = log(pdf(d, k))
 #      Calculating binomial probabilities when the trial probabilities are unequal, 
 #      Journal of Statistical Computation and Simulation, 14:2, 125-131, DOI: 10.1080/00949658208810534 
 #
-function poissonbinomial_pdf(p::AbstractArray{T,1}) where {T <: Real}
-  n = length(p)
-  S = zeros(T, n+1)
-  S[1] = 1-p[1]
-  S[2] = p[1]
-  @inbounds for col in 2:n
-    for r in 1:col
-        row = col - r + 1 
-        S[row+1] = (1-p[col])*S[row+1] + p[col] * S[row]
+function poissonbinomial_pdf(p)
+    S = zeros(eltype(p), length(p) + 1)
+    S[1] = 1
+    @inbounds for (col, p_col) in enumerate(p)
+        q_col = 1 - p_col
+        for row in col:(-1):1
+            S[row + 1] = q_col * S[row + 1] + p_col * S[row]
+        end
+        S[1] *= q_col
     end
-    S[1] *= 1-p[col]
-  end
-  return S
+    return S
 end
 
 # Computes the pdf of a poisson-binomial random variable using

test/poissonbinomial.jl

@@ -29,6 +29,7 @@ naive_sol = naive_pb(p)
 
 @test Distributions.poissonbinomial_pdf_fft(p) ≈ naive_sol
 @test Distributions.poissonbinomial_pdf(p) ≈ naive_sol
+@test Distributions.poissonbinomial_pdf(Tuple(p)) ≈ naive_sol
 
 @test Distributions.poissonbinomial_pdf_fft(p) ≈ Distributions.poissonbinomial_pdf(p)
 
@@ -46,23 +47,23 @@ for (p, n) in [(0.8, 6), (0.5, 10), (0.04, 20)]
     @test maximum(d) == n
     @test extrema(d) == (0, n)
     @test ntrials(d) == n
-    @test entropy(d)  ≈ entropy(dref)
-    @test median(d)   ≈ median(dref)
-    @test mean(d)     ≈ mean(dref)
-    @test var(d)      ≈ var(dref)
-    @test kurtosis(d) ≈ kurtosis(dref)
-    @test skewness(d) ≈ skewness(dref)
+    @test @inferred(entropy(d))  ≈ entropy(dref)
+    @test @inferred(median(d))   ≈ median(dref)
+    @test @inferred(mean(d))     ≈ mean(dref)
+    @test @inferred(var(d))      ≈ var(dref)
+    @test @inferred(kurtosis(d)) ≈ kurtosis(dref)
+    @test @inferred(skewness(d)) ≈ skewness(dref)
 
     for t=0:5
-        @test mgf(d, t) ≈ mgf(dref, t)
-        @test cf(d, t)  ≈ cf(dref, t)
+        @test @inferred(mgf(d, t)) ≈ mgf(dref, t)
+        @test @inferred(cf(d, t))  ≈ cf(dref, t)
     end
     for i=0.1:0.1:.9
-        @test quantile(d, i) ≈ quantile(dref, i)
+        @test @inferred(quantile(d, i)) ≈ quantile(dref, i)
     end
     for i=0:n
-        @test isapprox(cdf(d, i), cdf(dref, i), atol=1e-15)
-        @test isapprox(pdf(d, i), pdf(dref, i), atol=1e-15)
+        @test isapprox(@inferred(cdf(d, i)), cdf(dref, i), atol=1e-15)
+        @test isapprox(@inferred(pdf(d, i)), pdf(dref, i), atol=1e-15)
     end
 
 end
@@ -88,11 +89,11 @@ for (n₁, n₂, n₃, p₁, p₂, p₃) in [(10, 10, 10, 0.1, 0.5, 0.9),
     pmf2 = pdf.(b2, support(b2))
     pmf3 = pdf.(b3, support(b3))
 
-    @test mean(d) ≈ (mean(b1) + mean(b2) + mean(b3))
-    @test var(d)  ≈ (var(b1) + var(b2) + var(b3))
+    @test @inferred(mean(d)) ≈ (mean(b1) + mean(b2) + mean(b3))
+    @test @inferred(var(d))  ≈ (var(b1) + var(b2) + var(b3))
     for t=0:5
-        @test mgf(d, t) ≈ (mgf(b1, t) * mgf(b2, t) * mgf(b3, t))
-        @test cf(d, t)  ≈ (cf(b1, t) * cf(b2, t) * cf(b3, t))
+        @test @inferred(mgf(d, t)) ≈ (mgf(b1, t) * mgf(b2, t) * mgf(b3, t))
+        @test @inferred(cf(d, t))  ≈ (cf(b1, t) * cf(b2, t) * cf(b3, t))
     end
 
     for k=0:n
@@ -104,7 +105,7 @@ for (n₁, n₂, n₃, p₁, p₂, p₃) in [(10, 10, 10, 0.1, 0.5, 0.9),
             end
             m += pmf1[i+1] * mc
         end
-        @test isapprox(pdf(d, k), m, atol=1e-15)
+        @test isapprox(@inferred(pdf(d, k)), m, atol=1e-15)
     end
 end