def load_sites_and_obtain_their_grid_locations(wrf_in,sensors_file,sensor_file_type): # open the sensor dataset if(sensor_file_type == 'CSV'): sens = pd.read_csv(sensors_file,usecols=['long','lat','sensor_name']) elif(sensor_file_type == 'RDATA'): metadata = pyreadr.read_r(sensors_file.as_posix()) sens = metadata['AURN_metadata'][['site_id','latitude','longitude']].drop_duplicates() else: print('Sensor file type not recognised: {}'.format(sensor_file_type)) print('Should be CSV or RDATA, please amend as appropriate.') quit() # get the indexes from the wrf file with salem.open_wrf_dataset(wrf_in) as ds: sens['iarray'],sens['jarray'] = ds.salem.grid.transform(sens['longitude'],sens['latitude'],nearest=True) #%% check to make sure that all stations are within our model domain - drop those that aren't if any(sens['iarray']>ds.dims['west_east']) or any(sens['jarray']>ds.dims['south_north']): print('dropping these stations outside of the model domain:') print(sens[sens['jarray']>ds.dims['south_north']]) print(sens[sens['iarray']>ds.dims['west_east']]) sens = sens[sens['jarray']<=ds.dims['south_north']] sens = sens[sens['iarray']<=ds.dims['west_east']] return(sens)
def mask_lake(data_dir, shp): geo_path = '/home/zzhzhao/Model/tests/test-15/WPS/geo_em.d02.nc' lu = salem.open_wrf_dataset(os.path.join( data_dir, geo_path))['LU_INDEX'].isel(time=0) lu_lake = lu.salem.roi(shape=shp) mask = xr.where(lu_lake.notnull(), True, False) return mask
def test_cartopy_polar(): import cartopy fig = plt.figure(figsize=(8, 8)) ods = open_wrf_dataset(get_demo_file('geo_em_d02_polarstereo.nc')) ods = ods.isel(time=0) ax = plt.subplot(2, 2, 1) smap = ods.salem.get_map() smap.plot(ax=ax) p = ods.salem.cartopy() ax = plt.subplot(2, 2, 2, projection=p) ax.coastlines(resolution='50m') ax.gridlines() ax.set_extent(ods.salem.grid.extent, crs=p) ds = ods.salem.subset(corners=((-52.8, 70.11), (-52.8, 70.11)), margin=12) ax = plt.subplot(2, 2, 3) smap = ds.HGT_M.salem.quick_map(ax=ax, cmap='Oranges') ax.scatter(ds.XLONG_M, ds.XLAT_M, s=5, transform=smap.transform(ax=ax)) p = ds.salem.cartopy() ax = plt.subplot(2, 2, 4, projection=p) ds.HGT_M.plot.imshow(ax=ax, transform=p, cmap='Oranges') ax.coastlines(resolution='50m') ax.gridlines() ax.scatter(ds.XLONG_M, ds.XLAT_M, transform=cartopy.crs.PlateCarree(), s=5) return fig
def test_lookup_transform(): dsw = open_wrf_dataset(get_demo_file('wrfout_d01.nc')) dse = open_xr_dataset(get_demo_file('era_interim_tibet.nc')) out = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len) fig, ax = plt.subplots(1, 1) out.salem.quick_map(ax=ax) return fig
def mask_lake(data_dir, shp, testname, domain): # geo_path = f'/home/zzhzhao/Model/tests/{testname}/WPS/geo_em.d{domain:0>2d}.nc' geo_path = f'/home/zzhzhao/Model/tests/test-25/WPS/geo_em.d{domain:0>2d}.nc' lu = salem.open_wrf_dataset(os.path.join( data_dir, geo_path))['LU_INDEX'].isel(time=0) lu_lake = lu.salem.roi(shape=shp) mask = xr.where(lu_lake.notnull(), True, False) return mask
def test_lookup_transform(): dsw = open_wrf_dataset(get_demo_file('wrfout_d01.nc')) dse = open_xr_dataset(get_demo_file('era_interim_tibet.nc')) out = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len) fig, ax = plt.subplots(1, 1) sm = out.salem.get_map() sm.set_data(out) sm.set_geometry(dsw.salem.grid.extent_as_polygon(), edgecolor='r', linewidth=2) sm.visualize(ax=ax) return fig
def test_plot_on_map(): import salem from salem.utils import get_demo_file ds = salem.open_wrf_dataset(get_demo_file('wrfout_d01.nc')) t2_sub = ds.salem.subset(corners=((77., 20.), (97., 35.)), crs=salem.wgs84).T2.isel(time=2) shdf = salem.read_shapefile(get_demo_file('world_borders.shp')) shdf = shdf.loc[shdf['CNTRY_NAME'].isin( ['Nepal', 'Bhutan'])] # GeoPandas' GeoDataFrame t2_sub = t2_sub.salem.subset(shape=shdf, margin=2) # add 2 grid points t2_roi = t2_sub.salem.roi(shape=shdf) fig, ax = plt.subplots(1, 1) t2_roi.salem.quick_map(ax=ax) plt.tight_layout() return fig
def test_cartopy(): import cartopy fig = plt.figure(figsize=(8, 11)) ods = open_wrf_dataset(get_demo_file('wrfout_d01.nc')) ax = plt.subplot(3, 2, 1) smap = ods.salem.get_map() smap.plot(ax=ax) p = ods.salem.cartopy() ax = plt.subplot(3, 2, 2, projection=p) ax.coastlines() ax.add_feature(cartopy.feature.BORDERS, linestyle='-') ax.set_extent(ods.salem.grid.extent, crs=p) ds = ods.isel(west_east=slice(22, 28), south_north=slice(0, 6)) ds = ds.T2C.mean(dim='time', keep_attrs=True) ax = plt.subplot(3, 2, 3) smap = ds.salem.quick_map(ax=ax) ax.scatter(ds.lon, ds.lat, transform=smap.transform(ax=ax)) p = ds.salem.cartopy() ax = plt.subplot(3, 2, 4, projection=p) ds.plot.imshow(ax=ax, transform=p) ax.coastlines() ax.scatter(ds.lon, ds.lat, transform=cartopy.crs.PlateCarree()) ds = ods.isel(west_east=slice(80, 86), south_north=slice(80, 86)) ds = ds.T2C.mean(dim='time', keep_attrs=True) ax = plt.subplot(3, 2, 5) smap = ds.salem.quick_map(ax=ax, factor=1) ax.scatter(ds.lon, ds.lat, transform=smap.transform(ax=ax)) p = ds.salem.cartopy() ax = plt.subplot(3, 2, 6, projection=p) ds.plot.imshow(ax=ax, transform=p) ax.coastlines() ax.scatter(ds.lon, ds.lat, transform=cartopy.crs.PlateCarree()) return fig
# -*- coding: utf-8 -*- """ ==================== Customize Salem maps ==================== How to change the look of a Map? """ import salem import matplotlib.pyplot as plt # get the map from a WRF file ds = salem.open_wrf_dataset(salem.get_demo_file('wrfout_d01.nc')) smap = ds.salem.get_map(countries=False) # Change the country borders smap.set_shapefile(countries=True, color='C3', linewidths=2) # Add oceans and lakes smap.set_shapefile(oceans=True) smap.set_shapefile(rivers=True) smap.set_shapefile(lakes=True, facecolor='blue', edgecolor='blue') # Change the lon-lat countour setting smap.set_lonlat_contours(add_ytick_labels=False, interval=5, linewidths=1.5, linestyles='-', colors='C1')
# -*- coding: utf-8 -*- """ ============= Plot overlays ============= Add contours and wind arrows to a salem plot """ import salem import numpy as np import matplotlib.pyplot as plt # get the data at the latest time step ds = salem.open_wrf_dataset(salem.get_demo_file('wrfout_d01.nc')).isel(time=-1) # get the wind data at 10000 m a.s.l. u = ds.salem.wrf_zlevel('U', 10000.) v = ds.salem.wrf_zlevel('V', 10000.) ws = ds.salem.wrf_zlevel('WS', 10000.) # get the axes ready f, ax = plt.subplots() # plot the salem map background, make countries in grey smap = ds.salem.get_map(countries=False) smap.set_shapefile(countries=True, color='grey') smap.plot(ax=ax) # transform the coordinates to the map reference system and contour the data
def mask_lake(data_dir, shp): wrf_file = [f for f in os.listdir(data_dir) if f[9]=='2'][0] lu = salem.open_wrf_dataset(os.path.join(data_dir, wrf_file))['LU_INDEX'].isel(time=0) lu_lake = lu.salem.roi(shape=shp) mask = xr.where(lu_lake.notnull(), True, False) return mask
def wrf_to_tell_counties( wrf_file: str, wrf_variables: List[str], precisions: List[int], county_shapefile: str = './Geolocation/tl_2020_us_county/tl_2020_us_county.shp', weight_and_mapping_file: str = './grid_cell_to_county_weight.parquet', output_directory: str = './County_Output_Files', output_filename_suffix: str = '_County_Mean_Meteorology', n_jobs: int = -1, ) -> None: """ Aggregate WRF output data to county level using area weighted average. :param str wrf_file: path to the WRF output file to aggregate :param list(str) wrf_variables: list of variables to aggregate :param list(int) precisions: list of precisions corresponding to the variables to aggregate :param str county_shapefile: path to a shapefile (.shp) with county geometries :param str weight_and_mapping_file: path to read or write a weights file which maps WRF grid cell to county weight :param str output_directory: path to which output should be written :param str output_filename_suffix: string to append to the timestamp for the output file name :param int n_jobs: number of time slices to process in parallel """ begin_time = datetime.datetime.now() if not isfile(wrf_file): raise FileNotFoundError('No file to process, exiting...') wrf = salem.open_wrf_dataset(wrf_file) # if there's not already a mapping file, create one if not isfile(weight_and_mapping_file): # using the first file and time: # * get the crs # * create the mapping of cell index to county and weight wrf_crs = wrf.pyproj_srs wrf_df = wrf[wrf_variables].isel(time=0).to_dataframe().reset_index( drop=True) wrf_df = gpd.GeoDataFrame( wrf_df, geometry=wrf.salem.grid.to_geometry().geometry).set_crs(wrf_crs) wrf_df['cell_index'] = wrf_df.index.values # load the counties and reproject to WRF projection counties = gpd.read_file(county_shapefile)[["GEOID", "geometry" ]].rename(columns={ 'GEOID': 'FIPS', }).to_crs(wrf_crs) # find the intersection between counties and wrf cells try: intersection = gpd.overlay(counties, wrf_df, how='intersection') except ValueError as e: raise ValueError(f''' The intersection of county geometry and WRF grid cells resulted in invalid geometry. Please double check the geometry. {str(e)} ''') # weight by the intersection area intersection['area'] = intersection.area intersection['weight'] = ( intersection['area'] / intersection[['FIPS', 'area' ]].groupby('FIPS').area.transform('sum')) # reuse this mapping for the remaining files and slices mapping = intersection[['cell_index', 'FIPS', 'weight']] mapping.to_parquet(weight_and_mapping_file) # create the first output file write_output_file( compute_county_weighted_mean( intersection, wrf_variables, precisions, ), wrf_df.time.iloc[0], output_directory, output_filename_suffix, ) t_start = 1 else: mapping = pd.read_parquet(weight_and_mapping_file) t_start = 0 # create the remaining output for each time slice in each file Parallel(n_jobs=n_jobs)(delayed(process_time_slice)( wrf[wrf_variables].isel(time=i).to_dataframe().reset_index(drop=True), wrf_variables, precisions, mapping, output_directory, output_filename_suffix, ) for i in range(wrf.time.shape[0])[t_start:]) print('Elapsed time = ', datetime.datetime.now() - begin_time)
def main(): # constants BIGFONT = 22 MIDFONT = 18 SMFONT = 16 FIG_WIDTH = 15.0 FIG_HEIGHT = 10.0 #cases=["ERA5_C2008", "ERA5_TY2001", "ERA5_WRFROMS", "ERA5_WRF"] #cases=["ERA5_TY2001", "ERA5_WRFROMS"] ctrl_case = 'ERA5_WRFROMS' sen_case = 'ERA5_TY2001' #line_libs=['b.','g*','r^','k+'] #line_libs=['b.','b*','g.','g*','r^','k+'] line_libs = ['b.', 'g.', 'r.', 'k.'] wrf_root = '/disk/v092.yhuangci/lzhenn/1911-COAWST/' i_dom = 2 strt_time_str = '201809152200' end_time_str = '201809160600' box_R = 80 strt_time_obj = datetime.datetime.strptime(strt_time_str, '%Y%m%d%H%M') end_time_obj = datetime.datetime.strptime(end_time_str, '%Y%m%d%H%M') dateparse = lambda x: datetime.datetime.strptime(x, '%Y%m%d%H0000') # read track data ctrl tc_info_fn = wrf_root + '/' + ctrl_case + '/trck.' + ctrl_case + '.d0' + str( i_dom) df_tc_info = pd.read_csv(tc_info_fn, sep='\s+', parse_dates=True, index_col='timestamp', header=0, date_parser=dateparse) df_tc_info_ctrl = df_tc_info[((df_tc_info.index >= strt_time_obj) & (df_tc_info.index <= end_time_obj))] # read track data sen tc_info_fn = wrf_root + '/' + sen_case + '/trck.' + sen_case + '.d0' + str( i_dom) df_tc_info = pd.read_csv(tc_info_fn, sep='\s+', parse_dates=True, index_col='timestamp', header=0, date_parser=dateparse) df_tc_info_sen = df_tc_info[((df_tc_info.index >= strt_time_obj) & (df_tc_info.index <= end_time_obj))] tc_lat_ctrl = df_tc_info_ctrl['lat'] tc_lon_ctrl = df_tc_info_ctrl['lon'] tc_lat_sen = df_tc_info_sen['lat'] tc_lon_sen = df_tc_info_sen['lon'] # read raw input ds = salem.open_wrf_dataset('/disk/v092.yhuangci/lzhenn/1911-COAWST/' + ctrl_case + '/wrfout_d02') ds = ds.sel(time=slice(strt_time_obj, end_time_obj)) var1_ctrl = ds['LH'] # heat exch varmask_ctrl = ds['LANDMASK'] var1_ctrl.values = np.where(varmask_ctrl.values == 1, np.nan, var1_ctrl.values) # get rid off land points idx_ctrl = get_closest_idx(var1_ctrl.lat, var1_ctrl.lon, tc_lat_ctrl.values, tc_lon_ctrl.values) var1_ctrl_box_comp = box_composite(var1_ctrl.values, box_R, idx_ctrl) # nparray inout # read raw input ds = salem.open_wrf_dataset('/disk/v092.yhuangci/lzhenn/1911-COAWST/' + sen_case + '/wrfout_d02') ds = ds.sel(time=slice(strt_time_obj, end_time_obj)) var1_sen = ds['LH'] # heat exch varmask_sen = ds['LANDMASK'] var1_sen.values = np.where(varmask_sen.values == 1, np.nan, var1_sen.values) # get rid off land points idx_sen = get_closest_idx(var1_sen.lat, var1_sen.lon, tc_lat_sen.values, tc_lon_sen.values) var1_sen_box_comp = box_composite(var1_sen.values, box_R, idx_sen) # nparray inout var1_diff = var1_sen_box_comp - var1_ctrl_box_comp fig, ax = plt.subplots() # fig.subplots_adjust(left=0.08, bottom=0.18, right=0.99, top=0.92, wspace=None, hspace=None) im = ax.imshow(var1_diff) # plt.legend(loc='best', fontsize=SMFONT) plt.xlabel('distanceX', fontsize=SMFONT) plt.ylabel('distanceY', fontsize=SMFONT) plt.xticks(fontsize=SMFONT) plt.yticks(fontsize=SMFONT) plt.title('LH Diff, Box Composite', fontsize=BIGFONT) # fig.set_size_inches(FIG_WIDTH, FIG_HEIGHT) fig.savefig('../fig/boxdiff_LHF.png')
obs_source = 'china_measurements' if year in ['2014', '2015', '2016', '2017', '2020']: obs_filename = f'/nobackup/WRFChem/measurements/china_measurements_corrected/df_obs_summary_{year}.csv' elif use_openaq: obs_source = 'openaq' if year in ['2013', '2014', '2015', '2016', '2017', '2020']: obs_filename = f'/nobackup/WRFChem/openaq/csv/openaq_data_{year}_noduplicates.csv' elif year in ['2018', '2019']: if int(month) < 7: obs_filename = f'/nobackup/WRFChem/openaq/csv/openaq_data_{year}_part1_noduplicates.csv' else: obs_filename = f'/nobackup/WRFChem/openaq/csv/openaq_data_{year}_part2_noduplicates.csv' df_obs = pd.read_csv(obs_filename, index_col="date.utc", parse_dates=True) ds_wrfout = salem.open_wrf_dataset(wrfout_file) year_month_day = str(ds_wrfout.time.values)[2:12] hour = str(ds_wrfout.time.values)[13:15] df_obs_dt = df_obs[year_month_day].loc[df_obs[year_month_day].index.hour == int(hour)] for country in countries: df_obs_dt_country = df_obs_dt.loc[df_obs_dt.country == country] for parameter in parameters.keys(): ds_wrfout_parameter = ds_wrfout[parameters[parameter]].isel( bottom_top=0) have_weights = len(glob.glob(f'{path}/bilinear*')) != 0 regridder = xe.Regridder(ds_wrfout_parameter, grid_rectilinear_global,