def test_lat_lon_grid_deltas_mismatched_shape():
"""Test for lat_lon_grid_deltas for variable grid."""
lat = np.arange(40, 50, 2.5)
lon = np.array([[-100., -97.5, -95., -92.5], [-100., -97.5, -95., -92.5],
[-100., -97.5, -95., -92.5], [-100., -97.5, -95., -92.5]])
with pytest.raises(ValueError):
lat_lon_grid_deltas(lon, lat)
def test_lat_lon_grid_deltas_mismatched_shape():
"""Test for lat_lon_grid_deltas for variable grid."""
lat = np.arange(40, 50, 2.5)
lon = np.array([[-100., -97.5, -95., -92.5],
[-100., -97.5, -95., -92.5],
[-100., -97.5, -95., -92.5],
[-100., -97.5, -95., -92.5]])
with pytest.raises(ValueError):
lat_lon_grid_deltas(lon, lat)
def calcgw_gfs(v, lat, lon):
height, lats, lons = v.data(lat1=lat - gw_gfs_margin_deg,
lat2=lat + gw_gfs_margin_deg,
lon1=lon - gw_gfs_margin_deg,
lon2=lon + gw_gfs_margin_deg)
i = np.searchsorted(lats[:, 0], lat)
if abs(lats[i + 1, 0] - lat) < abs(lats[i, 0] - lat):
i = i + 1
j = np.searchsorted(lons[0, :], lon)
if abs(lons[0, i + 1] - lon) < abs(lons[0, i] - lon):
j = j + 1
#print('level', v.level, 'height', height[i,j], lats[i,j], lons[i,j])
# Set up some constants based on our projection, including the Coriolis parameter and
# grid spacing, converting lon/lat spacing to Cartesian
f = mpcalc.coriolis_parameter(np.deg2rad(lats)).to('1/s')
dx, dy = mpcalc.lat_lon_grid_deltas(lons, lats)
res_km = (dx[i, j] + dy[i, j]).magnitude / 2000.
# Smooth height data. Sigma=1.5 for gfs 0.5deg
height = ndimage.gaussian_filter(height, sigma=1.5 * 50 / res_km, order=0)
# In MetPy 0.5, geostrophic_wind() assumes the order of the dimensions is (X, Y),
# so we need to transpose from the input data, which are ordered lat (y), lon (x).
# Once we get the components,transpose again so they match our original data.
geo_wind_u, geo_wind_v = mpcalc.geostrophic_wind(height * units.m, f, dx,
dy)
return height[i, j], geo_wind_u[i, j], geo_wind_v[i, j]
def advection(variable, u_da, v_da, units_wind='m/s'):
"""
Calcula la advección en base a un dataarray de u y un dataarray de v
en m/s """
lats = u_da.lat
lons = u_da.lon
u = u_da.values * units['m/s']
v = v_da.values * units['m/s']
dx, dy = mpcalc.lat_lon_grid_deltas(lons, lats)
return mpcalc.advection(variable, [u, v],
(dx, dy), dim_order='yx')
def compute_vorticity(dset, uvar='10u', vvar='10v'):
dx, dy = mpcalc.lat_lon_grid_deltas(dset['lon'], dset['lat'])
vort = mpcalc.vorticity(dset[uvar], dset[vvar], dx[None, :, :],
dy[None, :, :])
vort = xr.DataArray(vort.magnitude,
coords=dset[uvar].coords,
attrs={
'standard_name': 'vorticity',
'units': vort.units
},
name='vort')
return xr.merge([dset, vort])
def compute_convergence(dset, uvar='10u', vvar='10v'):
dx, dy = mpcalc.lat_lon_grid_deltas(dset['lon'], dset['lat'])
conv = -mpcalc.divergence(dset[uvar], dset[vvar], dx[None, :, :],
dy[None, :, :])
conv = xr.DataArray(conv.magnitude,
coords=dset[uvar].coords,
attrs={
'standard_name': 'convergence',
'units': conv.units
},
name='conv')
return xr.merge([dset, conv])
def test_lat_lon_grid_deltas_1d():
"""Test for lat_lon_grid_deltas for variable grid."""
lat = np.arange(40, 50, 2.5)
lon = np.arange(-100, -90, 2.5)
dx, dy = lat_lon_grid_deltas(lon, lat)
dx_truth = np.array([[212943.5585, 212943.5585, 212943.5585],
[204946.2305, 204946.2305, 204946.2305],
[196558.8269, 196558.8269, 196558.8269],
[187797.3216, 187797.3216, 187797.3216]]) * units.meter
dy_truth = np.array([[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857]]) * units.meter
assert_almost_equal(dx, dx_truth, 4)
assert_almost_equal(dy, dy_truth, 4)
def test_lat_lon_grid_deltas_1d():
"""Test for lat_lon_grid_deltas for variable grid."""
lat = np.arange(40, 50, 2.5)
lon = np.arange(-100, -90, 2.5)
dx, dy = lat_lon_grid_deltas(lon, lat)
dx_truth = np.array([[212943.5585, 212943.5585, 212943.5585],
[204946.2305, 204946.2305, 204946.2305],
[196558.8269, 196558.8269, 196558.8269],
[187797.3216, 187797.3216, 187797.3216]]) * units.meter
dy_truth = np.array([[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857]]) * units.meter
assert_almost_equal(dx, dx_truth, 4)
assert_almost_equal(dy, dy_truth, 4)
def calc_fronts(u, v, q, t, lon, lat, date_list):
'''
Parse era5 variables to kinematics(), to calculate various thermal and kinematic front parameters.
12 are computed, but for brevity, only four are returned for now.
'''
#Using MetPy, derive equivalent potential temperature
ta_unit = units.units.K * t
dp_unit = mpcalc.dewpoint_from_specific_humidity(
q * units.units.dimensionless, ta_unit, 850 * units.units.hectopascals)
thetae = np.array(
mpcalc.equivalent_potential_temperature(850 * units.units.hectopascals,
ta_unit, dp_unit))
#From the lat-lon information accompanying the ERA5 data, resconstruct 2d coordinates and grid spacing (dx, dy)
x, y = np.meshgrid(lon, lat)
dx, dy = mpcalc.lat_lon_grid_deltas(x, y)
#Derive various kinematic/thermal diagnostics for each time step
kinemats = [
kinematics(u[i], v[i], thetae[i], dx, dy, y, smooth=True, sigma=2)
for i in np.arange(len(date_list))
]
F = [kinemats[i][0] for i in np.arange(len(date_list))
] #Frontogenesis function (degC / 100 km / 3 hr)
Fn = [kinemats[i][1] for i in np.arange(len(date_list))
] #Frontogenetical function (deg C / 100 km / 3 hr)
Fs = [kinemats[i][2] for i in np.arange(len(date_list))
] #Rotational component of frontogensis (deg C/ 100 km / 3 hr)
icon = [kinemats[i][3] for i in np.arange(len(date_list))
] #Instantaneous contraction rate (s^-1 * 1e5)
vgt = [kinemats[i][4] for i in np.arange(len(date_list))
] #Horizontal velovity gradient tensor magnitude (s^-1 * 1e5)
conv = [kinemats[i][5]
for i in np.arange(len(date_list))] #Convergence (s^-1 * 1e5)
vo = [kinemats[i][6]
for i in np.arange(len(date_list))] #Relative vorticity (s^-1 * 1e5)
tfp = [kinemats[i][7] for i in np.arange(len(date_list))
] #Thermal front parameter (km^-2)
mag_te = [kinemats[i][8] for i in np.arange(len(date_list))
] #Magnitude of theta-e gradient (100 km ^-1)
v_f = [kinemats[i][9]
for i in np.arange(len(date_list))] #Advection of TFP (m/s)
thetae = [kinemats[i][10] for i in np.arange(len(date_list))
] #Theta-e (K), may be smoothed depending on arguments
cond = [kinemats[i][11]
for i in np.arange(len(date_list))] #Extra condition
return [thetae, mag_te, tfp, v_f]
def test_lat_lon_grid_deltas_geod_kwargs():
"""Test that geod kwargs are overridden by users #774."""
lat = np.arange(40, 50, 2.5)
lon = np.arange(-100, -90, 2.5)
dx, dy = lat_lon_grid_deltas(lon, lat, a=4370997)
dx_truth = np.array([[146095.76101984, 146095.76101984, 146095.76101984],
[140608.9751528, 140608.9751528, 140608.9751528],
[134854.56713287, 134854.56713287, 134854.56713287],
[128843.49645823, 128843.49645823, 128843.49645823]]) * units.meter
dy_truth = np.array([[190720.72311199, 190720.72311199, 190720.72311199, 190720.72311199],
[190720.72311199, 190720.72311199, 190720.72311199, 190720.72311199],
[190720.72311199, 190720.72311199, 190720.72311199,
190720.72311199]]) * units.meter
assert_almost_equal(dx, dx_truth, 4)
assert_almost_equal(dy, dy_truth, 4)
def test_lat_lon_grid_deltas_extra_dimensions():
"""Test for lat_lon_grid_deltas with extra leading dimensions."""
lon, lat = np.meshgrid(np.arange(-100, -90, 2.5), np.arange(40, 50, 2.5))
lat = lat[None, None]
lon = lon[None, None]
dx_truth = np.array([[[[212943.5585, 212943.5585, 212943.5585],
[204946.2305, 204946.2305, 204946.2305],
[196558.8269, 196558.8269, 196558.8269],
[187797.3216, 187797.3216, 187797.3216]]]]) * units.meter
dy_truth = (np.array([[[[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857]]]]) *
units.meter)
dx, dy = lat_lon_grid_deltas(lon, lat)
assert_almost_equal(dx, dx_truth, 4)
assert_almost_equal(dy, dy_truth, 4)
def test_lat_lon_grid_deltas_extra_dimensions():
"""Test for lat_lon_grid_deltas with extra leading dimensions."""
lon, lat = np.meshgrid(np.arange(-100, -90, 2.5), np.arange(40, 50, 2.5))
lat = lat[None, None]
lon = lon[None, None]
dx_truth = np.array([[[[212943.5585, 212943.5585, 212943.5585],
[204946.2305, 204946.2305, 204946.2305],
[196558.8269, 196558.8269, 196558.8269],
[187797.3216, 187797.3216, 187797.3216]]]]) * units.meter
dy_truth = (np.array([[[[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857]]]])
* units.meter)
dx, dy = lat_lon_grid_deltas(lon, lat)
assert_almost_equal(dx, dx_truth, 4)
assert_almost_equal(dy, dy_truth, 4)
def pv_data():
"""Test data for all PV testing."""
u = np.array([[100, 90, 80, 70], [90, 80, 70, 60], [80, 70, 60, 50],
[70, 60, 50, 40]]) * units('m/s')
v = np.zeros_like(u) * units('m/s')
lats = np.array([[40, 40, 40, 40], [40.1, 40.1, 40.1, 40.1],
[40.2, 40.2, 40.2, 40.2], [40.3, 40.3, 40.3, 40.3]
]) * units.degrees
lons = np.array([[40, 39.9, 39.8, 39.7], [40, 39.9, 39.8, 39.7],
[40, 39.9, 39.8, 39.7], [40, 39.9, 39.8, 39.7]
]) * units.degrees
dx, dy = lat_lon_grid_deltas(lons, lats)
return u, v, lats, lons, dx, dy
def divergence(self, level):
"""
Uses a metpy function to calculate wind divergence in a given level.
Uses predefined fuctions to obtain wind and grid data.
Returns a np.array.
"""
# Grab lat/lon values
lat = self.data['lat'].values
lon = self.data['lon'].values
# Compute dx and dy spacing for use in divergence calculation
dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat)
# Extract wind components
uwnd = self.u_wind(level)
vwnd = self.v_wind(level)
# Use MetPy to compute the divergence for the given wind level
div = mpcalc.divergence(uwnd, vwnd, dx, dy)
return np.array(div * 100000)
def geostr_met(array_2d_geoheight, array_1d_lat, array_1d_lon) :
'''
Calculate Geostrophic wind map (2d) from 2d array area.
https://unidata.github.io/python-gallery/examples/Ageostrophic_Wind_Example.html
Parameters
----------
array_2d_geoheight : read numpy array [m]
array_1d_lat : read numpy array [deg]
array_1d_lon : read numpy array [deg]
Returns
-------
array_2d : return interplated and extrapolated 2d array
'''
import numpy as np
import metpy.calc as mpcalc
from metpy.units import units
# Combine 1D latitude and longitudes into a 2D grid of locations
lon_2d, lat_2d = np.meshgrid(array_1d_lon, array_1d_lat)
# Set up some constants based on our projection, including the Coriolis parameter and
# grid spacing, converting lon/lat spacing to Cartesian
f = mpcalc.coriolis_parameter(np.deg2rad(lat_2d)).to('1/s')
dx, dy = mpcalc.lat_lon_grid_deltas(lon_2d + 360, lat_2d)
dx, dy = np.array(dx), np.array(dy)
dy *= -1
# In MetPy 0.5, geostrophic_wind() assumes the order of the dimensions is (X, Y),
# so we need to transpose from the input data, which are ordered lat (y), lon (x).
# Once we get the components,transpose again so they match our original data.
geo_wind_u, geo_wind_v = mpcalc.geostrophic_wind(array_2d_geoheight.data * units.m, f, dx, dy)
return(geo_wind_u, geo_wind_v)
def test_lat_lon_grid_deltas_2d(flip_order):
"""Test for lat_lon_grid_deltas for variable grid with negative delta distances."""
lat = np.arange(40, 50, 2.5)
lon = np.arange(-100, -90, 2.5)
dx_truth = np.array([[212943.5585, 212943.5585, 212943.5585],
[204946.2305, 204946.2305, 204946.2305],
[196558.8269, 196558.8269, 196558.8269],
[187797.3216, 187797.3216, 187797.3216]]) * units.meter
dy_truth = np.array([[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857]]) * units.meter
if flip_order:
lon = lon[::-1]
lat = lat[::-1]
dx_truth = -1 * dx_truth[::-1]
dy_truth = -1 * dy_truth[::-1]
lon, lat = np.meshgrid(lon, lat)
dx, dy = lat_lon_grid_deltas(lon, lat)
assert_almost_equal(dx, dx_truth, 4)
assert_almost_equal(dy, dy_truth, 4)
def test_lat_lon_grid_deltas_2d(flip_order):
"""Test for lat_lon_grid_deltas for variable grid with negative delta distances."""
lat = np.arange(40, 50, 2.5)
lon = np.arange(-100, -90, 2.5)
dx_truth = np.array([[212943.5585, 212943.5585, 212943.5585],
[204946.2305, 204946.2305, 204946.2305],
[196558.8269, 196558.8269, 196558.8269],
[187797.3216, 187797.3216, 187797.3216]]) * units.meter
dy_truth = np.array([[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857],
[277987.1857, 277987.1857, 277987.1857, 277987.1857]]) * units.meter
if flip_order:
lon = lon[::-1]
lat = lat[::-1]
dx_truth = -1 * dx_truth[::-1]
dy_truth = -1 * dy_truth[::-1]
lon, lat = np.meshgrid(lon, lat)
dx, dy = lat_lon_grid_deltas(lon, lat)
assert_almost_equal(dx, dx_truth, 4)
assert_almost_equal(dy, dy_truth, 4)
def compute_spacing(dset):
dx, dy = mpcalc.lat_lon_grid_deltas(dset['lon'], dset['lat'])
dx = xr.DataArray(dx.magnitude,
dims=['y1', 'x1'],
attrs={
'standard_name': 'x grid spacing',
'units': dx.units
},
name='dx')
dy = xr.DataArray(dy.magnitude,
dims=['y2', 'x2'],
attrs={
'standard_name': 'y grid spacing',
'units': dx.units
},
name='dy')
out = xr.merge([dset, dx, dy])
out.attrs = dset.attrs
return out
def grid_area_average(var, rad, lon, lat):
"""Performs horizontal area-averaging of a field in latitude/longitude format.
refer to
https://github.com/tomerburg/metlib/blob/master/diagnostics/area_average.py
https://github.com/tomerburg/metlib/blob/master/diagnostics/area_average_sample.ipynb
Parameters
----------
var : (M, N) ndarray
Variable to perform area averaging on. Can be 2, 3 or 4 dimensions. If 2D, coordinates must
be lat/lon. If using additional dimensions, area-averaging will only be performed on the last
2 dimensions, assuming those are latitude and longitude.
rad : `pint.Quantity`
The radius over which to perform the spatial area-averaging.
lon : array-like
Array of longitudes defining the grid
lat : array-like
Array of latitudes defining the grid
Returns
-------
(M, N) ndarray
Area-averaged quantity, returned in the same dimensions as passed.
Notes
-----
This function was originally provided by Matthew Janiga and Philippe Papin using a Fortran wrapper for NCL,
and converted to python with further efficiency modifications by Tomer Burg, with permission from the original
authors.
This function assumes that the last 2 dimensions of var are ordered as (....,lat,lon).
Examples
import xarray as xr
from metpy.units import units
run_date = "20190106"
init = "1200"
url = f"http://thredds.ucar.edu/thredds/dodsC/grib/NCEP/GFS/Global_0p25deg/GFS_Global_0p25deg_{run_date}_{init}.grib2"
data = xr.open_dataset(url)
data_subset = data.isel(time2=0)
g = data_subset['Geopotential_height_isobaric'].sel(isobaric=pres_level)
lat = data_subset.lat.values
lon = data_subset.lon.values
radius = 500.0 * units('kilometers')
g_avg = area_average(g,radius,lon,lat)
"""
#convert radius to kilometers
rad = rad.to('kilometers')
#res = distance in km of dataset resolution, at the equator
londiff = lon[1] - lon[0]
latdiff = lat[1] - lat[0]
lat_0 = 0.0 - (latdiff / 2.0)
lat_1 = 0.0 + (latdiff / 2.0)
dx, dy = calc.lat_lon_grid_deltas(np.array([lon[0], lon[1]]),
np.array([lat_0, lat_1]))
dx = dx.to('km')
res = int((dx[0].magnitude + dx[1].magnitude) / 2.0) * units('km')
#---------------------------------------------------------------------
#Error checks
#Check to make sure latitudes increase
reversed_lat = 0
if lat[1] < lat[0]:
reversed_lat = 1
#Reverse latitude array
lat = lat[::-1]
#Determine which axis of variable array to reverse
lat_dim = len(var.shape) - 2
var = np.flip(var, lat_dim)
#Check to ensure input array has 2, 3 or 4 dimensions
var_dims = np.shape(var)
if len(var_dims) not in [2, 3, 4]:
print("only 2D, 3D and 4D arrays allowed")
return
#---------------------------------------------------------------------
#Prepare for computation
#Number of points in circle (with buffer)
box = int((rad / res) + 2)
#Define empty average array
var_avg = np.zeros((var.shape))
#Convert lat and lon arrays to 2D
nlat = len(lat)
nlon = len(lon)
lon2d, lat2d = np.meshgrid(lon, lat)
RPD = 0.0174532925
lat2d = lat2d * RPD
lon2d = lon2d * RPD
#Define radius of earth in km
eqrm = 6378.137
#Create mask for elements of array that are outside of the box
mask = np.zeros((lon2d.shape))
nbox = (2 * box + 1) * (2 * box + 1)
mask[box:nlat - box, box:nlon - box] = 1
mask[mask == 0] = np.nan
#Calculate area-averaging depending on the dimension sizes
if len(var_dims) == 2:
var_avg = _calcavg(var.magnitude, var_avg, lon2d, lat2d, nlon, nlat,
rad.magnitude, box, eqrm) * mask
elif len(var_dims) == 3:
for t in range(var_dims[0]):
var_avg[t, :, :] = _calcavg(var[t, :, :].magnitude,
var_avg[t, :, :], lon2d, lat2d, nlon,
nlat, rad.magnitude, box, eqrm) * mask
elif len(var_dims) == 4:
for t in range(var_dims[0]):
for l in range(var_dims[1]):
var_avg[t, l, :, :] = _calcavg(
var[t, l, :, :].magnitude, var_avg[t, l, :, :], lon2d,
lat2d, nlon, nlat, rad.magnitude, box, eqrm) * mask
#If latitude is reversed, then flip it back to its original order
if reversed_lat == 1:
lat_dim = len(var.shape) - 2
var_avg = np.flip(var_avg, lat_dim)
#Return area-averaged array with the same units as the input variable
return var_avg * var.units
def Crosssection_Wind_Temp_RH(
initial_time=None,
fhour=24,
levels=[1000, 950, 925, 900, 850, 800, 700, 600, 500, 400, 300, 200],
day_back=0,
model='ECMWF',
output_dir=None,
st_point=[43.5, 111.5],
ed_point=[33, 125.0],
map_extent=[70, 140, 15, 55],
h_pos=[0.125, 0.665, 0.25, 0.2]):
# micaps data directory
try:
data_dir = [
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='RH',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='UGRD',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='VGRD',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='TMP',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='HGT',
lvl='500'),
utl.Cassandra_dir(data_type='surface',
data_source=model,
var_name='PSFC')
]
except KeyError:
raise ValueError('Can not find all directories needed')
# get filename
if (initial_time != None):
filename = utl.model_filename(initial_time, fhour)
else:
filename = utl.filename_day_back_model(day_back=day_back, fhour=fhour)
# retrieve data from micaps server
rh = get_model_3D_grid(directory=data_dir[0][0:-1],
filename=filename,
levels=levels,
allExists=False)
if rh is None:
return
rh = rh.metpy.parse_cf().squeeze()
u = get_model_3D_grid(directory=data_dir[1][0:-1],
filename=filename,
levels=levels,
allExists=False)
if u is None:
return
u = u.metpy.parse_cf().squeeze()
v = get_model_3D_grid(directory=data_dir[2][0:-1],
filename=filename,
levels=levels,
allExists=False)
if v is None:
return
v = v.metpy.parse_cf().squeeze()
v2 = get_model_3D_grid(directory=data_dir[2][0:-1],
filename=filename,
levels=levels,
allExists=False)
if v2 is None:
return
v2 = v2.metpy.parse_cf().squeeze()
t = get_model_3D_grid(directory=data_dir[3][0:-1],
filename=filename,
levels=levels,
allExists=False)
if t is None:
return
t = t.metpy.parse_cf().squeeze()
gh = get_model_grid(data_dir[4], filename=filename)
psfc = get_model_grid(data_dir[5], filename=filename)
psfc = psfc.metpy.parse_cf().squeeze()
mask1 = ((psfc['lon'] >= t['lon'].values.min()) &
(psfc['lon'] <= t['lon'].values.max()) &
(psfc['lat'] >= t['lat'].values.min()) &
(psfc['lat'] <= t['lat'].values.max()))
t2, psfc_bdcst = xr.broadcast(t['data'], psfc['data'].where(mask1,
drop=True))
mask2 = (psfc_bdcst > -10000)
psfc_bdcst = psfc_bdcst.where(mask2, drop=True)
#psfc_bdcst=psfc_bdcst.metpy.parse_cf().squeeze()
if t is None:
return
resolution = u['lon'][1] - u['lon'][0]
x, y = np.meshgrid(u['lon'], u['lat'])
dx, dy = mpcalc.lat_lon_grid_deltas(u['lon'], u['lat'])
#rh=rh.rename(dict(lat='latitude',lon='longitude'))
cross = cross_section(rh, st_point, ed_point)
cross_rh = cross.set_coords(('lat', 'lon'))
cross = cross_section(u, st_point, ed_point)
cross_u = cross.set_coords(('lat', 'lon'))
cross = cross_section(v, st_point, ed_point)
cross_v = cross.set_coords(('lat', 'lon'))
cross_psfc = cross_section(psfc_bdcst, st_point, ed_point)
#cross_psfc=cross.set_coords(('lat', 'lon'))
cross_u['data'].attrs['units'] = units.meter / units.second
cross_v['data'].attrs['units'] = units.meter / units.second
cross_u['t_wind'], cross_v['n_wind'] = mpcalc.cross_section_components(
cross_u['data'], cross_v['data'])
cross = cross_section(t, st_point, ed_point)
cross_Temp = cross.set_coords(('lat', 'lon'))
cross_Td = mpcalc.dewpoint_rh(cross_Temp['data'].values * units.celsius,
cross_rh['data'].values * units.percent)
rh, pressure = xr.broadcast(cross_rh['data'], cross_Temp['level'])
cross_terrain = pressure - cross_psfc
crossection_graphics.draw_Crosssection_Wind_Temp_RH(
cross_rh=cross_rh,
cross_Temp=cross_Temp,
cross_u=cross_u,
cross_v=cross_v,
cross_terrain=cross_terrain,
gh=gh,
h_pos=h_pos,
st_point=st_point,
ed_point=ed_point,
levels=levels,
map_extent=map_extent,
model=model,
output_dir=output_dir)
# Nearly all of the calculations in `metpy.calc` will accept DataArrays by converting them
# into their corresponding unit arrays. While this may often work without any issues, we must
# keep in mind that because the calculations are working with unit arrays and not DataArrays:
#
# - The calculations will return unit arrays rather than DataArrays
# - Broadcasting must be taken care of outside of the calculation, as it would only recognize
# dimensions by order, not name
#
# Also, some of the units used in CF conventions (such as 'degrees_north') are not recognized
# by pint, so we must implement a workaround.
#
# As an example, we calculate geostropic wind at 500 hPa below:
lat, lon = xr.broadcast(y, x)
f = mpcalc.coriolis_parameter(lat.values * units.degrees)
dx, dy = mpcalc.lat_lon_grid_deltas(lon.values, lat.values)
heights = data['height'].loc[time[0]].loc[{vertical.name: 500.}]
u_geo, v_geo = mpcalc.geostrophic_wind(heights, f, dx, dy, dim_order='yx')
print(u_geo)
print(v_geo)
#########################################################################
# Plotting
# --------
#
# Like most meteorological data, we want to be able to plot these data. DataArrays can be used
# like normal numpy arrays in plotting code, or we can use some of xarray's plotting
# functionality.
#
# (More detail beyond the following can be found at `xarray's plotting reference
# <http://xarray.pydata.org/en/stable/plotting.html>`_.)
# Create a clean datetime object for plotting based on time of Geopotential heights
vtime = datetime.strptime(str(ds.time.data[0].astype('datetime64[ms]')),
'%Y-%m-%dT%H:%M:%S.%f')
######################################################################
# MetPy Absolute Vorticity Calculation
# ------------------------------------
#
# This code first uses MetPy to calcualte the grid deltas (sign aware) to
# use for derivative calculations with the funtcion
# ``lat_lon_grid_deltas()`` and then calculates ``absolute_vorticity()``
# using the wind components, grid deltas, and latitude values.
#
# Calculate grid spacing that is sign aware to use in absolute vorticity calculation
dx, dy = mpcalc.lat_lon_grid_deltas(lons, lats)
# Calculate absolute vorticity from MetPy function
avor_500 = mpcalc.absolute_vorticity(uwnd_500,
vwnd_500,
dx,
dy,
lats * units.degrees,
dim_order='yx')
######################################################################
# Map Creation
# ------------
#
# This next set of code creates the plot and draws contours on a Lambert
# Conformal map centered on -100 E longitude. The main view is over the
def Crosssection_Wind_Theta_e_Qv(
initial_time=None,
fhour=24,
levels=[1000, 950, 925, 900, 850, 800, 700, 600, 500, 400, 300, 200],
day_back=0,
model='ECMWF',
output_dir=None,
st_point=[20, 120.0],
ed_point=[50, 130.0],
map_extent=[70, 140, 15, 55],
h_pos=[0.125, 0.665, 0.25, 0.2]):
# micaps data directory
try:
data_dir = [
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='RH',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='UGRD',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='VGRD',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='TMP',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='HGT',
lvl='500')
]
except KeyError:
raise ValueError('Can not find all directories needed')
# get filename
if (initial_time != None):
filename = utl.model_filename(initial_time, fhour)
else:
filename = utl.filename_day_back_model(day_back=day_back, fhour=fhour)
# retrieve data from micaps server
rh = get_model_3D_grid(directory=data_dir[0][0:-1],
filename=filename,
levels=levels,
allExists=False)
if rh is None:
return
rh = rh.metpy.parse_cf().squeeze()
u = get_model_3D_grid(directory=data_dir[1][0:-1],
filename=filename,
levels=levels,
allExists=False)
if u is None:
return
u = u.metpy.parse_cf().squeeze()
v = get_model_3D_grid(directory=data_dir[2][0:-1],
filename=filename,
levels=levels,
allExists=False)
if v is None:
return
v = v.metpy.parse_cf().squeeze()
v2 = get_model_3D_grid(directory=data_dir[2][0:-1],
filename=filename,
levels=levels,
allExists=False)
if v2 is None:
return
v2 = v2.metpy.parse_cf().squeeze()
t = get_model_3D_grid(directory=data_dir[3][0:-1],
filename=filename,
levels=levels,
allExists=False)
if t is None:
return
t = t.metpy.parse_cf().squeeze()
gh = get_model_grid(data_dir[4], filename=filename)
if t is None:
return
resolution = u['lon'][1] - u['lon'][0]
x, y = np.meshgrid(u['lon'], u['lat'])
dx, dy = mpcalc.lat_lon_grid_deltas(u['lon'], u['lat'])
for ilvl in levels:
u2d = u.sel(level=ilvl)
#u2d['data'].attrs['units']=units.meter/units.second
v2d = v.sel(level=ilvl)
#v2d['data'].attrs['units']=units.meter/units.second
absv2d = mpcalc.absolute_vorticity(
u2d['data'].values * units.meter / units.second,
v2d['data'].values * units.meter / units.second, dx, dy,
y * units.degree)
if (ilvl == levels[0]):
absv3d = v2
absv3d['data'].loc[dict(level=ilvl)] = np.array(absv2d)
else:
absv3d['data'].loc[dict(level=ilvl)] = np.array(absv2d)
absv3d['data'].attrs['units'] = absv2d.units
#rh=rh.rename(dict(lat='latitude',lon='longitude'))
cross = cross_section(rh, st_point, ed_point)
cross_rh = cross.set_coords(('lat', 'lon'))
cross = cross_section(u, st_point, ed_point)
cross_u = cross.set_coords(('lat', 'lon'))
cross = cross_section(v, st_point, ed_point)
cross_v = cross.set_coords(('lat', 'lon'))
cross_u['data'].attrs['units'] = units.meter / units.second
cross_v['data'].attrs['units'] = units.meter / units.second
cross_u['t_wind'], cross_v['n_wind'] = mpcalc.cross_section_components(
cross_u['data'], cross_v['data'])
cross = cross_section(t, st_point, ed_point)
cross_t = cross.set_coords(('lat', 'lon'))
cross = cross_section(absv3d, st_point, ed_point)
cross_Td = mpcalc.dewpoint_rh(cross_t['data'].values * units.celsius,
cross_rh['data'].values * units.percent)
rh, pressure = xr.broadcast(cross_rh['data'], cross_t['level'])
Qv = mpcalc.specific_humidity_from_dewpoint(cross_Td, pressure)
cross_Qv = xr.DataArray(np.array(Qv) * 1000.,
coords=cross_rh['data'].coords,
dims=cross_rh['data'].dims,
attrs={'units': units('g/kg')})
Theta_e = mpcalc.equivalent_potential_temperature(
pressure, cross_t['data'].values * units.celsius, cross_Td)
cross_Theta_e = xr.DataArray(np.array(Theta_e),
coords=cross_rh['data'].coords,
dims=cross_rh['data'].dims,
attrs={'units': Theta_e.units})
crossection_graphics.draw_Crosssection_Wind_Theta_e_Qv(
cross_Qv=cross_Qv,
cross_Theta_e=cross_Theta_e,
cross_u=cross_u,
cross_v=cross_v,
gh=gh,
h_pos=h_pos,
st_point=st_point,
ed_point=ed_point,
levels=levels,
map_extent=map_extent,
output_dir=output_dir)
def PV_Div_uv(initTime=None,
fhour=6,
day_back=0,
model='ECMWF',
map_ratio=14 / 9,
zoom_ratio=20,
cntr_pnt=[104, 34],
levels=[
1000, 950, 925, 900, 850, 800, 700, 600, 500, 400, 300, 250,
200, 100
],
lvl_ana=250,
Global=False,
south_China_sea=True,
area=None,
city=False,
output_dir=None,
data_source='MICAPS',
**kwargs):
# micaps data directory
if (area != None):
south_China_sea = False
# micaps data directory
if (data_source == 'MICAPS'):
try:
data_dir = [
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='RH',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='UGRD',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='VGRD',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='TMP',
lvl=''),
utl.Cassandra_dir(data_type='high',
data_source=model,
var_name='HGT',
lvl='')
]
except KeyError:
raise ValueError('Can not find all directories needed')
# get filename
if (initTime != None):
filename = utl.model_filename(initTime, fhour)
else:
filename = utl.filename_day_back_model(day_back=day_back,
fhour=fhour)
# retrieve data from micaps server
rh = MICAPS_IO.get_model_3D_grid(directory=data_dir[0][0:-1],
filename=filename,
levels=levels,
allExists=False)
if rh is None:
return
u = MICAPS_IO.get_model_3D_grid(directory=data_dir[1][0:-1],
filename=filename,
levels=levels,
allExists=False)
if u is None:
return
v = MICAPS_IO.get_model_3D_grid(directory=data_dir[2][0:-1],
filename=filename,
levels=levels,
allExists=False)
if v is None:
return
t = MICAPS_IO.get_model_3D_grid(directory=data_dir[3][0:-1],
filename=filename,
levels=levels,
allExists=False)
if t is None:
return
if (data_source == 'CIMISS'):
# get filename
if (initTime != None):
filename = utl.model_filename(initTime, fhour, UTC=True)
else:
filename = utl.filename_day_back_model(day_back=day_back,
fhour=fhour,
UTC=True)
try:
# retrieve data from CIMISS server
rh = CMISS_IO.cimiss_model_3D_grid(
data_code=utl.CMISS_data_code(data_source=model,
var_name='RHU'),
init_time_str='20' + filename[0:8],
valid_time=fhour,
fcst_levels=levels,
fcst_ele="RHU",
units='%')
if rh is None:
return
u = CMISS_IO.cimiss_model_3D_grid(
data_code=utl.CMISS_data_code(data_source=model,
var_name='WIU'),
init_time_str='20' + filename[0:8],
valid_time=fhour,
fcst_levels=levels,
fcst_ele="WIU",
units='m/s')
if u is None:
return
v = CMISS_IO.cimiss_model_3D_grid(
data_code=utl.CMISS_data_code(data_source=model,
var_name='WIV'),
init_time_str='20' + filename[0:8],
valid_time=fhour,
fcst_levels=levels,
fcst_ele="WIV",
units='m/s')
if v is None:
return
t = CMISS_IO.cimiss_model_3D_grid(
data_code=utl.CMISS_data_code(data_source=model,
var_name='TEM'),
init_time_str='20' + filename[0:8],
valid_time=fhour,
fcst_levels=levels,
fcst_ele="TEM",
units='K')
if t is None:
return
t['data'].values = t['data'].values - 273.15
t['data'].attrs['units'] = 'C'
except KeyError:
raise ValueError('Can not find all data needed')
if (area != None):
cntr_pnt, zoom_ratio = utl.get_map_area(area_name=area)
map_extent = [0, 0, 0, 0]
map_extent[0] = cntr_pnt[0] - zoom_ratio * 1 * map_ratio
map_extent[1] = cntr_pnt[0] + zoom_ratio * 1 * map_ratio
map_extent[2] = cntr_pnt[1] - zoom_ratio * 1
map_extent[3] = cntr_pnt[1] + zoom_ratio * 1
delt_x = (map_extent[1] - map_extent[0]) * 0.2
delt_y = (map_extent[3] - map_extent[2]) * 0.1
#+ to solve the problem of labels on all the contours
mask1 = (rh['lon'] >
map_extent[0] - delt_x) & (rh['lon'] < map_extent[1] + delt_x) & (
rh['lat'] > map_extent[2] - delt_y) & (rh['lat'] <
map_extent[3] + delt_y)
mask2 = (u['lon'] >
map_extent[0] - delt_x) & (u['lon'] < map_extent[1] + delt_x) & (
u['lat'] > map_extent[2] - delt_y) & (u['lat'] <
map_extent[3] + delt_y)
mask3 = (t['lon'] >
map_extent[0] - delt_x) & (t['lon'] < map_extent[1] + delt_x) & (
t['lat'] > map_extent[2] - delt_y) & (t['lat'] <
map_extent[3] + delt_y)
#- to solve the problem of labels on all the contours
rh = rh.where(mask1, drop=True)
u = u.where(mask2, drop=True)
v = v.where(mask2, drop=True)
t = t.where(mask3, drop=True)
uv = xr.merge([u.rename({'data': 'u'}), v.rename({'data': 'v'})])
lats = np.squeeze(rh['lat'].values)
lons = np.squeeze(rh['lon'].values)
pres = np.array(levels) * 100 * units('Pa')
tmpk = mpcalc.smooth_n_point(
(t['data'].values.squeeze() + 273.15), 9, 2) * units('kelvin')
thta = mpcalc.potential_temperature(pres[:, None, None], tmpk)
uwnd = mpcalc.smooth_n_point(u['data'].values.squeeze(), 9,
2) * units.meter / units.second
vwnd = mpcalc.smooth_n_point(v['data'].values.squeeze(), 9,
2) * units.meter / units.second
dx, dy = mpcalc.lat_lon_grid_deltas(lons, lats)
# Comput the PV on all isobaric surfaces
pv_raw = mpcalc.potential_vorticity_baroclinic(
thta, pres[:, None, None], uwnd, vwnd, dx[None, :, :], dy[None, :, :],
lats[None, :, None] * units('degrees'))
div_raw = mpcalc.divergence(uwnd,
vwnd,
dx[None, :, :],
dy[None, :, :],
dim_order='yx')
# prepare data
idx_z1 = list(pres.m).index(((lvl_ana * units('hPa')).to(pres.units)).m)
pv = rh.copy(deep=True)
pv['data'].values = np.array(pv_raw).reshape(
np.append(1,
np.array(pv_raw).shape))
pv['data'].attrs['units'] = str(pv_raw.units)
pv.attrs['model'] = model
pv = pv.where(pv['level'] == lvl_ana, drop=True)
div = u.copy(deep=True)
div['data'].values = np.array(div_raw).reshape(
np.append(1,
np.array(div_raw).shape))
div['data'].attrs['units'] = str(div_raw.units)
div = div.where(div['level'] == lvl_ana, drop=True)
uv = uv.where(uv['level'] == lvl_ana, drop=True)
synoptic_graphics.draw_PV_Div_uv(pv=pv,
uv=uv,
div=div,
map_extent=map_extent,
regrid_shape=20,
city=city,
south_China_sea=south_China_sea,
output_dir=output_dir,
Global=Global)
def draw_vort_high(uwind,
vwind,
lon,
lat,
vort=None,
gh=None,
skip_vector=None,
smooth_factor=1.0,
map_region=None,
title_kwargs={},
outfile=None):
"""
Draw high vorticity.
Args:
uwind (np.array): u wind component, 2D array, [nlat, nlon]
vwind (np.array): v wind component, 2D array, [nlat, nlon]
lon (np.array): longitude, 1D array, [nlon]
lat (np.array): latitude, 1D array, [nlat]
vort (np.array, optional): vorticity component, 2D array, [nlat, nlon]
gh (np.array): geopotential height, 2D array, [nlat, nlon]
skip_vector (integer): skip grid number for vector plot
smooth_factor (float): smooth factor for vorticity, larger for smoother.
map_region (list or tuple): the map region limit, [lonmin, lonmax, latmin, latmax]
title_kwargs (dictionaly, optional): keyword arguments for _get_title function.
"""
# check default parameters
if skip_vector is None:
skip_vector = util.get_skip_vector(lon, lat, map_region)
# put data into fields
wind_field = util.minput_2d_vector(uwind,
vwind,
lon,
lat,
skip=skip_vector)
if vort is None:
dx, dy = calc.lat_lon_grid_deltas(lon, lat)
vort = calc.vorticity(uwind * units.meter / units.second,
vwind * units.meter / units.second, dx, dy)
vort = ndimage.gaussian_filter(vort, sigma=smooth_factor,
order=0) * 10**5
vort_field = util.minput_2d(vort, lon, lat, {
'long_name': 'vorticity',
'units': 's-1'
})
if gh is not None:
gh_feild = util.minput_2d(gh, lon, lat, {
'long_name': 'height',
'units': 'gpm'
})
#
# set up visual parameters
#
plots = []
# Setting the coordinates of the geographical area
if map_region is None:
china_map = map_set.get_mmap(name='CHINA_CYLINDRICAL',
subpage_frame_thickness=5)
else:
china_map = map_set.get_mmap(name='CHINA_REGION_CYLINDRICAL',
map_region=map_region,
subpage_frame_thickness=5)
plots.append(china_map)
# Background Coaslines
coastlines = map_set.get_mcoast(name='COAST_FILL')
plots.append(coastlines)
# Define the shading contour
vort_contour = magics.mcont(
legend='on',
contour_level_selection_type='level_list',
contour_level_list=[
-200.0, -100.0, -75.0, -50.0, -30.0, -20.0, -15.0, -13.0, -11.0,
-9.0, -7.0, -5.0, -3.0, -1.0, 1.0, 3.0, 5.0, 7.0, 9.0, 11.0, 13.0,
15.0, 20.0, 30.0, 50.0, 75.0, 100.0, 200.0
],
contour_shade='on',
contour="off",
contour_shade_method='area_fill',
contour_shade_colour_method='list',
contour_shade_colour_list=[
"rgb(0,0,0.3)", "rgb(0,0,0.5)", "rgb(0,0,0.7)", "rgb(0,0,0.9)",
"rgb(0,0.15,1)", "rgb(0,0.3,1)", "rgb(0,0.45,1)", "rgb(0,0.6,1)",
"rgb(0,0.75,1)", "rgb(0,0.85,1)", "rgb(0.2,0.95,1)",
"rgb(0.45,1,1)", "rgb(0.75,1,1)", "none", "rgb(1,1,0)",
"rgb(1,0.9,0)", "rgb(1,0.8,0)", "rgb(1,0.7,0)", "rgb(1,0.6,0)",
"rgb(1,0.5,0)", "rgb(1,0.4,0)", "rgb(1,0.3,0)", "rgb(1,0.15,0)",
"rgb(0.9,0,0)", "rgb(0.7,0,0)", "rgb(0.5,0,0)", "rgb(0.3,0,0)"
],
contour_reference_level=8.,
contour_highlight='off',
contour_hilo='off',
contour_label='off')
plots.extend([vort_field, vort_contour])
# Define the wind vector
wind_vector = common._get_wind_flags()
plots.extend([wind_field, wind_vector])
# Define the simple contouring for gh
if gh is not None:
gh_contour = common._get_gh_contour()
plots.extend([gh_feild, gh_contour])
# Add a legend
legend = common._get_legend(china_map, title="Vorticity [s^-1]")
plots.append(legend)
# Add the title
title_kwargs = check_kwargs(title_kwargs, 'head',
"500hPa Relative vorticity | Wind | GH")
title = common._get_title(**title_kwargs)
plots.append(title)
# Add china province
china_coastlines = map_set.get_mcoast(name='PROVINCE')
plots.append(china_coastlines)
# final plot
return util.magics_plot(plots, outfile)
winds = winds.transpose('month', 'realiz', 'latitude', 'longitude')
winds_clm = winds.mean(dim='realiz')
for i in np.arange(7):
var_WPV_EN = np.mean(hgt.z.values[i, index_WPV_EN, :, :], axis=0)
var_WPV_LN = np.mean(hgt.z.values[i, index_WPV_LN, :, :], axis=0)
var_normal = np.mean(hgt.z.values[i, :, :, :], axis=0)
var_SPV_EN = np.mean(hgt.z.values[i, index_SPV_EN, :, :], axis=0)
var_SPV_LN = np.mean(hgt.z.values[i, index_SPV_LN, :, :], axis=0)
var_WPV_all = np.mean(hgt.z.values[i, index_WPV_all.values, :, :], axis=0)
var_SPV_all = np.mean(hgt.z.values[i, index_SPV_all.values, :, :], axis=0)
px_WPV, py_WPV, lat = plumb_flux.ComputePlumbFluxes(
winds_clm.u.values[i, :, :],
winds_clm.v.values[i, :, :], var_WPV_all - var_normal,
np.zeros_like(var_WPV_EN), hgt.latitude.values, hgt.longitude.values)
dx, dy = calc.lat_lon_grid_deltas(hgt.longitude.values, lat)
div_WPV_all = calc.divergence(px_WPV, py_WPV, dx, dy)
px_SPV, py_SPV, lat = plumb_flux.ComputePlumbFluxes(
winds_clm.u.values[i, :, :],
winds_clm.v.values[i, :, :], var_SPV_all - var_normal,
np.zeros_like(var_WPV_EN), hgt.latitude.values, hgt.longitude.values)
div_SPV_all = calc.divergence(px_SPV, py_SPV, dx, dy)
px_WPV_EN, py_WPV_EN, lat = plumb_flux.ComputePlumbFluxes(
winds_clm.u.values[i, :, :],
winds_clm.v.values[i, :, :], var_WPV_EN - var_normal,
np.zeros_like(var_WPV_EN), hgt.latitude.values, hgt.longitude.values)
div_WPV_EN = calc.divergence(px_WPV_EN, py_WPV_EN, dx, dy)
px_SPV_EN, py_SPV_EN, lat = plumb_flux.ComputePlumbFluxes(
# Calculations
# ------------
#
# Most of the calculations in `metpy.calc` will accept DataArrays by converting them
# into their corresponding unit arrays. While this may often work without any issues, we must
# keep in mind that because the calculations are working with unit arrays and not DataArrays:
#
# - The calculations will return unit arrays rather than DataArrays
# - Broadcasting must be taken care of outside of the calculation, as it would only recognize
# dimensions by order, not name
#
# As an example, we calculate geostropic wind at 500 hPa below:
lat, lon = xr.broadcast(y, x)
f = mpcalc.coriolis_parameter(lat)
dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat, initstring=data_crs.proj4_init)
heights = data['height'].loc[time[0]].loc[{vertical.name: 500.}]
u_geo, v_geo = mpcalc.geostrophic_wind(heights, f, dx, dy)
print(u_geo)
print(v_geo)
#########################################################################
# Also, a limited number of calculations directly support xarray DataArrays or Datasets (they
# can accept *and* return xarray objects). Right now, this includes
#
# - Derivative functions
# - ``first_derivative``
# - ``second_derivative``
# - ``gradient``
# - ``laplacian``
# - Cross-section functions
extend=extend,
transform=proj_ccrs)
#Add a color bar
cbar = plt.colorbar(cs, cax=ax3, shrink=0.75, pad=0.01, ticks=clevs)
print("Filled contours for 250-hPa wind")
#--------------------------------------------------------------------------------------------------------
# 500-hPa smoothed vorticity
#--------------------------------------------------------------------------------------------------------
#Get the data for this variable
u = data['u'].sel(lev=500)
v = data['v'].sel(lev=500)
dx, dy = calc.lat_lon_grid_deltas(lon, lat)
vort = calc.vorticity(u, v, dx, dy)
smooth_vort = smooth(vort, 5.0) * 10**5
#Specify contour settings
clevs = np.arange(2, 20, 1)
cmap = plt.cm.autumn_r
extend = "max"
#Contour fill this variable
norm = col.BoundaryNorm(clevs, cmap.N)
cs = ax.contourf(lon,
lat,
smooth_vort,
clevs,
cmap=cmap,
# Nearly all of the calculations in `metpy.calc` will accept DataArrays by converting them
# into their corresponding unit arrays. While this may often work without any issues, we must
# keep in mind that because the calculations are working with unit arrays and not DataArrays:
#
# - The calculations will return unit arrays rather than DataArrays
# - Broadcasting must be taken care of outside of the calculation, as it would only recognize
# dimensions by order, not name
#
# Also, some of the units used in CF conventions (such as 'degrees_north') are not recognized
# by pint, so we must implement a workaround.
#
# As an example, we calculate geostropic wind at 500 hPa below:
lat, lon = xr.broadcast(y, x)
f = mpcalc.coriolis_parameter(lat.values * units.degrees)
dx, dy = mpcalc.lat_lon_grid_deltas(lon.values, lat.values)
heights = data['height'].loc[time[0]].loc[{vertical.name: 500.}]
u_geo, v_geo = mpcalc.geostrophic_wind(heights, f, dx, dy, dim_order='yx')
print(u_geo)
print(v_geo)
#########################################################################
# Plotting
# --------
#
# Like most meteorological data, we want to be able to plot these data. DataArrays can be used
# like normal numpy arrays in plotting code, or we can use some of xarray's plotting
# functionality.
#
# (More detail beyond the following can be found at `xarray's plotting reference
# <http://xarray.pydata.org/en/stable/plotting.html>`_.)
def metpy_read_wrf_cross(fname, plevels, tidx_in, lons_out, lats_out, start,
end):
ds = xr.open_dataset(fname).metpy.parse_cf().squeeze()
ds = ds.isel(Time=tidx)
print(ds)
ds1 = xr.Dataset()
p = units.Quantity(to_np(ds.p), 'hPa')
z = units.Quantity(to_np(ds.z), 'meter')
u = units.Quantity(to_np(ds.u), 'm/s')
v = units.Quantity(to_np(ds.v), 'm/s')
w = units.Quantity(to_np(ds.w), 'm/s')
tk = units.Quantity(to_np(ds.tk), 'kelvin')
th = units.Quantity(to_np(ds.th), 'kelvin')
eth = units.Quantity(to_np(ds.eth), 'kelvin')
wspd = units.Quantity(to_np(ds.wspd), 'm/s')
omega = units.Quantity(to_np(ds.omega), 'Pa/s')
plevels_unit = plevels * units.hPa
z, u, v, w, tk, th, eth, wspd, omega = log_interpolate_1d(plevels_unit,
p,
z,
u,
v,
w,
tk,
th,
eth,
wspd,
omega,
axis=0)
coords, dims = [plevs, ds.lat.values,
ds.lon.values], ["level", "lat", "lon"]
for name, var in zip(
['z', 'u', 'v', 'w', 'tk', 'th', 'eth', 'wspd', 'omega'],
[z, u, v, w, tk, th, eth, wspd, omega]):
#g = ndimage.gaussian_filter(var, sigma=3, order=0)
ds1[name] = xr.DataArray(to_np(mpcalc.smooth_n_point(var, 9)),
coords=coords,
dims=dims)
dx, dy = mpcalc.lat_lon_grid_deltas(ds.lon.values * units('degrees_E'),
ds.lat.values * units('degrees_N'))
# Calculate temperature advection using metpy function
for i, plev in enumerate(plevs):
adv = mpcalc.advection(eth[i, :, :], [u[i, :, :], v[i, :, :]],
(dx, dy),
dim_order='yx') * units('K/sec')
adv = ndimage.gaussian_filter(adv, sigma=3, order=0) * units('K/sec')
ds1['eth_adv_{:03d}'.format(plev)] = xr.DataArray(
np.array(adv),
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
div = mpcalc.divergence(u[i, :, :], v[i, :, :], dx, dy, dim_order='yx')
div = ndimage.gaussian_filter(div, sigma=3, order=0) * units('1/sec')
ds1['div_{:03d}'.format(plev)] = xr.DataArray(
np.array(div),
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
ds1['accrain'] = xr.DataArray(ds.accrain.values,
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
eth2 = mpcalc.equivalent_potential_temperature(
ds.slp.values * units.hPa, ds.t2m.values * units('K'),
ds.td2.values * units('celsius'))
ds1['eth2'] = xr.DataArray(eth2,
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
#ds1['sst'] = xr.DataArray(ndimage.gaussian_filter(ds.sst.values, sigma=3, order=0)-273.15, coords=[ds.lat.values,ds.lon.values], dims=["lat","lon"])
ds1 = ds1.metpy.parse_cf().squeeze()
cross = cross_section(ds1, start, end).set_coords(('lat', 'lon'))
cross.u.attrs['units'] = 'm/s'
cross.v.attrs['units'] = 'm/s'
cross['t_wind'], cross['n_wind'] = mpcalc.cross_section_components(
cross['u'], cross['v'])
weights = np.cos(np.deg2rad(ds.lat))
ds_weighted = ds.weighted(weights)
weighted = ds_weighted.mean(("lat"))
return ds1, cross, weighted
v1 = analysis.variables["VGRD_P0_L100_GLL0"][lev_index, :, :]
v = v1[lat_box1:lat_box2, lon_box1:lon_box2]
del v1
# create 2d lat and lon
lat2d = np.zeros((len(lat), len(lon)))
lon2d = np.zeros((len(lat), len(lon)))
for i in range(0, len(lon)):
lat2d[:, i] = lat
for i in range(0, len(lat)):
lon2d[i, :] = lon
dx, dy = mpcalc.lat_lon_grid_deltas(lon2d, lat2d)
div1 = mpcalc.divergence(u, v, dx, dy)
div = np.array(div1)
# open workspace for analysis plot
wks_type = "png"
wks = ngl.open_wks(
wks_type,
"GFSanalysis_%s_%s_divergence_%shPa" % (region, init_dt[0:10], lev_hPa))
# define resources for analysis plot
res = ngl.Resources()
def main():
load_start = dt.datetime.now()
#Try parsing arguments using argparse
parser = argparse.ArgumentParser(description='wrf non-parallel convective diagnostics processer')
parser.add_argument("-m",help="Model name",required=True)
parser.add_argument("-r",help="Region name (default is aus)",default="aus")
parser.add_argument("-t1",help="Time start YYYYMMDDHH",required=True)
parser.add_argument("-t2",help="Time end YYYYMMDDHH",required=True)
parser.add_argument("-e", help="CMIP5 experiment name (not required if using era5, erai or barra)", default="")
parser.add_argument("--ens", help="CMIP5 ensemble name (not required if using era5, erai or barra)", default="r1i1p1")
parser.add_argument("--group", help="CMIP6 modelling group name", default="")
parser.add_argument("--project", help="CMIP6 modelling intercomparison project", default="CMIP")
parser.add_argument("--ver6hr", help="Version on al33 for 6hr data", default="")
parser.add_argument("--ver3hr", help="Version on al33 for 3hr data", default="")
parser.add_argument("--issave",help="Save output (True or False, default is False)", default="False")
parser.add_argument("--outname",help="Name of saved output. In the form *outname*_*t1*_*t2*.nc. Default behaviour is the model name",default=None)
parser.add_argument("--al33",help="Should data be gathered from al33? Default is False, and data is gathered from r87. If True, then group is required",default="False")
args = parser.parse_args()
#Parse arguments from cmd line and set up inputs (date region model)
model = args.m
region = args.r
t1 = args.t1
t2 = args.t2
issave = args.issave
al33 = args.al33
if args.outname==None:
out_name = model
else:
out_name = args.outname
experiment = args.e
ensemble = args.ens
group = args.group
project = args.project
ver6hr = args.ver6hr
ver3hr = args.ver3hr
if region == "sa_small":
start_lat = -38; end_lat = -26; start_lon = 132; end_lon = 142
elif region == "aus":
start_lat = -44.525; end_lat = -9.975; start_lon = 111.975; end_lon = 156.275
elif region == "global":
start_lat = -70; end_lat = 70; start_lon = -180; end_lon = 179.75
else:
raise ValueError("INVALID REGION\n")
domain = [start_lat,end_lat,start_lon,end_lon]
try:
time = [dt.datetime.strptime(t1,"%Y%m%d%H"),dt.datetime.strptime(t2,"%Y%m%d%H")]
except:
raise ValueError("INVALID START OR END TIME. SHOULD BE YYYYMMDDHH\n")
if issave=="True":
issave = True
elif issave=="False":
issave = False
else:
raise ValueError("\n INVALID ISSAVE...SHOULD BE True OR False")
if al33=="True":
al33 = True
elif al33=="False":
al33 = False
else:
raise ValueError("\n INVALID al33...SHOULD BE True OR False")
#Load data
print("LOADING DATA...")
if model in ["ACCESS1-0","ACCESS1-3","GFDL-CM3","GFDL-ESM2M","CNRM-CM5","MIROC5",\
"MRI-CGCM3","IPSL-CM5A-LR","IPSL-CM5A-MR","GFDL-ESM2G","bcc-csm1-1","MIROC-ESM",\
"BNU-ESM"]:
#Check that t1 and t2 are in the same year
year = np.arange(int(t1[0:4]), int(t2[0:4])+1)
ta, hur, hgt, terrain, p_3d, ps, ua, va, uas, vas, tas, ta2d, tp, lon, lat, \
date_list = read_cmip(model, experiment, \
ensemble, year, domain, cmip_ver=5, al33=al33, group=group, ver6hr=ver6hr, ver3hr=ver3hr)
p = np.zeros(p_3d[0,:,0,0].shape)
tp = tp.astype("float32", order="C")
elif model in ["ACCESS-ESM1-5", "ACCESS-CM2"]:
year = np.arange(int(t1[0:4]), int(t2[0:4])+1)
ta, hur, hgt, terrain, p_3d, ps, ua, va, uas, vas, tas, ta2d, lon, lat, \
date_list = read_cmip(model, experiment,\
ensemble, year, domain, cmip_ver=6, group=group, project=project)
p = np.zeros(p_3d[0,:,0,0].shape)
else:
raise ValueError("Model not recognised")
ta = ta.astype("float32", order="C")
hur = hur.astype("float32", order="C")
hgt = hgt.astype("float32", order="C")
terrain = terrain.astype("float32", order="C")
p = p.astype("float32", order="C")
ps = ps.astype("float32", order="C")
ua = ua.astype("float32", order="C")
va = va.astype("float32", order="C")
uas = uas.astype("float32", order="C")
vas = vas.astype("float32", order="C")
tas= tas.astype("float32", order="C")
ta2d = ta2d.astype("float32", order="C")
lon = lon.astype("float32", order="C")
lat = lat.astype("float32", order="C")
gc.collect()
#This param list was originally given to AD for ERA5 global lightning report
#param = np.array(["mu_cape", "eff_cape","ncape","mu_cin", "muq", "s06", "s0500", "lr700_500", "mhgt", "ta500","tp","cp","laplacian","t_totals"])
#This param list is intended for application to GCMs based on AD ERA5 lightning report
param = np.array(["mu_cape","s06","laplacian","t_totals","ta850","ta500","dp850","tp",\
"muq","lr700_500","mhgt","z500"])
#Set output array
output_data = np.zeros((ps.shape[0], ps.shape[1], ps.shape[2], len(param)))
#Assign p levels to a 3d array, with same dimensions as input variables (ta, hgt, etc.)
#If the 3d p-lvl array already exists, then declare the variable "mdl_lvl" as true.
try:
p_3d;
mdl_lvl = True
full_p3d = p_3d
except:
mdl_lvl = False
p_3d = np.moveaxis(np.tile(p,[ta.shape[2],ta.shape[3],1]),[0,1,2],[1,2,0]).\
astype(np.float32)
print("LOAD TIME..."+str(dt.datetime.now()-load_start))
tot_start = dt.datetime.now()
for t in np.arange(0,ta.shape[0]):
output = np.zeros((1, ps.shape[1], ps.shape[2], len(param)))
cape_start = dt.datetime.now()
print(date_list[t])
if mdl_lvl:
p_3d = full_p3d[t]
dp = get_dp(hur=hur[t], ta=ta[t], dp_mask = False)
#Insert surface arrays, creating new arrays with "sfc" prefix
sfc_ta = np.insert(ta[t], 0, tas[t], axis=0)
sfc_hgt = np.insert(hgt[t], 0, terrain, axis=0)
sfc_dp = np.insert(dp, 0, ta2d[t], axis=0)
sfc_p_3d = np.insert(p_3d, 0, ps[t], axis=0)
sfc_ua = np.insert(ua[t], 0, uas[t], axis=0)
sfc_va = np.insert(va[t], 0, vas[t], axis=0)
#Sort by ascending p
a,temp1,temp2 = np.meshgrid(np.arange(sfc_p_3d.shape[0]) , np.arange(sfc_p_3d.shape[1]),\
np.arange(sfc_p_3d.shape[2]))
sort_inds = np.flip(np.lexsort([np.swapaxes(a,1,0),sfc_p_3d],axis=0), axis=0)
sfc_hgt = np.take_along_axis(sfc_hgt, sort_inds, axis=0)
sfc_dp = np.take_along_axis(sfc_dp, sort_inds, axis=0)
sfc_p_3d = np.take_along_axis(sfc_p_3d, sort_inds, axis=0)
sfc_ua = np.take_along_axis(sfc_ua, sort_inds, axis=0)
sfc_va = np.take_along_axis(sfc_va, sort_inds, axis=0)
sfc_ta = np.take_along_axis(sfc_ta, sort_inds, axis=0)
#Calculate q and wet bulb for pressure level arrays with surface values
sfc_ta_unit = units.units.degC*sfc_ta
sfc_dp_unit = units.units.degC*sfc_dp
sfc_p_unit = units.units.hectopascals*sfc_p_3d
sfc_hur_unit = mpcalc.relative_humidity_from_dewpoint(sfc_ta_unit, sfc_dp_unit)*\
100*units.units.percent
sfc_q_unit = mpcalc.mixing_ratio_from_relative_humidity(sfc_hur_unit,\
sfc_ta_unit,sfc_p_unit)
sfc_q = np.array(sfc_q_unit)
#Now get most-unstable CAPE (max CAPE in vertical, ensuring parcels used are AGL)
cape3d = wrf.cape_3d(sfc_p_3d,sfc_ta+273.15,\
sfc_q,sfc_hgt,\
terrain,ps[t],\
True,meta=False, missing=0)
cape = cape3d.data[0]
cin = cape3d.data[1]
lfc = cape3d.data[2]
lcl = cape3d.data[3]
el = cape3d.data[4]
#Mask values which are below the surface and above 350 hPa AGL
cape[(sfc_p_3d > ps[t]) | (sfc_p_3d<(ps[t]-350))] = np.nan
cin[(sfc_p_3d > ps[t]) | (sfc_p_3d<(ps[t]-350))] = np.nan
lfc[(sfc_p_3d > ps[t]) | (sfc_p_3d<(ps[t]-350))] = np.nan
lcl[(sfc_p_3d > ps[t]) | (sfc_p_3d<(ps[t]-350))] = np.nan
el[(sfc_p_3d > ps[t]) | (sfc_p_3d<(ps[t]-350))] = np.nan
#Get maximum (in the vertical), and get cin, lfc, lcl for the same parcel
mu_cape_inds = np.tile(np.nanargmax(cape,axis=0), (cape.shape[0],1,1))
mu_cape = np.take_along_axis(cape, mu_cape_inds, 0)[0]
muq = np.take_along_axis(sfc_q, mu_cape_inds, 0)[0] * 1000
#Calculate other parameters
#Thermo
thermo_start = dt.datetime.now()
lr700_500 = get_lr_p(ta[t], p_3d, hgt[t], 700, 500)
melting_hgt = get_t_hgt(sfc_ta,np.copy(sfc_hgt),0,terrain)
melting_hgt = np.where((melting_hgt < 0) | (np.isnan(melting_hgt)), 0, melting_hgt)
ta500 = get_var_p_lvl(np.copy(sfc_ta), sfc_p_3d, 500)
ta850 = get_var_p_lvl(np.copy(sfc_ta), sfc_p_3d, 850)
dp850 = get_var_p_lvl(np.copy(sfc_dp), sfc_p_3d, 850)
v_totals = ta850 - ta500
c_totals = dp850 - ta500
t_totals = v_totals + c_totals
#Winds
winds_start = dt.datetime.now()
s06 = get_shear_hgt(sfc_ua, sfc_va, np.copy(sfc_hgt), 0, 6000, terrain)
#Laplacian
x, y = np.meshgrid(lon,lat)
dx, dy = mpcalc.lat_lon_grid_deltas(x,y)
if mdl_lvl:
z500 = get_var_p_lvl(hgt[t], p_3d, 500)
laplacian = np.array(mpcalc.laplacian(z500,deltas=[dy,dx])*1e9)
else:
z500 = np.squeeze(hgt[t,p==500])
laplacian = np.array(mpcalc.laplacian(uniform_filter(z500, 4),deltas=[dy,dx])*1e9)
#Fill output
output = fill_output(output, t, param, ps, "mu_cape", mu_cape)
output = fill_output(output, t, param, ps, "muq", muq)
output = fill_output(output, t, param, ps, "s06", s06)
output = fill_output(output, t, param, ps, "lr700_500", lr700_500)
output = fill_output(output, t, param, ps, "ta500", ta500)
output = fill_output(output, t, param, ps, "ta850", ta850)
output = fill_output(output, t, param, ps, "dp850", dp850)
output = fill_output(output, t, param, ps, "mhgt", melting_hgt)
output = fill_output(output, t, param, ps, "tp", tp[t])
output = fill_output(output, t, param, ps, "laplacian", laplacian)
output = fill_output(output, t, param, ps, "t_totals", t_totals)
output = fill_output(output, t, param, ps, "z500", z500)
output_data[t] = output
print("SAVING DATA...")
param_out = []
for param_name in param:
temp_data = output_data[:,:,:,np.where(param==param_name)[0][0]]
param_out.append(temp_data)
#If the mhgt variable is zero everywhere, then it is likely that data has not been read.
#In this case, all values are missing, set to zero.
for t in np.arange(param_out[0].shape[0]):
if param_out[np.where(param=="mhgt")[0][0]][t].max() == 0:
for p in np.arange(len(param_out)):
param_out[p][t] = np.nan
if issave:
save_netcdf(region, model, out_name, date_list, lat, lon, param, param_out, \
out_dtype = "f4", compress=True)
print(dt.datetime.now() - tot_start)
def horizontal_map(variable_name, date, start_hour,
end_hour, pressure_level=False,
subset=False, initiation=False,
save=False, gif=False):
'''This function plots the chosen variable for the analysis
of the initiation environment on a horizontal (2D) map. Supported variables for plotting
procedure are updraft, reflectivity, helicity, pw, cape, cin, ctt, temperature_surface,
wind_shear, updraft_reflectivity, rh, omega, pvo, avo, theta_e, water_vapor, uv_wind and
divergence.'''
### Predefine some variables ###
# Get the list of all needed wrf files
data_dir = '/scratch3/thomasl/work/data/casestudy_baden/'
# Define save directory
save_dir = '/scratch3/thomasl/work/retrospective_part' '/casestudy_baden/horizontal_maps/'
# Change extent of plot
subset_extent = [6.2, 9.4, 46.5, 48.5]
# Set the location of the initiation of the thunderstorm
initiation_location = CoordPair(lat=47.25, lon=7.85)
# 2D variables:
if variable_name == 'updraft':
variable_name = 'W_UP_MAX'
title_name = 'Maximum Z-Wind Updraft'
colorbar_label = 'Max Z-Wind Updraft [$m$ $s^-$$^1$]'
save_name = 'updraft'
variable_min = 0
variable_max = 30
# Check if a certain pressure_level was defined.
if pressure_level != False:
sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name))
elif variable_name == 'reflectivity':
variable_name = 'REFD_MAX'
title_name = 'Maximum Derived Radar Reflectivity'
colorbar_label = 'Maximum Derived Radar Reflectivity [$dBZ$]'
save_name = 'reflectivity'
variable_min = 0
variable_max = 75
# Check if a certain pressure_level was defined.
if pressure_level != False:
sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name))
elif variable_name == 'helicity':
variable_name = 'UP_HELI_MAX'
title_name = 'Maximum Updraft Helicity'
colorbar_label = 'Maximum Updraft Helicity [$m^{2}$ $s^{-2}$]'
save_name = 'helicity'
variable_min = 0
variable_max = 140
# Check if a certain pressure_level was defined.
if pressure_level != False:
sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name))
elif variable_name == 'pw':
title_name = 'Precipitable Water'
colorbar_label = 'Precipitable Water [$kg$ $m^{-2}$]'
save_name = 'pw'
variable_min = 0
variable_max = 50
# Check if a certain pressure_level was defined.
if pressure_level != False:
sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name))
elif variable_name == 'cape':
variable_name = 'cape_2d'
title_name = 'CAPE'
colorbar_label = 'Convective Available Potential Energy' '[$J$ $kg^{-1}$]'
save_name = 'cape'
variable_min = 0
variable_max = 3000
# Check if a certain pressure_level was defined.
if pressure_level != False:
sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name))
elif variable_name == 'cin':
variable_name = 'cape_2d'
title_name = 'CIN'
colorbar_label = 'Convective Inhibition [$J$ $kg^{-1}$]'
save_name = 'cin'
variable_min = 0
variable_max = 100
# Check if a certain pressure_level was defined.
if pressure_level != False:
sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name))
elif variable_name == 'ctt':
title_name = 'Cloud Top Temperature'
colorbar_label = 'Cloud Top Temperature [$K$]'
save_name = 'cct'
variable_min = 210
variable_max = 300
# Check if a certain pressure_level was defined.
if pressure_level != False:
sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name))
elif variable_name == 'temperature_surface':
variable_name = 'T2'
title_name = 'Temperature @ 2 m'
colorbar_label = 'Temperature [$K$]'
save_name = 'temperature_surface'
variable_min = 285
variable_max = 305
# Check if a certain pressure_level was defined.
if pressure_level != False:
sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name))
elif variable_name == 'wind_shear':
variable_name = 'slp'
title_name = 'SLP, Wind @ 850hPa, Wind @ 500hPa\n' 'and 500-850hPa Vertical Wind Shear'
save_name = 'wind_shear'
variable_min = 1000
variable_max = 1020
# Check if a certain pressure_level was defined.
if pressure_level != False:
sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name))
elif variable_name == 'updraft_reflectivity':
variable_name = 'W_UP_MAX'
title_name = 'Updraft and Reflectivity'
colorbar_label = 'Max Z-Wind Updraft [$m$ $s^-$$^1$]'
save_name = 'updraft_reflectivity'
variable_min = 0
variable_max = 30
# Check if a certain pressure_level was defined.
if pressure_level != False:
sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name))
# 3D variables:
elif variable_name == 'rh':
title_name = 'Relative Humidity'
colorbar_label = 'Relative Humidity [$pct$]'
save_name = 'rh'
variable_min = 0
variable_max = 100
# Check if a certain pressure_level was defined.
if pressure_level == False:
sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name))
elif variable_name == 'omega':
title_name = 'Vertical Motion'
colorbar_label = 'Omega [$Pa$ $s^-$$^1$]'
save_name = 'omega'
variable_min = -50
variable_max = 50
# Check if a certain pressure_level was defined.
if pressure_level == False:
sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name))
elif variable_name == 'pvo':
title_name = 'Potential Vorticity'
colorbar_label = 'Potential Vorticity [$PVU$]'
save_name = 'pvo'
variable_min = -1
variable_max = 9
# Check if a certain pressure_level was defined.
if pressure_level == False:
sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name))
elif variable_name == 'avo':
title_name = 'Absolute Vorticity'
colorbar_label = 'Absolute Vorticity [$10^{-5}$' '$s^{-1}$]'
save_name = 'avo'
variable_min = -250
variable_max = 250
# Check if a certain pressure_level was defined.
if pressure_level == False:
sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name))
elif variable_name == 'theta_e':
title_name = 'Theta-E'
colorbar_label = 'Theta-E [$K$]'
save_name = 'theta_e'
variable_min = 315
variable_max = 335
# Check if a certain pressure_level was defined.
if pressure_level == False:
sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name))
elif variable_name == 'water_vapor':
variable_name = 'QVAPOR'
title_name = 'Water Vapor Mixing Ratio'
colorbar_label = 'Water Vapor Mixing Ratio [$g$ $kg^{-1}$]'
save_name = 'water_vapor'
variable_min = 5
variable_max = 15
# Check if a certain pressure_level was defined.
if pressure_level == False:
sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name))
elif variable_name == 'uv_wind':
variable_name = 'wspd_wdir'
title_name = 'Wind Speed and Direction'
colorbar_label = 'Wind Speed [$m$ $s^{-1}$]'
save_name = 'uv_wind'
variable_min = 0
variable_max = 10
# Check if a certain pressure_level was defined.
if pressure_level == False:
sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name))
elif variable_name == 'divergence':
variable_name = 'ua'
title_name = 'Horizontal Wind Divergence'
colorbar_label = 'Divergence [$10^{-6}$ $s^{-1}$]'
save_name = 'divergence'
variable_min = -2.5
variable_max = 2.5
# Check if a certain pressure_level was defined.
if pressure_level == False:
sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name))
# Make a list of all wrf files in data directory
wrflist = list()
for (dirpath, dirnames, filenames) in os.walk(data_dir):
wrflist += [os.path.join(dirpath, file) for file in filenames]
### Plotting Iteration ###
# Iterate over a list of hourly timesteps
time = list()
for i in range(start_hour, end_hour):
time = str(i).zfill(2)
# Iterate over all 5 minutes steps of hour
for j in range(0, 60, 5):
minutes = str(j).zfill(2)
# Load the netCDF files out of the wrflist
ncfile = [Dataset(x) for x in wrflist
if x.endswith('{}_{}:{}:00'.format(date, time, minutes))]
# Load variable(s)
if title_name == 'CAPE':
variable = getvar(ncfile, variable_name)[0,:]
elif title_name == 'CIN':
variable = getvar(ncfile, variable_name)[1,:]
elif variable_name == 'ctt':
variable = getvar(ncfile, variable_name, units='K')
elif variable_name == 'wspd_wdir':
variable = getvar(ncfile, variable_name)[0,:]
elif variable_name == 'QVAPOR':
variable = getvar(ncfile, variable_name)*1000 # convert to g/kg
else:
variable = getvar(ncfile, variable_name)
if variable_name == 'slp':
slp = variable.squeeze()
ua = getvar(ncfile, 'ua')
va = getvar(ncfile, 'va')
p = getvar(ncfile, 'pressure')
u_wind850 = interplevel(ua, p, 850)
v_wind850 = interplevel(va, p, 850)
u_wind850 = u_wind850.squeeze()
v_wind850 = v_wind850.squeeze()
u_wind500 = interplevel(ua, p, 500)
v_wind500 = interplevel(va, p, 500)
u_wind500 = u_wind500.squeeze()
v_wind500 = v_wind500.squeeze()
slp = ndimage.gaussian_filter(slp, sigma=3, order=0)
# Interpolating 3d data to a horizontal pressure level
if pressure_level != False:
p = getvar(ncfile, 'pressure')
variable_pressure = interplevel(variable, p,
pressure_level)
variable = variable_pressure
if variable_name == 'wspd_wdir':
ua = getvar(ncfile, 'ua')
va = getvar(ncfile, 'va')
u_pressure = interplevel(ua, p, pressure_level)
v_pressure = interplevel(va, p, pressure_level)
elif title_name == 'Updraft and Reflectivity':
reflectivity = getvar(ncfile, 'REFD_MAX')
elif title_name == 'Difference in Theta-E values':
variable = getvar(ncfile, variable_name)
p = getvar(ncfile, 'pressure')
variable_pressure1 = interplevel(variable, p, '950')
variable_pressure2 = interplevel(variable, p, '950')
elif variable_name == 'ua':
va = getvar(ncfile, 'va')
p = getvar(ncfile, 'pressure')
v_pressure = interplevel(va, p, pressure_level)
u_wind = variable.squeeze()
v_wind = v_pressure.squeeze()
u_wind.attrs['units']='meters/second'
v_wind.attrs['units']='meters/second'
lats, lons = latlon_coords(variable)
lats = lats.squeeze()
lons = lons.squeeze()
dx, dy = mpcalc.lat_lon_grid_deltas(to_np(lons), to_np(lats))
divergence = mpcalc.divergence(u_wind, v_wind, dx, dy, dim_order='yx')
divergence = divergence*1e3
# Define cart projection
lats, lons = latlon_coords(variable)
cart_proj = ccrs.LambertConformal(central_longitude=8.722206,
central_latitude=46.73585)
bounds = geo_bounds(wrfin=ncfile)
# Create figure
fig = plt.figure(figsize=(15, 10))
if variable_name == 'slp':
fig.patch.set_facecolor('k')
ax = plt.axes(projection=cart_proj)
### Set map extent ###
domain_extent = [3.701088, 13.814863, 43.85472,49.49499]
if subset == True:
ax.set_extent([subset_extent[0],subset_extent[1],
subset_extent[2],subset_extent[3]],
ccrs.PlateCarree())
else:
ax.set_extent([domain_extent[0]+0.7,domain_extent[1]-0.7,
domain_extent[2]+0.1,domain_extent[3]-0.1],
ccrs.PlateCarree())
# Plot contour of variables
levels_num = 11
levels = np.linspace(variable_min, variable_max, levels_num)
# Creating new colormap for diverging colormaps
def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100):
new_cmap = LinearSegmentedColormap.from_list(
'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval,
b=maxval),
cmap(np.linspace(minval, maxval, n)))
return new_cmap
cmap = plt.get_cmap('RdYlBu')
if title_name == 'CIN':
cmap = ListedColormap(sns.cubehelix_palette(levels_num-1,
start=.5, rot=-.75, reverse=True))
variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable),
levels=levels, transform=ccrs.PlateCarree(), extend='max',
cmap=cmap)
initiation_color = 'r*'
elif variable_name == 'ctt':
cmap = ListedColormap(sns.cubehelix_palette(levels_num-1,
start=.5, rot=-.75, reverse=True))
variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable),
levels=levels, transform=ccrs.PlateCarree(), extend='both',
cmap=cmap)
initiation_color = 'r*'
elif variable_name == 'pvo':
cmap = plt.get_cmap('RdYlBu_r')
new_cmap = truncate_colormap(cmap, 0.05, 0.9)
new_norm = DivergingNorm(vmin=-1., vcenter=2., vmax=10)
variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable),
levels=levels, transform=ccrs.PlateCarree(),
cmap=new_cmap, extend='both', norm=new_norm)
initiation_color = 'k*'
elif variable_name == 'avo':
cmap = plt.get_cmap('RdYlBu_r')
new_cmap = truncate_colormap(cmap, 0.05, 0.9)
new_norm = DivergingNorm(vmin=variable_min, vcenter=0, vmax=variable_max)
variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable),
levels=levels, transform=ccrs.PlateCarree(),
cmap=new_cmap, extend='both', norm=new_norm)
initiation_color = 'k*'
elif variable_name == 'omega':
new_cmap = truncate_colormap(cmap, 0.05, 0.9)
new_norm = DivergingNorm(vmin=variable_min, vcenter=0, vmax=variable_max)
variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable),
levels=levels, transform=ccrs.PlateCarree(),
cmap=new_cmap, extend='both', norm=new_norm)
initiation_color = 'k*'
elif variable_name == 'ua':
new_cmap = truncate_colormap(cmap, 0.05, 0.9)
new_norm = DivergingNorm(vmin=variable_min, vcenter=0, vmax=variable_max)
variable_plot = plt.contourf(to_np(lons), to_np(lats), divergence,
levels=levels, transform=ccrs.PlateCarree(),
cmap=new_cmap, extend='both', norm=new_norm)
initiation_color = 'k*'
elif variable_name == 'UP_HELI_MAX' or variable_name == 'W_UP_MAX' or variable_name == 'QVAPOR':
cmap = ListedColormap(sns.cubehelix_palette(levels_num-1,
start=.5, rot=-.75))
variable_plot = plt.contourf(to_np(lons), to_np(lats),
to_np(variable), levels=levels, extend='max',
transform=ccrs.PlateCarree(),cmap=cmap)
initiation_color = 'r*'
elif variable_name == 'theta_e' or variable_name == 't2':
cmap = ListedColormap(sns.cubehelix_palette(levels_num-1,
start=.5, rot=-.75))
variable_plot = plt.contourf(to_np(lons), to_np(lats),
to_np(variable), levels=levels, extend='both',
transform=ccrs.PlateCarree(),cmap=cmap)
initiation_color = 'r*'
elif variable_name == 'REFD_MAX':
levels = np.arange(5., 75., 5.)
dbz_rgb = np.array([[4,233,231],
[1,159,244], [3,0,244],
[2,253,2], [1,197,1],
[0,142,0], [253,248,2],
[229,188,0], [253,149,0],
[253,0,0], [212,0,0],
[188,0,0],[248,0,253],
[152,84,198]], np.float32) / 255.0
dbz_cmap, dbz_norm = from_levels_and_colors(levels, dbz_rgb,
extend='max')
variable_plot = plt.contourf(to_np(lons), to_np(lats),
to_np(variable), levels=levels, extend='max',
transform=ccrs.PlateCarree(), cmap=dbz_cmap,
norm=dbz_norm)
initiation_color = 'r*'
elif variable_name == 'slp':
ax.background_patch.set_fill(False)
wslice = slice(1, None, 12)
# Plot 850-hPa wind vectors
vectors850 = ax.quiver(to_np(lons)[wslice, wslice],
to_np(lats)[wslice, wslice],
to_np(u_wind850)[wslice, wslice],
to_np(v_wind850)[wslice, wslice],
headlength=4, headwidth=3, scale=400, color='gold',
label='850mb wind', transform=ccrs.PlateCarree(),
zorder=2)
# Plot 500-hPa wind vectors
vectors500 = ax.quiver(to_np(lons)[wslice, wslice],
to_np(lats)[wslice, wslice],
to_np(u_wind500)[wslice, wslice],
to_np(v_wind500)[wslice, wslice],
headlength=4, headwidth=3, scale=400,
color='cornflowerblue', zorder=2,
label='500mb wind', transform=ccrs.PlateCarree())
# Plot 500-850 shear
shear = ax.quiver(to_np(lons[wslice, wslice]),
to_np(lats[wslice, wslice]),
to_np(u_wind500[wslice, wslice]) -
to_np(u_wind850[wslice, wslice]),
to_np(v_wind500[wslice, wslice]) -
to_np(v_wind850[wslice, wslice]),
headlength=4, headwidth=3, scale=400,
color='deeppink', zorder=2,
label='500-850mb shear', transform=ccrs.PlateCarree())
contour = ax.contour(to_np(lons), to_np(lats), slp, levels=levels,
colors='lime', linewidths=2, alpha=0.5, zorder=1,
transform=ccrs.PlateCarree())
ax.clabel(contour, fontsize=12, inline=1, inline_spacing=4, fmt='%i')
# Add a legend
ax.legend(('850mb wind', '500mb wind', '500-850mb shear'), loc=4)
# Manually set colors for legend
legend = ax.get_legend()
legend.legendHandles[0].set_color('gold')
legend.legendHandles[1].set_color('cornflowerblue')
legend.legendHandles[2].set_color('deeppink')
initiation_color = 'w*'
else:
cmap = ListedColormap(sns.cubehelix_palette(10,
start=.5, rot=-.75))
variable_plot = plt.contourf(to_np(lons), to_np(lats),
to_np(variable), levels=levels,
transform=ccrs.PlateCarree(),cmap=cmap)
initiation_color = 'r*'
# Plot reflectivity contours with colorbar
if title_name == 'Updraft and Reflectivity':
dbz_levels = np.arange(35., 75., 5.)
dbz_rgb = np.array([[253,248,2],
[229,188,0], [253,149,0],
[253,0,0], [212,0,0],
[188,0,0],[248,0,253],
[152,84,198]], np.float32) / 255.0
dbz_cmap, dbz_norm = from_levels_and_colors(dbz_levels, dbz_rgb,
extend='max')
contours = plt.contour(to_np(lons), to_np(lats),
to_np(reflectivity),
levels=dbz_levels,
transform=ccrs.PlateCarree(),
cmap=dbz_cmap, norm=dbz_norm,
linewidths=1)
cbar_refl = mpu.colorbar(contours, ax, orientation='horizontal', aspect=10,
shrink=.5, pad=0.05)
cbar_refl.set_label('Maximum Derived Radar Reflectivity' '[$dBZ$]', fontsize=12.5)
colorbar_lines = cbar_refl.ax.get_children()
colorbar_lines[0].set_linewidths([10]*5)
# Add wind quivers for every 10th data point
if variable_name == 'wspd_wdir':
plt.quiver(to_np(lons[::10,::10]), to_np(lats[::10,::10]),
to_np(u_pressure[::10, ::10]),
to_np(v_pressure[::10, ::10]),
transform=ccrs.PlateCarree())
# Plot colorbar
if variable_name == 'slp':
pass
else:
cbar = mpu.colorbar(variable_plot, ax, orientation='vertical', aspect=40,
shrink=.05, pad=0.05)
cbar.set_label(colorbar_label, fontsize=15)
cbar.set_ticks(levels)
# Add borders and coastlines
if variable_name == 'slp':
ax.add_feature(cfeature.BORDERS.with_scale('10m'),
edgecolor='white', linewidth=2)
ax.add_feature(cfeature.COASTLINE.with_scale('10m'),
edgecolor='white', linewidth=2)
else:
ax.add_feature(cfeature.BORDERS.with_scale('10m'),
linewidth=0.8)
ax.add_feature(cfeature.COASTLINE.with_scale('10m'),
linewidth=0.8)
### Add initiation location ###
if initiation == True:
ax.plot(initiation_location.lon, initiation_location.lat,
initiation_color, markersize=20, transform=ccrs.PlateCarree())
# Add gridlines
lon = np.arange(0, 20, 1)
lat = np.arange(40, 60, 1)
gl = ax.gridlines(xlocs=lon, ylocs=lat, zorder=3)
# Add tick labels
mpu.yticklabels(lat, ax=ax, fontsize=12.5)
mpu.xticklabels(lon, ax=ax, fontsize=12.5)
# Make nicetime
file_name = '{}wrfout_d02_{}_{}:{}:00'.format(data_dir,
date, time, minutes)
xr_file = xr.open_dataset(file_name)
nicetime = pd.to_datetime(xr_file.QVAPOR.isel(Time=0).XTIME.values)
nicetime = nicetime.strftime('%Y-%m-%d %H:%M')
# Add plot title
if pressure_level != False:
ax.set_title('{} @ {} hPa'.format(title_name, pressure_level),
loc='left', fontsize=15)
ax.set_title('Valid time: {} UTC'.format(nicetime),
loc='right', fontsize=15)
else:
if variable_name == 'slp':
ax.set_title(title_name, loc='left', fontsize=15, color='white')
ax.set_title('Valid time: {} UTC'.format(nicetime),
loc='right', fontsize=15, color='white')
else:
ax.set_title(title_name, loc='left', fontsize=20)
ax.set_title('Valid time: {} UTC'.format(nicetime),
loc='right', fontsize=15)
plt.show()
### Save figure ###
if save == True:
if pressure_level != False:
if subset == True:
fig.savefig('{}/{}/horizontal_map_{}_subset_{}_{}_{}:{}.png'.format(
save_dir, save_name, save_name, pressure_level, date, time,
minutes), bbox_inches='tight', dpi=300)
else:
fig.savefig('{}/{}/horizontal_map_{}_{}_{}_{}:{}.png'.format(
save_dir, save_name, save_name, pressure_level, date, time,
minutes), bbox_inches='tight', dpi=300)
else:
if subset == True:
fig.savefig('{}/{}/horizontal_map_{}_subset_{}_{}:{}.png'.format(
save_dir, save_name, save_name, date, time, minutes),
bbox_inches='tight', dpi=300, facecolor=fig.get_facecolor())
else:
fig.savefig('{}/{}/horizontal_map_{}_{}_{}:{}.png'.format(
save_dir, save_name, save_name, date, time, minutes),
bbox_inches='tight', dpi=300, facecolor=fig.get_facecolor())
### Make a GIF from the plots ###
if gif == True:
# Predifine some variables
gif_data_dir = save_dir + save_name
gif_save_dir = '{}gifs/'.format(save_dir)
gif_save_name = 'horizontal_map_{}.gif'.format(save_name)
# GIF creating procedure
os.chdir(gif_data_dir)
image_folder = os.fsencode(gif_data_dir)
filenames = []
for file in os.listdir(image_folder):
filename = os.fsdecode(file)
if filename.endswith( ('.png') ):
filenames.append(filename)
filenames.sort()
images = list(map(lambda filename: imageio.imread(filename),
filenames))
imageio.mimsave(os.path.join(gif_save_dir + gif_save_name),
images, duration = 0.50)
# Calculations
# ------------
#
# Most of the calculations in `metpy.calc` will accept DataArrays by converting them
# into their corresponding unit arrays. While this may often work without any issues, we must
# keep in mind that because the calculations are working with unit arrays and not DataArrays:
#
# - The calculations will return unit arrays rather than DataArrays
# - Broadcasting must be taken care of outside of the calculation, as it would only recognize
# dimensions by order, not name
#
# As an example, we calculate geostropic wind at 500 hPa below:
lat, lon = xr.broadcast(y, x)
f = mpcalc.coriolis_parameter(lat)
dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat, initstring=data_crs.proj4_init)
heights = data['height'].loc[time[0]].loc[{vertical.name: 500.}]
u_geo, v_geo = mpcalc.geostrophic_wind(heights, f, dx, dy)
print(u_geo)
print(v_geo)
#########################################################################
# Also, a limited number of calculations directly support xarray DataArrays or Datasets (they
# can accept *and* return xarray objects). Right now, this includes
#
# - Derivative functions
# - ``first_derivative``
# - ``second_derivative``
# - ``gradient``
# - ``laplacian``
# - Cross-section functions
def getData(self, time, model_vars, mdl2stnd, previous_data=None):
'''
Name:
awips_model_base
Purpose:
A function to get data from NAM40 model to create HDWX products
Inputs:
request : A DataAccessLayer request object
time : List of datatime(s) for data to grab
model_vars : Dictionary with variables/levels to get
mdl2stnd : Dictionary to convert from model variable names
to standardized names
Outputs:
Returns a dictionary containing all data
Keywords:
previous_data : Dictionary with data from previous time step
'''
log = logging.getLogger(__name__)
# Set up function for logger
initTime, fcstTime = get_init_fcst_times(time[0])
data = {
'model': self._request.getLocationNames()[0],
'initTime': initTime,
'fcstTime': fcstTime
}
# Initialize empty dictionary
log.info('Attempting to download {} data'.format(data['model']))
for var in model_vars: # Iterate over variables in the vars list
log.debug('Getting: {}'.format(var))
self._request.setParameters(*model_vars[var]['parameters'])
# Set parameters for the download request
self._request.setLevels(*model_vars[var]['levels'])
# Set levels for the download request
response = DAL.getGridData(self._request, time) # Request the data
for res in response: # Iterate over all data request responses
varName = res.getParameter()
# Get name of the variable in the response
varLvl = res.getLevel()
# Get level of the variable in the response
varName = mdl2stnd[varName]
# Convert variable name to local standarized name
if varName not in data:
data[varName] = {}
# If variable name NOT in data dictionary, initialize new dictionary under key
data[varName][varLvl] = res.getRawData()
# Add data under level name
try: # Try to
unit = units(res.getUnit())
# Get units and convert to MetPy units
except: # On exception
unit = '?'
# Set units to ?
else: # If get units success
data[varName][varLvl] *= unit
# Get data and create MetPy quantity by multiplying by units
log.debug(
'Got data for:\n Var: {}\n Lvl: {}\n Unit: {}'.format(
varName, varLvl, unit))
data['lon'], data['lat'] = res.getLatLonCoords()
# Get latitude and longitude values
data['lon'] *= units('degree')
# Add units of degree to longitude
data['lat'] *= units('degree')
# Add units of degree to latitude
# Absolute vorticity
dx, dy = lat_lon_grid_deltas(data['lon'], data['lat'])
# Get grid spacing in x and y
uTag = mdl2stnd[model_vars['wind']['parameters'][0]]
# Get initial tag name for u-wind
vTag = mdl2stnd[model_vars['wind']['parameters'][1]]
# Get initial tag name for v-wind
if (uTag in data) and (
vTag in data): # If both tags are in the data structure
data['abs_vort'] = {}
# Add absolute vorticity key
for lvl in model_vars['wind'][
'levels']: # Iterate over all leves in the wind data
if (lvl in data[uTag]) and (
lvl in data[vTag]
): # If given level in both u- and v-wind dictionaries
log.debug('Computing absolute vorticity at {}'.format(lvl))
data['abs_vort'][ lvl ] = \
absolute_vorticity( data[uTag][lvl], data[vTag][lvl],
dx, dy, data['lat'] )
# Compute absolute vorticity
# 1000 MB equivalent potential temperature
if ('temperature' in data) and (
'dewpoint'
in data): # If temperature AND depoint data were downloaded
data['theta_e'] = {}
T, Td = 'temperature', 'dewpoint'
if ('1000.0MB' in data[T]) and (
'1000.0MB' in data[Td]
): # If temperature AND depoint data were downloaded
log.debug(
'Computing equivalent potential temperature at 1000 hPa')
data['theta_e']['1000.0MB'] = equivalent_potential_temperature(
1000.0 * units('hPa'), data[T]['1000.0MB'],
data[Td]['1000.0MB'])
return data
# MLCAPE
log.debug('Computing mixed layer CAPE')
T_lvl = list(data[T].keys())
Td_lvl = list(data[Td].keys())
levels = list(set(T_lvl).intersection(Td_lvl))
levels = [float(lvl.replace('MB', '')) for lvl in levels]
levels = sorted(levels, reverse=True)
nLvl = len(levels)
if nLvl > 0:
log.debug(
'Found {} matching levels in temperature and dewpoint data'
.format(nLvl))
nLat, nLon = data['lon'].shape
data['MLCAPE'] = np.zeros((
nLat,
nLon,
), dtype=np.float32) * units('J/kg')
TT = np.zeros((
nLvl,
nLat,
nLon,
), dtype=np.float32) * units('degC')
TTd = np.zeros((
nLvl,
nLat,
nLon,
), dtype=np.float32) * units('degC')
log.debug('Sorting temperature and dewpoint data by level')
for i in range(nLvl):
key = '{:.1f}MB'.format(levels[i])
TT[i, :, :] = data[T][key].to('degC')
TTd[i, :, :] = data[Td][key].to('degC')
levels = np.array(levels) * units.hPa
depth = 100.0 * units.hPa
log.debug('Iterating over grid boxes to compute MLCAPE')
for j in range(nLat):
for i in range(nLon):
try:
_, T_parc, Td_parc = mixed_parcel(
levels,
TT[:, j, i],
TTd[:, j, i],
depth=depth,
interpolate=False,
)
profile = parcel_profile(levels, T_parc, Td_parc)
cape, cin = cape_cin(levels, TT[:, j, i],
TTd[:, j, i], profile)
except:
log.warning(
'Failed to compute MLCAPE for lon/lat: {}; {}'.
format(data['lon'][j, i], data['lat'][j, i]))
else:
data['MLCAPE'][j, i] = cape
return data
def calc_param_wrf_par(it):
#Have copied this function here from calc_param_wrf_par, to use global arrays
t, param = it
wg = False
p_3d = np.moveaxis(np.tile(p, [ta.shape[0], ta.shape[2], ta.shape[3], 1]),
[0, 1, 2, 3], [0, 2, 3, 1])
param = np.array(param)
param_out = [0] * (len(param))
for i in np.arange(0, len(param)):
param_out[i] = np.empty((len(lat), len(lon)))
if len(param) != len(np.unique(param)):
ValueError("Each parameter can only appear once in parameter list")
print(date_list[t])
start = dt.datetime.now()
hur_unit = units.percent * hur[t, :, :, :]
ta_unit = units.degC * ta[t, :, :, :]
dp_unit = units.degC * dp[t, :, :, :]
p_unit = units.hectopascals * p_3d[t, :, :, :]
q_unit = mpcalc.mixing_ratio_from_relative_humidity(hur_unit,\
ta_unit,p_unit)
theta_unit = mpcalc.potential_temperature(p_unit, ta_unit)
q = np.array(q_unit)
ml_inds = ((p_3d[t] <= ps[t]) & (p_3d[t] >= (ps[t] - 100)))
ml_ta_avg = np.ma.masked_where(~ml_inds, ta[t]).mean(axis=0).data
ml_q_avg = np.ma.masked_where(~ml_inds, q).mean(axis=0).data
ml_hgt_avg = np.ma.masked_where(~ml_inds, hgt[t]).mean(axis=0).data
ml_p3d_avg = np.ma.masked_where(~ml_inds, p_3d[t]).mean(axis=0).data
ml_ta_arr = np.insert(ta[t], 0, ml_ta_avg, axis=0)
ml_q_arr = np.insert(q, 0, ml_q_avg, axis=0)
ml_hgt_arr = np.insert(hgt[t], 0, ml_hgt_avg, axis=0)
ml_p3d_arr = np.insert(p_3d[t], 0, ml_p3d_avg, axis=0)
a,temp1,temp2 = np.meshgrid(np.arange(ml_p3d_arr.shape[0]) ,\
np.arange(ml_p3d_arr.shape[1]), np.arange(ml_p3d_arr.shape[2]))
sort_inds = np.flipud(
np.lexsort([np.swapaxes(a, 1, 0), ml_p3d_arr], axis=0))
ml_ta_arr = np.take_along_axis(ml_ta_arr, sort_inds, axis=0)
ml_p3d_arr = np.take_along_axis(ml_p3d_arr, sort_inds, axis=0)
ml_hgt_arr = np.take_along_axis(ml_hgt_arr, sort_inds, axis=0)
ml_q_arr = np.take_along_axis(ml_q_arr, sort_inds, axis=0)
cape3d_mlavg = wrf.cape_3d(ml_p3d_arr,ml_ta_arr + 273.15,\
ml_q_arr,ml_hgt_arr,terrain,ps[t,:,:],False,meta=False,missing=0)
ml_cape = np.ma.masked_where(~((ml_ta_arr==ml_ta_avg) & (ml_p3d_arr==ml_p3d_avg)),\
cape3d_mlavg.data[0]).max(axis=0).filled(0)
ml_cin = np.ma.masked_where(~((ml_ta_arr==ml_ta_avg) & (ml_p3d_arr==ml_p3d_avg)),\
cape3d_mlavg.data[1]).max(axis=0).filled(0)
cape3d = wrf.cape_3d(p_3d[t,:,:,:],ta[t,:,:,:]+273.15,q,hgt[t,:,:,:],terrain,ps[t,:,:],\
True,meta=False,missing=0)
cape = cape3d.data[0]
cin = cape3d.data[1]
cape[p_3d[t] > ps[t] - 25] = np.nan
cin[p_3d[t] > ps[t] - 25] = np.nan
mu_cape_inds = np.nanargmax(cape, axis=0)
mu_cape = mu_cape_inds.choose(cape)
mu_cin = mu_cape_inds.choose(cin)
cape_2d = wrf.cape_2d(p_3d[t,:,:,:],ta[t,:,:,:]+273.15,q\
,hgt[t,:,:,:],terrain,ps[t,:,:],True,meta=False,missing=0)
lcl = cape_2d[2].data
lfc = cape_2d[3].data
del hur_unit, dp_unit, theta_unit, ml_inds, ml_ta_avg, ml_q_avg, \
ml_hgt_avg, ml_p3d_avg, ml_ta_arr, ml_q_arr, ml_hgt_arr, ml_p3d_arr, a, temp1, temp2,\
sort_inds, cape3d_mlavg, cape3d, cape, cin, cape_2d
if "relhum850-500" in param:
param_ind = np.where(param == "relhum850-500")[0][0]
param_out[param_ind] = get_mean_var_p(hur[t], p, 850, 500)
if "relhum1000-700" in param:
param_ind = np.where(param == "relhum1000-700")[0][0]
param_out[param_ind] = get_mean_var_p(hur[t], p, 1000, 700)
if "mu_cape" in param:
param_ind = np.where(param == "mu_cape")[0][0]
param_out[param_ind] = mu_cape
if "ml_cape" in param:
param_ind = np.where(param == "ml_cape")[0][0]
param_out[param_ind] = ml_cape
if "s06" in param:
param_ind = np.where(param == "s06")[0][0]
s06 = get_shear_hgt(ua[t],va[t],hgt[t],0,6000,\
uas[t],vas[t])
param_out[param_ind] = s06
if "s03" in param:
param_ind = np.where(param == "s03")[0][0]
s03 = get_shear_hgt(ua[t],va[t],hgt[t],0,3000,\
uas[t],vas[t])
param_out[param_ind] = s03
if "s01" in param:
param_ind = np.where(param == "s01")[0][0]
s01 = get_shear_hgt(ua[t],va[t],hgt[t],0,1000,\
uas[t],vas[t])
param_out[param_ind] = s01
if "s0500" in param:
param_ind = np.where(param == "s0500")[0][0]
param_out[param_ind] = get_shear_hgt(ua[t],va[t],hgt[t],0,500,\
uas[t],vas[t])
if "lr1000" in param:
param_ind = np.where(param == "lr1000")[0][0]
lr1000 = get_lr_hgt(ta[t], hgt[t], 0, 1000)
param_out[param_ind] = lr1000
if "mu_cin" in param:
param_ind = np.where(param == "mu_cin")[0][0]
param_out[param_ind] = mu_cin
if "lcl" in param:
param_ind = np.where(param == "lcl")[0][0]
temp_lcl = np.copy(lcl)
temp_lcl[temp_lcl <= 0] = np.nan
param_out[param_ind] = temp_lcl
if "ml_cin" in param:
param_ind = np.where(param == "ml_cin")[0][0]
param_out[param_ind] = ml_cin
if "srh01" in param:
param_ind = np.where(param == "srh01")[0][0]
srh01 = get_srh(ua[t], va[t], hgt[t], 1000, True, 850, 700, p)
param_out[param_ind] = srh01
if "srh03" in param:
srh03 = get_srh(ua[t], va[t], hgt[t], 3000, True, 850, 700, p)
param_ind = np.where(param == "srh03")[0][0]
param_out[param_ind] = srh03
if "srh06" in param:
param_ind = np.where(param == "srh06")[0][0]
srh06 = get_srh(ua[t], va[t], hgt[t], 6000, True, 850, 700, p)
param_out[param_ind] = srh06
if "ship" in param:
if "s06" not in param:
raise NameError("To calculate ship, s06 must be included")
param_ind = np.where(param == "ship")[0][0]
muq = mu_cape_inds.choose(q)
ship = get_ship(mu_cape, np.copy(muq), ta[t], ua[t], va[t], hgt[t], p,
np.copy(s06))
param_out[param_ind] = ship
if "mmp" in param:
param_ind = np.where(param == "mmp")[0][0]
param_out[param_ind] = get_mmp(ua[t],va[t],uas[t],vas[t],\
mu_cape,ta[t],hgt[t])
if "scp" in param:
if "srh03" not in param:
raise NameError("To calculate ship, srh03 must be included")
param_ind = np.where(param == "scp")[0][0]
scell_pot = get_supercell_pot(mu_cape,ua[t],va[t],hgt[t],ta_unit,p_unit,\
q_unit,srh03)
param_out[param_ind] = scell_pot
if "stp" in param:
if "srh01" not in param:
raise NameError("To calculate stp, srh01 must be included")
param_ind = np.where(param == "stp")[0][0]
stp = get_tornado_pot(ml_cape,np.copy(lcl),np.copy(ml_cin),ua[t],va[t],p_3d[t],hgt[t],p,\
np.copy(srh01))
param_out[param_ind] = stp
if "vo10" in param:
param_ind = np.where(param == "vo10")[0][0]
x, y = np.meshgrid(lon, lat)
dx, dy = mpcalc.lat_lon_grid_deltas(x, y)
vo10 = get_vo(uas[t], vas[t], dx, dy)
param_out[param_ind] = vo10
if "conv10" in param:
param_ind = np.where(param == "conv10")[0][0]
x, y = np.meshgrid(lon, lat)
dx, dy = mpcalc.lat_lon_grid_deltas(x, y)
param_out[param_ind] = get_conv(uas[t], vas[t], dx, dy)
if "conv1000-850" in param:
levs = np.where((p <= 1001) & (p >= 849))[0]
param_ind = np.where(param == "conv1000-850")[0][0]
x, y = np.meshgrid(lon, lat)
dx, dy = mpcalc.lat_lon_grid_deltas(x, y)
param_out[param_ind] = \
np.mean(np.stack([get_conv(ua[t,i],va[t,i],dx,dy) for i in levs]),axis=0)
if "conv800-600" in param:
levs = np.where((p <= 801) & (p >= 599))[0]
param_ind = np.where(param == "conv800-600")[0][0]
x, y = np.meshgrid(lon, lat)
dx, dy = mpcalc.lat_lon_grid_deltas(x, y)
param_out[param_ind] = \
np.mean(np.stack([get_conv(ua[t,i],va[t,i],dx,dy) for i in levs]),axis=0)
if "non_sc_stp" in param:
if "vo10" not in param:
raise NameError("To calculate non_sc_stp, vo must be included")
if "lr1000" not in param:
raise NameError("To calculate non_sc_stp, lr1000 must be included")
param_ind = np.where(param == "non_sc_stp")[0][0]
non_sc_stp = get_non_sc_tornado_pot(ml_cape,ml_cin,np.copy(lcl),ua[t],va[t],\
uas[t],vas[t],p_3d[t],ta[t],hgt[t],p,vo10,lr1000)
param_out[param_ind] = non_sc_stp
if "cape*s06" in param:
param_ind = np.where(param == "cape*s06")[0][0]
cs6 = ml_cape * np.power(s06, 1.67)
param_out[param_ind] = cs6
if "td850" in param:
param_ind = np.where(param == "td850")[0][0]
td850 = get_td_diff(ta[t], dp[t], p_3d[t], 850)
param_out[param_ind] = td850
if "td800" in param:
param_ind = np.where(param == "td800")[0][0]
param_out[param_ind] = get_td_diff(ta[t], dp[t], p_3d[t], 800)
if "td950" in param:
param_ind = np.where(param == "td950")[0][0]
param_out[param_ind] = get_td_diff(ta[t], dp[t], p_3d[t], 950)
if "wg" in param:
try:
param_ind = np.where(param == "wg")[0][0]
param_out[param_ind] = wg[t]
except ValueError:
print("wg field expected, but not parsed")
if "dcape" in param:
param_ind = np.where(param == "dcape")[0][0]
dcape = np.nanmax(get_dcape(p_3d[t], ta[t], hgt[t], p, ps[t]), axis=0)
param_out[param_ind] = dcape
if "mlm" in param:
param_ind = np.where(param == "mlm")[0][0]
mlm_u, mlm_v = get_mean_wind(ua[t], va[t], hgt[t], 800, 600, False,
None, "plevels", p)
mlm = np.sqrt(np.square(mlm_u) + np.square(mlm_v))
param_out[param_ind] = mlm
if "dlm" in param:
param_ind = np.where(param == "dlm")[0][0]
dlm_u, dlm_v = get_mean_wind(ua[t], va[t], hgt[t], 1000, 500, False,
None, "plevels", p)
dlm = np.sqrt(np.square(dlm_u) + np.square(dlm_v))
param_out[param_ind] = dlm
if "dlm+dcape" in param:
param_ind = np.where(param == "dlm+dcape")[0][0]
dlm_dcape = dlm + np.sqrt(2 * dcape)
param_out[param_ind] = dlm_dcape
if "mlm+dcape" in param:
param_ind = np.where(param == "mlm+dcape")[0][0]
mlm_dcape = mlm + np.sqrt(2 * dcape)
param_out[param_ind] = mlm_dcape
if "dcape*cs6" in param:
param_ind = np.where(param == "dcape*cs6")[0][0]
param_out[param_ind] = (dcape / 980.) * (cs6 / 20000)
if "dlm*dcape*cs6" in param:
param_ind = np.where(param == "dlm*dcape*cs6")[0][0]
param_out[param_ind] = (dlm_dcape / 30.) * (cs6 / 20000)
if "mlm*dcape*cs6" in param:
param_ind = np.where(param == "mlm*dcape*cs6")[0][0]
param_out[param_ind] = (mlm_dcape / 30.) * (cs6 / 20000)
if "dcp" in param:
param_ind = np.where(param == "dcp")[0][0]
param_out[param_ind] = (dcape / 980) * (mu_cape /
2000) * (s06 / 10) * (dlm / 8)
if "mf" in param:
param_ind = np.where(param == "mf")[0][0]
mf = ((ml_cape > 120) & (dcape > 350) & (mlm < 26))
mf = mf * 1.0
param_out[param_ind] = mf
if "sf" in param:
param_ind = np.where(param == "sf")[0][0]
sf = ((s06 >= 30) & (dcape < 500) & (mlm >= 26))
sf = sf * 1.0
param_out[param_ind] = sf
if "cond" in param:
param_ind = np.where(param == "cond")[0][0]
cond = (sf == 1.0) | (mf == 1.0)
cond = cond * 1.0
param_out[param_ind] = cond
return param_out
