From 3395fd2b6bf79486deeb680fce6ca008dfaa04ac Mon Sep 17 00:00:00 2001 From: Matthias Cuntz Date: Mon, 22 Mar 2021 00:29:44 +0100 Subject: [PATCH] New function add_cyclic that allows adding a cyclic point with n-dimensional longitudes and latitudes. It checks if a cyclic point is already present using the function has_cyclic. --- docs/source/contributors.rst | 1 + examples/scalar_data/wrapping_global.py | 91 ++++++ lib/cartopy/tests/test_util.py | 352 +++++++++++++++++++++++- lib/cartopy/util.py | 317 ++++++++++++++++++++- 4 files changed, 756 insertions(+), 5 deletions(-) create mode 100644 examples/scalar_data/wrapping_global.py diff --git a/docs/source/contributors.rst b/docs/source/contributors.rst index 2ab85ac60..769c0731f 100644 --- a/docs/source/contributors.rst +++ b/docs/source/contributors.rst @@ -42,6 +42,7 @@ the package wouldn't be as rich or diverse as it is today: * Kacper Makuch * Stephane Raynaud * John Krasting + * Matthias Cuntz Thank you! diff --git a/examples/scalar_data/wrapping_global.py b/examples/scalar_data/wrapping_global.py new file mode 100644 index 000000000..d4ef32b9c --- /dev/null +++ b/examples/scalar_data/wrapping_global.py @@ -0,0 +1,91 @@ +""" +Adding a cyclic point to help with wrapping of global data +========================================================== + +Cartopy represents data in Cartesian projected coordinates, meaning that 350 +degrees longitude, is not just 10 degrees away from 0 degrees as it is when +represented in spherical coordinates. This means that the plotting methods will +not plot data between the last and the first longitude. + +To help with this, the data and longitude/latitude coordinate arrays can be +expanded with a cyclic point to close this gap. The routine `add_cyclic` +repeats the last data column. It can also add the first longitude plus the +cyclic keyword (defaults to 360) to the end of the longitude array so that the +data values at the ending longitudes will be closed to the wrap point. + +""" +import numpy as np +import matplotlib.pyplot as plt +import cartopy.crs as ccrs +import cartopy.util as cutil + + +def main(): + + # data with longitude centers from 0 to 360 + nlon = 24 + nlat = 12 + # 7.5, 22.5, ..., 337.5, 352.5 + dlon = 360//nlon + lon = np.linspace(dlon/2., 360.-dlon/2., nlon) +# -82.5, -67.5, ..., 67.5, 82.5 + dlat = 180//nlat + lat = np.linspace(-90.+dlat/2., 90.-dlat/2., nlat) + # 0, 1, ..., 10, 11, 11, 10, ..., 1, 0 + data = np.concatenate((np.arange(nlon // 2), + np.arange(nlon // 2)[::-1])) + data = np.tile(data, nlat).reshape((nlat, nlon)) + + fig = plt.figure() + + # plot with central longitude 180 + ax1 = fig.add_subplot(2, 2, 1, + projection=ccrs.Robinson(central_longitude=180)) + ax1.set_title("1d longitudes, central longitude=180", + fontsize='small') + ax1.set_global() + ax1.contourf(lon, lat, data, + transform=ccrs.PlateCarree(), cmap='RdBu') + ax1.coastlines() + + # plot with central longitude 0 + ax2 = fig.add_subplot(2, 2, 2, + projection=ccrs.Robinson(central_longitude=0)) + ax2.set_title("1d longitudes, central longitude=0", + fontsize='small') + ax2.set_global() + ax2.contourf(lon, lat, data, + transform=ccrs.PlateCarree(), cmap='RdBu') + ax2.coastlines() + + # add cyclic points to data and longitudes + # latitudes are unchanged in 1-dimension + cdata, clon, clat = cutil.add_cyclic(data, lon, lat) + ax3 = fig.add_subplot(2, 2, 3, + projection=ccrs.Robinson(central_longitude=180)) + ax3.set_title("Cyclic 1d longitudes, central longitude=180", + fontsize='small') + ax3.set_global() + ax3.contourf(clon, clat, cdata, + transform=ccrs.PlateCarree(), cmap='RdBu') + ax3.coastlines() + + # add_cyclic also works with 2-dimensional data + # Cyclic points are added to data, longitudes, and latitudes to + # ensure the dimensions of the returned arrays are all the same shape. + lon2d, lat2d = np.meshgrid(lon, lat) + cdata, clon2d, clat2d = cutil.add_cyclic(data, lon2d, lat2d) + ax4 = fig.add_subplot(2, 2, 4, + projection=ccrs.Robinson(central_longitude=0)) + ax4.set_title("Cyclic 2d longitudes, central longitude=0", + fontsize='small') + ax4.set_global() + ax4.contourf(clon2d, clat2d, cdata, + transform=ccrs.PlateCarree(), cmap='RdBu') + ax4.coastlines() + + plt.show() + + +if __name__ == '__main__': + main() diff --git a/lib/cartopy/tests/test_util.py b/lib/cartopy/tests/test_util.py index 7e6492a8c..2ba6fab5b 100644 --- a/lib/cartopy/tests/test_util.py +++ b/lib/cartopy/tests/test_util.py @@ -9,7 +9,7 @@ from numpy.testing import assert_array_equal import pytest -from cartopy.util import add_cyclic_point +from cartopy.util import add_cyclic_point, add_cyclic, has_cyclic class Test_add_cyclic_point: @@ -64,3 +64,353 @@ def test_invalid_coord_size(self): def test_invalid_axis(self): with pytest.raises(ValueError): add_cyclic_point(self.data2d, axis=-3) + + +class TestAddCyclic: + """ + Test def add_cyclic(data, x=None, y=None, axis=-1, + cyclic=360, precision=1e-4): + - variations of data, x, and y with and without axis keyword + - different units of x - cyclic keyword + - detection of cyclic points - precision keyword + - error catching + """ + + @classmethod + def setup_class(cls): + # 2d and 4d data + cls.data2d = np.ones([3, 6]) * np.arange(6) + cls.data4d = np.ones([4, 6, 2, 3]) * \ + np.arange(6)[..., np.newaxis, np.newaxis] + # 1d lat (5) and lon (6) + # len(lat) != data.shape[0] + # len(lon) == data.shape[1] + cls.lons = np.arange(0, 360, 60) + cls.lats = np.arange(-90, 90, 180/5) + # 2d lat and lon + cls.lon2d, cls.lat2d = np.meshgrid(cls.lons, cls.lats) + # 3d lat and lon but with different 3rd dimension (4) as 4d data (2) + cls.lon3d = np.repeat(cls.lon2d, 4).reshape((*cls.lon2d.shape, 4)) + cls.lat3d = np.repeat(cls.lat2d, 4).reshape((*cls.lat2d.shape, 4)) + # cyclic data, lon, lat + cls.c_data2d = np.concatenate((cls.data2d, cls.data2d[:, :1]), axis=1) + cls.c_data4d = np.concatenate((cls.data4d, cls.data4d[:, :1]), axis=1) + cls.c_lons = np.concatenate((cls.lons, np.array([360.]))) + cls.c_lon2d = np.concatenate( + (cls.lon2d, np.full((cls.lon2d.shape[0], 1), 360.)), + axis=1) + cls.c_lon3d = np.concatenate( + (cls.lon3d, + np.full((cls.lon3d.shape[0], 1, cls.lon3d.shape[2]), 360.)), + axis=1) + cls.c_lats = cls.lats + cls.c_lat2d = np.concatenate((cls.lat2d, cls.lat2d[:, -1:]), axis=1) + cls.c_lat3d = np.concatenate((cls.lat3d, cls.lat3d[:, -1:, :]), axis=1) + + def test_data_only(self): + '''Test only data no x given''' + c_data = add_cyclic(self.data2d) + assert_array_equal(c_data, self.c_data2d) + + def test_data_only_ignore_y(self): + '''Test y given but no x''' + c_data = add_cyclic(self.data2d, y=self.lat2d) + assert_array_equal(c_data, self.c_data2d) + + def test_data_and_x_1d(self): + '''Test data 2d and x 1d''' + c_data, c_lons = add_cyclic(self.data2d, x=self.lons) + assert_array_equal(c_data, self.c_data2d) + assert_array_equal(c_lons, self.c_lons) + + def test_data_and_x_2d(self): + '''Test data and x 2d; no keyword name for x''' + c_data, c_lons = add_cyclic(self.data2d, self.lon2d) + assert_array_equal(c_data, self.c_data2d) + assert_array_equal(c_lons, self.c_lon2d) + + def test_data_and_x_y_1d(self): + '''Test data and x and y 1d''' + c_data, c_lons, c_lats = add_cyclic(self.data2d, x=self.lons, + y=self.lats) + assert_array_equal(c_data, self.c_data2d) + assert_array_equal(c_lons, self.c_lons) + assert_array_equal(c_lats, self.c_lats) + + def test_data_and_x_1d_y_2d(self): + '''Test data and x 1d and y 2d''' + c_data, c_lons, c_lats = add_cyclic(self.data2d, x=self.lons, + y=self.lat2d) + assert_array_equal(c_data, self.c_data2d) + assert_array_equal(c_lons, self.c_lons) + assert_array_equal(c_lats, self.c_lat2d) + + def test_data_and_x_y_2d(self): + '''Test data, x, and y 2d; no keyword name for x and y''' + c_data, c_lons, c_lats = add_cyclic(self.data2d, + self.lon2d, + self.lat2d) + assert_array_equal(c_data, self.c_data2d) + assert_array_equal(c_lons, self.c_lon2d) + assert_array_equal(c_lats, self.c_lat2d) + + def test_has_cyclic_1d(self): + '''Test detection of cyclic point 1d''' + c_data, c_lons = add_cyclic(self.c_data2d, x=self.c_lons) + assert_array_equal(c_data, self.c_data2d) + assert_array_equal(c_lons, self.c_lons) + + def test_has_cyclic_2d(self): + '''Test detection of cyclic point 2d''' + c_data, c_lons = add_cyclic(self.c_data2d, x=self.c_lon2d) + assert_array_equal(c_data, self.c_data2d) + assert_array_equal(c_lons, self.c_lon2d) + + def test_has_cyclic_2d_full(self): + '''Test detection of cyclic point 2d including y''' + c_data, c_lons, c_lats = add_cyclic(self.c_data2d, x=self.c_lon2d, + y=self.c_lat2d) + assert_array_equal(c_data, self.c_data2d) + assert_array_equal(c_lons, self.c_lon2d) + assert_array_equal(c_lats, self.c_lat2d) + + def test_data_only_with_axis(self): + '''Test axis keyword data only''' + c_data = add_cyclic(self.data4d, axis=1) + assert_array_equal(c_data, self.c_data4d) + + def test_data_and_x_with_axis_1d(self): + '''Test axis keyword data 4d, x 1d''' + c_data, c_lons = add_cyclic(self.data4d, x=self.lons, axis=1) + assert_array_equal(c_data, self.c_data4d) + assert_array_equal(c_lons, self.c_lons) + + def test_data_and_x_with_axis_2d(self): + '''Test axis keyword data 4d, x 2d''' + c_data, c_lons = add_cyclic(self.data4d, x=self.lon2d, + axis=1) + assert_array_equal(c_data, self.c_data4d) + assert_array_equal(c_lons, self.c_lon2d) + + def test_data_and_x_with_axis_3d(self): + '''Test axis keyword data 4d, x 3d''' + c_data, c_lons = add_cyclic(self.data4d, x=self.lon3d, + axis=1) + assert_array_equal(c_data, self.c_data4d) + assert_array_equal(c_lons, self.c_lon3d) + + def test_data_and_x_y_with_axis_2d(self): + '''Test axis keyword data 4d, x and y 2d''' + c_data, c_lons, c_lats = add_cyclic(self.data4d, + x=self.lon2d, + y=self.lat2d, + axis=1) + assert_array_equal(c_data, self.c_data4d) + assert_array_equal(c_lons, self.c_lon2d) + assert_array_equal(c_lats, self.c_lat2d) + + def test_data_and_x_y_with_axis_3d(self): + '''Test axis keyword data 4d, x and y 3d''' + c_data, c_lons, c_lats = add_cyclic(self.data4d, + x=self.lon3d, + y=self.lat3d, + axis=1) + assert_array_equal(c_data, self.c_data4d) + assert_array_equal(c_lons, self.c_lon3d) + assert_array_equal(c_lats, self.c_lat3d) + + def test_data_and_x_y_with_axis_nd(self): + '''Test axis keyword data 4d, x 3d and y 2d''' + c_data, c_lons, c_lats = add_cyclic(self.data4d, + x=self.lon3d, + y=self.lat2d, + axis=1) + assert_array_equal(c_data, self.c_data4d) + assert_array_equal(c_lons, self.c_lon3d) + assert_array_equal(c_lats, self.c_lat2d) + + def test_masked_data(self): + '''Test masked data''' + new_data = ma.masked_less(self.data2d, 3) + c_data = add_cyclic(new_data) + r_data = ma.concatenate((self.data2d, self.data2d[:, :1]), axis=1) + assert_array_equal(c_data, r_data) + assert ma.is_masked(c_data) + + def test_masked_data_and_x_y_2d(self): + '''Test masked data and x''' + new_data = ma.masked_less(self.data2d, 3) + new_lon = ma.masked_less(self.lon2d, 2) + c_data, c_lons, c_lats = add_cyclic(new_data, + x=new_lon, + y=self.lat2d) + r_data = ma.concatenate((self.data2d, self.data2d[:, :1]), axis=1) + r_lons = np.concatenate((self.lon2d, + np.full((self.lon2d.shape[0], 1), 360.)), + axis=1) + assert_array_equal(c_data, r_data) + assert_array_equal(c_lons, r_lons) + assert_array_equal(c_lats, self.c_lat2d) + assert ma.is_masked(c_data) + assert ma.is_masked(c_lons) + assert not ma.is_masked(c_lats) + + def test_cyclic(self): + '''Test cyclic keyword with axis data 4d, x 3d and y 2d''' + new_lons = np.deg2rad(self.lon3d) + new_lats = np.deg2rad(self.lat2d) + c_data, c_lons, c_lats = add_cyclic(self.data4d, x=new_lons, + y=new_lats, axis=1, + cyclic=np.deg2rad(360.)) + r_lons = np.concatenate( + (new_lons, + np.full((new_lons.shape[0], 1, new_lons.shape[2]), + np.deg2rad(360.))), + axis=1) + r_lats = np.concatenate((new_lats, new_lats[:, -1:]), axis=1) + assert_array_equal(c_data, self.c_data4d) + assert_array_equal(c_lons, r_lons) + assert_array_equal(c_lats, r_lats) + + def test_cyclic_has_cyclic(self): + '''Test detection of cyclic point with cyclic keyword''' + new_lons = np.deg2rad(self.lon2d) + new_lats = np.deg2rad(self.lat2d) + r_data = np.concatenate((self.data2d, self.data2d[:, :1]), axis=1) + r_lons = np.concatenate( + (new_lons, + np.full((new_lons.shape[0], 1), np.deg2rad(360.))), + axis=1) + r_lats = np.concatenate((new_lats, new_lats[:, -1:]), axis=1) + c_data, c_lons, c_lats = add_cyclic(r_data, x=r_lons, + y=r_lats, + + cyclic=np.deg2rad(360.)) + assert_array_equal(c_data, self.c_data2d) + assert_array_equal(c_lons, r_lons) + assert_array_equal(c_lats, r_lats) + + def test_precision_has_cyclic(self): + '''Test precision keyword detecting cyclic point''' + r_data = np.concatenate((self.data2d, self.data2d[:, :1]), axis=1) + r_lons = np.concatenate((self.lons, np.array([360.+1e-3]))) + c_data, c_lons = add_cyclic(r_data, x=r_lons, precision=1e-2) + assert_array_equal(c_data, r_data) + assert_array_equal(c_lons, r_lons) + + def test_precision_has_cyclic_no(self): + '''Test precision keyword detecting no cyclic point''' + new_data = np.concatenate((self.data2d, self.data2d[:, :1]), axis=1) + new_lons = np.concatenate((self.lons, np.array([360.+1e-3]))) + c_data, c_lons = add_cyclic(new_data, x=new_lons, precision=2e-4) + r_data = np.concatenate((new_data, new_data[:, :1]), axis=1) + r_lons = np.concatenate((new_lons, np.array([360.]))) + assert_array_equal(c_data, r_data) + assert_array_equal(c_lons, r_lons) + + def test_invalid_x_dimensionality(self): + '''Catch wrong x dimensions''' + with pytest.raises(ValueError): + c_data, c_lons = add_cyclic(self.data2d, x=self.lon3d) + + def test_invalid_y_dimensionality(self): + '''Catch wrong y dimensions''' + with pytest.raises(ValueError): + c_data, c_lons, c_lats = add_cyclic(self.data2d, + x=self.lon2d, + y=self.lat3d) + + def test_invalid_x_size_1d(self): + '''Catch wrong x size 1d''' + with pytest.raises(ValueError): + c_data, c_lons = add_cyclic(self.data2d, + x=self.lons[:-1]) + + def test_invalid_x_size_2d(self): + '''Catch wrong x size 2d''' + with pytest.raises(ValueError): + c_data, c_lons = add_cyclic(self.data2d, + x=self.lon2d[:, :-1]) + + def test_invalid_x_size_3d(self): + '''Catch wrong x size 3d''' + with pytest.raises(ValueError): + c_data, c_lons = add_cyclic(self.data4d, + x=self.lon3d[:, :-1, :], axis=1) + + def test_invalid_y_size(self): + '''Catch wrong y size 2d''' + with pytest.raises(ValueError): + c_data, c_lons, c_lats = add_cyclic( + self.data2d, x=self.lon2d, y=self.lat2d[:, 1:]) + + def test_invalid_axis(self): + '''Catch wrong axis keyword''' + with pytest.raises(ValueError): + add_cyclic(self.data2d, axis=-3) + + +class TestHasCyclic: + """ + Test def has_cyclic(x, axis=-1, cyclic=360, precision=1e-4): + - variations of x with and without axis keyword + - different unit of x - cyclic keyword + - detection of cyclic points - precision keyword + """ + + # 1d lon (6), lat (5) + lons = np.arange(0, 360, 60) + lats = np.arange(-90, 90, 180/5) + # 2d lon, lat + lon2d, lat2d = np.meshgrid(lons, lats) + # 3d lon + lon3d = np.repeat(lon2d, 4).reshape((*lon2d.shape, 4)) + # cyclic lon 1d, 2d, 3d + c_lons = np.concatenate((lons, np.array([360.]))) + c_lon2d = np.concatenate( + (lon2d, + np.full((lon2d.shape[0], 1), 360.)), + axis=1) + c_lon3d = np.concatenate( + (lon3d, + np.full((lon3d.shape[0], 1, lon3d.shape[2]), 360.)), + axis=1) + + @pytest.mark.parametrize( + "lon, clon", + [(lons, c_lons), + (lon2d, c_lon2d)]) + def test_data(self, lon, clon): + '''Test lon is not cyclic, clon is cyclic''' + assert not has_cyclic(lon) + assert has_cyclic(clon) + + @pytest.mark.parametrize( + "lon, clon, axis", + [(lons, c_lons, 0), + (lon2d, c_lon2d, 1), + (ma.masked_inside(lon2d, 100, 200), + ma.masked_inside(c_lon2d, 100, 200), + 1)]) + def test_data_axis(self, lon, clon, axis): + '''Test lon is not cyclic, clon is cyclic, with axis keyword''' + assert not has_cyclic(lon, axis=axis) + assert has_cyclic(clon, axis=axis) + + def test_3d_axis(self): + '''Test 3d with axis keyword, no keyword name for axis''' + assert has_cyclic(self.c_lon3d, 1) + assert not has_cyclic(self.lon3d, 1) + + def test_3d_axis_cyclic(self): + '''Test 3d with axis and cyclic keywords''' + new_clons = np.deg2rad(self.c_lon3d) + new_lons = np.deg2rad(self.lon3d) + assert has_cyclic(new_clons, axis=1, cyclic=np.deg2rad(360.)) + assert not has_cyclic(new_lons, axis=1, cyclic=np.deg2rad(360.)) + + def test_1d_precision(self): + '''Test 1d with precision keyword''' + new_clons = np.concatenate((self.lons, np.array([360.+1e-3]))) + assert has_cyclic(new_clons, precision=1e-2) + assert not has_cyclic(new_clons, precision=2e-4) diff --git a/lib/cartopy/util.py b/lib/cartopy/util.py index 46b5b78c9..e93e190b9 100644 --- a/lib/cartopy/util.py +++ b/lib/cartopy/util.py @@ -8,7 +8,6 @@ cartopy. """ - import numpy as np import numpy.ma as ma @@ -22,11 +21,11 @@ def add_cyclic_point(data, coord=None, axis=-1): ---------- data An n-dimensional array of data to add a cyclic point to. - coord: optional + coord : optional A 1-dimensional array which specifies the coordinate values for the dimension the cyclic point is to be added to. The coordinate values must be regularly spaced. Defaults to None. - axis: optional + axis : optional Specifies the axis of the data array to add the cyclic point to. Defaults to the right-most axis. @@ -41,7 +40,7 @@ def add_cyclic_point(data, coord=None, axis=-1): Examples -------- Adding a cyclic point to a data array, where the cyclic dimension is - the right-most dimension + the right-most dimension. >>> import numpy as np >>> data = np.ones([5, 6]) * np.arange(6) @@ -91,3 +90,313 @@ def add_cyclic_point(data, coord=None, axis=-1): else: return_value = new_data, new_coord return return_value + + +def _add_cyclic_data(data, axis=-1): + """ + Add a cyclic point to a data array. + + Parameters + ---------- + data : ndarray + An n-dimensional array of data to add a cyclic point to. + axis : int, optional + Specifies the axis of the data array to add the cyclic point to. + Defaults to the right-most axis. + + Returns + ------- + The data array with a cyclic point added. + + """ + slicer = [slice(None)] * data.ndim + try: + slicer[axis] = slice(0, 1) + except IndexError: + raise ValueError( + 'The specified axis does not correspond to an array dimension.') + npc = np.ma if np.ma.is_masked(data) else np + return npc.concatenate((data, data[tuple(slicer)]), axis=axis) + + +def _add_cyclic_x(x, axis=-1, cyclic=360): + """ + Add a cyclic point to a x/longitude coordinate array. + + Parameters + ---------- + x : ndarray + An array which specifies the x-coordinate values for + the dimension the cyclic point is to be added to. + axis : int, optional + Specifies the axis of the x-coordinate array to add the cyclic point + to. Defaults to the right-most axis. + cyclic : float, optional + Width of periodic domain (default: 360) + + Returns + ------- + The coordinate array ``x`` with a cyclic point added. + + """ + npc = np.ma if np.ma.is_masked(x) else np + # get cyclic x-coordinates + # cx is the code from basemap (addcyclic) + # https://github.com/matplotlib/basemap/blob/master/lib/mpl_toolkits/basemap/__init__.py + cx = (np.take(x, [0], axis=axis) + + cyclic * np.sign(np.diff(np.take(x, [0, -1], axis=axis), + axis=axis))) + # basemap ensures that the values do not exceed cyclic + # (next code line). We do not do this to deal with rotated grids that + # might have values not exactly 0. + # cx = npc.where(cx <= cyclic, cx, np.mod(cx, cyclic)) + return npc.concatenate((x, cx), axis=axis) + + +def has_cyclic(x, axis=-1, cyclic=360, precision=1e-4): + """ + Check if x/longitude coordinates already have a cyclic point. + + Checks all differences between the first and last + x-coordinates along ``axis`` to be less than ``precision``. + + Parameters + ---------- + x : ndarray + An array with the x-coordinate values to be checked for cyclic points. + axis : int, optional + Specifies the axis of the ``x`` array to be checked. + Defaults to the right-most axis. + cyclic : float, optional + Width of periodic domain (default: 360). + precision : float, optional + Maximal difference between first and last x-coordinate to detect + cyclic point (default: 1e-4). + + Returns + ------- + True if a cyclic point was detected along the given axis, + False otherwise. + + Examples + -------- + Check for cyclic x-coordinate in one dimension. + >>> import numpy as np + >>> lons = np.arange(0, 360, 60) + >>> clons = np.arange(0, 361, 60) + >>> print(has_cyclic(lons)) + False + >>> print(has_cyclic(clons)) + True + + Check for cyclic x-coordinate in two dimensions. + >>> lats = np.arange(-90, 90, 30) + >>> lon2d, lat2d = np.meshgrid(lons, lats) + >>> clon2d, clat2d = np.meshgrid(clons, lats) + >>> print(has_cyclic(lon2d)) + False + >>> print(has_cyclic(clon2d)) + True + + """ + npc = np.ma if np.ma.is_masked(x) else np + # transform to 0-cyclic, assuming e.g. -180 to 180 if any < 0 + x1 = np.mod(npc.where(x < 0, x + cyclic, x), cyclic) + dd = np.diff(np.take(x1, [0, -1], axis=axis), axis=axis) + if npc.all(np.abs(dd) < precision): + return True + else: + return False + + +def add_cyclic(data, x=None, y=None, axis=-1, + cyclic=360, precision=1e-4): + """ + Add a cyclic point to an array and optionally corresponding + x/longitude and y/latitude coordinates. + + Checks all differences between the first and last + x-coordinates along ``axis`` to be less than ``precision``. + + Parameters + ---------- + data : ndarray + An n-dimensional array of data to add a cyclic point to. + x : ndarray, optional + An n-dimensional array which specifies the x-coordinate values + for the dimension the cyclic point is to be added to, i.e. normally the + longitudes. Defaults to None. + + If ``x`` is given then *add_cyclic* checks if a cyclic point is + already present by checking all differences between the first and last + coordinates to be less than ``precision``. + No point is added if a cyclic point was detected. + + If ``x`` is 1- or 2-dimensional, ``x.shape[-1]`` must equal + ``data.shape[axis]``, otherwise ``x.shape[axis]`` must equal + ``data.shape[axis]``. + y : ndarray, optional + An n-dimensional array with the values of the y-coordinate, i.e. + normally the latitudes. + The cyclic point simply copies the last column. Defaults to None. + + No cyclic point is added if ``y`` is 1-dimensional. + If ``y`` is 2-dimensional, ``y.shape[-1]`` must equal + ``data.shape[axis]``, otherwise ``y.shape[axis]`` must equal + ``data.shape[axis]``. + axis : int, optional + Specifies the axis of the arrays to add the cyclic point to, + i.e. axis with changing x-coordinates. Defaults to the right-most axis. + cyclic : int or float, optional + Width of periodic domain (default: 360). + precision : float, optional + Maximal difference between first and last x-coordinate to detect + cyclic point (default: 1e-4). + + Returns + ------- + cyclic_data + The data array with a cyclic point added. + cyclic_x + The x-coordinate with a cyclic point, only returned if the ``x`` + keyword was supplied. + cyclic_y + The y-coordinate with the last column of the cyclic axis duplicated, + only returned if ``x`` was 2- or n-dimensional and the ``y`` + keyword was supplied. + + Examples + -------- + Adding a cyclic point to a data array, where the cyclic dimension is + the right-most dimension. + >>> import numpy as np + >>> data = np.ones([5, 6]) * np.arange(6) + >>> cyclic_data = add_cyclic(data) + >>> print(cyclic_data) # doctest: +NORMALIZE_WHITESPACE + [[0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.]] + + Adding a cyclic point to a data array and an associated x-coordinate. + >>> lons = np.arange(0, 360, 60) + >>> cyclic_data, cyclic_lons = add_cyclic(data, x=lons) + >>> print(cyclic_data) # doctest: +NORMALIZE_WHITESPACE + [[0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.]] + >>> print(cyclic_lons) + [ 0 60 120 180 240 300 360] + + Adding a cyclic point to a data array and an associated 2-dimensional + x-coordinate. + >>> lons = np.arange(0, 360, 60) + >>> lats = np.arange(-90, 90, 180/5) + >>> lon2d, lat2d = np.meshgrid(lons, lats) + >>> cyclic_data, cyclic_lon2d = add_cyclic(data, x=lon2d) + >>> print(cyclic_data) # doctest: +NORMALIZE_WHITESPACE + [[0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.]] + >>> print(cyclic_lon2d) + [[ 0 60 120 180 240 300 360] + [ 0 60 120 180 240 300 360] + [ 0 60 120 180 240 300 360] + [ 0 60 120 180 240 300 360] + [ 0 60 120 180 240 300 360]] + + Adding a cyclic point to a data array and the associated 2-dimensional + x- and y-coordinates. + >>> lons = np.arange(0, 360, 60) + >>> lats = np.arange(-90, 90, 180/5) + >>> lon2d, lat2d = np.meshgrid(lons, lats) + >>> cyclic_data, cyclic_lon2d, cyclic_lat2d = add_cyclic( + ... data, x=lon2d, y=lat2d) + >>> print(cyclic_data) # doctest: +NORMALIZE_WHITESPACE + [[0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.] + [0. 1. 2. 3. 4. 5. 0.]] + >>> print(cyclic_lon2d) + [[ 0 60 120 180 240 300 360] + [ 0 60 120 180 240 300 360] + [ 0 60 120 180 240 300 360] + [ 0 60 120 180 240 300 360] + [ 0 60 120 180 240 300 360]] + >>> print(cyclic_lat2d) + [[-90. -90. -90. -90. -90. -90. -90.] + [-54. -54. -54. -54. -54. -54. -54.] + [-18. -18. -18. -18. -18. -18. -18.] + [ 18. 18. 18. 18. 18. 18. 18.] + [ 54. 54. 54. 54. 54. 54. 54.]] + + Not adding a cyclic point if cyclic point detected in x. + >>> lons = np.arange(0, 361, 72) + >>> lats = np.arange(-90, 90, 180/5) + >>> lon2d, lat2d = np.meshgrid(lons, lats) + >>> cyclic_data, cyclic_lon2d, cyclic_lat2d = add_cyclic( + ... data, x=lon2d, y=lat2d) + >>> print(cyclic_data) # doctest: +NORMALIZE_WHITESPACE + [[0. 1. 2. 3. 4. 5.] + [0. 1. 2. 3. 4. 5.] + [0. 1. 2. 3. 4. 5.] + [0. 1. 2. 3. 4. 5.] + [0. 1. 2. 3. 4. 5.]] + >>> print(cyclic_lon2d) + [[ 0 72 144 216 288 360] + [ 0 72 144 216 288 360] + [ 0 72 144 216 288 360] + [ 0 72 144 216 288 360] + [ 0 72 144 216 288 360]] + >>> print(cyclic_lat2d) + [[-90. -90. -90. -90. -90. -90.] + [-54. -54. -54. -54. -54. -54.] + [-18. -18. -18. -18. -18. -18.] + [ 18. 18. 18. 18. 18. 18.] + [ 54. 54. 54. 54. 54. 54.]] + + """ + if x is None: + return _add_cyclic_data(data, axis=axis) + # if x was given + if x.ndim > 2: + xaxis = axis + else: + xaxis = -1 + if x.shape[xaxis] != data.shape[axis]: + estr = (f'x.shape[{xaxis}] does not match the size of the' + f' corresponding dimension of the data array:' + f' x.shape[{xaxis}] = {x.shape[xaxis]},' + f' data.shape[{axis}] = {data.shape[axis]}.') + raise ValueError(estr) + if has_cyclic(x, axis=xaxis, cyclic=cyclic, precision=precision): + if y is None: + return data, x + # if y was given + return data, x, y + # if not has_cyclic, add cyclic points to data and x + out_data = _add_cyclic_data(data, axis=axis) + out_x = _add_cyclic_x(x, axis=xaxis, cyclic=cyclic) + if y is None: + return out_data, out_x + # if y was given + if y.ndim == 1: + return out_data, out_x, y + if y.ndim > 2: + yaxis = axis + else: + yaxis = -1 + if y.shape[yaxis] != data.shape[axis]: + estr = (f'y.shape[{yaxis}] does not match the size of the' + f' corresponding dimension of the data array:' + f' y.shape[{yaxis}] = {y.shape[yaxis]},' + f' data.shape[{axis}] = {data.shape[axis]}.') + raise ValueError(estr) + out_y = _add_cyclic_data(y, axis=yaxis) + return out_data, out_x, out_y