File size: 8,226 Bytes
989c6ea
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
#
# 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