syrte
1/7/2019 - 6:20 PM

Not tested yet.

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