diff --git a/lib/iris/analysis/_grid_angles.py b/lib/iris/analysis/_grid_angles.py new file mode 100644 index 0000000000..1746834a65 --- /dev/null +++ b/lib/iris/analysis/_grid_angles.py @@ -0,0 +1,390 @@ +# (C) British Crown Copyright 2010 - 2018, Met Office +# +# This file is part of Iris. +# +# Iris is free software: you can redistribute it and/or modify it under +# the terms of the GNU Lesser General Public License as published by the +# Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# Iris is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with Iris. If not, see . +""" +Code to implement vector rotation by angles, and inferring gridcell angles +from coordinate points and bounds. + +""" +from __future__ import (absolute_import, division, print_function) +from six.moves import (filter, input, map, range, zip) # noqa + +import numpy as np + +import iris + + +def _3d_xyz_from_latlon(lon, lat): + """ + Return locations of (lon, lat) in 3D space. + + Args: + + * lon, lat: (arrays in degrees) + + Returns: + + xyz : (array, dtype=float64) + cartesian coordinates on a unit sphere. Dimension 0 maps x,y,z. + + """ + lon1 = np.deg2rad(lon).astype(np.float64) + lat1 = np.deg2rad(lat).astype(np.float64) + + x = np.cos(lat1) * np.cos(lon1) + y = np.cos(lat1) * np.sin(lon1) + z = np.sin(lat1) + + result = np.concatenate([array[np.newaxis] for array in (x, y, z)]) + + return result + + +def _latlon_from_xyz(xyz): + """ + Return arrays of lons+lats angles from xyz locations. + + Args: + + * xyz: (array) + positions array, of dims (3, ), where index 0 maps x/y/z. + + Returns: + + lonlat : (array) + spherical angles, of dims (2, ), in radians. + Dim 0 maps longitude, latitude. + + """ + lons = np.arctan2(xyz[1], xyz[0]) + axial_radii = np.sqrt(xyz[0] * xyz[0] + xyz[1] * xyz[1]) + lats = np.arctan2(xyz[2], axial_radii) + return np.array([lons, lats]) + + +def _angle(p, q, r): + """ + Return angle (in _radians_) of grid wrt local east. + Anticlockwise +ve, as usual. + {P, Q, R} are consecutive points in the same row, + eg {v(i,j),f(i,j),v(i+1,j)}, or {T(i-1,j),T(i,j),T(i+1,j)} + Calculate dot product of PR with lambda_hat at Q. + This gives us cos(required angle). + Disciminate between +/- angles by comparing latitudes of P and R. + p, q, r, are all 2-element arrays [lon, lat] of angles in degrees. + + """ +# old_style = True + old_style = False + if old_style: + mid_lons = np.deg2rad(q[0]) + + pr = _3d_xyz_from_latlon(r[0], r[1]) - _3d_xyz_from_latlon(p[0], p[1]) + pr_norm = np.sqrt(np.sum(pr**2, axis=0)) + pr_top = pr[1] * np.cos(mid_lons) - pr[0] * np.sin(mid_lons) + + index = pr_norm == 0 + pr_norm[index] = 1 + + cosine = np.maximum(np.minimum(pr_top / pr_norm, 1), -1) + cosine[index] = 0 + + psi = np.arccos(cosine) * np.sign(r[1] - p[1]) + psi[index] = np.nan + else: + # Calculate unit vectors. + midpt_lons, midpt_lats = q[0], q[1] + lmb_r, phi_r = (np.deg2rad(arr) for arr in (midpt_lons, midpt_lats)) + phi_hatvec_x = -np.sin(phi_r) * np.cos(lmb_r) + phi_hatvec_y = -np.sin(phi_r) * np.sin(lmb_r) + phi_hatvec_z = np.cos(phi_r) + shape_xyz = (1,) + midpt_lons.shape + phi_hatvec = np.concatenate([arr.reshape(shape_xyz) + for arr in (phi_hatvec_x, + phi_hatvec_y, + phi_hatvec_z)]) + lmb_hatvec_z = np.zeros(midpt_lons.shape) + lmb_hatvec_y = np.cos(lmb_r) + lmb_hatvec_x = -np.sin(lmb_r) + lmb_hatvec = np.concatenate([arr.reshape(shape_xyz) + for arr in (lmb_hatvec_x, + lmb_hatvec_y, + lmb_hatvec_z)]) + + pr = _3d_xyz_from_latlon(r[0], r[1]) - _3d_xyz_from_latlon(p[0], p[1]) + + # Dot products to form true-northward / true-eastward projections. + pr_cmpt_e = np.sum(pr * lmb_hatvec, axis=0) + pr_cmpt_n = np.sum(pr * phi_hatvec, axis=0) + psi = np.arctan2(pr_cmpt_n, pr_cmpt_e) + + # TEMPORARY CHECKS: + # ensure that the two unit vectors are perpendicular. + dotprod = np.sum(phi_hatvec * lmb_hatvec, axis=0) + assert np.allclose(dotprod, 0.0) + # ensure that the vector components carry the original magnitude. + mag_orig = np.sum(pr * pr) + mag_rot = np.sum(pr_cmpt_e * pr_cmpt_e) + np.sum(pr_cmpt_n * pr_cmpt_n) + rtol = 1.e-3 + check = np.allclose(mag_rot, mag_orig, rtol=rtol) + if not check: + print (mag_rot, mag_orig) + assert np.allclose(mag_rot, mag_orig, rtol=rtol) + + return psi + + +def gridcell_angles(x, y=None, cell_angle_boundpoints='mid-lhs, mid-rhs'): + """ + Calculate gridcell orientations for an arbitrary 2-dimensional grid. + + The input grid is defined by two 2-dimensional coordinate arrays with the + same dimensions (ny, nx), specifying the geolocations of a 2D mesh. + + Input values may be coordinate points (ny, nx) or bounds (ny, nx, 4). + However, if points, the edges in the X direction are assumed to be + connected by wraparound. + + Input can be either two arrays, two coordinates, or a single cube + containing two suitable coordinates identified with the 'x' and'y' axes. + + Args: + + The inputs (x [,y]) can be any of the folliwing : + + * x (:class:`~iris.cube.Cube`): + a grid cube with 2D X and Y coordinates, identified by 'axis'. + The coordinates must be 2-dimensional with the same shape. + The two dimensions represent grid dimensions in the order Y, then X. + + * x, y (:class:`~iris.coords.Coord`): + X and Y coordinates, specifying grid locations on the globe. + The coordinates must be 2-dimensional with the same shape. + The two dimensions represent grid dimensions in the order Y, then X. + If there is no coordinate system, they are assumed to be true + longitudes and latitudes. Units must convertible to 'degrees'. + + * x, y (2-dimensional arrays of same shape (ny, nx)): + longitude and latitude cell center locations, in degrees. + The two dimensions represent grid dimensions in the order Y, then X. + + * x, y (3-dimensional arrays of same shape (ny, nx, 4)): + longitude and latitude cell bounds, in degrees. + The first two dimensions are grid dimensions in the order Y, then X. + The last index maps cell corners anticlockwise from bottom-left. + + Optional Args: + + * cell_angle_boundpoints (string): + Controls which gridcell bounds locations are used to calculate angles, + if the inputs are bounds or bounded coordinates. + Valid values are 'lower-left, lower-right', which takes the angle from + the lower left to the lower right corner, and 'mid-lhs, mid-rhs' which + takes an angles between the average of the left-hand and right-hand + pairs of corners. The default is 'mid-lhs, mid-rhs'. + + Returns: + + angles : (2-dimensional cube) + + Cube of angles of grid-x vector from true Eastward direction for + each gridcell, in radians. + It also has longitude and latitude coordinates. If coordinates + were input the output has identical ones : If the input was 2d + arrays, the output coords have no bounds; or, if the input was 3d + arrays, the output coords have bounds and centrepoints which are + the average of the 4 bounds. + + """ + cube = None + if hasattr(x, 'add_aux_coord'): + # Passed a cube : extract 'x' and ;'y' axis coordinates. + cube = x # Save for later checking. + x, y = cube.coord(axis='x'), cube.coord(axis='y') + + # Now should have either 2 coords or 2 arrays. + if not hasattr(x, 'shape') and hasattr(y, 'shape'): + msg = ('Inputs (x,y) must have array shape property.' + 'Got type(x)={} and type(y)={}.') + raise ValueError(msg.format(type(x), type(y))) + + x_coord, y_coord = None, None + if hasattr(x, 'bounds') and hasattr(y, 'bounds'): + x_coord, y_coord = x.copy(), y.copy() + x_coord.convert_units('degrees') + y_coord.convert_units('degrees') + if x_coord.ndim != 2 or y_coord.ndim != 2: + msg = ('Coordinate inputs must have 2-dimensional shape. ', + 'Got x-shape of {} and y-shape of {}.') + raise ValueError(msg.format(x_coord.shape, y_coord.shape)) + if x_coord.shape != y_coord.shape: + msg = ('Coordinate inputs must have same shape. ', + 'Got x-shape of {} and y-shape of {}.') + raise ValueError(msg.format(x_coord.shape, y_coord.shape)) +# NOTE: would like to check that dims are in correct order, but can't do that +# if there is no cube. +# TODO: **document** -- another input format requirement +# x_dims, y_dims = (cube.coord_dims(co) for co in (x_coord, y_coord)) +# if x_dims != (0, 1) or y_dims != (0, 1): +# msg = ('Coordinate inputs must map to cube dimensions (0, 1). ', +# 'Got x-dims of {} and y-dims of {}.') +# raise ValueError(msg.format(x_dims, y_dims)) + if x_coord.has_bounds() and y_coord.has_bounds(): + x, y = x_coord.bounds, y_coord.bounds + else: + x, y = x_coord.points, y_coord.points + + elif hasattr(x, 'bounds') or hasattr(y, 'bounds'): + # One was a Coord, and the other not ? + is_and_not = ('x', 'y') + if hasattr(y, 'bounds'): + is_and_not = reversed(is_and_not) + msg = 'Input {!r} is a Coordinate, but {!r} is not.' + raise ValueError(*is_and_not) + + # Now have either 2 points arrays or 2 bounds arrays. + # Construct (lhs, mid, rhs) where these represent 3 adjacent points with + # increasing longitudes. + if x.ndim == 2: + # PROBLEM: we can't use this if data is not full-longitudes, + # i.e. rhs of array must connect to lhs (aka 'circular' coordinate). + # But we have no means of checking that ? + + # Use previous + subsequent points along longitude-axis as references. + # NOTE: we also have no way to check that dim #2 really is the 'X' dim. + mid = np.array([x, y]) + lhs = np.roll(mid, 1, 2) + rhs = np.roll(mid, -1, 2) + if not x_coord: + # Create coords for result cube : with no bounds. + y_coord = iris.coords.AuxCoord(x, standard_name='latitude', + units='degrees') + x_coord = iris.coords.AuxCoord(y, standard_name='longitude', + units='degrees') + else: + # Get lhs and rhs locations by averaging top+bottom each side. + # NOTE: so with bounds, we *don't* need full circular longitudes. + xyz = _3d_xyz_from_latlon(x, y) + angle_boundpoints_vals = {'mid-lhs, mid-rhs': '03_to_12', + 'lower-left, lower-right': '0_to_1'} + bounds_pos = angle_boundpoints_vals.get(cell_angle_boundpoints) + if bounds_pos == '0_to_1': + lhs_xyz = xyz[..., 0] + rhs_xyz = xyz[..., 1] + elif bounds_pos == '03_to_12': + lhs_xyz = 0.5 * (xyz[..., 0] + xyz[..., 3]) + rhs_xyz = 0.5 * (xyz[..., 1] + xyz[..., 2]) + else: + msg = ('unrecognised cell_angle_boundpoints of "{}", ' + 'must be one of {}') + raise ValueError(msg.format(cell_angle_boundpoints, + list(angle_boundpoints_vals.keys()))) + if not x_coord: + # Create bounded coords for result cube. + # Use average lhs+rhs points in 3d to get 'mid' points, as coords + # with no points are not allowed. + mid_xyz = 0.5 * (lhs_xyz + rhs_xyz) + mid_latlons = _latlon_from_xyz(mid_xyz) + # Create coords with given bounds, and averaged centrepoints. + x_coord = iris.coords.AuxCoord( + points=mid_latlons[0], bounds=x, + standard_name='longitude', units='degrees') + y_coord = iris.coords.AuxCoord( + points=mid_latlons[1], bounds=y, + standard_name='latitude', units='degrees') + # Convert lhs and rhs points back to latlon form -- IN DEGREES ! + lhs = np.rad2deg(_latlon_from_xyz(lhs_xyz)) + rhs = np.rad2deg(_latlon_from_xyz(rhs_xyz)) + # mid is coord.points, whether input or made up. + mid = np.array([x_coord.points, y_coord.points]) + + # Do the angle calcs, and return as a suitable cube. + angles = _angle(lhs, mid, rhs) + result = iris.cube.Cube(angles, + long_name='gridcell_angle_from_true_east', + units='radians') + result.add_aux_coord(x_coord, (0, 1)) + result.add_aux_coord(y_coord, (0, 1)) + return result + + +def true_vectors_from_grid_vectors(u_cube, v_cube, + grid_angles_cube=None, + grid_angles_kwargs=None): + """ + Rotate distance vectors from grid-oriented to true-latlon-oriented. + + .. Note:: + + This operation overlaps somewhat in function with + :func:`iris.analysis.cartography.rotate_winds`. + However, that routine only rotates vectors according to transformations + between coordinate systems. + This function, by contrast, can rotate vectors by arbitrary angles. + Most commonly, the angles are estimated solely from grid sampling + points, using :func:`gridcell_angles` : This allows operation on + complex meshes defined by two-dimensional coordinates, such as most + ocean grids. + + Args: + + * u_cube, v_cube : (cube) + Cubes of grid-u and grid-v vector components. + Units should be differentials of true-distance, e.g. 'm/s'. + + Optional args: + + * grid_angles_cube : (cube) + gridcell orientation angles. + Units must be angular, i.e. can be converted to 'radians'. + If not provided, grid angles are estimated from 'u_cube' using the + :func:`gridcell_angles` method. + + * grid_angles_kwargs : (dict or None) + Additional keyword args to be passed to the :func:`gridcell_angles` + method, if it is used. + + Returns: + + true_u, true_v : (cube) + Cubes of true-north oriented vector components. + Units are same as inputs. + + .. Note:: + + Vector magnitudes will always be the same as the inputs. + + """ + u_out, v_out = (cube.copy() for cube in (u_cube, v_cube)) + if not grid_angles_cube: + grid_angles_kwargs = grid_angles_kwargs or {} + grid_angles_cube = gridcell_angles(u_cube, **grid_angles_kwargs) + gridangles = grid_angles_cube.copy() + gridangles.convert_units('radians') + uu, vv, aa = (cube.data for cube in (u_out, v_out, gridangles)) + mags = np.sqrt(uu*uu + vv*vv) + angs = np.arctan2(vv, uu) + aa + uu, vv = mags * np.cos(angs), mags * np.sin(angs) + + # Promote all to masked arrays, and also apply mask at bad (NaN) angles. + mask = np.isnan(aa) + for cube in (u_out, v_out, aa): + if hasattr(cube.data, 'mask'): + mask |= cube.data.mask + u_out.data = np.ma.masked_array(uu, mask=mask) + v_out.data = np.ma.masked_array(vv, mask=mask) + + return u_out, v_out diff --git a/lib/iris/analysis/cartography.py b/lib/iris/analysis/cartography.py index 4b68dcc949..16c99c83a8 100644 --- a/lib/iris/analysis/cartography.py +++ b/lib/iris/analysis/cartography.py @@ -1,4 +1,4 @@ -# (C) British Crown Copyright 2010 - 2017, Met Office +# (C) British Crown Copyright 2010 - 2018, Met Office # # This file is part of Iris. # @@ -37,6 +37,9 @@ import iris.coord_systems import iris.exceptions from iris.util import _meshgrid +from ._grid_angles import ( + gridcell_angles, + true_vectors_from_grid_vectors as rotate_grid_vectors) # This value is used as a fall-back if the cube does not define the earth diff --git a/lib/iris/tests/unit/analysis/cartography/test_gridcell_angles.py b/lib/iris/tests/unit/analysis/cartography/test_gridcell_angles.py new file mode 100644 index 0000000000..de00d2bc76 --- /dev/null +++ b/lib/iris/tests/unit/analysis/cartography/test_gridcell_angles.py @@ -0,0 +1,392 @@ +# (C) British Crown Copyright 2018, Met Office +# +# This file is part of Iris. +# +# Iris is free software: you can redistribute it and/or modify it under +# the terms of the GNU Lesser General Public License as published by the +# Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# Iris is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with Iris. If not, see . +""" +Unit tests for the function +:func:`iris.analysis.cartography.gridcell_angles`. + +""" +from __future__ import (absolute_import, division, print_function) +from six.moves import (filter, input, map, range, zip) # noqa + +# Import iris.tests first so that some things can be initialised before +# importing anything else. +import iris.tests as tests + +import numpy as np +import numpy.ma as ma + +import cartopy.crs as ccrs +from iris.cube import Cube +from iris.coords import DimCoord, AuxCoord +import iris.coord_systems +from iris.analysis.cartography import unrotate_pole + +from iris.analysis.cartography import (gridcell_angles, + rotate_grid_vectors) + +import matplotlib.pyplot as plt +from orca_utils.plot_testing.blockplot_from_bounds import blockplot_2dll + + +def _rotated_grid_sample(pole_lat=15, pole_lon=-180, + lon_bounds=np.linspace(-30, 30, 6, endpoint=True), + lat_bounds=np.linspace(-30, 30, 6, endpoint=True)): + # Calculate *true* lat_bounds+lon_bounds for the rotated grid. + lon_bounds = np.array(lon_bounds, dtype=float) + lat_bounds = np.array(lat_bounds, dtype=float) + # Construct centrepoints. + lons = 0.5 * (lon_bounds[:-1] + lon_bounds[1:]) + lats = 0.5 * (lat_bounds[:-1] + lat_bounds[1:]) + # Convert all to full 2d arrays. + lon_bounds, lat_bounds = np.meshgrid(lon_bounds, lat_bounds) + lons, lats = np.meshgrid(lons, lats) + # Calculate true lats+lons for all points. + lons_true_bds, lats_true_bds = unrotate_pole(lon_bounds, lat_bounds, + pole_lon, pole_lat) + lons_true, lats_true = unrotate_pole(lons, lats, pole_lon, pole_lat) + # Make the 'unified' bounds into contiguous (ny, nx, 4) arrays. + def expand_unified_bds(bds): + ny, nx = bds.shape + bds_4 = np.zeros((ny - 1, nx - 1, 4)) + bds_4[:, :, 0] = bds[:-1, :-1] + bds_4[:, :, 1] = bds[:-1, 1:] + bds_4[:, :, 2] = bds[1:, 1:] + bds_4[:, :, 3] = bds[1:, :-1] + return bds_4 + + lon_true_bds4, lat_true_bds4 = (expand_unified_bds(bds) + for bds in (lons_true_bds, lats_true_bds)) + # Make these into a 2d-latlon grid for a cube + cube = Cube(np.zeros(lon_true_bds4.shape[:-1])) + co_x = AuxCoord(lons_true, bounds=lon_true_bds4, + standard_name='longitude', units='degrees') + co_y = AuxCoord(lats_true, bounds=lat_true_bds4, + standard_name='latitude', units='degrees') + cube.add_aux_coord(co_x, (0, 1)) + cube.add_aux_coord(co_y, (0, 1)) + return cube + + +class TestGridcellAngles(tests.IrisTest): + def _singlecell_30deg_cube(self, x0=90., y0=0., dx=20., dy=10.): + x_pts = np.array([[x0]]) + y_pts = np.array([[y0]]) + x_bds = x0 + dx * np.array([[[-1., 1, 0.5, -1.5]]]) +# self.assertArrayAllClose(x_bds, np.array([[[70., 110, 100, 60]]])) + y_bds = y0 + dy * np.array([[[-1., 1, 3, 1]]]) +# self.assertArrayAllClose(y_bds, np.array([[[-10., 10, 30, 10]]])) + co_x = AuxCoord(points=x_pts, bounds=x_bds, + standard_name='longitude', units='degrees') + co_y = AuxCoord(points=y_pts, bounds=y_bds, + standard_name='latitude', units='degrees') + cube = Cube(np.zeros((1, 1))) + cube.add_aux_coord(co_x, (0, 1)) + cube.add_aux_coord(co_y, (0, 1)) + return cube + + def _singlecell_diamond_cube(self, x0=90., y0=0., dy=10., dx_eq=None): + if dx_eq is None: + dx_eq = dy + x_pts = np.array([[x0]]) + y_pts = np.array([[y0]]) + dx = dx_eq / np.cos(np.deg2rad(y0)) + x_bds = np.array([[[x0, x0 + dx, x0, x0 - dx]]]) + y_bds = np.array([[[y0 - dy, y0, y0 + dy, y0]]]) + co_x = AuxCoord(points=x_pts, bounds=x_bds, + standard_name='longitude', units='degrees') + co_y = AuxCoord(points=y_pts, bounds=y_bds, + standard_name='latitude', units='degrees') + cube = Cube(np.zeros((1, 1))) + cube.add_aux_coord(co_x, (0, 1)) + cube.add_aux_coord(co_y, (0, 1)) + return cube + + def test_single_cell_equatorial(self): + plt.switch_backend('tkagg') + plt.figure(figsize=(10,10)) + ax = plt.axes(projection=ccrs.Mercator()) +# ax = plt.axes(projection=ccrs.NorthPolarStereo()) +# ax = plt.axes(projection=ccrs.Orthographic(central_longitude=90., +# central_latitude=30.)) + + lon0 = 90.0 + dy = 2.0 + y_0, y_n, ny = -80, 80, 9 + angles = [] + for lat in np.linspace(y_0, y_n, ny): + cube = self._singlecell_diamond_cube(x0=lon0, y0=lat, dy=dy) + angles_cube = gridcell_angles(cube, +# cell_angle_boundpoints='mid-lhs, mid-rhs') + cell_angle_boundpoints='lower-left, lower-right') + tmp_cube = angles_cube.copy() + tmp_cube.convert_units('degrees') + print('') + print(lat) + co_x, co_y = (cube.coord(axis=ax) for ax in ('x', 'y')) + print() + print(' at : {}, {}'.format(co_x.points[0, 0], co_y.points[0, 0])) + print(' x-bds:') + print(co_x.bounds) + print(' y-bds:') + print(co_y.bounds) + angle = tmp_cube.data[0, 0] + angles.append(angle) + print(angle) + blockplot_2dll(cube) + + ax.coastlines() + ax.set_global() + + # Plot constant NEly (45deg) arrows. + xx = np.array([lon0] * ny) + yy = np.linspace(y_0, y_n, ny) - dy + uu = np.array([1.0] * ny) + plt.quiver(xx, yy, + uu, np.cos(np.deg2rad(yy)), + zorder=2, color='red', + scale_units='xy', + transform=ccrs.PlateCarree()) + + # Also plot returned angles. + angles_arr_rad = np.deg2rad(angles) + u_arr = uu * np.cos(angles_arr_rad) + v_arr = uu * np.sin(angles_arr_rad) * np.cos(np.deg2rad(yy)) + + plt.quiver(xx, yy, + u_arr, + v_arr, + zorder=2, color='magenta', + scale_units='xy', + width=0.005, + scale=0.2e-6, +# width=0.5, + transform=ccrs.PlateCarree()) + + plt.show() + + + def test_values(self): + # Construct a rotated-pole grid and check angle calculation. + testcube = _rotated_grid_sample() + + cell_angle_boundpoints = 'mid-lhs, mid-rhs' +# cell_angle_boundpoints = 'lower-left, lower-right' +# cell_angle_boundpoints = 'garble' + angles_cube = gridcell_angles( + testcube, + cell_angle_boundpoints=cell_angle_boundpoints) + angles_cube.convert_units('radians') + + # testing phase... + print(np.rad2deg(angles_cube.data)) + + import matplotlib.pyplot as plt + plt.switch_backend('tkagg') + +# plot_map = 'north_polar_stereographic' +# plot_map = 'plate_carree' +# plot_map = 'mercator' + plot_map = 'north_polar_orthographic' + if plot_map == 'plate_carree': + scale = 0.1 + map_proj = ccrs.PlateCarree() + elif plot_map == 'mercator': + scale = 3.0e-6 + map_proj = ccrs.Mercator() + map_proj._threshold *= 0.01 + elif plot_map == 'north_polar_orthographic': + scale = 3.0e-6 + map_proj = ccrs.Orthographic(central_longitude=0.0, + central_latitude=90.0,) + map_proj._threshold *= 0.01 + elif plot_map == 'north_polar_stereographic': + scale = 3.0e-6 + map_proj = ccrs.NorthPolarStereo() + else: + assert 0 + + ax = plt.axes(projection=map_proj) + data_proj = ccrs.PlateCarree() + + deg_scale = 10.0 + +# angles = 'uv' + angles = 'xy' + + ax.coastlines() + ax.gridlines() + for i_bnd in range(4): + color = ['black', 'red', 'blue', 'magenta'][i_bnd] + plt.plot(testcube.coord('longitude').bounds[..., i_bnd], + testcube.coord('latitude').bounds[..., i_bnd], + '+', markersize=10., markeredgewidth=2., + markerfacecolor=color, markeredgecolor=color, + transform=data_proj) + + + # Show plain 0,1 + 1,0 (PlateCarree) vectors unrotated at the given points. + pts_shape = testcube.coord('longitude').shape + ny, nx = pts_shape + u0 = np.ones(pts_shape) + v0 = np.zeros(pts_shape) + u1 = v0.copy() + v1 = u0.copy() + + x0s = testcube.coord('longitude').points + y0s = testcube.coord('latitude').points + yscale = np.cos(np.deg2rad(y0s)) + plt.quiver(x0s, y0s, u0, v0 * yscale, + color='blue', width=0.005, + headwidth=2., # headlength=1.0, headaxislength=0.7, + angles=angles, + scale_units='xy', scale=scale, + transform=data_proj) + plt.quiver(x0s, y0s, u1, v1 * yscale, + color='red', width=0.005, + headwidth=2., # headlength=1.0, headaxislength=0.7, + angles=angles, + scale_units='xy', scale=scale, + transform=data_proj) + + # Add 45deg arrows (NEly), still on a PlateCarree map. + plt.quiver(x0s, y0s, v1, v1 * yscale, + color='green', width=0.005, + headwidth=2., # headlength=1.0, headaxislength=0.7, + angles=angles, + scale_units='xy', scale=scale, + transform=data_proj) + + + + # + # Repeat the above plotting short lines INSTEAD of quiver. + # + u0d = x0s + deg_scale * u0 + v0d = y0s + deg_scale * v0 + u1d = x0s + deg_scale * u1 + v1d = y0s + deg_scale * v1 + u2d = x0s + deg_scale * u0 + v2d = y0s + deg_scale * v1 + for iy in range(ny): + for ix in range(nx): + plt.plot([x0s[iy, ix], u0d[iy, ix]], + [y0s[iy, ix], v0d[iy, ix]], + ':', color='blue', linewidth=0.5, + transform=data_proj) + plt.plot([x0s[iy, ix], u1d[iy, ix]], + [y0s[iy, ix], v1d[iy, ix]], + ':', color='red', linewidth=0.5, + transform=data_proj) + plt.plot([x0s[iy, ix], u2d[iy, ix]], + [y0s[iy, ix], v2d[iy, ix]], + ':', color='green', linewidth=0.5, + transform=data_proj) + + + # Overplot BL-BR and BL-TL lines from the cell bounds. + co_lon, co_lat = [testcube.coord(name).copy() + for name in ('longitude', 'latitude')] + for co in (co_lon, co_lat): + co.convert_units('degrees') + lon_bds, lat_bds = [co.bounds for co in (co_lon, co_lat)] +# ny, nx = lon_bds.shape[:-1] + for iy in range(ny): + for ix in range(nx): + x0, y0 = lon_bds[iy, ix, 0], lat_bds[iy, ix, 0] + x1, y1 = lon_bds[iy, ix, 1], lat_bds[iy, ix, 1] + x2, y2 = lon_bds[iy, ix, 3], lat_bds[iy, ix, 3] + plt.plot([x0, x1], [y0, y1], 'x-', + color='orange', + transform=data_proj) + plt.plot([x0, x2], [y0, y2], 'x-', + color='orange', linestyle='--', + transform=data_proj) + + # Plot U0, rotated by cell angles, also at cell bottom-lefts. + u0_cube, u1_cube, v0_cube, v1_cube = [testcube.copy(data=aa) + for aa in (u0, v0, u1, v1)] + u0r_cube, v0r_cube = rotate_grid_vectors( + u0_cube, v0_cube, grid_angles_cube=angles_cube) + u0r, v0r = [cube.data for cube in (u0r_cube, v0r_cube)] + + xbl, ybl = lon_bds[..., 0], lat_bds[..., 0] + # + # Replace quiver here with delta-based lineplot + # + urd = xbl + deg_scale * u0r + vrd = ybl + deg_scale * v0r * yscale + for iy in range(ny): + for ix in range(nx): + plt.plot([xbl[iy, ix], urd[iy, ix]], + [ybl[iy, ix], vrd[iy, ix]], + ':', color='magenta', linewidth=2.5, + transform=data_proj) + # Show this is the SAME as lineplot + plt.quiver(xbl, ybl, u0r, v0r * yscale, + color='magenta', width=0.01, + headwidth=1.2, # headlength=1.0, headaxislength=0.7, + angles=angles, + scale_units='xy', scale=scale, + transform=data_proj) + + plt.suptitle('angles from "{}"'.format(cell_angle_boundpoints)) + +# # Also draw small lines pointing at the correct (TRUE, not ) angle. +# ny, nx = x0s.shape +# size_degrees = 1.0 +# angles = angles_cube.copy() +# angles.convert_units('radians') +# angles = angles.data +# lats = testcube.coord('latitude').copy() +# lats.convert_units('radians') +# lats = lats.points +# dxs = size_degrees * u0.copy() #* np.cos(angles) +# dys = size_degrees * u0.copy() # / np.sqrt(np.cos(lats)) +# x1s = x0s + dxs +# y1s = y0s + dys +## for iy in range(ny): +## for ix in range(nx): +## plt.plot([x0s[iy, ix], x1s[iy, ix]], +## [y0s[iy, ix], y1s[iy, ix]], +## 'o-', markersize=4., markeredgewidth=0., +## color='green', # scale_units='xy', scale=scale, +## transform=data_proj) +# plt.quiver(x0s, y0s, dxs, dys, +# color='green', linewidth=0.2, +# angles=angles, +# scale_units='xy', scale=scale * 0.6, +# transform=data_proj) + + + + ax.set_global() + plt.show() + + angles_cube.convert_units('degrees') + + self.assertArrayAllClose( + angles_cube.data, + [[33.421, 17.928, 0., -17.928, -33.421], + [41.981, 24.069, 0., -24.069, -41.981], + [56.624, 37.809, 0., -37.809, -56.624], + [79.940, 74.227, 0., -74.227, -79.940], + [107.313, 126.361, -180., -126.361, -107.313]], + atol=0.002) + + +if __name__ == "__main__": + tests.main()