From cc47342f149d353c63c9ddf579cf82195e4cd1c5 Mon Sep 17 00:00:00 2001
From: jg-854 <33372922+jg-854@users.noreply.github.com>
Date: Sun, 24 Jan 2021 13:43:03 +0000
Subject: [PATCH] Added Maximum Likelihood Estimation for Beta

Included unit tests to check functionality works.
Updated documentation too.
---
 docs/src/fit.md                   |  1 +
 src/univariate/continuous/beta.jl | 30 +++++++++++++++++++++++++++++-
 test/fit.jl                       |  6 ++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/docs/src/fit.md b/docs/src/fit.md
index b371d0daa..24e8ee1dc 100644
--- a/docs/src/fit.md
+++ b/docs/src/fit.md
@@ -38,6 +38,7 @@ The `fit_mle` method has been implemented for the following distributions:
 **Univariate:**
 
 - [`Bernoulli`](@ref)
+- [`Beta`](@ref)
 - [`Binomial`](@ref)
 - [`Categorical`](@ref)
 - [`DiscreteUniform`](@ref)
diff --git a/src/univariate/continuous/beta.jl b/src/univariate/continuous/beta.jl
index d1f0451ae..615aa767c 100644
--- a/src/univariate/continuous/beta.jl
+++ b/src/univariate/continuous/beta.jl
@@ -199,8 +199,36 @@ function rand(rng::AbstractRNG, d::Beta{T}) where T
 end
 
 #### Fit model
+"""
+    fit_mle(::Type{<:Beta}, x::AbstractArray{T})
 
-# TODO: add MLE method (should be similar to Dirichlet)
+Maximum Likelihood Estimate of `Beta` Distribution via Newton's Method
+"""
+function fit_mle(::Type{<:Beta}, x::AbstractArray{T};
+    maxiter::Int=1000, tol::Float64=1e-14) where T<:Real
+
+    α₀,β₀ = params(fit(Beta,x)) #initial guess of parameters
+    g₁ = mean(log.(x))
+    g₂ = mean(log.(one(T) .- x))
+    θ= [α₀ ; β₀ ]
+
+    converged = false
+    t=0
+    while !converged && t < maxiter #newton method
+        t+=1
+        temp1 = digamma(θ[1]+θ[2])
+        temp2 = trigamma(θ[1]+θ[2])
+        grad = [g₁+temp1-digamma(θ[1])
+               temp1+g₂-digamma(θ[2])]
+        hess = [temp2-trigamma(θ[1]) temp2
+                temp2 temp2-trigamma(θ[2])]
+        Δθ = hess\grad #newton step
+        θ .-= Δθ
+        converged = dot(Δθ,Δθ) < 2*tol #stopping criterion
+    end
+
+    return Beta(θ[1], θ[2])
+end
 
 """
     fit(::Type{<:Beta}, x::AbstractArray{T})
diff --git a/test/fit.jl b/test/fit.jl
index 897ba6b24..54be85887 100644
--- a/test/fit.jl
+++ b/test/fit.jl
@@ -69,6 +69,12 @@ end
         @test isa(d, dist)
         @test isapprox(d.α, 1.3, atol=0.1)
         @test isapprox(d.β, 3.7, atol=0.1)
+
+        d = fit_mle(dist, func[2](dist(1.3, 3.7), N))
+        @test isa(d, dist)
+        @test isapprox(d.α, 1.3, atol=0.1)
+        @test isapprox(d.β, 3.7, atol=0.1)
+
     end
 end