diff --git a/diagonals/examples/examples-moc.py b/diagonals/examples/examples-moc.py index 70587182e96b71176398fb0dc75982db2098a157..9d9277ed65f3d119d309fe242559125a9e2fe20b 100644 --- a/diagonals/examples/examples-moc.py +++ b/diagonals/examples/examples-moc.py @@ -57,7 +57,7 @@ def load_masks(): global_mask.data[..., -1] = 0.0 basins['Global_Ocean'] = global_mask.data.astype(np.float32) cubes = iris.load(REGIONS_FILE) - for cube in cubes[0:2]: + for cube in cubes: name = cube.name() if name in ('nav_lat', 'nav_lon', 'Caspian_Sea', 'Global_Ocean'): continue diff --git a/diagonals/mesh_helpers/nemo.py b/diagonals/mesh_helpers/nemo.py index 63fd4b766bb9597d523bf45d2da8beaaf7c41795..2896c8422265cca6fec57d68b8872d3cb4ca7e5c 100644 --- a/diagonals/mesh_helpers/nemo.py +++ b/diagonals/mesh_helpers/nemo.py @@ -1,7 +1,7 @@ import numpy as np import iris -from numba import vectorize +from numba import vectorize, float32 class Nemo(): @@ -59,6 +59,14 @@ class Nemo(): def get_mesh_var(self, var, dtype): return iris.load_cube(self.mesh_file, var).data.astype(dtype) + def get_region_mask(self, min_lat, max_lat, min_lon, max_lon, + dtype=np.float32): + lat = self.get_mesh_var('nav_lat', dtype) + lon = self.get_mesh_var('nav_lon', dtype) + mask = (_generate_mask(min_lat, max_lat, lat) + * _generate_mask(min_lon, max_lon, lon)) + return mask + @vectorize(['float32(float32, float32)'], target='cpu') def _get_area(e1, e2): @@ -96,3 +104,12 @@ def _get_volume(e1, e2, e3): Masked array containing the volume for the whole grid. """ return e1 * e2 * e3 + + +@vectorize(['float32(float32, float32, float32)'], nopython=True) +def _generate_mask(min, max, data): + if data < min: + return 0 + if data > max: + return 0 + return 1