def test_generate_grid_ds():
# simple case...just the dims
axis_dims = {'X': 'lon', 'Y': 'lat', 'Z': 'z'}
axis_coords = {'X': 'llon', 'Y': 'llat', 'Z': 'zz'}
ds_old = ds_original.copy()
ds_new = generate_grid_ds(ds_old,
axis_dims,
boundary_discontinuity={
'lon': 360,
'lat': 180
},
pad={'z': 'auto'})
assert_equal(ds_new,
ds_out_left.drop(['llon_left', 'llat_left', 'zz_left']))
# TODO why are they not identical ? assert identical fails
ds_new = generate_grid_ds(ds_original,
axis_dims,
axis_coords,
boundary_discontinuity={
'lon': 360,
'lat': 180,
'llon': 360,
'llat': 180
},
pad={
'z': 'auto',
'zz': 'auto'
})
assert_equal(ds_new, ds_out_left)
def recreate_grid_simple(ds, lon_name="lon", lat_name="lat"):
ds_full = generate_grid_ds(ds, {
"X": "x",
"Y": "y"
},
position=("center", "right"))
grid = Grid(ds_full, periodic=["X"])
ds_full = recreate_metrics(ds, grid)
return ds_full
def load_gos_data(gos_filenames):
#import xgcm
from xgcm import Grid
from xgcm.autogenerate import generate_grid_ds
# ====== load in all .nc files and combine into one xarray dataset
gos_map = xr.open_mfdataset(gos_filenames)
gos_map = gos_map.rename({'latitude': 'lat'}).rename({'longitude': 'lon'})
gos_select = gos_map #.sel(time='2016-11-19',lon=slice(10,16),lat=slice(-28,-24))
#gos_map.ugos
#dx = gos_map.lon.diff('lon')
#gos_map['rel_vort'] = gos_map.vgos.diff('lon')/gos_map.lon.diff('lon')
#gos_select = gos_map #gos_map.sel(time='2016-11-19',lon=slice(10,16),lat=slice(-28,-24))
# create grid for interpolation, differencing
#grid = xgcm.Grid(gos_select)
# for Satellite data:
# https://xgcm.readthedocs.io/en/latest/autogenerate_examples.html
ds_full = generate_grid_ds(gos_select, {'X': 'lon', 'Y': 'lat'})
ds_full.vgos
grid = Grid(ds_full, periodic=['X'])
# compute the difference (in degrees) along the longitude and latitude for both the cell center and the cell face
# need to specify the boundary_discontinutity in order to avoid the introduction of artefacts at the boundary
dlong = grid.diff(ds_full.lon, 'X', boundary_discontinuity=360)
dlonc = grid.diff(ds_full.lon_left, 'X', boundary_discontinuity=360)
#dlonc_wo_discontinuity = grid.diff(ds_full.lon_left, 'X')
dlatg = grid.diff(ds_full.lat, 'Y', boundary='fill', fill_value=np.nan)
dlatc = grid.diff(ds_full.lat_left,
'Y',
boundary='fill',
fill_value=np.nan)
# converted into approximate cartesian distances on a globe.
ds_full.coords['dxg'], ds_full.coords['dyg'] = dll_dist(
dlong, dlatg, ds_full.lon, ds_full.lat)
ds_full.coords['dxc'], ds_full.coords['dyc'] = dll_dist(
dlonc, dlatc, ds_full.lon, ds_full.lat)
# Relative vorticity: ζ = ∂ v/∂ x – ∂ u/∂ y
ds_full['dv_dx'] = grid.diff(ds_full.vgos, 'X') / ds_full.dxg
ds_full['du_dy'] = grid.diff(
ds_full.ugos, 'Y', boundary='fill', fill_value=np.nan) / ds_full.dyg
dv_dx = grid.interp(ds_full['dv_dx'],
'Y',
boundary='fill',
fill_value=np.nan) # get dv_dx and du_dy on same grid
du_dy = grid.interp(ds_full['du_dy'],
'X',
boundary='fill',
fill_value=np.nan)
ds_full['Rel_Vort'] = dv_dx - du_dy
# Vorticity Rossby Number = ζ / f
ds_full['Ro'] = ds_full.Rel_Vort / coriolis(ds_full.Rel_Vort.lat_left)
return ds_full
def test_generate_grid_ds():
# This needs more cases
axis_dims = {"X": "lon", "Y": "lat", "Z": "z"}
axis_coords = {"X": "llon", "Y": "llat", "Z": "zz"}
position = ("center", "outer")
boundary_discontinuity = {"lon": 360, "llon": 360, "lat": 180, "llat": 180}
pad = {"z": "auto", "zz": "auto"}
ds = generate_grid_ds(ds_original_left, axis_dims, axis_coords, position,
boundary_discontinuity, pad)
assert_equal(ds, ds_out_outer)
def test_generate_grid_ds():
# This needs more cases
axis_dims = {'X': 'lon', 'Y': 'lat', 'Z': 'z'}
axis_coords = {'X': 'llon', 'Y': 'llat', 'Z': 'zz'}
position = ('center', 'outer')
boundary_discontinuity = {'lon': 360, 'llon': 360, 'lat': 180, 'llat': 180}
pad = {'z': 'auto', 'zz': 'auto'}
ds = generate_grid_ds(ds_original_left, axis_dims, axis_coords, position,
boundary_discontinuity, pad)
assert_equal(ds, ds_out_outer)
def create_full_grid(base_ds, grid_dict=None):
"""Generate a full xgcm-compatible dataset from a reference datasets `base_ds`.
This dataset should be representing a tracer fields, e.g. the cell center.
Parameters
----------
base_ds : xr.Dataset
The reference ('base') datasets, assumed to be at the tracer position/cell center
grid_dict : dict, optional
Dictionary with info about the grid staggering.
Must be encoded using the base_ds attrs (e.g. {'model_name':{'axis_shift':{'X':'left',...}}}).
If deactivated (default), will load from the internal database for CMIP6 models, by default None
Returns
-------
xr.Dataset
xgcm compatible dataset
"""
# load dict with grid shift info for each axis
if grid_dict is None:
ff = open(grid_spec, "r")
grid_dict = yaml.safe_load(ff)
ff.close()
source_id = base_ds.attrs["source_id"]
grid_label = base_ds.attrs["grid_label"]
# if source_id not in dict, and grid label is gn, warn and ask to submit an issue
try:
axis_shift = grid_dict[source_id][grid_label]["axis_shift"]
except KeyError:
warnings.warn(
f"Could not find the source_id/grid_label ({source_id}/{grid_label}) combo in `grid_dict`, returning `None`. Please submit an issue to github: https://github.com/jbusecke/cmip6_preprocessing/issues"
)
return None
position = {k: ("center", axis_shift[k]) for k in axis_shift.keys()}
axis_dict = {"X": "x", "Y": "y"}
ds_grid = generate_grid_ds(base_ds,
axis_dict,
position=position,
boundary_discontinuity={"X": 360})
# TODO: man parse lev and lev_bounds as center and outer dims.
# I should also be able to do this with `generate_grid_ds`, but here we
# have the `lev_bounds` with most models, so that is probably more reliable.
# cheapest solution right now
if "lev" in ds_grid.dims:
ds_grid["lev"].attrs["axis"] = "Z"
return ds_grid
def test_generate_grid_ds():
# simple case...just the dims
axis_dims = {'X': 'lon', 'Y': 'lat', 'Z': 'z'}
axis_coords = {'X': 'llon', 'Y': 'llat', 'Z': 'zz'}
ds_old = ds_original.copy()
ds_new = generate_grid_ds(ds_old, axis_dims,
boundary_discontinuity={'lon': 360, 'lat': 180},
pad={'z': 'auto'})
assert_equal(ds_new, ds_out_left.drop(['llon_left',
'llat_left',
'zz_left']))
# TODO why are they not identical ? assert identical fails
ds_new = generate_grid_ds(ds_original,
axis_dims,
axis_coords,
boundary_discontinuity={'lon': 360,
'lat': 180,
'llon': 360,
'llat': 180},
pad={'z': 'auto', 'zz': 'auto'})
assert_equal(ds_new, ds_out_left)
def recreate_full_grid_old(ds,
lon_bound_name='lon_bounds',
lat_bound_name='lat_bounds'):
ds_full = generate_grid_ds(ds, {
'X': 'x',
'Y': 'y'
},
position=('center', 'right'))
grid = Grid(ds_full)
# Derive distances at u point
dlon = 0 # ill interpolate that later
dlat = ds[lat_bound_name].isel(vertex=1) - ds[lat_bound_name].isel(
vertex=2)
# interpolate the centered position
lat = (ds[lat_bound_name].isel(vertex=1) +
ds[lat_bound_name].isel(vertex=2)) / 2
lon = (ds[lon_bound_name].isel(vertex=1) +
ds[lon_bound_name].isel(vertex=2)) / 2
_, dy = dll_dist(dlon, dlat, lon, lat)
# strip coords and rename dims
ds_full.coords['dye'] = (['y', 'x_right'], dy.data)
# Derive distances at v point
dlon = ds[lon_bound_name].isel(vertex=3) - ds[lon_bound_name].isel(
vertex=2)
print(dlon)
# special check for dlon
temp = dlon.load().data
temp[temp < 0] = temp[temp < 0] + 360.0
dlon.data = temp
dlat = 0 # ill interpolate that later
# interpolate the centered position
lat = (ds[lat_bound_name].isel(vertex=3) +
ds[lat_bound_name].isel(vertex=2)) / 2
lon = (ds[lon_bound_name].isel(vertex=3) +
ds[lon_bound_name].isel(vertex=2)) / 2
dx, _ = dll_dist(dlon, dlat, lon, lat)
# strip coords and rename dims
ds_full.coords['dxn'] = (['y_right', 'x'], dx.data)
# interpolate the missing metrics
ds_full.coords['dxt'] = grid.interp(ds_full.coords['dxn'], 'Y')
ds_full.coords['dxne'] = grid.interp(ds_full.coords['dxn'], 'X')
ds_full.coords['dyt'] = grid.interp(ds_full.coords['dye'], 'X')
ds_full.coords['dyne'] = grid.interp(ds_full.coords['dye'], 'Y')
return ds_full
def recreate_full_grid(ds, lon_name="lon", lat_name="lat"):
ds_full = generate_grid_ds(ds, {
"X": "x",
"Y": "y"
},
position=("center", "right"))
grid = Grid(ds_full, periodic=['X'])
# infer dx at eastern bound from tracer points
lon0, lon1 = grid.axes["X"]._get_neighbor_data_pairs(
ds_full.lon.load(), "right")
lat0 = lat1 = ds_full.lat.load().data
dx = distance(lon0, lat0, lon1, lat1)
ds_full.coords["dxe"] = xr.DataArray(dx,
coords=grid.interp(ds_full.lon,
"X").coords)
# infer dy at northern bound from tracer points
lat0, lat1 = grid.axes["Y"]._get_neighbor_data_pairs(
ds_full.lat.load(), "right", boundary="extrapolate")
lon0 = lon1 = ds_full.lon.load().data
dy = distance(lon0, lat0, lon1, lat1)
ds_full.coords["dyn"] = xr.DataArray(dy,
coords=grid.interp(
ds_full.lat,
"Y",
boundary="extrapolate").coords)
# now simply interpolate all the other metrics
ds_full.coords['dxt'] = grid.interp(ds_full.coords['dxe'], 'X')
ds_full.coords['dxne'] = grid.interp(ds_full.coords['dxe'],
'Y',
boundary="extrapolate")
ds_full.coords['dxn'] = grid.interp(ds_full.coords['dxt'],
'Y',
boundary="extrapolate")
ds_full.coords['dyt'] = grid.interp(ds_full.coords['dyn'],
'Y',
boundary="extrapolate")
ds_full.coords['dyne'] = grid.interp(ds_full.coords['dyn'], 'X')
ds_full.coords['dye'] = grid.interp(ds_full.coords['dyt'], 'X')
ds_full.coords['area_t'] = ds_full.coords['dxt'] * ds_full.coords['dyt']
ds_full.coords['area_e'] = ds_full.coords['dxe'] * ds_full.coords['dye']
ds_full.coords['area_ne'] = ds_full.coords['dxne'] * ds_full.coords['dyne']
ds_full.coords['area_n'] = ds_full.coords['dxn'] * ds_full.coords['dyn']
# should i return the coords to dask?
return ds_full
def merge_variables_on_staggered_grid(data_dict,
modelname,
tracer_ref='thetao',
u_ref='uo',
v_ref='vo',
plot=True,
verbose=False):
"""Parses datavariables according to their staggered grid position.
Should also work for gr variables, which are assumed to be on an A-grid."""
# extract reference dataarrays (those need to be in there)
tracer = data_dict[tracer_ref]
u = data_dict[u_ref]
v = data_dict[v_ref]
# determine grid type
if tracer.lon.equals(u.lon):
grid_type = 'A'
else:
if u.lon.equals(v.lon):
grid_type = 'B'
else:
grid_type = 'C'
print('Grid Type: %s detected' % grid_type)
if grid_type == 'A':
# this should also work with interpolated and obs datasets
# Just merge everything together
ds_combined = xr.merge([v for v in data_dict.values()])
ds_full = generate_grid_ds(ds_combined, {
"X": "x",
"Y": "y"
},
position=("center", "left"))
else:
# now determine the axis shift
lon = {}
lat = {}
lon['tracer'] = tracer.lon
lon['u'] = u.lon
lon['v'] = v.lon
lat['tracer'] = tracer.lat
lat['u'] = u.lat
lat['v'] = v.lat
# vizualize the position
if plot:
ref_idx = 3
plt.figure()
for vi, var in enumerate(['tracer', 'u', 'v']):
plt.plot(lon[var].isel(x=ref_idx, y=ref_idx),
lat[var].isel(x=ref_idx, y=ref_idx),
marker='*',
markersize=25)
plt.text(lon[var].isel(x=ref_idx, y=ref_idx),
lat[var].isel(x=ref_idx, y=ref_idx),
var,
ha='center',
va='center')
plt.title('Staggered Grid visualizaton')
plt.show()
if verbose:
print('Determine grid shift')
lon_diff = lon['tracer'] - lon['u']
# elinate large values due to boundry disc
lon_diff = lon_diff.where(abs(lon_diff) < 180)
lat_diff = lat['tracer'] - lat['v']
position = dict()
for axis, diff in zip(['X', 'Y'], [lon_diff, lat_diff]):
if np.sign(diff.mean().load()) < 0:
position[axis] = ('center', 'left')
else:
position[axis] = ('center', 'right')
if verbose:
print('Regenerate grid')
ds_full = generate_grid_ds(tracer, {
"X": "x",
"Y": "y"
},
position=position)
# now sort all other variables in accordingly
def rename(da,
da_tracer,
da_u,
da_v,
grid_type,
position,
verbose=False):
# check with which variable the lon and lat agree
rename_dict = {
'B': {
'u': {
'x': 'x_' + position['X'][1],
'y': 'y_' + position['Y'][1],
'lon': 'lon_ne',
'lat': 'lat_ne',
'vertices_latitude': 'vertices_latitude_ne',
'vertices_longitude': 'vertices_longitude_ne',
},
'v': {
'x': 'x_' + position['X'][1],
'y': 'y_' + position['Y'][1],
'lon': 'lon_ne',
'lat': 'lat_ne',
'vertices_latitude': 'vertices_latitude_ne',
'vertices_longitude': 'vertices_longitude_ne',
},
},
'C': {
'u': {
'x': 'x_' + position['X'][1],
'lon': 'lon_e',
'lat': 'lat_e',
'vertices_latitude': 'vertices_latitude_e',
'vertices_longitude': 'vertices_longitude_e',
},
'v': {
'y': 'y_' + position['Y'][1],
'lon': 'lon_n',
'lat': 'lat_n',
'vertices_latitude': 'vertices_latitude_n',
'vertices_longitude': 'vertices_longitude_n',
},
},
}
loc = []
for data, name in zip([da_tracer, da_u, da_v],
['tracer', 'u', 'v']):
if da.lon.equals(data.lon):
loc.append(name)
if len(loc) != 1:
if grid_type == 'B' and set(loc) == set(['u', 'v']):
loc = ['u']
else:
raise RuntimeError('somthing went wrong')
loc = loc[0]
if loc != 'tracer':
re_dict = {
k: v
for k, v in rename_dict[grid_type][loc].items()
if k in da.variables
}
da = da.rename(re_dict)
return da
if verbose:
print('Renaming and Merging')
for k, da in data_dict.items():
# print('Merging: %s' %k)
da_renamed = rename(da,
tracer,
u,
v,
grid_type,
position,
verbose=verbose)
# # parse all the coordinate values from the reconstructed dataset to the new dataarray
# or the merging will create intermediate steps
for co in da_renamed.coords:
if co in ds_full.coords:
da_renamed.coords[co] = ds_full.coords[co]
if verbose:
print('Merge')
# place all new data_variables into the dataset
for dvar in da_renamed.data_vars:
if dvar not in ds_full.data_vars:
ds_full[dvar] = da_renamed[dvar]
# ds_full = xr.merge([ds_full, da_renamed])
return ds_full
def merge_variables_on_staggered_grid(
data_dict,
modelname,
tracer_ref="thetao",
u_ref="uo",
v_ref="vo",
plot=True,
verbose=False,
):
"""Parses datavariables according to their staggered grid position.
Should also work for gr variables, which are assumed to be on an A-grid."""
if any([not a in data_dict.keys() for a in [tracer_ref, u_ref, v_ref]]):
print(
"NON-REFERENCE MODE. This should just be used for a bunch of variables on the same grid"
)
grid_type = "A"
else:
# extract reference dataarrays (those need to be in there)
tracer = data_dict[tracer_ref]
u = data_dict[u_ref]
v = data_dict[v_ref]
# determine grid type
if tracer.lon.equals(u.lon):
grid_type = "A"
else:
if u.lon.equals(v.lon):
grid_type = "B"
else:
grid_type = "C"
print("Grid Type: %s detected" % grid_type)
if grid_type == "A":
# this should also work with interpolated and obs datasets
# Just merge everything together
ds_combined = xr.merge([v for v in data_dict.values()])
ds_full = generate_grid_ds(ds_combined, {
"X": "x",
"Y": "y"
},
position=("center", "right"))
else:
# now determine the axis shift
lon = {}
lat = {}
lon["tracer"] = tracer.lon
lon["u"] = u.lon
lon["v"] = v.lon
lat["tracer"] = tracer.lat
lat["u"] = u.lat
lat["v"] = v.lat
# vizualize the position
if plot:
ref_idx = 3
plt.figure()
for vi, var in enumerate(["tracer", "u", "v"]):
plt.plot(
lon[var].isel(x=ref_idx, y=ref_idx),
lat[var].isel(x=ref_idx, y=ref_idx),
marker="*",
markersize=25,
)
plt.text(
lon[var].isel(x=ref_idx, y=ref_idx),
lat[var].isel(x=ref_idx, y=ref_idx),
var,
ha="center",
va="center",
)
plt.title("Staggered Grid visualizaton")
plt.show()
if verbose:
print("Determine grid shift")
lon_diff = lon["tracer"] - lon["u"]
# elinate large values due to boundry disc
lon_diff = lon_diff.where(abs(lon_diff) < 180)
lat_diff = lat["tracer"] - lat["v"]
position = dict()
for axis, diff in zip(["X", "Y"], [lon_diff, lat_diff]):
if np.sign(diff.mean().load()) < 0:
position[axis] = ("center", "left")
else:
position[axis] = ("center", "right")
if verbose:
print("Regenerate grid")
ds_full = generate_grid_ds(tracer, {
"X": "x",
"Y": "y"
},
position=position)
if verbose:
print("Renaming and Merging")
for k, da in data_dict.items():
# print('Merging: %s' %k)
da_renamed = rename(da,
tracer,
u,
v,
grid_type,
position,
verbose=verbose)
# parse all the coordinate values from the reconstructed
# dataset to the new dataarray
# or the merging will create intermediate steps
for co in da_renamed.coords:
if co in ds_full.coords:
da_renamed.coords[co] = ds_full.coords[co]
if verbose:
print("Merge")
# place all new data_variables into the dataset
for dvar in da_renamed.data_vars:
if dvar not in ds_full.data_vars:
ds_full[dvar] = da_renamed[dvar]
# ds_full = xr.merge([ds_full, da_renamed])
return ds_full
def add_latlon_metrics(dset, dims=None):
"""
Infer 2D metrics (latitude/longitude) from gridded data file.
Parameters
----------
dset : xarray.Dataset
A dataset open from a file
dims : dict
Dimension pair in a dict, e.g., {'lat':'latitude', 'lon':'longitude'}
Return
-------
dset : xarray.Dataset
Input dataset with appropriated metrics added
grid : xgcm.Grid
The grid with appropriated metrics
"""
lon, lat = None, None
if dims is None:
for dim in dimXList:
if dim in dset or dim in dset.coords:
lon = dim
break
for dim in dimYList:
if dim in dset or dim in dset.coords:
lat = dim
break
if lon is None or lat is None:
raise Exception('unknown dimension names in dset, should be in ' +
str(dimXList + dimYList))
else:
lon, lat = dims['lon'], dims['lat']
ds = generate_grid_ds(dset, {'X': lon, 'Y': lat})
coords = ds.coords
if __is_periodic(coords[lon], 360.0):
periodic = 'X'
else:
periodic = []
grid = Grid(ds, periodic=periodic)
na = np.nan
if 'X' in periodic:
dlonG = grid.diff(ds[lon], 'X', boundary_discontinuity=360)
dlonC = grid.diff(ds[lon + '_left'], 'X', boundary_discontinuity=360)
else:
dlonG = grid.diff(ds[lon], 'X', boundary='fill', fill_value=na)
dlonC = grid.diff(ds[lon + '_left'],
'X',
boundary='fill',
fill_value=na)
dlatG = grid.diff(ds[lat], 'Y', boundary='fill', fill_value=na)
dlatC = grid.diff(ds[lat + '_left'], 'Y', boundary='fill', fill_value=na)
coords['dxG'], coords['dyG'] = __dll_dist(dlonG, dlatG, ds[lon], ds[lat])
coords['dxC'], coords['dyC'] = __dll_dist(dlonC, dlatC, ds[lon], ds[lat])
coords['rAc'] = ds['dyC'] * ds['dxC']
metrics = {
('X', ): ['dxG', 'dxC'], # X distances
('Y', ): ['dyG', 'dyC'], # Y distances
('X', 'Y'): ['rAc']
}
grid._assign_metrics(metrics)
return ds, grid
def test_recreate_metrics(xshift, yshift, z_axis):
# reconstruct all the metrics by hand and compare to inferred output
# * For now this is a regular lon lat grid. Might need to add some tests for more complex grids.
# Then again. This will not do a great job for those....
# create test dataset
ds = _test_data(z_axis=z_axis)
# TODO: generalize so this also works with e.g. zonal average sections (which dont have a X axis)
coord_dict = {"X": "x", "Y": "y"}
if z_axis:
coord_dict["Z"] = "lev"
ds_full = generate_grid_ds(
ds, coord_dict, position={"X": ("center", xshift), "Y": ("center", yshift)},
)
grid = Grid(ds_full)
ds_metrics, metrics_dict = recreate_metrics(ds_full, grid)
if z_axis:
# Check that the bound values are intact (previously those got alterd due to unexpected behaviour of .assign_coords())
assert "bnds" in ds_metrics.lev_bounds.dims
# compute the more complex metrics (I could wrap this into a function I guess?)
lon0, lon1 = grid.axes["X"]._get_neighbor_data_pairs(ds.lon.load(), xshift)
lat0, lat1 = grid.axes["X"]._get_neighbor_data_pairs(ds.lat.load(), xshift)
dx_gx_expected = distance(lon0, lat0, lon1, lat1)
lon0, lon1 = grid.axes["Y"]._get_neighbor_data_pairs(ds.lon.load(), yshift)
lat0, lat1 = grid.axes["Y"]._get_neighbor_data_pairs(ds.lat.load(), yshift)
dy_gy_expected = distance(lon0, lat0, lon1, lat1)
# corner metrics
# dx
if yshift == "left":
# dx
lon0, lon1 = grid.axes["X"]._get_neighbor_data_pairs(
_interp_vertex_to_bounds(ds_metrics.lon_verticies, "y").isel(bnds=0),
xshift,
)
lat0, lat1 = grid.axes["X"]._get_neighbor_data_pairs(
ds_metrics.lat_bounds.isel(bnds=0), xshift
)
elif yshift == "right":
lon0, lon1 = grid.axes["X"]._get_neighbor_data_pairs(
_interp_vertex_to_bounds(ds_metrics.lon_verticies, "y").isel(bnds=1),
xshift,
)
lat0, lat1 = grid.axes["X"]._get_neighbor_data_pairs(
ds_metrics.lat_bounds.isel(bnds=1), xshift
)
dx_gxgy_expected = distance(lon0, lat0, lon1, lat1)
# dy
if xshift == "left":
# dx
lat0, lat1 = grid.axes["Y"]._get_neighbor_data_pairs(
_interp_vertex_to_bounds(ds_metrics.lat_verticies, "x").isel(bnds=0),
yshift,
)
lon0, lon1 = grid.axes["Y"]._get_neighbor_data_pairs(
ds_metrics.lon_bounds.isel(bnds=0), yshift
)
elif xshift == "right":
lat0, lat1 = grid.axes["Y"]._get_neighbor_data_pairs(
_interp_vertex_to_bounds(ds_metrics.lat_verticies, "x").isel(bnds=1),
yshift,
)
lon0, lon1 = grid.axes["Y"]._get_neighbor_data_pairs(
ds_metrics.lon_bounds.isel(bnds=1), yshift
)
dy_gxgy_expected = distance(lon0, lat0, lon1, lat1)
if xshift == "left":
vertex_points = [0, 1]
else:
vertex_points = [2, 3]
lon0, lon1 = (
ds_metrics.lon_verticies.isel(vertex=vertex_points[0]),
ds_metrics.lon_verticies.isel(vertex=vertex_points[1]),
)
lat0, lat1 = (
ds_metrics.lat_verticies.isel(vertex=vertex_points[0]),
ds_metrics.lat_verticies.isel(vertex=vertex_points[1]),
)
dy_gx_expected = distance(lon0, lat0, lon1, lat1)
if yshift == "left":
vertex_points = [0, 3]
else:
vertex_points = [1, 2]
lon0, lon1 = (
ds_metrics.lon_verticies.isel(vertex=vertex_points[0]),
ds_metrics.lon_verticies.isel(vertex=vertex_points[1]),
)
lat0, lat1 = (
ds_metrics.lat_verticies.isel(vertex=vertex_points[0]),
ds_metrics.lat_verticies.isel(vertex=vertex_points[1]),
)
dx_gy_expected = distance(lon0, lat0, lon1, lat1)
if z_axis:
dz_t_expected = ds.lev_bounds.diff("bnds").squeeze().data
else:
dz_t_expected = None
for var, expected in [
("dz_t", dz_t_expected),
(
"dx_t",
distance(
ds_metrics.lon_bounds.isel(bnds=0).data,
ds_metrics.lat.data,
ds_metrics.lon_bounds.isel(bnds=1).data,
ds_metrics.lat.data,
),
),
(
"dy_t",
distance(
ds_metrics.lon.data,
ds_metrics.lat_bounds.isel(bnds=0).data,
ds_metrics.lon.data,
ds_metrics.lat_bounds.isel(bnds=1).data,
),
),
("dx_gx", dx_gx_expected),
("dy_gy", dy_gy_expected),
("dy_gx", dy_gx_expected),
("dx_gy", dx_gy_expected),
("dy_gxgy", dy_gxgy_expected),
("dx_gxgy", dx_gxgy_expected),
]:
if expected is not None:
print(var)
control = ds_metrics[var].data
if expected.shape != control.shape:
control = control.T
np.testing.assert_allclose(control, expected)
if z_axis:
assert set(["X", "Y", "Z"]).issubset(set(metrics_dict.keys()))
else:
assert set(["X", "Y"]).issubset(set(metrics_dict.keys()))
assert not "Z" in list(metrics_dict.keys())
def load_wrf(date_in, wrf_nc_dir):
'''Load wrf data
Because loading large files could cost much time,
please use "frames_per_outfile=1" in "namelist.input"
to generate wrfout* files.
'''
# get all wrf output files in the same day
f_wrf_pattern = os.path.join(wrf_nc_dir,
date_in.strftime('wrfout_*_%Y-%m-%d_*'))
wrf_list = glob.glob(f_wrf_pattern)
# omit the directory and get "yyyy-mm-dd_hh:mm:ss"
wrf_dates = [
datetime.strptime(
ntpath.basename(name).split('_', maxsplit=2)[-1],
'%Y-%m-%d_%H:%M:%S') for name in wrf_list
]
# get the index of the closest datetime
wrf_index = wrf_dates.index(min(wrf_dates, key=lambda d: abs(d - date_in)))
wrf_file = wrf_list[wrf_index]
# read selected wrf data
logging.info(' ' * 4 + f'Reading {wrf_file} ...')
wrf = xr.open_dataset(wrf_file)
wrf.attrs['wrfchem_filename'] = os.path.basename(wrf_file)
# generate lon_bounds and lat_bounds
wrf = wrf.rename({'XLONG': 'lon', 'XLAT': 'lat'}).isel(Time=0)
wrf = generate_grid_ds(wrf, {
'X': 'west_east',
'Y': 'south_north'
},
position=('center', 'outer'))
# convert from potential temperature to absolute temperature (K)
# http://mailman.ucar.edu/pipermail/wrf-users/2010/001896.html
# http://mailman.ucar.edu/pipermail/wrf-users/2013/003117.html
wrf['T'] = (wrf['T'] + 300) * ((wrf['P'] + wrf['PB']) / 1e5)**0.2865
# generate grid and bounds
grid_ds = Grid(wrf, periodic=False)
bnd = 'extrapolate'
wrf.coords['lon_b'] = grid_ds.interp(grid_ds.interp(wrf['lon'],
'X',
boundary=bnd,
fill_value=np.nan),
'Y',
boundary=bnd,
fill_value=np.nan)
wrf.coords['lat_b'] = grid_ds.interp(grid_ds.interp(wrf['lat'],
'X',
boundary=bnd,
fill_value=np.nan),
'Y',
boundary=bnd,
fill_value=np.nan)
logging.info(' ' * 8 + 'Finish reading')
return wrf_file, wrf