import numpy as np
from ...util.distributions import sample_kde
from ...util.log import Handle
from ...util.math import flattengrid, linrng_, linspc_, logrng_, logspc_
from ...util.plot.grid import bin_centres_to_edges
logger = Handle(__name__)
[docs]class DensityGrid(object):
def __init__(
self, x, y, extent=None, bins=50, logx=False, logy=False, coverage_scale=1.2
):
"""
Build a grid of x-y coordinates for use in evaluating KDE functions.
Parameters
-----------
x : :class:`np.ndarray`
y : :class:`np.ndarray`
extent : :class:`list`
Optionally-specified extent for the grid in the form (xmin, xmax, ymin, ymax).
bins : :class:`int` | :class:`tuple`
Number of bins for the grid. Can optionally specify
a tuple with (xbins, ybins).
logx : :class:`bool`
Whether to use a log-spaced index for the x dimension of the grid.
logy: :class:`bool`
Whether to use a log-spaced index for the y dimension of the grid.
coverage_scale : :class:`float`
Multiplier for the range of the grid relative to the data. If >1, the grid
will extend beyond the data.
"""
# limits
self.logx = logx
self.logy = logy
if not isinstance(bins, int):
assert len(bins) == 2 # x-y bins
self.xbins, self.ybins = bins
else:
self.xbins, self.ybins = bins, bins
self.coverage_scale = coverage_scale
if extent is None:
self.xmin, self.xmax, self.ymin, self.ymax = self.extent_from_xy(x, y)
else:
self.xmin, self.xmax, self.ymin, self.ymax = extent
# validation
self.xmin, self.xmax = min([self.xmin, self.xmax]), max(
[self.xmin, self.xmax]
)
self.ymin, self.ymax = min([self.ymin, self.ymax]), max(
[self.ymin, self.ymax]
)
self.xstep = self.get_xstep()
self.ystep = self.get_ystep()
if self.logx:
assert self.xmin > 0.0
assert (self.xmin / self.xstep) > 0.0
if self.logy:
assert self.ymin > 0.0
assert (self.ymin / self.ystep) > 0.0
self.calculate_grid()
[docs] def calculate_grid(self):
self.get_centre_grid()
self.get_edge_grid()
[docs] def get_ystep(self):
if self.logy:
return (self.ymax / self.ymin) / self.ybins
else:
return (self.ymax - self.ymin) / self.ybins
[docs] def get_xstep(self):
if self.logx:
return (self.xmax / self.xmin) / self.xbins
else:
return (self.xmax - self.xmin) / self.xbins
[docs] def extent_from_xy(self, x, y, coverage_scale=None):
cov = coverage_scale or self.coverage_scale
expand_grid = (cov - 1.0) / 2
return [
*[linrng_, logrng_][self.logx](x, exp=expand_grid),
*[linrng_, logrng_][self.logy](y, exp=expand_grid),
]
[docs] def get_xrange(self):
return self.xmin, self.xmax
[docs] def get_yrange(self):
return self.ymin, self.ymax
[docs] def get_extent(self):
return [*self.get_xrange(), *self.get_yrange()]
[docs] def get_range(self):
return [[*self.get_xrange()], [*self.get_yrange()]]
[docs] def update_grid_centre_ticks(self):
if self.logx:
self.grid_xc = logspc_(self.xmin, self.xmax, 1.0, self.xbins)
else:
self.grid_xc = linspc_(self.xmin, self.xmax, self.xstep, self.xbins)
if self.logy:
self.grid_yc = logspc_(self.ymin, self.ymax, 1.0, self.ybins)
else:
self.grid_yc = linspc_(self.ymin, self.ymax, self.ystep, self.ybins)
[docs] def update_grid_edge_ticks(self):
self.update_grid_centre_ticks()
if self.logx:
self.grid_xe = np.exp(bin_centres_to_edges(np.log(np.sort(self.grid_xc))))
else:
self.grid_xe = bin_centres_to_edges(np.sort(self.grid_xc))
if self.logy:
self.grid_ye = np.exp(bin_centres_to_edges(np.log(np.sort(self.grid_yc))))
else:
self.grid_ye = bin_centres_to_edges(np.sort(self.grid_yc))
[docs] def get_centre_grid(self):
self.update_grid_centre_ticks()
self.grid_xci, self.grid_yci = np.meshgrid(self.grid_xc, self.grid_yc)
[docs] def get_edge_grid(self):
self.update_grid_edge_ticks()
self.grid_xei, self.grid_yei = np.meshgrid(self.grid_xe, self.grid_ye)
[docs] def get_hex_extent(self):
if self.logx:
xex = [np.log(self.xmin / self.xstep), np.log(self.xmax * self.xstep)]
else:
xex = [self.xmin - self.xstep, self.xmax + self.xstep]
if self.logy:
yex = [np.log(self.ymin / self.ystep), np.log(self.ymax * self.ystep)]
else:
yex = [self.ymin - self.ystep, self.ymax + self.ystep]
return xex + yex
[docs] def kdefrom(
self,
xy,
xtransform=lambda x: x,
ytransform=lambda x: x,
mode="centres",
bw_method=None,
):
"""
Take an x-y array and sample a KDE on the grid.
"""
arr = xy.copy()
# generate x grid over range spanned by log(x)
arr[:, 0] = xtransform(xy[:, 0])
# generate y grid over range spanned by log(y)
arr[:, 1] = ytransform(xy[:, 1])
if mode == "centres":
assert np.isfinite(self.grid_xc).all() and np.isfinite(self.grid_yc).all()
zi = sample_kde(
arr,
flattengrid(
np.meshgrid(xtransform(self.grid_xc), ytransform(self.grid_yc))
),
bw_method=bw_method,
)
zi = zi.reshape(self.grid_xci.shape)
elif mode == "edges":
assert np.isfinite(self.grid_xe).all() and np.isfinite(self.grid_ye).all()
zi = sample_kde(
arr,
flattengrid(
np.meshgrid(xtransform(self.grid_xe), ytransform(self.grid_ye))
),
bw_method=bw_method,
)
zi = zi.reshape(self.grid_xei.shape)
else:
raise NotImplementedError("Valid modes are 'centres' and 'edges'.")
return zi