From 432adf56a0a255834218a6018f499d7138135066 Mon Sep 17 00:00:00 2001
From: Ian Thomas <ianthomas23@gmail.com>
Date: Wed, 15 Jun 2022 17:24:33 +0100
Subject: [PATCH] Correctly orient (y, x) arrays

---
 datashader/tests/test_utils.py            | 63 ++++++++++++++++++++++-
 datashader/transfer_functions/__init__.py |  6 +--
 datashader/utils.py                       | 20 ++++---
 3 files changed, 78 insertions(+), 11 deletions(-)

diff --git a/datashader/tests/test_utils.py b/datashader/tests/test_utils.py
index 22ec218b8..70f53691f 100644
--- a/datashader/tests/test_utils.py
+++ b/datashader/tests/test_utils.py
@@ -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():
@@ -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])
diff --git a/datashader/transfer_functions/__init__.py b/datashader/transfer_functions/__init__.py
index 7c8b6ca3f..f48a81d39 100755
--- a/datashader/transfer_functions/__init__.py
+++ b/datashader/transfer_functions/__init__.py
@@ -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
@@ -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:
@@ -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()
diff --git a/datashader/utils.py b/datashader/utils.py
index 5fa39c729..95f6ce1d3 100644
--- a/datashader/utils.py
+++ b/datashader/utils.py
@@ -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:]
@@ -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
@@ -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