def test_aliases():
# OceanDataset
assert od_in.aliases is None
# Wrong type
test_dict = 1
with pytest.raises(TypeError) as e:
od_in.set_aliases(test_dict)
assert str(e.value) == "`aliases` must be dict"
# Inhibit setter
test_dict = {var: var.lower() for var in od_in.dataset.variables}
with pytest.raises(AttributeError) as e:
od_in.aliases = test_dict
assert str(e.value) == _setter_error_message('aliases')
# Check setter
od_alias = OceanDataset(od_in.dataset.rename(test_dict))
od_out = od_in.set_aliases(test_dict)
assert set(test_dict.keys()).issubset(od_out._ds.variables)
assert set(test_dict.values()).issubset(od_out.dataset.variables)
# Check reopen
od_reopen = OceanDataset(od_out.dataset)
assert od_out.dataset.equals(od_reopen.dataset)
assert od_out._ds.equals(od_reopen._ds)
def test_OceanDataset(dataset):
if not isinstance(dataset, xr.Dataset):
# Raise error if wrong format
with pytest.raises(TypeError):
OceanDataset(dataset)
else:
# Just test that copy and repr don't raise errors
new_od = OceanDataset(dataset)
copy.copy(new_od)
repr(new_od)
def test_ver_sec(od_in, varName, contourName):
plt.close()
if "mooring_dist" in od_in.dataset.variables:
ds = od_in.dataset.drop_vars("mooring_dist")
od_in = OceanDataset(ds)
contour_kwargs = {"levels": 10}
clabel_kwargs = {"fontsize": 10}
meanAxes = True
intAxes = False
else:
contour_kwargs = None
clabel_kwargs = None
meanAxes = False
intAxes = True
ax = vertical_section(
od_in,
varName=varName,
contourName=contourName,
meanAxes=meanAxes,
intAxes=intAxes,
contour_kwargs=contour_kwargs,
clabel_kwargs=clabel_kwargs,
)
assert isinstance(ax, plt.Axes)
def test_vertical_cutout(od, dropAxes, add_Vbdr):
# Cover both case first index and any index
for i in range(2):
if i == 0:
ZRange = 0
elif i == 1:
ZRange = od.dataset["Zp1"].mean()
new_od = od.subsample.cutout(ZRange=ZRange,
dropAxes=dropAxes,
add_Vbdr=add_Vbdr)
if dropAxes is True and add_Vbdr is False:
assert (len(new_od.dataset["Zp1"]) == len(new_od.dataset["Z"]) ==
len(new_od.dataset["Zu"]) == len(
new_od.dataset["Zl"]) == 1)
elif dropAxes is False:
assert (len(new_od.dataset["Zp1"]) - 1 == len(new_od.dataset["Z"])
== len(new_od.dataset["Zu"]) == len(new_od.dataset["Zl"]))
assert (new_od.dataset["Zp1"].isel(
Zp1=slice(None, -1)).values == new_od.dataset["Zl"].values
).all()
assert (new_od.dataset["Zp1"].isel(
Zp1=slice(1, None)).values == new_od.dataset["Zu"].values
).all()
# Test Z, Zp1 only
new_ds = xr.Dataset({
var: od.dataset[var]
for var in od.dataset.variables
if "Z" in od.dataset[var].dims or "Zp1" in od.dataset[var].dims
})
new_ds.attrs = od.dataset.attrs
new_od = OceanDataset(new_ds)
new_od = new_od.subsample.cutout(ZRange=ZRange,
dropAxes=dropAxes,
add_Vbdr=add_Vbdr)
def test_import_MITgcm_rect_bin():
"""
Testing Magaldi and Haine, 2015
"""
# Dataset
ds = datasets['MITgcm_rect_bin']
# From open_oceandataset
od = OceanDataset(ds).import_MITgcm_rect_bin()
check_coords = {
'Y': 'center',
'Yp1': 'right',
'X': 'center',
'Xp1': 'right',
'Z': 'center',
'Zp1': 'outer',
'Zu': 'right',
'Zl': 'left',
'time': 'outer',
'time_midp': 'center'
}
# All coordinates
assert set(od.dataset.variables).issubset(od.dataset.coords)
# Check new dimensions
assert np.array_equal(od.dataset['XC'].isel(Y=0).values,
od.dataset['X'].values)
assert np.array_equal(od.dataset['XG'].isel(Yp1=0).values,
od.dataset['Xp1'].values)
assert np.array_equal(od.dataset['XU'].isel(Y=0).values,
od.dataset['Xp1'].values)
assert np.array_equal(od.dataset['XV'].isel(Yp1=0).values,
od.dataset['X'].values)
assert np.array_equal(od.dataset['YC'].isel(X=0).values,
od.dataset['Y'].values)
assert np.array_equal(od.dataset['YG'].isel(Xp1=0).values,
od.dataset['Yp1'].values)
assert np.array_equal(od.dataset['YU'].isel(Xp1=0).values,
od.dataset['Y'].values)
assert np.array_equal(od.dataset['YV'].isel(X=0).values,
od.dataset['Yp1'].values)
# Check grid
for axis in od.grid.axes.keys():
for pos in od.grid.axes[axis].coords.keys():
coord = od.grid.axes[axis].coords[pos].name
assert pos == check_coords[coord]
# Check time
od_time = pd.to_datetime(od.dataset['time'].values).to_julian_date()
od_time_midp = pd.to_datetime(
od.dataset['time_midp'].values).to_julian_date()
my_time = pd.to_datetime(ds['time'].values).to_julian_date()
my_time_midp = (my_time[:-1] + my_time[1:]) / 2
assert np.array_equal(od_time, my_time)
assert np.array_equal(od_time_midp, my_time_midp)
def test_time_resampling(od, timeFreq, sampMethod):
# Warning due to tiset(['1, 20'])me_midp
with pytest.warns(UserWarning):
new_od = od.subsample.cutout(timeFreq=timeFreq, sampMethod=sampMethod)
if timeFreq == "12H":
assert len(new_od.dataset["time"]) == np.ceil(len(od.dataset["time"]) / 2)
elif timeFreq == "6H":
assert len(new_od.dataset["time"]) == len(od.dataset["time"])
elif timeFreq == "1D" and sampMethod == "mean":
# Test time only
new_ds = new_od.dataset
new_ds = new_ds.drop_vars("time_midp")
new_od = OceanDataset(new_ds)
new_od.subsample.cutout(timeFreq="2D", sampMethod=sampMethod)
def test_time_resampling(od, timeFreq, sampMethod):
# Warning due to tiset(['1, 20'])me_midp
with pytest.warns(UserWarning):
new_od = od.subsample.cutout(timeFreq=timeFreq, sampMethod=sampMethod)
if timeFreq == '12H':
assert (len(new_od.dataset['time']) == np.ceil(
len(od.dataset['time']) / 2))
elif timeFreq == '6H':
assert (len(new_od.dataset['time']) == len(od.dataset['time']))
elif timeFreq == '1D' and sampMethod == 'mean':
with pytest.warns(RuntimeWarning):
# Test time only
new_ds = new_od.dataset
new_ds = new_ds.drop('time_midp')
new_od = OceanDataset(new_ds)
new_od.subsample.cutout(timeFreq='2D', sampMethod=sampMethod)
def test_particles(od, cartesian, varList, kwargs):
this_od = od
if cartesian:
this_od = this_od.set_parameters({"rSphere": None})
# Create 10 random paths
times = this_od.dataset["time"]
n_parts = 10
Ypart = np.empty((len(times), n_parts))
Xpart = np.empty((len(times), n_parts))
Zpart = np.zeros((len(times), n_parts))
for p in range(n_parts):
Ypart[:, p] = np.random.choice(this_od.dataset["Y"], len(times))
Xpart[:, p] = np.random.choice(this_od.dataset["X"], len(times))
# Extract particles
# Warning due to time_midp
with pytest.warns(UserWarning):
new_od = this_od.subsample.particle_properties(times=times,
Ypart=Ypart,
Xpart=Xpart,
Zpart=Zpart,
varList=varList,
**kwargs)
assert_array_equal(np.float32(new_od.dataset["XC"].values),
np.float32(Xpart))
assert_array_equal(np.float32(new_od.dataset["YC"].values),
np.float32(Ypart))
# Test without midp
new_ds = xr.Dataset({
var: this_od.dataset[var]
for var in this_od.dataset.variables
if "time_midp" not in this_od.dataset[var].dims
})
new_ds.attrs = this_od.dataset.attrs
new_od = OceanDataset(new_ds).set_grid_coords({"time": {
"time": None
}},
overwrite=False)
new_od.subsample.particle_properties(times=times,
Ypart=Ypart,
Xpart=Xpart,
Zpart=Zpart,
varList=varList)
def test_grid():
# OceanDataset
assert od_in.grid is None
# Check grid
od_out = od_in.set_grid_coords({'X': {'X': None, 'Xp1': 0.5}})
assert od_out.grid is not None
assert od_out.grid.axes['X']._periodic is False
assert od_out.grid.axes['X'].coords['center'].name == 'X'
if 'outer' in od_out.grid.axes['X'].coords:
assert od_out.grid.axes['X'].coords['outer'].name == 'Xp1'
elif 'right' in od_out.grid.axes['X'].coords:
assert od_out.grid.axes['X'].coords['right'].name == 'Xp1'
else: assert False
od_out = od_out.set_grid_periodic(['X'])
assert od_out.grid.axes['X']._periodic is True
# Check wrong name
with pytest.warns(UserWarning):
od_in.set_grid_coords({'X': {'X': None, 'Xp1': 0.5, 'test': -0.5}}).grid
# Inhibit setter
with pytest.raises(AttributeError):
od_in.grid = 1
with pytest.raises(AttributeError):
od_in._grid = 1
# Check aliases
test_dict = {var: var.lower() for var in od_in.dataset.variables}
od_alias = OceanDataset(od_in.dataset.rename(test_dict))
od_out = od_in.set_aliases(test_dict).set_grid_coords({'X': {'x': None, 'xp1': 0.5}})
assert od_out.grid.axes['X'].coords['center'].name == 'x'
assert od_out._grid.axes['X'].coords['center'].name == 'X'
if 'outer' in od_out.grid.axes['X'].coords:
assert od_out.grid.axes['X'].coords['outer'].name == 'xp1'
assert od_out._grid.axes['X'].coords['outer'].name == 'Xp1'
elif 'right' in od_out.grid.axes['X'].coords:
assert od_out.grid.axes['X'].coords['right'].name == 'xp1'
assert od_out._grid.axes['X'].coords['right'].name == 'Xp1'
else: assert False
# Directory
Datadir = './oceanspy/tests/Data/'
# Test oceandataset
MITgcm_curv_nc = open_oceandataset.from_netcdf('{}MITgcm_curv_nc.nc'
''.format(Datadir))
MITgcm_rect_bin = open_oceandataset.from_netcdf('{}MITgcm_rect_bin.nc'
''.format(Datadir))
MITgcm_rect_nc = open_oceandataset.from_netcdf('{}MITgcm_rect_nc.nc'
''.format(Datadir))
# =======
# CUTOUT
# =======
od = MITgcm_curv_nc
moor_od = OceanDataset(od.dataset.expand_dims('mooring'))
Ywarn = od.dataset['YG'].min() - 1
Xwarn = od.dataset['XG'].min() - 1
Zwarn = od.dataset['Zp1'].min() - 1
Twarn = od.dataset['time'].min() - 1
YRange = [od.dataset['YG'].min(), od.dataset['YG'].max()]
XRange = [od.dataset['XG'].min(), od.dataset['XG'].max()]
ZRange = [od.dataset['Zp1'].min(), od.dataset['Zp1'].max()]
timeRange = [od.dataset['time'].values[0], od.dataset['time'].values[-1]]
dropAxes = ['Y', 'X', 'Z', 'time']
dropwarn = 'mooring'
@pytest.mark.parametrize("od, YRange, XRange, ZRange, timeRange, dropAxes",
[(od, YRange, XRange, ZRange, timeRange, dropwarn),
(od, Ywarn, XRange, ZRange, timeRange, ['time']),
def test_import_MITgcm_curv():
"""
Testing exp_Arctic_Control
"""
# Dataset
ds = datasets['MITgcm_curv_nc']
# From open_oceandataset
od = OceanDataset(ds).import_MITgcm_curv_nc()
check_coords = {
'Y': 'center',
'Yp1': 'outer',
'X': 'center',
'Xp1': 'outer',
'Z': 'center',
'Zp1': 'outer',
'Zu': 'right',
'Zl': 'left',
'time': 'outer',
'time_midp': 'center'
}
# All coordinates
assert set(od.dataset.variables).issubset(od.dataset.coords)
# Check grid
for axis in od.grid.axes.keys():
for pos in od.grid.axes[axis].coords.keys():
coord = od.grid.axes[axis].coords[pos].name
assert pos == check_coords[coord]
# Check new dimensions
assert np.array_equal(
np.float32(od.dataset['XU'].values),
np.float32(
(od.dataset['XG'].values[:-1, :] + od.dataset['XG'].values[1:, :])
/ 2))
assert np.array_equal(
np.float32(od.dataset['YU'].values),
np.float32(
(od.dataset['YG'].values[:-1, :] + od.dataset['YG'].values[1:, :])
/ 2))
assert np.array_equal(
np.float32(od.dataset['XV'].values),
np.float32(
(od.dataset['XG'].values[:, :-1] + od.dataset['XG'].values[:, 1:])
/ 2))
assert np.array_equal(
np.float32(od.dataset['YV'].values),
np.float32(
(od.dataset['YG'].values[:, :-1] + od.dataset['YG'].values[:, 1:])
/ 2))
# Check time
od_time = pd.to_datetime(od.dataset['time'].values).to_julian_date()
od_time_midp = pd.to_datetime(
od.dataset['time_midp'].values).to_julian_date()
my_time = pd.to_datetime(ds['time'].values).to_julian_date()
my_time_midp = (my_time[:-1] + my_time[1:]) / 2
assert np.array_equal(od_time, my_time)
assert np.array_equal(od_time_midp, my_time_midp)
DEFAULT_PARAMETERS,
OCEANSPY_AXES,
OceanDataset,
open_oceandataset,
)
# Directory
Datadir = "./oceanspy/tests/Data/"
od = open_oceandataset.from_netcdf("{}MITgcm_rect_nc.nc" "".format(Datadir))
ECCO_url = "{}catalog_ECCO.yaml".format(Datadir)
ECCOod = open_oceandataset.from_catalog("LLC", ECCO_url)
# Remove global attributes
ds = od.dataset
ds.attrs = {}
clean_od = OceanDataset(ds)
# Aliased od
ds = od.dataset
aliases = {var: var + "_alias" for var in ds.data_vars}
ds = ds.rename(aliases)
alias_od = OceanDataset(ds).set_aliases(aliases)
# Grid only wihtout time_midp
ds = clean_od.dataset
ds = ds.drop_vars(ds.data_vars)
ds = ds.drop_vars(["Y", "time_midp"])
aliases = {dim: dim + "_alias" for dim in ds.dims}
ds_aliases = ds.rename(aliases)
nomidp_od = OceanDataset(ds)
alias_nomidp_od = OceanDataset(ds_aliases).set_aliases(aliases)
T = []
for i in range(len(ds["time_midp"])):
T = T + [np.datetime64(t0) + np.timedelta64(int(i * step * 1.0e3), "ms")]
ds["time_midp"] = np.array(T, dtype="datetime64") + np.timedelta64(
int(0.5 * step * 1.0e3), "ms")
# deltas
for var in ["drF", "dxC", "dyC", "dxF", "dyF", "dxG", "dyG", "dxV", "dyU"]:
ds[var] = xr.full_like(ds[var], step)
for var in ["rA", "rAw", "rAs", "rAz"]:
ds[var] = xr.full_like(ds[var], step**2)
for var in ["HFacC", "HFacW", "HFacS"]:
ds[var] = xr.ones_like(ds[var])
# Recreate oceandataset
od4calc = OceanDataset(ds)
# Gradient
sinX = xr.zeros_like(od4calc.dataset["Temp"]) + np.sin(od4calc.dataset["XC"])
sinY = xr.zeros_like(od4calc.dataset["Temp"]) + np.sin(od4calc.dataset["YC"])
sinZ = xr.zeros_like(od4calc.dataset["Temp"]) + np.sin(od4calc.dataset["Z"])
sintime = xr.zeros_like(od4calc.dataset["Temp"]) + np.sin(
(od4calc.dataset["time"] - od4calc.dataset["time"][0]) /
np.timedelta64(1, "s"))
sintime.attrs = od4calc.dataset["time"].attrs
cosX = xr.zeros_like(od4calc.dataset["U"]) + np.cos(od4calc.dataset["XU"])
cosY = xr.zeros_like(od4calc.dataset["V"]) + np.cos(od4calc.dataset["YV"])
cosZ = xr.zeros_like(od4calc.dataset["W"]) + np.cos(od4calc.dataset["Zl"])
costime = xr.zeros_like(od4calc.dataset["oceSPtnd"]) + np.cos(
(od4calc.dataset["time_midp"] - od4calc.dataset["time_midp"][0]) /
# Directory
Datadir = "./oceanspy/tests/Data/"
# Test oceandataset
MITgcm_curv_nc = open_oceandataset.from_netcdf("{}MITgcm_curv_nc.nc"
"".format(Datadir))
MITgcm_rect_bin = open_oceandataset.from_netcdf("{}MITgcm_rect_bin.nc"
"".format(Datadir))
MITgcm_rect_nc = open_oceandataset.from_netcdf("{}MITgcm_rect_nc.nc"
"".format(Datadir))
# =======
# CUTOUT
# =======
od = MITgcm_curv_nc
moor_od = OceanDataset(od.dataset.expand_dims("mooring"))
Ywarn = od.dataset["YG"].min() - 1
Xwarn = od.dataset["XG"].min() - 1
Zwarn = od.dataset["Zp1"].min() - 1
Twarn = od.dataset["time"].min() - 1
YRange = [od.dataset["YG"].min(), od.dataset["YG"].max()]
XRange = [od.dataset["XG"].min(), od.dataset["XG"].max()]
ZRange = [od.dataset["Zp1"].min(), od.dataset["Zp1"].max()]
timeRange = [od.dataset["time"].values[0], od.dataset["time"].values[-1]]
dropAxes = ["Y", "X", "Z", "time"]
dropwarn = "mooring"
@pytest.mark.parametrize(
"od, YRange, XRange, ZRange, timeRange, dropAxes",
[
def test_grid_coords():
# OceanDataset
assert od_in.grid_coords is None
# Wrong type
test_dict = 1
with pytest.raises(TypeError) as e:
od_in.set_grid_coords(test_dict)
assert str(e.value) == "`grid_coords` must be dict"
# Inhibit setter
test_dict = {'dim0': 1}
with pytest.raises(AttributeError) as e:
od_in.grid_coords = test_dict
assert str(e.value) == _setter_error_message('grid_coords')
# Wrong axis
with pytest.raises(ValueError) as e:
od_out = od_in.set_grid_coords(test_dict)
assert str(e.value) == (_wrong_axes_error_message(list(test_dict.keys())))
# Wrong type
test_dict = {'X': 1}
good_test_dict = {'Y': {'Y': None, 'Yp1': 0.5}}
with pytest.raises(TypeError) as e:
od_out = od_in.set_grid_coords(test_dict)
assert str(e.value) == "Invalid grid_coords. grid_coords example: {}".format(good_test_dict)
# Wrong shift
test_dict = {'X': {'X': 1}}
list_shift = [0.5, None, -0.5]
with pytest.raises(ValueError) as e:
od_out = od_in.set_grid_coords(test_dict)
assert str(e.value) == ("[{}] not a valid c_grid_axis_shift."
" Available options are {}").format(1, list_shift)
# Check set
od_out = od_in.set_grid_coords(good_test_dict)
assert od_out.grid_coords == good_test_dict
assert od_in.grid_coords != od_out.grid_coords
# Check overwrite
with pytest.raises(ValueError) as e:
od_out = od_out.set_grid_coords(good_test_dict)
assert str(e.value) == "[grid_coords] "+OVERWRITE_ERROR_MESSAGE
test_dict = {'Y': {'Y': 0.5, 'Yp1': None},
'X': {'X': None, 'Xp1': 0.5}}
assert od_out.set_grid_coords(test_dict, overwrite=False).grid_coords == {**od_out.grid_coords, **test_dict}
assert od_out.set_grid_coords(test_dict, overwrite=True).grid_coords == test_dict
# Midp
assert od_in.set_grid_coords({'X': {'X': None}} , add_midp=True).dataset == od_in.dataset
od_out = od_in.set_grid_coords({'X': {'X': 0.5}} , add_midp=True)
assert 'X_midp' in od_out.dataset.variables
assert np.array_equal(od_out.dataset['X_midp'].values,
(od_in.dataset['X'].values[:-1] + od_in.dataset['X'].values[1:])/2)
# Check aliases
od_alias = OceanDataset(od_in.dataset.rename({'X': 'x'}))
od_alias = od_alias.set_aliases({'X': 'x'})
od_out = od_alias.set_grid_coords({'X': {'x': 0.5}} , add_midp=True)
assert 'X_midp' in od_out._ds.variables
assert 'x_midp' in od_out.dataset.variables
import xarray as xr
import numpy as np
import pandas as pd
import sys
from oceanspy import OceanDataset
from oceanspy._oceandataset import _setter_error_message, _wrong_axes_error_message
from .datasets import datasets
from oceanspy import DEFAULT_PARAMETERS, AVAILABLE_PARAMETERS, TYPE_PARAMETERS, OCEANSPY_AXES
OVERWRITE_ERROR_MESSAGE = "has been previously set: `overwrite` must be bool"
# Drop attributes
ds_in = datasets['MITgcm_rect_bin']
ds_in.attrs = {}
od_in = OceanDataset(ds_in)
def test_name():
# OceanDataset
assert od_in.name is None
# Wrong type
test_str = 1
with pytest.raises(TypeError) as e:
od_in.set_name(test_str)
assert str(e.value) == "`name` must be str"
# Inhibit setter
test_str = 'test_name'.lower()
with pytest.raises(AttributeError) as e:
# Time Dimension
time = xr.DataArray(pd.date_range('2000-01-01', freq='M', periods=self.NT), dims = 'time')
return xr.Dataset({'X' : X, 'Xp1': Xp1,
'Y' : Y, 'Yp1': Yp1,
'Z' : Z, 'Zp1': Zp1, 'Zu': Zu, 'Zl': Zl,
'YC' : YC, 'XC' : XC,
'YG' : YG, 'XG' : XG,
'time': time, })
datasets = {'MITgcm_rect_nc' : Datasets().MITgcm_rect_nc(),
'MITgcm_rect_bin': Datasets().MITgcm_rect_bin(),
'MITgcm_curv_nc' : Datasets().MITgcm_curv_nc()}
oceandatasets = {'MITgcm_rect_nc' : OceanDataset(datasets['MITgcm_rect_nc']).import_MITgcm_rect_nc(),
'MITgcm_rect_bin': OceanDataset(datasets['MITgcm_rect_bin']).import_MITgcm_rect_bin(),
'MITgcm_curv_nc' : OceanDataset(datasets['MITgcm_curv_nc']).import_MITgcm_curv_nc()}
aliased_ods = {}
for od_name in oceandatasets:
dataset = oceandatasets[od_name].dataset
aliases = {var: 'alias_'+var for var in dataset.variables}
dataset = dataset.rename(aliases)
aliased_ods[od_name] = OceanDataset(dataset).set_aliases(aliases)
sin_od = OceanDataset(Datasets().sinusoidal()).set_grid_coords({'Y' : {'Y': None, 'Yp1': 0.5},
'X' : {'X': None, 'Xp1': 0.5},
'Z' : {'Z': None, 'Zp1': 0.5, 'Zu': 0.5, 'Zl': -0.5}})
