def test_interp_c_to_g_periodic(periodic_1d):
"""Interpolate from c grid to g grid."""
ds = periodic_1d
data_c = np.sin(ds['XC'])
# np.roll(np.arange(5), 1) --> [4, 0, 1, 2, 3]
# positive roll shifts left
data_expected = 0.5 * (data_c.values + np.roll(data_c.values, 1))
grid = Grid(ds)
data_g = grid.interp(data_c, 'X')
# check that the dimensions are right
assert data_g.dims == ('XG', )
xr.testing.assert_equal(data_g.XG, ds.XG)
assert len(data_g.XG) == len(data_g)
# check that the values are right
np.testing.assert_allclose(data_g.values, data_expected)
# try the same with chunks
data_c = np.sin(ds['XC'])
data_c = data_c.chunk(10)
data_g = grid.interp(data_c, 'X')
np.testing.assert_allclose(data_g.values, data_expected)
def calculate_momentum_budget(ds):
grid = Grid(ds)
combo = xr.Dataset()
combo["u"] = ds.u
combo["v"] = ds.v
combo["du_dx"] = grid.diff(ds.u, "X") / ds.dxtn
combo["du_dy"] = grid.diff(ds.u, "Y") / ds.dyte
combo["u_du_dx"] = grid.interp(-combo["du_dx"] * grid.interp(ds.u, "X"),
"Y")
combo["v_du_dy"] = grid.interp(-combo["du_dy"] * grid.interp(ds.v, "Y"),
"X")
combo["hor"] = combo["u_du_dx"] + combo["v_du_dy"]
combo["hor"].attrs[
"long_name"] = "Zonal Velocity tendency due to hor divergence of momentum"
combo["hor"].attrs["units"] = "m/s^(-2)"
# Add tracer and vertical vel in there to get all relavant. Then drop again
combo["wt"] = ds.wt # for now just to include 'sw_ocean'
combo["temp"] = ds.temp
combo = combo.drop(["wt", "temp"])
return combo
def add_MITgcm_missing_metrics(dset, periodic=None):
"""
Infer missing metrics from MITgcm output files.
Parameters
----------
dset : xarray.Dataset
A dataset open from a file
Return
-------
dset : xarray.Dataset
Input dataset with appropriated metrics added
grid : xgcm.Grid
The grid with appropriated metrics
"""
coords = dset.coords
grid = Grid(dset, periodic=periodic)
BCx = 'periodic' if grid.axes['X']._periodic else 'fill'
BCy = 'periodic' if grid.axes['Y']._periodic else 'fill'
if 'drW' not in coords: # vertical cell size at u point
coords['drW'] = dset.hFacW * dset.drF
if 'drS' not in coords: # vertical cell size at v point
coords['drS'] = dset.hFacS * dset.drF
if 'drC' not in coords: # vertical cell size at tracer point
coords['drC'] = dset.hFacC * dset.drF
if 'dxF' not in coords:
coords['dxF'] = grid.interp(dset.dxC, 'X', boundary=BCx)
if 'dyF' not in coords:
coords['dyF'] = grid.interp(dset.dyC, 'Y', boundary=BCy)
if 'dxV' not in coords:
coords['dxV'] = grid.interp(dset.dxG, 'X', boundary=BCx)
if 'dyU' not in coords:
coords['dyU'] = grid.interp(dset.dyG, 'Y', boundary=BCy)
if 'hFacZ' not in coords:
coords['hFacZ'] = grid.interp(dset.hFacS, 'X', boundary=BCx)
if 'maskZ' not in coords:
coords['maskZ'] = coords['hFacZ']
# Calculate vertical distances located on the cellboundary
# ds.coords['dzC'] = grid.diff(ds.depth, 'Z', boundary='extrapolate')
# Calculate vertical distances located on the cellcenter
# ds.coords['dzT'] = grid.diff(ds.depth_left, 'Z', boundary='extrapolate')
metrics = {
('X', ): ['dxG', 'dxF', 'dxC', 'dxV'], # X distances
('Y', ): ['dyG', 'dyF', 'dyC', 'dyU'], # Y distances
('Z', ): ['drW', 'drS', 'drC', 'drF'], # Z distances
('X', 'Y'): ['rAw', 'rAs', 'rA', 'rAz']
} # Areas
grid._assign_metrics(metrics)
return dset, grid
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 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 interpolated_velocities(u, v):
"""Interpolate the staggered horizontal velocities to the cell centers
Use's Ryan Abernathy's xgcm package.
Parameters
----------
u, v : xr.DataArray
staggered velocities.
Returns
-------
uint, vint : xr.DataArray
interpolated velocites on (xc, yc) grid
"""
# setup grid
xl = u.x.values
xc = xl + (xl[1] - xl[0]) / 2
yl = u.y.values
yc = yl + (yl[1] - yl[0]) / 2
u = u.rename({'x': 'xl', 'y': 'yc'}).assign_coords(xl=xl, yc=yc)
v = v.rename({'x': 'xc', 'y': 'yl'}).assign_coords(yl=yl, xc=xc)
# combine the data
ds = xr.Dataset({'u': u, 'v': v})
# create grid object
coords = {
'x': {
'center': 'xc',
'left': 'xl'
},
'y': {
'center': 'yc',
'left': 'yl'
}
}
grid = Grid(ds, periodic=['x'], coords=coords)
# use grid to interpolate
uint = grid.interp(ds.u, axis='x')
vint = grid.interp(ds.v, axis='y', boundary='extend')
return uint, vint
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 test_calculate_rel_vorticity():
datadict = datasets()
coords = datadict["coords"]
ds_b = datadict["B"]
grid_b = Grid(ds_b, coords=coords)
ds_c = datadict["C"]
grid_c = Grid(ds_c, coords=coords)
test_b = (grid_b.diff(grid_b.interp(ds_b.v * ds_b.dy_ne, "Y"), "X") -
grid_b.diff(grid_b.interp(ds_b.u * ds_b.dx_ne, "X"),
"Y")) / ds_b.area_t
zeta_b = calculate_rel_vorticity(
grid_b,
ds_b.u,
ds_b.v,
ds_b.dx_ne,
ds_b.dy_ne,
ds_b.area_t,
gridtype=None,
)
test_c = (grid_c.diff(ds_c.v * ds_c.dy_n, "X") -
grid_c.diff(ds_c.u * ds_c.dx_e, "Y")) / ds_c.area_ne
zeta_c = calculate_rel_vorticity(
grid_c,
ds_c.u,
ds_c.v,
ds_c.dx_e,
ds_c.dy_n,
ds_c.area_ne,
gridtype=None,
)
assert_allclose(test_b, zeta_b)
assert_allclose(test_c, zeta_c)
with pytest.raises(RuntimeError):
zeta_c = calculate_rel_vorticity(
grid_b,
ds_c.u,
ds_c.v,
ds_c.dx_n, # wrong coordinate
ds_c.dy_n,
ds_c.area_ne,
gridtype=None,
)
def test_interp_g_to_c_periodic(periodic_1d):
"""Interpolate from c grid to g grid."""
ds = periodic_1d
# a linear gradient in the ni direction
data_g = np.sin(ds['XG'])
data_expected = 0.5 * (data_g.values + np.roll(data_g.values, -1))
grid = Grid(ds)
data_c = grid.interp(data_g, 'X')
# check that the dimensions are right
assert data_c.dims == ('XC', )
xr.testing.assert_equal(data_c.XC, ds.XC)
assert len(data_c.XC) == len(data_c)
# check that the values are right
np.testing.assert_allclose(data_c.values, data_expected)
def test_interp_g_to_c_nonperiodic(nonperiodic_1d):
"""Interpolate from g grid to c grid."""
ds = nonperiodic_1d
# a linear gradient in the ni direction
grad = 0.43
data_u = grad * ds['ni_u']
data_expected = 0.5 * (data_u.values[1:] + data_u.values[:-1])
grid = Grid(ds, x_periodic=False)
data_c = grid.interp(data_u, 'X')
# check that the dimensions are right
assert data_c.dims == ('ni', )
xr.testing.assert_equal(data_c.ni, ds.ni)
assert len(data_c.ni) == len(data_c)
# check that the values are right
np.testing.assert_allclose(data_c.values, data_expected)
def add_split_tendencies(ds):
"""Reconstructs various fluxes and tendencies (x-y_x) from the monthly averaged output. Hardcoded to o2 atm."""
ds = ds.copy()
rho = 1035 #reference density
grid = Grid(ds)
# This should be calculated upstream (see: add_vertical_spacing), it is possibly totally wrong
if 'dzwt' not in ds.coords:
print(
'Vertical spacing for vertical vel cell is approximated!!! Use with caution'
)
# ds.coords['dzwt'] = grid.interp(ds['dzt'], 'Z')
# This avoids computation. Somehow things are triggered when these fields are coordinates.
ds['dzwt'] = grid.interp(ds['dzt'], 'Z')
# These should be less critical, but should be given nontheless
if 'dxte' not in ds.coords:
print('Spacing for `dxte` is approximated!!! Use with caution')
ds.coords['dxte'] = grid.interp(ds['dxt'], 'X')
if 'dytn' not in ds.coords:
print('Spacing for `dytn` is approximated!!! Use with caution')
ds.coords['dytn'] = grid.interp(ds['dyt'], 'Y')
# Calculate thickness weighted mass transports from velocities according to MOM5 formulation
uhrho_et, vhrho_nt, wrhot = reconstruct_hrho_trans(ds.u, ds.v, ds.wt,
ds.dzt * rho,
ds.dzu * rho, grid, rho)
# Reconstruct the flux terms
ds['o2_xflux_adv_recon'], ds['o2_yflux_adv_recon'], ds['o2_zflux_adv_recon'] = \
approximate_transport_op(uhrho_et, vhrho_nt, wrhot, ds['o2'], grid, boundary='extend')
# Calculate tendencies from all advective fluxes
for suffix in ['', '_recon']:
ds['o2_advection_x'+suffix], \
ds['o2_advection_y'+suffix], \
ds['o2_advection_z'+suffix] = tend_from_fluxes(ds['o2_xflux_adv'+suffix],
ds['o2_yflux_adv'+suffix],
ds['o2_zflux_adv'+suffix], grid)
# Reconstruct tendency terms seperately for changes in tracer and changes in transport
udiv, vdiv, wdiv = tend_from_fluxes(uhrho_et * grid._ds.dyte,
vhrho_nt * grid._ds.dxtn,
wrhot * grid._ds.area_t, grid)
ds['o2_advection_x_recon_vel'] = ds['o2'] * udiv
ds['o2_advection_y_recon_vel'] = ds['o2'] * vdiv
ds['o2_advection_z_recon_vel'] = ds['o2'] * wdiv
ds['o2_advection_x_recon_tracer'] = -grid.interp(
grid.diff(ds['o2'], 'X') / grid._ds.dxte * uhrho_et, 'X')
ds['o2_advection_y_recon_tracer'] = -grid.interp(
grid.diff(ds['o2'], 'Y') / grid._ds.dytn * vhrho_nt, 'Y')
ds['o2_advection_z_recon_tracer'] = -grid.interp(
grid.diff(ds['o2'], 'Z') / grid._ds.dzwt * wrhot, 'Z')
# # some interpolations... might want to remap these in the future, to get more accurate estimates
# u = grid.interp(ds['u'], 'Y')
# v = grid.interp(ds['v'], 'X')
# wt = ds['wt']
# # dxte = grid.interp(ds.dxu, 'Y')
# # dytn = grid.interp(ds.dyu, 'X')
# # o2_flux reconstructed from tracer and velocity field (vel*tracer*dyt*dzt*rho)
# # Are the values actually interpolated or do they take the center tracer value? Read up in MOM5 manual and correct if needed.
# # This will need some more advanced testing...in the end we cannot really reproduce the complex advection scheme, but it is worth trying
# # # to get as close as possible.
# # ds['o2_xflux_adv_recon'] = grid.interp(ds['o2'], 'X') * u * ds.dzt * ds.dyte * rho
# # ds['o2_yflux_adv_recon'] = grid.interp(ds['o2'], 'Y') * v * ds.dzt * ds.dxtn * rho
# # ds['o2_zflux_adv_recon'] = grid.interp(ds['o2'], 'Z') * wt * ds.dxt * ds.dyt * rho
# # Reconstruct the advective tendencies (as (tracer* s^-1) * dzt * rho)
# # also not sure about the numerics here...this implements finite difference approach...which mom used for some variables but not all...
# ds['o2_advection_x_recon_full'] = - (grid.diff(ds['o2_xflux_adv_recon'], 'X') / ds.area_t)
# ds['o2_advection_x_recon_du'] = - (ds['o2'] * grid.diff(u, 'X') / ds.dxt * ds.dzt * rho)
# ds['o2_advection_x_recon_do2'] = - (u * grid.diff(ds['o2'], 'X') / dxte * ds.dzt * rho)
# ds['o2_advection_y_recon_full'] = - (grid.diff(ds['o2_yflux_adv_recon'], 'Y') / ds.area_t)
# ds['o2_advection_y_recon_dv'] = - (ds['o2'] * grid.diff(v, 'Y') / ds.dyt * ds.dzt * rho)
# ds['o2_advection_y_recon_do2'] = - (v * grid.diff(ds['o2'], 'Y') / dytn * ds.dzt * rho)
# ds['o2_advection_z_recon_full'] = (grid.diff(ds['o2_zflux_adv_recon'], 'Z') / ds.area_t)
# ds['o2_advection_z_recon_dwt'] = (ds['o2'] * grid.diff(wt, 'Z') * rho)
# ds['o2_advection_z_recon_do2'] = (wt * grid.diff(ds['o2'], 'Z') * rho)
return ds
def read_topography(dictArgs, coords=None, point_type="t"):
"""Returns topography field based on the values of the CLI
arguments and the presence of intake catalogs.
Parameters
----------
dictArgs : dict, optional
dictionary of arguments obtained from the CLI parser, by default None
coords : tuple, optional
target xarray coordinates
point_type : str, optional
Requested grid type of t|q|u|v, by default "t"
Returns
-------
xarray.DataArray
topography array
"""
point_type = point_type.upper()
verbose = dictArgs["verbose"] if "verbose" in dictArgs else False
if dictArgs["topog"] is not None:
if verbose:
print("Using optional topg file for depth field")
ds = xr.open_dataset(dictArgs["topog"])
elif dictArgs["static"] is not None:
if verbose:
print("Using optional static file for depth field.")
ds = xr.open_dataset(dictArgs["static"])
elif dictArgs["gridspec"] is not None:
if verbose:
print("Using optional gridspec tar file for depth field.")
tar = tf.open(dictArgs["gridspec"])
ds = extract_from_tar(tar, "ocean_topog.nc")
elif dictArgs["platform"] is not None and dictArgs["config"] is not None:
if verbose:
print(
f"Using {dictArgs['platform']} {dictArgs['config']} intake catalog for depth field."
)
cat = open_intake_catalog(dictArgs["platform"], dictArgs["config"])
ds = cat["topog"].to_dask()
if "deptho" in list(ds.variables):
depth = ds.deptho
elif "depth" in list(ds.variables):
depth = ds.depth
coords = xr.Dataset(coords)
xedge = "outer" if (len(coords.xq) == len(coords.xh) + 1) else "right"
yedge = "outer" if (len(coords.yq) == len(coords.yh) + 1) else "right"
out_grid = Grid(
coords,
coords={
"X": {
"center": "xh",
xedge: "xq"
},
"Y": {
"center": "yh",
yedge: "yq"
}
},
periodic=["X"],
)
if point_type == "V":
depth = out_grid.interp(depth, "Y", boundary="fill")
elif point_type == "U":
depth = out_grid.interp(depth, "X", boundary="fill")
return depth
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
def get_roms(path):
ds = xr.open_mfdataset(path,
concat_dim='ocean_time',
combine='nested',
data_vars='minimal',
coords='minimal',
parallel=True,
chunks={
'ocean_time': 1,
})
ds = ds.rename({
'eta_u': 'eta_rho',
'xi_v': 'xi_rho',
'xi_psi': 'xi_u',
'eta_psi': 'eta_v'
})
coords = {
'X': {
'center': 'xi_rho',
'inner': 'xi_u'
},
'Y': {
'center': 'eta_rho',
'inner': 'eta_v'
},
'Z': {
'center': 's_rho',
'outer': 's_w'
}
}
grid = Grid(ds, coords=coords, periodic=[])
Zo_rho = (ds.hc * ds.s_rho + ds.Cs_r * ds.h) / (ds.hc + ds.h)
z_rho = Zo_rho * (ds.zeta + ds.h) + ds.zeta
Zo_w = (ds.hc * ds.s_w + ds.Cs_w * ds.h) / (ds.hc + ds.h)
z_w = Zo_w * (ds.zeta + ds.h) + ds.zeta
ds.coords['z_w'] = z_w.where(ds.mask_rho,
0).transpose('ocean_time', 's_w', 'eta_rho',
'xi_rho')
ds.coords['z_rho'] = z_rho.where(ds.mask_rho,
0).transpose('ocean_time', 's_rho',
'eta_rho', 'xi_rho')
ds['pm_v'] = grid.interp(ds.pm, 'Y')
ds['pn_u'] = grid.interp(ds.pn, 'X')
ds['pm_u'] = grid.interp(ds.pm, 'X')
ds['pn_v'] = grid.interp(ds.pn, 'Y')
ds['pm_psi'] = grid.interp(grid.interp(ds.pm, 'Y'),
'X') # at psi points (eta_v, xi_u)
ds['pn_psi'] = grid.interp(grid.interp(ds.pn, 'X'),
'Y') # at psi points (eta_v, xi_u)
ds['dx'] = 1 / ds.pm
ds['dx_u'] = 1 / ds.pm_u
ds['dx_v'] = 1 / ds.pm_v
ds['dx_psi'] = 1 / ds.pm_psi
ds['dy'] = 1 / ds.pn
ds['dy_u'] = 1 / ds.pn_u
ds['dy_v'] = 1 / ds.pn_v
ds['dy_psi'] = 1 / ds.pn_psi
ds['dz'] = grid.diff(ds.z_w, 'Z', boundary='fill')
ds['dz_w'] = grid.diff(ds.z_rho, 'Z', boundary='fill')
ds['dz_u'] = grid.interp(ds.dz, 'X')
ds['dz_w_u'] = grid.interp(ds.dz_w, 'X')
ds['dz_v'] = grid.interp(ds.dz, 'Y')
ds['dz_w_v'] = grid.interp(ds.dz_w, 'Y')
ds['dA'] = ds.dx * ds.dy
metrics = {
('X', ): ['dx', 'dx_u', 'dx_v', 'dx_psi'], # X distances
('Y', ): ['dy', 'dy_u', 'dy_v', 'dy_psi'], # Y distances
('Z', ): ['dz', 'dz_u', 'dz_v', 'dz_w', 'dz_w_u',
'dz_w_v'], # Z distances
('X', 'Y'): ['dA'] # Areas
}
grid = Grid(ds, coords=coords, metrics=metrics, periodic=[])
return ds, grid
#ds1 = xr.merge([ds1, dsgrid])
# Compute vorticity and divergence
vorticity = (-grid.diff(ds.U * ds.dxC, 'Y', boundary='fill') +
grid.diff(ds.V * ds.dyC, 'X', boundary='fill')) / ds.rAz
divergence = (grid.diff(ds.U * ds.dxC, 'X', boundary='fill') +
grid.diff(ds.V * ds.dyC, 'Y', boundary='fill')) / ds.rA
# Select data for a small region
lat1, lat2 = (29., 41.)
# 10x10 deg box near gulf stream
lon1, lon2 = (-121., -109.)
# 1 deg should be avoided on each side for the final data
u1 = grid.interp(ds.U, 'X', boundary='fill')
v1 = grid.interp(ds.V, 'Y', boundary='fill')
vort = grid.interp(grid.interp(vorticity, 'X', boundary='fill'),
'Y',
boundary='fill')
#w1 = grid1.interp(ds1.W, 'Z', boundary='fill')
#w1 = w1.sel({k = [1,3,5,10,30]})
#eta = ds.Eta
#eta = eta.where((ds.YC > lat1) & (ds.YC < lat2) & (ds.XC > lon1) & (ds.XC < lon2), drop = True)
u1 = u1.where(
(ds.YC > lat1) & (ds.YC < lat2) & (ds.XC > lon1) & (ds.XC < lon2),
drop=True)
v1 = v1.where(
(ds.YC > lat1) & (ds.YC < lat2) & (ds.XC > lon1) & (ds.XC < lon2),
drop=True)
