def test_divergence_asym(): """Test divergence calculation with a complicated field.""" u = np.array([[2, 4, 8], [0, 2, 2], [4, 6, 8]]) * units('m/s') v = np.array([[6, 4, 8], [2, 6, 0], [2, 2, 6]]) * units('m/s') c = divergence(u, v, 1 * units.meters, 2 * units.meters, dim_order='yx') true_c = np.array([[-2, 5.5, -2.5], [2., 0.5, -1.5], [3., -1.5, 8.5]]) / units.sec assert_array_equal(c, true_c) # Now try for xy ordered c = divergence(u.T, v.T, 1 * units.meters, 2 * units.meters, dim_order='xy') assert_array_equal(c, true_c.T)
def test_zero_divergence(): """Test divergence calculation when zeros should be returned.""" a = np.arange(3) u = np.c_[a, a, a] * units('m/s') c = divergence(u, u.T, 1 * units.meter, 1 * units.meter, dim_order='xy') true_c = 2. * np.ones_like(u) / units.sec assert_array_equal(c, true_c)
def test_divergence(): """Test divergence for simple case.""" a = np.arange(3) u = np.c_[a, a, a] * units('m/s') c = divergence(u, u, 1 * units.meter, 1 * units.meter, dim_order='xy') true_c = np.ones_like(u) / units.sec assert_array_equal(c, true_c)
def windProducts(x, y, u, v): dx, dy = mpcalc.lat_lon_grid_spacing(x, y) deformation = numpy.zeros(u.shape) convergence = numpy.zeros(u.shape) for i in range(u.shape[0]): shear, stretch = mpcalc.shearing_stretching_deformation( u[i, :, :], v[i, :, :], dx, dy) deformation[i] = numpy.sqrt(shear**2 + stretch**2) convergence[i] = -mpcalc.divergence(u[i, :, :], v[i, :, :], dx, dy) return deformation, convergence
def kinematics(u, v, thetae, dx, dy, lats, smooth=False, sigma=1): ''' Use metpy functions to calculate various kinematics, given 2d numpy arrays as inputs. X and y grid spacing, as well as a 2d array of latitudes is also required. Option to smooth the thetae field before calculating front diagnostics (recommended), with addition "sigma" parameter controlling the smoothness (see scipy docs) ''' if smooth: thetae = gaussian_filter(thetae, sigma) #Kinematics ddy_thetae = mpcalc.first_derivative(thetae, delta=dy, axis=0) ddx_thetae = mpcalc.first_derivative(thetae, delta=dx, axis=1) mag_thetae = np.sqrt(ddx_thetae**2 + ddy_thetae**2) div = mpcalc.divergence(u, v, dx, dy) strch_def = mpcalc.stretching_deformation(u, v, dx, dy) shear_def = mpcalc.shearing_deformation(u, v, dx, dy) tot_def = mpcalc.total_deformation(u, v, dx, dy) psi = 0.5 * np.arctan2(shear_def, strch_def) beta = np.arcsin( (-ddx_thetae * np.cos(psi) - ddy_thetae * np.sin(psi)) / mag_thetae) vo = mpcalc.vorticity(u, v, dx, dy) conv = -div * 1e5 F = np.array(0.5 * mag_thetae * (tot_def * np.cos(2 * beta) - div) * 1.08e4 * 1e5) Fn = np.array(0.5 * mag_thetae * (div - tot_def * np.cos(2 * beta)) * 1.08e4 * 1e5) Fs = np.array(0.5 * mag_thetae * (vo + tot_def * np.sin(2 * beta)) * 1.08e4 * 1e5) icon = np.array(0.5 * (tot_def - div) * 1e5) vgt = np.array(np.sqrt(div**2 + vo**2 + tot_def**2) * 1e5) #TFP ddx_thetae_scaled = ddx_thetae / mag_thetae ddy_thetae_scaled = ddy_thetae / mag_thetae ddy_mag_thetae = -1 * mpcalc.first_derivative(mag_thetae, delta=dy, axis=0) ddx_mag_thetae = -1 * mpcalc.first_derivative(mag_thetae, delta=dx, axis=1) tfp = ddx_mag_thetae * ddx_thetae_scaled + ddy_mag_thetae * ddy_thetae_scaled ddy_tfp = mpcalc.first_derivative(tfp, delta=dy, axis=0) ddx_tfp = mpcalc.first_derivative(tfp, delta=dx, axis=1) mag_tfp = np.sqrt(ddx_tfp**2 + ddy_tfp**2) v_f = (u * units.units("m/s")) * (ddx_tfp / mag_tfp) + ( v * units.units("m/s")) * (ddy_tfp / mag_tfp) #Extra condition ddy_ddy_mag_te = mpcalc.first_derivative(-1 * ddy_mag_thetae, delta=dy, axis=0) ddx_ddx_mag_te = mpcalc.first_derivative(-1 * ddx_mag_thetae, delta=dx, axis=1) cond = ddy_ddy_mag_te + ddx_ddx_mag_te return [F, Fn, Fs, icon, vgt, conv, vo*1e5, \ np.array(tfp.to("km^-2")*(100*100)), np.array(mag_thetae.to("km^-1"))*100, np.array(v_f), thetae, cond]
def compute_convergence(dset, uvar='10u', vvar='10v'): dx, dy = mpcalc.lat_lon_grid_deltas(dset['lon'], dset['lat']) conv = -mpcalc.divergence(dset[uvar], dset[vvar], dx[None, :, :], dy[None, :, :]) conv = xr.DataArray(conv.magnitude, coords=dset[uvar].coords, attrs={ 'standard_name': 'convergence', 'units': conv.units }, name='conv') return xr.merge([dset, conv])
def divergence(self, level): """ Uses a metpy function to calculate wind divergence in a given level. Uses predefined fuctions to obtain wind and grid data. Returns a np.array. """ # Grab lat/lon values lat = self.data['lat'].values lon = self.data['lon'].values # Compute dx and dy spacing for use in divergence calculation dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat) # Extract wind components uwnd = self.u_wind(level) vwnd = self.v_wind(level) # Use MetPy to compute the divergence for the given wind level div = mpcalc.divergence(uwnd, vwnd, dx, dy) return np.array(div * 100000)
winds_clm = winds.mean(dim='realiz') for i in np.arange(7): var_WPV_EN = np.mean(hgt.z.values[i, index_WPV_EN, :, :], axis=0) var_WPV_LN = np.mean(hgt.z.values[i, index_WPV_LN, :, :], axis=0) var_normal = np.mean(hgt.z.values[i, :, :, :], axis=0) var_SPV_EN = np.mean(hgt.z.values[i, index_SPV_EN, :, :], axis=0) var_SPV_LN = np.mean(hgt.z.values[i, index_SPV_LN, :, :], axis=0) var_WPV_all = np.mean(hgt.z.values[i, index_WPV_all.values, :, :], axis=0) var_SPV_all = np.mean(hgt.z.values[i, index_SPV_all.values, :, :], axis=0) px_WPV, py_WPV, lat = plumb_flux.ComputePlumbFluxes( winds_clm.u.values[i, :, :], winds_clm.v.values[i, :, :], var_WPV_all - var_normal, np.zeros_like(var_WPV_EN), hgt.latitude.values, hgt.longitude.values) dx, dy = calc.lat_lon_grid_deltas(hgt.longitude.values, lat) div_WPV_all = calc.divergence(px_WPV, py_WPV, dx, dy) px_SPV, py_SPV, lat = plumb_flux.ComputePlumbFluxes( winds_clm.u.values[i, :, :], winds_clm.v.values[i, :, :], var_SPV_all - var_normal, np.zeros_like(var_WPV_EN), hgt.latitude.values, hgt.longitude.values) div_SPV_all = calc.divergence(px_SPV, py_SPV, dx, dy) px_WPV_EN, py_WPV_EN, lat = plumb_flux.ComputePlumbFluxes( winds_clm.u.values[i, :, :], winds_clm.v.values[i, :, :], var_WPV_EN - var_normal, np.zeros_like(var_WPV_EN), hgt.latitude.values, hgt.longitude.values) div_WPV_EN = calc.divergence(px_WPV_EN, py_WPV_EN, dx, dy) px_SPV_EN, py_SPV_EN, lat = plumb_flux.ComputePlumbFluxes( winds_clm.u.values[i, :, :],
ax.set_extent([WLON, ELON, SLAT, NLAT]) ax.add_feature(cfeature.COASTLINE.with_scale('50m'), linewidth=0.5) ax.add_feature(cfeature.LAKES.with_scale('50m'), linewidth=0.5) ax.add_feature(cfeature.STATES, linewidth=0.5) ax.add_feature(cfeature.BORDERS, linewidth=0.5) return ax # ============================================================================= # FIG #1: 250: JET STREAM, GEOPOTENTIAL HEIGHT, DIVERGENCE # ============================================================================= H250 = HGT_DATA.variables['hgt'][TIME_INDEX, 8, :, :] U250 = UWND_DATA.variables['uwnd'][TIME_INDEX, 8, :, :] * units('m/s') V250 = VWND_DATA.variables['vwnd'][TIME_INDEX, 8, :, :] * units('m/s') SPEED250 = mpcalc.get_wind_speed(U250, V250) DIV250 = mpcalc.divergence(U250, V250, DX, DY, dim_order='YX') DIV250 = (DIV250 * (units('1/s'))) # ============================================================================= # FIG #2: 500: VORTICITY, GEOPOTENTIAL HEIGHT, VORTICITY ADVECTION # ============================================================================= H500 = HGT_DATA.variables['hgt'][TIME_INDEX, 5, :, :] U500 = UWND_DATA.variables['uwnd'][TIME_INDEX, 5, :, :] * units('m/s') V500 = VWND_DATA.variables['vwnd'][TIME_INDEX, 5, :, :] * units('m/s') DX, DY = mpcalc.lat_lon_grid_spacing(LON, LAT) VORT500 = mpcalc.vorticity(U500, V500, DX, DY, dim_order='YX') VORT500 = (VORT500 * (units('1/s'))) VORT_ADV500 = mpcalc.advection(VORT500, [U500, V500], (DX, DY), dim_order='yx') # ============================================================================= # FIG #3: 700: Q-VECTORS+CONVERGENCE, GEOPOTENTIAL HEIGHT # ============================================================================= H700 = HGT_DATA.variables['hgt'][TIME_INDEX, 3, :, :]
def metpy_read_wrf_cross(fname, plevels, tidx_in, lons_out, lats_out, start, end): ds = xr.open_dataset(fname).metpy.parse_cf().squeeze() ds = ds.isel(Time=tidx) print(ds) ds1 = xr.Dataset() p = units.Quantity(to_np(ds.p), 'hPa') z = units.Quantity(to_np(ds.z), 'meter') u = units.Quantity(to_np(ds.u), 'm/s') v = units.Quantity(to_np(ds.v), 'm/s') w = units.Quantity(to_np(ds.w), 'm/s') tk = units.Quantity(to_np(ds.tk), 'kelvin') th = units.Quantity(to_np(ds.th), 'kelvin') eth = units.Quantity(to_np(ds.eth), 'kelvin') wspd = units.Quantity(to_np(ds.wspd), 'm/s') omega = units.Quantity(to_np(ds.omega), 'Pa/s') plevels_unit = plevels * units.hPa z, u, v, w, tk, th, eth, wspd, omega = log_interpolate_1d(plevels_unit, p, z, u, v, w, tk, th, eth, wspd, omega, axis=0) coords, dims = [plevs, ds.lat.values, ds.lon.values], ["level", "lat", "lon"] for name, var in zip( ['z', 'u', 'v', 'w', 'tk', 'th', 'eth', 'wspd', 'omega'], [z, u, v, w, tk, th, eth, wspd, omega]): #g = ndimage.gaussian_filter(var, sigma=3, order=0) ds1[name] = xr.DataArray(to_np(mpcalc.smooth_n_point(var, 9)), coords=coords, dims=dims) dx, dy = mpcalc.lat_lon_grid_deltas(ds.lon.values * units('degrees_E'), ds.lat.values * units('degrees_N')) # Calculate temperature advection using metpy function for i, plev in enumerate(plevs): adv = mpcalc.advection(eth[i, :, :], [u[i, :, :], v[i, :, :]], (dx, dy), dim_order='yx') * units('K/sec') adv = ndimage.gaussian_filter(adv, sigma=3, order=0) * units('K/sec') ds1['eth_adv_{:03d}'.format(plev)] = xr.DataArray( np.array(adv), coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) div = mpcalc.divergence(u[i, :, :], v[i, :, :], dx, dy, dim_order='yx') div = ndimage.gaussian_filter(div, sigma=3, order=0) * units('1/sec') ds1['div_{:03d}'.format(plev)] = xr.DataArray( np.array(div), coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['accrain'] = xr.DataArray(ds.accrain.values, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) eth2 = mpcalc.equivalent_potential_temperature( ds.slp.values * units.hPa, ds.t2m.values * units('K'), ds.td2.values * units('celsius')) ds1['eth2'] = xr.DataArray(eth2, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) #ds1['sst'] = xr.DataArray(ndimage.gaussian_filter(ds.sst.values, sigma=3, order=0)-273.15, coords=[ds.lat.values,ds.lon.values], dims=["lat","lon"]) ds1 = ds1.metpy.parse_cf().squeeze() cross = cross_section(ds1, start, end).set_coords(('lat', 'lon')) cross.u.attrs['units'] = 'm/s' cross.v.attrs['units'] = 'm/s' cross['t_wind'], cross['n_wind'] = mpcalc.cross_section_components( cross['u'], cross['v']) weights = np.cos(np.deg2rad(ds.lat)) ds_weighted = ds.weighted(weights) weighted = ds_weighted.mean(("lat")) return ds1, cross, weighted
def gh_uv_div(initTime=None, fhour=6, day_back=0,model='GRAPES_GFS', gh_lev=500,uv_lev=850, map_ratio=14/9,zoom_ratio=20,cntr_pnt=[104,34], south_China_sea=True,area =None,city=False,output_dir=None,data_source='MICAPS',**kwargs): if(area != None): south_China_sea=False # micaps data directory if(data_source=='MICAPS'): try: data_dir = [utl.Cassandra_dir(data_type='high',data_source=model,var_name='HGT',lvl=gh_lev), utl.Cassandra_dir(data_type='high',data_source=model,var_name='UGRD',lvl=uv_lev), utl.Cassandra_dir(data_type='high',data_source=model,var_name='VGRD',lvl=uv_lev), utl.Cassandra_dir(data_type='surface',data_source=model,var_name='PSFC')] except KeyError: raise ValueError('Can not find all directories needed') # get filename if(initTime != None): filename = utl.model_filename(initTime, fhour) else: filename=utl.filename_day_back_model(day_back=day_back,fhour=fhour) # retrieve data from micaps server gh = MICAPS_IO.get_model_grid(data_dir[0], filename=filename) if gh is None: return u = MICAPS_IO.get_model_grid(data_dir[1], filename=filename) if u is None: return v = MICAPS_IO.get_model_grid(data_dir[2], filename=filename) if v is None: return psfc = MICAPS_IO.get_model_grid(data_dir[3], filename=filename) init_time = gh.coords['forecast_reference_time'].values if(data_source =='CIMISS'): # get filename if(initTime != None): filename = utl.model_filename(initTime, fhour,UTC=True) else: filename=utl.filename_day_back_model(day_back=day_back,fhour=fhour,UTC=True) # retrieve data from CIMISS server try: gh=CMISS_IO.cimiss_model_by_time('20'+filename[0:8],valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='GPH'), levattrs={'long_name':'pressure_level', 'units':'hPa', '_CoordinateAxisType':'-'}, fcst_level=gh_lev, fcst_ele="GPH", units='gpm') if gh is None: return gh['data'].values=gh['data'].values/10. u=CMISS_IO.cimiss_model_by_time('20'+filename[0:8],valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='WIU'), levattrs={'long_name':'pressure_level', 'units':'hPa', '_CoordinateAxisType':'-'}, fcst_level=uv_lev, fcst_ele="WIU", units='m/s') if u is None: return v=CMISS_IO.cimiss_model_by_time('20'+filename[0:8],valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='WIV'), levattrs={'long_name':'pressure_level', 'units':'hPa', '_CoordinateAxisType':'-'}, fcst_level=uv_lev, fcst_ele="WIV", units='m/s') if v is None: return psfc=CMISS_IO.cimiss_model_by_time('20'+filename[0:8], valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='PRS'), fcst_level=0, fcst_ele="PRS", units='Pa') psfc['data']=psfc['data']/100. except KeyError: raise ValueError('Can not find all data needed') # prepare data if(area != None): cntr_pnt,zoom_ratio=utl.get_map_area(area_name=area) map_extent,delt_x,delt_y=utl.get_map_extent(cntr_pnt, zoom_ratio, map_ratio) #+ to solve the problem of labels on all the contours dx,dy=mpcalc.lat_lon_grid_deltas(u['lon'].values.squeeze(),u['lat'].values.squeeze()) div=mpcalc.divergence(u['data'].values.squeeze()*units('m/s'), v['data'].values.squeeze()*units('m/s'), dx, dy, dim_order='yx') div_xr=u.copy(deep=True) div_xr['data'].values=div.magnitude[np.newaxis,np.newaxis,:,:] div_xr['data'].values=gaussian_filter(div_xr['data'].values,1) gh=utl.mask_terrian(gh_lev,psfc,gh) u=utl.mask_terrian(uv_lev,psfc,u) v=utl.mask_terrian(uv_lev,psfc,v) div_xr=utl.mask_terrian(uv_lev,psfc,div_xr) gh=utl.cut_xrdata(map_extent, gh, delt_x=delt_x, delt_y=delt_y) u=utl.cut_xrdata(map_extent, u, delt_x=delt_x, delt_y=delt_y) v=utl.cut_xrdata(map_extent, v, delt_x=delt_x, delt_y=delt_y) div_xr=utl.cut_xrdata(map_extent, div_xr, delt_x=delt_x, delt_y=delt_y) gh.attrs['model']=model div_xr.attrs['units']='s$^{-1}$' uv=xr.merge([u.rename({'data': 'u'}),v.rename({'data': 'v'})]) #- to solve the problem of labels on all the contours dynamic_graphics.draw_gh_uv_div( div=div_xr, gh=gh, uv=uv, map_extent=map_extent, regrid_shape=20, city=city,south_China_sea=south_China_sea, output_dir=output_dir)
def horizontal_map(variable_name, date, start_hour, end_hour, pressure_level=False, subset=False, initiation=False, save=False, gif=False): '''This function plots the chosen variable for the analysis of the initiation environment on a horizontal (2D) map. Supported variables for plotting procedure are updraft, reflectivity, helicity, pw, cape, cin, ctt, temperature_surface, wind_shear, updraft_reflectivity, rh, omega, pvo, avo, theta_e, water_vapor, uv_wind and divergence.''' ### Predefine some variables ### # Get the list of all needed wrf files data_dir = '/scratch3/thomasl/work/data/casestudy_baden/' # Define save directory save_dir = '/scratch3/thomasl/work/retrospective_part' '/casestudy_baden/horizontal_maps/' # Change extent of plot subset_extent = [6.2, 9.4, 46.5, 48.5] # Set the location of the initiation of the thunderstorm initiation_location = CoordPair(lat=47.25, lon=7.85) # 2D variables: if variable_name == 'updraft': variable_name = 'W_UP_MAX' title_name = 'Maximum Z-Wind Updraft' colorbar_label = 'Max Z-Wind Updraft [$m$ $s^-$$^1$]' save_name = 'updraft' variable_min = 0 variable_max = 30 # Check if a certain pressure_level was defined. if pressure_level != False: sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name)) elif variable_name == 'reflectivity': variable_name = 'REFD_MAX' title_name = 'Maximum Derived Radar Reflectivity' colorbar_label = 'Maximum Derived Radar Reflectivity [$dBZ$]' save_name = 'reflectivity' variable_min = 0 variable_max = 75 # Check if a certain pressure_level was defined. if pressure_level != False: sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name)) elif variable_name == 'helicity': variable_name = 'UP_HELI_MAX' title_name = 'Maximum Updraft Helicity' colorbar_label = 'Maximum Updraft Helicity [$m^{2}$ $s^{-2}$]' save_name = 'helicity' variable_min = 0 variable_max = 140 # Check if a certain pressure_level was defined. if pressure_level != False: sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name)) elif variable_name == 'pw': title_name = 'Precipitable Water' colorbar_label = 'Precipitable Water [$kg$ $m^{-2}$]' save_name = 'pw' variable_min = 0 variable_max = 50 # Check if a certain pressure_level was defined. if pressure_level != False: sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name)) elif variable_name == 'cape': variable_name = 'cape_2d' title_name = 'CAPE' colorbar_label = 'Convective Available Potential Energy' '[$J$ $kg^{-1}$]' save_name = 'cape' variable_min = 0 variable_max = 3000 # Check if a certain pressure_level was defined. if pressure_level != False: sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name)) elif variable_name == 'cin': variable_name = 'cape_2d' title_name = 'CIN' colorbar_label = 'Convective Inhibition [$J$ $kg^{-1}$]' save_name = 'cin' variable_min = 0 variable_max = 100 # Check if a certain pressure_level was defined. if pressure_level != False: sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name)) elif variable_name == 'ctt': title_name = 'Cloud Top Temperature' colorbar_label = 'Cloud Top Temperature [$K$]' save_name = 'cct' variable_min = 210 variable_max = 300 # Check if a certain pressure_level was defined. if pressure_level != False: sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name)) elif variable_name == 'temperature_surface': variable_name = 'T2' title_name = 'Temperature @ 2 m' colorbar_label = 'Temperature [$K$]' save_name = 'temperature_surface' variable_min = 285 variable_max = 305 # Check if a certain pressure_level was defined. if pressure_level != False: sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name)) elif variable_name == 'wind_shear': variable_name = 'slp' title_name = 'SLP, Wind @ 850hPa, Wind @ 500hPa\n' 'and 500-850hPa Vertical Wind Shear' save_name = 'wind_shear' variable_min = 1000 variable_max = 1020 # Check if a certain pressure_level was defined. if pressure_level != False: sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name)) elif variable_name == 'updraft_reflectivity': variable_name = 'W_UP_MAX' title_name = 'Updraft and Reflectivity' colorbar_label = 'Max Z-Wind Updraft [$m$ $s^-$$^1$]' save_name = 'updraft_reflectivity' variable_min = 0 variable_max = 30 # Check if a certain pressure_level was defined. if pressure_level != False: sys.exit('The variable {} is a 2D variable. ' 'Definition of a pressure_level for ' 'plotting process is not required.'.format(variable_name)) # 3D variables: elif variable_name == 'rh': title_name = 'Relative Humidity' colorbar_label = 'Relative Humidity [$pct$]' save_name = 'rh' variable_min = 0 variable_max = 100 # Check if a certain pressure_level was defined. if pressure_level == False: sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name)) elif variable_name == 'omega': title_name = 'Vertical Motion' colorbar_label = 'Omega [$Pa$ $s^-$$^1$]' save_name = 'omega' variable_min = -50 variable_max = 50 # Check if a certain pressure_level was defined. if pressure_level == False: sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name)) elif variable_name == 'pvo': title_name = 'Potential Vorticity' colorbar_label = 'Potential Vorticity [$PVU$]' save_name = 'pvo' variable_min = -1 variable_max = 9 # Check if a certain pressure_level was defined. if pressure_level == False: sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name)) elif variable_name == 'avo': title_name = 'Absolute Vorticity' colorbar_label = 'Absolute Vorticity [$10^{-5}$' '$s^{-1}$]' save_name = 'avo' variable_min = -250 variable_max = 250 # Check if a certain pressure_level was defined. if pressure_level == False: sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name)) elif variable_name == 'theta_e': title_name = 'Theta-E' colorbar_label = 'Theta-E [$K$]' save_name = 'theta_e' variable_min = 315 variable_max = 335 # Check if a certain pressure_level was defined. if pressure_level == False: sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name)) elif variable_name == 'water_vapor': variable_name = 'QVAPOR' title_name = 'Water Vapor Mixing Ratio' colorbar_label = 'Water Vapor Mixing Ratio [$g$ $kg^{-1}$]' save_name = 'water_vapor' variable_min = 5 variable_max = 15 # Check if a certain pressure_level was defined. if pressure_level == False: sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name)) elif variable_name == 'uv_wind': variable_name = 'wspd_wdir' title_name = 'Wind Speed and Direction' colorbar_label = 'Wind Speed [$m$ $s^{-1}$]' save_name = 'uv_wind' variable_min = 0 variable_max = 10 # Check if a certain pressure_level was defined. if pressure_level == False: sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name)) elif variable_name == 'divergence': variable_name = 'ua' title_name = 'Horizontal Wind Divergence' colorbar_label = 'Divergence [$10^{-6}$ $s^{-1}$]' save_name = 'divergence' variable_min = -2.5 variable_max = 2.5 # Check if a certain pressure_level was defined. if pressure_level == False: sys.exit('The variable {} is a 3D variable. ' 'Definition of a pressure_level for ' 'plotting process is required.'.format(variable_name)) # Make a list of all wrf files in data directory wrflist = list() for (dirpath, dirnames, filenames) in os.walk(data_dir): wrflist += [os.path.join(dirpath, file) for file in filenames] ### Plotting Iteration ### # Iterate over a list of hourly timesteps time = list() for i in range(start_hour, end_hour): time = str(i).zfill(2) # Iterate over all 5 minutes steps of hour for j in range(0, 60, 5): minutes = str(j).zfill(2) # Load the netCDF files out of the wrflist ncfile = [Dataset(x) for x in wrflist if x.endswith('{}_{}:{}:00'.format(date, time, minutes))] # Load variable(s) if title_name == 'CAPE': variable = getvar(ncfile, variable_name)[0,:] elif title_name == 'CIN': variable = getvar(ncfile, variable_name)[1,:] elif variable_name == 'ctt': variable = getvar(ncfile, variable_name, units='K') elif variable_name == 'wspd_wdir': variable = getvar(ncfile, variable_name)[0,:] elif variable_name == 'QVAPOR': variable = getvar(ncfile, variable_name)*1000 # convert to g/kg else: variable = getvar(ncfile, variable_name) if variable_name == 'slp': slp = variable.squeeze() ua = getvar(ncfile, 'ua') va = getvar(ncfile, 'va') p = getvar(ncfile, 'pressure') u_wind850 = interplevel(ua, p, 850) v_wind850 = interplevel(va, p, 850) u_wind850 = u_wind850.squeeze() v_wind850 = v_wind850.squeeze() u_wind500 = interplevel(ua, p, 500) v_wind500 = interplevel(va, p, 500) u_wind500 = u_wind500.squeeze() v_wind500 = v_wind500.squeeze() slp = ndimage.gaussian_filter(slp, sigma=3, order=0) # Interpolating 3d data to a horizontal pressure level if pressure_level != False: p = getvar(ncfile, 'pressure') variable_pressure = interplevel(variable, p, pressure_level) variable = variable_pressure if variable_name == 'wspd_wdir': ua = getvar(ncfile, 'ua') va = getvar(ncfile, 'va') u_pressure = interplevel(ua, p, pressure_level) v_pressure = interplevel(va, p, pressure_level) elif title_name == 'Updraft and Reflectivity': reflectivity = getvar(ncfile, 'REFD_MAX') elif title_name == 'Difference in Theta-E values': variable = getvar(ncfile, variable_name) p = getvar(ncfile, 'pressure') variable_pressure1 = interplevel(variable, p, '950') variable_pressure2 = interplevel(variable, p, '950') elif variable_name == 'ua': va = getvar(ncfile, 'va') p = getvar(ncfile, 'pressure') v_pressure = interplevel(va, p, pressure_level) u_wind = variable.squeeze() v_wind = v_pressure.squeeze() u_wind.attrs['units']='meters/second' v_wind.attrs['units']='meters/second' lats, lons = latlon_coords(variable) lats = lats.squeeze() lons = lons.squeeze() dx, dy = mpcalc.lat_lon_grid_deltas(to_np(lons), to_np(lats)) divergence = mpcalc.divergence(u_wind, v_wind, dx, dy, dim_order='yx') divergence = divergence*1e3 # Define cart projection lats, lons = latlon_coords(variable) cart_proj = ccrs.LambertConformal(central_longitude=8.722206, central_latitude=46.73585) bounds = geo_bounds(wrfin=ncfile) # Create figure fig = plt.figure(figsize=(15, 10)) if variable_name == 'slp': fig.patch.set_facecolor('k') ax = plt.axes(projection=cart_proj) ### Set map extent ### domain_extent = [3.701088, 13.814863, 43.85472,49.49499] if subset == True: ax.set_extent([subset_extent[0],subset_extent[1], subset_extent[2],subset_extent[3]], ccrs.PlateCarree()) else: ax.set_extent([domain_extent[0]+0.7,domain_extent[1]-0.7, domain_extent[2]+0.1,domain_extent[3]-0.1], ccrs.PlateCarree()) # Plot contour of variables levels_num = 11 levels = np.linspace(variable_min, variable_max, levels_num) # Creating new colormap for diverging colormaps def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100): new_cmap = LinearSegmentedColormap.from_list( 'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval, b=maxval), cmap(np.linspace(minval, maxval, n))) return new_cmap cmap = plt.get_cmap('RdYlBu') if title_name == 'CIN': cmap = ListedColormap(sns.cubehelix_palette(levels_num-1, start=.5, rot=-.75, reverse=True)) variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable), levels=levels, transform=ccrs.PlateCarree(), extend='max', cmap=cmap) initiation_color = 'r*' elif variable_name == 'ctt': cmap = ListedColormap(sns.cubehelix_palette(levels_num-1, start=.5, rot=-.75, reverse=True)) variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable), levels=levels, transform=ccrs.PlateCarree(), extend='both', cmap=cmap) initiation_color = 'r*' elif variable_name == 'pvo': cmap = plt.get_cmap('RdYlBu_r') new_cmap = truncate_colormap(cmap, 0.05, 0.9) new_norm = DivergingNorm(vmin=-1., vcenter=2., vmax=10) variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable), levels=levels, transform=ccrs.PlateCarree(), cmap=new_cmap, extend='both', norm=new_norm) initiation_color = 'k*' elif variable_name == 'avo': cmap = plt.get_cmap('RdYlBu_r') new_cmap = truncate_colormap(cmap, 0.05, 0.9) new_norm = DivergingNorm(vmin=variable_min, vcenter=0, vmax=variable_max) variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable), levels=levels, transform=ccrs.PlateCarree(), cmap=new_cmap, extend='both', norm=new_norm) initiation_color = 'k*' elif variable_name == 'omega': new_cmap = truncate_colormap(cmap, 0.05, 0.9) new_norm = DivergingNorm(vmin=variable_min, vcenter=0, vmax=variable_max) variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable), levels=levels, transform=ccrs.PlateCarree(), cmap=new_cmap, extend='both', norm=new_norm) initiation_color = 'k*' elif variable_name == 'ua': new_cmap = truncate_colormap(cmap, 0.05, 0.9) new_norm = DivergingNorm(vmin=variable_min, vcenter=0, vmax=variable_max) variable_plot = plt.contourf(to_np(lons), to_np(lats), divergence, levels=levels, transform=ccrs.PlateCarree(), cmap=new_cmap, extend='both', norm=new_norm) initiation_color = 'k*' elif variable_name == 'UP_HELI_MAX' or variable_name == 'W_UP_MAX' or variable_name == 'QVAPOR': cmap = ListedColormap(sns.cubehelix_palette(levels_num-1, start=.5, rot=-.75)) variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable), levels=levels, extend='max', transform=ccrs.PlateCarree(),cmap=cmap) initiation_color = 'r*' elif variable_name == 'theta_e' or variable_name == 't2': cmap = ListedColormap(sns.cubehelix_palette(levels_num-1, start=.5, rot=-.75)) variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable), levels=levels, extend='both', transform=ccrs.PlateCarree(),cmap=cmap) initiation_color = 'r*' elif variable_name == 'REFD_MAX': levels = np.arange(5., 75., 5.) dbz_rgb = np.array([[4,233,231], [1,159,244], [3,0,244], [2,253,2], [1,197,1], [0,142,0], [253,248,2], [229,188,0], [253,149,0], [253,0,0], [212,0,0], [188,0,0],[248,0,253], [152,84,198]], np.float32) / 255.0 dbz_cmap, dbz_norm = from_levels_and_colors(levels, dbz_rgb, extend='max') variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable), levels=levels, extend='max', transform=ccrs.PlateCarree(), cmap=dbz_cmap, norm=dbz_norm) initiation_color = 'r*' elif variable_name == 'slp': ax.background_patch.set_fill(False) wslice = slice(1, None, 12) # Plot 850-hPa wind vectors vectors850 = ax.quiver(to_np(lons)[wslice, wslice], to_np(lats)[wslice, wslice], to_np(u_wind850)[wslice, wslice], to_np(v_wind850)[wslice, wslice], headlength=4, headwidth=3, scale=400, color='gold', label='850mb wind', transform=ccrs.PlateCarree(), zorder=2) # Plot 500-hPa wind vectors vectors500 = ax.quiver(to_np(lons)[wslice, wslice], to_np(lats)[wslice, wslice], to_np(u_wind500)[wslice, wslice], to_np(v_wind500)[wslice, wslice], headlength=4, headwidth=3, scale=400, color='cornflowerblue', zorder=2, label='500mb wind', transform=ccrs.PlateCarree()) # Plot 500-850 shear shear = ax.quiver(to_np(lons[wslice, wslice]), to_np(lats[wslice, wslice]), to_np(u_wind500[wslice, wslice]) - to_np(u_wind850[wslice, wslice]), to_np(v_wind500[wslice, wslice]) - to_np(v_wind850[wslice, wslice]), headlength=4, headwidth=3, scale=400, color='deeppink', zorder=2, label='500-850mb shear', transform=ccrs.PlateCarree()) contour = ax.contour(to_np(lons), to_np(lats), slp, levels=levels, colors='lime', linewidths=2, alpha=0.5, zorder=1, transform=ccrs.PlateCarree()) ax.clabel(contour, fontsize=12, inline=1, inline_spacing=4, fmt='%i') # Add a legend ax.legend(('850mb wind', '500mb wind', '500-850mb shear'), loc=4) # Manually set colors for legend legend = ax.get_legend() legend.legendHandles[0].set_color('gold') legend.legendHandles[1].set_color('cornflowerblue') legend.legendHandles[2].set_color('deeppink') initiation_color = 'w*' else: cmap = ListedColormap(sns.cubehelix_palette(10, start=.5, rot=-.75)) variable_plot = plt.contourf(to_np(lons), to_np(lats), to_np(variable), levels=levels, transform=ccrs.PlateCarree(),cmap=cmap) initiation_color = 'r*' # Plot reflectivity contours with colorbar if title_name == 'Updraft and Reflectivity': dbz_levels = np.arange(35., 75., 5.) dbz_rgb = np.array([[253,248,2], [229,188,0], [253,149,0], [253,0,0], [212,0,0], [188,0,0],[248,0,253], [152,84,198]], np.float32) / 255.0 dbz_cmap, dbz_norm = from_levels_and_colors(dbz_levels, dbz_rgb, extend='max') contours = plt.contour(to_np(lons), to_np(lats), to_np(reflectivity), levels=dbz_levels, transform=ccrs.PlateCarree(), cmap=dbz_cmap, norm=dbz_norm, linewidths=1) cbar_refl = mpu.colorbar(contours, ax, orientation='horizontal', aspect=10, shrink=.5, pad=0.05) cbar_refl.set_label('Maximum Derived Radar Reflectivity' '[$dBZ$]', fontsize=12.5) colorbar_lines = cbar_refl.ax.get_children() colorbar_lines[0].set_linewidths([10]*5) # Add wind quivers for every 10th data point if variable_name == 'wspd_wdir': plt.quiver(to_np(lons[::10,::10]), to_np(lats[::10,::10]), to_np(u_pressure[::10, ::10]), to_np(v_pressure[::10, ::10]), transform=ccrs.PlateCarree()) # Plot colorbar if variable_name == 'slp': pass else: cbar = mpu.colorbar(variable_plot, ax, orientation='vertical', aspect=40, shrink=.05, pad=0.05) cbar.set_label(colorbar_label, fontsize=15) cbar.set_ticks(levels) # Add borders and coastlines if variable_name == 'slp': ax.add_feature(cfeature.BORDERS.with_scale('10m'), edgecolor='white', linewidth=2) ax.add_feature(cfeature.COASTLINE.with_scale('10m'), edgecolor='white', linewidth=2) else: ax.add_feature(cfeature.BORDERS.with_scale('10m'), linewidth=0.8) ax.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8) ### Add initiation location ### if initiation == True: ax.plot(initiation_location.lon, initiation_location.lat, initiation_color, markersize=20, transform=ccrs.PlateCarree()) # Add gridlines lon = np.arange(0, 20, 1) lat = np.arange(40, 60, 1) gl = ax.gridlines(xlocs=lon, ylocs=lat, zorder=3) # Add tick labels mpu.yticklabels(lat, ax=ax, fontsize=12.5) mpu.xticklabels(lon, ax=ax, fontsize=12.5) # Make nicetime file_name = '{}wrfout_d02_{}_{}:{}:00'.format(data_dir, date, time, minutes) xr_file = xr.open_dataset(file_name) nicetime = pd.to_datetime(xr_file.QVAPOR.isel(Time=0).XTIME.values) nicetime = nicetime.strftime('%Y-%m-%d %H:%M') # Add plot title if pressure_level != False: ax.set_title('{} @ {} hPa'.format(title_name, pressure_level), loc='left', fontsize=15) ax.set_title('Valid time: {} UTC'.format(nicetime), loc='right', fontsize=15) else: if variable_name == 'slp': ax.set_title(title_name, loc='left', fontsize=15, color='white') ax.set_title('Valid time: {} UTC'.format(nicetime), loc='right', fontsize=15, color='white') else: ax.set_title(title_name, loc='left', fontsize=20) ax.set_title('Valid time: {} UTC'.format(nicetime), loc='right', fontsize=15) plt.show() ### Save figure ### if save == True: if pressure_level != False: if subset == True: fig.savefig('{}/{}/horizontal_map_{}_subset_{}_{}_{}:{}.png'.format( save_dir, save_name, save_name, pressure_level, date, time, minutes), bbox_inches='tight', dpi=300) else: fig.savefig('{}/{}/horizontal_map_{}_{}_{}_{}:{}.png'.format( save_dir, save_name, save_name, pressure_level, date, time, minutes), bbox_inches='tight', dpi=300) else: if subset == True: fig.savefig('{}/{}/horizontal_map_{}_subset_{}_{}:{}.png'.format( save_dir, save_name, save_name, date, time, minutes), bbox_inches='tight', dpi=300, facecolor=fig.get_facecolor()) else: fig.savefig('{}/{}/horizontal_map_{}_{}_{}:{}.png'.format( save_dir, save_name, save_name, date, time, minutes), bbox_inches='tight', dpi=300, facecolor=fig.get_facecolor()) ### Make a GIF from the plots ### if gif == True: # Predifine some variables gif_data_dir = save_dir + save_name gif_save_dir = '{}gifs/'.format(save_dir) gif_save_name = 'horizontal_map_{}.gif'.format(save_name) # GIF creating procedure os.chdir(gif_data_dir) image_folder = os.fsencode(gif_data_dir) filenames = [] for file in os.listdir(image_folder): filename = os.fsdecode(file) if filename.endswith( ('.png') ): filenames.append(filename) filenames.sort() images = list(map(lambda filename: imageio.imread(filename), filenames)) imageio.mimsave(os.path.join(gif_save_dir + gif_save_name), images, duration = 0.50)
del v1 # create 2d lat and lon lat2d = np.zeros((len(lat), len(lon))) lon2d = np.zeros((len(lat), len(lon))) for i in range(0, len(lon)): lat2d[:, i] = lat for i in range(0, len(lat)): lon2d[i, :] = lon dx, dy = mpcalc.lat_lon_grid_deltas(lon2d, lat2d) div1 = mpcalc.divergence(u, v, dx, dy) div = np.array(div1) # open workspace for analysis plot wks_type = "png" wks = ngl.open_wks( wks_type, "GFSanalysis_%s_%s_divergence_%shPa" % (region, init_dt[0:10], lev_hPa)) # define resources for analysis plot res = ngl.Resources() res.nglDraw = False res.nglFrame = False
def Crosssection_Wind_Theta_e_div( initTime=None, fhour=24,lw_ratio=[16,9], levels=[1000, 950, 925, 900, 850, 800, 700,600,500,400,300,200], day_back=0,model='GRAPES_GFS',data_source='MICAPS', output_dir=None, st_point = [20, 120.0], ed_point = [50, 130.0], map_extent=[70,140,15,55], h_pos=[0.125, 0.665, 0.25, 0.2],**kwargs): # micaps data directory if(data_source == 'MICAPS'): try: data_dir = [utl.Cassandra_dir(data_type='high',data_source=model,var_name='RH',lvl=''), utl.Cassandra_dir(data_type='high',data_source=model,var_name='UGRD',lvl=''), utl.Cassandra_dir(data_type='high',data_source=model,var_name='VGRD',lvl=''), utl.Cassandra_dir(data_type='high',data_source=model,var_name='TMP',lvl=''), utl.Cassandra_dir(data_type='high',data_source=model,var_name='HGT',lvl='500'), utl.Cassandra_dir(data_type='surface',data_source=model,var_name='PSFC')] except KeyError: raise ValueError('Can not find all directories needed') # get filename if(initTime != None): filename = utl.model_filename(initTime, fhour) else: filename=utl.filename_day_back_model(day_back=day_back,fhour=fhour) # retrieve data from micaps server rh=MICAPS_IO.get_model_3D_grid(directory=data_dir[0][0:-1],filename=filename,levels=levels, allExists=False) rh = rh.metpy.parse_cf().squeeze() u=MICAPS_IO.get_model_3D_grid(directory=data_dir[1][0:-1],filename=filename,levels=levels, allExists=False) u = u.metpy.parse_cf().squeeze() v=MICAPS_IO.get_model_3D_grid(directory=data_dir[2][0:-1],filename=filename,levels=levels, allExists=False) v = v.metpy.parse_cf().squeeze() v2=MICAPS_IO.get_model_3D_grid(directory=data_dir[2][0:-1],filename=filename,levels=levels, allExists=False) v2 = v2.metpy.parse_cf().squeeze() t=MICAPS_IO.get_model_3D_grid(directory=data_dir[3][0:-1],filename=filename,levels=levels, allExists=False) t = t.metpy.parse_cf().squeeze() gh=MICAPS_IO.get_model_grid(data_dir[4], filename=filename) psfc=get_model_grid(data_dir[5], filename=filename) if(data_source == 'CIMISS'): # get filename if(initTime != None): filename = utl.model_filename(initTime, fhour,UTC=True) else: filename=utl.filename_day_back_model(day_back=day_back,fhour=fhour,UTC=True) try: rh=CMISS_IO.cimiss_model_3D_grid(init_time_str='20'+filename[0:8],valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='RHU'), fcst_levels=levels, fcst_ele="RHU", units='%') if rh is None: return u=CMISS_IO.cimiss_model_3D_grid(init_time_str='20'+filename[0:8],valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='WIU'), fcst_levels=levels, fcst_ele="WIU", units='m/s') if u is None: return v=CMISS_IO.cimiss_model_3D_grid(init_time_str='20'+filename[0:8],valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='WIV'), fcst_levels=levels, fcst_ele="WIV", units='m/s') if v is None: return v2=CMISS_IO.cimiss_model_3D_grid(init_time_str='20'+filename[0:8],valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='WIV'), fcst_levels=levels, fcst_ele="WIV", units='m/s') if v2 is None: return t=CMISS_IO.cimiss_model_3D_grid(init_time_str='20'+filename[0:8],valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='TEM'), fcst_levels=levels, fcst_ele="TEM", units='K') if t is None: return t['data'].values=t['data'].values-273.15 gh=CMISS_IO.cimiss_model_by_time('20'+filename[0:8],valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='GPH'), fcst_level=500, fcst_ele="GPH", units='gpm') if gh is None: return gh['data'].values=gh['data'].values/10. psfc=CMISS_IO.cimiss_model_by_time('20'+filename[0:8], valid_time=fhour, data_code=utl.CMISS_data_code(data_source=model,var_name='PRS'), fcst_level=0, fcst_ele="PRS", units='Pa') psfc['data']=psfc['data']/100. except KeyError: raise ValueError('Can not find all data needed') rh = rh.metpy.parse_cf().squeeze() u = u.metpy.parse_cf().squeeze() v = v.metpy.parse_cf().squeeze() v2 = v2.metpy.parse_cf().squeeze() t = t.metpy.parse_cf().squeeze() psfc=psfc.metpy.parse_cf().squeeze() resolution=u['lon'][1]-u['lon'][0] x,y=np.meshgrid(u['lon'], u['lat']) # +form 3D psfc mask1 = ( (psfc['lon']>=t['lon'].values.min())& (psfc['lon']<=t['lon'].values.max())& (psfc['lat']>=t['lat'].values.min())& (psfc['lat']<=t['lat'].values.max()) ) t2,psfc_bdcst=xr.broadcast(t['data'],psfc['data'].where(mask1, drop=True)) mask2=(psfc_bdcst > -10000) psfc_bdcst=psfc_bdcst.where(mask2, drop=True) # -form 3D psfc dx,dy=mpcalc.lat_lon_grid_deltas(u['lon'],u['lat']) for ilvl in levels: u2d=u.sel(level=ilvl) v2d=v.sel(level=ilvl) # absv2d=mpcalc.absolute_vorticity(u2d['data'].values*units.meter/units.second, # v2d['data'].values*units.meter/units.second,dx,dy,y*units.degree) div2d=mpcalc.divergence(u2d['data'].values.squeeze()*units('m/s'), v2d['data'].values.squeeze()*units('m/s'), dx, dy, dim_order='yx') if(ilvl == levels[0]): div3d = v2.copy() div3d['data'].loc[dict(level=ilvl)]=np.array(div2d) else: div3d['data'].loc[dict(level=ilvl)]=np.array(div2d) div3d['data'].attrs['units']=div2d.units #rh=rh.rename(dict(lat='latitude',lon='longitude')) cross = cross_section(rh, st_point, ed_point) cross_rh=cross.set_coords(('lat', 'lon')) cross = cross_section(u, st_point, ed_point) cross_u=cross.set_coords(('lat', 'lon')) cross = cross_section(v, st_point, ed_point) cross_v=cross.set_coords(('lat', 'lon')) cross_psfc = cross_section(psfc_bdcst, st_point, ed_point) cross_u['data'].attrs['units']=units.meter/units.second cross_v['data'].attrs['units']=units.meter/units.second cross_u['t_wind'], cross_v['n_wind'] = mpcalc.cross_section_components(cross_u['data'],cross_v['data']) cross = cross_section(t, st_point, ed_point) cross_t=cross.set_coords(('lat', 'lon')) cross = cross_section(div3d, st_point, ed_point) cross_div3d=cross.set_coords(('lat', 'lon')) cross_Td = mpcalc.dewpoint_rh(cross_t['data'].values*units.celsius, cross_rh['data'].values* units.percent) rh,pressure = xr.broadcast(cross_rh['data'],cross_t['level']) pressure.attrs['units']='hPa' Theta_e=mpcalc.equivalent_potential_temperature(pressure, cross_t['data'].values*units.celsius, cross_Td) cross_terrain=pressure-cross_psfc cross_Theta_e = xr.DataArray(np.array(Theta_e), coords=cross_rh['data'].coords, dims=cross_rh['data'].dims, attrs={'units': Theta_e.units}) crossection_graphics.draw_Crosssection_Wind_Theta_e_div( cross_div3d=cross_div3d, cross_Theta_e=cross_Theta_e, cross_u=cross_u, cross_v=cross_v,cross_terrain=cross_terrain,gh=gh, h_pos=h_pos,st_point=st_point,ed_point=ed_point, levels=levels,map_extent=map_extent,lw_ratio=lw_ratio, output_dir=output_dir)
## first, smooth everything pretty heavily: tmpk_700 = ndimage.gaussian_filter(tmpk_700, sigma=6, order=0) * units('kelvin') uwnd_700 = ndimage.gaussian_filter(uwnd_700, sigma=6, order=0) * units('m/s') vwnd_700 = ndimage.gaussian_filter(vwnd_700, sigma=6, order=0) * units('m/s') # less smoothing for the height hght_700 = ndimage.gaussian_filter(hght_700, sigma=2, order=0) * units('m') if plot_700qvect: # Compute the Q-vector components uqvect, vqvect = mpcalc.q_vector(uwnd_700, vwnd_700, tmpk_700, 700*units.hPa, dx, dy) # Compute the divergence of the Q-vectors calculated above q_div = -2*mpcalc.divergence(uqvect, vqvect, dx, dy, dim_order='yx') ## smooth it too q_div = ndimage.gaussian_filter(q_div, sigma=4, order=0) * units('meter/(kilogram second)') #Map Creation print('plotting 700 q-vectors') # Set Projection of Data datacrs = ccrs.PlateCarree() # Set Projection of Plot plotcrs = ccrs.LambertConformal(central_latitude=[30, 60], central_longitude=-100) fig = plt.figure(1, figsize=(20,16))
hgt = xr.open_dataset(PATH_DATA + FILE_HGT_S4) hgt = hgt - hgt.mean(dim='longitude') winds = xr.open_dataset(PATH_DATA + FILE_WINDS, chunks={'latitude':10}) winds = winds.transpose('month', 'realiz', 'latitude', 'longitude') winds_clm = winds.mean(dim='realiz') for i in np.arange(0, 7): var_normal = np.mean(hgt.z.values[i, :, :, :], axis=0) px, py, lat = plumb_flux.ComputePlumbFluxes(winds_clm.u.values[i, :, :], winds_clm.v.values[i, :, :], hgt.z.values[i, :, :, :], var_normal, hgt.latitude.values, hgt.longitude.values, 20000) dx, dy = calc.lat_lon_grid_deltas(hgt.longitude.values, lat) #div = calc.divergence(px, py, dx, dy) px_normal, py_normal = np.mean(px, axis=0), np.mean(py, axis=0) var_ninio_WPV = np.mean(hgt.z.values[i, index_ninio_WPV, :, :], axis=0) px_ninio_WPV, py_ninio_WPV = np.mean(px[index_ninio_WPV, :, :], axis=0) - px_normal, np.mean(py[index_ninio_WPV, :, :], axis=0) - py_normal div_ninio_WPV = calc.divergence(px_ninio_WPV, py_ninio_WPV, dx, dy) var_ninia_WPV = np.mean(hgt.z.values[i, index_ninia_WPV, :, :], axis=0) px_ninia_WPV, py_ninia_WPV = np.mean(px[index_ninia_WPV, :, :], axis=0) - px_normal, np.mean(py[index_ninia_WPV, :, :], axis=0) - py_normal div_ninia_WPV = calc.divergence(px_ninia_WPV, py_ninia_WPV, dx, dy) var_ninio_SPV = np.mean(hgt.z.values[i, index_ninio_SPV, :, :], axis=0) px_ninio_SPV, py_ninio_SPV = np.mean(px[index_ninio_SPV, :, :], axis=0) - px_normal, np.mean(py[index_ninio_SPV, :, :], axis=0) - py_normal div_ninio_SPV = calc.divergence(px_ninio_SPV, py_ninio_SPV, dx, dy) var_ninia_SPV = np.mean(hgt.z.values[i, index_ninia_SPV, :, :], axis=0) px_ninia_SPV, py_ninia_SPV = np.mean(px[index_ninia_SPV, :, :], axis=0) - px_normal, np.mean(py[index_ninia_SPV, :, :], axis=0) - py_normal div_ninia_SPV = calc.divergence(px_ninia_SPV, py_ninia_SPV, dx, dy)
def PV_Div_uv(initTime=None, fhour=6, day_back=0, model='ECMWF', map_ratio=14 / 9, zoom_ratio=20, cntr_pnt=[104, 34], levels=[ 1000, 950, 925, 900, 850, 800, 700, 600, 500, 400, 300, 250, 200, 100 ], lvl_ana=250, Global=False, south_China_sea=True, area=None, city=False, output_dir=None, data_source='MICAPS', **kwargs): # micaps data directory if (area != None): south_China_sea = False # micaps data directory if (data_source == 'MICAPS'): try: data_dir = [ utl.Cassandra_dir(data_type='high', data_source=model, var_name='RH', lvl=''), utl.Cassandra_dir(data_type='high', data_source=model, var_name='UGRD', lvl=''), utl.Cassandra_dir(data_type='high', data_source=model, var_name='VGRD', lvl=''), utl.Cassandra_dir(data_type='high', data_source=model, var_name='TMP', lvl=''), utl.Cassandra_dir(data_type='high', data_source=model, var_name='HGT', lvl='') ] except KeyError: raise ValueError('Can not find all directories needed') # get filename if (initTime != None): filename = utl.model_filename(initTime, fhour) else: filename = utl.filename_day_back_model(day_back=day_back, fhour=fhour) # retrieve data from micaps server rh = MICAPS_IO.get_model_3D_grid(directory=data_dir[0][0:-1], filename=filename, levels=levels, allExists=False) if rh is None: return u = MICAPS_IO.get_model_3D_grid(directory=data_dir[1][0:-1], filename=filename, levels=levels, allExists=False) if u is None: return v = MICAPS_IO.get_model_3D_grid(directory=data_dir[2][0:-1], filename=filename, levels=levels, allExists=False) if v is None: return t = MICAPS_IO.get_model_3D_grid(directory=data_dir[3][0:-1], filename=filename, levels=levels, allExists=False) if t is None: return if (data_source == 'CIMISS'): # get filename if (initTime != None): filename = utl.model_filename(initTime, fhour, UTC=True) else: filename = utl.filename_day_back_model(day_back=day_back, fhour=fhour, UTC=True) try: # retrieve data from CIMISS server rh = CMISS_IO.cimiss_model_3D_grid( data_code=utl.CMISS_data_code(data_source=model, var_name='RHU'), init_time_str='20' + filename[0:8], valid_time=fhour, fcst_levels=levels, fcst_ele="RHU", units='%') if rh is None: return u = CMISS_IO.cimiss_model_3D_grid( data_code=utl.CMISS_data_code(data_source=model, var_name='WIU'), init_time_str='20' + filename[0:8], valid_time=fhour, fcst_levels=levels, fcst_ele="WIU", units='m/s') if u is None: return v = CMISS_IO.cimiss_model_3D_grid( data_code=utl.CMISS_data_code(data_source=model, var_name='WIV'), init_time_str='20' + filename[0:8], valid_time=fhour, fcst_levels=levels, fcst_ele="WIV", units='m/s') if v is None: return t = CMISS_IO.cimiss_model_3D_grid( data_code=utl.CMISS_data_code(data_source=model, var_name='TEM'), init_time_str='20' + filename[0:8], valid_time=fhour, fcst_levels=levels, fcst_ele="TEM", units='K') if t is None: return t['data'].values = t['data'].values - 273.15 t['data'].attrs['units'] = 'C' except KeyError: raise ValueError('Can not find all data needed') if (area != None): cntr_pnt, zoom_ratio = utl.get_map_area(area_name=area) map_extent = [0, 0, 0, 0] map_extent[0] = cntr_pnt[0] - zoom_ratio * 1 * map_ratio map_extent[1] = cntr_pnt[0] + zoom_ratio * 1 * map_ratio map_extent[2] = cntr_pnt[1] - zoom_ratio * 1 map_extent[3] = cntr_pnt[1] + zoom_ratio * 1 delt_x = (map_extent[1] - map_extent[0]) * 0.2 delt_y = (map_extent[3] - map_extent[2]) * 0.1 #+ to solve the problem of labels on all the contours mask1 = (rh['lon'] > map_extent[0] - delt_x) & (rh['lon'] < map_extent[1] + delt_x) & ( rh['lat'] > map_extent[2] - delt_y) & (rh['lat'] < map_extent[3] + delt_y) mask2 = (u['lon'] > map_extent[0] - delt_x) & (u['lon'] < map_extent[1] + delt_x) & ( u['lat'] > map_extent[2] - delt_y) & (u['lat'] < map_extent[3] + delt_y) mask3 = (t['lon'] > map_extent[0] - delt_x) & (t['lon'] < map_extent[1] + delt_x) & ( t['lat'] > map_extent[2] - delt_y) & (t['lat'] < map_extent[3] + delt_y) #- to solve the problem of labels on all the contours rh = rh.where(mask1, drop=True) u = u.where(mask2, drop=True) v = v.where(mask2, drop=True) t = t.where(mask3, drop=True) uv = xr.merge([u.rename({'data': 'u'}), v.rename({'data': 'v'})]) lats = np.squeeze(rh['lat'].values) lons = np.squeeze(rh['lon'].values) pres = np.array(levels) * 100 * units('Pa') tmpk = mpcalc.smooth_n_point( (t['data'].values.squeeze() + 273.15), 9, 2) * units('kelvin') thta = mpcalc.potential_temperature(pres[:, None, None], tmpk) uwnd = mpcalc.smooth_n_point(u['data'].values.squeeze(), 9, 2) * units.meter / units.second vwnd = mpcalc.smooth_n_point(v['data'].values.squeeze(), 9, 2) * units.meter / units.second dx, dy = mpcalc.lat_lon_grid_deltas(lons, lats) # Comput the PV on all isobaric surfaces pv_raw = mpcalc.potential_vorticity_baroclinic( thta, pres[:, None, None], uwnd, vwnd, dx[None, :, :], dy[None, :, :], lats[None, :, None] * units('degrees')) div_raw = mpcalc.divergence(uwnd, vwnd, dx[None, :, :], dy[None, :, :], dim_order='yx') # prepare data idx_z1 = list(pres.m).index(((lvl_ana * units('hPa')).to(pres.units)).m) pv = rh.copy(deep=True) pv['data'].values = np.array(pv_raw).reshape( np.append(1, np.array(pv_raw).shape)) pv['data'].attrs['units'] = str(pv_raw.units) pv.attrs['model'] = model pv = pv.where(pv['level'] == lvl_ana, drop=True) div = u.copy(deep=True) div['data'].values = np.array(div_raw).reshape( np.append(1, np.array(div_raw).shape)) div['data'].attrs['units'] = str(div_raw.units) div = div.where(div['level'] == lvl_ana, drop=True) uv = uv.where(uv['level'] == lvl_ana, drop=True) synoptic_graphics.draw_PV_Div_uv(pv=pv, uv=uv, div=div, map_extent=map_extent, regrid_shape=20, city=city, south_China_sea=south_China_sea, output_dir=output_dir, Global=Global)
# Use MetPy to compute the baroclinic potential vorticity on all isobaric # levels and other variables # # Compute dx and dy spacing for use in vorticity calculation dx, dy = mpcalc.lat_lon_grid_deltas(lons, lats) # Comput the PV on all isobaric surfaces pv = mpcalc.potential_vorticity_baroclinic( thta, pres[:, None, None], uwnd, vwnd, dx[None, :, :], dy[None, :, :], lats[None, :, None] * units('degrees')) # Use MetPy to compute the divergence on the pressure surfaces div = mpcalc.divergence(uwnd, vwnd, dx[None, :, :], dy[None, :, :], dim_order='yx') # Find the index value for the 250-hPa surface i250 = list(pres.m).index(((250 * units('hPa')).to(pres.units)).m) ###################################################################### # Map Creation # ------------ # # This next set of code creates the plot and draws contours on a Lambert # Conformal map centered on -100 E longitude. The main view is over the # CONUS with isobaric PV map with PV contoured every 1 PVU and divergence # colorshaded. #
def PV_Div_uv(initial_time=None, fhour=6, day_back=0,model='ECMWF', map_ratio=19/9,zoom_ratio=20,cntr_pnt=[102,34], levels=[1000, 950, 925, 900, 850, 800, 700,600,500,400,300,250,200,100],lvl_ana=250, Global=False, south_China_sea=True,area = '全国',city=False,output_dir=None ): # micaps data directory try: data_dir = [utl.Cassandra_dir(data_type='high',data_source=model,var_name='RH',lvl=''), utl.Cassandra_dir(data_type='high',data_source=model,var_name='UGRD',lvl=''), utl.Cassandra_dir(data_type='high',data_source=model,var_name='VGRD',lvl=''), utl.Cassandra_dir(data_type='high',data_source=model,var_name='TMP',lvl=''), utl.Cassandra_dir(data_type='high',data_source=model,var_name='HGT',lvl='')] except KeyError: raise ValueError('Can not find all directories needed') # get filename if(initial_time != None): filename = utl.model_filename(initial_time, fhour) else: filename=utl.filename_day_back_model(day_back=day_back,fhour=fhour) # retrieve data from micaps server rh=get_model_3D_grid(directory=data_dir[0][0:-1],filename=filename,levels=levels, allExists=False) if rh is None: return u=get_model_3D_grid(directory=data_dir[1][0:-1],filename=filename,levels=levels, allExists=False) if u is None: return v=get_model_3D_grid(directory=data_dir[2][0:-1],filename=filename,levels=levels, allExists=False) if v is None: return t=get_model_3D_grid(directory=data_dir[3][0:-1],filename=filename,levels=levels, allExists=False) if t is None: return # get filename if(initial_time != None): filename = utl.model_filename(initial_time, fhour) else: filename=utl.filename_day_back_model(day_back=day_back,fhour=fhour) lats = np.squeeze(rh['lat'].values) lons = np.squeeze(rh['lon'].values) pres = np.array(levels)*100 * units('Pa') tmpk = mpcalc.smooth_n_point(t['data'].values.squeeze(), 9, 2)*units('degC') thta = mpcalc.potential_temperature(pres[:, None, None], tmpk) uwnd = mpcalc.smooth_n_point(u['data'].values.squeeze(), 9, 2)*units.meter/units.second vwnd = mpcalc.smooth_n_point(v['data'].values.squeeze(), 9, 2)*units.meter/units.second dx, dy = mpcalc.lat_lon_grid_deltas(lons, lats) # Comput the PV on all isobaric surfaces pv = mpcalc.potential_vorticity_baroclinic(thta, pres[:, None, None], uwnd, vwnd, dx[None, :, :], dy[None, :, :], lats[None, :, None] * units('degrees')) div = mpcalc.divergence(uwnd, vwnd, dx[None, :, :], dy[None, :, :], dim_order='yx') # prepare data idx_z1 = list(pres.m).index(((lvl_ana * units('hPa')).to(pres.units)).m) if(area != None): cntr_pnt,zoom_ratio=utl.get_map_area(area_name=area) map_extent=[0,0,0,0] map_extent[0]=cntr_pnt[0]-zoom_ratio*1*map_ratio map_extent[1]=cntr_pnt[0]+zoom_ratio*1*map_ratio map_extent[2]=cntr_pnt[1]-zoom_ratio*1 map_extent[3]=cntr_pnt[1]+zoom_ratio*1 delt_x=(map_extent[1]-map_extent[0])*0.2 delt_y=(map_extent[3]-map_extent[2])*0.1 #+ to solve the problem of labels on all the contours idx_x1 = np.where((lons > map_extent[0]-delt_x) & (lons < map_extent[1]+delt_x)) idx_y1 = np.where((lats > map_extent[2]-delt_y) & (lats < map_extent[3]+delt_y)) #- to solve the problem of labels on all the contours init_time = u.coords['forecast_reference_time'].values pv = { 'lon': lons[idx_x1], 'lat': lats[idx_y1], 'data': np.array(pv)[idx_z1,idx_y1[0][0]:(idx_y1[0][-1]+1),idx_x1[0][0]:(idx_x1[0][-1]+1)], 'lev':str(lvl_ana), 'model':model, 'fhour':fhour, 'init_time':init_time} uv = { 'lon': lons[idx_x1], 'lat': lats[idx_y1], 'udata': np.array(uwnd)[idx_z1,idx_y1[0][0]:(idx_y1[0][-1]+1),idx_x1[0][0]:(idx_x1[0][-1]+1)], 'vdata': np.array(vwnd)[idx_z1,idx_y1[0][0]:(idx_y1[0][-1]+1),idx_x1[0][0]:(idx_x1[0][-1]+1)], 'lev':str(lvl_ana)} div = { 'lon': lons[idx_x1], 'lat': lats[idx_y1], 'data': np.array(div)[idx_z1,idx_y1[0][0]:(idx_y1[0][-1]+1),idx_x1[0][0]:(idx_x1[0][-1]+1)], 'lev':str(lvl_ana)} synoptic_graphics.draw_PV_Div_uv( pv=pv, uv=uv, div=div, map_extent=map_extent, regrid_shape=20, city=city,south_China_sea=south_China_sea, output_dir=output_dir,Global=Global)
def metpy_temp_adv(fname, plevels, tidx): ds = xr.open_dataset(fname).metpy.parse_cf().squeeze() ds = ds.isel(Time=tidx) print(ds) dx, dy = mpcalc.lat_lon_grid_deltas(ds.lon.values * units('degrees_E'), ds.lat.values * units('degrees_N')) plevels_unit = plevels * units.hPa ds1 = xr.Dataset() p = units.Quantity(to_np(ds.p), 'hPa') z = units.Quantity(to_np(ds.z), 'meter') u = units.Quantity(to_np(ds.u), 'm/s') v = units.Quantity(to_np(ds.v), 'm/s') w = units.Quantity(to_np(ds.w), 'm/s') tk = units.Quantity(to_np(ds.tk), 'kelvin') th = units.Quantity(to_np(ds.th), 'kelvin') eth = units.Quantity(to_np(ds.eth), 'kelvin') wspd = units.Quantity(to_np(ds.wspd), 'm/s') omega = units.Quantity(to_np(ds.omega), 'Pa/s') z, u, v, w, tk, th, eth, wspd, omega = log_interpolate_1d(plevels_unit, p, z, u, v, w, tk, th, eth, wspd, omega, axis=0) coords, dims = [plevs, ds.lat.values, ds.lon.values], ["level", "lat", "lon"] for name, var in zip( ['z', 'u', 'v', 'w', 'tk', 'th', 'eth', 'wspd', 'omega'], [z, u, v, w, tk, th, eth, wspd, omega]): #g = ndimage.gaussian_filter(var, sigma=3, order=0) ds1[name] = xr.DataArray(to_np(var), coords=coords, dims=dims) # Calculate temperature advection using metpy function for i, plev in enumerate(plevs): uqvect, vqvect = mpcalc.q_vector(u[i, :, :], v[i, :, :], th[i, :, :], plev * units.hPa, dx, dy) #uqvect, vqvect = mpcalc.q_vector(u[i,:,:], v[i,:,:], th.to('degC')[i,:,:], plev*units.hPa, dx, dy) q_div = -2 * mpcalc.divergence(uqvect, vqvect, dx, dy, dim_order='yx') ds1['uq_{:03d}'.format(plev)] = xr.DataArray( np.array(uqvect), coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['vq_{:03d}'.format(plev)] = xr.DataArray( np.array(vqvect), coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['q_div_{:03d}'.format(plev)] = xr.DataArray( np.array(q_div), coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) # adv = mpcalc.advection(tk[i,:,:], [u[i,:,:], v[i,:,:]], (dx, dy), dim_order='yx') * units('K/sec') # adv = ndimage.gaussian_filter(adv, sigma=3, order=0) * units('K/sec') # ds1['tk_adv_{:03d}'.format(plev)] = xr.DataArray(np.array(adv), coords=[ds.lat.values,ds.lon.values], dims=["lat","lon"]) # # adv = mpcalc.advection(th[i,:,:], [u[i,:,:], v[i,:,:]], (dx, dy), dim_order='yx') * units('K/sec') # adv = ndimage.gaussian_filter(adv, sigma=3, order=0) * units('K/sec') # ds1['th_adv_{:03d}'.format(plev)] = xr.DataArray(np.array(adv), coords=[ds.lat.values,ds.lon.values], dims=["lat","lon"]) # adv = mpcalc.advection(eth[i, :, :], [u[i, :, :], v[i, :, :]], (dx, dy), dim_order='yx') * units('K/sec') adv = ndimage.gaussian_filter(adv, sigma=3, order=0) * units('K/sec') ds1['eth_adv_{:03d}'.format(plev)] = xr.DataArray( np.array(adv), coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) div = mpcalc.divergence(u[i, :, :], v[i, :, :], dx, dy, dim_order='yx') div = ndimage.gaussian_filter(div, sigma=3, order=0) * units('1/sec') ds1['div_{:03d}'.format(plev)] = xr.DataArray( np.array(div), coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['Time'] = ds.Time eth2 = mpcalc.equivalent_potential_temperature( ds.slp.values * units.hPa, ds.t2m.values * units('K'), ds.td2.values * units('celsius')) ds1['eth2'] = xr.DataArray(eth2, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['sst'] = xr.DataArray( ndimage.gaussian_filter(ds.sst.values, sigma=3, order=0) - 273.15, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['t2m'] = xr.DataArray(ds.t2m.values, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['th2'] = xr.DataArray(ds.th2.values, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['mask'] = xr.DataArray(ds.mask.values, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['u10'] = xr.DataArray(ds.u10.values, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['v10'] = xr.DataArray(ds.v10.values, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['slp'] = xr.DataArray(ds.slp.values, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['rain'] = xr.DataArray(ds.rain.values, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) ds1['accrain'] = xr.DataArray(ds.accrain.values, coords=[ds.lat.values, ds.lon.values], dims=["lat", "lon"]) #ds1['latent'] = xr.DataArray(ds.latent.values, coords=[ds.lat.values,ds.lon.values], dims=["lat","lon"]) #ds1['acclatent'] = xr.DataArray(ds.acclatent.values, coords=[ds.lat.values,ds.lon.values], dims=["lat","lon"]) return ds1
def main(): load_start = dt.datetime.now() #Try parsing arguments using argparse parser = argparse.ArgumentParser( description='wrf non-parallel convective diagnostics processer') parser.add_argument("-m", help="Model name", required=True) parser.add_argument("-r", help="Region name (default is aus)", default="aus") parser.add_argument("-t1", help="Time start YYYYMMDDHH", required=True) parser.add_argument("-t2", help="Time end YYYYMMDDHH", required=True) parser.add_argument( "-e", help= "CMIP5 experiment name (not required if using era5, erai or barra)", default="") parser.add_argument( "--barpa_forcing_mdl", help="BARPA forcing model (erai or ACCESS1-0). Default erai.", default="erai") parser.add_argument( "--ens", help="CMIP5 ensemble name (not required if using era5, erai or barra)", default="r1i1p1") parser.add_argument("--group", help="CMIP6 modelling group name", default="") parser.add_argument("--project", help="CMIP6 modelling intercomparison project", default="CMIP") parser.add_argument("--ver6hr", help="Version on al33 for 6hr data", default="") parser.add_argument("--ver3hr", help="Version on al33 for 3hr data", default="") parser.add_argument("--issave", help="Save output (True or False, default is False)", default="False") parser.add_argument( "--ub4", help= "Where to get era5 data. Default True for ub4 project, otherwise rt52", default="True") parser.add_argument( "--outname", help= "Name of saved output. In the form *outname*_*t1*_*t2*.nc. Default behaviour is the model name", default=None) parser.add_argument( "--is_dcape", help="Should DCAPE be calculated? (1 or 0. Default is 1)", default=1) parser.add_argument( "--al33", help= "Should data be gathered from al33? Default is False, and data is gathered from r87. If True, then group is required", default="False") parser.add_argument( "--delta_t", help= "Time step spacing for ERA5 data, in hours. Default is one the minimum spacing (1 hour)", default="1") parser.add_argument( "--era5_interp", help= "Horizontally interpolate model data before calculating convective parameters", default="False") args = parser.parse_args() #Parse arguments from cmd line and set up inputs (date region model) model = args.m region = args.r t1 = args.t1 t2 = args.t2 issave = args.issave ub4 = args.ub4 al33 = args.al33 if args.outname == None: out_name = model else: out_name = args.outname is_dcape = args.is_dcape barpa_forcing_mdl = args.barpa_forcing_mdl experiment = args.e ensemble = args.ens group = args.group project = args.project ver6hr = args.ver6hr ver3hr = args.ver3hr delta_t = int(args.delta_t) era5_interp = args.era5_interp if region == "sa_small": start_lat = -38 end_lat = -26 start_lon = 132 end_lon = 142 elif region == "aus": start_lat = -44.525 end_lat = -9.975 start_lon = 111.975 end_lon = 156.275 elif region == "global": start_lat = -70 end_lat = 70 start_lon = -180 end_lon = 179.75 else: raise ValueError("INVALID REGION\n") domain = [start_lat, end_lat, start_lon, end_lon] try: time = [ dt.datetime.strptime(t1, "%Y%m%d%H"), dt.datetime.strptime(t2, "%Y%m%d%H") ] except: raise ValueError("INVALID START OR END TIME. SHOULD BE YYYYMMDDHH\n") if era5_interp == "True": era5_interp = True elif era5_interp == "False": era5_interp = False else: raise ValueError("\n INVALID era5_interp...SHOULD BE True OR False") if ub4 == "True": ub4 = True elif ub4 == "False": ub4 = False else: raise ValueError("\n INVALID ub4...SHOULD BE True OR False") if issave == "True": issave = True elif issave == "False": issave = False else: raise ValueError("\n INVALID ISSAVE...SHOULD BE True OR False") if al33 == "True": al33 = True elif al33 == "False": al33 = False else: raise ValueError("\n INVALID al33...SHOULD BE True OR False") #Load data print("LOADING DATA...") if model == "erai": ta,temp1,hur,hgt,terrain,p,ps,wap,ua,va,uas,vas,tas,ta2d,\ cp,tp,wg10,mod_cape,lon,lat,date_list = \ read_erai(domain,time) cp = cp.astype("float32", order="C") tp = tp.astype("float32", order="C") mod_cape = mod_cape.astype("float32", order="C") elif model == "era5": if ub4: ta,temp1,hur,hgt,terrain,p,ps,ua,va,uas,vas,tas,ta2d,\ cp,wg10,mod_cape,lon,lat,date_list = \ read_era5(domain,time,delta_t=delta_t) else: ta,temp1,hur,hgt,terrain,p,ps,ua,va,uas,vas,tas,ta2d,\ cp,tp,wg10,mod_cape,lon,lat,date_list = \ read_era5_rt52(domain,time,delta_t=delta_t) cp = cp.astype("float32", order="C") tp = tp.astype("float32", order="C") mod_cape = mod_cape.astype("float32", order="C") wap = np.zeros(hgt.shape) elif model == "barra": ta,temp1,hur,hgt,terrain,p,ps,wap,ua,va,uas,vas,tas,ta2d,wg10,lon,lat,date_list = \ read_barra(domain,time) elif model == "barra_fc": ta,temp1,hur,hgt,terrain,p,ps,wap,ua,va,uas,vas,tas,ta2d,wg10,lon,lat,date_list = \ read_barra_fc(domain,time) elif model == "barpa": ta,hur,hgt,terrain,p,ps,ua,va,uas,vas,tas,ta2d,wg10,lon,lat,date_list = \ read_barpa(domain, time, experiment, barpa_forcing_mdl, ensemble) wap = np.zeros(hgt.shape) temp1 = None elif model == "barra_ad": wg10,temp2,ta,temp1,hur,hgt,terrain,p,ps,wap,ua,va,uas,vas,tas,ta2d,lon,lat,date_list = \ read_barra_ad(domain, time, False) elif model in ["ACCESS1-0","ACCESS1-3","GFDL-CM3","GFDL-ESM2M","CNRM-CM5","MIROC5",\ "MRI-CGCM3","IPSL-CM5A-LR","IPSL-CM5A-MR","GFDL-ESM2G","bcc-csm1-1","MIROC-ESM",\ "BNU-ESM"]: #Check that t1 and t2 are in the same year year = np.arange(int(t1[0:4]), int(t2[0:4]) + 1) ta, hur, hgt, terrain, p_3d, ps, ua, va, uas, vas, tas, ta2d, tp, lon, lat, \ date_list = read_cmip(model, experiment, \ ensemble, year, domain, cmip_ver=5, al33=al33, group=group, ver6hr=ver6hr, ver3hr=ver3hr) wap = np.zeros(hgt.shape) wg10 = np.zeros(ps.shape) mod_cape = np.zeros(ps.shape) p = np.zeros(p_3d[0, :, 0, 0].shape) #date_list = pd.to_datetime(date_list).to_pydatetime() temp1 = None tp = tp.astype("float32", order="C") elif model in ["ACCESS-ESM1-5", "ACCESS-CM2"]: year = np.arange(int(t1[0:4]), int(t2[0:4]) + 1) ta, hur, hgt, terrain, p_3d, ps, ua, va, uas, vas, tas, ta2d, lon, lat, \ date_list = read_cmip(model, experiment,\ ensemble, year, domain, cmip_ver=6, group=group, project=project) wap = np.zeros(hgt.shape) wg10 = np.zeros(ps.shape) p = np.zeros(p_3d[0, :, 0, 0].shape) #date_list = pd.to_datetime(date_list).to_pydatetime() temp1 = None else: raise ValueError("Model not recognised") del temp1 ta = ta.astype("float32", order="C") hur = hur.astype("float32", order="C") hgt = hgt.astype("float32", order="C") terrain = terrain.astype("float32", order="C") p = p.astype("float32", order="C") ps = ps.astype("float32", order="C") wap = wap.astype("float32", order="C") ua = ua.astype("float32", order="C") va = va.astype("float32", order="C") uas = uas.astype("float32", order="C") vas = vas.astype("float32", order="C") tas = tas.astype("float32", order="C") ta2d = ta2d.astype("float32", order="C") wg10 = wg10.astype("float32", order="C") lon = lon.astype("float32", order="C") lat = lat.astype("float32", order="C") gc.collect() param = np.array([ "mu_cape", "mu_cin", "muq", "s06", "s0500", "lr700_500", "mhgt", "ta500", "tp" ]) if model in ["erai", "era5"]: param = np.concatenate([param, ["mod_cape"]]) #Option to interpolate to the ERA5 grid if era5_interp: #Interpolate model data to the ERA5 grid from era5_read import get_lat_lon_rt52 as get_era5_lat_lon era5_lon, era5_lat = get_era5_lat_lon() era5_lon_ind = np.where((era5_lon >= domain[2]) & (era5_lon <= domain[3]))[0] era5_lat_ind = np.where((era5_lat >= domain[0]) & (era5_lat <= domain[1]))[0] era5_lon = era5_lon[era5_lon_ind] era5_lat = era5_lat[era5_lat_ind] terrain = interp_era5(terrain, lon, lat, era5_lon, era5_lat, d3=False) #Set output array output_data = np.zeros( (ps.shape[0], era5_lat.shape[0], era5_lon.shape[0], len(param))) else: output_data = np.zeros( (ps.shape[0], ps.shape[1], ps.shape[2], len(param))) #Assign p levels to a 3d array, with same dimensions as input variables (ta, hgt, etc.) #If the 3d p-lvl array already exists, then declare the variable "mdl_lvl" as true. try: p_3d mdl_lvl = True full_p3d = p_3d except: mdl_lvl = False if era5_interp: p_3d = np.moveaxis(np.tile(p,[ta.shape[2],ta.shape[3],1]),[0,1,2],[1,2,0]).\ astype(np.float32) else: p_3d = np.moveaxis(np.tile(p,[era5_lat.shape[0],era5_lon.shape[0],1]),[0,1,2],[1,2,0]).\ astype(np.float32) print("LOAD TIME..." + str(dt.datetime.now() - load_start)) tot_start = dt.datetime.now() for t in np.arange(0, ta.shape[0]): cape_start = dt.datetime.now() if era5_interp: ta_t = interp_era5(ta[t], lon, lat, era5_lon, era5_lat, d3=True) hur_t = interp_era5(hur[t], lon, lat, era5_lon, era5_lat, d3=True) hgt_t = interp_era5(hgt[t], lon, lat, era5_lon, era5_lat, d3=True) ps_t = interp_era5(ps[t], lon, lat, era5_lon, era5_lat, d3=False) wap_t = interp_era5(wap[t], lon, lat, era5_lon, era5_lat, d3=True) ua_t = interp_era5(ua[t], lon, lat, era5_lon, era5_lat, d3=True) va_t = interp_era5(va[t], lon, lat, era5_lon, era5_lat, d3=True) uas_t = interp_era5(uas[t], lon, lat, era5_lon, era5_lat, d3=False) vas_t = interp_era5(vas[t], lon, lat, era5_lon, era5_lat, d3=False) tas_t = interp_era5(tas[t], lon, lat, era5_lon, era5_lat, d3=False) ta2d_t = interp_era5(ta2d[t], lon, lat, era5_lon, era5_lat, d3=False) tp_t = interp_era5(tp[t], lon, lat, era5_lon, era5_lat, d3=False) mod_cape_t = interp_era5(mod_cape[t], lon, lat, era5_lon, era5_lat, d3=False) else: ta_t = ta[t] hur_t = hur[t] hgt_t = hgt[t] ps_t = ps[t] wap_t = wap[t] ua_t = ua[t] va_t = va[t] uas_t = uas[t] vas_t = vas[t] tas_t = tas[t] ta2d_t = ta2d[t] tp_t = tp[t] mod_cape_t = mod_cape[t] print(date_list[t]) output = np.zeros((1, ps_t.shape[0], ps_t.shape[1], len(param))) if mdl_lvl: if era5_interp: p_3d = interp_era5(full_p3d[t], lon, lat, era5_lon, era5_lat, d3=True) else: p_3d = full_p3d[t] dp = get_dp(hur=hur_t, ta=ta_t, dp_mask=False) #Insert surface arrays, creating new arrays with "sfc" prefix sfc_ta = np.insert(ta_t, 0, tas_t, axis=0) sfc_hgt = np.insert(hgt_t, 0, terrain, axis=0) sfc_dp = np.insert(dp, 0, ta2d_t, axis=0) sfc_p_3d = np.insert(p_3d, 0, ps_t, axis=0) sfc_ua = np.insert(ua_t, 0, uas_t, axis=0) sfc_va = np.insert(va_t, 0, vas_t, axis=0) sfc_wap = np.insert(wap_t, 0, np.zeros(vas_t.shape), axis=0) #Sort by ascending p a,temp1,temp2 = np.meshgrid(np.arange(sfc_p_3d.shape[0]) , np.arange(sfc_p_3d.shape[1]),\ np.arange(sfc_p_3d.shape[2])) sort_inds = np.flip(np.lexsort([np.swapaxes(a, 1, 0), sfc_p_3d], axis=0), axis=0) sfc_hgt = np.take_along_axis(sfc_hgt, sort_inds, axis=0) sfc_dp = np.take_along_axis(sfc_dp, sort_inds, axis=0) sfc_p_3d = np.take_along_axis(sfc_p_3d, sort_inds, axis=0) sfc_ua = np.take_along_axis(sfc_ua, sort_inds, axis=0) sfc_va = np.take_along_axis(sfc_va, sort_inds, axis=0) sfc_ta = np.take_along_axis(sfc_ta, sort_inds, axis=0) #Calculate q and wet bulb for pressure level arrays with surface values sfc_ta_unit = units.units.degC * sfc_ta sfc_dp_unit = units.units.degC * sfc_dp sfc_p_unit = units.units.hectopascals * sfc_p_3d hur_unit = mpcalc.relative_humidity_from_dewpoint(ta_t*units.units.degC, dp*units.units.degC)*\ 100*units.units.percent q_unit = mpcalc.mixing_ratio_from_relative_humidity(hur_unit,\ ta_t*units.units.degC,np.array(p_3d)*units.units.hectopascals) sfc_hur_unit = mpcalc.relative_humidity_from_dewpoint(sfc_ta_unit, sfc_dp_unit)*\ 100*units.units.percent sfc_q_unit = mpcalc.mixing_ratio_from_relative_humidity(sfc_hur_unit,\ sfc_ta_unit,sfc_p_unit) sfc_theta_unit = mpcalc.potential_temperature(sfc_p_unit, sfc_ta_unit) sfc_thetae_unit = mpcalc.equivalent_potential_temperature( sfc_p_unit, sfc_ta_unit, sfc_dp_unit) sfc_thetae = np.array(mpcalc.equivalent_potential_temperature(ps_t*units.units.hectopascals,tas_t*units.units.degC,\ ta2d_t*units.units.degC)) sfc_q = np.array(sfc_q_unit) sfc_hur = np.array(sfc_hur_unit) #sfc_wb = np.array(wrf.wetbulb( sfc_p_3d*100, sfc_ta+273.15, sfc_q, units="degC")) #Use getcape.f90 #cape_gb_mu1, cape_gb_mu4 = getcape_driver(sfc_p_3d, sfc_ta, sfc_dp, ps_t) #Now get most-unstable CAPE (max CAPE in vertical, ensuring parcels used are AGL) cape3d = wrf.cape_3d(sfc_p_3d,sfc_ta+273.15,\ sfc_q,sfc_hgt,\ terrain,ps_t,\ True,meta=False, missing=0) cape = cape3d.data[0] cin = cape3d.data[1] lfc = cape3d.data[2] lcl = cape3d.data[3] el = cape3d.data[4] #Mask values which are below the surface and above 350 hPa AGL cape[(sfc_p_3d > ps_t) | (sfc_p_3d < (ps_t - 350))] = np.nan cin[(sfc_p_3d > ps_t) | (sfc_p_3d < (ps_t - 350))] = np.nan lfc[(sfc_p_3d > ps_t) | (sfc_p_3d < (ps_t - 350))] = np.nan lcl[(sfc_p_3d > ps_t) | (sfc_p_3d < (ps_t - 350))] = np.nan el[(sfc_p_3d > ps_t) | (sfc_p_3d < (ps_t - 350))] = np.nan #Get maximum (in the vertical), and get cin, lfc, lcl for the same parcel mu_cape_inds = np.tile(np.nanargmax(cape, axis=0), (cape.shape[0], 1, 1)) mu_cape = np.take_along_axis(cape, mu_cape_inds, 0)[0] mu_cin = np.take_along_axis(cin, mu_cape_inds, 0)[0] mu_lfc = np.take_along_axis(lfc, mu_cape_inds, 0)[0] mu_lcl = np.take_along_axis(lcl, mu_cape_inds, 0)[0] mu_el = np.take_along_axis(el, mu_cape_inds, 0)[0] muq = np.take_along_axis(sfc_q, mu_cape_inds, 0)[0] * 1000 #Calculate other parameters #Thermo thermo_start = dt.datetime.now() lr700_500 = get_lr_p(ta_t, p_3d, hgt_t, 700, 500) melting_hgt = get_t_hgt(sfc_ta, np.copy(sfc_hgt), 0, terrain) melting_hgt = np.where((melting_hgt < 0) | (np.isnan(melting_hgt)), 0, melting_hgt) ta500 = get_var_p_lvl(np.copy(sfc_ta), sfc_p_3d, 500) ta925 = get_var_p_lvl(np.copy(sfc_ta), sfc_p_3d, 925) ta850 = get_var_p_lvl(np.copy(sfc_ta), sfc_p_3d, 850) ta700 = get_var_p_lvl(np.copy(sfc_ta), sfc_p_3d, 700) rho = mpcalc.density( np.array(sfc_p_3d) * (units.units.hectopascal), sfc_ta * units.units.degC, sfc_q_unit) rho925 = np.array(get_var_p_lvl(np.array(rho), sfc_p_3d, 925)) rho850 = np.array(get_var_p_lvl(np.array(rho), sfc_p_3d, 850)) rho700 = np.array(get_var_p_lvl(np.array(rho), sfc_p_3d, 700)) #Winds winds_start = dt.datetime.now() s06 = get_shear_hgt(sfc_ua, sfc_va, np.copy(sfc_hgt), 0, 6000, terrain) s0500 = get_shear_p(ua_t, va_t, p_3d, "sfc", np.array([500]), p_3d, uas=uas_t, vas=vas_t)[0] #WAP if model in ["erai", "era5"]: sfc_w = mpcalc.vertical_velocity( wap_t * (units.units.pascal / units.units.second),\ np.array(p_3d) * (units.units.hectopascal), \ ta_t * units.units.degC, q_unit) w925 = np.array(get_var_p_lvl(np.array(sfc_w), p_3d, 925)) w850 = np.array(get_var_p_lvl(np.array(sfc_w), p_3d, 850)) w700 = np.array(get_var_p_lvl(np.array(sfc_w), p_3d, 700)) #Convergence if era5_interp: x, y = np.meshgrid(era5_lon, era5_lat) else: x, y = np.meshgrid(lon, lat) dx, dy = mpcalc.lat_lon_grid_deltas(x, y) u925 = np.array(get_var_p_lvl(np.copy(sfc_ua), sfc_p_3d, 925)) u850 = np.array(get_var_p_lvl(np.copy(sfc_ua), sfc_p_3d, 850)) u700 = np.array(get_var_p_lvl(np.copy(sfc_ua), sfc_p_3d, 700)) v925 = np.array(get_var_p_lvl(np.copy(sfc_va), sfc_p_3d, 925)) v850 = np.array(get_var_p_lvl(np.copy(sfc_va), sfc_p_3d, 850)) v700 = np.array(get_var_p_lvl(np.copy(sfc_va), sfc_p_3d, 700)) conv925 = -1e5 * np.array( mpcalc.divergence(u925 * (units.units.meter / units.units.second), v925 * (units.units.meter / units.units.second), dx, dy)) conv850 = -1e5 * np.array( mpcalc.divergence(u850 * (units.units.meter / units.units.second), v850 * (units.units.meter / units.units.second), dx, dy)) conv700 = -1e5 * np.array( mpcalc.divergence(u700 * (units.units.meter / units.units.second), v700 * (units.units.meter / units.units.second), dx, dy)) #CS6 mucs6 = mu_cape * np.power(s06, 1.67) #Fill output output = fill_output(output, t, param, ps, "mu_cape", mu_cape) output = fill_output(output, t, param, ps, "mu_cin", mu_cin) output = fill_output(output, t, param, ps, "muq", muq) output = fill_output(output, t, param, ps, "s06", s06) output = fill_output(output, t, param, ps, "s0500", s0500) output = fill_output(output, t, param, ps, "lr700_500", lr700_500) output = fill_output(output, t, param, ps, "ta500", ta500) output = fill_output(output, t, param, ps, "mhgt", melting_hgt) output = fill_output(output, t, param, ps, "tp", tp_t) if (model == "erai") | (model == "era5"): output = fill_output(output, t, param, ps, "mod_cape", mod_cape_t) output_data[t] = output print("SAVING DATA...") param_out = [] for param_name in param: temp_data = output_data[:, :, :, np.where(param == param_name)[0][0]] param_out.append(temp_data) #If the mhgt variable is zero everywhere, then it is likely that data has not been read. #In this case, all values are missing, set to zero. for t in np.arange(param_out[0].shape[0]): if param_out[np.where(param == "mhgt")[0][0]][t].max() == 0: for p in np.arange(len(param_out)): param_out[p][t] = np.nan if issave: if era5_interp: save_netcdf(region, model, out_name, date_list, era5_lat, era5_lon, param, param_out, \ out_dtype = "f4", compress=True) else: save_netcdf(region, model, out_name, date_list, lat, lon, param, param_out, \ out_dtype = "f4", compress=True) print(dt.datetime.now() - tot_start)
geo_mag = mpcalc.wind_speed(ugeo, vgeo) grad_mag = geo_mag.m - (geo_mag.m**2) / (f.magnitude * R) ugrad, vgrad = mpcalc.wind_components(grad_mag * units('m/s'), wdir) # Calculate Ageostrophic wind uageo = ugrad - ugeo vageo = vgrad - vgeo # Compute QVectors uqvect, vqvect = mpcalc.q_vector(ugeo, vgeo, T * units.degC, 500 * units.hPa, dx, dy) # Calculate divergence of the ageostrophic wind div = mpcalc.divergence(uageo, vageo, dx, dy, dim_order='yx') # Calculate Relative Vorticity Advection relvor = mpcalc.vorticity(ugeo, vgeo, dx, dy, dim_order='yx') adv = mpcalc.advection(relvor, (ugeo, vgeo), (dx, dy), dim_order='yx') ###################################################################### # Create figure containing Geopotential Heights, Temperature, Divergence # of the Ageostrophic Wind, Relative Vorticity Advection (shaded), # geostrphic wind barbs, and Q-vectors. # fig = plt.figure(figsize=(10, 10)) ax = plt.subplot(111) # Plot Geopotential Height Contours
import numpy as np import metpy.calc as mpcalc import matplotlib.pyplot as plt start = timeit.default_timer() data_path = Path('/Users/mret0001/Desktop/') glob_pattern_2d = 'v900_*.nc' v_files = sorted([str(f) for f in data_path.rglob(f'{glob_pattern_2d}')]) glob_pattern_2d = 'u900_*.nc' u_files = sorted([str(f) for f in data_path.rglob(f'{glob_pattern_2d}')]) for ufile, vfile in zip(u_files, v_files): date = ufile[-24:-11] u_wind = xr.open_dataarray(home+f'/Desktop/u900_{date}_reggrid.nc') v_wind = xr.open_dataarray(home+f'/Desktop/v900_{date}_reggrid.nc') div = xr.DataArray(mpcalc.divergence(u_wind, v_wind).squeeze()) ds = div.to_dataset(name='div') ds['div'].attrs['height'] = "900 hPa" ds['div'].attrs['units'] = "1/s" ds['div'].attrs['long_name'] = "Divergence" ds.to_netcdf(home+f'/Desktop/div900_{date}_reggrid.nc') del ds, u_wind, v_wind stop = timeit.default_timer() print(f'This script needed {stop-start} seconds.')
def metpy_read_wrf(tidx, froms, froms0, fwrf, lons_out, lats_out): wrfin = Dataset(fwrf) ds = xr.Dataset() slp = getvar(wrfin, "slp", timeidx=tidx) ds['slp'] = smooth2d(slp, 3, cenweight=4) lats, lons = latlon_coords(slp) ds['lats'] = lats ds['lons'] = lons landmask = extract_vars(wrfin, timeidx=tidx, varnames=["LANDMASK"]).get("LANDMASK") u10 = extract_vars(wrfin, timeidx=tidx, varnames=["U10"]).get("U10") v10 = extract_vars(wrfin, timeidx=tidx, varnames=["V10"]).get("V10") ds['u10'] = u10.where(landmask == 0) ds['v10'] = v10.where(landmask == 0) latent = extract_vars(wrfin, timeidx=tidx, varnames=["LH"]).get("LH") ds['latent'] = smooth2d(latent, 3, cenweight=4) #latent.where(landmask==0) t2m = extract_vars(wrfin, timeidx=tidx, varnames=["T2"]).get("T2") - 273.15 ds['t2m'] = smooth2d(t2m, 3, cenweight=4) sst = extract_vars(wrfin, timeidx=tidx, varnames=["SST"]).get("SST") ds['sst'] = sst.where(landmask == 0) romsin = xr.open_dataset(froms) romsin = romsin.rename({"lat_rho": "lat", "lon_rho": "lon"}) romsin = romsin.isel(ocean_time=tidx) ds['sst_5m'] = romsin.isel(z_r=0).temp ds['water_temp'] = romsin.temp ds['water_ucur'] = romsin.ucur / 100 ds['water_vcur'] = romsin.vcur / 100 ds.water_ucur.attrs['units'] = 'm/s' ds.water_vcur.attrs['units'] = 'm/s' romsin = xr.open_dataset(froms0) romsin = romsin.rename({"lat_rho": "lat", "lon_rho": "lon"}) romsin = romsin.isel(ocean_time=tidx) ds['h'] = romsin.h mld = get_oml_depth(froms0, t_in=tidx) mld = smooth2d(mld, 3, cenweight=4) ds['oml_depth'] = xr.DataArray(mld, ds.sst_5m.coords, ds.sst_5m.dims, ds.sst_5m.attrs) ds['oml_depth2'] = ds.oml_depth.where(ds.h > 20) ds = ds.drop(['XLONG', 'XLAT', 'XTIME', 'Time']) ds = ds.rename({'south_north': 'eta_rho', 'west_east': 'xi_rho'}) interp_method = 'bilinear' ds_out = xr.Dataset({ 'lat': (['lat'], lats_out), 'lon': (['lon'], lons_out) }) regridder = xe.Regridder(ds, ds_out, interp_method) regridder.clean_weight_file() ds = regridder(ds) ds = ds.squeeze() dxy = 10000. ds = ds.metpy.parse_cf().squeeze() utau, vtau = wind_stress(to_np(ds.u10), to_np(ds.v10)) ds['u_tau'] = xr.DataArray(utau, ds.u10.coords, ds.u10.dims, ds.u10.attrs) ds['v_tau'] = xr.DataArray(vtau, ds.v10.coords, ds.v10.dims, ds.v10.attrs) curl = mpcalc.vorticity(utau * units('m/s'), utau * units('m/s'), dx=dxy * units.meter, dy=dxy * units.meter) ds['wind_stress_curl'] = xr.DataArray(np.array(curl), ds.u10.coords, ds.u10.dims, ds.u10.attrs) div = [] for z in range(len(ds.z_r)): div0 = mpcalc.divergence(ds.water_ucur.isel(z_r=z) * units('m/s'), ds.water_vcur.isel(z_r=z) * units('m/s'), dx=dxy * units.meter, dy=dxy * units.meter) div.append(div0) ds['cur_div'] = xr.DataArray(div, ds.water_ucur.coords, ds.water_ucur.dims, ds.water_ucur.attrs) div = mpcalc.divergence(ds.water_ucur.isel(z_r=2) * units('m/s'), ds.water_vcur.isel(z_r=2) * units('m/s'), dx=dxy * units.meter, dy=dxy * units.meter) ds['surf_cur_div'] = xr.DataArray(np.array(div), ds.u10.coords, ds.u10.dims, ds.u10.attrs) print(ds) return ds
def ComputePlumbFluxesDecomp(uclm, vclm, z, zclm, lat, lon, lev): #input: model climatology of u,v, z #z for the event to analyze #restringe domain to latitudes south to 0 (to avoid problems with sin (0°)) z = z[:, lat < 0, :] zclm = zclm[lat < 0, :] uclm = uclm[lat < 0, :] vclm = vclm[lat < 0, :] lat = lat[lat < 0] [nrealiz, nlats, nlons] = np.shape(z) #get dimensions #earth radius a = 6400000 coslat = np.cos(lat * 3.14 / 180) sinlat = np.sin(lat * 3.14 / 180) #Coriolis parameter f = 2 * 7.29e-5 * sinlat #gravity g = 9.8 # magnitude of basic state wind speed magU = np.sqrt(np.add(np.power(uclm, 2), np.power(vclm, 2))) #QG stream function psia_clim = g / np.transpose(np.tile(f, (nlons, 1))) * zclm psia = g / np.transpose(np.tile(f, (nrealiz, nlons, 1)), (0, 2, 1)) * z psia_m = np.mean(psia, axis=0) #anomalies with respecto to the ensemble mean psiaa = psia - psia_m psia_m = psia_m - psia_clim #psia_clim = psia_clim - np.transpose(np.tile(np.mean(psia_clim, axis=1), (nlons, 1))) #psi derivatives dpsiamdlon = c_diff(psia_m, lon * 3.14 / 180, 1) ddpsiamdlonlon = c_diff(dpsiamdlon, lon * 3.14, 1) dpsiamdlat = c_diff(psia_m, lat * 3.14 / 180, 0) ddpsiamdlatlat = c_diff(dpsiamdlat, lat * 3.14 / 180, 0) ddpsiamdlatlon = c_diff(dpsiamdlat, lon * 3.14 / 180, 1) dpsiacdlon = c_diff(psia_clim, lon * 3.14 / 180, 1) ddpsiacdlonlon = c_diff(dpsiacdlon, lon * 3.14, 1) dpsiacdlat = c_diff(psia_clim, lat * 3.14 / 180, 0) ddpsiacdlatlat = c_diff(dpsiacdlat, lat * 3.14 / 180, 0) ddpsiacdlatlon = c_diff(dpsiacdlat, lon * 3.14 / 180, 1) dpsidlon = c_diff(psiaa, lon * 3.14 / 180, 2) ddpsidlonlon = c_diff(dpsidlon, lon * 3.14, 2) dpsidlat = c_diff(psiaa, lat * 3.14 / 180, 1) ddpsidlatlat = c_diff(dpsidlat, lat * 3.14 / 180, 1) ddpsidlatlon = c_diff(dpsidlat, lon * 3.14 / 180, 2) #linear terms termx1l = 2 * dpsiacdlon * dpsiamdlon - psia_clim * ddpsiamdlonlon - psia_m * ddpsiacdlonlon termxyl = dpsiacdlon * dpsiamdlat + dpsiamdlon * dpsiacdlat -\ psia_clim * ddpsiamdlatlon - psia_m * ddpsiacdlatlon termy1l = 2 * dpsiacdlat * dpsiamdlat - psia_clim * ddpsiamdlatlat - psia_m * ddpsiacdlatlat # "p" is normalized by 1000hPa coefcos = np.transpose(np.tile(coslat, (nlons, 1))) coeff = coefcos * (lev / 100000) / (2 * magU) #coeff = 1 / (2 * magU) #x-component pxl = coeff / (a * a * coefcos) * (uclm * termx1l / coefcos + vclm * termxyl) #px = coeff * (uclm * termxu + vclm * termxv) #y-component pyl = coeff / (a * a) * (uclm / coefcos * termxyl + vclm * termy1l) #py = coeff * ( uclm * termxv + vclm * termyv) #non-linear terms termx1nl = dpsiamdlon**2 - psia_m * ddpsiamdlonlon termxynl = dpsiamdlat * dpsiamdlon - psia_m * ddpsiamdlatlon termy1nl = dpsiamdlat**2 - psia_m * ddpsiamdlatlat #x-component pxnl = coeff / (a * a * coefcos) * (uclm * termx1nl / coefcos + vclm * termxynl) #px = coeff * (uclm * termxu + vclm * termxv) #y-component pynl = coeff / (a * a) * (uclm / coefcos * termxynl + vclm * termy1nl) #py = coeff * ( uclm * termxv + vclm * termyv) #compute divergence dx, dy = calc.lat_lon_grid_deltas(lon, lat) divl = calc.divergence(pxl, pyl, dx, dy) divnl = calc.divergence(pxnl, pynl, dx, dy) pxem = pxl + pxnl pyem = pyl + pynl divem = calc.divergence(pxem, pyem, dx, dy) var = { 'pxl': pxl, 'pyl': pyl, 'pxnl': pxnl, 'pynl': pynl, 'divl': divl, 'divnl': divnl, 'pxem': pxem, 'pyem': pyem, 'divem': divem, 'lat': lat } return var