diff --git a/verde/tests/test_utils.py b/verde/tests/test_utils.py index 0f8f8722e..0d6119b99 100644 --- a/verde/tests/test_utils.py +++ b/verde/tests/test_utils.py @@ -5,11 +5,12 @@ import numpy as np import numpy.testing as npt +import xarray as xr from scipy.spatial import cKDTree # pylint: disable=no-name-in-module import pytest from ..coordinates import grid_coordinates -from ..utils import parse_engine, dummy_jit, kdtree +from ..utils import parse_engine, dummy_jit, kdtree, grid_to_table from .. import utils @@ -52,3 +53,25 @@ def test_kdtree(): npt.assert_allclose(dist, 0.1) if not use_pykdtree: assert isinstance(tree, cKDTree) + + +def test_grid_to_table_order(): + "Check that coordinates are in the right order when converting to tables" + lon, lat = grid_coordinates(region=(1, 10, -10, -1), shape=(3, 4)) + data = lon ** 2 + # If the DataArray is created with coords in an order that doesn't match + # the dims (which is valid), we were getting it wrong because we were + # relying on the order of the coords instead of dims. This test would have + # caught that bug. + grid = xr.DataArray( + data=data, + coords={"longitude": lon[0, :], "latitude": lat[:, 0]}, + dims=("latitude", "longitude"), + ).to_dataset(name="field") + table = grid_to_table(grid) + true_lat = [-10, -10, -10, -10, -5.5, -5.5, -5.5, -5.5, -1, -1, -1, -1] + true_lon = [1, 4, 7, 10, 1, 4, 7, 10, 1, 4, 7, 10] + true_field = [1, 16, 49, 100, 1, 16, 49, 100, 1, 16, 49, 100] + npt.assert_allclose(true_lat, table.latitude) + npt.assert_allclose(true_lon, table.longitude) + npt.assert_allclose(true_field, table.field) diff --git a/verde/utils.py b/verde/utils.py index 42ee7359e..bb24c7ac9 100644 --- a/verde/utils.py +++ b/verde/utils.py @@ -237,6 +237,11 @@ def grid_to_table(grid): ... coords=(np.arange(4), np.arange(5, 10)), ... dims=['northing', 'easting'] ... ) + >>> print(temperature.values) + [[ 0 1 2 3 4] + [ 5 6 7 8 9] + [10 11 12 13 14] + [15 16 17 18 19]] >>> grid = xr.Dataset({"temperature": temperature}) >>> table = grid_to_table(grid) >>> list(sorted(table.columns)) @@ -267,23 +272,17 @@ def grid_to_table(grid): [20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39] """ - coordinate_names = [*grid.coords.keys()] - coord_north = grid.coords[coordinate_names[0]].values - coord_east = grid.coords[coordinate_names[1]].values - coordinates = [i.ravel() for i in np.meshgrid(coord_east, coord_north)] - coord_dict = { - coordinate_names[0]: coordinates[1], - coordinate_names[1]: coordinates[0], - } - variable_name = [*grid.data_vars.keys()] - variable_data = grid.to_array().values - variable_arrays = variable_data.reshape( - len(variable_name), int(len(variable_data.ravel()) / len(variable_name)) - ) - var_dict = dict(zip(variable_name, variable_arrays)) - coord_dict.update(var_dict) - data = pd.DataFrame(coord_dict) - return data + data_names = list(grid.data_vars.keys()) + data_arrays = [grid[name].values.ravel() for name in data_names] + coordinate_names = list(grid[data_names[0]].dims) + north = grid.coords[coordinate_names[0]].values + east = grid.coords[coordinate_names[1]].values + # Need to flip the coordinates because the names are in northing and + # easting order + coordinates = [i.ravel() for i in np.meshgrid(east, north)][::-1] + data_dict = dict(zip(coordinate_names, coordinates)) + data_dict.update(dict(zip(data_names, data_arrays))) + return pd.DataFrame(data_dict) def kdtree(coordinates, use_pykdtree=True, **kwargs):