def test_3d(get_global_ds):
"""check that vertical coordinate"""
ds = get_global_ds
grid = ecco_v4_py.get_llc_grid(ds)
maskK = ds['maskC']
maskL = grid.interp(maskK, 'Z', to='left', boundary='fill')
maskU = grid.interp(maskK, 'Z', to='right', boundary='fill')
maskKp1 = ds.maskC.isel(k=0).broadcast_like(ds.k_p1)
for mask in [maskK, maskL, maskU, maskKp1]:
ecco_v4_py.get_basin_mask('atl', mask)
def test_ext_basin_masks(get_global_ds):
"""make sure we can make the extended masks
"""
ds = get_global_ds
ext_names = ['atlExt', 'pacExt', 'indExt']
individual_names = [[
'atl', 'mexico', 'hudson', 'med', 'north', 'baffin', 'gin'
], ['pac', 'bering', 'okhotsk', 'japan', 'eastChina'],
['ind', 'southChina', 'java', 'timor', 'red', 'gulf']]
for ext, ind in zip(ext_names, individual_names):
maskE = ecco_v4_py.get_basin_mask(ext, ds.maskC.isel(k=0))
maskI = ecco_v4_py.get_basin_mask(ind, ds.maskC.isel(k=0))
assert np.all(maskE == maskI)
def test_notimplemented(get_test_ds):
ds = get_test_ds
with pytest.raises(NotImplementedError):
if len(ds.tile) < 13:
ecco_v4_py.get_basin_mask('atl', ds.maskC)
else:
ecco_v4_py.get_basin_mask('atl', ds.sel(tile=0).maskC)
ecco_v4_py.get_basin_mask('atl', (1. * ds.maskC).diff(dim='i'))
ecco_v4_py.get_basin_mask('atl', (1. * ds.maskC).diff(dim='j'))
def test_meridional_stf(get_test_ds, lats, basin, doFlip):
"""compute a meridional streamfunction"""
ds = get_test_ds
grid = ecco_v4_py.get_llc_grid(ds)
ds['U'], ds['V'] = get_fake_vectors(ds['U'], ds['V'])
ds = ds.rename({'U': 'UVELMASS', 'V': 'VVELMASS'})
if basin is None or len(ds.tile) == 13:
trsp = ecco_v4_py.calc_meridional_stf(ds,
lats,
doFlip=doFlip,
basin_name=basin,
grid=grid)
basinW = ds['maskW']
basinS = ds['maskS']
if basin is not None:
basinW = ecco_v4_py.get_basin_mask(basin, basinW)
basinS = ecco_v4_py.get_basin_mask(basin, basinS)
lats = [lats] if np.isscalar(lats) else lats
for lat in lats:
maskW, maskS = ecco_v4_py.vector_calc.get_latitude_masks(
lat, ds['YC'], grid)
trspx = (ds['drF'] * ds['dyG'] *
np.abs(maskW)).where(basinW).sum(dim=['i_g', 'j', 'tile'])
trspy = (ds['drF'] * ds['dxG'] *
np.abs(maskS)).where(basinS).sum(dim=['i', 'j_g', 'tile'])
test = trsp.sel(lat=lat).psi_moc.squeeze().reset_coords(drop=True)
expected = (1e-6 * (trspx + trspy)).reset_coords(drop=True)
if doFlip:
expected = expected.isel(k=slice(None, None, -1))
expected = expected.cumsum(dim='k')
if doFlip:
expected = -1 * expected.isel(k=slice(None, None, -1))
xr.testing.assert_allclose(test, expected)
else:
with pytest.raises(NotImplementedError):
trsp = ecco_v4_py.calc_meridional_stf(ds,
lats,
doFlip=doFlip,
basin_name=basin,
grid=grid)
def test_meridional_trsp(get_test_ds, myfunc, tfld, xflds, yflds, factor, lats,
basin):
"""compute a transport"""
ds = get_test_ds
grid = ecco_v4_py.get_llc_grid(ds)
ds['U'], ds['V'] = get_fake_vectors(ds['U'], ds['V'])
for fx, fy in zip(xflds, yflds):
ds[fx] = ds['U'].copy()
ds[fy] = ds['V'].copy()
if basin is None or len(ds.tile) == 13:
trsp = myfunc(ds, lats, basin_name=basin, grid=grid)
basinW = ds['maskW']
basinS = ds['maskS']
if basin is not None:
basinW = ecco_v4_py.get_basin_mask(basin, basinW)
basinS = ecco_v4_py.get_basin_mask(basin, basinS)
lats = [lats] if np.isscalar(lats) else lats
expx = (ds['drF'] * ds['dyG']).copy(
) if tfld == 'vol_trsp_z' else 2. * xr.ones_like(ds['hFacW'])
expy = (ds['drF'] * ds['dxG']).copy(
) if tfld == 'vol_trsp_z' else 2. * xr.ones_like(ds['hFacS'])
for lat in lats:
maskW, maskS = ecco_v4_py.vector_calc.get_latitude_masks(
lat, ds['YC'], grid)
trspx = (expx *
np.abs(maskW)).where(basinW).sum(dim=['i_g', 'j', 'tile'])
trspy = (expy *
np.abs(maskS)).where(basinS).sum(dim=['i', 'j_g', 'tile'])
test = trsp.sel(lat=lat)[tfld].squeeze().reset_coords(drop=True)
expected = (factor * (trspx + trspy)).reset_coords(drop=True)
xr.testing.assert_allclose(test, expected)
else:
with pytest.raises(NotImplementedError):
trsp = myfunc(ds, lats, basin_name=basin, grid=grid)
def test_each_basin_masks(get_global_ds):
"""make sure we can make the basin masks
"""
ds = get_global_ds
all_basins = ecco_v4_py.read_llc_to_tiles(os.path.join(
_test_dir, '..', '..', 'binary_data'),
'basins.data',
less_output=True)
ext_names = ['atlExt', 'pacExt', 'indExt']
for i, basin in enumerate(ecco_v4_py.get_available_basin_names(), start=1):
mask = ecco_v4_py.get_basin_mask(basin, ds.maskC.isel(k=0))
assert np.all(mask.values == (all_basins == i))
# user_lon_0=0,\
# show_colorbar=True, show_grid_labels=True, show_grid_lines=True)
maskC = ds_ecco.maskC
maskW = ds_ecco.maskW
maskS = ds_ecco.maskS
plt.figure(figsize=(12,5))
ecco.plot_tiles(dome_maskC+maskC.isel(k=0), cmap='viridis')
plt.show()
lat_maskW = grid.diff(dome_maskC,'X',boundary='fill')
lat_maskS = grid.diff(dome_maskC,'Y',boundary='fill')
atl_maskW = ecco.get_basin_mask(basin_name='atlExt', mask=maskW.isel(k=0))
atl_maskS = ecco.get_basin_mask(basin_name='atlExt', mask=maskS.isel(k=0))
plt.figure(figsize=(12,5))
ecco.plot_proj_to_latlon_grid(ds_ecco.XC, ds_ecco.YC, atl_maskW, projection_type='robin', cmap='viridis', user_lon_0=0)
plt.show()
mvt = ecco.calc_meridional_vol_trsp(ds_ecco,lat_vals=lat, basin_name='atlExt')
mht = ecco.calc_meridional_heat_trsp(ds_ecco,lat_vals=lat, basin_name='atlExt')
mst = ecco.calc_meridional_salt_trsp(ds_ecco,lat_vals=lat, basin_name='atlExt')
plt.figure(figsize=(12,5))
plt.plot(mvt.time,mvt.vol_trsp)
plt.ylabel('Meridional Volume Transport (Sv)')
plt.show()
plt.savefig('mvt_26N_Altlantic.png')
def test_bin_dir_is_here(get_global_ds, hide_bin_dir):
hide_bin_dir
ds = get_global_ds
with pytest.raises(OSError):
ecco_v4_py.get_basin_mask('atl', ds.maskC.isel(k=0))