Repository: ESMValGroup/ESMValCore

Issue:
Altitude <--> pressure conversion in `extract_levels` fails for masked arrays
**Describe the bug**
The [conversion between altitudes and pressure levels](https://github.com/ESMValGroup/ESMValCore/blob/af5d6ccdc6643c0e416918ff472a087b8c75e9f1/esmvalcore/preprocessor/_regrid.py#L930-L938) in the `extract_levels` preprocessor fails if the input arrays are masked arrays:

```bash
...
File "/home/b/b309141/repos/ESMValCore/esmvalcore/preprocessor/_regrid.py", line 935, in extract_levels
    add_plev_from_altitude(cube)
  File "/home/b/b309141/repos/ESMValCore/esmvalcore/cmor/_fixes/shared.py", line 67, in add_plev_from_altitude
    pressure_points = altitude_to_pressure(height_coord.core_points())
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/scipy/interpolate/_polyint.py", line 79, in __call__
    x, x_shape = self._prepare_x(x)
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/scipy/interpolate/_polyint.py", line 91, in _prepare_x
    x = _asarray_validated(x, check_finite=False, as_inexact=True)
  File "/work/bd0854/b309141/mambaforge/envs/esm/lib/python3.10/site-packages/scipy/_lib/_util.py", line 250, in _asarray_validated
    raise ValueError('masked arrays are not supported')
ValueError: masked arrays are not supported
```

I will come up with a fix for this.

Discussion:

Patch:
diff --git a/esmvalcore/cmor/_fixes/icon/_base_fixes.py b/esmvalcore/cmor/_fixes/icon/_base_fixes.py
index d7b3a496a..f2502ac0a 100644
--- a/esmvalcore/cmor/_fixes/icon/_base_fixes.py
+++ b/esmvalcore/cmor/_fixes/icon/_base_fixes.py
@@ -219,6 +219,7 @@ class IconFix(NativeDatasetFix):
         # Note: we use a lock here to prevent multiple processes from
         # downloading the file simultaneously and to ensure that other
         # processes wait until the download has finished
+        self.CACHE_DIR.mkdir(parents=True, exist_ok=True)
         lock = FileLock(self.CACHE_DIR / f"{grid_name}.lock")
         with lock:
             grid_path = self.CACHE_DIR / grid_name
@@ -241,7 +242,6 @@ class IconFix(NativeDatasetFix):
                 grid_url,
                 grid_path,
             )
-            self.CACHE_DIR.mkdir(parents=True, exist_ok=True)
             with requests.get(grid_url, stream=True,
                               timeout=self.TIMEOUT) as response:
                 response.raise_for_status()
diff --git a/esmvalcore/cmor/_fixes/shared.py b/esmvalcore/cmor/_fixes/shared.py
index 3883520b7..0e4dfead0 100644
--- a/esmvalcore/cmor/_fixes/shared.py
+++ b/esmvalcore/cmor/_fixes/shared.py
@@ -5,6 +5,7 @@ from functools import lru_cache
 
 import dask.array as da
 import iris
+import numpy as np
 import pandas as pd
 from cf_units import Unit
 from iris import NameConstraint
@@ -46,6 +47,33 @@ def add_aux_coords_from_cubes(cube, cubes, coord_dict):
         cubes.remove(coord_cube)
 
 
