EPPT-2721 contrails output categorical data (#2209)

* Function to output categorical (integer) array from input boolean arrays

* Unit test

* Initial function to generate categorical cube

* Initial test for categorical cube output

* Categorical cube test returns expected shape and data

* Check cube contents and metadata in test

* todo

* Refactoring of class to reduce argument lists, increase use of self, and tidy up docstrings

* Changes to unit tests to account for function signature changes. Because SVP ice table is now used within the persistency function, this required changing the parametrisation of the combination test. Every unique combination (16 total) now cannot be tested easily, so we have settled for most of the combinations.

* Array shape checking in process_from_arrays

* Cube input error checks. Minor refactoring of input array checks.

* minor edit

* Update categorical metadata

* Category check added to cube output unit test

* remove print

* Cube input check test

* Initial version of test to check that pressure, temperature and relative humidity inputs can produce the expected result from an Appleman diagram.

* Modify test to show in-progress. To be completed once work from contrail persistency fix has been merged.

* Added contrail class to API

* Changed categorical output dtype to int32

* End-to-end test that checks values of output categorical cube, from input temperature and humidity cubes

* Changed hard-coded freezing temperature to use absolute zero constant

Author: ryan-cocking-mo

improver/api/__init__.py

@@ -37,6 +37,7 @@
     "CloudCondensationLevel": "improver.psychrometric_calculations.cloud_condensation_level",
     "CloudTopTemperature": "improver.psychrometric_calculations.cloud_top_temperature",
     "Combine": "improver.cube_combiner",
+    "CondensationTrailFormation": "improver.psychrometric_calculations.condensation_trails",
     "ConstructReliabilityCalibrationTables": "improver.calibration.reliability_calibration",
     "ContinuousRankedProbabilityScoreMinimisers": "improver.calibration.emos_calibration",
     "ContrailHeightExtractor": "improver.psychrometric_calculations.condensation_trails",

improver/psychrometric_calculations/condensation_trails.py

@@ -7,18 +7,28 @@
 from typing import Tuple, Union
 
 import numpy as np
+from iris.coords import DimCoord
 from iris.cube import Cube, CubeList
 
 from improver import BasePlugin
-from improver.constants import EARTH_REPSILON
+from improver.categorical.utilities import categorical_attributes
+from improver.constants import ABSOLUTE_ZERO, EARTH_REPSILON
 from improver.generate_ancillaries.generate_svp_derivative_table import (
     SaturatedVapourPressureDerivativeTable,
     SaturatedVapourPressureTable,
 )
+from improver.metadata.utilities import (
+    create_new_diagnostic_cube,
+    generate_mandatory_attributes,
+)
 from improver.psychrometric_calculations.psychrometric_calculations import (
     calculate_svp_in_air,
 )
 from improver.utilities.common_input_handle import as_cubelist
+from improver.utilities.cube_manipulation import (
+    add_coordinate_to_cube,
+    get_dim_coord_names,
+)
 
 
 class CondensationTrailFormation(BasePlugin):
@@ -48,6 +58,8 @@ class CondensationTrailFormation(BasePlugin):
     engine_mixing_ratios = None
     critical_temperatures = None
     critical_intercepts = None
+    nonpersistent_contrails = None
+    persistent_contrails = None
 
     def __init__(self, engine_contrail_factors: list = [3e-5, 3.4e-5, 3.9e-5]):
         """Initialises the Class
@@ -208,7 +220,7 @@ def _critical_temperatures_and_intercepts_for_given_contrail_factor(
         )
         return critical_temperature, critical_intercept
 
-    def _calculate_critical_temperatures_and_intercepts(self):
+    def _calculate_critical_temperatures_and_intercepts(self) -> None:
         """Calculate the critical temperatures and intercepts on pressure levels for all engine contrail factors."""
         self.critical_temperatures = np.zeros(
             self.engine_mixing_ratios.shape[:2] + self.relative_humidity.shape[1:],
@@ -218,11 +230,9 @@ def _calculate_critical_temperatures_and_intercepts(self):
             self.engine_mixing_ratios.shape[:2], dtype=np.float32
         )
 
-        svp_table = SaturatedVapourPressureTable(
-            183.15, 253.15, water_only=True
-        ).process()
+        svp_table = SaturatedVapourPressureTable(water_only=True).process()
         svp_derivative_table = SaturatedVapourPressureDerivativeTable(
-            183.15, 253.15, water_only=True
+            water_only=True
         ).process()
 
         for i, engine_mixing_ratio_for_contrail_factor in enumerate(
@@ -236,22 +246,11 @@ def _calculate_critical_temperatures_and_intercepts(self):
                 )
             )
 
-    def _calculate_contrail_persistency(
-        self,
-        saturated_vapour_pressure_ice: np.ndarray,
-    ) -> Tuple[np.ndarray, np.ndarray]:
+    def _calculate_contrail_persistency(self) -> None:
         """
         Apply four conditions to determine whether non-persistent or persistent contrails will form.
 
         .. include:: extended_documentation/psychrometric_calculations/condensation_trails/formation_conditions.rst
-
-        Args:
-            saturated_vapour_pressure_ice: The saturated vapour pressure with respect to ice, on pressure
-                levels. Pressure is the leading axis (Pa).
-
-        Returns:
-            Two boolean arrays that state whether 'non-persistent' or 'persistent' contrails will form, respectively.
-            Array axes are [contrail factor, pressure level, latitude, longitude].
         """
 
         def reshape_and_broadcast(arr, target_shape):
@@ -271,6 +270,12 @@ def reshape_and_broadcast(arr, target_shape):
             self.critical_intercepts, self.critical_temperatures.shape
         )
 
+        # saturated vapour pressure with respect to ice, on pressure levels
+        svp_table = SaturatedVapourPressureTable(ice_only=True).process()
+        saturated_vapour_pressure_ice = np.interp(
+            self.temperature, svp_table.coord("air_temperature").points, svp_table.data
+        )
+
         # Condition 1
         vapour_pressure_above_threshold = (
             self.local_vapour_pressure[np.newaxis]
@@ -284,20 +289,79 @@ def reshape_and_broadcast(arr, target_shape):
         # Condition 3
         air_is_saturated = self.local_vapour_pressure > saturated_vapour_pressure_ice
         # Condition 4
-        temperature_below_freezing = self.temperature < 273.15
+        temperature_below_freezing = self.temperature < abs(ABSOLUTE_ZERO)
 
-        nonpersistent_contrails = (
+        # Boolean arrays that are true when the specific contrail type will form.
+        # Array axes are [contrail factor, pressure level, latitude, longitude].
+        self.nonpersistent_contrails = (
             vapour_pressure_above_threshold
             & temperature_below_threshold
             & ~(air_is_saturated & temperature_below_freezing)
         )
-        persistent_contrails = (
+        self.persistent_contrails = (
             vapour_pressure_above_threshold
             & temperature_below_threshold
             & air_is_saturated
             & temperature_below_freezing
         )
-        return nonpersistent_contrails, persistent_contrails
+
+    def _boolean_to_categorical(self) -> np.ndarray:
+        """
+        Combine two boolean arrays of contrail persistency into a single categorical array of contrail formation.
+
+        Returns:
+            Array of categorical (integer) data, where 0 = no contrails, 1 = non-persistent contrails and 2 = persistent
+            contrails.
+        """
+        categorical = np.where(
+            self.nonpersistent_contrails & ~self.persistent_contrails, 1, 0
+        )
+        categorical = np.where(
+            ~self.nonpersistent_contrails & self.persistent_contrails, 2, categorical
+        )
+        return categorical
+
+    def _create_contrail_formation_cube(
+        self, categorical_data: np.ndarray, template_cube: Cube
+    ) -> Cube:
+        """
+        Create a contrail formation cube, populated with categorical data.
+
+        Args:
+            categorical_data: Categorical (integer) data of contrail formation. Leading axes are [contrail factor,
+                pressure level].
+            template_cube: Cube from which to derive dimensions, coordinates and mandatory attributes.
+
+        Returns:
+            Categorical cube of contrail formation, where 0 = no contrails, 1 = non-persistent contrails and
+            2 = persistent contrails. Has the same shape as categorical_data.
+        """
+        contrail_factor_coord = DimCoord(
+            points=self._engine_contrail_factors,
+            var_name="engine_contrail_factor",
+            units="kg kg-1 K-1",
+        )
+        template_cube = add_coordinate_to_cube(
+            template_cube, new_coord=contrail_factor_coord
+        )
+        mandatory_attributes = generate_mandatory_attributes([template_cube])
+
+        decision_tree = {
+            "meta": {"name": "contrail_type"},
+            "None": {"leaf": 0},
+            "Non-persistent": {"leaf": 1},
+            "Persistent": {"leaf": 2},
+        }
+        optional_attributes = categorical_attributes(decision_tree, "contrail_type")
+
+        return create_new_diagnostic_cube(
+            name="contrail_type",
+            units="1",
+            template_cube=template_cube,
+            mandatory_attributes=mandatory_attributes,
+            optional_attributes=optional_attributes,
+            data=categorical_data.astype(np.int32),
+        )
 
     def process_from_arrays(
         self,
@@ -306,7 +370,7 @@ def process_from_arrays(
         pressure_levels: np.ndarray,
     ) -> np.ndarray:
         """
-        Main entry point of this class for data as Numpy arrays
+        Main entry point of this class for data as Numpy arrays.
 
         Process the temperature, humidity and pressure data to calculate the
         contrails data.
@@ -317,9 +381,29 @@ def process_from_arrays(
             pressure_levels (np.ndarray): Pressure levels (Pa).
 
         Returns:
-            np.ndarray: The calculated engine mixing ratios on pressure levels (Pa/K).
-            This is a placeholder until the full contrail formation logic is implemented.
+            Categorical (integer) array of contrail formation
+
+            - 0 = no contrails
+            - 1 = non-persistent contrails
+            - 2 = persistent contrails
+
+            Array axes are [contrail factor, pressure level, latitude, longitude], where latitude and longitude are
+            only included if present in the temperature and relative humidity input arrays.
         """
+        arrays = (temperature, relative_humidity, pressure_levels)
+        if arrays[2].ndim != 1:
+            raise ValueError(f"Expected 1D pressure array, got {arrays[2].ndim}D.")
+        if arrays[0].shape != arrays[1].shape:
+            raise ValueError(
+                f"Temperature and relative humidity arrays must have same shape:"
+                f"  {arrays[0].shape}\n  {arrays[1].shape}"
+            )
+        if arrays[0].shape[0] != arrays[2].size or arrays[1].shape[0] != arrays[2].size:
+            raise ValueError(
+                f"Leading axes of arrays must match:"
+                f"  {arrays[0].shape}\n  {arrays[1].shape}\n  {arrays[2].shape}"
+            )
+
         self.temperature = temperature
         self.relative_humidity = relative_humidity
         self.pressure_levels = pressure_levels
@@ -330,7 +414,8 @@ def process_from_arrays(
             self.pressure_levels
         )
         self._calculate_critical_temperatures_and_intercepts()
-        return self.engine_mixing_ratios
+        self._calculate_contrail_persistency()
+        return self._boolean_to_categorical()
 
     def process(self, *cubes: Union[Cube, CubeList]) -> Cube:
         """
@@ -344,32 +429,65 @@ def process(self, *cubes: Union[Cube, CubeList]) -> Cube:
                     Cube of the relative humidity on pressure levels.
 
         Returns:
-            Cube of heights above sea level at which contrails will form.
+            Categorical (integer) cube of contrail formation
+
+            - 0 = no contrails
+            - 1 = non-persistent contrails
+            - 2 = persistent contrails
+
+            Cube dimensions are [contrail factor, pressure level, latitude, longitude], where latitude and longitude are
+            only included if present in the input cubes.
         """
+        # Extract input cubes
         cubes = as_cubelist(*cubes)
-        (temperature_cube, humidity_cube) = CubeList(cubes).extract(
-            ["air_temperature", "relative_humidity"]
-        )
+        names_to_extract = ["air_temperature", "relative_humidity"]
+        if len(cubes) != len(names_to_extract):
+            raise ValueError(
+                f"Expected {len(names_to_extract)} cubes, got {len(cubes)}."
+            )
+        try:
+            (temperature_cube, humidity_cube) = CubeList(cubes).extract(
+                names_to_extract
+            )
+        except Exception as e:
+            raise ValueError(
+                f"Could not extract names '{names_to_extract}' from cubelist."
+            ) from e
+
+        # Check cube dimensions are equal
+        if (
+            get_dim_coord_names(temperature_cube) != get_dim_coord_names(humidity_cube)
+            or temperature_cube.shape != humidity_cube.shape
+        ):
+            raise ValueError(
+                f"Cube dimensional coordinates must match:"
+                f"  {temperature_cube.summary(True, 25)}"
+                f"  {humidity_cube.summary(True, 25)}"
+            )
+
         temperature_cube.convert_units("K")
         humidity_cube.convert_units("kg kg-1")
 
-        # Get the pressure levels from the first cube
+        # Get pressure levels
         pressure_coord = temperature_cube.coord("pressure")
         pressure_coord.convert_units("Pa")
 
+        if "pressure".casefold() != get_dim_coord_names(temperature_cube)[0].casefold():
+            raise ValueError(
+                f"Pressure must be the leading axis (got '{get_dim_coord_names(temperature_cube)}')."
+            )
+
         # Calculate contrail formation using numpy arrays
-        _ = self.process_from_arrays(
+        contrail_formation_data = self.process_from_arrays(
             temperature_cube.data, humidity_cube.data, pressure_coord.points
         )
 
-        # Placeholder return to silence my type checker
-        return_cube = Cube(
-            self.engine_mixing_ratios,
-            long_name="engine_mixing_ratios",
-            units="Pa K-1",
+        # Create output cube using contrail formation data
+        contrail_formation_cube = self._create_contrail_formation_cube(
+            contrail_formation_data, temperature_cube
         )
 
-        return return_cube
+        return contrail_formation_cube
 
 
 class ContrailHeightExtractor(BasePlugin):