Not tested yet.
import numpy as np
from sklearn.neighbors import KDTree
from numba import vectorize
@vectorize(["float64(float64, float64, float64, float64, float64, float64, float64)"])
def binorm(xi, yi, x, y, xvar, yvar, xyvar=0):
X, Y = xi - x, yi - y
det = xvar * yvar - xyvar**2
p = (np.exp(-(X**2 * yvar + Y**2 * xvar - 2 * X * Y * xyvar) / (2 * det)) / (2 * np.pi * np.sqrt(det)))
return p
def binorm_grid(xi, yi, x, y, xvar, yvar, xyvar=0):
xi, yi = xi[..., None, None], yi[..., None, :, None]
p = binorm(xi, yi, x, y, xvar, yvar, xyvar)
return p.mean(-1)
class AdapKDE2D:
def __init__(self, x, y, n_eps=2, n_ngb=None, scale=None):
"""
scale: tuple of length 2
Naively, one may use scale = [x.std(), y.std()].
Only the ratio of scale matters. None for equal scale.
"""
npts = len(x)
if scale is None:
pts = np.vstack([x, y]).T
else:
pts = np.vstack([x / scale[0], y / scale[1]]).T
if n_ngb is None:
n_ngb = int(np.sqrt(npts))
eps = KDTree(pts).query(pts, n_ngb)[0].T[-1] / np.sqrt(n_ngb / np.pi)
kern = n_eps * eps
if scale is None:
xsig, ysig = kern, kern
else:
xsig, ysig = kern * scale[0], kern * scale[1]
self.x = x
self.y = y
self.xsig = xsig
self.ysig = ysig
self.xvar = xsig**2
self.yvar = ysig**2
def density(self, xi, yi):
p = binorm(xi, yi, self.x, self.y, self.xvar, self.yvar)
return p
def densisty_mesh(self, xi, yi):
p = binorm_grid(xi, yi, self.x, self.y, self.xvar, self.yvar)
return p