def test_latitude_masks(get_test_ds):
"""run through lats, and ensure we're grabbing the closest
"south-grid-cell-location" (whether that's in the x or y direction!)
to each latitude value"""
ds = get_test_ds
grid = get_llc_grid(ds)
yW = grid.interp(ds['YC'], 'X', boundary='fill')
yS = grid.interp(ds['YC'], 'Y', boundary='fill')
wetW = ds['maskW'].isel(k=0)
wetS = ds['maskS'].isel(k=0)
dLat = 0.5 # is this robust?
nx = 90
for lat in np.arange(-89, 89, 10):
print('lat: ', lat)
maskW, maskS = vector_calc.get_latitude_masks(lat, ds['YC'], grid)
maskW = maskW.where((maskW != 0) & wetW, 0.)
maskS = maskS.where((maskS != 0) & wetS, 0.)
arctic = int((len(maskW.tile) - 1) / 2)
assert not (maskW > 0).sel(tile=slice(arctic + 1, None)).any()
assert not (maskS < 0).sel(tile=slice(arctic - 1)).any()
assert (yW - lat < dLat).where(ds['maskW'].isel(k=0)
& (maskW != 0)).all().values
assert (yS - lat < dLat).where(ds['maskS'].isel(k=0)
& (maskS != 0)).all().values
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_optional_grid(get_test_ds):
"""simple, make sure we can optionally provide grid..."""
ds = get_test_ds
grid = get_llc_grid(ds)
uX = xr.ones_like(ds['U'].isel(k=0)).load()
vY = xr.ones_like(ds['V'].isel(k=0)).load()
u1, v1 = vector_calc.UEVNfromUXVY(uX, vY, ds)
u2, v2 = vector_calc.UEVNfromUXVY(uX, vY, ds, grid)
assert (u1 == u2).all()
assert (v1 == v2).all()
def test_separate_coords(get_test_ds, myfunc, myarg):
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'})
myarg['grid'] = grid
expected = myfunc(ds, **myarg)
coords = ds.coords.to_dataset().reset_coords()
ds = ds.reset_coords(drop=True)
test = myfunc(ds, coords=coords, **myarg)
xr.testing.assert_allclose(test['psi_moc'].reset_coords(drop=True),
expected['psi_moc'].reset_coords(drop=True))
def test_separate_coords(get_test_ds, myfunc, fld, xflds, yflds, lat):
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()
expected = myfunc(ds, lat, grid=grid)
coords = ds.coords.to_dataset().reset_coords()
ds = ds.reset_coords(drop=True)
test = myfunc(ds, lat, coords=coords, grid=grid)
xr.testing.assert_equal(test[fld].reset_coords(drop=True),
expected[fld].reset_coords(drop=True))
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_trsp_masking(get_test_ds, lat):
"""make sure internal masking is legit"""
ds = get_test_ds
grid = ecco_v4_py.get_llc_grid(ds)
ds['U'], ds['V'] = get_fake_vectors(ds['U'], ds['V'])
ds['U'] = ds['U'].where(ds['maskW'], 0.)
ds['V'] = ds['V'].where(ds['maskS'], 0.)
expected = ecco_v4_py.meridional_trsp_at_depth(ds['U'], ds['V'], lat, ds)
coords = ds[['Z', 'YC', 'XC', 'dyG', 'dxG', 'time']].copy()
coords.attrs = ds.attrs.copy()
ds = ds.reset_coords(drop=True)
test = ecco_v4_py.meridional_trsp_at_depth(ds['U'], ds['V'], lat, coords)
xr.testing.assert_equal(test['trsp_z'].reset_coords(drop=True),
expected['trsp_z'].reset_coords(drop=True))
def test_section_trsp(get_test_ds, myfunc, tfld, xflds, yflds, factor, args,
mask, error):
"""compute a volume transport,
within the lat/lon portion of the domain"""
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()
myargs = args.copy()
if mask:
myargs['maskW'], myargs[
'maskS'] = ecco_v4_py.vector_calc.get_latitude_masks(
30, ds['YC'], grid)
else:
myargs['maskW'] = None
myargs['maskS'] = None
if error is None:
trsp = myfunc(ds, grid=grid, **myargs)
maskW, maskS = ecco_v4_py.calc_section_trsp._parse_section_trsp_inputs(
ds, grid=grid, **myargs)
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'])
trspx = (expx * np.abs(maskW)).where(
ds['maskW']).sum(dim=['i_g', 'j', 'tile'])
trspy = (expy * np.abs(maskS)).where(
ds['maskS']).sum(dim=['i', 'j_g', 'tile'])
test = trsp[tfld].squeeze().reset_coords(drop=True)
expected = (factor * (trspx + trspy)).reset_coords(drop=True)
xr.testing.assert_equal(test, expected)
else:
with pytest.raises(error):
trsp = myfunc(ds, **myargs)
def test_section_stf(get_test_ds, args, mask, error, doFlip):
"""compute streamfunction across section"""
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'})
myargs = args.copy()
if mask:
myargs['maskW'], myargs[
'maskS'] = ecco_v4_py.vector_calc.get_latitude_masks(
30, ds['YC'], grid)
else:
myargs['maskW'] = None
myargs['maskS'] = None
if error is None:
trsp = ecco_v4_py.calc_section_stf(ds,
doFlip=doFlip,
grid=grid,
**myargs)
maskW, maskS = ecco_v4_py.calc_section_trsp._parse_section_trsp_inputs(
ds, **myargs)
trspx = (ds['drF'] * ds['dyG'] * np.abs(maskW)).where(
ds['maskW']).sum(dim=['i_g', 'j', 'tile'])
trspy = (ds['drF'] * ds['dxG'] * np.abs(maskS)).where(
ds['maskS']).sum(dim=['i', 'j_g', 'tile'])
test = trsp.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(error):
trsp = ecco_v4_py.calc_section_stf(ds, **myargs)
def test_latitude_mask(get_test_ds):
"""run through lats, and ensure we're grabbing the closest
"south-grid-cell-location" (whether that's in the x or y direction!)
to each latitude value"""
ds = get_test_ds
grid = get_llc_grid(ds)
wetC = ds['maskC'].isel(k=0)
dLat = 0.5 # is this robust?
nx = 90
for lat in np.arange(-89, 89, 10):
print('lat: ', lat)
maskC = scalar_calc.get_latitude_mask(lat, ds['YC'], grid)
maskC = maskC.where((maskC != 0) & wetC, 0.)
assert (ds['YC'] - lat < dLat).where((maskC != 0)).all().values
def test_trsp_masking(get_test_ds, section_name):
"""make sure internal masking is legit"""
ds = get_test_ds
grid = ecco_v4_py.get_llc_grid(ds)
ds['U'], ds['V'] = get_fake_vectors(ds['U'], ds['V'])
ds['U'] = ds['U'].where(ds['maskW'], 0.)
ds['V'] = ds['V'].where(ds['maskS'], 0.)
pt1, pt2 = ecco_v4_py.get_section_endpoints(section_name)
_, maskW, maskS = ecco_v4_py.get_section_line_masks(pt1, pt2, ds)
expected = ecco_v4_py.section_trsp_at_depth(ds['U'], ds['V'], maskW, maskS,
ds)
ds = ds.drop_vars(['maskW', 'maskS'])
test = ecco_v4_py.section_trsp_at_depth(ds['U'], ds['V'], maskW, maskS)
xr.testing.assert_equal(test['trsp_z'].reset_coords(drop=True),
expected['trsp_z'].reset_coords(drop=True))
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)
print(ecco_monthly_mean.time.isel(time=[0, -1]).values)
# In[11]:
# each monthly mean record is bookended by a snapshot.
#we should have one more snapshot than monthly mean record
print('number of monthly mean records: ', len(ecco_monthly_mean.time))
print('number of monthly snapshot records: ', len(ecco_monthly_snaps.time))
# ## Create the xgcm 'grid' object
#
# the xgcm grid object makes it easy to make flux divergence calculations across different tiles of the lat-lon-cap grid.
# In[12]:
ecco_xgcm_grid = ecco.get_llc_grid(ecco_grid)
ecco_xgcm_grid
# ## Calculate LHS: $\eta$ time tendency: $G_{\text{total tendency}}$
#
# We calculate the monthly-averaged time tendency of ``ETAN`` by differencing monthly ``ETAN`` snapshots.
# Subtract the numpy arrays $\eta(t+1)$ - $\eta(t)$. This operation gives us $\Delta$ ``ETAN`` $/ \Delta$ t (month) records.
# In[13]:
num_months = len(ecco_monthly_snaps.time)
G_total_tendency_month = ecco_monthly_snaps.ETAN.isel(
time=range(1, num_months)).values - ecco_monthly_snaps.ETAN.isel(
time=range(0, num_months - 1)).values
# The result is a numpy array of 264 months
# Recall that for tiles 7-12, the y-dimension actually runs East-West. # Therefore, we want to compute a finite difference in the x-dimension on these tiles to get the latitude band at 26$^\circ$N. # For tiles 1-5, we clearly want to difference in the y-dimension. # Things get more complicated on tile 6. # Here we make the [xgcm Grid object](https://xgcm.readthedocs.io/en/latest/api.html#grid) which allows us to compute finite differences in simple one liners. # This object understands how each of the tiles on the LLC grid connect, because we provide that information. # To see under the hood, checkout the utility function [get_llc_grid](https://github.com/ECCO-GROUP/ECCOv4-py/blob/master/ecco_v4_py/ecco_utils.py) where these connections are defined. # In[10]: ecco_ds # In[11]: grid = ecco.get_llc_grid(ecco_ds) # In[12]: lat_maskW = grid.diff(dome_maskC, 'X', boundary='fill') lat_maskS = grid.diff(dome_maskC, 'Y', boundary='fill') # In[13]: plt.figure(figsize=(12, 6)) ecco.plot_tiles(lat_maskW + maskW.isel(k=0), cmap='viridis') plt.show() # In[14]: plt.figure(figsize=(12, 6))
def test_get_grid(get_test_ds):
"""make sure we can make a grid ... that's it"""
grid = ecco_v4_py.get_llc_grid(get_test_ds)
'THETA': 'THETA_snp'
})
])
# ### Create the xgcm 'grid' object
#
# The `xgcm` 'grid' object is used to calculate the flux divergences across different tiles of the lat-lon-cap grid and the time derivatives from ``THETA`` snapshots
# In[17]:
# Change time axis of the snapshot variables
ds.time_snp.attrs['c_grid_axis_shift'] = 0.5
# In[18]:
grid = ecco.get_llc_grid(ds)
# ### Number of seconds in each month
# The xgcm `grid` object includes information on the time axis, such that we can use it to get $\Delta t$, which is the time span between the beginning and end of each month (in seconds).
# In[19]:
delta_t = grid.diff(ds.time_snp, 'T', boundary='fill', fill_value=np.nan)
# Convert to seconds
delta_t = delta_t.astype('f4') / 1e9
# ## Calculate total tendency of $\theta$ ($G^{\theta}_\textrm{total}$)
#
# We calculate the monthly-averaged time tendency of ``THETA`` by differencing monthly ``THETA`` snapshots. Remember that we need to include a scaling factor due to the nonlinear free surface formulation. Thus, we need to use snapshots of both `ETAN` and `THETA` to evaluate $s^*\theta$.
dask_chunk=True)
date_last_record = \
ecco.extract_yyyy_mm_dd_hh_mm_ss_from_datetime64(ds_ecco_monthly_snaps.time[-1].values)
# load monthly mean data
ds_ecco_monthly_mean = \
ecco.recursive_load_ecco_var_from_years_nc(dir_monthly,\
vars_to_load=['ETAN',\
'oceFWflx',\
'UVELMASS',\
'VVELMASS',\
'WVELMASS'],\
years_to_load='all',\
dask_chunk=True)
grid_ecco_xgcm = ecco.get_llc_grid(ds_grid)
# calculte total tendency
num_months = len(ds_ecco_monthly_mean.time)
G_total_tendency_month = ds_ecco_monthly_mean.ETAN.isel(time=range(1,num_months)).values - ds_ecco_monthly_mean.ETAN.isel(time=range(0,num_months-1)).values
tmp = ds_ecco_monthly_mean.ETAN.isel(time=range(0,num_months-1)).copy(deep=True)
tmp.name = 'G_total_tendency_month'
tmp.values = G_total_tendency_month
G_total_tendency_month = tmp
t = []
for year in range(year_start,year_end+1):
for month in range(1,13):
t = np.append(t,date(year,month,1).toordinal())
seconds_per_month = 3600 * 24 * (t[1:]-t[0:len(t)-1])