Handle land mask

source/aerobulk/flux.py

@@ -2,6 +2,7 @@
 
 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 +14,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,
@@ -69,16 +77,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(
@@ -145,17 +157,20 @@ 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
+    # 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, rad_sw, rad_lw)
     )
-    return ql, qh, taux, tauy, t_s, evap
+
+    out_data = aeroskin.mod_aerobulk_wrapper_skin.aerobulk_model_skin(
+        algo, zt, zu, *args_shrunk, niter
+    )
+
+    return tuple(unshrink_arr(o, sst.shape, ocean_index) for o in out_data)
 
 
 def input_and_output_check(func):

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)
+        if use_xr:
+            arr = xr.DataArray(arr)
+        if land_mask:
+            arr[
+                multi_indices[0], multi_indices[1], :
+            ] = np.nan  # add NaNs to mimic land mask
+        if chunks:
+            arr = arr.chunk(chunks)
+
+        # 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
+        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,40 @@
+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("skin_correction", [True, False])
+def test_land_mask(skin_correction):
+    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, "ecmwf", 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, "ecmwf", 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)