def test_values_and_endianness(all_mds_datadirs):
"""Make sure we read all the grid variables."""
dirname, expected = all_mds_datadirs
if expected['geometry'] == 'llc' and (dask.__version__ < '0.11.2'):
pytest.xfail("LLC value tests require fixed dask")
# default endianness
ds = xmitgcm.open_mdsdataset(dirname,
iters=None,
read_grid=True,
swap_dims=False,
geometry=expected['geometry'])
# now reverse endianness
ds_le = xmitgcm.open_mdsdataset(dirname,
iters=None,
read_grid=True,
endian='<',
swap_dims=False,
geometry=expected['geometry'])
for vname, (idx, val) in expected['expected_values'].items():
np.testing.assert_allclose(ds[vname].values[idx], val)
# dask arrays that have been concatenated revert to native endianness
# https://github.com/dask/dask/issues/1647
if ds[vname].dtype.byteorder == '>':
val_le = ds[vname].values.newbyteorder('<')[idx]
np.testing.assert_allclose(ds_le[vname].values[idx], val_le)
def test_multiple_iters(multidim_mds_datadirs):
"""Test ability to load multiple iters into a single dataset."""
dirname, expected = multidim_mds_datadirs
# first try specifying the iters
ds = xmitgcm.open_mdsdataset(
dirname, read_grid=False, geometry=expected['geometry'],
iters=expected['all_iters'],
prefix=expected['prefixes'])
assert list(ds.iter.values) == expected['all_iters']
# now infer the iters, should be the same
ds2 = xmitgcm.open_mdsdataset(
dirname, read_grid=False, iters='all', geometry=expected['geometry'],
prefix=expected['prefixes'])
assert ds.equals(ds2)
# In the test datasets, there is no PHL.0000000000.data file.
# By default we infer the prefixes from the first iteration number, so this
# leads to an error.
# (Need to specify iters because there is some diagnostics output with
# weird iterations numbers present in some experiments.)
with pytest.raises(IOError):
ds = xmitgcm.open_mdsdataset(
dirname, read_grid=False, iters=expected['all_iters'],
geometry=expected['geometry'])
# now hide all the PH and PHL files: should be able to infer prefixes fine
missing_files = [os.path.basename(f)
for f in glob(os.path.join(dirname, 'PH*.0*data'))]
with hide_file(dirname, *missing_files):
ds = xmitgcm.open_mdsdataset(
dirname, read_grid=False, iters=expected['all_iters'],
geometry=expected['geometry'])
def get_compressed_level_index(grid_dir,
index_fname='llc4320_compressed_level_index.nc',
geometry='llc'):
''' Some doc
'''
#
ds = xm.open_mdsdataset(
'',
grid_dir=grid_dir,
iters=None,
geometry=geometry,
read_grid=True,
default_dtype=np.dtype('>f4'),
ignore_unknown_vars=True) #, extra_metadata=llc4320)
# get shape
#nz, nface, ny, nx = ds.hFacC.shape
#shape = (1, nface, ny, nx)
try:
ds_index = xr.open_dataset(grid_dir + index_fname)
except OSError:
# compute and save mask indices
print('Create llc4320_compressed_level_index.nc in grid_dir')
ds_index = ((ds.reset_coords()[['hFacC', 'hFacW', 'hFacS']] > 0).sum(
axis=(1, 2, 3)))
ds_index.coords['k'] = ds.k
ds_index.load().to_netcdf(grid_dir + index_fname)
print('done')
return ds_index, ds
def test_open_dataset_2D_diags(all_mds_datadirs):
# convert 3D fields with only 2D diagnostic output
# https://github.com/xgcm/xmitgcm/issues/140
dirname, expected = all_mds_datadirs
shape = expected['shape']
nz = shape[0]
ny, nx = shape[-2:]
shape_2d = shape[1:]
dims_2d = ('j', 'i')
if expected['geometry'] == 'llc':
dims_2d = ('face', ) + dims_2d
ny = nx * shape[-3]
dims_3d = dims_2d if nz == 1 else ('k', ) + dims_2d
dims_2d = ('time', ) + dims_2d
dims_3d = ('time', ) + dims_3d
it = expected['test_iternum']
kwargs = dict(iters=it,
geometry=expected['geometry'],
read_grid=False,
swap_dims=False)
to_hide = ['T.%010d.meta' % it, 'T.%010d.data' % it]
with hide_file(dirname, *to_hide):
ldir = py.path.local(dirname)
old_prefix = 'Eta.%010d' % it
new_prefix = 'T.%010d' % it
for suffix in ['.data', '.meta']:
lp = ldir.join(old_prefix + suffix)
lp.copy(ldir.join(new_prefix + suffix))
ds = xmitgcm.open_mdsdataset(dirname, prefix=['T'], **kwargs)
def test_llc_dims(llc_mds_datadirs, method, with_refdate):
"""Check that the LLC file dimensions are correct."""
dirname, expected = llc_mds_datadirs
if with_refdate:
kwargs = {
'delta_t': expected['delta_t'],
'ref_date': expected['ref_date']
}
else:
kwargs = {}
ds = xmitgcm.open_mdsdataset(dirname,
iters=expected['test_iternum'],
geometry=expected['geometry'],
llc_method=method,
**kwargs)
nz, nface, ny, nx = expected['shape']
nt = 1
assert ds.dims['face'] == 13
assert ds.rA.dims == ('face', 'j', 'i')
assert ds.rA.values.shape == (nface, ny, nx)
assert ds.U.dims == ('time', 'k', 'face', 'j', 'i_g')
assert ds.U.values.shape == (nt, nz, nface, ny, nx)
assert ds.V.dims == ('time', 'k', 'face', 'j_g', 'i')
assert ds.V.values.shape == (nt, nz, nface, ny, nx)
print(ds.U.chunks)
if method == "smallchunks":
assert ds.U.chunks == (nt * (1, ), nz * (1, ), nface * (1, ), (ny, ),
(nx, ))
def test_llc_extra_metadata(llc_mds_datadirs, method):
"""Check that the LLC reads properly when using extra_metadata."""
dirname, expected = llc_mds_datadirs
nz, nface, ny, nx = expected['shape']
nt = 1
llc = xmitgcm.utils.get_extra_metadata(domain='llc', nx=nx)
ds = xmitgcm.open_mdsdataset(dirname,
iters=expected['test_iternum'],
geometry=expected['geometry'],
llc_method=method,
extra_metadata=llc)
assert ds.dims['face'] == 13
assert ds.rA.dims == ('face', 'j', 'i')
assert ds.rA.values.shape == (nface, ny, nx)
assert ds.U.dims == ('time', 'k', 'face', 'j', 'i_g')
assert ds.U.values.shape == (nt, nz, nface, ny, nx)
assert ds.V.dims == ('time', 'k', 'face', 'j_g', 'i')
assert ds.V.values.shape == (nt, nz, nface, ny, nx)
if method == "smallchunks":
assert ds.U.chunks == (nt * (1, ), nz * (1, ), nface * (1, ), (ny, ),
(nx, ))
def test_open_mdsdataset_minimal(all_mds_datadirs):
"""Create a minimal xarray object with only dimensions in it."""
dirname, expected = all_mds_datadirs
ds = xmitgcm.open_mdsdataset(dirname,
iters=None,
read_grid=False,
swap_dims=False,
geometry=expected['geometry'])
# the expected dimensions of the dataset
eshape = expected['shape']
if len(eshape) == 3:
nz, ny, nx = eshape
nface = None
elif len(eshape) == 4:
nz, nface, ny, nx = eshape
else:
raise ValueError("Invalid expected shape")
coords = {
'i': np.arange(nx),
'i_g': np.arange(nx),
# 'i_z': np.arange(nx),
'j': np.arange(ny),
'j_g': np.arange(ny),
# 'j_z': np.arange(ny),
'k': np.arange(nz),
'k_u': np.arange(nz),
'k_l': np.arange(nz),
'k_p1': np.arange(nz + 1)
}
if nface is not None:
coords['face'] = np.arange(nface)
if 'layers' in expected:
for layer_name, n_layers in expected['layers'].items():
for suffix, offset in zip(['bounds', 'center', 'interface'],
[0, -1, -2]):
dimname = 'l' + layer_name[0] + '_' + suffix[0]
index = np.arange(n_layers + offset)
coords[dimname] = index
ds_expected = xr.Dataset(coords=coords)
assert ds_expected.equals(ds)
# check that the datatypes are correct
for key in coords:
assert ds[key].dtype == np.int64
# check for comodo metadata needed by xgcm
assert ds['i'].attrs['axis'] == 'X'
assert ds['i_g'].attrs['axis'] == 'X'
assert ds['i_g'].attrs['c_grid_axis_shift'] == -0.5
assert ds['j'].attrs['axis'] == 'Y'
assert ds['j_g'].attrs['axis'] == 'Y'
assert ds['j_g'].attrs['c_grid_axis_shift'] == -0.5
assert ds['k'].attrs['axis'] == 'Z'
assert ds['k_l'].attrs['axis'] == 'Z'
assert ds['k_l'].attrs['c_grid_axis_shift'] == -0.5
def test_llc_facets_2d_to_compact(llc_mds_datadirs):
from xmitgcm.utils import llc_facets_2d_to_compact, get_extra_metadata
from xmitgcm.utils import rebuild_llc_facets, read_raw_data
from xmitgcm.utils import write_to_binary
from xmitgcm import open_mdsdataset
dirname, expected = llc_mds_datadirs
# open dataset
ds = open_mdsdataset(dirname,
iters=expected['test_iternum'],
geometry=expected['geometry'])
nt, nfaces, ny, nx = expected['shape']
md = get_extra_metadata(domain=expected['geometry'], nx=nx)
# split in facets
facets = rebuild_llc_facets(ds['XC'], md)
flatdata = llc_facets_2d_to_compact(facets, md)
# compare with raw data
raw = read_raw_data(dirname + '/XC.data', np.dtype('>f'), (nfaces, ny, nx))
flatraw = raw.flatten()
assert len(flatdata) == len(flatraw)
assert flatdata.min() == flatraw.min()
assert flatdata.max() == flatraw.max()
# write new file
write_to_binary(flatdata, 'tmp.bin', dtype=np.dtype('f'))
md5new = file_md5_checksum('tmp.bin')
md5old = file_md5_checksum(dirname + '/XC.data')
assert md5new == md5old
os.remove('tmp.bin')
def plot_yield(name_file,
data_dir='./',
save=False,
e=2,
f=1.0,
t=0.0,
tstp=-1):
data = open_mdsdataset(data_dir, prefix=[name_file], chunks=None)
sig1 = data['SIsig1'][tstp]
sig2 = data['SIsig2'][tstp]
sig1f = sig1.values.flatten()
sig2f = sig2.values.flatten()
print len(sig1f)
plt.figure(1)
ax = plt.gca()
elli = Ellipse(xy=((-f + t) / 2, (-f + t) / 2),
width=(f + t) / e * np.sqrt(2),
height=(f + t) * np.sqrt(2),
angle=-45,
edgecolor='b',
fc='None',
lw=0.5)
ax.add_patch(elli)
plt.plot(sig1f, sig2f, 'r.', markersize=1)
xlim = ax.get_xlim()
ylim = ax.get_ylim()
ax.set(xlim=xlim, ylim=ylim)
cordx = [min(xlim[0], ylim[0]), min(xlim[1], ylim[1])]
cordy = [min(xlim[0], ylim[0]), min(xlim[1], ylim[1])]
plt.plot(cordx, cordy, '-k', lw=1)
coradx = [-min(-xlim[0], ylim[1]), min(xlim[1], -ylim[0])]
corady = [min(-xlim[0], ylim[1]), -min(xlim[1], -ylim[0])]
plt.plot(coradx, corady, '-k', lw=1)
plt.axvline(linewidth=1, color='k')
plt.axhline(linewidth=1, color='k')
plt.title(r" Yield curve ")
plt.axes().set_aspect('equal')
plt.grid()
if save == True:
plt.savefig("yield_curve.png", bbox_inches='tight', dpi=200)
plt.show()
# plt.close('all')
return None
def test_levels_diagnostics(mds_datadirs_with_inputfiles):
dirname, expected = mds_datadirs_with_inputfiles
for diagname, (levels, (idx, value)) in expected['diag_levels'].items():
ds = xmitgcm.open_mdsdataset(dirname,
prefix=[diagname],
levels=levels,
geometry=expected['geometry'])
assert ds['Zl'].values[idx] == value
with pytest.warns(UserWarning, match='nz will be ignored'):
xmitgcm.open_mdsdataset(dirname,
prefix=[diagname],
levels=levels,
geometry=expected['geometry'],
nz=12)
def plot_comp(name_file,
dirs,
varnames,
loga=0.,
beg=0,
end=-1,
step=1,
dpi=100):
data_dir1 = dirs[0]
data1 = open_mdsdataset(data_dir1, prefix=[name_file])
data_dir2 = dirs[1]
data2 = open_mdsdataset(data_dir2, prefix=[name_file])
if loga == 0.:
loga = np.full((len(varnames)), False, dtype=bool)
if beg == 0 and end == -1 and step == 1:
nbr = data1.dims['time']
times = xrange(nbr)
else:
var_cho = np.array(var[beg:end + 1:step, :])
for ivar in range(len(varnames)):
name_var = varnames[ivar]
var_dir = str('./plots_comp_' + name_var + '_' +
str(data1.coords['iter'][beg].values) + '_' +
str(data1.coords['iter'][end].values))
var_dirp = str(var_dir + '/')
if os.path.exists(var_dir):
shutil.rmtree(var_dir)
os.mkdir(var_dir)
var_max = 0.9
var_min = 1.1
for tstep in times:
plot_comp_forloop(data1, data2, name_var, tstep, loga[ivar],
var_dirp, var_min, var_max, dpi)
return None
def test_drc_length(all_mds_datadirs):
"""Test that open_mdsdataset is adding an extra level to drC if it has length nr"""
dirname, expected = all_mds_datadirs
#Only older versions of the gcm have len(drC) = nr, so force len(drC) = nr for the test
copyfile(os.path.join(dirname, 'DRF.data'), os.path.join(dirname, 'DRC.data'))
copyfile(os.path.join(dirname, 'DRF.meta'), os.path.join(dirname, 'DRC.meta'))
ds = xmitgcm.open_mdsdataset(
dirname, iters=None, read_grid=True,
geometry=expected['geometry'])
assert len(ds.drC)==(len(ds.drF)+1)
def test_read_grid(all_mds_datadirs):
"""Make sure we read all the grid variables."""
dirname, expected = all_mds_datadirs
ds = xmitgcm.open_mdsdataset(dirname,
iters=None,
read_grid=True,
geometry=expected['geometry'])
for vname in _EXPECTED_GRID_VARS:
assert vname in ds
def mitgcm_Movie(ddir,
prefix=["tracer_snapshots"],
maskname="hFacC",
varname="TRAC01",
clim=[-1, 35]):
ds = open_mdsdataset(ddir, prefix=prefix, swap_dims=False)
ds = ds[varname].where(ds[maskname] == 1)
run = os.path.basename(os.path.normpath(ddir))
odir = ddir + "/movie"
Movie(ds, odir, clim=clim, moviename=run)
def test_open_with_ref_date_grid(mds_datadirs_with_refdate, method):
"""With ref_date, without grid."""
dirname, expected = mds_datadirs_with_refdate
ds = xmitgcm.open_mdsdataset(dirname,
iters='all',
prefix=['S'],
ref_date=expected['ref_date'],
read_grid=method,
delta_t=expected['delta_t'],
geometry=expected['geometry'])
def test_read_grid(all_mds_datadirs):
"""Make sure we read all the grid variables."""
dirname, expected = all_mds_datadirs
ds = xmitgcm.open_mdsdataset(
dirname, iters=None, read_grid=True,
geometry=expected['geometry'])
for vname in _EXPECTED_GRID_VARS:
assert vname in ds
# actually load the data, to check for dask-related errors
ds.load()
def test_prefixes(all_mds_datadirs):
"""Make sure we read all the grid variables."""
dirname, expected = all_mds_datadirs
prefixes = ['U', 'V', 'W', 'T', 'S', 'PH'] # , 'PHL', 'Eta']
iters = [expected['test_iternum']]
ds = xmitgcm.open_mdsdataset(
dirname, iters=iters, prefix=prefixes,
read_grid=False, geometry=expected['geometry'])
for p in prefixes:
assert p in ds
def plot_yield_vid(name_file, data_dir='./', e=2, f=1.0, t=0.0):
FFMpegWriter = manimation.writers['ffmpeg']
metadata = dict(title='yield', artist='dringeis')
writer = FFMpegWriter(fps=30, metadata=metadata, bitrate=-1)
data = open_mdsdataset(data_dir, prefix=[name_file], chunks=None)
sig1 = data['SIsig1']
sig2 = data['SIsig2']
fig = plt.figure()
ax = fig.gca()
elli = Ellipse(xy=((-f + t) / 2, (-f + t) / 2),
width=(f + t) / e * np.sqrt(2),
height=(f + t) * np.sqrt(2),
angle=-45,
edgecolor='b',
fc='None',
lw=0.5)
ax.add_patch(elli)
im, = plt.plot([], [], 'r.', markersize=1)
xlim = ax.get_xlim()
ylim = ax.get_ylim()
ax.set(xlim=xlim, ylim=ylim)
cordx = [min(xlim[0], ylim[0]), min(xlim[1], ylim[1])]
cordy = [min(xlim[0], ylim[0]), min(xlim[1], ylim[1])]
plt.plot(cordx, cordy, '-k', lw=0.5)
coradx = [-min(-xlim[0], ylim[1]), min(xlim[1], -ylim[0])]
corady = [min(-xlim[0], ylim[1]), -min(xlim[1], -ylim[0])]
plt.plot(coradx, corady, '-k', lw=0.5)
plt.axvline(linewidth=0.5, color='k')
plt.axhline(linewidth=0.5, color='k')
plt.title(r" Yield curve ")
plt.axes().set_aspect('equal')
plt.grid()
with writer.saving(fig, "yield.mp4", 200):
for t in xrange(data.dims['time']):
sig1f = sig1[t, :, :].values.flatten()
sig2f = sig2[t, :, :].values.flatten()
im.set_data(sig1f, sig2f)
writer.grab_frame()
return None
def test_swap_dims(all_mds_datadirs):
"""See if we can swap dimensions."""
dirname, expected = all_mds_datadirs
kwargs = dict(iters=None, read_grid=True, geometry=expected['geometry'])
expected_dims = ['XC', 'XG', 'YC', 'YG', 'Z', 'Zl', 'Zp1', 'Zu']
# make sure we never swap if not reading grid
assert 'i' in xmitgcm.open_mdsdataset(dirname,
iters=None,
read_grid=False,
geometry=expected['geometry'])
if expected['geometry'] in ('llc', 'curvilinear'):
# make sure swapping is not the default
ds = xmitgcm.open_mdsdataset(dirname, **kwargs)
assert 'i' in ds
# and is impossible
with pytest.raises(ValueError) as excinfo:
ds = xmitgcm.open_mdsdataset(dirname, swap_dims=True, **kwargs)
else:
# make sure swapping *IS* the default
assert 'i' not in xmitgcm.open_mdsdataset(dirname, **kwargs)
ds = xmitgcm.open_mdsdataset(dirname,
geometry=expected['geometry'],
iters=None,
read_grid=True,
swap_dims=True,
grid_vars_to_coords=True)
# check for comodo metadata needed by xgcm
assert ds['XC'].attrs['axis'] == 'X'
assert ds['XG'].attrs['axis'] == 'X'
assert ds['XG'].attrs['c_grid_axis_shift'] == -0.5
assert ds['YC'].attrs['axis'] == 'Y'
assert ds['YG'].attrs['axis'] == 'Y'
assert ds['YG'].attrs['c_grid_axis_shift'] == -0.5
assert ds['Z'].attrs['axis'] == 'Z'
assert ds['Zl'].attrs['axis'] == 'Z'
assert ds['Zl'].attrs['c_grid_axis_shift'] == -0.5
# add extra layers dimensions if needed
if 'layers' in expected:
for layer_name in expected['layers']:
extra_dims = [
'layer_' + layer_name + suffix
for suffix in ['_bounds', '_center', '_interface']
]
expected_dims += extra_dims
assert list(ds.dims.keys()) == expected_dims
# make sure swapping works with multiple iters
ds = xmitgcm.open_mdsdataset(dirname,
geometry=expected['geometry'],
prefix=['S'])
#print(ds)
ds.load()
assert 'XC' in ds['S'].dims
assert 'YC' in ds['S'].dims
def test_tload_with_ref_date_swap(mds_datadirs_with_refdate, method):
"""With ref_date, without grid."""
dirname, expected = mds_datadirs_with_refdate
ds = xmitgcm.open_mdsdataset(dirname,
iters='all',
prefix=['S'],
ref_date=expected['ref_date'],
read_grid=True,
delta_t=expected['delta_t'],
swap_dims=method,
geometry=expected['geometry'])
ds.time.load()
def load_masks_from_mds(grid_dir):
import xmitgcm
ds = xmitgcm.open_mdsdataset(grid_dir, iters=None, geometry='llc')
points = ['C', 'W', 'S']
masks = [
ds['hfac' + point].reset_coords(drop=True).rename('mask' + point)
for point in points
]
ds_mask = xr.merge(masks)
for c in ds_mask.coords:
ds_mask[c] = ds_mask[c].astype('i2')
return ds_mask.chunk({'k': 1, 'face': -1})
def test_mask_values(all_mds_datadirs):
"""Test that open_mdsdataset generates binary masks with correct values"""
dirname, expected = all_mds_datadirs
ds = xmitgcm.open_mdsdataset(
dirname, iters=None, read_grid=True,
geometry=expected['geometry'])
hFac_list = ['hFacC', 'hFacW', 'hFacS']
mask_list = ['maskC', 'maskW', 'maskS']
for hFac, mask in zip(hFac_list, mask_list):
xr.testing.assert_equal(ds[hFac] * ds[mask], ds[hFac])
def test_serialize_nonstandard_calendar(multidim_mds_datadirs, tmp_path):
dirname, expected = multidim_mds_datadirs
ref_date = '2680-01-01 00:00:00'
calendar = '360_day'
ds = xmitgcm.open_mdsdataset(dirname,
iters='all',
prefix=['S'],
ref_date=ref_date,
calendar=calendar,
read_grid=False,
delta_t=expected['delta_t'],
geometry=expected['geometry'])
ds.to_zarr(tmp_path / 'test.zarr')
def test_date_parsing(mds_datadirs_with_refdate):
"""Verify that time information is decoded properly."""
dirname, expected = mds_datadirs_with_refdate
ds = xmitgcm.open_mdsdataset(dirname,
iters='all',
prefix=['S'],
ref_date=expected['ref_date'],
read_grid=False,
delta_t=expected['delta_t'],
geometry=expected['geometry'])
for i, date in expected['expected_time']:
assert ds.time[i].values == date
def test_extra_variables(all_mds_datadirs):
"""Test that open_mdsdataset reads extra_variables correctly"""
dirname, expected = all_mds_datadirs
extra_variable_data = dict(
# use U,V to test with attrs (including mate)
testdataU=dict(dims=['k', 'j', 'i_g'],
attrs=dict(standard_name='sea_water_x_velocity',
mate='testdataV',
long_name='Zonal Component of Velocity',
units='m s-1')),
testdataV=dict(dims=['k', 'j_g', 'i'],
attrs=dict(standard_name='sea_water_y_velocity',
mate='testdataU',
long_name='Meridional Component of Velocity',
units='m s-1')),
# use T to test without attrs
testdataT=dict(dims=['k', 'j', 'i'], attrs=dict()))
for var in ["U", "V", "T"]:
input_dir = os.path.join(
dirname, '{}.{:010d}'.format(var, expected['test_iternum']))
test_dir = os.path.join(
dirname, 'testdata{}.{:010d}'.format(var,
expected['test_iternum']))
if input_dir + ".meta" not in os.listdir(
) or input_dir + ".data" not in os.listdir():
return
copyfile(input_dir + ".meta", test_dir + ".meta")
copyfile(input_dir + ".data", test_dir + ".data")
assert test_dir + ".meta" not in os.listdir(
), f"{var} did not copy meta!"
assert test_dir + ".data" not in os.listdir(
), f"{var} did not copy data!"
ds = xmitgcm.open_mdsdataset(dirname,
read_grid=False,
iters=expected['test_iternum'],
geometry=expected['geometry'],
prefix=list(extra_variable_data.keys()),
extra_variables=extra_variable_data)
for var in extra_variable_data.keys():
assert var in ds
if 'mate' in ds[var].attrs:
mate = ds[var].attrs['mate']
assert ds[mate].attrs['mate'] == var
def test_separate_grid_dir(all_mds_datadirs):
"""Make sure we can have the grid files in a separate directory."""
dirname, expected = all_mds_datadirs
prefixes = ['U', 'V', 'W', 'T', 'S', 'PH'] # , 'PHL', 'Eta']
iters = [expected['test_iternum']]
with hide_file(dirname,
*['XC.meta', 'XC.data', 'RC.meta', 'RC.data']) as grid_dir:
ds = xmitgcm.open_mdsdataset(
dirname, grid_dir=grid_dir, iters=iters, prefix=prefixes,
read_grid=False, geometry=expected['geometry'])
for p in prefixes:
assert p in ds
def test_ref_date(mds_datadirs_with_refdate, swap_dims, read_grid, load):
"""With ref_date, without grid."""
dirname, expected = mds_datadirs_with_refdate
if expected['geometry']=='llc' and swap_dims:
pytest.skip("can't swap_dims with geometry=='llc'")
ds = xmitgcm.open_mdsdataset(dirname, iters='all', prefix=['S'],
ref_date=expected['ref_date'],
delta_t=expected['delta_t'],
geometry=expected['geometry'],
read_grid=read_grid, swap_dims=swap_dims)
if load:
ds.time.load()
def calc_mean_EK(start=100, end=1000, step=100):
times = np.arange(start, end + 1, step, dtype=np.int)
var_list = ['U', 'V', 'W']
ds2 = open_mdsdataset('.',
prefix=['CN', 'SN', 'XC', 'YC'],
iters=100,
geometry='curvilinear',
ignore_unknown_vars=True)
XC = np.array(ds2.XC)
YC = np.array(ds2.YC)
#EK = np.zeros(XC.shape)
for i in range(len(times)):
time = times[i]
print("TIME: ", time / times[-1])
ds = open_mdsdataset('.',
prefix=var_list,
iters=time,
geometry='curvilinear',
ignore_unknown_vars=True,
default_dtype=np.float)
U = np.array(ds.U)[0, :, :]
V = np.array(ds.V)[0, :, :]
W = np.array(ds.W)[0, :, :]
U2 = np.sqrt(U * U + V * V + W * W)
if i == 0:
EK = U2.copy()
else:
EK += U2
return EK
def test_date_parsing(mds_datadirs_with_refdate):
"""Verify that time information is decoded properly."""
dirname, expected = mds_datadirs_with_refdate
ds = xmitgcm.open_mdsdataset(dirname, iters='all', prefix=['S'],
ref_date=expected['ref_date'], read_grid=False,
delta_t=expected['delta_t'],
geometry=expected['geometry'])
for i, date in expected['expected_time']:
assert ds.time[i].values == date
# since time was decoded, this encoding should be removed from attributes
assert 'units' not in ds.time.attrs
assert 'calendar' not in ds.time.attrs
def test_diagnostics(mds_datadirs_with_diagnostics):
"""Try reading dataset with diagnostics output."""
dirname, expected = mds_datadirs_with_diagnostics
diag_prefix, expected_diags = expected['diagnostics']
ds = xmitgcm.open_mdsdataset(dirname,
read_grid=False,
iters=expected['test_iternum'],
prefix=[diag_prefix],
geometry=expected['geometry'])
for diagname in expected_diags:
assert diagname in ds
# check vector mates
if 'mate' in ds[diagname].attrs:
mate = ds[diagname].attrs['mate']
assert ds[mate].attrs['mate'] == diagname
