Implement gradlogpartition for Exponential Family Distributions

docs/src/interface.md

@@ -38,11 +38,13 @@ isproper
 getbasemeasure
 getsufficientstatistics
 getlogpartition
+getgradlogpartition
 getfisherinformation
 getsupport
 basemeasure
 sufficientstatistics
 logpartition
+gradlogpartition
 fisherinformation
 isbasemeasureconstant
 ConstantBaseMeasure

src/common.jl

@@ -44,3 +44,5 @@ function binomial_prod(n, p, x)
         end
     end
 end
+
+mvdigamma(η,p) = sum( digamma(η + (one(d) - d)/2) for d=1:p)

src/distributions/bernoulli.jl

@@ -94,6 +94,11 @@ getlogpartition(::NaturalParametersSpace, ::Type{Bernoulli}) = (η) -> begin
     return -log(logistic(-η₁))
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Bernoulli}) = (η) -> begin
+    (η₁,) = unpack_parameters(Bernoulli, η)
+    return SA[logistic(η₁)]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{Bernoulli}) = (η) -> begin
     (η₁,) = unpack_parameters(Bernoulli, η)
     f = logistic(-η₁)

src/distributions/beta.jl

@@ -62,6 +62,16 @@ getlogpartition(::NaturalParametersSpace, ::Type{Beta}) = (η) -> begin
     return logbeta(η₁ + one(η₁), η₂ + one(η₂))
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Beta}) = (η) -> begin
+    (η₁, η₂) = unpack_parameters(Beta, η)
+    η₁p = η₁ + one(η₁)
+    η₂p = η₂ + one(η₂)
+    ηsum = η₁p + η₂p
+    dig = digamma(ηsum)
+
+    return SA[digamma(η₁p) - dig, digamma(η₂p) - dig]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{Beta}) = (η) -> begin
     (η₁, η₂) = unpack_parameters(Beta, η)
     psia = trigamma(η₁ + one(η₁))

src/distributions/binomial.jl

@@ -92,6 +92,11 @@ getlogpartition(::NaturalParametersSpace, ::Type{Binomial}, ntrials) = (η) -> b
     return ntrials * log1pexp(η₁)
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Binomial}, ntrials) = (η) -> begin
+    (η₁,) = unpack_parameters(Binomial, η)
+    return SA[ntrials*exp(η₁) / (one(η₁) + exp(η₁))]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{Binomial}, ntrials) = (η) -> begin
     (η₁,) = unpack_parameters(Binomial, η)
     aux = logistic(η₁)

src/distributions/categorical.jl

@@ -79,6 +79,19 @@ getlogpartition(::NaturalParametersSpace, ::Type{Categorical}, conditioner) =
         return logsumexp(η)
     end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Categorical}, conditioner) = 
+    (η) -> begin
+        if (conditioner !== length(η))
+            throw(
+                DimensionMismatch(
+                    "Cannot evaluate the logparition of the `Categorical` with `conditioner = $(conditioner)` and vector of natural parameters `η = $(η)`"
+                )
+            )
+        end
+        sumη = vmapreduce(exp, +, η)
+        return vmap(d->exp(d)/sumη ,η)
+    end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{Categorical}, conditioner) =
     (η) -> begin
         if (conditioner !== length(η))

src/distributions/chi_squared.jl

@@ -64,6 +64,11 @@ getlogpartition(::NaturalParametersSpace, ::Type{Chisq}) = (η) -> begin
     return loggamma(η1 + o) + (η1 + o) * logtwo
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Chisq}) = (η) -> begin
+    (η1,) = unpack_parameters(Chisq, η)
+    return SA[digamma(η1 + one(η1)) + logtwo]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{Chisq}) = (η) -> begin
     (η1,) = unpack_parameters(Chisq, η)
     SA[trigamma(η1 + one(η1));;]

src/distributions/dirichlet.jl

