official-source/DLWP-CS/DLWP/plot/util.py

#
# Copyright (c) 2017-18 Jonathan Weyn <jweyn@uw.edu>
#
# See the file LICENSE for your rights.
#

"""
Plotting utilities.
"""

import numpy as np
from matplotlib.colors import LinearSegmentedColormap, ListedColormap


def rotate_vector_r(basemap, uin, vin, lons, lats, returnxy=False):
    """
    I got this method from StackOverflow at some point, but didn't document its origin. Work is not my own.

    Similar to Basemap.rotate_vector except that it returns map-projected
    vectors to horizontal grid vectors. The following is the docstring
    for Basemap.rotate_vector. Note this function takes the basemap instance
    as an argument.

    Rotate a vector field (``uin,vin``) on a rectilinear grid
    with longitudes = ``lons`` and latitudes = ``lats`` from
    geographical (lat/lon) into map projection (x/y) coordinates.

    Differs from transform_vector in that no interpolation is done.
    The vector is returned on the same grid, but rotated into
    x,y coordinates.

    The input vector field is defined in spherical coordinates (it
    has eastward and northward components) while the output
    vector field is rotated to map projection coordinates (relative
    to x and y). The magnitude of the vector is preserved.

    .. tabularcolumns:: |l|L|

    ==============   ====================================================
    Arguments        Description
    ==============   ====================================================
    uin, vin         input vector field on a lat/lon grid.
    lons, lats       Arrays containing longitudes and latitudes
                     (in degrees) of input data in increasing order.
                     For non-cylindrical projections (those other than
                     ``cyl``, ``merc``, ``cyl``, ``gall`` and ``mill``) lons
                     must fit within range -180 to 180.
    ==============   ====================================================

    Returns ``uout, vout`` (rotated vector field).
    If the optional keyword argument
    ``returnxy`` is True (default is False),
    returns ``uout,vout,x,y`` (where ``x,y`` are the map projection
    coordinates of the grid defined by ``lons,lats``).
    """

    # if lons,lats are 1d and uin,vin are 2d, and
    # lats describes 1st dim of uin,vin, and
    # lons describes 2nd dim of uin,vin, make lons,lats 2d
    # with meshgrid.
    if lons.ndim == lats.ndim == 1 and uin.ndim == vin.ndim == 2 and \
            uin.shape[1] == vin.shape[1] == lons.shape[0] and \
            uin.shape[0] == vin.shape[0] == lats.shape[0]:
        lons, lats = np.meshgrid(lons, lats)
    else:
        if not lons.shape == lats.shape == uin.shape == vin.shape:
            raise TypeError("shapes of lons,lats and uin,vin don't match")
    x, y = basemap(lons, lats)
    # rotate from geographic to map coordinates.

    # Map the (lon, lat) vector in the complex plane.
    uvc = uin + 1j * vin
    uvmag = np.abs(uvc)
    theta = np.angle(uvc)

    # Define a displacement (dlon, dlat) that moves all
    # positions (lons, lats) a small distance in the
    # direction of the original vector.
    dc = 1E-5 * np.exp(theta * 1j)
    dlat = dc.imag * np.cos(np.radians(lats))
    dlon = dc.real

    # Deal with displacements that overshoot the North or South Pole.
    farnorth = np.abs(lats + dlat) >= 90.0
    somenorth = farnorth.any()
    if somenorth:
        dlon[farnorth] *= -1.0
        dlat[farnorth] *= -1.0

    # Add displacement to original location and find the native coordinates.
    lon1 = lons + dlon
    lat1 = lats + dlat
    xn, yn = basemap(lon1, lat1)

    # Determine the angle of the displacement in the native coordinates.
    vecangle = np.arctan2(yn - y, xn - x)
    if somenorth:
        vecangle[farnorth] += np.pi

    # Reverse the direction of vecangle about the original vector
    vecdiff = vecangle - theta
    vecangle -= 2. * vecdiff

    # Compute the x-y components of the original vector.
    uvcout = uvmag * np.exp(1j * vecangle)
    uout = uvcout.real
    vout = uvcout.imag

    if returnxy:
        return uout, vout, x, y
    else:
        return uout, vout


