Source code for iris.analysis._grid_angles

# Copyright Iris contributors
# This file is part of Iris and is released under the BSD license.
# See LICENSE in the root of the repository for full licensing details.

"""Implement vector rotation by angles.

Code to implement vector rotation by angles, and inferring gridcell angles
from coordinate points and bounds.


import as ccrs
import numpy as np

import iris

def _3d_xyz_from_latlon(lon, lat):
    """Return locations of (lon, lat) in 3D space.

    lon, lat : float array
        Arrays of longitudes and latitudes, in degrees.
        Both the same shape.

    array of dtype=float64
        Cartesian coordinates on a unit sphere.
        Shape is (3, <input-shape>).
        The x / y / z coordinates are in xyz[0 / 1 / 2].

    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.

    xyz : array
        Array of 3-D cartesian coordinates.
        Shape (3, <input_points_dimensions>).
        x / y / z values are in xyz[0 / 1 / 2].

        Longitude and latitude position angles, in degrees.
        Shape (2, <input_points_dimensions>).
        The longitudes / latitudes are in lonlat[0 / 1].

    lons = np.rad2deg(np.arctan2(xyz[1], xyz[0]))
    radii = np.sqrt(np.sum(xyz * xyz, axis=0))
    lats = np.rad2deg(np.arcsin(xyz[2] / radii))
    return np.array([lons, lats])

