From 5311dff7bc00b5aad0f2fb0dd655fb8545ac50cf Mon Sep 17 00:00:00 2001
From: Nikos Ignatiadis <nikos.ignatiadis01@gmail.com>
Date: Wed, 31 Aug 2022 01:31:46 +0300
Subject: [PATCH] add beta binomial distribution (#234)

* add beta binomial distribution

* run JuliaFormatter

* Apply suggestions from code review

Co-authored-by: Chad Scherrer <chad.scherrer@gmail.com>

* remove proxy, bump version

Co-authored-by: Chad Scherrer <chad.scherrer@gmail.com>
---
 Project.toml                      |  2 +-
 src/MeasureTheory.jl              |  1 +
 src/parameterized/betabinomial.jl | 37 +++++++++++++++++++++++++++++++
 test/runtests.jl                  | 15 ++++++++-----
 4 files changed, 49 insertions(+), 6 deletions(-)
 create mode 100644 src/parameterized/betabinomial.jl

diff --git a/Project.toml b/Project.toml
index 175faf41..02f6a0e0 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "MeasureTheory"
 uuid = "eadaa1a4-d27c-401d-8699-e962e1bbc33b"
 authors = ["Chad Scherrer <chad.scherrer@gmail.com> and contributors"]
-version = "0.17.2"
+version = "0.17.3"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
diff --git a/src/MeasureTheory.jl b/src/MeasureTheory.jl
index 5be30634..a4534d14 100644
--- a/src/MeasureTheory.jl
+++ b/src/MeasureTheory.jl
@@ -155,6 +155,7 @@ include("parameterized/binomial.jl")
 include("parameterized/multinomial.jl")
 include("parameterized/lkj-cholesky.jl")
 include("parameterized/negativebinomial.jl")
+include("parameterized/betabinomial.jl")
 include("parameterized/gamma.jl")
 include("parameterized/snedecorf.jl")
 include("parameterized/inverse-gaussian.jl")
diff --git a/src/parameterized/betabinomial.jl b/src/parameterized/betabinomial.jl
new file mode 100644
index 00000000..00d48947
--- /dev/null
+++ b/src/parameterized/betabinomial.jl
@@ -0,0 +1,37 @@
+# Beta-Binomial distribution
+
+export BetaBinomial
+import Base
+using SpecialFunctions
+
+@parameterized BetaBinomial(n, α, β)
+
+basemeasure(d::BetaBinomial) = CountingMeasure()
+
+testvalue(::BetaBinomial) = 0
+
+@kwstruct BetaBinomial(n, α, β)
+
+function Base.rand(
+    rng::AbstractRNG,
+    ::Type{T},
+    d::BetaBinomial{(:n, :α, :β)},
+) where {T}
+    k = rand(rng, T, Beta(d.α, d.β))
+    return rand(rng, T, Binomial(d.n, k))
+end
+
+@inline function insupport(d::BetaBinomial, x)
+    isinteger(x) && 0 ≤ x ≤ d.n
+end
+
+@inline function logdensity_def(d::BetaBinomial{(:n, :α, :β)}, y)
+    (n, α, β) = (d.n, d.α, d.β)
+    logbinom = -log1p(n) - logbeta(y + 1, n - y + 1)
+    lognum = logbeta(y + α, n - y + β)
+    logdenom = logbeta(α, β)
+    return logbinom + lognum - logdenom
+end
+
+asparams(::Type{<:BetaBinomial}, ::StaticSymbol{:α}) = asℝ₊
+asparams(::Type{<:BetaBinomial}, ::StaticSymbol{:β}) = asℝ₊
diff --git a/test/runtests.jl b/test/runtests.jl
index aaea8646..ae5cd693 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -54,6 +54,7 @@ test_measures = Any[
     Bernoulli(0.2)
     Beta(2, 3)
     Binomial(10, 0.3)
+    BetaBinomial(10, 2, 3)
     Cauchy()
     Dirichlet(ones(3))
     Exponential()
@@ -277,7 +278,7 @@ end
 
 @testset "Product of Diracs" begin
     x = randn(3)
-    t = as(productmeasure(Dirac.(x))) 
+    t = as(productmeasure(Dirac.(x)))
     @test transform(t, []) == x
 end
 
@@ -297,7 +298,7 @@ end
 
 #     chain = Chain(kernel, μ)
 
-#     dyniterate(iter::TimeLift, ::Nothing) = dyniterate(iter, 0=>nothing) 
+#     dyniterate(iter::TimeLift, ::Nothing) = dyniterate(iter, 0=>nothing)
 #     tr1 = trace(TimeLift(chain), nothing, u -> u[1] > 15)
 #     tr2 = trace(TimeLift(rand(Random.GLOBAL_RNG, chain)), nothing, u -> u[1] > 15)
 #     collect(Iterators.take(chain, 10))
@@ -348,8 +349,8 @@ end
     # NOTE: The `test_broken` below are mostly because of the change to `Affine`.
     # For example, `Normal{(:μ,:σ)}` is now `Affine{(:μ,:σ), Normal{()}}`.
     # The problem is not really with these measures, but with the tests
-    # themselves. 
-    # 
+    # themselves.
+    #
     # We should instead probably be doing e.g.
     # `D = typeof(Normal(μ=0.3, σ=4.1))`
 
@@ -371,6 +372,10 @@ end
         @test repro(Beta, (:α, :β))
     end
 
+    @testset "BetaBinomial" begin
+        @test repro(BetaBinomial, (:n, :α, :β), (n = 10,))
+    end
+
     @testset "Cauchy" begin
         @test_broken repro(Cauchy, (:μ, :σ))
     end
@@ -652,7 +657,7 @@ end
     end
 
     x = rand(d)
-    
+
     @test logdensityof(d, x) isa Real
 end