def N2(rho, grid, rho0=1025.0, sboundary="fill", sfill_value=np.nan):
"""Calculate buoyancy frequency squared (vertical buoyancy gradient).
Inputs
------
rho: DataArray
Density [kg/m^3]
grid: xgcm.grid
Grid object associated with rho
rho0: int, float
Reference density [kg/m^3].
sboundary: string, optional
Passed to `grid` method calls; vertical boundary selection for
calculating z derivative.
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 fill value
associated with sboundary input.
From xgcm documentation:
The value to use in the boundary condition with `boundary='fill'`.
Returns
-------
DataArray of buoyancy frequency squared on rho/w grids.
Output is `[T,Z,Y,X]`.
Notes
-----
N2 = -g d(rho)/dz / rho0
Example usage
-------------
>>> xroms.N2(rho, grid)
"""
assert isinstance(rho, xr.DataArray), "rho must be DataArray"
drhodz = xroms.ddz(rho, grid, sboundary=sboundary, sfill_value=sfill_value)
var = -g * drhodz / rho0
var.attrs["name"] = "N2"
var.attrs[
"long_name"] = "buoyancy frequency squared, or vertical buoyancy gradient"
var.attrs["units"] = "1/s^2" # inherits grid
var.name = var.attrs["name"]
return var
def test_ddz():
testvars = ["salt", "u", "v", "z_w"]
for testvar in testvars:
acc = ds[testvar].xroms.ddz()
assert np.allclose(acc, xroms.ddz(ds[testvar], grid))
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.ddz(testvar)
assert np.allclose(acc, xroms.ddz(ds[testvar], grid))
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 dvdz(v, grid, sboundary="extend", sfill_value=None):
"""Calculate the eta component of vertical shear [1/s]
Inputs
------
v: DataArray
eta component of velocity [m/s]
grid: xgcm.grid
Grid object associated with v
sboundary: string, optional
Passed to `grid` method calls; vertical boundary selection for
calculating z derivative.
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 fill value
associated with sboundary input.
From xgcm documentation:
The value to use in the boundary condition with `boundary='fill'`.
Returns
-------
DataArray of eta component of vertical shear on v/w grids.
Output is `[T,Z,Y,X]`.
Notes
-----
v_z = ddz(v)
Wrapper of `ddz`
Example usage
-------------
>>> xroms.dvdz(v, grid)
"""
attrs = {
"name": "dvdz",
"long_name": "v component of vertical shear",
"units": "1/s",
"grid": grid,
}
return xroms.ddz(v,
grid,
attrs=attrs,
sboundary=sboundary,
sfill_value=sfill_value)
def test_ddz():
ddz = (salt[2] - salt[0]) / (z_rho[2] - z_rho[0])
assert np.allclose(xroms.ddz(ds.salt, grid)[0, 1, 0, 0], ddz)
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 ddz(
self,
hcoord=None,
scoord=None,
hboundary="extend",
hfill_value=None,
sboundary="extend",
sfill_value=None,
attrs=None,
):
"""Calculate d/dz for a variable.
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
calculating z derivative. This same value will be used for 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 fill value
associated with sboundary input.
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. For example:
`attrs={'name': 'varname', 'long_name': 'longvarname', 'units': 'units'}`
Returns
-------
DataArray of vertical derivative of variable with
attributes altered to reflect calculation.
Notes
-----
This will alter the number of points in the s dimension.
Example usage
-------------
>>> ds.salt.xroms.ddz()
"""
return xroms.ddz(
self.da,
self.da.attrs["grid"],
hcoord=hcoord,
scoord=scoord,
hboundary=hboundary,
hfill_value=hfill_value,
sboundary=sboundary,
sfill_value=sfill_value,
attrs=attrs,
)