cloudnetpy 1.49.9__py3-none-any.whl → 1.87.3__py3-none-any.whl

cloudnetpy/products/product_tools.py

@@ -1,65 +1,87 @@
 """General helper classes and functions for all products."""
-from collections import namedtuple
+
+from dataclasses import dataclass
+from os import PathLike
+from typing import NamedTuple
 
 import netCDF4
 import numpy as np
+import numpy.typing as npt
 from numpy import ma
+from numpy.typing import NDArray
 
 from cloudnetpy import constants, utils
 from cloudnetpy.categorize import atmos_utils
 from cloudnetpy.datasource import DataSource
 
-IceCoefficients = namedtuple("IceCoefficients", "K2liquid0 ZT T Z c")
 
+class IceCoefficients(NamedTuple):
+    """Coefficients for ice effective radius retrieval."""
 
-class CategorizeBits:
-    """Class holding information about category and quality bits.
+    K2liquid0: float
+    ZT: float
+    T: float
+    Z: float
+    c: float
 
-    Args:
-        categorize_file (str): Categorize file name.
 
-    Attributes:
-        category_bits (dict): Dictionary containing boolean fields for `droplet`,
-            `falling`, `cold`, `melting`, `aerosol`, `insect`.
+@dataclass
+class CategoryBits:
+    droplet: NDArray[np.bool_]
+    falling: NDArray[np.bool_]
+    freezing: NDArray[np.bool_]
+    melting: NDArray[np.bool_]
+    aerosol: NDArray[np.bool_]
+    insect: NDArray[np.bool_]
 
-        quality_bits (dict): Dictionary containing boolean fields for `radar`,
-            `lidar`, `clutter`, `molecular`, `attenuated` and `corrected`.
 
-    """
+@dataclass
+class QualityBits:
+    radar: NDArray[np.bool_]
+    lidar: NDArray[np.bool_]
+    clutter: NDArray[np.bool_]
+    molecular: NDArray[np.bool_]
+    attenuated_liquid: NDArray[np.bool_]
+    corrected_liquid: NDArray[np.bool_]
+    attenuated_rain: NDArray[np.bool_]
+    corrected_rain: NDArray[np.bool_]
+    attenuated_melting: NDArray[np.bool_]
+    corrected_melting: NDArray[np.bool_]
 
-    category_keys = (
-        "droplet",
-        "falling",
-        "cold",
-        "melting",
-        "aerosol",
-        "insect",
-    )
-
-    quality_keys = (
-        "radar",
-        "lidar",
-        "clutter",
-        "molecular",
-        "attenuated",
-        "corrected",
-    )
-
-    def __init__(self, categorize_file: str):
+
+class CategorizeBits:
+    def __init__(self, categorize_file: str | PathLike) -> None:
         self._categorize_file = categorize_file
-        self.category_bits = self._read_bits("category")
-        self.quality_bits = self._read_bits("quality")
+        self.category_bits = self._read_category_bits()
+        self.quality_bits = self._read_quality_bits()
 
-    def _read_bits(self, bit_type: str) -> dict:
-        """Converts bitfield into dictionary."""
+    def _read_category_bits(self) -> CategoryBits:
         with netCDF4.Dataset(self._categorize_file) as nc:
-            try:
-                bitfield = nc.variables[f"{bit_type}_bits"][:]
-            except KeyError as err:
-                raise KeyError from err
-            keys = getattr(CategorizeBits, f"{bit_type}_keys")
-            bits = {key: utils.isbit(bitfield, i) for i, key in enumerate(keys)}
-        return bits
+            bits = nc.variables["category_bits"][:]
+            return CategoryBits(
+                droplet=utils.isbit(bits, 0),
+                falling=utils.isbit(bits, 1),
+                freezing=utils.isbit(bits, 2),
+                melting=utils.isbit(bits, 3),
+                aerosol=utils.isbit(bits, 4),
+                insect=utils.isbit(bits, 5),
+            )
+
+    def _read_quality_bits(self) -> QualityBits:
+        with netCDF4.Dataset(self._categorize_file) as nc:
+            bits = nc.variables["quality_bits"][:]
+            return QualityBits(
+                radar=utils.isbit(bits, 0),
+                lidar=utils.isbit(bits, 1),
+                clutter=utils.isbit(bits, 2),
+                molecular=utils.isbit(bits, 3),
+                attenuated_liquid=utils.isbit(bits, 4),
+                corrected_liquid=utils.isbit(bits, 5),
+                attenuated_rain=utils.isbit(bits, 6),
+                corrected_rain=utils.isbit(bits, 7),
+                attenuated_melting=utils.isbit(bits, 8),
+                corrected_melting=utils.isbit(bits, 9),
+            )
 
 
 class ProductClassification(CategorizeBits):