@@ -69,6 +69,11 @@ getfisherinformation(::NaturalParametersSpace, ::Type{Dirichlet}) =
         return Diagonal(map(d -> trigamma(d + 1), η1)) - Ones{Float64}(n, n) * trigamma(sum(η1) + n)
     end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Dirichlet}) = (η) -> begin
+    (η1,) = unpack_parameters(Dirichlet, η)
+    return digamma.(η1 .+ 1) .- digamma(sum(η1 .+ 1))
+end
+
 # Mean parametrization
 
 getlogpartition(::MeanParametersSpace, ::Type{Dirichlet}) = (θ) -> begin
@@ -83,3 +88,8 @@ getfisherinformation(::MeanParametersSpace, ::Type{Dirichlet}) = (θ) -> begin
     n = length(α)
     return Diagonal(map(d -> trigamma(d), α)) - Ones{Float64}(n, n) * trigamma(sum(α))
 end
+
+getgradlogpartition(::MeanParametersSpace, ::Type{Dirichlet}) = (θ) -> begin
+    (α,) = unpack_parameters(Dirichlet, θ)
+    return digamma.(α) .- digamma(sum(α))
+end

src/distributions/exponential.jl

@@ -44,6 +44,11 @@ getlogpartition(::NaturalParametersSpace, ::Type{Exponential}) = (η) -> begin
     return -log(-η₁)
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Exponential}) = (η) -> begin
+    (η₁,) = unpack_parameters(Exponential, η)
+    return SA[-1/η₁]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{Exponential}) = (η) -> begin
     (η₁,) = unpack_parameters(Exponential, η)
     SA[inv(η₁^2);;]

src/distributions/gamma_family/gamma_family.jl

@@ -100,6 +100,11 @@ getfisherinformation(::NaturalParametersSpace, ::Type{Gamma}) = (η) -> begin
     SA[trigamma(η₁ + one(η₁)) -inv(η₂); -inv(η₂) (η₁+one(η₁))/(η₂^2)]
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Gamma}) = (η) -> begin
+    (η₁, η₂) = unpack_parameters(Gamma, η)
+    return SA[digamma(η₁ + one(η₁)) - log(-η₂), - (η₁ + one(η₁)) / η₂]
+end
+
 # Mean parametrization
 
 getlogpartition(::MeanParametersSpace, ::Type{Gamma}) = (θ) -> begin
@@ -114,3 +119,8 @@ getfisherinformation(::MeanParametersSpace, ::Type{Gamma}) = (θ) -> begin
         inv(scale) shape/abs2(scale)
     ]
 end
+
+getgradlogpartition(::MeanParametersSpace, ::Type{Gamma}) = (θ) -> begin
+    (shape, scale) = unpack_parameters(Gamma, θ)
+    return SA[digamma(shape) - log(scale), - shape / scale]
+end

src/distributions/gamma_inverse.jl

@@ -63,6 +63,13 @@ getlogpartition(::NaturalParametersSpace, ::Type{GammaInverse}) = (η) -> begin
     return loggamma(-η₁ - one(η₁)) - (-η₁ - one(η₁)) * log(-η₂)
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{GammaInverse}) = (η) -> begin
+    (η₁, η₂) = unpack_parameters(GammaInverse, η)
+    dη1 = -digamma(-η₁ - one(η₁)) + log(-η₂)
+    dη2 = - (-η₁ - one(η₁))/η₂
+    return SA[dη1, dη2]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{GammaInverse}) =
     (η) -> begin
         (η₁, η₂) = unpack_parameters(GammaInverse, η)

src/distributions/geometric.jl

@@ -49,6 +49,11 @@ getfisherinformation(::NaturalParametersSpace, ::Type{Geometric}) = (η) -> begi
     return SA[exp(η1) / (one(η1) - exp(η1))^2;;]
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Geometric}) = (η) -> begin
+    (η1,) = unpack_parameters(Geometric, η)
+    return SA[exp(η1) / (one(η1) - exp(η1));]
+end
+
 # Mean parametrization
 
 getlogpartition(::MeanParametersSpace, ::Type{Geometric}) = (θ) -> begin
