Making KernelDensity.pdf interface consistent with Distributions.pdf

.gitignore

@@ -1,2 +1,5 @@
 examples/.ipynb_checkpoints/*
-Manifest.toml
+Manifest.toml
+.vscode/settings.json
+tmp.jl
+Project.toml

src/KernelDensity.jl

@@ -12,6 +12,8 @@ export kde, kde_lscv, UnivariateKDE, BivariateKDE, InterpKDE, pdf
 
 abstract type AbstractKDE end
 
+Base.broadcastable(x::AbstractKDE) = Ref(x)
+
 include("univariate.jl")
 include("bivariate.jl")
 include("interp.jl")

src/interp.jl

@@ -1,12 +1,12 @@
 import Interpolations: interpolate, scale
 
+
 mutable struct InterpKDE{K,I} <: AbstractKDE
     kde::K
     itp::I
     InterpKDE{K,I}(kde::K, itp::I) where {K,I} = new{K,I}(kde, itp)
 end
 
-
 function InterpKDE(kde::UnivariateKDE, opts...)
     itp_u = interpolate(kde.density, opts...)
     itp = scale(itp_u, kde.x)
@@ -24,9 +24,20 @@ function InterpKDE(kde::BivariateKDE, opts...)
 end
 InterpKDE(kde::BivariateKDE) = InterpKDE(kde::BivariateKDE, BSpline(Quadratic(Line(OnGrid()))))
 
+
+# pdf interface implementation
+# it should be consistent with Distributions.pdf
+
+# any dimension
 pdf(ik::InterpKDE,x::Real...) = ik.itp(x...)
-pdf(ik::InterpKDE,xs::AbstractVector) = [ik.itp(x) for x in xs]
-pdf(ik::InterpKDE,xs::AbstractVector,ys::AbstractVector) = [ik.itp(x,y) for x in xs, y in ys]
+pdf(ik::InterpKDE, V::AbstractVector) = ik.itp(V...)
+pdf(ik::InterpKDE, M::AbstractArray{<:Real, N}) where N = reshape(mapslices(v->pdf(ik, v), M, dims = 1), size(M)[2:end])
 
+# 1 dimension
 pdf(k::UnivariateKDE,x) = pdf(InterpKDE(k),x)
+Base.Broadcast.broadcasted(::typeof(pdf),k::UnivariateKDE,xs) = Base.Broadcast.broadcasted(InterpKDE(k).itp, xs)
+
+# 2 dimensions
 pdf(k::BivariateKDE,x,y) = pdf(InterpKDE(k),x,y)
+pdf(ik::InterpKDE,xs::AbstractVector,ys::AbstractVector) = ik.itp.(xs,ys')
+pdf(k::BivariateKDE, M) = pdf(InterpKDE(k),M)

test/interp.jl

@@ -1,25 +1,50 @@
 using Test
 using KernelDensity
 
-X = randn(100)
-Y = randn(100)
-
-k = kde(X)
-@test pdf(k, k.x) ≈ k.density
-
-k = kde((X,Y))
-@test pdf(k, k.x, k.y) ≈ k.density
-
-# Try to evaluate the KDE outside the interpolation domain
-# The KDE is allowed to be zero, but it should not be greater than the exact solution
-k = kde([0.0], bandwidth=1.0)
-@test pdf(k, k.x) ≈ k.density
-@test pdf(k, -10.0) ≤ pdf(Normal(), -10.0)
-@test pdf(k, +10.0) ≤ pdf(Normal(), +10.0)
-
-k = kde(([0.0],[0.0]), bandwidth=(1.0, 1.0))
-@test pdf(k, k.x, k.y) ≈ k.density
-@test pdf(k, -10.0, 0.0) ≤ pdf(MvNormal(2, 1.0), [-10.0, 0.0])
-@test pdf(k, +10.0, 0.0) ≤ pdf(MvNormal(2, 1.0), [+10.0, 0.0])
-@test pdf(k, 0.0, -10.0) ≤ pdf(MvNormal(2, 1.0), [0.0, -10.0])
-@test pdf(k, 0.0, +10.0) ≤ pdf(MvNormal(2, 1.0), [0.0, +10.0])
+@testset "interpolation computation" begin
+    X = randn(100)
+    Y = randn(100)
+
+    k = kde(X)
+    @test pdf.(k, k.x) ≈ k.density
+
+    k = kde((X,Y))
+    @test pdf(k, k.x, k.y) ≈ k.density
+
+    # Try to evaluate the KDE outside the interpolation domain
+    # The KDE is allowed to be zero, but it should not be greater than the exact solution
+    k = kde([0.0], bandwidth=1.0)
+    @test pdf.(k, k.x) ≈ k.density
+    @test pdf(k, -10.0) ≤ pdf(Normal(), -10.0)
+    @test pdf(k, +10.0) ≤ pdf(Normal(), +10.0)
+
+    k = kde(([0.0],[0.0]), bandwidth=(1.0, 1.0))
+    @test pdf(k, k.x, k.y) ≈ k.density
+    @test pdf(k, -10.0, 0.0) ≤ pdf(MvNormal(2, 1.0), [-10.0, 0.0])
+    @test pdf(k, +10.0, 0.0) ≤ pdf(MvNormal(2, 1.0), [+10.0, 0.0])
+    @test pdf(k, 0.0, -10.0) ≤ pdf(MvNormal(2, 1.0), [0.0, -10.0])
+    @test pdf(k, 0.0, +10.0) ≤ pdf(MvNormal(2, 1.0), [0.0, +10.0])
+end
+
+@testset "pdf method interface" begin
+    k = kde([-1., 1.])
+    ik = InterpKDE(k)
+
+    @test pdf.(k, [0., 1.]) ≈ [pdf(k, 0.), pdf(k, 1.)] ≈ pdf.(k,[0., 1.]) ≈ [pdf(ik, 0.), pdf(ik, 1.)]
+    @test all( pdf.(k, (0., 1.)) .≈ (pdf(k, 0.), pdf(k, 1.)) )
+    @test pdf.(k, [0. 1.; 2. -1.]) ≈ [pdf(k, 0.) pdf(k, 1.); pdf(k, 2.) pdf(k, -1.)]
+
+    k2d = kde(([-1., 1.], [0., 1.]))
+    ik2d = InterpKDE(k2d)
+
+    @test pdf(k2d, [0.5, 0.1]) ≈ pdf(k2d, [0.5; 0.1]) ≈ pdf(k2d, 0.5, 0.1) ≈ pdf(ik2d, 0.5, 0.1)
+    @test pdf(k2d, [0.5 1.; 0.1 1.]) ≈ [pdf(ik2d, 0.5, 0.1), pdf(ik2d, 1., 1.)]
+    @static if VERSION >= v"1.1"
+        @test pdf(k2d, [0.5; 1. ;;; 0.1; 1.]) ≈ [pdf(ik2d, 0.5, 1.) pdf(ik2d, 0.1, 1.)]
+    else
+        M = zeros(2, 1, 2)
+        M[:,1,1] .= [0.5, 1]
+        M[:,1,2] .= [0.1, 1]
+        @test pdf(k2d, M) ≈ [pdf(ik2d, 0.5, 1.) pdf(ik2d, 0.1, 1.)]
+    end
+end