Vector plots 2

lib/iris/analysis/_grid_angles.py

@@ -24,6 +24,7 @@
 
 import numpy as np
 
+import cartopy.crs as ccrs
 import iris
 
 
@@ -33,12 +34,16 @@ def _3d_xyz_from_latlon(lon, lat):
 
     Args:
 
-    * lon, lat: (arrays in degrees)
+    * lon, lat: (float array)
+        Arrays of longitudes and latitudes, in degrees.
+        Both the same shape.
 
     Returns:
 
-        xyz : (array, dtype=float64)
-            cartesian coordinates on a unit sphere.  Dimension 0 maps x,y,z.
+    * xyz : (array, 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)
@@ -60,91 +65,87 @@ def _latlon_from_xyz(xyz):
     Args:
 
     * xyz: (array)
-        positions array, of dims (3, <others>), where index 0 maps x/y/z.
+        Array of 3-D cartesian coordinates.
+        Shape (3, <input_points_dimensions>).
+        x / y / z values are in xyz[0 / 1 / 2],
 
     Returns:
 
-        lonlat : (array)
-            spherical angles, of dims (2, <others>), in radians.
-            Dim 0 maps longitude, latitude.
+    * lonlat : (array)
+        longitude and latitude position angles, in degrees.
+        Shape (2, <input_points_dimensions>).
+        The longitudes / latitudes are in lonlat[0 / 1].
 
     """
-    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)
+    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):
     """
-    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.
+    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)}.
+
+    Method:
+
+        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.
+
+    Args:
+
+    * 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].
+
+    Returns:
+
+    * angle : (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>.
 
     """
-#    old_style = True
-    old_style = False
-    if old_style:
-        mid_lons = np.deg2rad(q[0])
+    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)
+    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
+    index = pr_norm == 0
+    pr_norm[index] = 1
 
-        cosine = np.maximum(np.minimum(pr_top / pr_norm, 1), -1)
-        cosine[index] = 0
+    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
+    psi = np.arccos(cosine) * np.sign(r[1] - p[1])
+    psi[index] = np.nan
+
+    return np.rad2deg(psi)
 
 
 def gridcell_angles(x, y=None, cell_angle_boundpoints='mid-lhs, mid-rhs'):
@@ -201,12 +202,18 @@ def gridcell_angles(x, y=None, cell_angle_boundpoints='mid-lhs, mid-rhs'):
         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.
+            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
@@ -216,55 +223,105 @@ def gridcell_angles(x, y=None, cell_angle_boundpoints='mid-lhs, mid-rhs'):
         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'):
+    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()
-        x_coord.convert_units('degrees')
-        y_coord.convert_units('degrees')
+
+        # 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. ',
+            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. ',
+            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 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(*is_and_not)
+        raise ValueError(msg.format(*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.
+    # 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:
-        # 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 ?
+        # 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.
 
-        # 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.
+        # 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)
@@ -275,9 +332,11 @@ def gridcell_angles(x, y=None, cell_angle_boundpoints='mid-lhs, mid-rhs'):
             x_coord = iris.coords.AuxCoord(y, standard_name='longitude',
                                            units='degrees')
     else:
-        # Get lhs and rhs locations by averaging top+bottom each side.
+        # 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)
@@ -294,8 +353,8 @@ def gridcell_angles(x, y=None, cell_angle_boundpoints='mid-lhs, mid-rhs'):
                                         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.
+            # 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.
@@ -305,28 +364,30 @@ def gridcell_angles(x, y=None, cell_angle_boundpoints='mid-lhs, mid-rhs'):
             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.
+        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='radians')
+                            units='degrees')
     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):
+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