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