move expected_loglik here from ApproximateGPs.jl (#70)

Moves `expected_loglik` from ApproximateGPs.jl to GPLikelihoods.jlApproximateGPs.

Changes from the original code in ApproximateGPs.jl:
- rename to `expected_loglikelihood` to be more explicit
- changed signature: argument order is now `quadrature`, `lik`, `q_f`, `y`
- rename `_default_quadrature` to `default_expectation_method`
- rename expectation structs to `DefaultExpectationMethod`, `AnalyticExpectation`, `GaussHermiteExpectation`, `MonteCarloExpectation`; they are not exported from here
- drop abstract base type (was not used for anything)
- drop default argument for GaussHermite / MonteCarlo quadrature points (moved into `DefaultExpectationMethod`)
- GaussHermiteExpectation now stores nodes&weights within struct, resolves #71
- move analytic definitions into files for each likelihood implementation

Note: bumps julia compat to 1.6.

Co-authored-by: willtebbutt <wct23@cam.ac.uk>

.github/workflows/CI.yml

@@ -22,7 +22,7 @@ jobs:
  matrix:
  version:
  - '1'
- - '1.3'
+ - '1.6'
  - 'nightly'
  os:
  - ubuntu-latest

Project.toml

@@ -1,20 +1,28 @@
 name = "GPLikelihoods"
 uuid = "6031954c-0455-49d7-b3b9-3e1c99afaf40"
 authors = ["JuliaGaussianProcesses Team"]
-version = "0.3.1"
+version = "0.4.0"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
+IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
+ChainRulesCore = "1.7"
 Distributions = "0.19, 0.20, 0.21, 0.22, 0.23, 0.24, 0.25"
+FastGaussQuadrature = "0.4"
 Functors = "0.1, 0.2"
 InverseFunctions = "0.1.2"
+IrrationalConstants = "0.1"
+SpecialFunctions = "1, 2"
 StatsFuns = "0.9.13"
-julia = "1.3"
+julia = "1.6"

src/GPLikelihoods.jl

@@ -34,6 +34,8 @@ include("links.jl")
 
 # Likelihoods
 abstract type AbstractLikelihood end
+
+include("expectations.jl")
 include("likelihoods/bernoulli.jl")
 include("likelihoods/categorical.jl")
 include("likelihoods/gaussian.jl")

src/expectations.jl

@@ -0,0 +1,117 @@
+using FastGaussQuadrature: gausshermite
+using SpecialFunctions: loggamma
+using ChainRulesCore: ChainRulesCore
+using IrrationalConstants: sqrt2, invsqrtπ
+
+struct DefaultExpectationMethod end
+
+struct AnalyticExpectation end
+
+struct GaussHermiteExpectation
+ xs::Vector{Float64}
+ ws::Vector{Float64}
+end
+GaussHermiteExpectation(n::Integer) = GaussHermiteExpectation(gausshermite(n)...)
+
+ChainRulesCore.@non_differentiable gausshermite(n)
+
+struct MonteCarloExpectation
+ n_samples::Int
+end
+
+default_expectation_method(_) = GaussHermiteExpectation(20)
+
+"""
+ expected_loglikelihood(
+ quadrature,
+ lik,
+ q_f::AbstractVector{<:Normal},
+ y::AbstractVector,
+ )
+
+This function computes the expected log likelihood:
+
+```math
+ ∫ q(f) log p(y | f) df
+```
+where `p(y | f)` is the process likelihood. This is described by `lik`, which should be a
+callable that takes `f` as input and returns a Distribution over `y` that supports
+`loglikelihood(lik(f), y)`.
+
+`q(f)` is an approximation to the latent function values `f` given by:
+```math
+ q(f) = ∫ p(f | u) q(u) du
+```
+where `q(u)` is the variational distribution over inducing points (see
+[`elbo`](@ref)). The marginal distributions of `q(f)` are given by `q_f`.
+
+`quadrature` determines which method is used to calculate the expected log
+likelihood - see [`elbo`](@ref) for more details.
+
+# Extended help
+
+`q(f)` is assumed to be an `MvNormal` distribution and `p(y | f)` is assumed to
+have independent marginals such that only the marginals of `q(f)` are required.
+"""
+expected_loglikelihood(quadrature, lik, q_f, y)
+
+"""
+ expected_loglikelihood(::DefaultExpectationMethod, lik, q_f::AbstractVector{<:Normal}, y::AbstractVector)
+
+The expected log likelihood, using the default quadrature method for the given likelihood.
+(The default quadrature method is defined by `default_expectation_method(lik)`, and should
+be the closed form solution if it exists, but otherwise defaults to Gauss-Hermite
+quadrature.)
+"""
+function expected_loglikelihood(
+ ::DefaultExpectationMethod, lik, q_f::AbstractVector{<:Normal}, y::AbstractVector
+)
+ quadrature = default_expectation_method(lik)
+ return expected_loglikelihood(quadrature, lik, q_f, y)
+end
+
+function expected_loglikelihood(
+ mc::MonteCarloExpectation, lik, q_f::AbstractVector{<:Normal}, y::AbstractVector
+)
+ # take `n_samples` reparameterised samples
+ f_μ = mean.(q_f)
+ fs = f_μ .+ std.(q_f) .* randn(eltype(f_μ), length(q_f), mc.n_samples)
+ lls = loglikelihood.(lik.(fs), y)
+ return sum(lls) / mc.n_samples
+end
+
+# Compute the expected_loglikelihood over a collection of observations and marginal distributions
+function expected_loglikelihood(
+ gh::GaussHermiteExpectation, lik, q_f::AbstractVector{<:Normal}, y::AbstractVector
+)
+ # Compute the expectation via Gauss-Hermite quadrature
+ # using a reparameterisation by change of variable
+ # (see e.g. en.wikipedia.org/wiki/Gauss%E2%80%93Hermite_quadrature)
+ return sum(Broadcast.instantiate(
+ Broadcast.broadcasted(y, q_f) do yᵢ, q_fᵢ # Loop over every pair
+ # of marginal distribution q(fᵢ) and observation yᵢ
+ expected_loglikelihood(gh, lik, q_fᵢ, yᵢ)
+ end,
+ ))
+end
+
+# Compute the expected_loglikelihood for one observation and a marginal distributions
+function expected_loglikelihood(gh::GaussHermiteExpectation, lik, q_f::Normal, y)
+ μ = mean(q_f)
+ σ̃ = sqrt2 * std(q_f)
+ return invsqrtπ * sum(Broadcast.instantiate(
+ Broadcast.broadcasted(gh.xs, gh.ws) do x, w # Loop over every
+ # pair of Gauss-Hermite point x with weight w
+ f = σ̃ * x + μ
+ loglikelihood(lik(f), y) * w
+ end,
+ ))
+end
+
+function expected_loglikelihood(
+ ::AnalyticExpectation, lik, q_f::AbstractVector{<:Normal}, y::AbstractVector
+)
+ return error(
+ "No analytic solution exists for $(typeof(lik)). Use `DefaultExpectationMethod`, `GaussHermiteExpectation` or `MonteCarloExpectation` instead.",
+ )
+end

src/likelihoods/exponential.jl

@@ -16,3 +16,15 @@ ExponentialLikelihood(l=exp) = ExponentialLikelihood(link(l))
 (l::ExponentialLikelihood)(f::Real) = Exponential(l.invlink(f))
 
 (l::ExponentialLikelihood)(fs::AbstractVector{<:Real}) = Product(map(l, fs))
+
+function expected_loglikelihood(
+ ::AnalyticExpectation,
+ ::ExponentialLikelihood{ExpLink},
+ q_f::AbstractVector{<:Normal},
+ y::AbstractVector{<:Real},
+)
+ f_μ = mean.(q_f)
+ return sum(-f_μ - y .* exp.((var.(q_f) / 2) .- f_μ))
+end
+
+default_expectation_method(::ExponentialLikelihood{ExpLink}) = AnalyticExpectation()

src/likelihoods/gamma.jl

@@ -21,3 +21,18 @@ GammaLikelihood(α::Real=1.0, l=exp) = GammaLikelihood(α, link(l))
 (l::GammaLikelihood)(f::Real) = Gamma(l.α, l.invlink(f))
 
 (l::GammaLikelihood)(fs::AbstractVector{<:Real}) = Product(map(l, fs))
+
+function expected_loglikelihood(
+ ::AnalyticExpectation,
+ lik::GammaLikelihood{ExpLink},
+ q_f::AbstractVector{<:Normal},
+ y::AbstractVector{<:Real},
+)
+ f_μ = mean.(q_f)
+ return sum(
+ (lik.α - 1) * log.(y) .- y .* exp.((var.(q_f) / 2) .- f_μ) .- lik.α * f_μ .-
+ loggamma(lik.α),
+ )
+end
+
+default_expectation_method(::GammaLikelihood{ExpLink}) = AnalyticExpectation()

src/likelihoods/gaussian.jl

@@ -23,6 +23,19 @@ GaussianLikelihood(σ²::Real) = GaussianLikelihood([σ²])
 
 (l::GaussianLikelihood)(fs::AbstractVector{<:Real}) = MvNormal(fs, first(l.σ²) * I)
 
+function expected_loglikelihood(
+ ::AnalyticExpectation,
+ lik::GaussianLikelihood,
+ q_f::AbstractVector{<:Normal},
+ y::AbstractVector{<:Real},
+)
+ return sum(
+ -0.5 * (log(2π) .+ log.(lik.σ²) .+ ((y .- mean.(q_f)) .^ 2 .+ var.(q_f)) / lik.σ²)
+ )
+end
+
+default_expectation_method(::GaussianLikelihood) = AnalyticExpectation()
+
 """
  HeteroscedasticGaussianLikelihood(l=exp)
 

src/likelihoods/poisson.jl

@@ -19,3 +19,15 @@ PoissonLikelihood(l=exp) = PoissonLikelihood(link(l))
 (l::PoissonLikelihood)(f::Real) = Poisson(l.invlink(f))
 
 (l::PoissonLikelihood)(fs::AbstractVector{<:Real}) = Product(map(l, fs))
+
+function expected_loglikelihood(
+ ::AnalyticExpectation,
+ ::PoissonLikelihood{ExpLink},
+ q_f::AbstractVector{<:Normal},
+ y::AbstractVector{<:Real},
+)
+ f_μ = mean.(q_f)
+ return sum((y .* f_μ) - exp.(f_μ .+ (var.(q_f) / 2)) - loggamma.(y .+ 1))
+end
+
+default_expectation_method(::PoissonLikelihood{ExpLink}) = AnalyticExpectation()

test/Project.toml

@@ -4,6 +4,7 @@ Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 Distributions = "0.19, 0.20, 0.21, 0.22, 0.23, 0.24, 0.25"

test/expectations.jl

@@ -0,0 +1,100 @@
+@testset "expectations" begin
+ # Test that the various methods of computing expectations return the same
+ # result.
+ rng = MersenneTwister(123456)
+ q_f = Normal.(zeros(10), ones(10))
+
+ likelihoods_to_test = [
+ ExponentialLikelihood(),
+ GammaLikelihood(),
+ PoissonLikelihood(),
+ GaussianLikelihood(),
+ ]
+
+ @testset "testing all analytic implementations" begin
+ # Test that we're not missing any analytic implementation in `likelihoods_to_test`!
+ implementation_types = [
+ (; quadrature=m.sig.types[2], lik=m.sig.types[3]) for
+ m in methods(GPLikelihoods.expected_loglikelihood)
+ ]
+ analytic_likelihoods = [
+ m.lik for m in implementation_types if
+ m.quadrature == GPLikelihoods.AnalyticExpectation && m.lik != Any
+ ]
+ for lik_type in analytic_likelihoods
+ lik_type_instances = filter(lik -> isa(lik, lik_type), likelihoods_to_test)
+ @test !isempty(lik_type_instances)
+ lik = first(lik_type_instances)
+ @test GPLikelihoods.default_expectation_method(lik) isa
+ GPLikelihoods.AnalyticExpectation
+ end
+ end
+
+ @testset "$(nameof(typeof(lik)))" for lik in likelihoods_to_test
+ methods = [
+ GaussHermiteExpectation(100),
+ MonteCarloExpectation(1e7),
+ GPLikelihoods.DefaultExpectationMethod(),
+ ]
+ def = GPLikelihoods.default_expectation_method(lik)
+ if def isa GPLikelihoods.AnalyticExpectation
+ push!(methods, def)
+ end
+ y = rand.(rng, lik.(zeros(10)))
+
+ results = map(m -> GPLikelihoods.expected_loglikelihood(m, lik, q_f, y), methods)
+ @test all(x -> isapprox(x, results[end]; atol=1e-6, rtol=1e-3), results)
+ end
+
+ @test GPLikelihoods.expected_loglikelihood(
+ MonteCarloExpectation(1), GaussianLikelihood(), q_f, zeros(10)
+ ) isa Real
+ @test GPLikelihoods.expected_loglikelihood(
+ GaussHermiteExpectation(1), GaussianLikelihood(), q_f, zeros(10)
+ ) isa Real
+ @test GPLikelihoods.default_expectation_method(θ -> Normal(0, θ)) isa
+ GaussHermiteExpectation
+
+ # see https://github.com/JuliaGaussianProcesses/ApproximateGPs.jl/issues/82
+ @testset "testing Zygote compatibility with GaussHermiteExpectation" begin
+ N = 10
+ gh = GaussHermiteExpectation(12)
+ μs = randn(rng, N)
+ σs = rand(rng, N)
+
+ # Test differentiation with variational parameters
+ for lik in likelihoods_to_test
+ y = rand.(rng, lik.(rand.(Normal.(μs, σs))))
+ gμ, glogσ = Zygote.gradient(μs, log.(σs)) do μ, logσ
+ GPLikelihoods.expected_loglikelihood(gh, lik, Normal.(μ, exp.(logσ)), y)
+ end
+ @test all(isfinite, gμ)
+ @test all(isfinite, glogσ)
+ end
+
+ # Test differentiation with likelihood parameters
+ # Test GaussianLikelihood parameter
+ σ = 1.0
+ y = randn(rng, N)
+ glogσ = only(
+ Zygote.gradient(log(σ)) do x
+ GPLikelihoods.expected_loglikelihood(
+ gh, GaussianLikelihood(exp(x)), Normal.(μs, σs), y
+ )
+ end,
+ )
+ @test isfinite(glogσ)
+
+ # Test GammaLikelihood parameter
+ α = 2.0
+ y = rand.(rng, Gamma.(α, rand(N)))
+ glogα = only(
+ Zygote.gradient(log(α)) do x
+ GPLikelihoods.expected_loglikelihood(
+ gh, GammaLikelihood(exp(x)), Normal.(μs, σs), y
+ )
+ end,
+ )
+ @test isfinite(glogα)
+ end
+end

test/runtests.jl

@@ -1,10 +1,12 @@
 using GPLikelihoods
+using GPLikelihoods: GaussHermiteExpectation, MonteCarloExpectation
 using GPLikelihoods.TestInterface: test_interface
 using Test
 using Random
 using Functors
 using Distributions
 using StatsFuns
+using Zygote
 
 @testset "GPLikelihoods.jl" begin
  include("links.jl")
@@ -17,4 +19,5 @@ using StatsFuns
  include("likelihoods/exponential.jl")
  include("likelihoods/negativebinomial.jl")
  end
+ include("expectations.jl")
 end