Improve the logpdf of NegativeBinomial

Project.toml

@@ -1,7 +1,7 @@
 name = "Distributions"
 uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
 authors = ["JuliaStats"]
-version = "0.25.70"
+version = "0.25.71"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/univariate/discrete/negativebinomial.jl

@@ -97,14 +97,16 @@ end
 
 # Implement native pdf and logpdf since it's relatively straight forward and allows for ForwardDiff
 function logpdf(d::NegativeBinomial, k::Real)
-    r = d.r * log(d.p) + k * log1p(-d.p)
-    if isone(d.p) && iszero(k)
-        return zero(r)
-    elseif !insupport(d, k)
-        return oftype(r, -Inf)
-    else
-        return r - log(k + d.r) - logbeta(d.r, k + 1)
+    r, p = params(d)
+    z = xlogy(r, p) + xlog1py(k, -p)
+
+    if iszero(k)
+        # in this case `logpdf(d, k) = z - log(k + r) - logbeta(r, k + 1) = z` analytically
+        # but unfortunately not numerically, so we handle this case separately to improve accuracy
+        return z
     end
+
+    return insupport(d, k) ? z - log(k + r) - logbeta(r, k + 1) : oftype(z, -Inf)
 end
 
 # cdf and quantile functions are more involved so we still rely on Rmath
@@ -140,35 +142,39 @@ cf(d::NegativeBinomial, t::Real) = laplace_transform(d, -t*im)
 
 # ChainRules definitions
 
+## Callable struct to fix type inference issues caused by captured values
+struct LogPDFNegativeBinomialPullback{D,T<:Real}
+    ∂r::T
+    ∂p::T
+end
+
+function (f::LogPDFNegativeBinomialPullback{D})(Δ) where {D}
+    Δr = Δ * f.∂r
+    Δp = Δ * f.∂p
+    Δd = ChainRulesCore.Tangent{D}(; r=Δr, p=Δp)
+    return ChainRulesCore.NoTangent(), Δd, ChainRulesCore.NoTangent()
+end
+
 function ChainRulesCore.rrule(::typeof(logpdf), d::NegativeBinomial, k::Real)
-    # Compute log probability
+    # Compute log probability (as in the definition of `logpdf(d, k)` above)
     r, p = params(d)
-    edgecase = isone(p) && iszero(k)
-    insupp = insupport(d, k)
-
-    # Primal computation
-    Ω = r * log(p) + k * log1p(-p)
-    if edgecase
-        Ω = zero(Ω)
-    elseif !insupp
-        Ω = oftype(Ω, -Inf)
+    z = xlogy(r, p) + xlog1py(k, -p)
+    if iszero(k)
+        Ω = z
+        ∂r = oftype(z, log(p))
+        ∂p = oftype(z, r/p)
+    elseif insupport(d, k)
+        Ω = z - log(k + r) - logbeta(r, k + 1)
+        ∂r = oftype(z, log(p) - inv(k + r) - digamma(r) + digamma(r + k + 1))
+        ∂p = oftype(z, r/p - k / (1 - p))
     else
-        Ω = Ω - log(k + r) - logbeta(r, k + 1)
+        Ω = oftype(z, -Inf)
+        ∂r = oftype(z, NaN)
+        ∂p = oftype(z, NaN)
     end
 
     # Define pullback
-    function logpdf_NegativeBinomial_pullback(Δ)
-        Δr = Δ * (log(p) - inv(k + r) - digamma(r) + digamma(r + k + 1))
-        Δp = Δ * (r / p - k / (1 - p))
-        if edgecase
-            Δp = oftype(Δp, Δ * r)
-        elseif !insupp
-            Δr = oftype(Δr, NaN)
-            Δp = oftype(Δp, NaN)
-        end
-        Δd = ChainRulesCore.Tangent{typeof(d)}(; r=Δr, p=Δp)
-        return ChainRulesCore.NoTangent(), Δd, ChainRulesCore.NoTangent()
-    end
+    logpdf_NegativeBinomial_pullback = LogPDFNegativeBinomialPullback{typeof(d),typeof(z)}(∂r, ∂p)
 
     return Ω, logpdf_NegativeBinomial_pullback
 end

test/univariate/discrete/negativebinomial.jl

@@ -2,13 +2,17 @@ using Distributions
 using Test, ForwardDiff
 using ChainRulesTestUtils
 using FiniteDifferences
+using StatsFuns
 
 # Currently, most of the tests for NegativeBinomial are in the "ref" folder.
 # Eventually, we might want to consolidate the tests here
 
 test_cgf(NegativeBinomial(10,0.5), (-1f0, -200.0,-1e6))
 test_cgf(NegativeBinomial(3,0.1),  (-1f0, -200.0,-1e6))
-mydiffp(r, p, k) = r/p - k/(1 - p)
+
+mydiffp(r, p, k) = iszero(k) ? r/p : r/p - k/(1 - p)
+mydiffr(r, p, k) = iszero(k) ? log(p) : log(p) - inv(k + r) - digamma(r) + digamma(r + k + 1)
+
 @testset "issue #1603" begin
     d = NegativeBinomial(4, 0.2)
     fdm = central_fdm(5, 1)
@@ -23,19 +27,29 @@ mydiffp(r, p, k) = r/p - k/(1 - p)
     @test fdm2(Base.Fix1(cf, d), 0) ≈ -m2
 end
 
-
 @testset "NegativeBinomial r=$r, p=$p, k=$k" for
     p in exp10.(-10:0) .- eps(), # avoid p==1 since it's not differentiable
         r in exp10.(range(-10, stop=2, length=25)),
             k in (0, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024)
 
     @test ForwardDiff.derivative(_p -> logpdf(NegativeBinomial(r, _p), k), p) ≈ mydiffp(r, p, k) rtol=1e-12 atol=1e-12
+    @test ForwardDiff.derivative(_r -> logpdf(NegativeBinomial(_r, p), k), r) ≈ mydiffr(r, p, k) rtol=1e-12 atol=1e-12
 end
 
 @testset "Check the corner case p==1" begin
-    @test logpdf(NegativeBinomial(0.5, 1.0), 0) === 0.0
-    @test logpdf(NegativeBinomial(0.5, 1.0), 1) === -Inf
-    @test all(iszero, rand(NegativeBinomial(rand(), 1.0), 10))
+    for r in randexp(10)
+        d = NegativeBinomial(r, 1.0)
+        @test @inferred(logpdf(d, 0)) === 0.0
+        @test @inferred(logpdf(d, -1)) === -Inf
+        @test @inferred(logpdf(d, 1)) === -Inf
+        @test all(iszero, rand(d, 10))
+    end
+end
+
+@testset "Check the corner case k==0" begin
+	for r in randexp(5), p in rand(5)
+        @test @inferred(logpdf(NegativeBinomial(r, p), 0)) === xlogy(r, p)
+    end
 end
 
 @testset "rrule: logpdf of NegativeBinomial" begin
@@ -59,3 +73,11 @@ end
     test_rrule(logpdf, dist, 0; fdm=fdm)
     test_rrule(logpdf, dist, 0.0 ⊢ ChainRulesTestUtils.NoTangent(); fdm=fdm)
 end
+
+@testset "issue #1582" begin
+    dp = mydiffp(1.0, 1.0, 0.0)
+    @test ForwardDiff.derivative(p -> logpdf(NegativeBinomial(1.0, p), 0.0), 1.0) == dp == 1.0
+
+    dr = mydiffr(1.0, 1.0, 0.0)
+    @test ForwardDiff.derivative(r -> logpdf(NegativeBinomial(r, 1.0), 0.0), 1.0) == dr == 0.0
+end