def radar_colormap():
    """
    Function to output a matplotlib color map object for reflectivity based on
    the National Weather Service color scheme.
    """

    nws_reflectivity_colors = [
        #    "#646464", # ND
        #    "#ccffff", # -30
        #    "#cc99cc", # -25
        #    "#996699", # -20
        #    "#663366", # -15
        #    "#cccc99", # -10
        #    "#999966", # -5
        #    "#646464", # 0
        "#ffffff",  # 0 white
        "#04e9e7",  # 5
        "#019ff4",  # 10
        "#0300f4",  # 15
        "#02fd02",  # 20
        "#01c501",  # 25
        "#008e00",  # 30
        "#fdf802",  # 35
        "#e5bc00",  # 40
        "#fd9500",  # 45
        "#fd0000",  # 50
        "#d40000",  # 55
        "#bc0000",  # 60
        "#f800fd",  # 65
        "#9854c6",  # 70
        #    "#fdfdfd" # 75
    ]
    return ListedColormap(nws_reflectivity_colors)


def blue_red_colormap(size=100, reverse=False, white_padding=1, ):
    size = size - (size % 4)
    ll = size // 4
    clist = []
    s1 = [0.5 + 0.5 / ll * x for x in range(ll)]
    s2 = [1.0 / ll * x for x in range(ll)]
    for x in s1:
        clist.append([0.0, 0.0, x])
    for x in s2:
        clist.append([x, x, 1.0])
    for x in range(white_padding):
        clist.append([1.0, 1.0, 1.0])
    for x in range(ll):
        clist.append([1.0, s2[-x - 1], s2[-x - 1]])
    for x in range(ll):
        clist.append([s1[-x - 1], 0.0, 0.0])
    if reverse:
        clist = clist[::-1]

    from matplotlib.colors import LinearSegmentedColormap
    cmap = LinearSegmentedColormap.from_list('BuRd', clist)
    return cmap


def rgb_colormap(color='blue', size=100, reverse=False, white_padding=1, ):
    ll = size - white_padding
    clist = []
    r = 1
    g = 1
    b = 1
    if color == 'red':
        r = 0.5
    elif color == 'green':
        g = 0.5
    elif color == 'blue':
        b = 0.5
    else:
        raise ValueError('Select "red", "green", or "blue" for color.')
    for x in range(white_padding):
        clist.append([1.0, 1.0, 1.0])
    for x in range(white_padding, size, 1):
        y = x - white_padding
        clist.append([1.0 - 1.0 * r * y / ll, 1.0 - 1.0 * g * y / ll, 1.0 - 1.0 * b * y / ll])
    if reverse:
        clist = clist[::-1]

    cmap = LinearSegmentedColormap.from_list(color, clist)
    return cmap


def shifted_color_map(cmap, start=0, midpoint=0.5, stop=1.0, name='shiftedcmap'):
    """
    Function to offset the "center" of a colormap. Useful for
    data with a negative min and positive max requiring the
    middle of the colormap's dynamic range to be at zero.

    Input
    -----
      cmap : The matplotlib colormap to be altered
      start : Offset from lowest point in the colormap's range.
          Defaults to 0.0 (no lower offset). Should be between
          0.0 and `midpoint`.
      midpoint : The new center of the colormap. Defaults to
          0.5 (no shift). Should be between 0.0 and 1.0. In
          general, this should be  1 - vmax/(vmax + abs(vmin)).
          For example, for a data range from -15.0 to +5.0 and
          the center of the colormap at 0.0, `midpoint`
          should be set to  1 - 5/(5 + 15)) or 0.75
      stop : Offset from highest point in the colormap's range.
          Defaults to 1.0 (no upper offset). Should be between
          `midpoint` and 1.0.
    """
    cdict = {
        'red': [],
        'green': [],
        'blue': [],
        'alpha': []
    }
    # regular index to compute the colors
    reg_index = np.linspace(start, stop, 257)
    # shifted index to match the data
    shift_index = np.hstack([
        np.linspace(0.0, midpoint, 128, endpoint=False),
        np.linspace(midpoint, 1.0, 129, endpoint=True)
    ])
    for rj, sj in zip(reg_index, shift_index):
        r, g, b, a = cmap(rj)
        cdict['red'].append((sj, r, r))
        cdict['green'].append((sj, g, g))
        cdict['blue'].append((sj, b, b))
        cdict['alpha'].append((sj, a, a))
    new_cmap = LinearSegmentedColormap(name, cdict)
    # plt.register_cmap(cmap=new_cmap)
    return new_cmap