def test_padding_mixed(self, boundary_width):
ds, coords, _ = datasets_grid_metric("C")
grid = Grid(ds, coords=coords)
data = ds.tracer
axis_padding_mapping = {"X": "periodic", "Y": "extend"}
method_mapping = {
"periodic": "wrap",
"extend": "edge",
}
# iterate over all axes
expected = data.copy(deep=True)
for ax, widths in boundary_width.items():
dim = grid._get_dims_from_axis(
data, ax)[0] # ? not entirely sure why this is a list
expected = expected.pad({dim: widths},
method_mapping[axis_padding_mapping[ax]])
# we intentionally strip all coords from padded arrays
expected = _strip_all_coords(expected)
result = pad(
data,
grid,
boundary=axis_padding_mapping,
boundary_width=boundary_width,
fill_value=None,
other_component=None,
)
xr.testing.assert_allclose(expected, result)
def test_padding_fill(self, boundary_width, fill_value):
ds, coords, _ = datasets_grid_metric("C")
grid = Grid(ds, coords=coords)
data = ds.tracer
# iterate over all axes
expected = data.copy(deep=True)
for ax, widths in boundary_width.items():
dim = grid._get_dims_from_axis(
data, ax)[0] # ? not entirely sure why this is a list
expected = expected.pad({dim: widths},
"constant",
constant_values=fill_value)
# we intentionally strip all coords from padded arrays
expected = _strip_all_coords(expected)
result = pad(
data,
grid,
boundary="fill",
boundary_width=boundary_width,
fill_value=fill_value,
other_component=None,
)
xr.testing.assert_allclose(expected, result)
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 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 test_interp_all():
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)
for var in ["u", "v", "tracer"]:
for ds, grid in zip([ds_b, ds_c], [grid_b, grid_c]):
for target, control_dims in zip(
["center", "right"],
[["xt", "yt", "time"], ["xu", "yu", "time"]],
):
print(ds)
print(grid)
ds_interp = interp_all(grid, ds, target=target)
assert set(ds_interp[var].dims) == set(control_dims)
assert set(ds_interp.coords) == set(ds.coords)
ds_interp_nocoords = interp_all(grid,
ds,
target=target,
keep_coords=False)
assert set(ds_interp_nocoords.coords) != set(ds.coords)
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 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 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 add_vertical_spacing(ds):
grid = Grid(ds)
ds.coords['dst'] = calculate_ds(ds, dim='st')
ds.coords['dswt'] = calculate_ds(ds, dim='sw')
ds.coords['dzt'] = calculate_dz(ds['eta_t'], ds['ht'], ds['dst'])
ds.coords['dzu'] = grid.min(grid.min(ds['dzt'], 'X'), 'Y')
# # Avoids suspected computation midway during the dask graph creation...Should probably raise an xarray issue about this.
# ds['dzt'] = calculate_dz(ds['eta_t'], ds['ht'], ds['dst'])
# ds['dzu'] = grid.min(grid.min(ds['dzt'], 'X'),'Y')
# # Lets try this as a workaround...
# for vv in ['dzt', 'dzu']:
# ds.coords[vv] = ds[vv]
# the dzwt value is dependent on the model version (finite vol vs. engergetically consistent?; See MOM5 manual section 10.4.2)
return ds
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 define_grid(ds, names):
"""build a xgcm.Grid object
Args:
ds (xarray.Dataset): dataset with model's grid
names (dict): dictionary containing dimensions/coordinates names
Returns:
xgcm.Grid: grid object
"""
x_center, y_center = names["x_center"], names["y_center"]
x_corner, y_corner = names["x_corner"], names["y_corner"]
qcoord = "outer" if is_symetric(ds, names) else "right"
grid = Grid(
ds,
coords={
"X": {"center": x_center, qcoord: x_corner},
"Y": {"center": y_center, qcoord: y_corner},
},
periodic=["X"],
)
return grid
def full_preprocessing(dat_dict, modelname,
tracer_ref="thetao",
u_ref="uo",
v_ref="vo",
plot=True,
verbose=False):
"""Fully preprocess data for one model ensemble .
The input needs to be a dictionary in the form:
{'<source_id>':{'<varname_a>':'<uri_a>', '<varname_b>':'<uri_b>', ...}}
"""
renaming_dict = cmip6_renaming_dict()
# homogenize the naming
dat_dict = {
var: cmip6_homogenization(data, renaming_dict[modelname])
for var, data in dat_dict.items()
}
# broadcast lon and lat values if they are 1d
if renaming_dict[modelname]["lon"] is None:
dat_dict = {var: broadcast_lonlat(data) for var, data in dat_dict.items()}
# merge all variables together on the correct staggered grid
ds = merge_variables_on_staggered_grid(
dat_dict, modelname, u_ref=u_ref, v_ref=v_ref,plot=plot, verbose=verbose
)
grid_temp = Grid(ds)
ds = recreate_metrics(ds, grid_temp)
return ds
def test_grid_repr(all_datasets):
grid = Grid(all_datasets)
print(grid)
r = repr(grid).split('\n')
assert r[0] == "<xgcm.Grid>"
# all datasets should have at least an X axis
assert r[1].startswith('X-axis:')
def lat_slice(ds, ref, model_processing=True):
ds = ds.copy()
if model_processing:
grid = Grid(ds)
ds = interp_all(grid, ds)
ds = ds.interp(xt_ocean=ref.xt_ocean)
return ds
def eq_mean(
ds,
roi=dict(yt_ocean=slice(-1, 1)),
model_processing=True,
coord_include=["dzt"],
):
"""Calculate the equatorial slices, defined by the weighted mean around the equator.
`model_processing` will add budget decompositions and interpolate all fields onto tracer grid."""
ds = ds.copy()
for co in coord_include:
if co in list(ds.variables):
ds[co + "_temp"] = ds[co]
ds = ds.drop(co)
# This works but seems clunky that I have to go through these lengths to assign any variable as a data_variable
if model_processing:
grid = Grid(ds)
ds = interp_all(grid, ds)
# This step is not conservative. Just a linear interpolation. But that only has a small effect on the velocities, since the budget terms are already on the tracer grid. Same for the other functions.
ds = ds.sel(**roi)
ds_mean = weighted_mean(ds, ds.dyt, dim="yt_ocean")
for co in coord_include:
tempname = co + "_temp"
if tempname in list(ds_mean.data_vars):
ds_mean.coords[co] = ds_mean[tempname]
ds_mean = ds_mean.drop(tempname)
ds_mean.attrs["averaged region"] = str(roi)
return ds_mean
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 test_get_axis_pos():
datadict = datasets()
ds = datadict["C"]
coords = datadict["coords"]
grid = Grid(ds, coords=coords)
assert _get_axis_pos(grid, "X", ds.u) == "right"
assert _get_axis_pos(grid, "X", ds.tracer) == "center"
assert _get_axis_pos(grid, "Z", ds.u) is 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_face_connections_right_left_same_axis(self, boundary_width,
ds_faces, fill_value):
face_connections = {
"face": {
0: {
"X": (None, (1, "X", False))
},
1: {
"X": ((0, "X", False), None)
},
}
}
grid = Grid(ds_faces, face_connections=face_connections)
data = ds_faces.data_c
# fill in zeros for y boundary width if not given
boundary_width["Y"] = boundary_width.get("Y", (0, 0))
# # restrict data here, so its easier to see the output
# data = data.isel(y=slice(0, 2), x=slice(0, 2))
data = data.reset_coords(drop=True).reset_index(data.dims, drop=True)
# prepad the arrays
face_0_padded, face_1_padded = _prepad_right_left_same_axis(
data, boundary_width, fill_value)
# then simply add the corresponding slice to each face according to the connection
face_0_expected = xr.concat(
[
face_0_padded,
face_1_padded.isel(x=slice(0, boundary_width["X"][1]))
],
dim="x",
)
face_1_expected = xr.concat(
[
face_0_padded.isel(x=slice(
-boundary_width["X"][0],
None if boundary_width["X"][0] > 0 else 0,
# this is a bit annoying. if the boundary width on this side is
# 0 I want nothing to be padded. but slice(0,None) pads the whole array...
)),
face_1_padded,
],
dim="x",
)
expected = xr.concat([face_0_expected, face_1_expected], dim="face")
result = pad(
data,
grid,
boundary_width=boundary_width,
boundary="fill",
fill_value=fill_value,
other_component=None,
)
xr.testing.assert_allclose(result, expected)
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 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 test_find_dim():
datadict = datasets()
ds = datadict["C"]
grid = Grid(ds)
assert _find_dim(grid, ds, "X") == ["xt", "xu"]
assert _find_dim(grid, ds, "Z") is None
assert _find_dim(grid, ds["timeseries"], "X") is None
assert _find_dim(grid, ds["timeseries"], "X") is None
assert _find_dim(grid, ds["tracer"], "X") == ["xt"]
assert _find_dim(grid, ds["u"], "X") == ["xu"]
def depth_slice(ds, st_ocean=250, model_processing=True):
"""Calculate a depth slice via interpolation.
`model_processing` will interpolate all fields onto tracer grid."""
ds = ds.copy()
if model_processing:
grid = Grid(ds)
ds = interp_all(grid, ds)
ds = ds.interp(st_ocean=st_ocean)
ds.attrs["depth_slice"] = str(st_ocean)
return ds
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_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 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_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_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 test_infer_gridtype():
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)
# This should fail(unkown gridtype)
ds_fail = datadict["fail_gridtype"]
grid_fail = Grid(ds_fail, coords=coords)
# This is not supported yet ('inner' and 'outer' dims)
coords2 = datadict["fail_coords"]
ds_fail2 = datadict["fail_dimtype"]
grid_fail2 = Grid(ds_fail2, coords=coords2)
assert _infer_gridtype(grid_b, ds_b.u, ds_b.v) == "B"
assert _infer_gridtype(grid_c, ds_c.u, ds_c.v) == "C"
with pytest.raises(RuntimeError, match=r"Gridtype not recognized *"):
_infer_gridtype(grid_fail, ds_fail.u, ds_fail.v)
with pytest.raises(RuntimeError): # , match=r'`inner` or `outer` *'
_infer_gridtype(grid_fail2, ds_fail2.u, ds_fail2.v)
def test_xgcm_weighted_mean(axis, metric_list, gridtype):
ds = datasets()[gridtype]
grid = Grid(ds)
a = xgcm_weighted_mean(grid, ds, axis, metric_list)
for var in ["tracer", "u", "v"]:
b = w_mean(grid, ds[var], axis, metric_list)
c = xgcm_weighted_mean(grid, ds[var], axis, metric_list)
assert_allclose(a[var], b)
assert_allclose(b, c)
for var in ["timeseries"]:
b = ds[var]
c = xgcm_weighted_mean(grid, ds[var], axis, metric_list)
assert_allclose(a[var], b)
assert_allclose(b, c)