Merge pull request #40 from paigem/handle-land-mask

Handle land mask

source/aerobulk/flux.py

@@ -1,7 +1,6 @@
-import functools
-
 import aerobulk.aerobulk.mod_aerobulk_wrap_noskin as aeronoskin
 import aerobulk.aerobulk.mod_aerobulk_wrap_skin as aeroskin
+import numpy as np
 import xarray as xr
 
 VALID_ALGOS = ["coare3p0", "coare3p6", "ecmwf", "ncar", "andreas"]
@@ -13,6 +12,13 @@ def _check_algo(algo, valids):
         raise ValueError(f"Algorithm {algo} not valid. Choose from {valids}.")
 
 
+# Unshrink the data (i.e. put land NaN values back in their correct locations)
+def unshrink_arr(shrunk_array, shape, ocean_index):
+    unshrunk_array = np.full(shape, np.nan)
+    unshrunk_array[ocean_index] = np.squeeze(shrunk_array)
+    return unshrunk_array
+
+
 def noskin_np(
     sst,
     t_zt,
@@ -27,6 +33,7 @@ def noskin_np(
 ):
     """Python wrapper for aerobulk without skin correction.
     !ATTENTION If input not provided in correct units, will crash.
+    !ATTENTION Missing values taken from NaN values in sst field. No input variables may have NaN values that are not reflected in the sst.
 
     Parameters
     ----------
@@ -69,16 +76,20 @@ def noskin_np(
     evap : numpy.array
         evaporation         [mm/s] aka [kg/m^2/s] (usually <0, as ocean loses water!)
     """
-    (
-        ql,
-        qh,
-        taux,
-        tauy,
-        evap,
-    ) = aeronoskin.mod_aerobulk_wrapper_noskin.aerobulk_model_noskin(
-        algo, zt, zu, sst, t_zt, hum_zt, u_zu, v_zu, slp, niter
+
+    # Define the land mask from the native SST land mask
+    ocean_index = np.where(~np.isnan(sst))
+
+    # Shrink the input data (i.e. remove all land points)
+    args_shrunk = tuple(
+        np.atleast_3d(a[ocean_index]) for a in (sst, t_zt, hum_zt, u_zu, v_zu, slp)
+    )
+
+    out_data = aeronoskin.mod_aerobulk_wrapper_noskin.aerobulk_model_noskin(
+        algo, zt, zu, *args_shrunk, niter
     )
-    return ql, qh, taux, tauy, evap
+
+    return tuple(unshrink_arr(o, sst.shape, ocean_index) for o in out_data)
 
 
 def skin_np(
@@ -97,6 +108,7 @@ def skin_np(
 ):
     """Python wrapper for aerobulk with skin correction.
     !ATTENTION If input not provided in correct units, will crash.
+    !ATTENTION Missing values taken from NaN values in sst field. No input variables may have NaN values that are not reflected in the sst.
 
     Parameters
     ----------
@@ -145,51 +157,22 @@ def skin_np(
     evap : numpy.array
         evaporation         [mm/s] aka [kg/m^2/s] (usually <0, as ocean loses water!)
     """
-    (
-        ql,
-        qh,
-        taux,
-        tauy,
-        t_s,
-        evap,
-    ) = aeroskin.mod_aerobulk_wrapper_skin.aerobulk_model_skin(
-        algo, zt, zu, sst, t_zt, hum_zt, u_zu, v_zu, slp, rad_sw, rad_lw, niter
-    )
-    return ql, qh, taux, tauy, t_s, evap
-
-
-def input_and_output_check(func):
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        # Check the input shape
-        test_arg = args[
-            0
-        ]  # assuming that all the input shapes are the same size. TODO: More thorough check
-        if len(test_arg.dims) < 3:
-            # TODO promote using expand_dims?
-            raise NotImplementedError(
-                f"Aerobulk-Python expects all input fields as 3D arrays. Found {len(test_arg.dims)} dimensions on input."
-            )
-        if len(test_arg.dims) > 4:
-            # TODO iterate over extra dims? Or reshape?
-            raise NotImplementedError(
-                f"Aerobulk-Python expects all input fields as 3D arrays. Found {len(test_arg.dims)} dimensions on input."
-            )
+    # Define the land mask from the native SST land mask
+    ocean_index = np.where(~np.isnan(sst))
 
-        out_vars = func(*args, **kwargs)
-
-        # TODO: Here we could 'un-reshape' or squeeze the output according to the logic above
+    # Shrink the input data (i.e. remove all land points)
+    args_shrunk = tuple(
+        np.atleast_3d(a[ocean_index])
+        for a in (sst, t_zt, hum_zt, u_zu, v_zu, slp, rad_sw, rad_lw)
+    )
 
-        if any(var.ndim != 3 for var in out_vars):
-            raise ValueError(
-                f"f2py returned result of unexpected shape. Got {[var.shape for var in out_vars]}"
-            )
-        return out_vars
+    out_data = aeroskin.mod_aerobulk_wrapper_skin.aerobulk_model_skin(
+        algo, zt, zu, *args_shrunk, niter
+    )
 
-    return wrapper
+    return tuple(unshrink_arr(o, sst.shape, ocean_index) for o in out_data)
 
 
-@input_and_output_check
 def noskin(
     sst, t_zt, hum_zt, u_zu, v_zu, slp=101000.0, algo="coare3p0", zt=2, zu=10, niter=6
 ):
@@ -198,6 +181,7 @@ def noskin(
     Warnings
     --------
     !ATTENTION If input not provided in the units shown in [] below the code will crash.
+    !ATTENTION Missing values taken from NaN values in sst field. No input variables may have NaN values that are not reflected in the sst.
 
     Parameters
     ----------
@@ -277,7 +261,6 @@ def noskin(
     return out_vars
 
 
-@input_and_output_check
 def skin(
     sst,
     t_zt,
@@ -298,6 +281,7 @@ def skin(
     Warnings
     --------
     !ATTENTION If input not provided in the units shown in [] below the code will crash.
+    !ATTENTION Missing values taken from NaN values in sst field. No input variables may have NaN values that are not reflected in the sst.
 
     Parameters
     ----------

source/fortran/mod_aerobulk_wrap_skin.pyf

@@ -8,7 +8,7 @@ python module mod_aerobulk_wrap_skin ! in
             use mod_aerobulk, only: aerobulk_init,aerobulk_bye
             use mod_aerobulk_compute
             subroutine aerobulk_model_skin(ni,nj,nt,calgo,zt,zu,sst,t_zt,hum_zt,u_zu,v_zu,slp,rad_sw,rad_lw,niter,ql,qh,tau_x,tau_y,t_s,evap) ! in :mod_aerobulk_wrap_skin:mod_aerobulk_wrap_skin.f90:mod_aerobulk_wrapper_skin
-                threadsafe
+
                 integer, optional,intent(in),check(shape(sst, 0) == ni),depend(sst) :: ni=shape(sst, 0)
                 integer, optional,intent(in),check(shape(sst, 1) == nj),depend(sst) :: nj=shape(sst, 1)
                 integer, optional,intent(in),check(shape(sst, 2) == nt),depend(sst) :: nt=shape(sst, 2)

tests/create_test_data.py

@@ -0,0 +1,51 @@
+from typing import Dict, Tuple
+
+import numpy as np
+import xarray as xr
+from numpy.random import default_rng
+
+
+def create_data(
+    shape: Tuple[int, ...],
+    chunks: Dict[str, int] = {},
+    skin_correction: bool = False,
+    order: str = "F",
+    use_xr=True,
+    land_mask=False,
+):
+    size = shape[0] * shape[1]
+    shape2d = (shape[0], shape[1])
+    rng = default_rng()
+    indices = rng.choice(size, size=int(size * 0.3), replace=False)
+    multi_indices = np.unravel_index(indices, shape2d)
+
+    def _arr(value, chunks, order):
+        arr = np.full(shape, value, order=order)
+        # adds random noise scaled by a percentage of the value
+        randomize_factor = 0.001
+        randomize_range = value * randomize_factor
+        noise = np.random.rand(*shape) * randomize_range
+        arr = arr + noise
+
+        if land_mask:
+            arr[
+                multi_indices[0], multi_indices[1], :
+            ] = np.nan  # add NaNs to mimic land mask
+        if use_xr:
+            arr = xr.DataArray(arr)
+            if chunks:
+                arr = arr.chunk(chunks)
+        return arr
+
+    sst = _arr(290.0, chunks, order)
+    t_zt = _arr(280.0, chunks, order)
+    hum_zt = _arr(0.001, chunks, order)
+    u_zu = _arr(1.0, chunks, order)
+    v_zu = _arr(-1.0, chunks, order)
+    slp = _arr(101000.0, chunks, order)
+    rad_sw = _arr(0.000001, chunks, order)
+    rad_lw = _arr(350.0, chunks, order)
+    if skin_correction:
+        return sst, t_zt, hum_zt, u_zu, v_zu, rad_sw, rad_lw, slp
+    else:
+        return sst, t_zt, hum_zt, u_zu, v_zu, slp

tests/test_flux_np.py

@@ -0,0 +1,52 @@
+import numpy as np
+import pytest
+from aerobulk.flux import noskin_np, skin_np
+from create_test_data import create_data
+
+"""Tests for the numpy land_mask wrapper"""
+
+
+@pytest.mark.parametrize(
+    "algo, skin_correction",
+    [
+        ("ecmwf", True),
+        ("ecmwf", False),
+        ("coare3p0", True),
+        ("coare3p0", False),
+        ("coare3p6", True),
+        ("coare3p6", False),
+        ("andreas", False),
+        ("ncar", False),
+    ],
+)
+def test_land_mask(skin_correction, algo):
+    shape = (2, 3, 4)
+    size = shape[0] * shape[1] * shape[2]
+
+    if skin_correction:
+        func = skin_np
+    else:
+        func = noskin_np
+
+    args = create_data(
+        shape,
+        chunks=False,
+        skin_correction=skin_correction,
+        use_xr=False,
+        land_mask=True,
+    )
+    out_data = func(*args, algo, 2, 10, 6)
+
+    # Check the location of all NaNs is correct
+    for o in out_data:
+        np.testing.assert_allclose(np.isnan(args[0]), np.isnan(o))
+
+    # Check that values of the unshrunk array are correct
+    for i in range(size):
+        index = np.unravel_index(i, shape)
+        if not np.isnan(out_data[0][index]):
+            single_inputs = tuple(np.atleast_3d(i[index]) for i in args)
+
+            single_outputs = func(*single_inputs, algo, 2, 10, 6)
+            for so, o in zip(single_outputs, out_data):
+                assert so == o[index]

tests/test_flux_xr.py

@@ -1,44 +1,11 @@
-from typing import Dict, Tuple
-
-import numpy as np
 import pytest
 import xarray as xr
 from aerobulk import noskin, skin
+from create_test_data import create_data
 
 """Tests for the xarray wrapper"""
 
 
-def create_data(
-    shape: Tuple[int, ...],
-    chunks: Dict[str, int] = {},
-    skin_correction: bool = False,
-    order: str = "F",
-):
-    def _arr(value, chunks, order):
-        arr = xr.DataArray(np.full(shape, value, order=order))
-
-        # adds random noise scaled by a percentage of the value
-        randomize_factor = 0.001
-        randomize_range = value * randomize_factor
-        arr = arr + np.random.rand(*shape) + randomize_range
-        if chunks:
-            arr = arr.chunk(chunks)
-        return arr
-
-    sst = _arr(290.0, chunks, order)
-    t_zt = _arr(280.0, chunks, order)
-    hum_zt = _arr(0.001, chunks, order)
-    u_zu = _arr(1.0, chunks, order)
-    v_zu = _arr(-1.0, chunks, order)
-    slp = _arr(101000.0, chunks, order)
-    rad_sw = _arr(0.000001, chunks, order)
-    rad_lw = _arr(350, chunks, order)
-    if skin_correction:
-        return sst, t_zt, hum_zt, u_zu, v_zu, rad_sw, rad_lw, slp
-    else:
-        return sst, t_zt, hum_zt, u_zu, v_zu, slp
-
-
 @pytest.mark.parametrize("algo", ["wrong"])
 def test_algo_error_noskin(algo):
     shape = (1, 1, 1)
@@ -99,3 +66,21 @@ def test_transpose_invariance_noskin(self, chunks, skin_correction, order):
                 ii.transpose("dim_0", "dim_1", "dim_2"),
                 iii.transpose("dim_0", "dim_1", "dim_2"),
             )
+
+
+@pytest.mark.parametrize("skin_correction", [True, False])
+def test_all_input_array_sizes_valid(skin_correction):
+    shapes = (
+        (3, 4),
+        (2, 3, 4),
+        (2, 3, 4, 5),
+    )  # create_data() only allows for inputs of 2 or more dimensions
+    data = (create_data(s, skin_correction=skin_correction) for s in shapes)
+    if skin_correction:
+        func = skin
+    else:
+        func = noskin
+    tuple(func(*d, "coare3p0", 2, 10, 6) for d in data)
+    assert (
+        1 == 1
+    )  # This line is always true, but verifies that the above line doesn't crash the Fortran code