def test_divergence_asym():
"""Test divergence calculation with a complicated field."""
u = np.array([[2, 4, 8], [0, 2, 2], [4, 6, 8]]) * units('m/s')
v = np.array([[6, 4, 8], [2, 6, 0], [2, 2, 6]]) * units('m/s')
c = divergence(u, v, 1 * units.meters, 2 * units.meters, dim_order='yx')
true_c = np.array([[-2, 5.5, -2.5], [2., 0.5, -1.5], [3., -1.5, 8.5]]) / units.sec
assert_array_equal(c, true_c)
# Now try for xy ordered
c = divergence(u.T, v.T, 1 * units.meters, 2 * units.meters, dim_order='xy')
assert_array_equal(c, true_c.T)
def test_zero_divergence():
"""Test divergence calculation when zeros should be returned."""
a = np.arange(3)
u = np.c_[a, a, a] * units('m/s')
c = divergence(u, u.T, 1 * units.meter, 1 * units.meter, dim_order='xy')
true_c = 2. * np.ones_like(u) / units.sec
assert_array_equal(c, true_c)
def test_divergence():
"""Test divergence for simple case."""
a = np.arange(3)
u = np.c_[a, a, a] * units('m/s')
c = divergence(u, u, 1 * units.meter, 1 * units.meter, dim_order='xy')
true_c = np.ones_like(u) / units.sec
assert_array_equal(c, true_c)
def windProducts(x, y, u, v):
dx, dy = mpcalc.lat_lon_grid_spacing(x, y)
deformation = numpy.zeros(u.shape)
convergence = numpy.zeros(u.shape)
for i in range(u.shape[0]):
shear, stretch = mpcalc.shearing_stretching_deformation(
u[i, :, :], v[i, :, :], dx, dy)
deformation[i] = numpy.sqrt(shear**2 + stretch**2)
convergence[i] = -mpcalc.divergence(u[i, :, :], v[i, :, :], dx, dy)
return deformation, convergence
def kinematics(u, v, thetae, dx, dy, lats, smooth=False, sigma=1):
'''
Use metpy functions to calculate various kinematics, given 2d numpy arrays as inputs. X and y grid spacing, as well as a 2d array of latitudes is also required.
Option to smooth the thetae field before calculating front diagnostics (recommended), with addition "sigma" parameter controlling the smoothness (see scipy docs)
'''
if smooth:
thetae = gaussian_filter(thetae, sigma)
#Kinematics
ddy_thetae = mpcalc.first_derivative(thetae, delta=dy, axis=0)
ddx_thetae = mpcalc.first_derivative(thetae, delta=dx, axis=1)
mag_thetae = np.sqrt(ddx_thetae**2 + ddy_thetae**2)
div = mpcalc.divergence(u, v, dx, dy)
strch_def = mpcalc.stretching_deformation(u, v, dx, dy)
shear_def = mpcalc.shearing_deformation(u, v, dx, dy)
tot_def = mpcalc.total_deformation(u, v, dx, dy)
psi = 0.5 * np.arctan2(shear_def, strch_def)
beta = np.arcsin(
(-ddx_thetae * np.cos(psi) - ddy_thetae * np.sin(psi)) / mag_thetae)
vo = mpcalc.vorticity(u, v, dx, dy)
conv = -div * 1e5
F = np.array(0.5 * mag_thetae * (tot_def * np.cos(2 * beta) - div) *
1.08e4 * 1e5)
Fn = np.array(0.5 * mag_thetae * (div - tot_def * np.cos(2 * beta)) *
1.08e4 * 1e5)
Fs = np.array(0.5 * mag_thetae * (vo + tot_def * np.sin(2 * beta)) *
1.08e4 * 1e5)
icon = np.array(0.5 * (tot_def - div) * 1e5)
vgt = np.array(np.sqrt(div**2 + vo**2 + tot_def**2) * 1e5)
#TFP
ddx_thetae_scaled = ddx_thetae / mag_thetae
ddy_thetae_scaled = ddy_thetae / mag_thetae
ddy_mag_thetae = -1 * mpcalc.first_derivative(mag_thetae, delta=dy, axis=0)
ddx_mag_thetae = -1 * mpcalc.first_derivative(mag_thetae, delta=dx, axis=1)
tfp = ddx_mag_thetae * ddx_thetae_scaled + ddy_mag_thetae * ddy_thetae_scaled
ddy_tfp = mpcalc.first_derivative(tfp, delta=dy, axis=0)
ddx_tfp = mpcalc.first_derivative(tfp, delta=dx, axis=1)
mag_tfp = np.sqrt(ddx_tfp**2 + ddy_tfp**2)
v_f = (u * units.units("m/s")) * (ddx_tfp / mag_tfp) + (
v * units.units("m/s")) * (ddy_tfp / mag_tfp)
#Extra condition
ddy_ddy_mag_te = mpcalc.first_derivative(-1 * ddy_mag_thetae,
delta=dy,
axis=0)
ddx_ddx_mag_te = mpcalc.first_derivative(-1 * ddx_mag_thetae,
delta=dx,
axis=1)
cond = ddy_ddy_mag_te + ddx_ddx_mag_te
return [F, Fn, Fs, icon, vgt, conv, vo*1e5, \
np.array(tfp.to("km^-2")*(100*100)), np.array(mag_thetae.to("km^-1"))*100, np.array(v_f), thetae, cond]
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 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)
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(
winds_clm.u.values[i, :, :],
ax.set_extent([WLON, ELON, SLAT, NLAT])
ax.add_feature(cfeature.COASTLINE.with_scale('50m'), linewidth=0.5)
ax.add_feature(cfeature.LAKES.with_scale('50m'), linewidth=0.5)
ax.add_feature(cfeature.STATES, linewidth=0.5)
ax.add_feature(cfeature.BORDERS, linewidth=0.5)
return ax
# =============================================================================
# FIG #1: 250: JET STREAM, GEOPOTENTIAL HEIGHT, DIVERGENCE
# =============================================================================
H250 = HGT_DATA.variables['hgt'][TIME_INDEX, 8, :, :]
U250 = UWND_DATA.variables['uwnd'][TIME_INDEX, 8, :, :] * units('m/s')
V250 = VWND_DATA.variables['vwnd'][TIME_INDEX, 8, :, :] * units('m/s')
SPEED250 = mpcalc.get_wind_speed(U250, V250)
DIV250 = mpcalc.divergence(U250, V250, DX, DY, dim_order='YX')
DIV250 = (DIV250 * (units('1/s')))
# =============================================================================
# FIG #2: 500: VORTICITY, GEOPOTENTIAL HEIGHT, VORTICITY ADVECTION
# =============================================================================
H500 = HGT_DATA.variables['hgt'][TIME_INDEX, 5, :, :]
U500 = UWND_DATA.variables['uwnd'][TIME_INDEX, 5, :, :] * units('m/s')
V500 = VWND_DATA.variables['vwnd'][TIME_INDEX, 5, :, :] * units('m/s')
DX, DY = mpcalc.lat_lon_grid_spacing(LON, LAT)
VORT500 = mpcalc.vorticity(U500, V500, DX, DY, dim_order='YX')
VORT500 = (VORT500 * (units('1/s')))
VORT_ADV500 = mpcalc.advection(VORT500, [U500, V500], (DX, DY), dim_order='yx')
# =============================================================================
# FIG #3: 700: Q-VECTORS+CONVERGENCE, GEOPOTENTIAL HEIGHT
# =============================================================================
H700 = HGT_DATA.variables['hgt'][TIME_INDEX, 3, :, :]
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
def gh_uv_div(initTime=None, fhour=6, day_back=0,model='GRAPES_GFS',
gh_lev=500,uv_lev=850,
map_ratio=14/9,zoom_ratio=20,cntr_pnt=[104,34],
south_China_sea=True,area =None,city=False,output_dir=None,data_source='MICAPS',**kwargs):
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='HGT',lvl=gh_lev),
utl.Cassandra_dir(data_type='high',data_source=model,var_name='UGRD',lvl=uv_lev),
utl.Cassandra_dir(data_type='high',data_source=model,var_name='VGRD',lvl=uv_lev),
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(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
gh = MICAPS_IO.get_model_grid(data_dir[0], filename=filename)
if gh is None:
return
u = MICAPS_IO.get_model_grid(data_dir[1], filename=filename)
if u is None:
return
v = MICAPS_IO.get_model_grid(data_dir[2], filename=filename)
if v is None:
return
psfc = MICAPS_IO.get_model_grid(data_dir[3], filename=filename)
init_time = gh.coords['forecast_reference_time'].values
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)
# retrieve data from CIMISS server
try:
gh=CMISS_IO.cimiss_model_by_time('20'+filename[0:8],valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='GPH'),
levattrs={'long_name':'pressure_level', 'units':'hPa', '_CoordinateAxisType':'-'},
fcst_level=gh_lev, fcst_ele="GPH", units='gpm')
if gh is None:
return
gh['data'].values=gh['data'].values/10.
u=CMISS_IO.cimiss_model_by_time('20'+filename[0:8],valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='WIU'),
levattrs={'long_name':'pressure_level', 'units':'hPa', '_CoordinateAxisType':'-'},
fcst_level=uv_lev, fcst_ele="WIU", units='m/s')
if u is None:
return
v=CMISS_IO.cimiss_model_by_time('20'+filename[0:8],valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='WIV'),
levattrs={'long_name':'pressure_level', 'units':'hPa', '_CoordinateAxisType':'-'},
fcst_level=uv_lev, fcst_ele="WIV", units='m/s')
if v is None:
return
psfc=CMISS_IO.cimiss_model_by_time('20'+filename[0:8], valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='PRS'),
fcst_level=0, fcst_ele="PRS", units='Pa')
psfc['data']=psfc['data']/100.
except KeyError:
raise ValueError('Can not find all data needed')
# prepare data
if(area != None):
cntr_pnt,zoom_ratio=utl.get_map_area(area_name=area)
map_extent,delt_x,delt_y=utl.get_map_extent(cntr_pnt, zoom_ratio, map_ratio)
#+ to solve the problem of labels on all the contours
dx,dy=mpcalc.lat_lon_grid_deltas(u['lon'].values.squeeze(),u['lat'].values.squeeze())
div=mpcalc.divergence(u['data'].values.squeeze()*units('m/s'),
v['data'].values.squeeze()*units('m/s'),
dx, dy, dim_order='yx')
div_xr=u.copy(deep=True)
div_xr['data'].values=div.magnitude[np.newaxis,np.newaxis,:,:]
div_xr['data'].values=gaussian_filter(div_xr['data'].values,1)
gh=utl.mask_terrian(gh_lev,psfc,gh)
u=utl.mask_terrian(uv_lev,psfc,u)
v=utl.mask_terrian(uv_lev,psfc,v)
div_xr=utl.mask_terrian(uv_lev,psfc,div_xr)
gh=utl.cut_xrdata(map_extent, gh, delt_x=delt_x, delt_y=delt_y)
u=utl.cut_xrdata(map_extent, u, delt_x=delt_x, delt_y=delt_y)
v=utl.cut_xrdata(map_extent, v, delt_x=delt_x, delt_y=delt_y)
div_xr=utl.cut_xrdata(map_extent, div_xr, delt_x=delt_x, delt_y=delt_y)
gh.attrs['model']=model
div_xr.attrs['units']='s$^{-1}$'
uv=xr.merge([u.rename({'data': 'u'}),v.rename({'data': 'v'})])
#- to solve the problem of labels on all the contours
dynamic_graphics.draw_gh_uv_div(
div=div_xr, gh=gh, uv=uv,
map_extent=map_extent, regrid_shape=20,
city=city,south_China_sea=south_China_sea,
output_dir=output_dir)
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)
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()
res.nglDraw = False
res.nglFrame = False
def Crosssection_Wind_Theta_e_div(
initTime=None, fhour=24,lw_ratio=[16,9],
levels=[1000, 950, 925, 900, 850, 800, 700,600,500,400,300,200],
day_back=0,model='GRAPES_GFS',data_source='MICAPS',
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],**kwargs):
# 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='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(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)
rh = rh.metpy.parse_cf().squeeze()
u=MICAPS_IO.get_model_3D_grid(directory=data_dir[1][0:-1],filename=filename,levels=levels, allExists=False)
u = u.metpy.parse_cf().squeeze()
v=MICAPS_IO.get_model_3D_grid(directory=data_dir[2][0:-1],filename=filename,levels=levels, allExists=False)
v = v.metpy.parse_cf().squeeze()
v2=MICAPS_IO.get_model_3D_grid(directory=data_dir[2][0:-1],filename=filename,levels=levels, allExists=False)
v2 = v2.metpy.parse_cf().squeeze()
t=MICAPS_IO.get_model_3D_grid(directory=data_dir[3][0:-1],filename=filename,levels=levels, allExists=False)
t = t.metpy.parse_cf().squeeze()
gh=MICAPS_IO.get_model_grid(data_dir[4], filename=filename)
psfc=get_model_grid(data_dir[5], filename=filename)
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:
rh=CMISS_IO.cimiss_model_3D_grid(init_time_str='20'+filename[0:8],valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='RHU'),
fcst_levels=levels, fcst_ele="RHU", units='%')
if rh is None:
return
u=CMISS_IO.cimiss_model_3D_grid(init_time_str='20'+filename[0:8],valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='WIU'),
fcst_levels=levels, fcst_ele="WIU", units='m/s')
if u is None:
return
v=CMISS_IO.cimiss_model_3D_grid(init_time_str='20'+filename[0:8],valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='WIV'),
fcst_levels=levels, fcst_ele="WIV", units='m/s')
if v is None:
return
v2=CMISS_IO.cimiss_model_3D_grid(init_time_str='20'+filename[0:8],valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='WIV'),
fcst_levels=levels, fcst_ele="WIV", units='m/s')
if v2 is None:
return
t=CMISS_IO.cimiss_model_3D_grid(init_time_str='20'+filename[0:8],valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='TEM'),
fcst_levels=levels, fcst_ele="TEM", units='K')
if t is None:
return
t['data'].values=t['data'].values-273.15
gh=CMISS_IO.cimiss_model_by_time('20'+filename[0:8],valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='GPH'),
fcst_level=500, fcst_ele="GPH", units='gpm')
if gh is None:
return
gh['data'].values=gh['data'].values/10.
psfc=CMISS_IO.cimiss_model_by_time('20'+filename[0:8], valid_time=fhour,
data_code=utl.CMISS_data_code(data_source=model,var_name='PRS'),
fcst_level=0, fcst_ele="PRS", units='Pa')
psfc['data']=psfc['data']/100.
except KeyError:
raise ValueError('Can not find all data needed')
rh = rh.metpy.parse_cf().squeeze()
u = u.metpy.parse_cf().squeeze()
v = v.metpy.parse_cf().squeeze()
v2 = v2.metpy.parse_cf().squeeze()
t = t.metpy.parse_cf().squeeze()
psfc=psfc.metpy.parse_cf().squeeze()
resolution=u['lon'][1]-u['lon'][0]
x,y=np.meshgrid(u['lon'], u['lat'])
# +form 3D psfc
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)
# -form 3D psfc
dx,dy=mpcalc.lat_lon_grid_deltas(u['lon'],u['lat'])
for ilvl in levels:
u2d=u.sel(level=ilvl)
v2d=v.sel(level=ilvl)
# absv2d=mpcalc.absolute_vorticity(u2d['data'].values*units.meter/units.second,
# v2d['data'].values*units.meter/units.second,dx,dy,y*units.degree)
div2d=mpcalc.divergence(u2d['data'].values.squeeze()*units('m/s'),
v2d['data'].values.squeeze()*units('m/s'),
dx, dy, dim_order='yx')
if(ilvl == levels[0]):
div3d = v2.copy()
div3d['data'].loc[dict(level=ilvl)]=np.array(div2d)
else:
div3d['data'].loc[dict(level=ilvl)]=np.array(div2d)
div3d['data'].attrs['units']=div2d.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_psfc = cross_section(psfc_bdcst, st_point, ed_point)
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(div3d, st_point, ed_point)
cross_div3d=cross.set_coords(('lat', 'lon'))
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'])
pressure.attrs['units']='hPa'
Theta_e=mpcalc.equivalent_potential_temperature(pressure,
cross_t['data'].values*units.celsius,
cross_Td)
cross_terrain=pressure-cross_psfc
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_div(
cross_div3d=cross_div3d, cross_Theta_e=cross_Theta_e, 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,lw_ratio=lw_ratio,
output_dir=output_dir)
## first, smooth everything pretty heavily:
tmpk_700 = ndimage.gaussian_filter(tmpk_700, sigma=6, order=0) * units('kelvin')
uwnd_700 = ndimage.gaussian_filter(uwnd_700, sigma=6, order=0) * units('m/s')
vwnd_700 = ndimage.gaussian_filter(vwnd_700, sigma=6, order=0) * units('m/s')
# less smoothing for the height
hght_700 = ndimage.gaussian_filter(hght_700, sigma=2, order=0) * units('m')
if plot_700qvect:
# Compute the Q-vector components
uqvect, vqvect = mpcalc.q_vector(uwnd_700, vwnd_700, tmpk_700, 700*units.hPa, dx, dy)
# Compute the divergence of the Q-vectors calculated above
q_div = -2*mpcalc.divergence(uqvect, vqvect, dx, dy, dim_order='yx')
## smooth it too
q_div = ndimage.gaussian_filter(q_div, sigma=4, order=0) * units('meter/(kilogram second)')
#Map Creation
print('plotting 700 q-vectors')
# Set Projection of Data
datacrs = ccrs.PlateCarree()
# Set Projection of Plot
plotcrs = ccrs.LambertConformal(central_latitude=[30, 60], central_longitude=-100)
fig = plt.figure(1, figsize=(20,16))
hgt = xr.open_dataset(PATH_DATA + FILE_HGT_S4)
hgt = hgt - hgt.mean(dim='longitude')
winds = xr.open_dataset(PATH_DATA + FILE_WINDS, chunks={'latitude':10})
winds = winds.transpose('month', 'realiz', 'latitude', 'longitude')
winds_clm = winds.mean(dim='realiz')
for i in np.arange(0, 7):
var_normal = np.mean(hgt.z.values[i, :, :, :], axis=0)
px, py, lat = plumb_flux.ComputePlumbFluxes(winds_clm.u.values[i, :, :], winds_clm.v.values[i, :, :], hgt.z.values[i, :, :, :], var_normal, hgt.latitude.values, hgt.longitude.values, 20000)
dx, dy = calc.lat_lon_grid_deltas(hgt.longitude.values, lat)
#div = calc.divergence(px, py, dx, dy)
px_normal, py_normal = np.mean(px, axis=0), np.mean(py, axis=0)
var_ninio_WPV = np.mean(hgt.z.values[i, index_ninio_WPV, :, :], axis=0)
px_ninio_WPV, py_ninio_WPV = np.mean(px[index_ninio_WPV, :, :], axis=0) - px_normal, np.mean(py[index_ninio_WPV, :, :], axis=0) - py_normal
div_ninio_WPV = calc.divergence(px_ninio_WPV, py_ninio_WPV, dx, dy)
var_ninia_WPV = np.mean(hgt.z.values[i, index_ninia_WPV, :, :], axis=0)
px_ninia_WPV, py_ninia_WPV = np.mean(px[index_ninia_WPV, :, :], axis=0) - px_normal, np.mean(py[index_ninia_WPV, :, :], axis=0) - py_normal
div_ninia_WPV = calc.divergence(px_ninia_WPV, py_ninia_WPV, dx, dy)
var_ninio_SPV = np.mean(hgt.z.values[i, index_ninio_SPV, :, :], axis=0)
px_ninio_SPV, py_ninio_SPV = np.mean(px[index_ninio_SPV, :, :], axis=0) - px_normal, np.mean(py[index_ninio_SPV, :, :], axis=0) - py_normal
div_ninio_SPV = calc.divergence(px_ninio_SPV, py_ninio_SPV, dx, dy)
var_ninia_SPV = np.mean(hgt.z.values[i, index_ninia_SPV, :, :], axis=0)
px_ninia_SPV, py_ninia_SPV = np.mean(px[index_ninia_SPV, :, :], axis=0) - px_normal, np.mean(py[index_ninia_SPV, :, :], axis=0) - py_normal
div_ninia_SPV = calc.divergence(px_ninia_SPV, py_ninia_SPV, dx, dy)
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)
# Use MetPy to compute the baroclinic potential vorticity on all isobaric
# levels and other variables
#
# Compute dx and dy spacing for use in vorticity calculation
dx, dy = mpcalc.lat_lon_grid_deltas(lons, lats)
# Comput the PV on all isobaric surfaces
pv = mpcalc.potential_vorticity_baroclinic(
thta, pres[:, None, None], uwnd, vwnd, dx[None, :, :], dy[None, :, :],
lats[None, :, None] * units('degrees'))
# Use MetPy to compute the divergence on the pressure surfaces
div = mpcalc.divergence(uwnd,
vwnd,
dx[None, :, :],
dy[None, :, :],
dim_order='yx')
# Find the index value for the 250-hPa surface
i250 = list(pres.m).index(((250 * units('hPa')).to(pres.units)).m)
######################################################################
# 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
# CONUS with isobaric PV map with PV contoured every 1 PVU and divergence
# colorshaded.
#
def PV_Div_uv(initial_time=None, fhour=6, day_back=0,model='ECMWF',
map_ratio=19/9,zoom_ratio=20,cntr_pnt=[102,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 = '全国',city=False,output_dir=None
):
# 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='')]
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
u=get_model_3D_grid(directory=data_dir[1][0:-1],filename=filename,levels=levels, allExists=False)
if u is None:
return
v=get_model_3D_grid(directory=data_dir[2][0:-1],filename=filename,levels=levels, allExists=False)
if v is None:
return
t=get_model_3D_grid(directory=data_dir[3][0:-1],filename=filename,levels=levels, allExists=False)
if t is None:
return
# 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)
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(), 9, 2)*units('degC')
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 = mpcalc.potential_vorticity_baroclinic(thta, pres[:, None, None], uwnd, vwnd,
dx[None, :, :], dy[None, :, :],
lats[None, :, None] * units('degrees'))
div = 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)
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
idx_x1 = np.where((lons > map_extent[0]-delt_x) &
(lons < map_extent[1]+delt_x))
idx_y1 = np.where((lats > map_extent[2]-delt_y) &
(lats < map_extent[3]+delt_y))
#- to solve the problem of labels on all the contours
init_time = u.coords['forecast_reference_time'].values
pv = {
'lon': lons[idx_x1],
'lat': lats[idx_y1],
'data': np.array(pv)[idx_z1,idx_y1[0][0]:(idx_y1[0][-1]+1),idx_x1[0][0]:(idx_x1[0][-1]+1)],
'lev':str(lvl_ana),
'model':model,
'fhour':fhour,
'init_time':init_time}
uv = {
'lon': lons[idx_x1],
'lat': lats[idx_y1],
'udata': np.array(uwnd)[idx_z1,idx_y1[0][0]:(idx_y1[0][-1]+1),idx_x1[0][0]:(idx_x1[0][-1]+1)],
'vdata': np.array(vwnd)[idx_z1,idx_y1[0][0]:(idx_y1[0][-1]+1),idx_x1[0][0]:(idx_x1[0][-1]+1)],
'lev':str(lvl_ana)}
div = {
'lon': lons[idx_x1],
'lat': lats[idx_y1],
'data': np.array(div)[idx_z1,idx_y1[0][0]:(idx_y1[0][-1]+1),idx_x1[0][0]:(idx_x1[0][-1]+1)],
'lev':str(lvl_ana)}
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 metpy_temp_adv(fname, plevels, tidx):
ds = xr.open_dataset(fname).metpy.parse_cf().squeeze()
ds = ds.isel(Time=tidx)
print(ds)
dx, dy = mpcalc.lat_lon_grid_deltas(ds.lon.values * units('degrees_E'),
ds.lat.values * units('degrees_N'))
plevels_unit = plevels * units.hPa
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')
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(var), coords=coords, dims=dims)
# Calculate temperature advection using metpy function
for i, plev in enumerate(plevs):
uqvect, vqvect = mpcalc.q_vector(u[i, :, :], v[i, :, :], th[i, :, :],
plev * units.hPa, dx, dy)
#uqvect, vqvect = mpcalc.q_vector(u[i,:,:], v[i,:,:], th.to('degC')[i,:,:], plev*units.hPa, dx, dy)
q_div = -2 * mpcalc.divergence(uqvect, vqvect, dx, dy, dim_order='yx')
ds1['uq_{:03d}'.format(plev)] = xr.DataArray(
np.array(uqvect),
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
ds1['vq_{:03d}'.format(plev)] = xr.DataArray(
np.array(vqvect),
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
ds1['q_div_{:03d}'.format(plev)] = xr.DataArray(
np.array(q_div),
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
# adv = mpcalc.advection(tk[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['tk_adv_{:03d}'.format(plev)] = xr.DataArray(np.array(adv), coords=[ds.lat.values,ds.lon.values], dims=["lat","lon"])
#
# adv = mpcalc.advection(th[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['th_adv_{:03d}'.format(plev)] = xr.DataArray(np.array(adv), coords=[ds.lat.values,ds.lon.values], dims=["lat","lon"])
#
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['Time'] = ds.Time
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['t2m'] = xr.DataArray(ds.t2m.values,
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
ds1['th2'] = xr.DataArray(ds.th2.values,
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
ds1['mask'] = xr.DataArray(ds.mask.values,
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
ds1['u10'] = xr.DataArray(ds.u10.values,
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
ds1['v10'] = xr.DataArray(ds.v10.values,
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
ds1['slp'] = xr.DataArray(ds.slp.values,
coords=[ds.lat.values, ds.lon.values],
dims=["lat", "lon"])
ds1['rain'] = xr.DataArray(ds.rain.values,
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"])
#ds1['latent'] = xr.DataArray(ds.latent.values, coords=[ds.lat.values,ds.lon.values], dims=["lat","lon"])
#ds1['acclatent'] = xr.DataArray(ds.acclatent.values, coords=[ds.lat.values,ds.lon.values], dims=["lat","lon"])
return ds1
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(
"--barpa_forcing_mdl",
help="BARPA forcing model (erai or ACCESS1-0). Default erai.",
default="erai")
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(
"--ub4",
help=
"Where to get era5 data. Default True for ub4 project, otherwise rt52",
default="True")
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(
"--is_dcape",
help="Should DCAPE be calculated? (1 or 0. Default is 1)",
default=1)
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")
parser.add_argument(
"--delta_t",
help=
"Time step spacing for ERA5 data, in hours. Default is one the minimum spacing (1 hour)",
default="1")
parser.add_argument(
"--era5_interp",
help=
"Horizontally interpolate model data before calculating convective parameters",
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
ub4 = args.ub4
al33 = args.al33
if args.outname == None:
out_name = model
else:
out_name = args.outname
is_dcape = args.is_dcape
barpa_forcing_mdl = args.barpa_forcing_mdl
experiment = args.e
ensemble = args.ens
group = args.group
project = args.project
ver6hr = args.ver6hr
ver3hr = args.ver3hr
delta_t = int(args.delta_t)
era5_interp = args.era5_interp
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 era5_interp == "True":
era5_interp = True
elif era5_interp == "False":
era5_interp = False
else:
raise ValueError("\n INVALID era5_interp...SHOULD BE True OR False")
if ub4 == "True":
ub4 = True
elif ub4 == "False":
ub4 = False
else:
raise ValueError("\n INVALID ub4...SHOULD BE True OR False")
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 == "erai":
ta,temp1,hur,hgt,terrain,p,ps,wap,ua,va,uas,vas,tas,ta2d,\
cp,tp,wg10,mod_cape,lon,lat,date_list = \
read_erai(domain,time)
cp = cp.astype("float32", order="C")
tp = tp.astype("float32", order="C")
mod_cape = mod_cape.astype("float32", order="C")
elif model == "era5":
if ub4:
ta,temp1,hur,hgt,terrain,p,ps,ua,va,uas,vas,tas,ta2d,\
cp,wg10,mod_cape,lon,lat,date_list = \
read_era5(domain,time,delta_t=delta_t)
else:
ta,temp1,hur,hgt,terrain,p,ps,ua,va,uas,vas,tas,ta2d,\
cp,tp,wg10,mod_cape,lon,lat,date_list = \
read_era5_rt52(domain,time,delta_t=delta_t)
cp = cp.astype("float32", order="C")
tp = tp.astype("float32", order="C")
mod_cape = mod_cape.astype("float32", order="C")
wap = np.zeros(hgt.shape)
elif model == "barra":
ta,temp1,hur,hgt,terrain,p,ps,wap,ua,va,uas,vas,tas,ta2d,wg10,lon,lat,date_list = \
read_barra(domain,time)
elif model == "barra_fc":
ta,temp1,hur,hgt,terrain,p,ps,wap,ua,va,uas,vas,tas,ta2d,wg10,lon,lat,date_list = \
read_barra_fc(domain,time)
elif model == "barpa":
ta,hur,hgt,terrain,p,ps,ua,va,uas,vas,tas,ta2d,wg10,lon,lat,date_list = \
read_barpa(domain, time, experiment, barpa_forcing_mdl, ensemble)
wap = np.zeros(hgt.shape)
temp1 = None
elif model == "barra_ad":
wg10,temp2,ta,temp1,hur,hgt,terrain,p,ps,wap,ua,va,uas,vas,tas,ta2d,lon,lat,date_list = \
read_barra_ad(domain, time, False)
elif 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)
wap = np.zeros(hgt.shape)
wg10 = np.zeros(ps.shape)
mod_cape = np.zeros(ps.shape)
p = np.zeros(p_3d[0, :, 0, 0].shape)
#date_list = pd.to_datetime(date_list).to_pydatetime()
temp1 = None
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)
wap = np.zeros(hgt.shape)
wg10 = np.zeros(ps.shape)
p = np.zeros(p_3d[0, :, 0, 0].shape)
#date_list = pd.to_datetime(date_list).to_pydatetime()
temp1 = None
else:
raise ValueError("Model not recognised")
del temp1
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")
wap = wap.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")
wg10 = wg10.astype("float32", order="C")
lon = lon.astype("float32", order="C")
lat = lat.astype("float32", order="C")
gc.collect()
param = np.array([
"mu_cape", "mu_cin", "muq", "s06", "s0500", "lr700_500", "mhgt",
"ta500", "tp"
])
if model in ["erai", "era5"]:
param = np.concatenate([param, ["mod_cape"]])
#Option to interpolate to the ERA5 grid
if era5_interp:
#Interpolate model data to the ERA5 grid
from era5_read import get_lat_lon_rt52 as get_era5_lat_lon
era5_lon, era5_lat = get_era5_lat_lon()
era5_lon_ind = np.where((era5_lon >= domain[2])
& (era5_lon <= domain[3]))[0]
era5_lat_ind = np.where((era5_lat >= domain[0])
& (era5_lat <= domain[1]))[0]
era5_lon = era5_lon[era5_lon_ind]
era5_lat = era5_lat[era5_lat_ind]
terrain = interp_era5(terrain, lon, lat, era5_lon, era5_lat, d3=False)
#Set output array
output_data = np.zeros(
(ps.shape[0], era5_lat.shape[0], era5_lon.shape[0], len(param)))
else:
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
if era5_interp:
p_3d = np.moveaxis(np.tile(p,[ta.shape[2],ta.shape[3],1]),[0,1,2],[1,2,0]).\
astype(np.float32)
else:
p_3d = np.moveaxis(np.tile(p,[era5_lat.shape[0],era5_lon.shape[0],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]):
cape_start = dt.datetime.now()
if era5_interp:
ta_t = interp_era5(ta[t], lon, lat, era5_lon, era5_lat, d3=True)
hur_t = interp_era5(hur[t], lon, lat, era5_lon, era5_lat, d3=True)
hgt_t = interp_era5(hgt[t], lon, lat, era5_lon, era5_lat, d3=True)
ps_t = interp_era5(ps[t], lon, lat, era5_lon, era5_lat, d3=False)
wap_t = interp_era5(wap[t], lon, lat, era5_lon, era5_lat, d3=True)
ua_t = interp_era5(ua[t], lon, lat, era5_lon, era5_lat, d3=True)
va_t = interp_era5(va[t], lon, lat, era5_lon, era5_lat, d3=True)
uas_t = interp_era5(uas[t], lon, lat, era5_lon, era5_lat, d3=False)
vas_t = interp_era5(vas[t], lon, lat, era5_lon, era5_lat, d3=False)
tas_t = interp_era5(tas[t], lon, lat, era5_lon, era5_lat, d3=False)
ta2d_t = interp_era5(ta2d[t],
lon,
lat,
era5_lon,
era5_lat,
d3=False)
tp_t = interp_era5(tp[t], lon, lat, era5_lon, era5_lat, d3=False)
mod_cape_t = interp_era5(mod_cape[t],
lon,
lat,
era5_lon,
era5_lat,
d3=False)
else:
ta_t = ta[t]
hur_t = hur[t]
hgt_t = hgt[t]
ps_t = ps[t]
wap_t = wap[t]
ua_t = ua[t]
va_t = va[t]
uas_t = uas[t]
vas_t = vas[t]
tas_t = tas[t]
ta2d_t = ta2d[t]
tp_t = tp[t]
mod_cape_t = mod_cape[t]
print(date_list[t])
output = np.zeros((1, ps_t.shape[0], ps_t.shape[1], len(param)))
if mdl_lvl:
if era5_interp:
p_3d = interp_era5(full_p3d[t],
lon,
lat,
era5_lon,
era5_lat,
d3=True)
else:
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)
sfc_wap = np.insert(wap_t, 0, np.zeros(vas_t.shape), 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
hur_unit = mpcalc.relative_humidity_from_dewpoint(ta_t*units.units.degC, dp*units.units.degC)*\
100*units.units.percent
q_unit = mpcalc.mixing_ratio_from_relative_humidity(hur_unit,\
ta_t*units.units.degC,np.array(p_3d)*units.units.hectopascals)
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_theta_unit = mpcalc.potential_temperature(sfc_p_unit, sfc_ta_unit)
sfc_thetae_unit = mpcalc.equivalent_potential_temperature(
sfc_p_unit, sfc_ta_unit, sfc_dp_unit)
sfc_thetae = np.array(mpcalc.equivalent_potential_temperature(ps_t*units.units.hectopascals,tas_t*units.units.degC,\
ta2d_t*units.units.degC))
sfc_q = np.array(sfc_q_unit)
sfc_hur = np.array(sfc_hur_unit)
#sfc_wb = np.array(wrf.wetbulb( sfc_p_3d*100, sfc_ta+273.15, sfc_q, units="degC"))
#Use getcape.f90
#cape_gb_mu1, cape_gb_mu4 = getcape_driver(sfc_p_3d, sfc_ta, sfc_dp, ps_t)
#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]
mu_cin = np.take_along_axis(cin, mu_cape_inds, 0)[0]
mu_lfc = np.take_along_axis(lfc, mu_cape_inds, 0)[0]
mu_lcl = np.take_along_axis(lcl, mu_cape_inds, 0)[0]
mu_el = np.take_along_axis(el, 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)
ta925 = get_var_p_lvl(np.copy(sfc_ta), sfc_p_3d, 925)
ta850 = get_var_p_lvl(np.copy(sfc_ta), sfc_p_3d, 850)
ta700 = get_var_p_lvl(np.copy(sfc_ta), sfc_p_3d, 700)
rho = mpcalc.density(
np.array(sfc_p_3d) * (units.units.hectopascal),
sfc_ta * units.units.degC, sfc_q_unit)
rho925 = np.array(get_var_p_lvl(np.array(rho), sfc_p_3d, 925))
rho850 = np.array(get_var_p_lvl(np.array(rho), sfc_p_3d, 850))
rho700 = np.array(get_var_p_lvl(np.array(rho), sfc_p_3d, 700))
#Winds
winds_start = dt.datetime.now()
s06 = get_shear_hgt(sfc_ua, sfc_va, np.copy(sfc_hgt), 0, 6000, terrain)
s0500 = get_shear_p(ua_t,
va_t,
p_3d,
"sfc",
np.array([500]),
p_3d,
uas=uas_t,
vas=vas_t)[0]
#WAP
if model in ["erai", "era5"]:
sfc_w = mpcalc.vertical_velocity( wap_t * (units.units.pascal / units.units.second),\
np.array(p_3d) * (units.units.hectopascal), \
ta_t * units.units.degC, q_unit)
w925 = np.array(get_var_p_lvl(np.array(sfc_w), p_3d, 925))
w850 = np.array(get_var_p_lvl(np.array(sfc_w), p_3d, 850))
w700 = np.array(get_var_p_lvl(np.array(sfc_w), p_3d, 700))
#Convergence
if era5_interp:
x, y = np.meshgrid(era5_lon, era5_lat)
else:
x, y = np.meshgrid(lon, lat)
dx, dy = mpcalc.lat_lon_grid_deltas(x, y)
u925 = np.array(get_var_p_lvl(np.copy(sfc_ua), sfc_p_3d, 925))
u850 = np.array(get_var_p_lvl(np.copy(sfc_ua), sfc_p_3d, 850))
u700 = np.array(get_var_p_lvl(np.copy(sfc_ua), sfc_p_3d, 700))
v925 = np.array(get_var_p_lvl(np.copy(sfc_va), sfc_p_3d, 925))
v850 = np.array(get_var_p_lvl(np.copy(sfc_va), sfc_p_3d, 850))
v700 = np.array(get_var_p_lvl(np.copy(sfc_va), sfc_p_3d, 700))
conv925 = -1e5 * np.array(
mpcalc.divergence(u925 * (units.units.meter / units.units.second),
v925 * (units.units.meter / units.units.second),
dx, dy))
conv850 = -1e5 * np.array(
mpcalc.divergence(u850 * (units.units.meter / units.units.second),
v850 * (units.units.meter / units.units.second),
dx, dy))
conv700 = -1e5 * np.array(
mpcalc.divergence(u700 * (units.units.meter / units.units.second),
v700 * (units.units.meter / units.units.second),
dx, dy))
#CS6
mucs6 = mu_cape * np.power(s06, 1.67)
#Fill output
output = fill_output(output, t, param, ps, "mu_cape", mu_cape)
output = fill_output(output, t, param, ps, "mu_cin", mu_cin)
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, "s0500", s0500)
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, "mhgt", melting_hgt)
output = fill_output(output, t, param, ps, "tp", tp_t)
if (model == "erai") | (model == "era5"):
output = fill_output(output, t, param, ps, "mod_cape", mod_cape_t)
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:
if era5_interp:
save_netcdf(region, model, out_name, date_list, era5_lat, era5_lon, param, param_out, \
out_dtype = "f4", compress=True)
else:
save_netcdf(region, model, out_name, date_list, lat, lon, param, param_out, \
out_dtype = "f4", compress=True)
print(dt.datetime.now() - tot_start)
geo_mag = mpcalc.wind_speed(ugeo, vgeo)
grad_mag = geo_mag.m - (geo_mag.m**2) / (f.magnitude * R)
ugrad, vgrad = mpcalc.wind_components(grad_mag * units('m/s'), wdir)
# Calculate Ageostrophic wind
uageo = ugrad - ugeo
vageo = vgrad - vgeo
# Compute QVectors
uqvect, vqvect = mpcalc.q_vector(ugeo, vgeo, T * units.degC, 500 * units.hPa,
dx, dy)
# Calculate divergence of the ageostrophic wind
div = mpcalc.divergence(uageo, vageo, dx, dy, dim_order='yx')
# Calculate Relative Vorticity Advection
relvor = mpcalc.vorticity(ugeo, vgeo, dx, dy, dim_order='yx')
adv = mpcalc.advection(relvor, (ugeo, vgeo), (dx, dy), dim_order='yx')
######################################################################
# Create figure containing Geopotential Heights, Temperature, Divergence
# of the Ageostrophic Wind, Relative Vorticity Advection (shaded),
# geostrphic wind barbs, and Q-vectors.
#
fig = plt.figure(figsize=(10, 10))
ax = plt.subplot(111)
# Plot Geopotential Height Contours
import numpy as np
import metpy.calc as mpcalc
import matplotlib.pyplot as plt
start = timeit.default_timer()
data_path = Path('/Users/mret0001/Desktop/')
glob_pattern_2d = 'v900_*.nc'
v_files = sorted([str(f) for f in data_path.rglob(f'{glob_pattern_2d}')])
glob_pattern_2d = 'u900_*.nc'
u_files = sorted([str(f) for f in data_path.rglob(f'{glob_pattern_2d}')])
for ufile, vfile in zip(u_files, v_files):
date = ufile[-24:-11]
u_wind = xr.open_dataarray(home+f'/Desktop/u900_{date}_reggrid.nc')
v_wind = xr.open_dataarray(home+f'/Desktop/v900_{date}_reggrid.nc')
div = xr.DataArray(mpcalc.divergence(u_wind, v_wind).squeeze())
ds = div.to_dataset(name='div')
ds['div'].attrs['height'] = "900 hPa"
ds['div'].attrs['units'] = "1/s"
ds['div'].attrs['long_name'] = "Divergence"
ds.to_netcdf(home+f'/Desktop/div900_{date}_reggrid.nc')
del ds, u_wind, v_wind
stop = timeit.default_timer()
print(f'This script needed {stop-start} seconds.')
def metpy_read_wrf(tidx, froms, froms0, fwrf, lons_out, lats_out):
wrfin = Dataset(fwrf)
ds = xr.Dataset()
slp = getvar(wrfin, "slp", timeidx=tidx)
ds['slp'] = smooth2d(slp, 3, cenweight=4)
lats, lons = latlon_coords(slp)
ds['lats'] = lats
ds['lons'] = lons
landmask = extract_vars(wrfin, timeidx=tidx,
varnames=["LANDMASK"]).get("LANDMASK")
u10 = extract_vars(wrfin, timeidx=tidx, varnames=["U10"]).get("U10")
v10 = extract_vars(wrfin, timeidx=tidx, varnames=["V10"]).get("V10")
ds['u10'] = u10.where(landmask == 0)
ds['v10'] = v10.where(landmask == 0)
latent = extract_vars(wrfin, timeidx=tidx, varnames=["LH"]).get("LH")
ds['latent'] = smooth2d(latent, 3, cenweight=4) #latent.where(landmask==0)
t2m = extract_vars(wrfin, timeidx=tidx, varnames=["T2"]).get("T2") - 273.15
ds['t2m'] = smooth2d(t2m, 3, cenweight=4)
sst = extract_vars(wrfin, timeidx=tidx, varnames=["SST"]).get("SST")
ds['sst'] = sst.where(landmask == 0)
romsin = xr.open_dataset(froms)
romsin = romsin.rename({"lat_rho": "lat", "lon_rho": "lon"})
romsin = romsin.isel(ocean_time=tidx)
ds['sst_5m'] = romsin.isel(z_r=0).temp
ds['water_temp'] = romsin.temp
ds['water_ucur'] = romsin.ucur / 100
ds['water_vcur'] = romsin.vcur / 100
ds.water_ucur.attrs['units'] = 'm/s'
ds.water_vcur.attrs['units'] = 'm/s'
romsin = xr.open_dataset(froms0)
romsin = romsin.rename({"lat_rho": "lat", "lon_rho": "lon"})
romsin = romsin.isel(ocean_time=tidx)
ds['h'] = romsin.h
mld = get_oml_depth(froms0, t_in=tidx)
mld = smooth2d(mld, 3, cenweight=4)
ds['oml_depth'] = xr.DataArray(mld, ds.sst_5m.coords, ds.sst_5m.dims,
ds.sst_5m.attrs)
ds['oml_depth2'] = ds.oml_depth.where(ds.h > 20)
ds = ds.drop(['XLONG', 'XLAT', 'XTIME', 'Time'])
ds = ds.rename({'south_north': 'eta_rho', 'west_east': 'xi_rho'})
interp_method = 'bilinear'
ds_out = xr.Dataset({
'lat': (['lat'], lats_out),
'lon': (['lon'], lons_out)
})
regridder = xe.Regridder(ds, ds_out, interp_method)
regridder.clean_weight_file()
ds = regridder(ds)
ds = ds.squeeze()
dxy = 10000.
ds = ds.metpy.parse_cf().squeeze()
utau, vtau = wind_stress(to_np(ds.u10), to_np(ds.v10))
ds['u_tau'] = xr.DataArray(utau, ds.u10.coords, ds.u10.dims, ds.u10.attrs)
ds['v_tau'] = xr.DataArray(vtau, ds.v10.coords, ds.v10.dims, ds.v10.attrs)
curl = mpcalc.vorticity(utau * units('m/s'),
utau * units('m/s'),
dx=dxy * units.meter,
dy=dxy * units.meter)
ds['wind_stress_curl'] = xr.DataArray(np.array(curl), ds.u10.coords,
ds.u10.dims, ds.u10.attrs)
div = []
for z in range(len(ds.z_r)):
div0 = mpcalc.divergence(ds.water_ucur.isel(z_r=z) * units('m/s'),
ds.water_vcur.isel(z_r=z) * units('m/s'),
dx=dxy * units.meter,
dy=dxy * units.meter)
div.append(div0)
ds['cur_div'] = xr.DataArray(div, ds.water_ucur.coords, ds.water_ucur.dims,
ds.water_ucur.attrs)
div = mpcalc.divergence(ds.water_ucur.isel(z_r=2) * units('m/s'),
ds.water_vcur.isel(z_r=2) * units('m/s'),
dx=dxy * units.meter,
dy=dxy * units.meter)
ds['surf_cur_div'] = xr.DataArray(np.array(div), ds.u10.coords,
ds.u10.dims, ds.u10.attrs)
print(ds)
return ds
def ComputePlumbFluxesDecomp(uclm, vclm, z, zclm, lat, lon, lev):
#input: model climatology of u,v, z
#z for the event to analyze
#restringe domain to latitudes south to 0 (to avoid problems with sin (0°))
z = z[:, lat < 0, :]
zclm = zclm[lat < 0, :]
uclm = uclm[lat < 0, :]
vclm = vclm[lat < 0, :]
lat = lat[lat < 0]
[nrealiz, nlats, nlons] = np.shape(z) #get dimensions
#earth radius
a = 6400000
coslat = np.cos(lat * 3.14 / 180)
sinlat = np.sin(lat * 3.14 / 180)
#Coriolis parameter
f = 2 * 7.29e-5 * sinlat
#gravity
g = 9.8
# magnitude of basic state wind speed
magU = np.sqrt(np.add(np.power(uclm, 2), np.power(vclm, 2)))
#QG stream function
psia_clim = g / np.transpose(np.tile(f, (nlons, 1))) * zclm
psia = g / np.transpose(np.tile(f, (nrealiz, nlons, 1)), (0, 2, 1)) * z
psia_m = np.mean(psia, axis=0)
#anomalies with respecto to the ensemble mean
psiaa = psia - psia_m
psia_m = psia_m - psia_clim
#psia_clim = psia_clim - np.transpose(np.tile(np.mean(psia_clim, axis=1), (nlons, 1)))
#psi derivatives
dpsiamdlon = c_diff(psia_m, lon * 3.14 / 180, 1)
ddpsiamdlonlon = c_diff(dpsiamdlon, lon * 3.14, 1)
dpsiamdlat = c_diff(psia_m, lat * 3.14 / 180, 0)
ddpsiamdlatlat = c_diff(dpsiamdlat, lat * 3.14 / 180, 0)
ddpsiamdlatlon = c_diff(dpsiamdlat, lon * 3.14 / 180, 1)
dpsiacdlon = c_diff(psia_clim, lon * 3.14 / 180, 1)
ddpsiacdlonlon = c_diff(dpsiacdlon, lon * 3.14, 1)
dpsiacdlat = c_diff(psia_clim, lat * 3.14 / 180, 0)
ddpsiacdlatlat = c_diff(dpsiacdlat, lat * 3.14 / 180, 0)
ddpsiacdlatlon = c_diff(dpsiacdlat, lon * 3.14 / 180, 1)
dpsidlon = c_diff(psiaa, lon * 3.14 / 180, 2)
ddpsidlonlon = c_diff(dpsidlon, lon * 3.14, 2)
dpsidlat = c_diff(psiaa, lat * 3.14 / 180, 1)
ddpsidlatlat = c_diff(dpsidlat, lat * 3.14 / 180, 1)
ddpsidlatlon = c_diff(dpsidlat, lon * 3.14 / 180, 2)
#linear terms
termx1l = 2 * dpsiacdlon * dpsiamdlon - psia_clim * ddpsiamdlonlon - psia_m * ddpsiacdlonlon
termxyl = dpsiacdlon * dpsiamdlat + dpsiamdlon * dpsiacdlat -\
psia_clim * ddpsiamdlatlon - psia_m * ddpsiacdlatlon
termy1l = 2 * dpsiacdlat * dpsiamdlat - psia_clim * ddpsiamdlatlat - psia_m * ddpsiacdlatlat
# "p" is normalized by 1000hPa
coefcos = np.transpose(np.tile(coslat, (nlons, 1)))
coeff = coefcos * (lev / 100000) / (2 * magU)
#coeff = 1 / (2 * magU)
#x-component
pxl = coeff / (a * a * coefcos) * (uclm * termx1l / coefcos +
vclm * termxyl)
#px = coeff * (uclm * termxu + vclm * termxv)
#y-component
pyl = coeff / (a * a) * (uclm / coefcos * termxyl + vclm * termy1l)
#py = coeff * ( uclm * termxv + vclm * termyv)
#non-linear terms
termx1nl = dpsiamdlon**2 - psia_m * ddpsiamdlonlon
termxynl = dpsiamdlat * dpsiamdlon - psia_m * ddpsiamdlatlon
termy1nl = dpsiamdlat**2 - psia_m * ddpsiamdlatlat
#x-component
pxnl = coeff / (a * a * coefcos) * (uclm * termx1nl / coefcos +
vclm * termxynl)
#px = coeff * (uclm * termxu + vclm * termxv)
#y-component
pynl = coeff / (a * a) * (uclm / coefcos * termxynl + vclm * termy1nl)
#py = coeff * ( uclm * termxv + vclm * termyv)
#compute divergence
dx, dy = calc.lat_lon_grid_deltas(lon, lat)
divl = calc.divergence(pxl, pyl, dx, dy)
divnl = calc.divergence(pxnl, pynl, dx, dy)
pxem = pxl + pxnl
pyem = pyl + pynl
divem = calc.divergence(pxem, pyem, dx, dy)
var = {
'pxl': pxl,
'pyl': pyl,
'pxnl': pxnl,
'pynl': pynl,
'divl': divl,
'divnl': divnl,
'pxem': pxem,
'pyem': pyem,
'divem': divem,
'lat': lat
}
return var
