def test_to_grid():
testvars = ["salt", "u", "v", "z_w"]
for testvar in testvars:
for scoord in ["s_rho", "s_w"]:
var = xroms.to_grid(ds[testvar], grid, hcoord="rho", scoord=scoord)
# check for correct dims
assert scoord in var.dims
assert "eta_rho" in var.dims
assert "xi_rho" in var.dims
var = xroms.to_grid(ds[testvar], grid, hcoord="u", scoord=scoord)
assert scoord in var.dims
assert "eta_rho" in var.dims
assert "xi_u" in var.dims
var = xroms.to_grid(ds[testvar], grid, hcoord="v", scoord=scoord)
assert scoord in var.dims
assert "eta_v" in var.dims
assert "xi_rho" in var.dims
var = xroms.to_grid(ds[testvar], grid, hcoord="psi", scoord=scoord)
assert scoord in var.dims
assert "eta_v" in var.dims
assert "xi_u" in var.dims
def test_to_grid():
testvars = ["salt", "u", "v", "z_w"]
for testvar in testvars:
for scoord in ["s_w", "s_rho"]:
for hcoord in ["rho", "u", "v", "psi"]:
acc = ds[testvar].xroms.to_grid(hcoord=hcoord, scoord=scoord)
assert np.allclose(
acc,
xroms.to_grid(ds[testvar],
grid,
hcoord=hcoord,
scoord=scoord))
acc.name == acc.attrs["name"]
acc.attrs["grid"] == ds.xroms.grid
items = [
"T",
"X",
"Y",
"Z",
"longitude",
"latitude",
"vertical",
"time",
]
assert set(items).issubset(acc.cf.get_valid_keys())
acc = ds.xroms.to_grid(testvar, hcoord=hcoord, scoord=scoord)
assert np.allclose(
acc,
xroms.to_grid(ds[testvar],
grid,
hcoord=hcoord,
scoord=scoord))
acc.name == acc.attrs["name"]
acc.attrs["grid"] == ds.xroms.grid
items = [
"T",
"X",
"Y",
"Z",
"longitude",
"latitude",
"vertical",
"time",
]
assert set(items).issubset(acc.cf.get_valid_keys())
def test_KE():
rho = xroms.density(temp, salt, z_rho)[:, np.newaxis, np.newaxis]
s2 = (u ** 2 + v ** 2)[np.newaxis, :, :]
KE = 0.5 * rho * s2
# xroms
s = xroms.speed(ds.u, ds.v, grid)
assert np.allclose(
xroms.to_grid(xroms.KE(ds.rho0, s), grid, hcoord="psi").mean(),
KE.mean(),
rtol=1e-2,
)
def hgrad(
q,
grid,
which="both",
z=None,
hcoord=None,
scoord=None,
hboundary="extend",
hfill_value=None,
sboundary="extend",
sfill_value=None,
attrs=None,
):
"""Return gradients of property q accounting for s coordinates.
Inputs
------
q: DataArray
Property to take gradients of.
grid: xgcm.grid
Grid object associated with q.
which: string, optional
Which components of gradient to return.
* 'both': return both components of hgrad.
* 'xi': return only xi-direction.
* 'eta': return only eta-direction.
z: DataArray, ndarray, optional
Depth [m]. If None, use z coordinate attached to q.
hcoord: string, optional
Name of horizontal grid to interpolate output to.
Options are 'rho', 'psi', 'u', 'v'.
scoord: string, optional
Name of vertical grid to interpolate output to.
Options are 's_rho', 's_w', 'rho', 'w'.
hboundary: string, optional
Passed to `grid` method calls; horizontal boundary selection
for calculating horizontal derivatives of q. This same value
will be used for all horizontal grid changes too.
From xgcm documentation:
A flag indicating how to handle boundaries:
* None: Do not apply any boundary conditions. Raise an error if
boundary conditions are required for the operation.
* 'fill': Set values outside the array boundary to fill_value
(i.e. a Neumann boundary condition.)
* 'extend': Set values outside the array to the nearest array
value. (i.e. a limited form of Dirichlet boundary condition.
hfill_value: float, optional
Passed to `grid` method calls; horizontal boundary selection
fill value.
From xgcm documentation:
The value to use in the boundary condition with `boundary='fill'`.
sboundary: string, optional
Passed to `grid` method calls; vertical boundary selection
for calculating horizontal derivatives of q. This same value will
be used for all vertical grid changes too.
From xgcm documentation:
A flag indicating how to handle boundaries:
* None: Do not apply any boundary conditions. Raise an error if
boundary conditions are required for the operation.
* 'fill': Set values outside the array boundary to fill_value
(i.e. a Neumann boundary condition.)
* 'extend': Set values outside the array to the nearest array
value. (i.e. a limited form of Dirichlet boundary condition.
sfill_value: float, optional
Passed to `grid` method calls; vertical boundary selection
fill value.
From xgcm documentation:
The value to use in the boundary condition with `boundary='fill'`.
attrs: dict, optional
Dictionary of attributes to add to resultant arrays. Requires that
q is DataArray.
Returns
-------
DataArray(s) of dqdxi and/or dqdeta, the gradients of q
in the xi- and eta-directions with attributes altered to reflect calculation.
Notes
-----
dqdxi = dqdx*dzdz - dqdz*dzdx
dqdeta = dqdy*dzdz - dqdz*dzdy
Derivatives are taken in the ROMS curvilinear grid native xi- and eta- directions.
These derivatives properly account for the fact that ROMS vertical coordinates are
s coordinates and therefore can vary in time and space.
The xi derivative will alter the number of points in the xi and s dimensions.
The eta derivative will alter the number of points in the eta and s dimensions.
Example usage
-------------
>>> dtempdxi, dtempdeta = xroms.hgrad(ds.temp, grid)
"""
assert isinstance(q, xr.DataArray), "var must be DataArray"
if z is None:
try:
coords = list(q.coords)
z_coord_name = coords[[coord[:2] == "z_"
for coord in coords].index(True)]
z = q[z_coord_name]
is3D = True
except:
# if we get here that means that q doesn't have z coords (like zeta)
is3D = False
else:
is3D = True
if which in ["both", "xi"]:
if is3D:
dqdx = grid.interp(
grid.derivative(q,
"X",
boundary=hboundary,
fill_value=hfill_value),
"Z",
boundary=sboundary,
fill_value=sfill_value,
)
dqdz = grid.interp(
grid.derivative(q,
"Z",
boundary=sboundary,
fill_value=sfill_value),
"X",
boundary=hboundary,
fill_value=hfill_value,
)
dzdx = grid.interp(
grid.derivative(z,
"X",
boundary=hboundary,
fill_value=hfill_value),
"Z",
boundary=sboundary,
fill_value=sfill_value,
)
dzdz = grid.interp(
grid.derivative(z,
"Z",
boundary=sboundary,
fill_value=sfill_value),
"X",
boundary=hboundary,
fill_value=hfill_value,
)
dqdxi = dqdx * dzdz - dqdz * dzdx
else: # 2D variables
dqdxi = grid.derivative(q,
"X",
boundary=hboundary,
fill_value=hfill_value)
if attrs is None and isinstance(q, xr.DataArray):
attrs = q.attrs.copy()
attrs["name"] = "d" + q.name + "dxi"
attrs["units"] = "1/m * " + attrs.setdefault("units", "units")
attrs[
"long_name"] = "horizontal xi derivative of " + attrs.setdefault(
"long_name", "var")
attrs["grid"] = grid
dqdxi = xroms.to_grid(
dqdxi,
grid,
hcoord=hcoord,
scoord=scoord,
attrs=attrs,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
if which in ["both", "eta"]:
if is3D:
dqdy = grid.interp(
grid.derivative(q,
"Y",
boundary=hboundary,
fill_value=hfill_value),
"Z",
boundary=sboundary,
fill_value=sfill_value,
)
dqdz = grid.interp(
grid.derivative(q,
"Z",
boundary=sboundary,
fill_value=sfill_value),
"Y",
boundary=hboundary,
fill_value=hfill_value,
)
dzdy = grid.interp(
grid.derivative(z,
"Y",
boundary=hboundary,
fill_value=hfill_value),
"Z",
boundary=sboundary,
fill_value=sfill_value,
)
dzdz = grid.interp(
grid.derivative(z,
"Z",
boundary=sboundary,
fill_value=sfill_value),
"Y",
boundary=hboundary,
fill_value=hfill_value,
)
dqdeta = dqdy * dzdz - dqdz * dzdy
else: # 2D variables
dqdeta = grid.derivative(q,
"Y",
boundary=hboundary,
fill_value=hfill_value)
if attrs is None and isinstance(q, xr.DataArray):
attrs = q.attrs.copy()
attrs["name"] = "d" + q.name + "deta"
attrs["units"] = "1/m * " + attrs.setdefault("units", "units")
attrs[
"long_name"] = "horizontal eta derivative of " + attrs.setdefault(
"long_name", "var")
attrs["grid"] = grid
dqdeta = xroms.to_grid(
dqdeta,
grid,
hcoord=hcoord,
scoord=scoord,
attrs=attrs,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
if which == "both":
return dqdxi, dqdeta
elif which == "xi":
return dqdxi
elif which == "eta":
return dqdeta
else:
print("nothing being returned from hgrad")
def ertel(
phi,
u,
v,
f,
grid,
hcoord="rho",
scoord="s_rho",
hboundary="extend",
hfill_value=None,
sboundary="extend",
sfill_value=None,
):
"""Calculate Ertel potential vorticity of phi.
Inputs
------
phi: DataArray
Conservative tracer. Usually this would be the buoyancy but
could be another approximately conservative tracer. The
buoyancy can be calculated as:
>>> xroms.buoyancy(temp, salt, 0)
and then input as `phi`.
u: DataArray
xi component of velocity [m/s]
v: DataArray
eta component of velocity [m/s]
f: DataArray
Coriolis parameter [1/s]
grid: xgcm.grid
Grid object associated with u, v
hcoord: string, optional.
Name of horizontal grid to interpolate output to.
Options are 'rho', 'psi', 'u', 'v'.
scoord: string, optional.
Name of vertical grid to interpolate output to.
Options are 's_rho', 's_w', 'rho', 'w'.
hboundary: string, optional
Passed to `grid` method calls; horizontal boundary selection
for calculating horizontal derivatives of phi and for calculating
relative vorticity. This same value will be used for all
horizontal grid changes too.
From xgcm documentation:
A flag indicating how to handle boundaries:
* None: Do not apply any boundary conditions. Raise an error if
boundary conditions are required for the operation.
* 'fill': Set values outside the array boundary to fill_value
(i.e. a Neumann boundary condition.)
* 'extend': Set values outside the array to the nearest array
value. (i.e. a limited form of Dirichlet boundary condition.
hfill_value: float, optional
Passed to `grid` method calls; horizontal boundary selection
fill value.
From xgcm documentation:
The value to use in the boundary condition with `boundary='fill'`.
sboundary: string, optional
Passed to `grid` method calls; vertical boundary selection
for calculating horizontal and vertical derivatives of phi, and
for calculating relative vorticity. This same value will be used for
all vertical grid changes too.
From xgcm documentation:
A flag indicating how to handle boundaries:
* None: Do not apply any boundary conditions. Raise an error if
boundary conditions are required for the operation.
* 'fill': Set values outside the array boundary to fill_value
(i.e. a Neumann boundary condition.)
* 'extend': Set values outside the array to the nearest array
value. (i.e. a limited form of Dirichlet boundary condition.
sfill_value: float, optional
Passed to `grid` method calls; vertical boundary selection
fill value.
From xgcm documentation:
The value to use in the boundary condition with `boundary='fill'`.
Returns
-------
DataArray of the Ertel potential vorticity for the input tracer.
Notes
-----
epv = -v_z * phi_x + u_z * phi_y + (f + v_x - u_y) * phi_z
This is not set up to accept different boundary choices for different variables.
Example usage:
>>> xroms.ertel(ds.dye_01, ds.u, ds.v, ds.f, ds.attrs['grid'], scoord='s_w');
"""
assert isinstance(phi, xr.DataArray), "phi must be DataArray"
assert isinstance(u, xr.DataArray), "u must be DataArray"
assert isinstance(v, xr.DataArray), "v must be DataArray"
assert isinstance(f, xr.DataArray), "f must be DataArray"
phi_xi, phi_eta = xroms.hgrad(
phi,
grid,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
phi_xi = xroms.to_grid(
phi_xi,
grid,
hcoord=hcoord,
scoord=scoord,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
phi_eta = xroms.to_grid(
phi_eta,
grid,
hcoord=hcoord,
scoord=scoord,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
phi_z = xroms.ddz(
phi,
grid,
hcoord=hcoord,
scoord=scoord,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
# vertical shear (horizontal components of vorticity)
u_z = xroms.ddz(
u,
grid,
hcoord=hcoord,
scoord=scoord,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
v_z = xroms.ddz(
v,
grid,
hcoord=hcoord,
scoord=scoord,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
# vertical component of vorticity
vort = relative_vorticity(
u,
v,
grid,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
vort = xroms.to_grid(
vort,
grid,
hcoord=hcoord,
scoord=scoord,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
# combine terms to get the ertel potential vorticity
epv = -v_z * phi_xi + u_z * phi_eta + (f + vort) * phi_z
attrs = {
"name": "ertel",
"long_name": "ertel potential vorticity",
"units": "tracer/(m*s)",
"grid": grid,
}
epv = xroms.to_grid(
epv,
grid,
hcoord=hcoord,
scoord=scoord,
attrs=attrs,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
return epv
def to_grid(
self,
hcoord=None,
scoord=None,
hboundary="extend",
hfill_value=None,
sboundary="extend",
sfill_value=None,
):
"""Implement grid changes.
Inputs
------
hcoord: string, optional.
Name of horizontal grid to interpolate output to.
Options are 'rho', 'psi', 'u', 'v'.
scoord: string, optional.
Name of vertical grid to interpolate output to.
Options are 's_rho', 's_w', 'rho', 'w'.
hboundary: string, optional
Passed to `grid` method calls; horizontal boundary selection
for grid changes.
From xgcm documentation:
A flag indicating how to handle boundaries:
* None: Do not apply any boundary conditions. Raise an error if
boundary conditions are required for the operation.
* 'fill': Set values outside the array boundary to fill_value
(i.e. a Neumann boundary condition.)
* 'extend': Set values outside the array to the nearest array
value. (i.e. a limited form of Dirichlet boundary condition.
hfill_value: float, optional
Passed to `grid` method calls; horizontal boundary selection
fill value.
From xgcm documentation:
The value to use in the boundary condition with `boundary='fill'`.
sboundary: string, optional
Passed to `grid` method calls; vertical boundary selection
for grid changes.
From xgcm documentation:
A flag indicating how to handle boundaries:
* None: Do not apply any boundary conditions. Raise an error if
boundary conditions are required for the operation.
* 'fill': Set values outside the array boundary to fill_value
(i.e. a Neumann boundary condition.)
* 'extend': Set values outside the array to the nearest array
value. (i.e. a limited form of Dirichlet boundary condition.
sfill_value: float, optional
Passed to `grid` method calls; vertical boundary selection
fill value.
From xgcm documentation:
The value to use in the boundary condition with `boundary='fill'`.
Returns
-------
DataArray interpolated onto hcoord horizontal and scoord
vertical grids.
Notes
-----
If var is already on selected grid, nothing happens.
Example usage
-------------
>>> ds.salt.xroms.to_grid(hcoord='rho', scoord='w')
"""
return xroms.to_grid(
self.da,
self.da.attrs["grid"],
hcoord=hcoord,
scoord=scoord,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
def to_grid(
self,
varname,
hcoord=None,
scoord=None,
hboundary="extend",
hfill_value=None,
sboundary="extend",
sfill_value=None,
):
"""Implement grid changes.
Inputs
------
varname: str
Name of variable in Dataset to operate on.
hcoord: string, optional.
Name of horizontal grid to interpolate output to.
Options are 'rho', 'psi', 'u', 'v'.
scoord: string, optional.
Name of vertical grid to interpolate output to.
Options are 's_rho', 's_w', 'rho', 'w'.
hboundary: string, optional
Passed to `grid` method calls; horizontal boundary selection
for grid changes.
From xgcm documentation:
A flag indicating how to handle boundaries:
* None: Do not apply any boundary conditions. Raise an error if
boundary conditions are required for the operation.
* 'fill': Set values outside the array boundary to fill_value
(i.e. a Neumann boundary condition.)
* 'extend': Set values outside the array to the nearest array
value. (i.e. a limited form of Dirichlet boundary condition.
hfill_value: float, optional
Passed to `grid` method calls; horizontal boundary selection
fill value.
From xgcm documentation:
The value to use in the boundary condition with `boundary='fill'`.
sboundary: string, optional
Passed to `grid` method calls; vertical boundary selection
for grid changes.
From xgcm documentation:
A flag indicating how to handle boundaries:
* None: Do not apply any boundary conditions. Raise an error if
boundary conditions are required for the operation.
* 'fill': Set values outside the array boundary to fill_value
(i.e. a Neumann boundary condition.)
* 'extend': Set values outside the array to the nearest array
value. (i.e. a limited form of Dirichlet boundary condition.
sfill_value: float, optional
Passed to `grid` method calls; vertical boundary selection
fill value.
From xgcm documentation:
The value to use in the boundary condition with `boundary='fill'`.
Returns
-------
DataArray interpolated onto hcoord horizontal and scoord
vertical grids.
Notes
-----
If var is already on selected grid, nothing happens.
Example usage
-------------
>>> ds.xroms.to_grid('salt', hcoord='rho', scoord='w')
"""
assert isinstance(
varname, str
), "varname should be a string of the name of a variable stored in the Dataset"
assert varname in self.ds, 'variable called "varname" must be in Dataset'
var = xroms.to_grid(
self.ds[varname],
self.grid,
hcoord=hcoord,
scoord=scoord,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
)
self._ds[var.name] = var
return self._ds[var.name]