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

cloudnetpy/model_evaluation/plotting/plot_tools.py

@@ -1,7 +1,7 @@
 import netCDF4
 import numpy as np
-from matplotlib import cm
-from matplotlib.colors import ListedColormap
+import numpy.typing as npt
+from matplotlib.axes import Axes
 from numpy import ma
 
 from cloudnetpy.model_evaluation.model_metadata import MODELS
@@ -12,17 +12,15 @@ def parse_wanted_names(
     name: str,
     model: str,
     variables: list | None = None,
+    *,
     advance: bool = False,
 ) -> tuple[list, list]:
-    """Returns standard and advection lists of product types to plot"""
-    if variables:
-        names = variables
-    else:
-        names = parse_dataset_keys(nc_file, name, advance, model)
+    """Returns standard and advection lists of product types to plot."""
+    names = variables or parse_dataset_keys(nc_file, name, advance=advance, model=model)
     standard_n = [n for n in names if name in n and "adv" not in n]
     standard_n = sort_model2first_element(standard_n, model)
     advection_n = [n for n in names if name in n and "adv" in n]
-    model_names = [n for n in names if f"{model}_" in n and not f"_{model}_" in n]
+    model_names = [n for n in names if f"{model}_" in n and f"_{model}_" not in n]
     for i, model_n in enumerate(model_names):
         advection_n.insert(0 + i, model_n)
     if len(advection_n) < len(standard_n):
@@ -33,18 +31,18 @@ def parse_wanted_names(
 
 
 def parse_dataset_keys(
-    nc_file: str, product: str, advance: bool, model: str = ""
+    nc_file: str,
+    product: str,
+    *,
+    advance: bool,
+    model: str,
 ) -> list:
     names = list(netCDF4.Dataset(nc_file).variables.keys())
     a_names = ["cirrus", "snow"]
     model_vars = []
     for n in names:
-        if model not in n:
-            model_vars.append(n)
-        elif model in n and product not in n:
+        if model not in n or (model in n and product not in n):
             model_vars.append(n)
-        else:
-            continue
     if not advance:
         for a in a_names:
             for n in names:
@@ -77,7 +75,7 @@ def sort_cycles(names: list, model: str) -> tuple[list, list]:
 
 
 def read_data_characters(nc_file: str, name: str, model: str) -> tuple:
-    """Gets dimensions and data for plotting"""
+    """Gets dimensions and data for plotting."""
     nc = netCDF4.Dataset(nc_file)
     data = nc.variables[name][:]
     data = mask_small_values(data, name)
@@ -85,10 +83,12 @@ def read_data_characters(nc_file: str, name: str, model: str) -> tuple:
     x = reshape_1d2nd(x, data)
     try:
         y = nc.variables[f"{model}_height"][:]
-    except KeyError:
+    except KeyError as err:
         model_info = MODELS[model]
         cycles = model_info.cycle
-        assert cycles is not None
+        if cycles is None:
+            msg = f"Invalid model: {model}"
+            raise RuntimeError(msg) from err
         cycles_split = [x.strip() for x in cycles.split(",")]
         cycle = [cycle for cycle in cycles_split if cycle in name]
         y = nc.variables[f"{model}_{cycle[0]}_height"][:]
@@ -96,13 +96,13 @@ def read_data_characters(nc_file: str, name: str, model: str) -> tuple:
     try:
         mask = y.mask
         if mask.any():
-            x, y, data = change2one_dim_axes(x, y, data)
+            x, y, data = change2one_dim_axes(ma.array(x), y, data)
     except AttributeError:
         return data, x, y
     return data, x, y
 
 
-def mask_small_values(data: ma.MaskedArray, name: str):
+def mask_small_values(data: ma.MaskedArray, name: str) -> ma.MaskedArray:
     data[data <= 0] = ma.masked
     if "lwc" in name:
         data[data < 1e-5] = ma.masked
@@ -111,35 +111,28 @@ def mask_small_values(data: ma.MaskedArray, name: str):
     return data
 
 
-def reshape_1d2nd(one_d: np.ndarray, two_d: np.ndarray) -> np.ndarray:
+def reshape_1d2nd(one_d: npt.NDArray, two_d: npt.NDArray) -> npt.NDArray:
     new_arr = np.zeros(two_d.shape)
     for i in range(len(two_d[0])):
         new_arr[:, i] = one_d
     return new_arr
 
 
-def create_segment_values(arrays: list) -> tuple:
-    # 0=no data, 1=model, 2=intersection, 3=observation
-    new_array = np.zeros(arrays[0].shape, dtype=int)
-    for i, array in enumerate(arrays):
-        new_array[~array] = new_array[~array] + i + 1
-    new_array[new_array == 2] = 4
-    new_array[new_array == 3] = 2
-    new_array[new_array == 4] = 3
-
-    colors = cm.get_cmap("YlGnBu", 256)
-    newcolors = colors(np.linspace(0, 1, 256))
-    # No data, model, both, observation
-    cmap = ListedColormap(["white", "khaki", newcolors[90], newcolors[140]])
-    return new_array, cmap
+def create_segment_values(model: ma.MaskedArray, obs: ma.MaskedArray) -> npt.NDArray:
+    new_array = np.zeros(model.shape, dtype=int)
+    new_array[model.mask & obs.mask] = 0  # No data
+    new_array[~model.mask & obs.mask] = 1  # Only model
+    new_array[~obs.mask & model.mask] = 3  # Only observation
+    new_array[(~model.mask == 1) & (~obs.mask == 1)] = 2  # Both
+    return new_array
 
 
-def set_yaxis(ax, max_y: float, min_y: float = 0.0):
+def set_yaxis(ax: Axes, max_y: float, min_y: float = 0.0) -> None:
     ax.set_ylim(min_y, max_y)
     ax.set_ylabel("Height (km)", fontsize=13)
 
 
-def rolling_mean(data: ma.MaskedArray, n: int = 4) -> np.ndarray:
+def rolling_mean(data: ma.MaskedArray, n: int = 4) -> npt.NDArray:
     mmr = []
     for i in range(len(data)):
         if not data[i : i + n].mask.all():
@@ -150,7 +143,9 @@ def rolling_mean(data: ma.MaskedArray, n: int = 4) -> np.ndarray:
 
 
 def change2one_dim_axes(
-    x: ma.MaskedArray, y: ma.MaskedArray, data: np.ndarray
+    x: ma.MaskedArray,
+    y: ma.MaskedArray,
+    data: npt.NDArray,
 ) -> tuple:
     # If any mask in x or y, change 2d to 1d axes values
     # Common shape need to match 2d data.