def calc_transports(ds, domain, fnt):
print('prepare grid')
DZTx = xr_DZ_xgcm(domain=domain, grid='T')
DZUx = xr_DZ_xgcm(domain=domain, grid='U')
(grid, ds_) = pop_tools.to_xgcm_grid_dataset(
ds) #, metrics=metrics, coords=coords)
# print(grid)
# print(ds_)
ds_['DZT'] = DZTx
ds_['DZU'] = DZUx
# print(ds_)
# grid = xgcm.Grid(ds_)
print('prepare mask')
MASK_tt = ds_.REGION_MASK
Atl_MASK_tt = xr.DataArray(data=np.in1d(MASK_tt, [6, 8, 9, 12]).reshape(
MASK_tt.shape),
dims=MASK_tt.dims,
coords=MASK_tt.coords)
print('prepare renaming')
rn = {'nlat_t': 'nlat_u', 'nlon_t': 'nlon_u'}
ac = {'nlat_u': ds_['nlat_u'].values, 'nlon_u': ds_['nlon_u'].values}
Atl_MASK_uu = Atl_MASK_tt.rename(rn).assign_coords(ac)
print('actual calculation')
ds_trans = calc_transports_pbc(grid=grid,
ds=ds_,
MASK_tt=Atl_MASK_tt,
MASK_uu=Atl_MASK_uu)
print('calculation done')
if fnt is None: return ds_trans
else:
ds_trans.to_netcdf(fnt)
print(f'created file: {fnt}')
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_to_xgcm_grid_dataset(ds, old_spatial_coords, axes):
grid, ds_new = pop_tools.to_xgcm_grid_dataset(ds, metrics=None)
assert isinstance(grid, xgcm.Grid)
assert set(axes) == set(grid.axes.keys())
new_spatial_coords = ['nlon_u', 'nlat_u', 'nlon_t', 'nlat_t']
assert set(new_spatial_coords).issubset(set(ds_new.coords))
assert not set(new_spatial_coords).intersection(set(ds.coords))
assert not set(old_spatial_coords).intersection(set(ds_new.coords))
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 xr_DZ_xgcm(domain, grid):
""" creates DZT/DZU 3D fields compatible with xgcm grids """
assert domain in ['ocn', 'ocn_low', 'ocn_rect']
assert grid in ['T', 'U']
fn = f'{path_samoc}/geometry/DZ{grid}_xgcm_{domain}.nc'
try: # file exists
# assert 1==0
assert os.path.exists(fn)
DZ = xr.open_dataarray(fn)
except: # create DZT/U
nlat_, nlon_ = f'nlat_{grid.lower()}', f'nlon_{grid.lower()}'
if domain == 'ocn': fe = file_ex_ocn_ctrl
elif domain == 'ocn_low': fe = file_ex_ocn_lpd
ds_ = xr.open_dataset(fe, decode_times=False)
(grid_, ds) = pop_tools.to_xgcm_grid_dataset(ds_)
H = ds[f'H{grid}']
KM = ds[f'KM{grid}']
DZ = xr.DataArray(
data=np.tile(ds.dz.values,
reps=(len(ds[nlon_].values), len(ds[nlat_].values),
1)), # imt,jmt,km
coords={
nlon_: ds[nlon_].values,
nlat_: ds[nlat_].values,
'z_t': ds['z_t'].values
},
dims=[nlon_, nlat_, 'z_t']).transpose() # -> km,jmt,imt
if domain == 'ocn': # in co
DZ[0, :, :] = xr.where(KM == -1, -1, DZ[0, :, :])
DZ[0, :, :] = xr.where(KM == 0, np.nan, DZ[0, :, :])
for k in np.arange(1, len(ds['z_t'])):
DZ[k, :, :] = xr.where(k >= KM - 1, np.nan, DZ[k, :, :])
DZ[k, :, :] = xr.where(k == KM - 1,
H - DZ.isel(z_t=slice(0, k)).sum('z_t'),
DZ[k, :, :])
# test that the total depth in DZ is the same as ds.HT/U
xr.testing.assert_allclose(DZ.sum('z_t'),
H.drop([f'{grid}LONG', f'{grid}LAT']),
rtol=0.01)
DZ.to_netcdf(fn)
return DZ
def calculate_AMOC_sigma_z(domain, ds, fn=None):
""" calculate the AMOC in depth and density space """
assert domain in ['ocn', 'ocn_low']
for q in ['PD', 'VVEL', 'DXT', 'DYT', 'DXU', 'DYU', 'REGION_MASK']:
assert q in ds
(grid, ds_) = pop_tools.to_xgcm_grid_dataset(ds)
ds_['DZU'] = xr_DZ_xgcm(domain=domain, grid='U')
metrics = {
('X'): ['DXT', 'DXU'], # X distances
('Y'): ['DYT', 'DYU'], # Y distances
('Z'): ['DZU'], # Z distances
}
coords = {
'X': {
'center': 'nlon_t',
'right': 'nlon_u'
},
'Y': {
'center': 'nlat_t',
'right': 'nlat_u'
},
'Z': {
'center': 'z_t',
'left': 'z_w_top',
'right': 'z_w_bot'
}
}
grid = xgcm.Grid(ds_, metrics=metrics, coords=coords)
print('merged annual datasets do not convert to U/T-lat/lons')
if 'nlat' in ds_.VVEL.dims:
rn = {'nlat': 'nlat_u', 'nlon': 'nlon_u'}
ac = {'nlat_u': ds_.nlat_u, 'nlon_u': ds_.nlon_u}
ds_['VVEL'] = ds_.VVEL.rename(rn).assign_coords()
if 'nlat' in ds_.PD.dims:
rn = {'nlat': 'nlat_t', 'nlon': 'nlon_t'}
ac = {'nlat_t': ds_.nlat_t, 'nlon_t': ds_.nlon_t}
ds_['PD'] = ds_.PD.rename(rn).assign_coords(ac)
print('interpolating density to UU point')
ds_['PD'] = grid.interp(grid.interp(ds_['PD'], 'X'), 'Y')
print('interpolating REGION_MASK to UU point')
fn_MASK = f'{path_prace}/MOC/AMOC_MASK_uu_{domain}.nc'
if os.path.exists(fn_MASK):
AMOC_MASK_uu = xr.open_dataarray(fn_MASK)
else:
MASK_uu = grid.interp(grid.interp(ds_.REGION_MASK, 'Y'), 'X')
AMOC_MASK_uu = xr.DataArray(np.in1d(
MASK_uu, [-12, 6, 7, 8, 9, 11, 12]).reshape(MASK_uu.shape),
dims=MASK_uu.dims,
coords=MASK_uu.coords)
AMOC_MASK_uu.to_netcdf(fn_MASK)
print('AMOC(y,z); [cm^3/s] -> [Sv]')
AMOC_yz = (grid.integrate(
grid.cumint(ds_.VVEL.where(AMOC_MASK_uu), 'Z', boundary='fill'), 'X') /
1e12)
# AMOC_yz = (ds_.VVEL*ds_.DZU*ds_.DXU).where(AMOC_MASK_uu).sum('nlon_u').cumsum('z_t')/1e12
AMOC_yz = AMOC_yz.rename({'z_w_top': 'z_t'}).assign_coords({'z_t': ds.z_t})
AMOC_yz.name = 'AMOC(y,z)'
print('AMOC(sigma_0,z); [cm^3/s] -> [Sv]')
if int(ds_.PD.isel(z_t=0).mean().values) == 0:
PD, PDbins = ds_.PD * 1000, np.arange(-10, 7, .05)
if int(ds_.PD.isel(z_t=0).mean().values) == 1:
PD, PDbins = (ds_.PD - 1) * 1000, np.arange(5, 33, .05)
print('histogram')
weights = ds_.VVEL.where(AMOC_MASK_uu) * ds_.DZU * ds_.DXU / 1e12
# ds_.PD.isel(z_t=0).plot()
AMOC_sz = histogram(PD, bins=[PDbins], dim=['z_t'],
weights=weights).sum('nlon_u',
skipna=True).cumsum('PD_bin').T
AMOC_sz.name = 'AMOC(y,PD)'
# output to file
if fn is not None: xr.merge([AMOC_yz, AMOC_sz]).to_netcdf(fn)
return AMOC_yz, AMOC_sz