def _angle(p, q, r):
    """Estimate grid-angles to true-Eastward direction.

    Estimate grid-angles to true-Eastward direction from positions in the same
    grid row, but at increasing column (grid-Eastward) positions.

    {P, Q, R} are locations of consecutive points in the same grid row.
    These could be successive points in a single grid, e.g.::

        {T(i-1,j), T(i,j), T(i+1,j)}

    or a mixture of U/V and T gridpoints if row positions are aligned, e.g::

        {v(i,j), f(i,j), v(i+1,j)}.


        Calculate dot product of a unit-vector parallel to P-->R, unit-scaled,
        with the unit eastward (true longitude) vector at Q.
        This value is cos(required angle).
        Discriminate between +/- angles by comparing latitudes of P and R.
        Return NaN where any P-->R are zero.

        .. note::

            This method assumes that the vector PR is parallel to the surface
            at the longitude of Q, as it uses the length of PR as the basis for
            the cosine ratio.
            That is only exact when Q is at the same longitude as the midpoint
            of PR, and this typically causes errors which grow with increasing
            gridcell angle.
            However, we retain this method because it reproduces the "standard"
            gridcell-orientation-angle arrays found in files output by the CICE
            model, which presumably uses an equivalent calculation.

    p, q, r : float array
        Arrays of angles, in degrees.
        All the same shape.
        Shape is (2, <input_points_dimensions>).
        Longitudes / latitudes are in array[0 / 1].

    float array
        Grid angles relative to true-East, in degrees.
        Positive when grid-East is anticlockwise from true-East.
        Shape is same as <input_points_dimensions>.

    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

    return np.rad2deg(psi)

[docs] 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. The inputs (x [,y]) can be different, see the parameters section. Parameters ---------- 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. cell_angle_boundpoints : str, default="mid-lhs, mid-rhs" 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 ------- 2-dimensional cube Cube of angles of grid-x vector from true Eastward direction for each gridcell, in degrees. It also has "true" longitude and latitude coordinates, with no coordinate system. When the input has coords, then the output ones are identical if the inputs are true-latlons, otherwise they are transformed true-latlon versions. When the input has bounded coords, then the output coords have matching bounds and centrepoints (possibly transformed). When the input is 2d arrays, or has unbounded coords, then the output coords have matching points and no bounds. When the input is 3d arrays, then the output coords have matching bounds, and the centrepoints are an average of the 4 boundpoints. """ 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") or not 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 and y are Coords. x_coord, y_coord = x.copy(), y.copy() # They must be angles : convert into degrees for coord in (x_coord, y_coord): if not coord.units.is_convertible("degrees"): msg = ( "Input X and Y coordinates must have angular " 'units. Got units of "{!s}" and "{!s}".' ) raise ValueError(msg.format(x_coord.units, y_coord.units)) 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)) if cube: x_dims, y_dims = (cube.coord_dims(co) for co in (x, y)) if x_dims != y_dims: msg = ( "X and Y coordinates must have the same cube " "dimensions. Got x-dims = {} and y-dims = {}." ) raise ValueError(msg.format(x_dims, y_dims)) cs = x_coord.coord_system if y_coord.coord_system != cs: msg = ( "Coordinate inputs must have same coordinate system. " "Got x of {} and y of {}." ) raise ValueError(msg.format(cs, y_coord.coord_system)) # Base calculation on bounds if we have them, or points as a fallback. 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 # Make sure these arrays are ordinary lats+lons, in degrees. if cs is not None: # Transform points into true lats + lons. crs_src = cs.as_cartopy_crs() crs_pc = ccrs.PlateCarree() def transform_xy_arrays(x, y): # Note: flatten, as transform_points is limited to 2D arrays. shape = x.shape x, y = (arr.flatten() for arr in (x, y)) pts = crs_pc.transform_points(crs_src, x, y) x = pts[..., 0].reshape(shape) y = pts[..., 1].reshape(shape) return x, y # Transform the main reference points into standard lats+lons. x, y = transform_xy_arrays(x, y) # Likewise replace the original coordinates with transformed ones, # because the calculation also needs the centrepoint values. xpts, ypts = (coord.points for coord in (x_coord, y_coord)) xbds, ybds = (coord.bounds for coord in (x_coord, y_coord)) xpts, ypts = transform_xy_arrays(xpts, ypts) xbds, ybds = transform_xy_arrays(xbds, ybds) x_coord = iris.coords.AuxCoord( points=xpts, bounds=xbds, standard_name="longitude", units="degrees", ) y_coord = iris.coords.AuxCoord( points=ypts, bounds=ybds, standard_name="latitude", units="degrees", ) 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(msg.format(*is_and_not)) # Now have either 2 points arrays (ny, nx) or 2 bounds arrays (ny, nx, 4). # Construct (lhs, mid, rhs) where these represent 3 points at increasing # grid-x indices (columns). # Also make suitable X and Y coordinates for the result cube. if x.ndim == 2: # Data is points arrays. # Use previous + subsequent points along grid-x-axis as references. # PROBLEM: we assume that the rhs connects to the lhs, so we should # really only use this if data is full-longitudes (as a 'circular' # coordinate). # This is mentioned in the docstring, but we have no general means of # checking it. # NOTE: we take the 2d grid as presented, so the second dimension is # the 'X' dim. Again, that is implicit + can't be checked. 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: # Data is bounds arrays. # Use gridcell corners at different grid-x positions as references. # NOTE: so with bounds, we *don't* need full circular longitudes. xyz = _3d_xyz_from_latlon(x, y) # Support two different choices of reference points locations. 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 of lhs+rhs points in 3d to get 'mid' points, # as coords without 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 = _latlon_from_xyz(lhs_xyz) rhs = _latlon_from_xyz(rhs_xyz) # 'mid' is coord.points, whether from input or just 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="degrees" ) result.add_aux_coord(x_coord, (0, 1)) result.add_aux_coord(y_coord, (0, 1)) return result
[docs] def rotate_grid_vectors(u_cube, v_cube, grid_angles_cube=None, grid_angles_kwargs=None): """Rotate distance vectors from grid-oriented to true-latlon-oriented. Can also rotate by arbitrary angles, if they are passed in. .. 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. Parameters ---------- 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'. grid_angles_cube : cube, optional 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, optional Additional keyword args to be passed to the :func:`gridcell_angles` method, if it is used. Returns ------- (cube, cube) Tuple of cubes of true-north oriented vector components. Units are same as inputs. Notes ----- Vector magnitudes will always be the same as the inputs. This function does not maintain laziness when called; it realises data. See more at :doc:`/userguide/real_and_lazy_data`. """ 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 = ( 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(, "mask"): mask |= =, mask=mask) =, mask=mask) return u_out, v_out