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_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
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,