@@ -60,3 +65,8 @@ getfisherinformation(::MeanParametersSpace, ::Type{Geometric}) = (θ) -> begin
     (p,) = unpack_parameters(Geometric, θ)
     return SA[one(p) / (p^2 * (one(p) - p));;]
 end
+
+getgradlogpartition(::MeanParametersSpace, ::Type{Geometric}) = (θ) -> begin
+    (p,) = unpack_parameters(Geometric, θ)
+    return SA[one(p) / (p^2 - p);]
+end

src/distributions/laplace.jl

@@ -168,6 +168,11 @@ getlogpartition(::NaturalParametersSpace, ::Type{Laplace}, _) = (η) -> begin
     return log(-2 / η₁)
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Laplace}, _) = (η) -> begin
+    (η₁,) = unpack_parameters(Laplace, η)
+    return SA[-inv(η₁);]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{Laplace}, _) = (η) -> begin
     (η₁,) = unpack_parameters(Laplace, η)
     return SA[inv(η₁^2);;]

src/distributions/lognormal.jl

@@ -53,6 +53,13 @@ getlogpartition(::NaturalParametersSpace, ::Type{LogNormal}) = (η) -> begin
     return -(η₁ + 1)^2 / (4η₂) - log(-2η₂) / 2
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{LogNormal}) = (η) -> begin
+    (η₁, η₂) = unpack_parameters(LogNormal, η) 
+    dη1 =  -(η₁ + 1)/(2η₂)
+    dη2 = (η₁ + 1)^2/(4η₂^2) - inv(η₂)/2
+    return SA[dη1, dη2]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{LogNormal}) =
     (η) -> begin
         (η₁, η₂) = unpack_parameters(LogNormal, η)
@@ -66,6 +73,13 @@ getlogpartition(::MeanParametersSpace, ::Type{LogNormal}) = (θ) -> begin
     return abs2(μ) / (2abs2(σ)) + log(σ)
 end
 
+getgradlogpartition(::MeanParametersSpace, ::Type{LogNormal}) = (θ) -> begin    
+    (μ, σ) = unpack_parameters(LogNormal, θ)
+    dμ = abs(μ) / (abs2(σ))
+    dσ = -abs2(μ) / (σ^3) + 1/σ
+    return SA[dμ, dσ]
+end
+
 getfisherinformation(::MeanParametersSpace, ::Type{LogNormal}) = (θ) -> begin
     (μ, σ) = unpack_parameters(LogNormal, θ)
     invσ² = inv(abs2(σ))

src/distributions/matrix_dirichlet.jl

@@ -153,6 +153,14 @@ getlogpartition(::NaturalParametersSpace, ::Type{MatrixDirichlet}) =
         )
     end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{MatrixDirichlet}) = 
+    (η) -> begin
+        (η1,) = unpack_parameters(MatrixDirichlet, η)
+        return vmapreduce(
+            d -> getgradlogpartition(NaturalParametersSpace(), Dirichlet)(convert(Vector, d)),vcat,
+            eachcol(η1))
+    end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{MatrixDirichlet}) =
     (η) -> begin
         (η1,) = unpack_parameters(MatrixDirichlet, η)
@@ -177,6 +185,14 @@ getlogpartition(::MeanParametersSpace, ::Type{MatrixDirichlet}) =
         )
     end
 
+getgradlogpartition(::MeanParametersSpace, ::Type{MatrixDirichlet}) = 
+    (θ) -> begin
+        (α,) = unpack_parameters(MatrixDirichlet, θ)
+        return vmapreduce(
+            d -> getgradlogpartition(NaturalParametersSpace(), Dirichlet)(convert(Vector, d)),vcat,
+            eachcol(α))
+    end
+
 getfisherinformation(::MeanParametersSpace, ::Type{MatrixDirichlet}) =
     (θ) -> begin
         (α,) = unpack_parameters(MatrixDirichlet, θ)

src/distributions/negative_binomial.jl