@@ -74,7 +96,7 @@ class ProductClassification(CategorizeBits):
 
     """
 
-    def __init__(self, categorize_file: str):
+    def __init__(self, categorize_file: str | PathLike) -> None:
         super().__init__(categorize_file)
         self.is_rain = get_is_rain(categorize_file)
 
@@ -84,79 +106,99 @@ class IceClassification(ProductClassification):
     Child of ProductClassification().
     """
 
-    def __init__(self, categorize_file: str):
+    def __init__(self, categorize_file: str | PathLike) -> None:
         super().__init__(categorize_file)
+        self._is_attenuated = self._find_attenuated()
+        self._is_corrected = self._find_corrected()
         self.is_ice = self._find_ice()
         self.would_be_ice = self._find_would_be_ice()
         self.corrected_ice = self._find_corrected_ice()
         self.uncorrected_ice = self._find_uncorrected_ice()
         self.ice_above_rain = self._find_ice_above_rain()
-        self.cold_above_rain = self._find_cold_above_rain()
+        self.clear_above_rain = self._find_clear_above_rain()
 
-    def _find_ice(self) -> np.ndarray:
+    def _find_clear_above_rain(self) -> npt.NDArray:
         return (
-            self.category_bits["falling"]
-            & self.category_bits["cold"]
-            & ~self.category_bits["melting"]
-            & ~self.category_bits["insect"]
+            utils.transpose(self.is_rain) * ~self.is_ice
+            & self.category_bits.freezing
+            & ~self.category_bits.melting
         )
 
