change bivariate to use FFTs, make inteface consistent

Author: simonbyrne

README.md

@@ -11,8 +11,8 @@ The main accessor function is `kde`:
 ```
 kde(data)
 ```
-will construct a `UnivariateKDE` object from the data. The optional keyword arguments are
-* `endpoints`: the lower and upper limits of the kde as a tuple. Due to the
+will construct a `UnivariateKDE` object from the real vector `data`. The optional keyword arguments are
+* `boundary`: the lower and upper limits of the kde as a tuple. Due to the
   fourier transforms used internally, there should be sufficient spacing to
   prevent wrap-around at the boundaries.
 * `npoints`: the number of interpolation points to use. The function uses
@@ -34,16 +34,23 @@ allows specifying the internal grid to use. Optional keyword arguments are
 kde(data, dist::Distribution)
 ```
 allows specifying the exact distribution to use as the kernel. Optional
-keyword arguments are `endpoints` and `npoints`.
+keyword arguments are `boundary` and `npoints`.
 
 ```
 kde(data, midpoints::Range, dist::Distribution)
 ```
 allows specifying both the distribution and grid.
 
-## To do
+### Bivariate
 
-* Use an in-place FFT.
-* Spline interpolation
-* Bias correction
-* Improve bandwidth selection
+The usage mirrors that of the univariate case, except that `data` is now
+either a tuple of vectors
+```
+kde((xdata, ydata))
+```
+or a matrix with two columns
+```
+kde(datamatrix)
+```
+Similarly, the optional arguments all now take tuple arguments:
+e.g. `boundary` now takes a tuple of tuples `((xlo,xhi),(ylo,yhi))`.

TODO

@@ -0,0 +1,6 @@
+* Interface to plotting libs
+* Use in-place FFT
+* Spline interpolation
+* Bias correction
+* Improve bandwidth selection (particularly for bivariate)
+* Allow use of arbitrary bivariate kernels

src/KDE.jl

@@ -3,8 +3,8 @@ module KDE
 using StatsBase
 using Distributions
 
-import Base: conv, FloatRange
-import StatsBase: RealVector
+import Base: conv
+import StatsBase: RealVector, RealMatrix
 import Distributions: twoπ
 
 export kde, UnivariateKDE, BivariateKDE

src/bivariate.jl

@@ -1,46 +1,114 @@
 # Store both grid and density for KDE over R2
-immutable BivariateKDE
-    x::Vector{Float64}
-    y::Vector{Float64}
+immutable BivariateKDE{Rx<:Range,Ry<:Range}
+    x::Rx
+    y::Ry
     density::Matrix{Float64}
 end
 
+function kernel_dist{D<:UnivariateDistribution}(::Type{D},w::(Real,Real))
+    kernel_dist(D,w[1]), kernel_dist(D,w[2])
+end
+function kernel_dist{Dx<:UnivariateDistribution,Dy<:UnivariateDistribution}(::Type{(Dx,Dy)},w::(Real,Real))
+    kernel_dist(Dx,w[1]), kernel_dist(Dy,w[2])
+end
 
-# Algorithm from MASS Chapter 5 for calculating 2D KDE
-function kde(x::RealVector, y::RealVector; width::Float64=NaN, resolution::Int=25)
-    n = length(x)
+# TODO: there are probably better choices.
+function default_bandwidth(data::(RealVector,RealVector))
+    default_bandwidth(data[1]), default_bandwidth(data[2])
+end
 
-    if length(y) != n
-        error("x and y must have the same length")
-    end
+# tabulate data for kde
+function tabulate(data::(RealVector, RealVector), midpoints::(Range, Range))
+    xdata, ydata = data
+    ndata = length(xdata)
+    length(ydata) == ndata || error("data vectors must be of same length")
 
-    if isnan(width)
-        h1 = kde_bandwidth(x)
-        h2 = kde_bandwidth(y)
-    else
-        h1 = width
-        h2 = width
+    xmid, ymid = midpoints
+    nx, ny = length(xmid), length(ymid)
+    sx, sy = step(xmid), step(ymid)
+
+    # Set up a grid for discretized data
+    grid = zeros(Float64, nx, ny)
+    ainc = 1.0 / (ndata*(sx*sy)^2)
+
+    # weighted discretization (cf. Jones and Lotwick)
+    for (x, y) in zip(xdata,ydata)
+        kx, ky = searchsortedfirst(xmid,x), searchsortedfirst(ymid,y)
+        jx, jy = kx-1, ky-1
+        if 1 <= jx <= nx && 1 <= jy <= ny
+            grid[jx,jy] += (xmid[kx]-x)*(ymid[ky]-y)*ainc
+            grid[kx,jy] += (x-xmid[jx])*(ymid[ky]-y)*ainc
+            grid[jx,ky] += (xmid[kx]-x)*(y-ymid[jy])*ainc
+            grid[kx,ky] += (x-xmid[jx])*(y-ymid[jy])*ainc
+        end
     end
 
-    min_x, max_x = extrema(x)
-    min_y, max_y = extrema(y)
+    # returns an un-convolved KDE
+    BivariateKDE(xmid, ymid, grid)
+end
+
+# convolution with product distribution of two univariates distributions
+function conv(k::BivariateKDE, dist::(UnivariateDistribution,UnivariateDistribution) )
+    # Transform to Fourier basis
+    Kx, Ky = size(k.density)
+    ft = rfft(k.density)
 
-    grid_x = [min_x:((max_x - min_x) / (resolution - 1)):max_x]
-    grid_y = [min_y:((max_y - min_y) / (resolution - 1)):max_y]
+    distx, disty = dist
 
-    mx = Array(Float64, resolution, n)
-    my = Array(Float64, resolution, n)
-    for i in 1:resolution
-        for j in 1:n
-            mx[i, j] = pdf(Normal(), (grid_x[i] - x[j]) / h1)
-            my[i, j] = pdf(Normal(), (grid_y[i] - y[j]) / h2)
+    # Convolve fft with characteristic function of kernel
+    cx = -twoπ/(step(k.x)*Kx)
+    cy = -twoπ/(step(k.y)*Ky)
+    for j = 1:size(ft,2)
+        for i = 1:size(ft,1)
+            ft[i,j] *= cf(distx,(i-1)*cx)*cf(disty,min(j-1,Ky-j+1)*cy)
         end
     end
 
-    z = A_mul_Bt(mx, my)
-    for i in 1:(resolution^2)
-        z[i] /= (n * h1 * h2)
-    end
+    # Invert the Fourier transform to get the KDE
+    BivariateKDE(k.x, k.y, irfft(ft, Kx))
+end
+
+typealias BivariateDistribution Union(MultivariateDistribution,(UnivariateDistribution,UnivariateDistribution))
+
+function kde(data::(RealVector, RealVector), midpoints::(Range, Range), dist::BivariateDistribution)
+    k = tabulate(data,midpoints)
+    conv(k,dist)
+end
+
+function kde(data::(RealVector, RealVector), dist::BivariateDistribution;
+             boundary::((Real,Real),(Real,Real)) = (kde_boundary(data[1],std(dist[1])),
+                                                     kde_boundary(data[2],std(dist[2]))),
+             npoints::(Int,Int)=(128,128))
+
+    xmid = kde_range(boundary[1],npoints[1])
+    ymid = kde_range(boundary[2],npoints[2])
+
+    kde(data,(xmid,ymid),dist)
+end
+
+function kde(data::(RealVector, RealVector), midpoints::(Range, Range);
+             bandwidth=default_bandwidth(data), kernel=Normal)
+
+    dist = kernel_dist(kernel,bandwidth)
+    kde(data,midpoints,dist)
+end
+
+function kde(data::(RealVector, RealVector);
+             bandwidth=default_bandwidth(data),
+             kernel=Normal,
+             boundary::((Real,Real),(Real,Real)) = (kde_boundary(data[1],bandwidth[1]),
+                                                     kde_boundary(data[2],bandwidth[2])),
+             npoints::(Int,Int)=(128,128))
+
+    dist = kernel_dist(kernel,bandwidth)
+    xmid = kde_range(boundary[1],npoints[1])
+    ymid = kde_range(boundary[2],npoints[2])
+
+    kde(data,(xmid,ymid),dist)
+end
 
-    return BivariateKDE(grid_x, grid_y, z)
+# matrix data
+function kde(data::RealMatrix,args...;kwargs...)
+    size(data,2) == 2 || error("Can only construct KDE from matrices with 2 columns.")
+    kde((data[:,1],data[:,2]),args...;kwargs...)
 end

src/univariate.jl

@@ -13,7 +13,7 @@ kernel_dist{D}(::Type{D},w::Real) = (s = w/std(D(0.0,1.0)); D(0.0,s))
 
 
 # Silverman's rule of thumb for KDE bandwidth selection
-function kde_bandwidth(data::Vector{Float64}, alpha::Float64 = 0.9)
+function default_bandwidth(data::Vector{Float64}, alpha::Float64 = 0.9)
     # Determine length of data
     ndata = length(data)
 
@@ -51,11 +51,19 @@ end
 
 # default kde range
 # Should extend enough beyond the data range to avoid cyclic correlation from the FFT
-function kde_range(data::RealVector, bandwidth::Real)
+function kde_boundary(data::RealVector, bandwidth::Real)
     lo, hi = extrema(data)
     lo - 4.0*bandwidth, hi + 4.0*bandwidth
 end
 
+# convert boundary and npoints to Range object
+function kde_range(boundary::(Real,Real), npoints::Int)
+    lo, hi = boundary
+    lo < hi || error("boundary (a,b) must have a < b")
+
+    step = (hi - lo) / (npoints-1)
+    lo:step:hi
+end
 
 # tabulate data for kde
 function tabulate(data::RealVector, midpoints::Range)
@@ -67,7 +75,7 @@ function tabulate(data::RealVector, midpoints::Range)
     grid = zeros(Float64, npoints)
     ainc = 1.0 / (ndata*s*s)
 
-    # weighted discetization (cf. Jones and Lotwick)
+    # weighted discretization (cf. Jones and Lotwick)
     for x in data
         k = searchsortedfirst(midpoints,x)
         j = k-1
@@ -81,11 +89,9 @@ function tabulate(data::RealVector, midpoints::Range)
     UnivariateKDE(midpoints, grid)
 end
 
-
-
 # convolve raw KDE with kernel
 # TODO: use in-place fft
-function conv(k::UnivariateKDE, dist::Distribution)
+function conv(k::UnivariateKDE, dist::UnivariateDistribution)
     # Transform to Fourier basis
     K = length(k.density)
     ft = rfft(k.density)
@@ -105,34 +111,29 @@ function conv(k::UnivariateKDE, dist::Distribution)
     UnivariateKDE(k.x, irfft(ft, K))
 end
 
-
-function kde(data::RealVector, midpoints::Range, dist::Distribution)
+# main kde interface methods
+function kde(data::RealVector, midpoints::Range, dist::UnivariateDistribution)
     k = tabulate(data, midpoints)
     conv(k,dist)
 end
 
-function kde(data::RealVector, dist::Distribution; 
-             endpoints::(Real,Real)=kde_range(data,std(dist)), npoints::Int=2048)
-    
-    lo, hi = endpoints
-    lo < hi || error("endpoints (a,b) must have a < b")
-
-    step = (hi - lo) / npoints
-    midpoints = lo:step:hi
+function kde(data::RealVector, dist::UnivariateDistribution;
+             boundary::(Real,Real)=kde_boundary(data,std(dist)), npoints::Int=2048)
 
+    midpoints = kde_range(boundary,npoints)
     kde(data,midpoints,dist)
 end
 
-function kde(data::RealVector, midpoints::Range; 
-            bandwidth=kde_bandwidth(data), kernel=Normal)
-    bandwidth <= 0.0 && error("Bandwidth must be positive")
+function kde(data::RealVector, midpoints::Range;
+            bandwidth=default_bandwidth(data), kernel=Normal)
+    bandwidth > 0.0 || error("Bandwidth must be positive")
     dist = kernel_dist(kernel,bandwidth)
     kde(data,midpoints,dist)
 end
 
-function kde(data::RealVector; bandwidth=kde_bandwidth(data), kernel=Normal, 
-             npoints::Int=2048, endpoints::(Real,Real)=kde_range(data,bandwidth))
-    bandwidth <= 0.0 && error("Bandwidth must be positive")
+function kde(data::RealVector; bandwidth=default_bandwidth(data), kernel=Normal,
+             npoints::Int=2048, boundary::(Real,Real)=kde_boundary(data,bandwidth))
+    bandwidth > 0.0 || error("Bandwidth must be positive")
     dist = kernel_dist(kernel,bandwidth)
-    kde(data,dist;endpoints=endpoints,npoints=npoints)
+    kde(data,dist;boundary=boundary,npoints=npoints)
 end