@@ -106,6 +106,11 @@ getlogpartition(::NaturalParametersSpace, ::Type{NegativeBinomial}, conditioner)
     return -conditioner * log(one(η1) - exp(η1))
 end
 
+getgradlogpartition(::NaturalParametersSpace,::Type{NegativeBinomial}, conditioner) = (η) -> begin
+    (η1,) = unpack_parameters(NegativeBinomial, η)
+    return SA[-conditioner*(-exp(η1)/(one(η1)-exp(η1)));]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{NegativeBinomial}, r) = (η) -> begin
     (η1,) = unpack_parameters(NegativeBinomial, η)
     return SA[r * exp(η1) / (one(η1) - exp(η1))^2;;]
@@ -118,6 +123,11 @@ getlogpartition(::MeanParametersSpace, ::Type{NegativeBinomial}, conditioner) =
     return -conditioner * log(one(p) - p)
 end
 
+getgradlogpartition(::MeanParametersSpace,::Type{NegativeBinomial}, conditioner) = (θ) -> begin
+    (p,) = unpack_parameters(NegativeBinomial, η)
+    return SA[conditioner*inv(one(p) - p);]
+end
+
 getfisherinformation(::MeanParametersSpace, ::Type{NegativeBinomial}, r) = (θ) -> begin
     (p,) = unpack_parameters(NegativeBinomial, θ)
     return SA[r / (p^2 * (one(p) - p));;]

src/distributions/normal_family/normal_family.jl

@@ -575,6 +575,12 @@ getlogpartition(::NaturalParametersSpace, ::Type{NormalMeanVariance}) = (η) ->
     return -abs2(η₁) / 4η₂ - log(-2η₂) / 2
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{NormalMeanVariance}) =
+    (η) -> begin
+        (η₁, η₂) = unpack_parameters(NormalMeanVariance, η)
+        return SA[-η₁ * inv(η₂*2), abs2(η₁) / ( 4 * abs2(η₂)) - 1 / (2 * η₂)]
+    end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{NormalMeanVariance}) =
     (η) -> begin
         (η₁, η₂) = unpack_parameters(NormalMeanVariance, η)
@@ -591,6 +597,12 @@ getlogpartition(::MeanParametersSpace, ::Type{NormalMeanVariance}) = (θ) -> beg
     return μ / 2σ² + log(sqrt(σ²))
 end
 
+getgradlogpartition(::MeanParametersSpace, ::Type{NormalMeanVariance}) =
+    (θ) -> begin
+        (μ, σ²) = unpack_parameters(NormalMeanVariance, θ)
+        return SA[μ / σ², - abs2(μ) / (2σ²^2) + 1 / σ²]
+    end
+
 getfisherinformation(::MeanParametersSpace, ::Type{NormalMeanVariance}) = (θ) -> begin
     (_, σ²) = unpack_parameters(NormalMeanVariance, θ)
     return SA[inv(σ²) 0; 0 inv(2 * (σ²^2))]
@@ -678,6 +690,13 @@ getlogpartition(::NaturalParametersSpace, ::Type{MvNormalMeanCovariance}) = (η)
     return (dot(η₁, Cinv, η₁) / 2 - (k * log(2) + l)) / 2
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{MvNormalMeanCovariance}) =
+    (η) -> begin
+        (η₁, η₂) = unpack_parameters(MvNormalMeanCovariance, η)
+        Cinv, _ = cholinv_logdet(-η₂)
+        return pack_parameters(MvNormalMeanCovariance, (0.5 * Cinv * η₁, 0.25 * Cinv * η₁ * η₁' * Cinv + 0.5 * Cinv))
+    end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{MvNormalMeanCovariance}) =
     (η) -> begin
         (η₁, η₂) = unpack_parameters(MvNormalMeanCovariance, η)

src/distributions/normal_gamma.jl

@@ -143,6 +143,18 @@ getlogpartition(::NaturalParametersSpace, ::Type{NormalGamma}) = (η) -> begin
     return loggamma(η3half) - log(-2η2) * (1 / 2) - (η3half) * log(-η4 + η1^2 / (4η2))
 end
 
+getgradlogpartition(::NaturalParametersSpace,::Type{NormalGamma}) = (η) -> begin  
+    (η1, η2, η3, η4) = unpack_parameters(NormalGamma, η)
+    η3half = η3 + (1 / 2)
+    c   = (-η4 + η1^2/(4η2))
+    dη1 = -η3half*((η1/(2η2)) / c) 
+    dη2 = -inv(η2)/2 - η3half*(-η1^2/(4η2^2) / c)
+    dη3 = digamma(η3half) - log(c)
+    dη4 = η3half /c
+
+    return SA[dη1, dη2, dη3, dη4]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{NormalGamma}) =
     (η) -> begin
         (η1, η2, η3, η4) = unpack_parameters(NormalGamma, η)

