def grid_calculations(grid: xgcm.Grid, ds_full: xr.core.dataset.Dataset):
# Compute difference (in degrees) along longitude and latitude for
# both cell center and left
dlong = grid.diff(ds_full.xt_ocean, 'X', boundary_discontinuity=360)
dlonc = grid.diff(ds_full.xt_ocean_left, 'X', boundary_discontinuity=360)
dlatg = grid.diff(ds_full.yt_ocean,
'Y',
boundary='fill',
fill_value=np.nan)
dlatc = grid.diff(ds_full.yt_ocean_left,
'Y',
boundary='fill',
fill_value=np.nan)
# Convert degrees to actual Cartesian distances on the Earth
# add distances to coordinates in data
ds_full.coords['dxg'], ds_full.coords['dyg'] = dll_dist(
dlong, dlatg, ds_full.xt_ocean, ds_full.yt_ocean)
ds_full.coords['dxc'], ds_full.coords['dyc'] = dll_dist(
dlonc, dlatc, ds_full.xt_ocean, ds_full.yt_ocean)
# Calculate area of each gridcell
ds_full.coords['area_c'] = ds_full.dxc * ds_full.dyc
# Fill nan values
dyg = ds_full.dyg.fillna(111000)
dyc = ds_full.dyc.fillna(111000)
ds_full.coords['dyg'] = dyg
ds_full.coords['dyc'] = dyc
def split_adv_budget(ds):
print('I think this is outdated....')
ds = ds.copy()
if 'o2_xflux_adv' in list(ds.data_vars):
grid = Grid(ds)
area = ds.area_t
div_x = -grid.diff(ds.o2_xflux_adv, 'X', boundary='fill') / area
div_y = -grid.diff(ds.o2_yflux_adv, 'Y', boundary='fill') / area
div_z = grid.diff(ds.o2_zflux_adv, 'Z', boundary='fill') / area
for data, name in zip([div_x, div_y, div_z],
['o2_advection_%s' % a
for a in ['x', 'y', 'z']]):
ds[name] = data
return ds
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 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 test_diff_c_to_g_periodic(periodic_1d):
ds = periodic_1d
# a linear gradient in the ni direction
data_c = np.sin(ds['XC'])
data_expected = data_c.values - np.roll(data_c.values, 1)
grid = Grid(ds)
data_g = grid.diff(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.diff(data_c, 'X')
np.testing.assert_allclose(data_g.values, data_expected)
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_diff_c_to_g_nonperiodic(nonperiodic_1d):
ds = nonperiodic_1d
# a linear gradient in the ni direction
grad = 0.24
data_c = grad * ds['ni']
data_expected = data_c.values[1:] - data_c.values[:-1]
grid = Grid(ds, x_periodic=False)
data_u = grid.diff(data_c, 'X')
# check that the dimensions are right
assert data_u.dims == ('ni_u', )
xr.testing.assert_equal(data_u.ni_u, ds.ni_u[1:-1])
assert len(data_u.ni_u) == len(data_u)
# check that the values are right
np.testing.assert_allclose(data_u.values, data_expected)
np.testing.assert_allclose(data_u.values, grad)
def test_diff_g_to_c_periodic(periodic_1d):
ds = periodic_1d
# a linear gradient in the ni direction
data_g = np.sin(ds['XG'])
# np.roll(np.arange(5), -1) --> [1, 2, 3, 4, 0]
# negative roll shifts right
data_expected = np.roll(data_g.values, -1) - data_g.values
#data_expected = np.cos(ds['XC']).values * (2*np.pi) / 100.
grid = Grid(ds)
data_c = grid.diff(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_diff_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 = data_u.values[1:] - data_u.values[:-1]
grid = Grid(ds, x_periodic=False)
data_c = grid.diff(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)
np.testing.assert_allclose(data_c.values, grad)
def calc_curl_stress(
ds,
ds_static,
varx="tauuo",
vary="tauvo",
areacello_bu="areacello_bu",
xdim="lon",
ydim="lat",
rho0=1035.0,
maskname="wet_c",
lonname="geolon_c",
latname="geolat_c",
):
"""Calculate curl of stress acting on surface of the ocean.
Parameters
----------
ds : xarray.Dataset dataset with tauuo and tauvo
ds_static : xarray.Dataset with grid values
varname : str, optional
Name of the tauuo and tauvo variables, by default "tauuo" and "tauvo"
area : str, optional
Name of the area variable, by default "areacello_bu"
xdim : str, optional
Name of the longitude coordinate, by default "lon"
ydim : str, optional
Name of the latitude coordinate, by default "lat"
rho0: float, optional
Reference density of seawater, by default 1035.0 kg/m3
maskname: str, optional
Name of land/sea mask, by default wet_c
lonname: str, optional
Name of longitude variable, by default geolon_c
latname: str, optional
Name of latitude variable, by default geolat_c
Returns
-------
xarray.DataArray stress_curl
curl of surface ocean stress
"""
area = ds_static[areacello_bu]
taux = ds[varx].mean(dim="time")
tauy = ds[vary].mean(dim="time")
# fill nan with 0.0 since want 0.0 values over land for the curl operation
taux = taux.fillna(0.0)
tauy = tauy.fillna(0.0)
grid = Grid(
ds_static,
coords={
"X": {
"center": "xh",
"outer": "xq"
},
"Y": {
"center": "yh",
"outer": "yq"
},
},
periodic=["X"],
)
stress_curl = -grid.diff(taux * ds_static.dxCu, "Y",
boundary="fill") + grid.diff(
tauy * ds_static.dyCv, "X", boundary="fill")
stress_curl = stress_curl / (area * rho0)
stress_curl = stress_curl.where(ds_static[maskname] == 1)
stress_curl = stress_curl.assign_coords({
lonname: ds_static[lonname],
latname: ds_static[latname]
})
return stress_curl
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 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 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
'X': {
'center': 'i',
'right': 'i_g'
},
'Y': {
'center': 'j',
'right': 'j_g'
}
})
#ds1 = model.get_dataset(varnames=['W'], type='latlon')
#grid1 = Grid(ds1, coords={'X': {'center': 'i', 'right': 'i_g'}, 'Y': {'center': 'j', 'right': 'j_g'}, 'Z': {'center': 'k', 'outer': 'k_l'}})
#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',