def test_vertical_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 momVort3
ds_out = vertical_relative_vorticity(od_in)
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.vertical_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_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 TS_diagram(od, display=True, FuncAnimation_kwargs=None, **kwargs):
"""
Animate TS diagrams.
Parameters
----------
od: OceanDataset
oceandataset used to plot.
display: bool
If True, display the animation.
FuncAnimation_kwargs: dict
Keyword arguments from :py:func:`matplotlib.animation.FuncAnimation`
**kwargs:
Keyword arguments from :py:func:`oceanspy.plot.TS_diagram`
Returns
-------
anim: matplotlib.animation.FuncAnimation
Animation object
See also
--------
oceanspy.plot.TS_diagram
"""
# Check parameters
_check_instance(
{
"od": od,
"display": display,
"FuncAnimation_kwargs": FuncAnimation_kwargs
},
{
"od": "oceanspy.OceanDataset",
"display": "bool",
"FuncAnimation_kwargs": ["type(None)", "dict"],
},
)
# Handle kwargs
if FuncAnimation_kwargs is None:
FuncAnimation_kwargs = {}
FuncAnimation_kwargs = dict(FuncAnimation_kwargs)
# Name of the plot_functions
plot_func = eval("_plot.TS_diagram")
# First cutout and get time
cutout_kwargs = kwargs.pop("cutout_kwargs", None)
if cutout_kwargs is not None:
od = od.subsample.cutout(**cutout_kwargs)
time = od._ds["time"]
# Check Temp and S
varList = ["Temp", "S"]
od = _compute._add_missing_variables(od, varList)
# Fix T and S axes
Tlim = kwargs.pop("Tlim", None)
Slim = kwargs.pop("Slim", None)
if Tlim is None:
cmap_params = _determine_cmap_params(od._ds["Temp"].values,
center=False)
Tlim = [cmap_params["vmin"], cmap_params["vmax"]]
if Slim is None:
cmap_params = _determine_cmap_params(od._ds["S"].values, center=False)
Slim = [cmap_params["vmin"], cmap_params["vmax"]]
kwargs["Tlim"] = Tlim
kwargs["Slim"] = Slim
# Fix density
dens = kwargs.pop("dens", None)
if dens is None:
t, s = _xr.broadcast(
_xr.DataArray(_np.linspace(Tlim[0], Tlim[-1], 100), dims=("t")),
_xr.DataArray(_np.linspace(Slim[0], Slim[-1], 100), dims=("s")),
)
odSigma0 = _ospy.OceanDataset(_xr.Dataset({"Temp": t, "S": s}))
odSigma0 = odSigma0.set_parameters(od.parameters)
odSigma0 = odSigma0.compute.potential_density_anomaly()
odSigma0._ds = odSigma0._ds.set_coords(["Temp", "S"])
# Freezing point
paramsList = ["tempFrz0", "dTempFrz_dS"]
params2use = {
par: od.parameters[par]
for par in od.parameters if par in paramsList
}
tempFrz0 = params2use["tempFrz0"]
dTempFrz_dS = params2use["dTempFrz_dS"]
freez_point = tempFrz0 + odSigma0._ds["S"] * dTempFrz_dS
# Extract Density
dens = odSigma0._ds["Sigma0"]
dens = dens.where(odSigma0._ds["Temp"] > freez_point)
kwargs["dens"] = dens
# Fix colorbar
colorName = kwargs.pop("colorName", None)
if colorName is not None:
# Add missing variables (use private)
_colorName = _rename_aliased(od, colorName)
od = _compute._add_missing_variables(od, _colorName)
# Extract color (use public)
color = od.dataset[colorName]
# Create colorbar (stolen from xarray)
cmap_kwargs = kwargs.pop("cmap_kwargs", None)
if cmap_kwargs is None:
cmap_kwargs = {}
cmap_kwargs = dict(cmap_kwargs)
cmap_kwargs["plot_data"] = color.values
kwargs["cmap_kwargs"] = _determine_cmap_params(**cmap_kwargs)
kwargs["colorName"] = colorName
# Remove ax
_ax_warning(kwargs)
# Animation
anim = _create_animation(od=od,
time=time,
plot_func=plot_func,
func_kwargs=kwargs,
display=display,
**FuncAnimation_kwargs)
return anim
def TS_diagram(od, display=True, FuncAnimation_kwargs=None,
**kwargs):
"""
Animate TS diagrams.
Parameters
----------
od: OceanDataset
oceandataset used to plot.
display: bool
If True, display the animation.
FuncAnimation_kwargs: dict
Keyword arguments from :py:func:`matplotlib.animation.FuncAnimation`
**kwargs:
Keyword arguments from :py:func:`oceanspy.plot.TS_diagram`
Returns
-------
anim: matplotlib.animation.FuncAnimation
Animation object
See also
--------
oceanspy.plot.TS_diagram
"""
# Check parameters
_check_instance({'od': od,
'display': display,
'FuncAnimation_kwargs': FuncAnimation_kwargs},
{'od': 'oceanspy.OceanDataset',
'display': 'bool',
'FuncAnimation_kwargs': ['type(None)', 'dict']})
# Handle kwargs
if FuncAnimation_kwargs is None:
FuncAnimation_kwargs = {}
FuncAnimation_kwargs = dict(FuncAnimation_kwargs)
# Name of the plot_functions
plot_func = eval('_plot.TS_diagram')
# First cutout and get time
cutout_kwargs = kwargs.pop('cutout_kwargs', None)
if cutout_kwargs is not None:
od = od.subsample.cutout(**cutout_kwargs)
time = od._ds['time']
# Check Temp and S
varList = ['Temp', 'S']
od = _compute._add_missing_variables(od, varList)
# Fix T and S axes
Tlim = kwargs.pop('Tlim', None)
Slim = kwargs.pop('Slim', None)
if Tlim is None:
cmap_params = _determine_cmap_params(od._ds['Temp'].values,
center=False)
Tlim = [cmap_params['vmin'], cmap_params['vmax']]
if Slim is None:
cmap_params = _determine_cmap_params(od._ds['S'].values,
center=False)
Slim = [cmap_params['vmin'], cmap_params['vmax']]
kwargs['Tlim'] = Tlim
kwargs['Slim'] = Slim
# Fix density
dens = kwargs.pop('dens', None)
if dens is None:
t, s = _xr.broadcast(_xr.DataArray(_np.linspace(Tlim[0],
Tlim[-1],
100), dims=('t')),
_xr.DataArray(_np.linspace(Slim[0],
Slim[-1],
100), dims=('s')))
odSigma0 = _ospy.OceanDataset(_xr.Dataset({'Temp': t,
'S': s}))
odSigma0 = odSigma0.set_parameters(od.parameters)
odSigma0 = odSigma0.compute.potential_density_anomaly()
odSigma0._ds = odSigma0._ds.set_coords(['Temp', 'S'])
# Freezing point
paramsList = ['tempFrz0', 'dTempFrz_dS']
params2use = {par: od.parameters[par]
for par in od.parameters
if par in paramsList}
tempFrz0 = params2use['tempFrz0']
dTempFrz_dS = params2use['dTempFrz_dS']
freez_point = tempFrz0 + odSigma0._ds['S']*dTempFrz_dS
# Extract Density
dens = odSigma0._ds['Sigma0']
dens = dens.where(odSigma0._ds['Temp'] > freez_point)
kwargs['dens'] = dens
# Fix colorbar
colorName = kwargs.pop('colorName', None)
if colorName is not None:
# Add missing variables (use private)
_colorName = _rename_aliased(od, colorName)
od = _compute._add_missing_variables(od, _colorName)
# Extract color (use public)
color = od.dataset[colorName]
# Create colorbar (stolen from xarray)
cmap_kwargs = kwargs.pop('cmap_kwargs', None)
if cmap_kwargs is None:
cmap_kwargs = {}
cmap_kwargs = dict(cmap_kwargs)
cmap_kwargs['plot_data'] = color.values
kwargs['cmap_kwargs'] = _determine_cmap_params(**cmap_kwargs)
kwargs['colorName'] = colorName
# Remove ax
_ax_warning(kwargs)
# Animation
anim = _create_animation(od=od,
time=time,
plot_func=plot_func,
func_kwargs=kwargs,
display=display,
**FuncAnimation_kwargs)
return anim
def TS_diagram(od,
Tlim=None,
Slim=None,
dens=None,
meanAxes=None,
colorName=None,
plotFreez=True,
ax=None,
cmap_kwargs=None,
contour_kwargs=None,
clabel_kwargs=None,
cutout_kwargs=None,
**kwargs):
"""
Plot temperature-salinity diagram.
Parameters
----------
od: OceanDataset
oceandataset used to plot.
Tlim: array_like with 2 elements
Temperature limits on the y axis.
If None, uses min and max values.
Slim: array_like with 2 elements
Salinity limits on the x axis.
If None, uses min and max values.
dens: xarray.DataArray
DataArray with densities used for isopycnals.
Must have coordinates (Temp, S).
In None, dens is inferred from Temp and S.
meanAxes: 1D array_like, str, or None
List of axes over which to apply weighted mean.
If None, don't average.
colorName: str, None
Name of the variable to use to color (e.g., Temp).
If None, uses plot insted of scatter (much faster)
plotFreez: bool
If True, plot freezing line in blue.
ax: matplotlib.pyplot.axes
If None, uses the current axis.
cmap_kwargs: dict
Keyword arguments for the colormap (same used by xarray)
contour_kwargs: dict
Keyword arguments for
:py:func:`matplotlib.pytplot.contour` (isopycnals)
clabel_kwargs: dict
Keyword arguments for
:py:func:`matplotlib.pytplot.clabel` (isopycnals)
cutout_kwargs: dict
Keyword arguments for
:py:func:`oceanspy.subsample.cutout`
**kwargs:
If colorName is None:
Kewyword arguments for :py:func:`matplotlib.pytplot.plot`
Otherwise,
kewyword arguments for :py:func:`matplotlib.pytplot.scatter`
Returns
-------
ax: matplotlib.pyplot.axes
Axes object.
References
----------
http://xarray.pydata.org/en/stable/plotting.html#introduction
See Also
--------
oceanspy.animate.TS_diagram
"""
# Check parameters
_check_instance(
{
'od': od,
'colorName': colorName,
'plotFreez': plotFreez,
'ax': ax,
'cmap_kwargs': cmap_kwargs,
'contour_kwargs': contour_kwargs,
'clabel_kwargs': clabel_kwargs,
'cutout_kwargs': cutout_kwargs,
'dens': dens
}, {
'od': 'oceanspy.OceanDataset',
'colorName': ['type(None)', 'str'],
'plotFreez': 'bool',
'ax': ['type(None)', 'matplotlib.pyplot.Axes'],
'cmap_kwargs': ['type(None)', 'dict'],
'contour_kwargs': ['type(None)', 'dict'],
'clabel_kwargs': ['type(None)', 'dict'],
'cutout_kwargs': ['type(None)', 'dict'],
'dens': ['type(None)', 'xarray.DataArray']
})
if Tlim is not None:
Tlim = _np.asarray(Tlim)
if Tlim.size != 2:
raise ValueError('`Tlim` must contain 2 elements')
Tlim = Tlim.reshape(2)
if Slim is not None:
Slim = _np.asarray(Slim)
if Slim.size != 2:
raise ValueError('`Slim` must contain 2 elements')
Slim = Slim.reshape(2)
if dens is not None and not set(['Temp', 'S']).issubset(dens.coords):
raise ValueError('`dens` must have coordinates (Temp, S)')
# Change None in empty dict
if cmap_kwargs is None:
cmap_kwargs = {}
if contour_kwargs is None:
contour_kwargs = {}
if clabel_kwargs is None:
clabel_kwargs = {}
if cutout_kwargs is None:
cutout_kwargs = {}
# Cutout first
if len(cutout_kwargs) != 0:
od = od.subsample.cutout(**cutout_kwargs)
# Check and extract T and S
varList = ['Temp', 'S']
od = _add_missing_variables(od, varList)
# Compute mean
if meanAxes is not None:
mean_ds = _compute.weighted_mean(od,
varNameList=['Temp', 'S'],
axesList=meanAxes,
storeWeights=False,
aliased=False)
T = mean_ds['w_mean_Temp'].rename('Temp')
S = mean_ds['w_mean_S'].rename('S')
lost_coords = list(set(od._ds['Temp'].dims) - set(T.coords))
else:
T = od._ds['Temp']
S = od._ds['S']
lost_coords = []
# Extract color field, and interpolate if needed
if colorName is not None:
# Add missing variables (use private)
_colorName = _rename_aliased(od, colorName)
od = _add_missing_variables(od, _colorName)
# Extract color (use public)
color = od.dataset[colorName]
if meanAxes is not None:
mean_ds = _compute.weighted_mean(od,
varNameList=_colorName,
axesList=meanAxes,
storeWeights=False,
aliased=False)
color = mean_ds['w_mean_' + _colorName].rename(_colorName)
else:
color = od.dataset[colorName]
grid = od.grid
dims2interp = [dim for dim in color.dims if dim not in T.dims]
# Interpolation
for dim in dims2interp:
for axis in od.grid.axes.keys():
if dim in [
od.grid.axes[axis].coords[k]
for k in od.grid.axes[axis].coords.keys()
]:
print('Interpolating [{}] along [{}]-axis.'
''.format(colorName, axis))
attrs = color.attrs
color = grid.interp(color,
axis,
to='center',
boundary='fill',
fill_value=_np.nan)
color.attrs = attrs
# Broadcast, in case color has different dimensions
T, S, color = _xr.broadcast(T, S, color)
# Compute density
T = T.persist()
S = S.persist()
if Tlim is None:
Tlim = [T.min().values, T.max().values]
if Slim is None:
Slim = [S.min().values, S.max().values]
if dens is None:
print('Isopycnals: ', end='')
tlin = _xr.DataArray(_np.linspace(Tlim[0], Tlim[-1], 100), dims=('t'))
slin = _xr.DataArray(_np.linspace(Slim[0], Slim[-1], 100), dims=('s'))
t, s = _xr.broadcast(tlin, slin)
odSigma0 = _ospy.OceanDataset(_xr.Dataset({'Temp': t, 'S': s}))
odSigma0 = odSigma0.set_parameters(od.parameters)
odSigma0 = odSigma0.compute.potential_density_anomaly()
odSigma0._ds = odSigma0._ds.set_coords(['Temp', 'S'])
# Freezing point
paramsList = ['tempFrz0', 'dTempFrz_dS']
params2use = {
par: od.parameters[par]
for par in od.parameters if par in paramsList
}
tempFrz0 = params2use['tempFrz0']
dTempFrz_dS = params2use['dTempFrz_dS']
freez_point = tempFrz0 + odSigma0._ds['S'] * dTempFrz_dS
# Extract Density
dens = odSigma0._ds['Sigma0'].where(odSigma0._ds['Temp'] > freez_point)
# Create axis
if ax is None:
ax = _plt.gca()
# Use plot if colorless (faster!), otherwise use scatter
if colorName is None:
default_kwargs = {'color': 'k', 'linestyle': 'None', 'marker': '.'}
kwargs = {**default_kwargs, **kwargs}
ax.plot(S.values.flatten(), T.values.flatten(), **kwargs)
else:
# Mask points out of axes
color = color.where(_np.logical_and(T > min(Tlim), T < max(Tlim)))
color = color.where(_np.logical_and(S > min(Slim), T < max(Slim)))
color = color.stack(all_dims=color.dims)
c = color.values
# Create colorbar (stolen from xarray)
cmap_kwargs['plot_data'] = c
cmap_params = _xr.plot.utils._determine_cmap_params(**cmap_kwargs)
extend = cmap_params.pop('extend')
_ = cmap_params.pop('levels')
kwargs = {**cmap_params, **kwargs}
# Scatter
sc = ax.scatter(S.values.flatten(), T.values.flatten(), c=c, **kwargs)
_plt.colorbar(sc,
label=_xr.plot.utils.label_from_attrs(color),
extend=extend)
# Plot isopycnals
t = dens['Temp']
s = dens['S']
default_contour_kwargs = {'colors': 'gray'}
contour_kwargs = {**default_contour_kwargs, **contour_kwargs}
CS = ax.contour(s.values, t.values, dens.values, **contour_kwargs)
ax.clabel(CS, **clabel_kwargs)
# Plot freezing point
if plotFreez:
paramsList = ['tempFrz0', 'dTempFrz_dS']
params2use = {
par: od.parameters[par]
for par in od.parameters if par in paramsList
}
tempFrz0 = params2use['tempFrz0']
dTempFrz_dS = params2use['dTempFrz_dS']
s = _np.unique(s.values.flatten())
ax.plot(s, tempFrz0 + s * dTempFrz_dS, 'b')
# Set labels and limits
ax.set_xlabel(_xr.plot.utils.label_from_attrs(S))
ax.set_ylabel(_xr.plot.utils.label_from_attrs(T))
ax.set_xlim(Slim)
ax.set_ylim(Tlim)
# Set title
title = []
all_coords = list(lost_coords) + list(T.coords)
skip_coords = ['X', 'Y', 'Xp1', 'Yp1']
if any([dim in od._ds.dims for dim in ['mooring', 'station', 'particle']]):
skip_coords = [
coord for coord in od._ds.coords if 'X' in coord or 'Y' in coord
]
for coord in all_coords:
if coord not in skip_coords:
if coord in list(lost_coords):
da = od._ds['Temp']
pref = '<'
suf = '>'
else:
da = T
pref = ''
suf = ''
rng = [da[coord].min().values, da[coord].max().values]
units = da[coord].attrs.pop('units', '')
if units.lower() == 'none':
units = ''
if 'time' in coord:
for i, v in enumerate(rng):
ts = _pd.to_datetime(str(v))
rng[i] = ts.strftime('%Y-%m-%d %r')
if rng[0] == rng[-1]:
rng = '{}'.format(rng[0])
else:
rng = 'from {} to {}'.format(rng[0], rng[1])
title = title + [
'{}{}{}: {} {}'
''.format(pref, coord, suf, rng, units)
]
ax.set_title('\n'.join(title))
return ax
def TS_diagram(
od,
Tlim=None,
Slim=None,
dens=None,
meanAxes=None,
colorName=None,
plotFreez=True,
ax=None,
cmap_kwargs=None,
contour_kwargs=None,
clabel_kwargs=None,
cutout_kwargs=None,
**kwargs
):
"""
Plot temperature-salinity diagram.
Parameters
----------
od: OceanDataset
oceandataset used to plot.
Tlim: array_like with 2 elements
Temperature limits on the y axis.
If None, uses min and max values.
Slim: array_like with 2 elements
Salinity limits on the x axis.
If None, uses min and max values.
dens: xarray.DataArray
DataArray with densities used for isopycnals.
Must have coordinates (Temp, S).
In None, dens is inferred from Temp and S.
meanAxes: 1D array_like, str, or None
List of axes over which to apply weighted mean.
If None, don't average.
colorName: str, None
Name of the variable to use to color (e.g., Temp).
If None, uses plot insted of scatter (much faster)
plotFreez: bool
If True, plot freezing line in blue.
ax: matplotlib.pyplot.axes
If None, uses the current axis.
cmap_kwargs: dict
Keyword arguments for the colormap (same used by xarray)
contour_kwargs: dict
Keyword arguments for
:py:func:`matplotlib.pytplot.contour` (isopycnals)
clabel_kwargs: dict
Keyword arguments for
:py:func:`matplotlib.pytplot.clabel` (isopycnals)
cutout_kwargs: dict
Keyword arguments for
:py:func:`oceanspy.subsample.cutout`
**kwargs:
If colorName is None:
Kewyword arguments for :py:func:`matplotlib.pytplot.plot`
Otherwise,
kewyword arguments for :py:func:`matplotlib.pytplot.scatter`
Returns
-------
ax: matplotlib.pyplot.axes
Axes object.
References
----------
http://xarray.pydata.org/en/stable/plotting.html#introduction
See Also
--------
oceanspy.animate.TS_diagram
"""
# Check parameters
_check_instance(
{
"od": od,
"colorName": colorName,
"plotFreez": plotFreez,
"ax": ax,
"cmap_kwargs": cmap_kwargs,
"contour_kwargs": contour_kwargs,
"clabel_kwargs": clabel_kwargs,
"cutout_kwargs": cutout_kwargs,
"dens": dens,
},
{
"od": "oceanspy.OceanDataset",
"colorName": ["type(None)", "str"],
"plotFreez": "bool",
"ax": ["type(None)", "matplotlib.pyplot.Axes"],
"cmap_kwargs": ["type(None)", "dict"],
"contour_kwargs": ["type(None)", "dict"],
"clabel_kwargs": ["type(None)", "dict"],
"cutout_kwargs": ["type(None)", "dict"],
"dens": ["type(None)", "xarray.DataArray"],
},
)
if Tlim is not None:
Tlim = _np.asarray(Tlim)
if Tlim.size != 2:
raise ValueError("`Tlim` must contain 2 elements")
Tlim = Tlim.reshape(2)
if Slim is not None:
Slim = _np.asarray(Slim)
if Slim.size != 2:
raise ValueError("`Slim` must contain 2 elements")
Slim = Slim.reshape(2)
if dens is not None and not set(["Temp", "S"]).issubset(dens.coords):
raise ValueError("`dens` must have coordinates (Temp, S)")
# Change None in empty dict
if cmap_kwargs is None:
cmap_kwargs = {}
cmap_kwargs = dict(cmap_kwargs)
if contour_kwargs is None:
contour_kwargs = {}
contour_kwargs = dict(contour_kwargs)
if clabel_kwargs is None:
clabel_kwargs = {}
clabel_kwargs = dict(clabel_kwargs)
if cutout_kwargs is None:
cutout_kwargs = {}
cutout_kwargs = dict(cutout_kwargs)
# Cutout first
if len(cutout_kwargs) != 0:
od = od.subsample.cutout(**cutout_kwargs)
# Check and extract T and S
varList = ["Temp", "S"]
od = _add_missing_variables(od, varList)
# Compute mean
if meanAxes is not None:
mean_ds = _compute.weighted_mean(
od,
varNameList=["Temp", "S"],
axesList=meanAxes,
storeWeights=False,
aliased=False,
)
T = mean_ds["w_mean_Temp"].rename("Temp")
S = mean_ds["w_mean_S"].rename("S")
lost_coords = list(set(od._ds["Temp"].dims) - set(T.coords))
else:
T = od._ds["Temp"]
S = od._ds["S"]
lost_coords = []
# Extract color field, and interpolate if needed
if colorName is not None:
# Add missing variables (use private)
_colorName = _rename_aliased(od, colorName)
od = _add_missing_variables(od, _colorName)
# Extract color (use public)
color = od.dataset[colorName]
if meanAxes is not None:
mean_ds = _compute.weighted_mean(
od,
varNameList=_colorName,
axesList=meanAxes,
storeWeights=False,
aliased=False,
)
color = mean_ds["w_mean_" + _colorName].rename(_colorName)
else:
color = od.dataset[colorName]
grid = od.grid
dims2interp = [dim for dim in color.dims if dim not in T.dims]
# Interpolation
for dim in dims2interp:
for axis in od.grid.axes.keys():
if dim in [
od.grid.axes[axis].coords[k]
for k in od.grid.axes[axis].coords.keys()
]:
print(
"Interpolating [{}] along [{}]-axis." "".format(colorName, axis)
)
attrs = color.attrs
color = grid.interp(
color, axis, to="center", boundary="fill", fill_value=_np.nan
)
color.attrs = attrs
# Broadcast, in case color has different dimensions
T, S, color = _xr.broadcast(T, S, color)
# Compute density
T = T.persist()
S = S.persist()
if Tlim is None:
Tlim = [T.min().values, T.max().values]
if Slim is None:
Slim = [S.min().values, S.max().values]
if dens is None:
print("Isopycnals: ", end="")
tlin = _xr.DataArray(_np.linspace(Tlim[0], Tlim[-1], 100), dims=("t"))
slin = _xr.DataArray(_np.linspace(Slim[0], Slim[-1], 100), dims=("s"))
t, s = _xr.broadcast(tlin, slin)
odSigma0 = _ospy.OceanDataset(_xr.Dataset({"Temp": t, "S": s}))
odSigma0 = odSigma0.set_parameters(od.parameters)
odSigma0 = odSigma0.compute.potential_density_anomaly()
odSigma0._ds = odSigma0._ds.set_coords(["Temp", "S"])
# Freezing point
paramsList = ["tempFrz0", "dTempFrz_dS"]
params2use = {
par: od.parameters[par] for par in od.parameters if par in paramsList
}
tempFrz0 = params2use["tempFrz0"]
dTempFrz_dS = params2use["dTempFrz_dS"]
freez_point = tempFrz0 + odSigma0._ds["S"] * dTempFrz_dS
# Extract Density
dens = odSigma0._ds["Sigma0"].where(odSigma0._ds["Temp"] > freez_point)
# Create axis
if ax is None:
ax = _plt.gca()
# Use plot if colorless (faster!), otherwise use scatter
if colorName is None:
default_kwargs = {"color": "k", "linestyle": "None", "marker": "."}
kwargs = {**default_kwargs, **kwargs}
ax.plot(S.values.flatten(), T.values.flatten(), **kwargs)
else:
# Mask points out of axes
color = color.where(_np.logical_and(T > min(Tlim), T < max(Tlim)))
color = color.where(_np.logical_and(S > min(Slim), T < max(Slim)))
color = color.stack(all_dims=color.dims)
c = color.values
# Create colorbar (stolen from xarray)
cmap_kwargs["plot_data"] = c
cmap_params = _xr.plot.utils._determine_cmap_params(**cmap_kwargs)
extend = cmap_params.pop("extend")
_ = cmap_params.pop("levels")
kwargs = {**cmap_params, **kwargs}
# Scatter
sc = ax.scatter(S.values.flatten(), T.values.flatten(), c=c, **kwargs)
_plt.colorbar(sc, label=_xr.plot.utils.label_from_attrs(color), extend=extend)
# Plot isopycnals
t = dens["Temp"]
s = dens["S"]
col_keys = ["colors", "cmap"]
default_contour_kwargs = {key: contour_kwargs.pop(key, None) for key in col_keys}
if all(default_contour_kwargs[key] is None for key in col_keys):
default_contour_kwargs["colors"] = "gray"
contour_kwargs = {**default_contour_kwargs, **contour_kwargs}
CS = ax.contour(s.values, t.values, dens.values, **contour_kwargs)
ax.clabel(CS, **clabel_kwargs)
# Plot freezing point
if plotFreez:
paramsList = ["tempFrz0", "dTempFrz_dS"]
params2use = {
par: od.parameters[par] for par in od.parameters if par in paramsList
}
tempFrz0 = params2use["tempFrz0"]
dTempFrz_dS = params2use["dTempFrz_dS"]
s = _np.unique(s.values.flatten())
ax.plot(s, tempFrz0 + s * dTempFrz_dS, "b")
# Set labels and limits
ax.set_xlabel(_xr.plot.utils.label_from_attrs(S))
ax.set_ylabel(_xr.plot.utils.label_from_attrs(T))
ax.set_xlim(Slim)
ax.set_ylim(Tlim)
# Set title
title = []
all_coords = list(lost_coords) + list(T.coords)
skip_coords = ["X", "Y", "Xp1", "Yp1"]
if any([dim in od._ds.dims for dim in ["mooring", "station", "particle"]]):
skip_coords = [coord for coord in od._ds.coords if "X" in coord or "Y" in coord]
for coord in all_coords:
if coord not in skip_coords:
if coord in list(lost_coords):
da = od._ds["Temp"]
pref = "<"
suf = ">"
else:
da = T
pref = ""
suf = ""
rng = [da[coord].min().values, da[coord].max().values]
units = da[coord].attrs.pop("units", "")
if units.lower() == "none":
units = ""
if "time" in coord:
for i, v in enumerate(rng):
ts = _pd.to_datetime(str(v))
rng[i] = ts.strftime("%Y-%m-%d %r")
if rng[0] == rng[-1]:
rng = "{}".format(rng[0])
else:
rng = "from {} to {}".format(rng[0], rng[1])
title = title + ["{}{}{}: {} {}" "".format(pref, coord, suf, rng, units)]
ax.set_title("\n".join(title))
return ax
# Directory
Datadir = "./oceanspy/tests/Data/"
# Create an oceandataset for testing calculus functions
od = open_oceandataset.from_netcdf("{}MITgcm_rect_nc.nc" "".format(Datadir))
# Create an oceandataset for testing calculus functions
od_curv = open_oceandataset.from_netcdf("{}MITgcm_curv_nc.nc"
"".format(Datadir))
# Aliased od
ds = od.dataset
aliases = {var: var + "_alias" for var in ds.data_vars}
ds = ds.rename(aliases)
alias_od = ospy.OceanDataset(ds).set_aliases(aliases)
# Budgets
od_bdg = open_oceandataset.from_netcdf("{}budgets.nc" "".format(Datadir))
@pytest.mark.parametrize("od_in", [od])
def test_missing_horizontal_spacing(od_in):
# Compute
od_in = od_in.subsample.cutout(varList=["dxC", "dxG", "dyC", "dyG"])
ds = missing_horizontal_spacing(od_in)
for varName in ds.variables:
var = ds[varName]
check = od.dataset[varName]
mask = xr.where(np.logical_or(check.isnull(), var.isnull()), 0, 1)
def TS_diagram(od, display=True, FuncAnimation_kwargs=None, **kwargs):
"""
Animate temperature-salinity diagram.
Parameters
----------
od: OceanDataset
oceandataset to check for missing variables
display: bool
display the animation in the notebook
FuncAnimation_kwargs: dict
Keyword arguments from matplotlib.animation.FuncAnimation
**kwargs:
Keyword arguments for plot.TS_diagram
Returns
-------
Animation object
See also
--------
plot.TS_diagram
"""
# Check input
if not isinstance(od, _ospy.OceanDataset):
raise TypeError('`od` must be OceanDataset')
if not isinstance(FuncAnimation_kwargs, (dict, type(None))):
raise TypeError('`FuncAnimation_kwargs` must be dict or None')
# Handle kwargs
if FuncAnimation_kwargs is None: FuncAnimation_kwargs = {}
# Name of the plot_functions
plot_func = eval('_plot.TS_diagram')
# First cutout and get time
cutout_kwargs = kwargs.pop('cutout_kwargs', None)
if cutout_kwargs is not None: od = od.subsample.cutout(**cutout_kwargs)
time = od._ds['time']
# Check Temp and S
varList = ['Temp', 'S']
od = _compute._add_missing_variables(od, varList)
# Fix T and S axes
Tlim = kwargs.pop('Tlim', None)
Slim = kwargs.pop('Slim', None)
if Tlim is None:
cmap_params = _xr.plot.utils._determine_cmap_params(
od._ds['Temp'].values, center=False)
Tlim = [cmap_params['vmin'], cmap_params['vmax']]
if Slim is None:
cmap_params = _xr.plot.utils._determine_cmap_params(od._ds['S'].values,
center=False)
Slim = [cmap_params['vmin'], cmap_params['vmax']]
kwargs['Tlim'] = Tlim
kwargs['Slim'] = Slim
# Fix density
dens = kwargs.pop('dens', None)
if dens is None:
t, s = _xr.broadcast(
_xr.DataArray(_np.linspace(Tlim[0], Tlim[-1], 100), dims=('t')),
_xr.DataArray(_np.linspace(Slim[0], Slim[-1], 100), dims=('s')))
odSigma0 = _ospy.OceanDataset(_xr.Dataset({
'Temp': t,
'S': s
})).set_parameters(od.parameters)
odSigma0 = odSigma0.compute.potential_density_anomaly()
odSigma0._ds = odSigma0._ds.set_coords(['Temp', 'S'])
# Freezing point
paramsList = ['tempFrz0', 'dTempFrz_dS']
params2use = {
par: od.parameters[par]
for par in od.parameters if par in paramsList
}
tempFrz0 = params2use['tempFrz0']
dTempFrz_dS = params2use['dTempFrz_dS']
freez_point = tempFrz0 + odSigma0._ds['S'] * dTempFrz_dS
# Extract Density
dens = odSigma0._ds['Sigma0'].where(odSigma0._ds['Temp'] > freez_point)
kwargs['dens'] = dens
# Fix colorbar
colorName = kwargs.pop('colorName', None)
if colorName is not None:
# Add missing variables (use private)
_colorName = _compute._rename_aliased(od, colorName)
od = _compute._add_missing_variables(od, _colorName)
# Extract color (use public)
color = od.dataset[colorName]
# Create colorbar (stolen from xarray)
cmap_kwargs = kwargs.pop('cmap_kwargs', None)
if cmap_kwargs is None: cmap_kwargs = {}
cmap_kwargs['plot_data'] = color.values
kwargs['cmap_kwargs'] = _xr.plot.utils._determine_cmap_params(
**cmap_kwargs)
kwargs['colorName'] = colorName
# Pop ax, it doesn't work for animation
ax = kwargs.pop('ax', None)
if ax is not None:
_warnings.warn(
"\n`ax` can not be provided for animations. "
"This function will use the current axis",
stacklevel=2)
# Animation
anim = _create_animation(od=od,
time=time,
plot_func=plot_func,
func_kwargs=kwargs,
display=display,
**FuncAnimation_kwargs)
return anim
