def test_bad_schema():
region_defs = {
'NAtl-STG': {
'match': {
'REGION_MASK': [6]
},
'bounds': {
'TLAT': [32.0, 42.0]
}
}
}
with pytest.raises(AssertionError):
pop_tools.region_mask_3d('POP_gx1v7', region_defs=region_defs)
def test_make_masks():
for grid in pop_tools.grid_defs.keys():
region_masks = pop_tools.list_region_masks(grid)
for region_mask in region_masks:
mask3d = pop_tools.region_mask_3d(grid, mask_name=region_mask)
assert isinstance(mask3d, xr.DataArray)
assert mask3d.dims == ('region', 'nlat', 'nlon')
def test_user_defined_mask():
region_defs = {
'NAtl-STG': [{
'match': {
'REGION_MASK': [6]
},
'bounds': {
'TLAT': [32.0, 42.0],
'TLONG': [310.0, 350.0]
},
}],
'NAtl-SPG': [{
'match': {
'REGION_MASK': [6]
},
'bounds': {
'TLAT': [50.0, 60.0],
'TLONG': [310.0, 350.0]
},
}],
}
mask3d = pop_tools.region_mask_3d('POP_gx1v7', region_defs=region_defs)
assert isinstance(mask3d, xr.DataArray)
assert mask3d.dims == ('region', 'nlat', 'nlon')
def compute_regional_integrated_MHT(MHT_vertical, basin, longitude, latitude,
gy) -> float:
'''
Integrates the vertically integrated MHT in a basin across a given latitude
Regions are defined by poptools. options are ['Black Sea', 'Baltic Sea', 'Red Sea', 'Southern Ocean', 'Pacific Ocean',
'Indian Ocean', 'Persian Gulf', 'Atlantic Ocean', 'Mediterranean Sea',
'Lab. Sea & Baffin Bay', 'GIN Seas', 'Arctic Ocean', 'Hudson Bay']
Grid must be gx1v6.
gy is the latitude of interest (where we calculate MHT)
MHT_vertical must be dimensions of time x lat x lon
'''
# get regions with pop tools
grid_name = 'POP_gx1v6'
ds = pop_tools.get_grid(grid_name)
TLAT = ds.TLAT
TLONG = ds.TLONG
# raise an error if basin is not an option
# raise an error if MHT_vertical is not right shape
#if (MHT_vertical.ndim < 2):
# raise ValueError("MHT_vertical is not at least rank-2")
#if (MHT_vertical.shape[2] != longitude.shape[0]):
# raise ValueError("MHT_vertical axis 1 is not shape of longitude")
#if (MHT_vertical.shape[1] != latitude.shape[0]):
# raise ValueError("MHT_vertical axis 2 is not shape of latitude")
# raise an error if latitude of choice is not in the basin
# select grid. region will be ones, all else will be zeros. can view regions at: https://pop-tools.readthedocs.io/en/latest/examples/re gion-mask.html#Alternative-region-masks
mask3d = pop_tools.region_mask_3d(grid_name,
mask_name='Pacific-Indian-Atlantic')
mask2d = mask3d.sel(region=basin)
# interpolate this to the grid we have
[xx, yy] = np.meshgrid(longitude, latitude)
m = griddata((TLONG.values.flatten(), TLAT.values.flatten()),
mask2d.values.flatten(), (xx, yy),
method='nearest')
for i in range(0, MHT_vertical.shape[0]): # loop over timesteps
MHT_vertical[i, ...] = MHT_vertical[
i, ...] * m # regions not in our basin become zero
# find latitude closest to the one we asked for, integrate across
idy = np.searchsorted(latitude, gy)
# find distance between points at this latitude
p1 = (latitude[idy], longitude[0])
p2 = (latitude[idy], longitude[1]
) # should be the same at any given longitude -> check
dx = geodesic(p1, p2).km * 1000 # turn into meters instead of km
# integrate across the latitude of choice. need to use cumsum because of nan's
#tmp=np.cumsum(MHT_vertical[:,idy_MHT,idx1:idxend],axis=1)
#MHT=tmp[:,-1]
MHT_vertical = MHT_vertical * dx
MHT = np.sum(MHT_vertical[:, idy, :], axis=1)
return MHT
def get_pop_region_mask_za(
mask_type='3d',
grid_name='POP_gx1v7',
):
"""return a region mask for zonal averaging"""
mask3d = pop_tools.region_mask_3d(grid_name,
mask_name='Pacific-Indian-Atlantic')
nregion = len(mask3d.region)
if mask_type.lower() == '3d':
return mask3d
elif mask_type.lower() == '2d':
mask2d = xr.full_like(mask3d.isel(region=0),
fill_value=0,
dtype=np.int32)
for i in range(
1, nregion
): # skip first index because "za" puts the global field in there
mask2d = xr.where(mask3d.isel(region=i) == 1, i, mask2d)
mask2d.name = 'REGION_MASK'
return mask2d
raise ValueError(
f'unknown mask type: {mask_type}\nexpecting either "2d" or "3d"')
def test_make_default_mask():
for grid in pop_tools.grid_defs.keys():
mask3d = pop_tools.region_mask_3d(grid)
sum_over_region = mask3d.sum('region').values.ravel()
np.testing.assert_equal(sum_over_region[sum_over_region != 0.0], 1.0)
np.testing.assert_equal(sum_over_region[sum_over_region != 1.0], 0.0)
