Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add HU, HT, KMU to get grid #54

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Add HU, HT, KMU to get grid #54

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4059fff
add HT, HU, KMU to grid output
bradyrx May 29, 2020
24f115e
add tests for HT, HU, KMU on get_grid
bradyrx May 29, 2020
376eadb
Use numba for KMU calculation
bradyrx May 29, 2020
c44a8c8
initialize KMU outside of function
bradyrx May 30, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 44 additions & 0 deletions pop_tools/grid.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,12 @@ def get_grid(grid_name, scrip=False):
assert kmt_flat.max() <= len(z_t), 'Max KMT > length z_t'
KMT = kmt_flat.reshape(grid_attrs['lateral_dims']).astype(np.int32)

# derive KMU
KMU = np.zeros_like(KMT)
for i in range(KMT.shape[0]):
for j in range(KMT.shape[1]):
KMU[i, j] = min(KMT[i, j], KMT[i - 1, j], KMT[i, j - 1], KMT[i - 1, j - 1])

# read REGION_MASK
region_mask_fname = INPUTDATA.fetch(grid_attrs['region_mask_fname'], downloader=downloader)
region_mask_flat = np.fromfile(region_mask_fname, dtype='>i4', count=-1)
Expand All @@ -140,6 +146,15 @@ def get_grid(grid_name, scrip=False):
), f'unexpected dims in region_mask file: {grid_attrs["region_mask_fname"]}'
REGION_MASK = region_mask_flat.reshape(grid_attrs['lateral_dims']).astype(np.int32)

# derive depth of columns of ocean T-points and U-points
KMT_reidx = KMT - 1
KMT_reidx[KMT_reidx == -1] = 0
Comment on lines +148 to +149
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This deals with the issues @klindsay28 mentioned about IndexErrors. Since the index from KMT is not zero-based.

Question: Would one anticipate the HT/HU output for land to be np.nan or 0? Currently it's 0.

HT = z_w[KMT_reidx]

KMU_reidx = KMU - 1
KMU_reidx[KMU_reidx == -1] = 0
HU = z_w[KMU_reidx]

# output dataset
dso = xr.Dataset()
if scrip:
Expand Down Expand Up @@ -267,6 +282,35 @@ def get_grid(grid_name, scrip=False):
},
)

dso['KMU'] = xr.DataArray(
KMU,
dims=('nlat', 'nlon'),
attrs={
'long_name': 'k Index of Deepest Grid Cell on U Grid',
'coordinates': 'ULONG ULAT',
},
)

dso['HT'] = xr.DataArray(
HT,
dims=('nlat', 'nlon'),
attrs={
'units': 'cm',
'long_name': 'depth of ocean column on T grid',
'coordinates': 'TLONG TLAT',
},
)

dso['HU'] = xr.DataArray(
HU,
dims=('nlat', 'nlon'),
attrs={
'units': 'cm',
'long_name': 'depth of ocean column on U grid',
'coordinates': 'ULONG ULAT',
},
)

dso['REGION_MASK'] = xr.DataArray(
REGION_MASK,
dims=('nlat', 'nlon'),
Expand Down
10 changes: 10 additions & 0 deletions tests/test_grid.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import os

import pytest
import xarray as xr

import pop_tools
Expand All @@ -26,3 +27,12 @@ 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)


@pytest.mark.parametrize('grid', pop_tools.grid_defs.keys())
def test_HT_HU_KMU_in_grid(grid):
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably unnecessary, but I like to add comprehensive testing for any PR. I know this is at odds with the fact that other variables aren't checked. Let me know if we should retain this.

print(grid)
ds = pop_tools.get_grid(grid)
assert 'HT' in ds, 'Missing variable HT'
assert 'HU' in ds, 'Missing variable HU'
assert 'KMU' in ds, 'Missing variable KMU'