-    def _find_would_be_ice(self) -> np.ndarray:
-        warm_falling = (
-            self.category_bits["falling"]
-            & ~self.category_bits["cold"]
-            & ~self.category_bits["insect"]
+    def _find_attenuated(self) -> npt.NDArray:
+        return (
+            self.quality_bits.attenuated_liquid
+            | self.quality_bits.attenuated_rain
+            | self.quality_bits.attenuated_melting
         )
-        return warm_falling | self.category_bits["melting"]
 
-    def _find_corrected_ice(self) -> np.ndarray:
+    def _find_corrected(self) -> npt.NDArray:
         return (
-            self.is_ice
-            & self.quality_bits["attenuated"]
-            & self.quality_bits["corrected"]
+            self.quality_bits.corrected_liquid
+            | self.quality_bits.corrected_rain
+            | self.quality_bits.corrected_melting
+        )
+
+    def _find_ice(self) -> npt.NDArray:
+        return (
+            self.category_bits.falling
+            & self.category_bits.freezing
+            & ~self.category_bits.melting
+            & ~self.category_bits.insect
+        )
+
+    def _find_would_be_ice(self) -> npt.NDArray:
+        warm_falling = (
+            self.category_bits.falling
+            & ~self.category_bits.freezing
+            & ~self.category_bits.insect
         )
+        return warm_falling | self.category_bits.melting
 
-    def _find_uncorrected_ice(self) -> np.ndarray:
+    def _find_corrected_ice(self) -> npt.NDArray:
+        return self.is_ice & self._is_attenuated & self._is_corrected
+
+    def _find_uncorrected_ice(self) -> npt.NDArray:
+        uncorrected_melting = (
+            self.quality_bits.attenuated_melting & ~self.quality_bits.corrected_melting
+        )
+        uncorrected_rain = (
+            self.quality_bits.attenuated_rain & ~self.quality_bits.corrected_rain
+        )
+        uncorrected_liquid = (
+            self.quality_bits.attenuated_liquid & ~self.quality_bits.corrected_liquid
+        )
         return (
             self.is_ice
-            & self.quality_bits["attenuated"]
-            & ~self.quality_bits["corrected"]
+            & self._is_attenuated
+            & (uncorrected_melting | uncorrected_rain | uncorrected_liquid)
         )
 
-    def _find_ice_above_rain(self) -> np.ndarray:
+    def _find_ice_above_rain(self) -> npt.NDArray:
         is_rain = utils.transpose(self.is_rain)
         return (self.is_ice * is_rain) == 1
 
-    def _find_cold_above_rain(self) -> np.ndarray:
-        is_cold = self.category_bits["cold"]
-        is_rain = utils.transpose(self.is_rain)
-        is_cold_rain = (is_cold * is_rain) == 1
-        return is_cold_rain & ~self.category_bits["melting"]
-
 
 class IceSource(DataSource):
     """Base class for different ice products."""
 
-    def __init__(self, categorize_file: str, product: str):
+    def __init__(self, categorize_file: str | PathLike, product: str) -> None:
         super().__init__(categorize_file)
-        self.wl_band = utils.get_wl_band(float(self.getvar("radar_frequency")))
-        self.temperature = get_temperature(categorize_file)
+        self.radar_frequency = float(self.getvar("radar_frequency"))
+        self.wl_band = utils.get_wl_band(self.radar_frequency)
+        self.temperature = _get_temperature(categorize_file)
         self.product = product
         self.coefficients = self._get_coefficients()
 
-    def append_main_variable_including_rain(
-        self, ice_classification: IceClassification
+    def append_icy_data(
+        self,
+        ice_classification: IceClassification,
     ) -> None:
         """Adds the main variable (including ice above rain)."""
-        data_including_rain = self._convert_z()
-        data_including_rain[~ice_classification.is_ice] = ma.masked
-        self.append_data(data_including_rain, f"{self.product}_inc_rain")
-
-    def append_main_variable(self, ice_classification: IceClassification) -> None:
-        """Adds the main variable (excluding rain)."""
-        data = ma.copy(self.data[f"{self.product}_inc_rain"][:])
-        data[ice_classification.ice_above_rain] = ma.masked
-        self.append_data(data, self.product)
+        data = self._convert_z()
+        data[~ice_classification.is_ice | ice_classification.uncorrected_ice] = (
+            ma.masked
+        )
+        self.append_data(data, f"{self.product}")
 
     def append_status(self, ice_classification: IceClassification) -> None:
         """Adds the status of retrieval."""
@@ -165,10 +207,9 @@ class IceSource(DataSource):
         is_data = ~data.mask
         retrieval_status[is_data] = 1
         retrieval_status[is_data & ice_classification.corrected_ice] = 3
-        retrieval_status[is_data & ice_classification.uncorrected_ice] = 2
         retrieval_status[~is_data & ice_classification.is_ice] = 4
-        retrieval_status[ice_classification.cold_above_rain] = 6
-        retrieval_status[ice_classification.ice_above_rain] = 5
+        retrieval_status[ice_classification.uncorrected_ice] = 2
+        retrieval_status[ice_classification.clear_above_rain] = 6
         retrieval_status[ice_classification.would_be_ice & (retrieval_status == 0)] = 7
         self.append_data(retrieval_status, f"{self.product}_retrieval_status")
 
@@ -178,18 +219,27 @@ class IceSource(DataSource):
         References:
             Hogan et.al. 2006, https://doi.org/10.1175/JAM2340.1
         """
+        msg = f"Unsupported band: {self.wl_band}"
         if self.product == "ier":
-            if self.wl_band == 0:
+            if self.wl_band == "Ka":
                 return IceCoefficients(0.878, -0.000205, -0.0015, 0.0016, -1.52)
-            return IceCoefficients(0.669, -0.000296, -0.00193, -0.000, -1.502)
-        if self.wl_band == 0:
+            if self.wl_band == "W":
+                return IceCoefficients(0.669, -0.000296, -0.00193, -0.000, -1.502)
+            raise ValueError(msg)
+        if self.wl_band == "Ka":
             return IceCoefficients(0.878, 0.000242, -0.0186, 0.0699, -1.63)
-        return IceCoefficients(0.669, 0.000580, -0.00706, 0.0923, -0.992)
+        if self.wl_band == "W":
+            return IceCoefficients(0.669, 0.000580, -0.00706, 0.0923, -0.992)
+        raise ValueError(msg)
 
-    def _convert_z(self, z_variable: str = "Z") -> np.ndarray:
+    def _convert_z(self, z_variable: str = "Z") -> npt.NDArray:
         """Calculates temperature weighted z, i.e. ice effective radius [m]."""
-        assert self.product in ("iwc", "ier")
-        assert z_variable in ("Z", "Z_sensitivity")
+        if self.product not in ("iwc", "ier"):
+            msg = f"Invalid product: {self.product}"
+            raise ValueError(msg)
+        if z_variable not in ("Z", "Z_sensitivity"):
+            msg = f"Invalid z_variable: {z_variable}"
+            raise ValueError(msg)
         temperature = (
             self.temperature if z_variable == "Z" else ma.mean(self.temperature, axis=0)
         )
@@ -212,39 +262,31 @@ class IceSource(DataSource):
 
     def _get_z_factor(self) -> float:
         """Returns empirical scaling factor for radar echo."""
-        return float(utils.lin2db(self.coefficients.K2liquid0 / 0.93))
-
+        k2 = np.array(self.coefficients.K2liquid0) / 0.93
+        return float(utils.lin2db(k2))
 
-def get_is_rain(filename: str) -> np.ndarray:
-    try:
-        rainfall_rate = read_nc_fields(filename, "rainfall_rate")
-    except KeyError:
-        rainfall_rate = read_nc_fields(filename, "rain_rate")
-    is_rain = rainfall_rate != 0
-    assert isinstance(is_rain, ma.MaskedArray)
-    is_rain[is_rain.mask] = True
-    return np.array(is_rain)
 
+def get_is_rain(filename: str | PathLike) -> npt.NDArray:
+    # TODO: Check that this is correct
+    with netCDF4.Dataset(filename) as nc:
+        for name in ["rain_detected", "rainfall_rate", "rain_rate"]:
+            if name in nc.variables:
+                data = nc.variables[name][:]
+                data = data != 0
+                data[data.mask] = True
+                return np.array(data)
+    msg = "No rain data found."
+    raise ValueError(msg)
 
-def read_nc_fields(nc_file: str, names: str | list) -> ma.MaskedArray | list:
-    """Reads selected variables from a netCDF file.
-
-    Args:
-        nc_file: netCDF file name.
-        names: Variables to be read, e.g. 'temperature' or ['ldr', 'lwp'].
 
-    Returns:
-        ndarray/list: Array in case of one variable passed as a string.
-        List of arrays otherwise.
-
-    """
-    names = [names] if isinstance(names, str) else names
+def read_nc_field(nc_file: str | PathLike, name: str) -> ma.MaskedArray:
     with netCDF4.Dataset(nc_file) as nc:
-        data = [nc.variables[name][:] for name in names]
-    return data[0] if len(data) == 1 else data
+        return nc.variables[name][:]
 
 
-def interpolate_model(cat_file: str, names: str | list) -> dict[str, np.ndarray]:
+def interpolate_model(
+    cat_file: str | PathLike, names: str | list
+) -> dict[str, npt.NDArray]:
     """Interpolates 2D model field into dense Cloudnet grid.
 
     Args:
@@ -257,8 +299,8 @@ def interpolate_model(cat_file: str, names: str | list) -> dict[str, np.ndarray]
 
     """
 
-    def _interp_field(var_name: str) -> np.ndarray:
-        values = read_nc_fields(
+    def _interp_field(var_name: str) -> npt.NDArray:
+        values = _read_nc_fields(
             cat_file,
             ["model_time", "model_height", var_name, "time", "height"],
         )
@@ -268,12 +310,12 @@ def interpolate_model(cat_file: str, names: str | list) -> dict[str, np.ndarray]
     return {name: _interp_field(name) for name in names}
 
 
-def get_temperature(categorize_file: str) -> np.ndarray:
+def _read_nc_fields(nc_file: str | PathLike, names: list[str]) -> list[ma.MaskedArray]:
+    with netCDF4.Dataset(nc_file) as nc:
+        return [nc.variables[name][:] for name in names]
+
+
+def _get_temperature(categorize_file: str | PathLike) -> npt.NDArray:
     """Returns interpolated temperatures in Celsius."""
     atmosphere = interpolate_model(categorize_file, "temperature")
     return atmos_utils.k2c(atmosphere["temperature"])
-
-
-def get_mwrpy_coeffs(nc_file: str) -> str:
-    with netCDF4.Dataset(nc_file) as nc:
-        return nc.mwrpy_coefficients