def test_shortcuts(od_in):
# Only use some variables
list_calc = [
"Temp",
"U",
"V",
"W",
"HFacC",
"HFacW",
"HFacS",
"drC",
"drF",
"dxC",
"dyC",
"dxF",
"dyF",
"dxG",
"dyG",
"dxV",
"dyU",
"rA",
"rAw",
"rAs",
"rAz",
]
od_in = od_in.subsample.cutout(varList=list_calc)
# Gradient
ds_out = gradient(od_in)
od_out = od_in.compute.gradient()
ds_out_IN_od_out(ds_out, od_out)
# Divergence
ds_out = divergence(od_in, iName="U", jName="V", kName="W")
od_out = od_in.compute.divergence(iName="U", jName="V", kName="W")
ds_out_IN_od_out(ds_out, od_out)
# Curl
ds_out = curl(od_in, iName="U", jName="V", kName="W")
od_out = od_in.compute.curl(iName="U", jName="V", kName="W")
ds_out_IN_od_out(ds_out, od_out)
# Laplacian
ds_out = laplacian(od_in, "Temp")
od_out = od_in.compute.laplacian(varNameList="Temp")
ds_out_IN_od_out(ds_out, od_out)
# Weighted mean
ds_out = weighted_mean(od_in)
od_out = od_in.compute.weighted_mean()
ds_out_IN_od_out(ds_out, od_out)
# Integral
ds_out = integral(od_in)
od_out = od_in.compute.integral()
ds_out_IN_od_out(ds_out, od_out)
def test_relative_vorticity(od_in):
# Extract and remove zeta
check = od_in._ds["momVort3"]
aliases = od.aliases
if aliases is None:
aliases = {}
vortName = aliases.pop("momVort3", None)
ds = od.dataset
if vortName is None:
ds = ds.drop_vars("momVort3")
else:
ds = ds.drop_vars(vortName)
od_in = ospy.OceanDataset(ds).set_aliases(aliases)
# Compute momVort1, momVort2, momVort3
ds_out = relative_vorticity(od_in)
varName = "momVort"
for i in range(3):
assert ds_out[varName + str(i + 1)].attrs["units"] == "s^-1"
long_name = "{}-component of relative vorticity".format(chr(105 + i))
assert ds_out[varName + str(i + 1)].attrs["long_name"] == long_name
# Check values
vort = curl(od_in, iName="U", jName="V", kName="W")
for i, curlName in enumerate(["dW_dY-dV_dZ", "dU_dZ-dW_dX",
"dV_dX-dU_dY"]):
assert_allclose(vort[curlName].values,
ds_out[varName + str(i + 1)].values)
var = ds_out["momVort3"]
# Mask and check
mask = xr.where(np.logical_or(check.isnull(), var.isnull()), 0, 1)
assert_allclose(
check.where(mask, drop=True).values,
var.where(mask, drop=True).values,
equal_nan=False,
)
# Test shortcut
od_out = od_in.compute.relative_vorticity()
ds_out_IN_od_out(ds_out, od_out)
def test_relative_vorticity(od_in):
# Extract and remove zeta
check = od_in._ds['momVort3']
aliases = od.aliases
if aliases is None:
aliases = {}
vortName = aliases.pop('momVort3', None)
ds = od.dataset
if vortName is None:
ds = ds.drop('momVort3')
else:
ds = ds.drop(vortName)
od_in = ospy.OceanDataset(ds).set_aliases(aliases)
# Compute momVort1, momVort2, momVort3
ds_out = relative_vorticity(od_in)
varName = 'momVort'
for i in range(3):
assert ds_out[varName+str(i+1)].attrs['units'] == 's^-1'
long_name = '{}-component of relative vorticity'.format(chr(105+i))
assert ds_out[varName+str(i+1)].attrs['long_name'] == long_name
# Check values
vort = curl(od_in, iName='U', jName='V', kName='W')
for i, curlName in enumerate(['dW_dY-dV_dZ',
'dU_dZ-dW_dX',
'dV_dX-dU_dY']):
assert_allclose(vort[curlName].values, ds_out[varName+str(i+1)].values)
var = ds_out['momVort3']
# Mask and check
mask = xr.where(np.logical_or(check.isnull(), var.isnull()), 0, 1)
assert_allclose(check.where(mask, drop=True).values,
var.where(mask, drop=True).values, equal_nan=False)
# Test shortcut
od_out = od_in.compute.relative_vorticity()
ds_out_IN_od_out(ds_out, od_out)
def test_curl(od):
velocities = [
[None, "sinVZ", "sinWY"],
["sinUZ", None, "sinWX"],
["sinUY", "sinVX", None],
["sinUY", "sinVX", "sinWX"],
]
for _, vels in enumerate(velocities):
# Add units
if None not in vels:
for varName in vels:
od._ds[varName].attrs["units"] = "m/s"
curl_ds = curl(od, iName=vels[0], jName=vels[1], kName=vels[2])
# sin' = cos
for var in curl_ds.data_vars:
var = curl_ds[var]
coords = {coord[0]: var[coord] for coord in var.coords}
coords["Z"], coords["Y"], coords["X"] = xr.broadcast(
coords["Z"], coords["Y"], coords["X"])
terms = var.name.split("-")
for i, term in enumerate(terms):
axis = term[-1]
terms[i] = np.cos(coords[axis])
check = terms[0] - terms[1]
mask = xr.where(np.logical_or(check.isnull(), var.isnull()), 0, 1)
# Assert using numpy
if None in vels:
assert_array_almost_equal(
check.where(mask, drop=True).values,
var.where(mask, drop=True).values,
7,
)
def test_curl_errors(od, iName, jName, kName):
with pytest.raises(ValueError):
curl(od, iName=iName, jName=jName, kName=kName)