def test_get_extra_metadata(domain, nx):
from xmitgcm.utils import get_extra_metadata
em = get_extra_metadata(domain=domain, nx=nx)
assert type(em) == dict
with pytest.raises(ValueError):
em = get_extra_metadata(domain='notinlist', nx=nx)
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 test_rebuild_llc_facets(domain, nx):
from xmitgcm.utils import rebuild_llc_facets, get_extra_metadata
md = get_extra_metadata(domain=domain, nx=nx)
nfaces = len(md['transpose_face'])
da = xarray.DataArray(np.empty((nfaces, nx, nx)), dims=['face', 'j', 'i'])
facets = rebuild_llc_facets(da, md)
for facet in range(5):
# test we get the original size
if md['facet_orders'][facet] == 'C':
expected_shape = (
md['ny_facets'][facet],
nx,
)
elif md['facet_orders'][facet] == 'F':
expected_shape = (
nx,
md['ny_facets'][facet],
)
if domain == 'aste' and facet == 1: # this facet is empty
pass
else:
assert facets['facet' + str(facet)].shape == expected_shape
def test_find_concat_dim_facet(domain, nx):
from xmitgcm.utils import find_concat_dim_facet, get_extra_metadata
md = get_extra_metadata(domain=domain, nx=nx)
nfacets = len(md['ny_facets'])
for facet in range(5):
da = xarray.DataArray(np.empty((nfacets, md['ny_facets'][facet], nx)),
dims=['face', 'j', 'i'])
concat_dim, non_concat_dim = find_concat_dim_facet(da, facet, md)
print(concat_dim, non_concat_dim)
if md['facet_orders'][facet] == 'C':
assert concat_dim == 'j'
assert non_concat_dim == 'i'
elif md['facet_orders'][facet] == 'F':
assert concat_dim == 'i'
assert non_concat_dim == 'j'
def get_test_ds(all_mds_datadirs):
"""make 2 tests when called, one with global, one with ASTE,
using fixture above"""
dirname, expected = all_mds_datadirs
kwargs = {}
if 'aste' in dirname:
kwargs['extra_metadata'] = get_extra_metadata('aste', 270)
kwargs['tiles_to_load'] = [0, 1, 2, 3, 4, 5]
kwargs['nx'] = 270
domain = 'aste'
else:
domain = 'global'
# read in array
ds = ecco.load_ecco_vars_from_mds(
dirname,
model_time_steps_to_load=expected['test_iternum'],
mds_files=['state_2d_set1', 'U', 'V', 'W', 'T', 'S'],
**kwargs)
ds.attrs['domain'] = domain
return ds
def test_get_grid_from_input(all_grid_datadirs, usedask):
from xmitgcm.utils import get_grid_from_input, get_extra_metadata
from xmitgcm.utils import read_raw_data
dirname, expected = all_grid_datadirs
md = get_extra_metadata(domain=expected['domain'], nx=expected['nx'])
ds = get_grid_from_input(dirname + '/' + expected['gridfile'],
geometry=expected['geometry'],
dtype=np.dtype('d'),
endian='>',
use_dask=usedask,
extra_metadata=md)
# test types
assert type(ds) == xarray.Dataset
assert type(ds['XC']) == xarray.core.dataarray.DataArray
if usedask:
ds.load()
# check all variables are in
expected_variables = [
'XC', 'YC', 'DXF', 'DYF', 'RAC', 'XG', 'YG', 'DXV', 'DYU', 'RAZ',
'DXC', 'DYC', 'RAW', 'RAS', 'DXG', 'DYG'
]
for var in expected_variables:
assert type(ds[var]) == xarray.core.dataarray.DataArray
assert ds[var].values.shape == expected['shape']
# check we don't leave points behind
if expected['geometry'] == 'llc':
nx = expected['nx'] + 1
nvars = len(expected_variables)
sizeofd = 8
grid = expected['gridfile']
grid1 = dirname + '/' + grid.replace('<NFACET>', '001')
grid2 = dirname + '/' + grid.replace('<NFACET>', '002')
grid3 = dirname + '/' + grid.replace('<NFACET>', '003')
grid4 = dirname + '/' + grid.replace('<NFACET>', '004')
grid5 = dirname + '/' + grid.replace('<NFACET>', '005')
size1 = os.path.getsize(grid1)
size2 = os.path.getsize(grid2)
size3 = os.path.getsize(grid3)
size4 = os.path.getsize(grid4)
size5 = os.path.getsize(grid5)
ny1 = int(size1 / sizeofd / nvars / nx)
ny2 = int(size2 / sizeofd / nvars / nx)
ny3 = int(size3 / sizeofd / nvars / nx)
ny4 = int(size4 / sizeofd / nvars / nx)
ny5 = int(size5 / sizeofd / nvars / nx)
xc1 = read_raw_data(grid1,
dtype=np.dtype('>d'),
shape=(ny1, nx),
partial_read=True)
xc2 = read_raw_data(grid2,
dtype=np.dtype('>d'),
shape=(ny2, nx),
partial_read=True)
xc3 = read_raw_data(grid3,
dtype=np.dtype('>d'),
shape=(ny3, nx),
partial_read=True)
xc4 = read_raw_data(grid4,
dtype=np.dtype('>d'),
shape=(ny4, nx),
order='F',
partial_read=True)
xc5 = read_raw_data(grid5,
dtype=np.dtype('>d'),
shape=(ny5, nx),
order='F',
partial_read=True)
yc1 = read_raw_data(grid1,
dtype=np.dtype('>d'),
shape=(ny1, nx),
partial_read=True,
offset=nx * ny1 * sizeofd)
yc2 = read_raw_data(grid2,
dtype=np.dtype('>d'),
shape=(ny2, nx),
partial_read=True,
offset=nx * ny2 * sizeofd)
yc3 = read_raw_data(grid3,
dtype=np.dtype('>d'),
shape=(ny3, nx),
partial_read=True,
offset=nx * ny3 * sizeofd)
yc4 = read_raw_data(grid4,
dtype=np.dtype('>d'),
shape=(ny4, nx),
order='F',
partial_read=True,
offset=nx * ny4 * sizeofd)
yc5 = read_raw_data(grid5,
dtype=np.dtype('>d'),
shape=(ny5, nx),
order='F',
partial_read=True,
offset=nx * ny5 * sizeofd)
xc = np.concatenate([
xc1[:-1, :-1].flatten(), xc2[:-1, :-1].flatten(),
xc3[:-1, :-1].flatten(), xc4[:-1, :-1].flatten(),
xc5[:-1, :-1].flatten()
])
yc = np.concatenate([
yc1[:-1, :-1].flatten(), yc2[:-1, :-1].flatten(),
yc3[:-1, :-1].flatten(), yc4[:-1, :-1].flatten(),
yc5[:-1, :-1].flatten()
])
xc_from_ds = ds['XC'].values.flatten()
yc_from_ds = ds['YC'].values.flatten()
assert xc.min() == xc_from_ds.min()
assert xc.max() == xc_from_ds.max()
assert yc.min() == yc_from_ds.min()
assert yc.max() == yc_from_ds.max()
# passing llc without metadata should fail
if expected['geometry'] == 'llc':
with pytest.raises(ValueError):
ds = get_grid_from_input(dirname + '/' + expected['gridfile'],
geometry=expected['geometry'],
dtype=np.dtype('d'),
endian='>',
use_dask=False,
extra_metadata=None)