fix warpgaussian logpdf error/making neal funnel logpdf working with mooncake

Author: zuhengxu

example/targets/banana.jl

@@ -1,7 +1,3 @@
-using Distributions, Random
-using Plots
-using IrrationalConstants
-
 """
     Banana{T<:Real}
 

example/targets/cross.jl

@@ -1,4 +1,3 @@
-using Distributions, Random
 """
     Cross(μ::Real=2.0, σ::Real=0.15)
 

example/targets/neal_funnel.jl

@@ -1,5 +1,3 @@
-using Distributions, Random
-
 """
     Funnel{T<:Real}
 
@@ -45,18 +43,19 @@ Funnel(dim::Int) = Funnel(dim, 0.0, 9.0)
 Base.length(p::Funnel) = p.dim
 Base.eltype(p::Funnel{T}) where {T<:Real} = T
 
-function Distributions._rand!(rng::AbstractRNG, p::Funnel, x::AbstractVecOrMat)
-    T = eltype(x)
+function Distributions._rand!(rng::AbstractRNG, p::Funnel{T}, x::AbstractVecOrMat{T}) where {T<:Real}
     d, μ, σ = p.dim, p.μ, p.σ
     d == size(x, 1) || error("Dimension mismatch")
     x[1, :] .= randn(rng, T, size(x, 2)) .* σ .+ μ
     x[2:end, :] .= randn(rng, T, d - 1, size(x, 2)) .* exp.(@view(x[1, :]) ./ 2)'
     return x
 end
 
-function Distributions._logpdf(p::Funnel, x::AbstractVector)
+function Distributions._logpdf(p::Funnel{T}, x::AbstractVector{T}) where {T<:Real}
     d, μ, σ = p.dim, p.μ, p.σ
-    lpdf1 = logpdf(Normal(μ, σ), x[1])
-    lpdfs = logpdf.(Normal.(zeros(T, d - 1), exp(x[1] / 2)), @view(x[2:end]))
-    return lpdf1 + sum(lpdfs)
+    x1 = x[1]
+    x2 = x[2:end]
+    lpdf_x1 = logpdf(Normal(μ, σ), x1)
+    lpdf_x2_given_1 = logpdf(MvNormal(zeros(T, d-1), exp(x1)I), x2)
+    return lpdf_x1 + lpdf_x2_given_1
 end

example/targets/warped_gaussian.jl

@@ -1,5 +1,3 @@
-using Distributions, Random, LinearAlgebra, IrrationalConstants
-
 """
     WarpedGauss{T<:Real}
 
@@ -39,11 +37,11 @@ WarpedGauss(σ1::T, σ2::T) where {T<:Real} = WarpedGauss{T}(σ1, σ2)
 WarpedGauss() = WarpedGauss(1.0, 0.12)
 
 Base.length(p::WarpedGauss) = 2
-Base.eltype(p::WarpedGauss{T}) where {T<:Real} = T
+Base.eltype(::WarpedGauss{T}) where {T<:Real} = T
 Distributions.sampler(p::WarpedGauss) = p
 
 # Define the transformation function φ and the inverse ϕ⁻¹ for the warped Gaussian distribution
-function ϕ!(p::WarpedGauss, z::AbstractVector)
+function ϕ!(::WarpedGauss, z::AbstractVector)
     length(z) == 2 || error("Dimension mismatch")
     x, y = z
     r = norm(z)
@@ -53,7 +51,7 @@ function ϕ!(p::WarpedGauss, z::AbstractVector)
     return z
 end
 
-function ϕ⁻¹(p::WarpedGauss, z::AbstractVector)
+function ϕ⁻¹(::WarpedGauss, z::AbstractVector)
     length(z) == 2 || error("Dimension mismatch")
     x, y = z
     r = norm(z)
@@ -71,7 +69,7 @@ end
 
 function Distributions._rand!(rng::AbstractRNG, p::WarpedGauss, x::AbstractVecOrMat)
     size(x, 1) == 2 || error("Dimension mismatch")
-    σ₁, σ₂ = p.σ₁, p.σ₂
+    σ₁, σ₂ = p.σ1, p.σ2
     randn!(rng, x)
     x .*= [σ₁, σ₂]
     for y in eachcol(x)
@@ -82,7 +80,7 @@ end
 
 function Distributions._logpdf(p::WarpedGauss, x::AbstractVector)
     size(x, 1) == 2 || error("Dimension mismatch")
-    σ₁, σ₂ = p.σ₁, p.σ₂
+    σ₁, σ₂ = p.σ1, p.σ2
     S = [σ₁, σ₂] .^ 2
     z, logJ = ϕ⁻¹(p, x)
     return -sum(z .^ 2 ./ S) / 2 - IrrationalConstants.log2π - log(σ₁) - log(σ₂) + logJ