Skip to content

Commit

Permalink
Correctly orient (y, x) arrays (#1095)
Browse files Browse the repository at this point in the history
  • Loading branch information
ianthomas23 authored Jun 16, 2022
1 parent de63b66 commit 99ee558
Show file tree
Hide file tree
Showing 3 changed files with 78 additions and 11 deletions.
63 changes: 62 additions & 1 deletion datashader/tests/test_utils.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
from __future__ import absolute_import
from datashape import dshape
import numpy as np
from xarray import DataArray

from datashader.utils import Dispatcher, apply, isreal
from datashader.utils import Dispatcher, apply, calc_res, isreal, orient_array


def test_Dispatcher():
Expand All @@ -27,7 +29,66 @@ def test_isreal():
assert not isreal('complex64')
assert not isreal('{x: int64, y: float64}')


def test_apply():
f = lambda a, b, c=1, d=2: a + b + c + d
assert apply(f, (1, 2,)) == 6
assert apply(f, (1, 2,), dict(c=3)) == 8


def test_calc_res():
x = [5, 7]
y = [0, 1]
z = [[0, 0], [0, 0]]
dims = ('y', 'x')

# x and y increasing
xarr = DataArray(z, coords=dict(x=x, y=y), dims=dims)
xres, yres = calc_res(xarr)
assert xres == 2
assert yres == -1

# x increasing, y decreasing
xarr = DataArray(z, coords=dict(x=x, y=y[::-1]), dims=dims)
xres, yres = calc_res(xarr)
assert xres == 2
assert yres == 1

# x decreasing, y increasing
xarr = DataArray(z, coords=dict(x=x[::-1], y=y), dims=dims)
xres, yres = calc_res(xarr)
assert xres == -2
assert yres == -1

# x and y decreasing
xarr = DataArray(z, coords=dict(x=x[::-1], y=y[::-1]), dims=dims)
xres, yres = calc_res(xarr)
assert xres == -2
assert yres == 1


def test_orient_array():
x = [5, 7]
y = [0, 1]
z = np.array([[0, 1], [2, 3]])
dims = ('y', 'x')

# x and y increasing
xarr = DataArray(z, coords=dict(x=x, y=y), dims=dims)
arr = orient_array(xarr)
assert np.array_equal(arr, z)

# x increasing, y decreasing
xarr = DataArray(z, coords=dict(x=x, y=y[::-1]), dims=dims)
arr = orient_array(xarr)
assert np.array_equal(arr, z[::-1])

# x decreasing, y increasing
xarr = DataArray(z, coords=dict(x=x[::-1], y=y), dims=dims)
arr = orient_array(xarr)
assert np.array_equal(arr, z[:, ::-1])

# x and y decreasing
xarr = DataArray(z, coords=dict(x=x[::-1], y=y[::-1]), dims=dims)
arr = orient_array(xarr)
assert np.array_equal(arr, z[::-1, ::-1])
6 changes: 3 additions & 3 deletions datashader/transfer_functions/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

from datashader.colors import rgb, Sets1to3
from datashader.composite import composite_op_lookup, over, validate_operator
from datashader.utils import nansum_missing, ngjit, orient_array
from datashader.utils import nansum_missing, ngjit

try:
import cupy
Expand Down Expand Up @@ -230,7 +230,7 @@ def _interpolate(agg, cmap, how, alpha, span, min_alpha, name, rescale_discrete_
raise ValueError("agg must be 2D")
interpolater = _normalize_interpolate_how(how)

data = orient_array(agg)
data = agg.data
if isinstance(data, da.Array):
data = data.compute()
else:
Expand Down Expand Up @@ -356,7 +356,7 @@ def _colorize(agg, color_key, how, alpha, span, min_alpha, name, color_baseline,

# Reorient array (transposing the category dimension first)
agg_t = agg.transpose(*((agg.dims[-1],)+agg.dims[:2]))
data = orient_array(agg_t).transpose([1, 2, 0])
data = agg_t.data.transpose([1, 2, 0])
if isinstance(data, da.Array):
data = data.compute()
color_data = data.copy()
Expand Down
20 changes: 13 additions & 7 deletions datashader/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ def nansum_missing(array, axis):

def calc_res(raster):
"""Calculate the resolution of xarray.DataArray raster and return it as the
two-tuple (xres, yres).
two-tuple (xres, yres). yres is positive if it is decreasing.
"""
h, w = raster.shape[-2:]
ydim, xdim = raster.dims[-2:]
Expand Down Expand Up @@ -251,6 +251,15 @@ def get_indices(start, end, coords, res):
return sidx, eidx


def _flip_array(array, xflip, yflip):
# array may have 2 or 3 dimensions, last one is x-dimension, last but one is y-dimension.
if yflip:
array = array[..., ::-1, :]
if xflip:
array = array[..., :, ::-1]
return array


def orient_array(raster, res=None, layer=None):
"""
Reorients the array to a canonical orientation depending on
Expand All @@ -277,12 +286,9 @@ def orient_array(raster, res=None, layer=None):
if layer is not None: array = array[layer-1]
r0zero = np.timedelta64(0, 'ns') if isinstance(res[0], np.timedelta64) else 0
r1zero = np.timedelta64(0, 'ns') if isinstance(res[1], np.timedelta64) else 0
if array.ndim == 2:
if res[0] < r0zero: array = array[:, ::-1]
if res[1] > r1zero: array = array[::-1]
else:
if res[0] < r0zero: array = array[:, :, ::-1]
if res[1] > r1zero: array = array[:, ::-1]
xflip = res[0] < r0zero
yflip = res[1] > r1zero
array = _flip_array(array, xflip, yflip)
return array


Expand Down

0 comments on commit 99ee558

Please sign in to comment.