Add Heteroscedastic Gaussian Likelihood (#17)

* Add Heteroscedastic Gaussian Likelihood

* Address code review

* Usa isa

* Update src/likelihoods/gaussian.jl

Co-authored-by: David Widmann <[email protected]>

* Edit docs, patch bump, style fix

Co-authored-by: David Widmann <[email protected]>

Author: sharanry

Project.toml

@@ -1,7 +1,7 @@
 name = "GPLikelihoods"
 uuid = "6031954c-0455-49d7-b3b9-3e1c99afaf40"
 authors = ["willtebbutt <[email protected]>"]
-version = "0.1.0"
+version = "0.1.1"
 
 [deps]
 AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"

src/GPLikelihoods.jl

@@ -7,7 +7,9 @@ using Functors
 
 import Distributions
 
-export GaussianLikelihood, PoissonLikelihood
+export GaussianLikelihood, 
+    HeteroscedasticGaussianLikelihood, 
+    PoissonLikelihood
 
 # Likelihoods
 include("likelihoods/gaussian.jl")

src/likelihoods/gaussian.jl

@@ -20,3 +20,21 @@ GaussianLikelihood() = GaussianLikelihood(1e-6)
 (l::GaussianLikelihood)(f::Real) = Normal(f, sqrt(l.σ²))
 
 (l::GaussianLikelihood)(fs::AbstractVector{<:Real}) = MvNormal(fs, sqrt(l.σ²))
+
+"""
+    HeteroscedasticGaussianLikelihood(σ²)
+
+Heteroscedastic Gaussian likelihood. 
+This is a Gaussian likelihood whose mean and the log of whose variance are functions of the
+latent process.
+
+```math
+    p(y|[f, g]) = Normal(y | f, exp(g))
+```
+On calling, this would return a normal distribution with mean `f` and variance `exp(g)`.
+"""
+struct HeteroscedasticGaussianLikelihood end
+
+(::HeteroscedasticGaussianLikelihood)(f::AbstractVector{<:Real}) = Normal(f[1], exp(f[2]))
+
+(::HeteroscedasticGaussianLikelihood)(fs::AbstractVector) = MvNormal(first.(fs), exp.(last.(fs)))

test/likelihoods/gaussian.jl

@@ -11,3 +11,25 @@
     @test length(rand(rng, lik(rand(rng, lfgp.fx)))) == 10
     @test keys(Functors.functor(lik)[1]) == (:σ²,)
 end
+
+@testset "HeteroscedasticGaussianLikelihood" begin
+    rng = MersenneTwister(123)
+    gp = GP(IndependentMOKernel(SqExponentialKernel()))
+    IN_DIM = 3
+    OUT_DIM = 2 # one for the mean the other for the log-standard deviation
+    N = 10
+    x = [rand(rng, IN_DIM) for _ in 1:N]
+    X = MOInput(x, OUT_DIM)
+    lik = HeteroscedasticGaussianLikelihood()
+    lgp = LatentGP(gp, lik, 1e-5)
+    lfgp = lgp(X)
+
+    Y = rand(rng, lfgp.fx)
+
+    y = [Y[[i + j*N for j in 0:(OUT_DIM - 1)]] for i in 1:N]
+    # Replace with mo_inverse_transform once it is merged
+
+    @test lik(y) isa Distribution
+    @test length(rand(rng, lik(y))) == 10
+    @test Functors.functor(lik)[1] == ()
+end