Constrain precision parameter during optimization

Project.toml

@@ -1,6 +1,6 @@
 name = "BetaRegression"
 uuid = "2339b9c3-daaf-4eaa-90d5-e8471159c344"
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"

src/BetaRegression.jl

@@ -317,13 +317,14 @@ function StatsAPI.score(b::BetaRegressionModel)
     ∂θ = zero(params(b))
     Tr = copy(η)
     @inbounds for i in eachindex(y, η)
-        ηᵢ = η[i]
-        μᵢ = linkinv(link, ηᵢ)
         yᵢ = y[i]
-        a = digamma((1 - μᵢ) * ϕ)
-        r = logit(yᵢ) - digamma(μᵢ * ϕ) + a
-        ∂θ[end] += μᵢ * r + log(1 - yᵢ) - a + ψϕ
-        Tr[i] = ϕ * r * mueta(link, ηᵢ)
+        μᵢ, omμᵢ, dμdη = inverselink(link, η[i])
+        ψp = digamma(ϕ * μᵢ)
+        ψq = digamma(ϕ * omμᵢ)
+        Δ = logit(yᵢ) - ψp + ψq   # logit(yᵢ) - 𝔼(logit(yᵢ))
+        z = log1p(-yᵢ) - ψq + ψϕ  # log(1 - yᵢ) - 𝔼(log(1 - yᵢ))
+        ∂θ[end] += fma(μᵢ, Δ, z)
+        Tr[i] = ϕ * Δ * dμdη
     end
     mul!(view(∂θ, 1:size(X, 2)), X', Tr)
     return ∂θ
@@ -333,14 +334,12 @@ end
 # Q for observed information (pg 10). `p = μ * ϕ` and `q = (1 - μ) * ϕ` are the beta
 # distribution parameters in the typical parameterization, `ψ′_` is `trigamma(_)`.
 function weightdiag(link, p, q, ψ′p, ψ′q, ϕ, yᵢ, ηᵢ, dμdη, expected)
-    w = ϕ * (ψ′p + ψ′q)
+    w = abs(ϕ) * (ψ′p + ψ′q)
     if expected
-        return sqrt(w) * dμdη
+        return sqrt(w) * abs(dμdη)
     else
         w *= dμdη^2
-        ystar = logit(yᵢ)
-        μstar = digamma(p) - digamma(q)
-        w += (ystar - μstar) * dmueta(link, ηᵢ)
+        w += (logit(yᵢ) - digamma(p) + digamma(q)) * dmueta(link, ηᵢ)
         return sqrt(w)
     end
 end
@@ -363,15 +362,14 @@ function 🐟(b::BetaRegressionModel, expected::Bool, inverse::Bool)
     γ = zero(ϕ)
     for i in eachindex(y, η, w)
         ηᵢ = η[i]
-        μᵢ = linkinv(link, ηᵢ)
+        μᵢ, omμᵢ, dμdη = inverselink(link, ηᵢ)
         p = μᵢ * ϕ
-        q = (1 - μᵢ) * ϕ
+        q = omμᵢ * ϕ
         ψ′p = trigamma(p)
         ψ′q = trigamma(q)
-        dμdη = mueta(link, ηᵢ)
         w[i] = weightdiag(link, p, q, ψ′p, ψ′q, ϕ, y[i], ηᵢ, dμdη, expected)
         Tc[i] = (ψ′p * p - ψ′q * q) * dμdη
-        γ += ψ′p * μᵢ^2 + ψ′q * (1 - μᵢ)^2 - ψ′ϕ
+        γ += ψ′p * μᵢ^2 + ψ′q * omμᵢ^2 - ψ′ϕ
     end
     Xᵀ = copy(adjoint(X))
     XᵀTc = Xᵀ * Tc
@@ -432,14 +430,16 @@ approximately zero. This is determined by `isapprox` using the specified `atol`
 """
 function StatsAPI.fit!(b::BetaRegressionModel; maxiter=100, atol=1e-8, rtol=1e-8)
     initialize!(b)
-    z = zero(params(b))
+    θ = params(b)
+    z = zero(θ)
     for iter in 1:maxiter
         U = score(b)
         checkfinite(U, iter)
         isapprox(U, z; atol, rtol) && return b  # converged!
         K = 🐟(b, true, true)
         checkfinite(K, iter)
-        mul!(params(b), K, U, true, true)
+        mul!(θ, K, U, true, true)
+        θ[end] = max(θ[end], eps(eltype(θ)))  # impose positivity constraint on ϕ
         linearpredictor!(b)
     end
     throw(ConvergenceException(maxiter))

test/runtests.jl

@@ -1,4 +1,5 @@
 using BetaRegression
+using Distributions
 using GLM
 using StatsBase
 using Test
@@ -212,13 +213,32 @@ end
     end
 end
 
-@testset "Resetting an invalid initial precision" begin
+@testset "Parameter constraints" begin
+    # Generating distribution and `rand(_, 10)` for several problematic cases. Without
+    # constraints placed on the parameters during optimization, the models built using
+    # these as their respective responses fail to converge. See issue #6.
+    d = [Beta(0.5, 0.5) => [0.9020980693394288, 0.055577211500829754, 0.23132559790498958,
+                            0.5813942170987118, 0.9709116084487788, 0.7754094004739907,
+                            0.05982031817793439, 0.8670342033149658, 0.683216406088941,
+                            0.141451701046685],
+         Beta(0.2, 0.2) => [0.023848292440454045, 0.9355547503109088, 0.9924242111793663,
+                            0.008946868197901494, 0.04010019793873337, 0.16627955105469605,
+                            0.999828498409377, 0.9999768604670911, 0.9930500996079485,
+                            0.9642281316220574],
+         Beta(0.5, 10) => [0.034541975300900515, 0.01939889173586482, 0.24100845407491417,
+                           0.0011108618208461997, 0.00937646060618697, 0.014267329875521517,
+                           0.04085538895184706, 0.016118136919340345, 0.008924027953777908,
+                           0.0760514673345253],
+         Beta(10, 0.3) => [0.9946023297631961, 0.9680165382504563, 0.999142668249286,
+                           0.9917649725155366, 0.9843468826146887, 0.9999997547187489,
+                           0.9943098787910513, 0.9991354297579114, 0.9521187943069395,
+                           0.9805165186163991]]
     X = ones(10, 1)
-    # Generated via `Beta(0.5, 0.5)`
-    y = [0.9020980693394288, 0.055577211500829754, 0.23132559790498958,
-         0.5813942170987118, 0.9709116084487788, 0.7754094004739907,
-         0.05982031817793439, 0.8670342033149658, 0.683216406088941,
-         0.141451701046685]
-    model = fit(BetaRegressionModel, X, y)
-    @test linkinv(Link(model), only(coef(model))) ≈ mean(y) rtol=0.05
+    @testset "Generated from $dist" for (dist, y) in d
+        model = fit(BetaRegressionModel, X, y)
+        @test linkinv(Link(model), only(coef(model))) ≈ mean(y) rtol=0.05
+        μ̂ = mean(fitted(model))
+        ϕ̂ = precision(model)
+        @test μ̂ * (1 - μ̂) / (1 + ϕ̂) ≈ var(y) rtol=0.5
+    end
 end