fix mgf of NegativeBinomial (#1604)

* fix mgf of NegativeBinomial * fix mgf of Geometric * Update src/univariate/discrete/geometric.jl Co-authored-by: David Widmann <[email protected]> * create test/.../geometric.jl * fix geometric cf * fix NegativeBinomial cf * fix * Update src/univariate/discrete/geometric.jl Co-authored-by: David Widmann <[email protected]> * Update src/univariate/discrete/geometric.jl Co-authored-by: David Widmann <[email protected]> * Update test/univariate/discrete/geometric.jl Co-authored-by: David Widmann <[email protected]> * Update test/univariate/discrete/negativebinomial.jl Co-authored-by: David Widmann <[email protected]> * Update test/univariates.jl Co-authored-by: David Widmann <[email protected]> * Update test/univariates.jl Co-authored-by: David Widmann <[email protected]> * improve cf, mgf or Geometric, NegativeBinomial * fix * fix Co-authored-by: David Widmann <[email protected]>
JuliaStats · Aug 16, 2022 · 5582c44 · 5582c44
1 parent aca3920
commit 5582c44
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 15 deletions.
diff --git a/src/univariate/discrete/geometric.jl b/src/univariate/discrete/geometric.jl
@@ -122,17 +122,12 @@ function invlogccdf(d::Geometric{T}, lp::Real) where T<:Real
     max(ceil(lp/log1p(-d.p)) - 1, zero(T))
 end
 
-function mgf(d::Geometric, t::Real)
+function laplace_transform(d::Geometric, t)
     p = succprob(d)
-    p / (expm1(-t) + p)
+    p / (p - (1 - p) * expm1(-t))
 end
-
-function cf(d::Geometric, t::Real)
-    p = succprob(d)
-    # replace with expm1 when complex version available
-    p / (exp(-t*im) - 1 + p)
-end
-
+mgf(d::Geometric, t::Real) = laplace_transform(d, -t)
+cf(d::Geometric, t::Real) = laplace_transform(d, -t*im)
 
 ### Sampling
 

diff --git a/src/univariate/discrete/negativebinomial.jl b/src/univariate/discrete/negativebinomial.jl
@@ -126,15 +126,13 @@ function rand(rng::AbstractRNG, d::NegativeBinomial)
     end
 end
 
-function mgf(d::NegativeBinomial, t::Real)
+function laplace_transform(d::NegativeBinomial, t)
     r, p = params(d)
-    return ((1 - p) * exp(t))^r / (1 - p * exp(t))^r
+    return laplace_transform(Geometric(p, check_args=false), t)^r
 end
 
-function cf(d::NegativeBinomial, t::Real)
-    r, p = params(d)
-    return (((1 - p) * cis(t)) / (1 - p * cis(t)))^r
-end
+mgf(d::NegativeBinomial, t::Real) = laplace_transform(d, -t)
+cf(d::NegativeBinomial, t::Real) = laplace_transform(d, -t*im)
 
 # ChainRules definitions
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -64,6 +64,7 @@ const tests = [
     "univariate/continuous/chernoff",
     "univariate_bounds", # extra file compared to /src
     "univariate/discrete/negativebinomial",
+    "univariate/discrete/geometric",
     "univariate/discrete/bernoulli",
     "univariate/discrete/soliton",
     "univariate/continuous/skewnormal",

diff --git a/test/univariate/discrete/geometric.jl b/test/univariate/discrete/geometric.jl
@@ -0,0 +1,18 @@
+using Distributions
+using Test
+using FiniteDifferences
+
+@testset "Geometric mgf and k vs k-1 parametrization #1604" begin
+    d = Geometric(0.2)
+    @test mgf(d, 0) == 1
+    @test cf(d, 0) == 1
+
+    fdm1 = central_fdm(5, 1)
+    @test fdm1(Base.Fix1(mgf, d), 0) ≈ mean(d)
+    @test fdm1(Base.Fix1(cf, d), 0) ≈ mean(d) * im
+
+    fdm2 = central_fdm(5, 2)
+    m2 = var(d) + mean(d)^2
+    @test fdm2(Base.Fix1(mgf, d), 0) ≈ m2
+    @test fdm2(Base.Fix1(cf, d), 0) ≈ -m2
+end
diff --git a/test/univariate/discrete/negativebinomial.jl b/test/univariate/discrete/negativebinomial.jl
@@ -7,6 +7,20 @@ using FiniteDifferences
 # Eventually, we might want to consolidate the tests here
 
 mydiffp(r, p, k) = r/p - k/(1 - p)
+@testset "issue #1603" begin
+    d = NegativeBinomial(4, 0.2)
+    fdm = central_fdm(5, 1)
+    @test fdm(Base.Fix1(mgf, d), 0) ≈ mean(d)
+    d = NegativeBinomial(1, 0.2)
+    @test fdm(Base.Fix1(mgf, d), 0) ≈ mean(d)
+    @test fdm(Base.Fix1(cf, d), 0) ≈ mean(d) * im
+
+    fdm2 = central_fdm(5, 2)
+    m2 = var(d) + mean(d)^2
+    @test fdm2(Base.Fix1(mgf, d), 0) ≈ m2
+    @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