def test_lateral_fill_np_array():
# generate psuedo-data
dx, dy = 0.05, 0.05
y, x = np.mgrid[slice(1, 3 + dy, dy), slice(1, 5 + dx, dx)]
z_orig = np.sin(x)**10 + np.cos(10 + y * x) * np.cos(x)
# construct mask and apply mask
valid_points = np.ones(z_orig.shape, dtype=np.bool)
valid_points = np.where(y < 0.5 * np.sin(5 * x) + 1.5, False, valid_points)
z_orig = np.where(~valid_points, np.nan, z_orig)
# add missing values
z_miss = z_orig.copy()
z_miss[:20, 62:] = np.nan
z_miss[15:18, 0:2] = 10.0
# compute lateral fill
z_fill = pop_tools.lateral_fill_np_array(z_miss, valid_points)
# load reference data
ref_data_file = DATASETS.fetch('lateral_fill_np_array_filled_ref.npz')
with np.load(ref_data_file) as data:
z_fill_ref = data['arr_0']
# assert that we match the reference solution
np.testing.assert_allclose(z_fill,
z_fill_ref,
atol=1e-14,
equal_nan=True,
verbose=True)
def test_get_grid_scrip():
ds_test = pop_tools.get_grid('POP_gx3v7', scrip=True)
ds_ref = xr.open_dataset(DATASETS.fetch('POP_gx3v7.nc'))
assert ds_compare(ds_test,
ds_ref,
assertion='allclose',
rtol=1e-14,
atol=1e-14)
def test_to_xgcm_grid_dataset_missing_xgcm():
from unittest import mock
with pytest.raises(ImportError):
with mock.patch.dict(sys.modules, {'xgcm': None}):
filepath = DATASETS.fetch('tend_zint_100m_Fe.nc')
ds = xr.open_dataset(filepath)
_, _ = pop_tools.to_xgcm_grid_dataset(ds, metrics=None)
def test_eos_ds_dask():
fname = DATASETS.fetch('cesm_pop_monthly.T62_g17.nc')
ds = xr.open_dataset(fname,
decode_times=False,
decode_coords=False,
chunks={'z_t': 20})
rho = pop_tools.eos(ds.SALT, ds.TEMP, depth=ds.z_t * 1e-2)
assert isinstance(rho, xr.DataArray)
def test_to_xgcm_grid_dataset(file):
filepath = DATASETS.fetch(file)
ds = xr.open_dataset(filepath)
grid, ds_new = pop_tools.to_xgcm_grid_dataset(ds, metrics=None)
assert isinstance(grid, xgcm.Grid)
assert set(['X', 'Y', 'Z']) == set(grid.axes.keys())
new_spatial_coords = set(['nlon_u', 'nlat_u', 'nlon_t', 'nlat_t'])
for coord in new_spatial_coords:
assert coord in ds_new.coords
assert coord not in ds.coords
def test_eos_xarray_2():
fname = DATASETS.fetch('cesm_pop_monthly.T62_g17.nc')
ds = xr.open_dataset(fname, decode_times=False, decode_coords=False)
rho, drhodS, drhodT = pop_tools.eos(ds.SALT,
ds.TEMP,
depth=ds.z_t * 1e-2,
return_coefs=True)
assert isinstance(rho, xr.DataArray)
assert isinstance(drhodS, xr.DataArray)
assert isinstance(drhodT, xr.DataArray)
def test_dataset_loading():
fname = DATASETS.fetch('tend_zint_100m_Fe.nc')
ds = xr.open_dataset(fname)
assert isinstance(ds, xr.Dataset)
import sys
import numpy as np
import pytest
import xarray as xr
import xgcm
import pop_tools
from pop_tools import DATASETS
ds_a = xr.open_dataset(DATASETS.fetch('tend_zint_100m_Fe.nc'), chunks={})
ds_b = xr.open_dataset(DATASETS.fetch('g.e20.G.TL319_t13.control.001_hfreq-coarsen.nc'), chunks={})
ds_c = ds_a[['TLAT', 'ULONG', 'ULAT', 'TLONG']].isel(nlat=slice(0, 2), nlon=slice(0, 2))
ds_c['anom'] = xr.DataArray(
np.arange(48).reshape(2, 2, 3, 2, 2),
dims=['S', 'L', 'M', 'nlat', 'nlon'],
coords={'S': [2000, 2001], 'L': [1, 2], 'M': [1, 2, 3]},
attrs={
'grid_loc': '2110',
'long_name': 'Dissolved Inorganic Iron Tendency Vertical Integral, 0-100m',
'cell_methods': 'time: mean',
'units': 'mmol/m^3 cm/s',
},
)
@pytest.mark.parametrize(
'ds, old_spatial_coords, axes',
[
(ds_a, ['nlat', 'nlon', 'z_w'], ['X', 'Y', 'Z']),
(ds_b, ['nlat', 'nlon', 'z_w'], ['X', 'Y', 'Z']),