import numba
import numpy as np
def loess_quantiles(x, xp, fp, q, *, bw):
"""Compute quantiles using local kernel approximations
:param q:
the quantiles to evaluate
:param bw:
the scales of the Gaussian kernel
"""
x, xp, fp, q, bw = (np.asarray(a) for a in (x, xp, fp, q, bw))
assert x.ndim == 2
assert xp.ndim == 2 and xp.shape[1] == x.shape[1]
assert fp.ndim == 1 and fp.shape[0] == xp.shape[0]
assert bw.ndim == 1 and bw.shape[0] == x.shape[1]
res = np.zeros((len(x), len(q)), dtype=fp.dtype)
sort_idx = np.argsort(fp)
_loess_quantiles(res, x, xp[sort_idx], fp[sort_idx], q, bw)
return res
@numba.jit(nopython=True, nogil=True)
def _loess_quantiles(res, x, xp, fp, q, bw):
bw = bw.reshape(1, -1)
for i in range(len(x)):
xi = x[i, :].reshape(1, -1)
dx = (xp[:, :] - xi) / bw
log_w = -0.5 * np.sum(dx ** 2.0, axis=1)
w = np.cumsum(np.exp(log_w - log_w.max()))
w = w / w[-1]
res[i, :] = np.interp(q, w, fp)