import math
from toolz import memoize
import numpy as np
from datashader.glyphs.glyph import Glyph
from datashader.resampling import infer_interval_breaks
from datashader.utils import isreal, ngjit, ngjit_parallel
import numba
from numba import cuda, prange
try:
import cupy
from datashader.transfer_functions._cuda_utils import cuda_args
except Exception:
cupy = None
cuda_args = None
def _cuda_mapper(mapper):
@cuda.jit
def kernel(in_array, out_array):
i, j = cuda.grid(2)
if i < out_array.shape[0] and j < out_array.shape[1]:
out_array[i, j] = mapper(in_array[i, j])
def cuda_map(in_array):
out_array = cupy.zeros(in_array.shape, dtype='float64')
in_array = cuda.to_device(in_array)
kernel[cuda_args(in_array.shape)](in_array, out_array)
return out_array
return cuda_map
class _QuadMeshLike(Glyph):
def __init__(self, x, y, name):
self.x = x
self.y = y
self.name = name
@property
def ndims(self):
return 2
@property
def inputs(self):
return (self.x, self.y, self.name)
def validate(self, in_dshape):
if not isreal(in_dshape.measure[str(self.x)]):
raise ValueError('x must be real')
elif not isreal(in_dshape.measure[str(self.y)]):
raise ValueError('y must be real')
elif not isreal(in_dshape.measure[str(self.name)]):
raise ValueError('aggregate value must be real')
@property
def x_label(self):
return self.x
@property
def y_label(self):
return self.y
[docs]class QuadMeshRectilinear(_QuadMeshLike):
def _compute_bounds_from_1d_centers(
self, xr_ds, dim, maybe_expand=False, orient=True
):
vals = xr_ds[dim].values
# Assume dimension is sorted in ascending or descending order
v0, v1, v_nm1, v_n = [
vals[i] for i in [0, 1, -2, -1]
]
# Check if we should swap order
if v_n < v0:
descending = True
v0, v1, v_nm1, v_n = v_n, v_nm1, v1, v0
else:
descending = False
bounds = (v0 - 0.5 * (v1 - v0), v_n + 0.5 * (v_n - v_nm1))
if not orient and descending:
# swap back to descending order
bounds = bounds[1], bounds[0]
if maybe_expand:
bounds = self.maybe_expand_bounds(bounds)
return bounds
def compute_x_bounds(self, xr_ds):
return self._compute_bounds_from_1d_centers(xr_ds, self.x, maybe_expand=True)
def compute_y_bounds(self, xr_ds):
return self._compute_bounds_from_1d_centers(xr_ds, self.y, maybe_expand=True)
def compute_bounds_dask(self, xr_ds):
return self.compute_x_bounds(xr_ds), self.compute_y_bounds(xr_ds)
def infer_interval_breaks(self, centers):
# Infer breaks for 1D array of centers
return infer_interval_breaks(centers)
@memoize
def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2,
_antialias_stage_2_funcs):
x_name = self.x
y_name = self.y
name = self.name
@ngjit
@self.expand_aggs_and_cols(append)
def perform_extend(i, j, xs, ys, *aggs_and_cols):
x0i, x1i = xs[i], xs[i + 1]
# Make sure x0 <= x1
if x0i > x1i:
x0i, x1i = x1i, x0i
# Make sure single pixel quads are represented
if x0i == x1i:
x1i += 1
y0i, y1i = ys[j], ys[j + 1]
# Make sure y0 <= y1
if y0i > y1i:
y0i, y1i = y1i, y0i
# Make sure single pixel quads are represented
if y0i == y1i:
y1i += 1
# x1i and y1i are not included in the iteration. this
# serves to avoid overlapping quads and it avoids the need
# for special handling of quads that end on exactly on the
# upper bound.
for xi in range(x0i, x1i):
for yi in range(y0i, y1i):
append(j, i, xi, yi, *aggs_and_cols)
@cuda.jit
@self.expand_aggs_and_cols(append)
def extend_cuda(xs, ys, *aggs_and_cols):
i, j = cuda.grid(2)
if i < (xs.shape[0] - 1) and j < (ys.shape[0] - 1):
perform_extend(i, j, xs, ys, *aggs_and_cols)
@ngjit
@self.expand_aggs_and_cols(append)
def extend_cpu(xs, ys, *aggs_and_cols):
for i in range(len(xs) - 1):
for j in range(len(ys) - 1):
perform_extend(i, j, xs, ys, *aggs_and_cols)
def extend(aggs, xr_ds, vt, bounds, x_breaks=None, y_breaks=None):
from datashader.core import LinearAxis
use_cuda = cupy and isinstance(xr_ds[name].data, cupy.ndarray)
# Build axis transform (mapper) functions
if use_cuda:
x_mapper2 = _cuda_mapper(x_mapper)
y_mapper2 = _cuda_mapper(y_mapper)
else:
x_mapper2 = x_mapper
y_mapper2 = y_mapper
# Convert from bin centers to interval edges
if x_breaks is None:
x_centers = xr_ds[x_name].values
if use_cuda:
x_centers = cupy.array(x_centers)
x_breaks = self.infer_interval_breaks(x_centers)
if y_breaks is None:
y_centers = xr_ds[y_name].values
if use_cuda:
y_centers = cupy.array(y_centers)
y_breaks = self.infer_interval_breaks(y_centers)
x0, x1, y0, y1 = bounds
xspan = x1 - x0
yspan = y1 - y0
if x_mapper is LinearAxis.mapper:
xscaled = (x_breaks - x0) / xspan
else:
xscaled = (x_mapper2(x_breaks) - x0) / xspan
if y_mapper is LinearAxis.mapper:
yscaled = (y_breaks - y0) / yspan
else:
yscaled = (y_mapper2(y_breaks) - y0) / yspan
xinds = np.where((xscaled >= 0) & (xscaled <= 1))[0]
yinds = np.where((yscaled >= 0) & (yscaled <= 1))[0]
if len(xinds) == 0 or len(yinds) == 0:
# Nothing to do
return
xm0, xm1 = max(xinds.min() - 1, 0), xinds.max() + 1
ym0, ym1 = max(yinds.min() - 1, 0), yinds.max() + 1
plot_height, plot_width = aggs[0].shape[:2]
# Downselect xs and ys and convert to int
xs = (xscaled[xm0:xm1 + 1] * plot_width).astype(int).clip(0, plot_width)
ys = (yscaled[ym0:ym1 + 1] * plot_height).astype(int).clip(0, plot_height)
# For input "column", down select to valid range
cols_full = info(xr_ds.transpose(y_name, x_name), aggs[0].shape[:2])
cols = tuple([c[ym0:ym1, xm0:xm1] for c in cols_full])
aggs_and_cols = aggs + cols
if use_cuda:
do_extend = extend_cuda[cuda_args(xr_ds[name].shape)]
else:
do_extend = extend_cpu
do_extend(xs, ys, *aggs_and_cols)
return extend
@ngjit
def build_scale_translate(out_size, out0, out1, src_size, src0, src1):
translate_y = src_size * (out0 - src0) / (src1 - src0)
scale_y = (src_size * (out1 - out0)) / (out_size * (src1 - src0))
return scale_y, translate_y
[docs]class QuadMeshRaster(QuadMeshRectilinear):
def is_upsample(self, source, x, y, name, x_range, y_range, out_w, out_h):
# Check upsampling in x
src_w = len(source[x])
if x_range is None:
upsample_width = out_w >= src_w
else:
out_x0, out_x1 = x_range
src_x0, src_x1 = self._compute_bounds_from_1d_centers(
source, x, maybe_expand=False, orient=False
)
src_xbinsize = math.fabs((src_x1 - src_x0) / src_w)
out_xbinsize = math.fabs((out_x1 - out_x0) / out_w)
upsample_width = src_xbinsize >= out_xbinsize
# Check upsampling in y
src_h = len(source[y])
if y_range is None:
upsample_height = out_h >= src_h
else:
out_y0, out_y1 = y_range
src_y0, src_y1 = self._compute_bounds_from_1d_centers(
source, y, maybe_expand=False, orient=False
)
src_ybinsize = math.fabs((src_y1 - src_y0) / src_h)
out_ybinsize = math.fabs((out_y1 - out_y0) / out_h)
upsample_height = src_ybinsize >= out_ybinsize
return upsample_width, upsample_height
@memoize
def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2,
_antialias_stage_2_funcs):
x_name = self.x
y_name = self.y
name = self.name
@ngjit_parallel
def upsample_cpu(
src_w, src_h, translate_x, translate_y, scale_x, scale_y,
offset_x, offset_y, out_w, out_h, agg, col
):
for out_j in prange(out_h):
src_j = int(math.floor(scale_y * (out_j + 0.5) + translate_y - offset_y))
for out_i in range(out_w):
src_i = int(math.floor(scale_x * (out_i + 0.5) + translate_x - offset_x))
if src_j < 0 or src_j >= src_h or src_i < 0 or src_i >= src_w:
agg[out_j, out_i] = np.nan
else:
agg[out_j, out_i] = col[src_j, src_i]
@cuda.jit
def upsample_cuda(
src_w, src_h, translate_x, translate_y, scale_x, scale_y,
offset_x, offset_y, out_w, out_h, agg, col
):
out_i, out_j = cuda.grid(2)
if out_i < out_w and out_j < out_h:
src_j = int(math.floor(scale_y * (out_j + 0.5) + translate_y - offset_y))
src_i = int(math.floor(scale_x * (out_i + 0.5) + translate_x - offset_x))
if src_j < 0 or src_j >= src_h or src_i < 0 or src_i >= src_w:
agg[out_j, out_i] = np.nan
else:
agg[out_j, out_i] = col[src_j, src_i]
@ngjit_parallel
@self.expand_aggs_and_cols(append)
def downsample_cpu(
src_w, src_h, translate_x, translate_y, scale_x, scale_y,
offset_x, offset_y, out_w, out_h, *aggs_and_cols
):
for out_j in prange(out_h):
src_j0 = int(max(
math.floor(scale_y * (out_j + 0.0) + translate_y - offset_y), 0
))
src_j1 = int(min(
math.floor(scale_y * (out_j + 1.0) + translate_y - offset_y), src_h
))
for out_i in range(out_w):
src_i0 = int(max(
math.floor(scale_x * (out_i + 0.0) + translate_x - offset_x), 0
))
src_i1 = int(min(
math.floor(scale_x * (out_i + 1.0) + translate_x - offset_x), src_w
))
for src_j in range(src_j0, src_j1):
for src_i in range(src_i0, src_i1):
append(src_j, src_i, out_i, out_j, *aggs_and_cols)
@cuda.jit
@self.expand_aggs_and_cols(append)
def downsample_cuda(
src_w, src_h, translate_x, translate_y, scale_x, scale_y,
offset_x, offset_y, out_w, out_h, *aggs_and_cols
):
out_i, out_j = cuda.grid(2)
if out_i < out_w and out_j < out_h:
src_j0 = max(
math.floor(scale_y * (out_j + 0.0) + translate_y - offset_y), 0
)
src_j1 = min(
math.floor(scale_y * (out_j + 1.0) + translate_y - offset_y), src_h
)
src_i0 = max(
math.floor(scale_x * (out_i + 0.0) + translate_x - offset_x), 0
)
src_i1 = min(
math.floor(scale_x * (out_i + 1.0) + translate_x - offset_x), src_w
)
for src_j in range(src_j0, src_j1):
for src_i in range(src_i0, src_i1):
append(src_j, src_i, out_i, out_j, *aggs_and_cols)
def extend(aggs, xr_ds, vt, bounds,
scale_x=None, scale_y=None, translate_x=None, translate_y=None,
offset_x=None, offset_y=None, src_xbinsize=None, src_ybinsize=None):
use_cuda = cupy and isinstance(xr_ds[name].data, cupy.ndarray)
# Compute output constants
out_h, out_w = aggs[0].shape
out_x0, out_x1, out_y0, out_y1 = bounds
out_xbinsize = math.fabs((out_x1 - out_x0) / out_w)
out_ybinsize = math.fabs((out_y1 - out_y0) / out_h)
# Compute source constants
xr_ds = xr_ds.transpose(y_name, x_name)
src_h, src_w = xr_ds[name].shape
if (scale_x is None or scale_y is None or
translate_x is None or translate_y is None or
offset_x is None or offset_y is None or
src_xbinsize is None or src_ybinsize is None ):
# Compute bin sizes from bounds
src_x0, src_x1 = self._compute_bounds_from_1d_centers(
xr_ds, x_name, maybe_expand=False, orient=False
)
src_y0, src_y1 = self._compute_bounds_from_1d_centers(
xr_ds, y_name, maybe_expand=False, orient=False
)
src_xbinsize = math.fabs((src_x1 - src_x0) / src_w)
src_ybinsize = math.fabs((src_y1 - src_y0) / src_h)
# Compute scale/translate
scale_y, translate_y = build_scale_translate(
out_h, out_y0, out_y1, src_h, src_y0, src_y1
)
scale_x, translate_x = build_scale_translate(
out_w, out_x0, out_x1, src_w, src_x0, src_x1
)
offset_x = offset_y = 0
# Build aggs_and_cols tuple
cols = info(xr_ds, aggs[0].shape[:2])
aggs_and_cols = tuple(aggs) + tuple(cols)
if src_h == 0 or src_w == 0 or out_h == 0 or out_w == 0:
# Nothing to do
return
elif src_xbinsize >= out_xbinsize and src_ybinsize >= out_ybinsize:
# Upsample
if use_cuda:
do_sampling = upsample_cuda[cuda_args((out_w, out_h))]
else:
do_sampling = upsample_cpu
return do_sampling(
src_w, src_h, translate_x, translate_y, scale_x, scale_y,
offset_x, offset_y, out_w, out_h, aggs[0], cols[0]
)
else:
# Downsample. Note that caller is responsible for making sure to not
# mix upsampling and downsampling.
if use_cuda:
do_sampling = downsample_cuda[cuda_args((out_w, out_h))]
else:
do_sampling = downsample_cpu
return do_sampling(
src_w, src_h, translate_x, translate_y, scale_x, scale_y,
offset_x, offset_y, out_w, out_h, *aggs_and_cols
)
return extend
[docs]class QuadMeshCurvilinear(_QuadMeshLike):
def compute_x_bounds(self, xr_ds):
x_breaks = self.infer_interval_breaks(xr_ds[self.x].values)
bounds = Glyph._compute_bounds_2d(x_breaks)
return self.maybe_expand_bounds(bounds)
def compute_y_bounds(self, xr_ds):
y_breaks = self.infer_interval_breaks(xr_ds[self.y].values)
bounds = Glyph._compute_bounds_2d(y_breaks)
return self.maybe_expand_bounds(bounds)
def compute_bounds_dask(self, xr_ds):
return self.compute_x_bounds(xr_ds), self.compute_y_bounds(xr_ds)
def infer_interval_breaks(self, centers):
# Infer breaks for 1D array of centers
breaks = infer_interval_breaks(centers, axis=1)
breaks = infer_interval_breaks(breaks, axis=0)
return breaks
@memoize
def _build_extend(self, x_mapper, y_mapper, info, append, _antialias_stage_2,
_antialias_stage_2_funcs):
x_name = self.x
y_name = self.y
name = self.name
@ngjit
@self.expand_aggs_and_cols(append)
def perform_extend(
i, j,
plot_height, plot_width, xs, ys, xverts, yverts,
yincreasing, eligible, intersect, *aggs_and_cols
):
# make array of quad x any vertices
xverts[0] = xs[j, i]
xverts[1] = xs[j, i + 1]
xverts[2] = xs[j + 1, i + 1]
xverts[3] = xs[j + 1, i]
xverts[4] = xverts[0]
yverts[0] = ys[j, i]
yverts[1] = ys[j, i + 1]
yverts[2] = ys[j + 1, i + 1]
yverts[3] = ys[j + 1, i]
yverts[4] = yverts[0]
# Compute the rectilinear bounding box around the quad and
# skip quad if there is no chance for it to intersect
# viewport
xmin = min(min(xverts[0], xverts[1]), min(xverts[2], xverts[3]))
if xmin >= plot_width:
return
xmin = max(xmin, 0)
xmax = max(max(xverts[0], xverts[1]), max(xverts[2], xverts[3]))
if xmax < 0:
return
xmax = min(xmax, plot_width)
ymin = min(min(yverts[0], yverts[1]), min(yverts[2], yverts[3]))
if ymin >= plot_height:
return
ymin = max(ymin, 0)
ymax = max(max(yverts[0], yverts[1]), max(yverts[2], yverts[3]))
if ymax < 0:
return
ymax = min(ymax, plot_height)
# Handle subpixel quads
if xmin == xmax or ymin == ymax:
# If either dimension is a single pixel, then render it
if xmin == xmax and xmax < plot_width:
xmax += 1
if ymin == ymax and ymax < plot_height:
ymax += 1
for yi in range(ymin, ymax):
for xi in range(xmin, xmax):
append(j, i, xi, yi, *aggs_and_cols)
return
# make yincreasing an array holding whether each edge is
# increasing vertically (+1), decreasing vertically (-1),
# or horizontal (0).
yincreasing[:] = 0
for k in range(4):
if yverts[k + 1] > yverts[k]:
yincreasing[k] = 1
elif yverts[k + 1] < yverts[k]:
yincreasing[k] = -1
for yi in range(ymin, ymax):
eligible[:] = 1
for xi in range(xmin, xmax):
intersect[:] = 0
# Test edges
for edge_i in range(4):
# Skip if we already know edge is ineligible
if not eligible[edge_i]:
continue
# Check if edge is fully to left of point. If
# so, we don't need to consider it again for
# this row.
if ((xverts[edge_i] < xi) and
(xverts[edge_i + 1] < xi)):
eligible[edge_i] = 0
continue
# Check if edge is fully above or below point.
# If so, we don't need to consider it again
# for this row.
if ((yverts[edge_i] > yi) ==
(yverts[edge_i + 1] > yi)):
eligible[edge_i] = 0
continue
# Now check if edge is to the right of point.
# A is vector from point to first vertex
ax = xverts[edge_i] - xi
ay = yverts[edge_i] - yi
# B is vector from point to second vertex
bx = xverts[edge_i + 1] - xi
by = yverts[edge_i + 1] - yi
# Compute cross product of B and A
bxa = (bx * ay - by * ax)
# If cross product has same sign as yincreasing
# then edge intersects to the right
intersect[edge_i] = (
bxa * yincreasing[edge_i] < 0
)
intersections = (
intersect[0] + intersect[1] + intersect[2] + intersect[3]
)
if intersections % 2 == 1:
# If odd number of intersections, point
# is inside quad
append(j, i, xi, yi, *aggs_and_cols)
@cuda.jit
@self.expand_aggs_and_cols(append)
def extend_cuda(plot_height, plot_width, xs, ys, *aggs_and_cols):
# # For consistency with CPU path, we initialize all arrays here
# # xverts/yverts arrays
xverts = cuda.local.array(5, dtype=numba.types.int32)
yverts = cuda.local.array(5, dtype=numba.types.int32)
#
# # Array holding whether each edge is increasing
# # vertically (+1), decreasing vertically (-1),
# # or horizontal (0).
yincreasing = cuda.local.array(4, dtype=numba.types.int8)
# # Array that will hold mask of whether edges are
# # eligible for intersection tests
eligible = cuda.local.array(4, dtype=numba.types.int8)
# # Array that will hold a mask of whether edges
# # intersect the ray to the right of test point
intersect = cuda.local.array(4, dtype=numba.types.int8)
i, j = cuda.grid(2)
if i < (xs.shape[0] - 1) and j < (ys.shape[0] - 1):
perform_extend(
i, j, plot_height, plot_width, xs, ys,
xverts, yverts, yincreasing, eligible, intersect,
*aggs_and_cols
)
@ngjit
@self.expand_aggs_and_cols(append)
def extend_cpu(plot_height, plot_width, xs, ys, *aggs_and_cols):
# For performance, we initialize all arrays once before the loop
# xverts/yverts arrays
xverts = np.zeros(5, dtype=np.int32)
yverts = np.zeros(5, dtype=np.int32)
# Array holding whether each edge is increasing
# vertically (+1), decreasing vertically (-1),
# or horizontal (0).
yincreasing = np.zeros(4, dtype=np.int8)
# Array that will hold mask of whether edges are
# eligible for intersection tests
eligible = np.ones(4, dtype=np.int8)
# Array that will hold a mask of whether edges
# intersect the ray to the right of test point
intersect = np.zeros(4, dtype=np.int8)
y_len, x_len, = xs.shape
for i in range(x_len - 1):
for j in range(y_len - 1):
perform_extend(
i, j, plot_height, plot_width, xs, ys,
xverts, yverts, yincreasing, eligible, intersect, *aggs_and_cols
)
def extend(aggs, xr_ds, vt, bounds, x_breaks=None, y_breaks=None):
from datashader.core import LinearAxis
use_cuda = cupy and isinstance(xr_ds[name].data, cupy.ndarray)
# Build axis transform (mapper) functions
if use_cuda:
x_mapper2 = _cuda_mapper(x_mapper)
y_mapper2 = _cuda_mapper(y_mapper)
else:
x_mapper2 = x_mapper
y_mapper2 = y_mapper
# Convert from bin centers to interval edges
if x_breaks is None:
x_centers = xr_ds[x_name].values
if use_cuda:
x_centers = cupy.array(x_centers)
x_breaks = self.infer_interval_breaks(x_centers)
if y_breaks is None:
y_centers = xr_ds[y_name].values
if use_cuda:
y_centers = cupy.array(y_centers)
y_breaks = self.infer_interval_breaks(y_centers)
# Scale x and y vertices into integer canvas coordinates
x0, x1, y0, y1 = bounds
xspan = x1 - x0
yspan = y1 - y0
if x_mapper is LinearAxis.mapper:
xscaled = (x_breaks - x0) / xspan
else:
xscaled = (x_mapper2(x_breaks) - x0) / xspan
if y_mapper is LinearAxis.mapper:
yscaled = (y_breaks - y0) / yspan
else:
yscaled = (y_mapper2(y_breaks) - y0) / yspan
plot_height, plot_width = aggs[0].shape[:2]
xs = (xscaled * plot_width).astype(int)
ys = (yscaled * plot_height).astype(int)
coord_dims = xr_ds.coords[x_name].dims
aggs_and_cols = aggs + info(xr_ds.transpose(*coord_dims), aggs[0].shape[:2])
if use_cuda:
do_extend = extend_cuda[cuda_args(xr_ds[name].shape)]
else:
do_extend = extend_cpu
do_extend(
plot_height, plot_width, xs, ys, *aggs_and_cols
)
return extend