src/distributions/pareto.jl

@@ -148,6 +148,11 @@ getlogpartition(::NaturalParametersSpace, ::Type{Pareto}, conditioner) = (η) ->
     return log(conditioner^(one(η1) + η1) / (-one(η1) - η1))
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Pareto}, conditioner) = (η) -> begin
+    (η1,) = unpack_parameters(Pareto, η)
+    return SA[log(conditioner) - one(η1)/(one(η1)+η1);]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{Pareto}, _) = (η) -> begin
     (η1,) = unpack_parameters(Pareto, η)
     return SA[1 / (-1 - η1)^2;;]
@@ -160,6 +165,11 @@ getlogpartition(::MeanParametersSpace, ::Type{Pareto}, conditioner) = (θ) -> be
     return -log(shape) - shape * log(conditioner)
 end
 
+getgradlogpartition(::MeanParametersSpace, ::Type{Pareto}, conditioner) = (θ) -> begin
+    (shape,) = unpack_parameters(Pareto, θ)
+    return SA[-inv(shape) - log(conditioner);]
+end
+
 getfisherinformation(::MeanParametersSpace, ::Type{Pareto}, conditioner) = (θ) -> begin
     (α,) = unpack_parameters(Pareto, θ)
     ### Below fisher information is problematic if α is larger than conditioner as Pareto 

src/distributions/poisson.jl

@@ -60,6 +60,11 @@ getlogpartition(::NaturalParametersSpace, ::Type{Poisson}) = (η) -> begin
     return exp(η1)
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Poisson}) = (η) -> begin
+    (η1,) = unpack_parameters(Poisson, η)
+    return SA[exp(η1)]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{Poisson}) = (η) -> begin
     (η1,) = unpack_parameters(Poisson, η)
     SA[exp(η1);;]

src/distributions/rayleigh.jl

@@ -56,6 +56,11 @@ getlogpartition(::NaturalParametersSpace, ::Type{Rayleigh}) = (η) -> begin
     return -log(-2 * η1)
 end
 
+getgradlogpartition(::NaturalParametersSpace, ::Type{Rayleigh}) = (η) -> begin
+    (η1, ) = unpack_parameters(Rayleigh, η)
+    return SA[-inv(η1);]
+end
+
 getfisherinformation(::NaturalParametersSpace, ::Type{Rayleigh}) = (η) -> begin
     (η1,) = unpack_parameters(Rayleigh, η)
     SA[inv(η1^2);;]
@@ -68,6 +73,11 @@ getlogpartition(::MeanParametersSpace, ::Type{Rayleigh}) = (θ) -> begin
     return 2 * log(σ)
 end
 
+getgradlogpartition(::MeanParametersSpace, ::Type{Rayleigh}) = (θ) -> begin
+    (σ,) = unpack_parameters(Rayleigh, θ)
+    return SA[2/σ;]
+end
+
 getfisherinformation(::MeanParametersSpace, ::Type{Rayleigh}) = (θ) -> begin
     (σ,) = unpack_parameters(Rayleigh, θ)
     return SA[4 / σ^2;;]