+def _map_on_filled(function, array):
+    """Map function on filled array."""
+    if array.size == 0:
+        return array
+
+    # We support dask and numpy arrays
+    # Note: numpy's dispatch mechanism does not work here because of the usage
+    # of `np.ma.filled` and `np.ma.masked_array`.
+    if isinstance(array, da.core.Array):
+        num_module = da
+    else:
+        num_module = np
+
+    mask = num_module.ma.getmaskarray(array)
+
+    # Fill masked values with dummy value (simply use first value in array for
+    # this to preserve dtype; these entries get re-masked later so the actual
+    # value does not matter). Note: `array.fill_value` is not defined for dask
+    # arrays, and `ma.filled` also works for regular (non-masked) arrays.
+    fill_value = num_module.ravel(array)[0]
+    array = num_module.ma.filled(array, fill_value)
+
+    # Apply function and return masked array
+    array = function(array)
+    return num_module.ma.masked_array(array, mask=mask)
+
+
 def add_plev_from_altitude(cube):
     """Add pressure level coordinate from altitude coordinate.
 
@@ -64,11 +92,15 @@ def add_plev_from_altitude(cube):
         if height_coord.units != 'm':
             height_coord.convert_units('m')
         altitude_to_pressure = get_altitude_to_pressure_func()
-        pressure_points = altitude_to_pressure(height_coord.core_points())
+        pressure_points = _map_on_filled(
+            altitude_to_pressure, height_coord.core_points()
+        )
         if height_coord.core_bounds() is None:
             pressure_bounds = None
         else:
-            pressure_bounds = altitude_to_pressure(height_coord.core_bounds())
+            pressure_bounds = _map_on_filled(
+                altitude_to_pressure, height_coord.core_bounds()
+            )
         pressure_coord = iris.coords.AuxCoord(pressure_points,
                                               bounds=pressure_bounds,
                                               standard_name='air_pressure',
@@ -98,11 +130,15 @@ def add_altitude_from_plev(cube):
         if plev_coord.units != 'Pa':
             plev_coord.convert_units('Pa')
         pressure_to_altitude = get_pressure_to_altitude_func()
-        altitude_points = pressure_to_altitude(plev_coord.core_points())
+        altitude_points = _map_on_filled(
+            pressure_to_altitude, plev_coord.core_points()
+        )
         if plev_coord.core_bounds() is None:
             altitude_bounds = None
         else:
-            altitude_bounds = pressure_to_altitude(plev_coord.core_bounds())
+            altitude_bounds = _map_on_filled(
+                pressure_to_altitude, plev_coord.core_bounds()
+            )
         altitude_coord = iris.coords.AuxCoord(altitude_points,
                                               bounds=altitude_bounds,
                                               standard_name='altitude',


Test patch:
diff --git a/tests/integration/cmor/_fixes/test_shared.py b/tests/integration/cmor/_fixes/test_shared.py
index e82005f59..92d2cfba2 100644
--- a/tests/integration/cmor/_fixes/test_shared.py
+++ b/tests/integration/cmor/_fixes/test_shared.py
@@ -1,4 +1,5 @@
 """Tests for shared functions for fixes."""
+import dask.array as da
 import iris
 import iris.coords
 import iris.cube
@@ -8,6 +9,7 @@ from cf_units import Unit
 from iris import NameConstraint
 
 from esmvalcore.cmor._fixes.shared import (
+    _map_on_filled,
     add_altitude_from_plev,
     add_aux_coords_from_cubes,
     add_plev_from_altitude,
@@ -103,18 +105,102 @@ def test_add_aux_coords_from_cubes(coord_dict, output):
         assert "got 2" in str(err.value)
 
 
+def test_map_on_filled_np_empty_array():
+    """Test `_map_on_filled` with empty numpy array."""
+    array_in = np.array([])
+    output = _map_on_filled(lambda x: x**2, array_in)
+    assert isinstance(output, np.ndarray)
+    assert not np.ma.isMaskedArray(output)
+    np.testing.assert_equal(output, [])
+
+
+def test_map_on_filled_np_no_mask():
+    """Test `_map_on_filled` with non-masked numpy array."""
+    array_in = np.array([1, 2, 3])
+    output = _map_on_filled(lambda x: x**2, array_in)
+    assert isinstance(output, np.ndarray)
+    assert np.ma.isMaskedArray(output)
+    np.testing.assert_equal(output, [1, 4, 9])
+
+
+def test_map_on_filled_np_mask():
+    """Test `_map_on_filled` with masked numpy array."""
+    array_in = np.ma.masked_equal([999.0, 4.0, 999.0], 999.0)
+    output = _map_on_filled(lambda x: x**2, array_in)
+    assert isinstance(output, np.ndarray)
+    assert np.ma.isMaskedArray(output)
+    np.testing.assert_allclose(output, [999.0, 16.0, 999.0])
+    np.testing.assert_equal(output.mask, [True, False, True])
+
+
+def test_map_on_filled_np_mask_not_used():
+    """Test `_map_on_filled` with masked numpy array."""
+    array_in = np.ma.masked_equal([2, 4, 5], 10)
+    output = _map_on_filled(lambda x: x**2, array_in)
+    assert isinstance(output, np.ndarray)
+    assert np.ma.isMaskedArray(output)
+    np.testing.assert_allclose(output, [4, 16, 25])
+    np.testing.assert_equal(output.mask, [False, False, False])
+
+
+def test_map_on_filled_da_empty_array():
+    """Test `_map_on_filled` with empty dask array."""
+    array_in = da.array([])
+    output = _map_on_filled(lambda x: x**2, array_in)
+    assert isinstance(output, da.core.Array)
+    output = output.compute()
+    assert not np.ma.isMaskedArray(output)
+    np.testing.assert_equal(output, [])
+
+
+def test_map_on_filled_da_no_mask():
+    """Test `_map_on_filled` with non-masked dask array."""
+    array_in = da.arange(3)
+    output = _map_on_filled(lambda x: x**2, array_in)
+    assert isinstance(output, da.core.Array)
+    output = output.compute()
+    assert np.ma.isMaskedArray(output)
+    np.testing.assert_equal(output, [0, 1, 4])
+
+
+def test_map_on_filled_da_mask():
+    """Test `_map_on_filled` with masked dask array."""
+    array_in = da.ma.masked_equal(da.arange(3) + 3, 3)
+    output = _map_on_filled(lambda x: x**2, array_in)
+    assert isinstance(output, da.core.Array)
+    output = output.compute()
+    assert np.ma.isMaskedArray(output)
+    np.testing.assert_equal(output, [3, 16, 25])
+    np.testing.assert_equal(output.mask, [True, False, False])
+
+
+def test_map_on_filled_da_mask_not_used():
+    """Test `_map_on_filled` with masked dask array."""
+    array_in = da.ma.masked_equal(da.arange(3), 10)
+    output = _map_on_filled(lambda x: x**2, array_in)
+    assert isinstance(output, da.core.Array)
+    output = output.compute()
+    assert np.ma.isMaskedArray(output)
+    np.testing.assert_allclose(output, [0, 1, 4])
+    np.testing.assert_equal(output.mask, [False, False, False])
+
+
 ALT_COORD = iris.coords.AuxCoord([0.0], bounds=[[-100.0, 500.0]],
                                  standard_name='altitude', units='m')
+ALT_COORD_MASKED = ALT_COORD.copy(np.ma.masked_equal([0.0], 0.0))
 ALT_COORD_NB = iris.coords.AuxCoord([0.0], standard_name='altitude', units='m')
 ALT_COORD_KM = iris.coords.AuxCoord([0.0], bounds=[[-0.1, 0.5]],
                                     var_name='alt', long_name='altitude',
                                     standard_name='altitude', units='km')
 P_COORD = iris.coords.AuxCoord([101325.0], bounds=[[102532.0, 95460.8]],
                                standard_name='air_pressure', units='Pa')
+P_COORD_MASKED = P_COORD.copy(np.ma.masked_equal([0.0], 0.0))
 P_COORD_NB = iris.coords.AuxCoord([101325.0], standard_name='air_pressure',
                                   units='Pa')
 CUBE_ALT = iris.cube.Cube([1.0], var_name='x',
                           aux_coords_and_dims=[(ALT_COORD, 0)])
+CUBE_ALT_MASKED = iris.cube.Cube([1.0], var_name='x',
+                                 aux_coords_and_dims=[(ALT_COORD_MASKED, 0)])
 CUBE_ALT_NB = iris.cube.Cube([1.0], var_name='x',
                              aux_coords_and_dims=[(ALT_COORD_NB, 0)])
 CUBE_ALT_KM = iris.cube.Cube([1.0], var_name='x',
@@ -123,6 +209,7 @@ CUBE_ALT_KM = iris.cube.Cube([1.0], var_name='x',
 
 TEST_ADD_PLEV_FROM_ALTITUDE = [
     (CUBE_ALT.copy(), P_COORD.copy()),
+    (CUBE_ALT_MASKED.copy(), P_COORD_MASKED.copy()),
     (CUBE_ALT_NB.copy(), P_COORD_NB.copy()),
     (CUBE_ALT_KM.copy(), P_COORD.copy()),
     (iris.cube.Cube(0.0), None),
@@ -142,7 +229,24 @@ def test_add_plev_from_altitude(cube, output):
     assert not cube.coords('air_pressure')
     add_plev_from_altitude(cube)
     air_pressure_coord = cube.coord('air_pressure')
-    assert air_pressure_coord == output
+    metadata_list = [
+        'var_name',
+        'standard_name',
+        'long_name',
+        'units',
+        'attributes',
+    ]
+    for attr in metadata_list:
+        assert getattr(air_pressure_coord, attr) == getattr(output, attr)
+    np.testing.assert_allclose(
+        air_pressure_coord.points, output.points, atol=1e-7
+    )
+    if output.bounds is None:
+        assert air_pressure_coord.bounds is None
+    else:
+        np.testing.assert_allclose(
+            air_pressure_coord.bounds, output.bounds, rtol=1e-3
+        )
     assert cube.coords('altitude')
 
 
@@ -152,6 +256,8 @@ P_COORD_HPA = iris.coords.AuxCoord([1013.25], bounds=[[1025.32, 954.60]],
                                    long_name='pressure', units='hPa')
 CUBE_PLEV = iris.cube.Cube([1.0], var_name='x',
                            aux_coords_and_dims=[(P_COORD, 0)])
+CUBE_PLEV_MASKED = iris.cube.Cube([1.0], var_name='x',
+                                  aux_coords_and_dims=[(P_COORD_MASKED, 0)])
 CUBE_PLEV_NB = iris.cube.Cube([1.0], var_name='x',
                               aux_coords_and_dims=[(P_COORD_NB, 0)])
 CUBE_PLEV_HPA = iris.cube.Cube([1.0], var_name='x',
@@ -160,6 +266,7 @@ CUBE_PLEV_HPA = iris.cube.Cube([1.0], var_name='x',
 
 TEST_ADD_ALTITUDE_FROM_PLEV = [
     (CUBE_PLEV.copy(), ALT_COORD.copy()),
+    (CUBE_PLEV_MASKED.copy(), ALT_COORD_MASKED.copy()),
     (CUBE_PLEV_NB.copy(), ALT_COORD_NB.copy()),
     (CUBE_PLEV_HPA.copy(), ALT_COORD.copy()),
     (iris.cube.Cube(0.0), None),