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

cloudnetpy/model_evaluation/plotting/plotting.py

@@ -1,18 +1,27 @@
 import logging
 import os
 import sys
+from datetime import date
 
+import matplotlib as mpl
 import matplotlib.pyplot as plt
+import netCDF4
 import numpy as np
+import numpy.typing as npt
+from matplotlib.axes import Axes
+from matplotlib.colorbar import Colorbar
+from matplotlib.colorizer import ColorizingArtist
+from matplotlib.colors import ListedColormap
 from matplotlib.patches import Patch
+from matplotlib.pyplot import Figure
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from numpy import ma
 
 import cloudnetpy.model_evaluation.plotting.plot_tools as p_tools
-import cloudnetpy.plotting.plotting as cloud_plt
 from cloudnetpy.model_evaluation.model_metadata import MODELS
-from cloudnetpy.model_evaluation.plotting.plot_meta import ATTRIBUTES
+from cloudnetpy.model_evaluation.plotting.plot_meta import ATTRIBUTES, PlotMeta
 from cloudnetpy.model_evaluation.statistics.statistical_methods import DayStatistics
+from cloudnetpy.plotting.plotting import Dimensions, get_log_cbar_tick_labels, lin2log
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
@@ -21,6 +30,7 @@ def generate_L3_day_plots(
     nc_file: str,
     product: str,
     model: str,
+    *,
     title: bool = True,
     var_list: list | None = None,
     fig_type: str | None = "group",
@@ -28,7 +38,7 @@ def generate_L3_day_plots(
     save_path: str | None = None,
     image_name: str | None = None,
     show: bool | None = False,
-):
+) -> list[Dimensions]:
     """Generate visualizations for level 3 dayscale products.
     With figure type visualizations can be subplot in group, pair, single or
     statistic of given product. In group fig_type all different methods are plot
@@ -38,6 +48,7 @@ def generate_L3_day_plots(
     Single fig_type will plot each product variable in a own figure.
     Statistic fig_type will plot select statistical method of all product method
     in same fig.
+
     Args:
         nc_file (str): Path to source file
         product (str): Name of product wanted to plot
@@ -63,6 +74,7 @@ def generate_L3_day_plots(
         In case of model cycles, cycles are visualized in their on figures same
         way as an individual model run would be visualized in its own in a group
         figure.
+
     Examples:
         >>> from cloudnetpy.model_evaluation.plotting.plotting
         import generate_L3_day_plots
@@ -76,17 +88,25 @@ def generate_L3_day_plots(
         >>> generate_L3_day_plots(l3_day_file, product, model,
         >>>                       fig_type='statistic', stats=['error'])
     """
+
+    def _check_cycle_names() -> None:
+        if not c_names:
+            raise AttributeError
+
     cls = __import__("plotting")
     model_info = MODELS[model]
     model_name = model_info.model_name
     name_set = p_tools.parse_wanted_names(nc_file, product, model, var_list)
-    for names in name_set:
+    unique_tuples = {tuple(lst) for lst in name_set}
+    name_set_unique = tuple(list(tup) for tup in unique_tuples)
+
+    dimensions = []
+    for names in name_set_unique:
         if len(names) > 0:
             try:
                 cycle_names, cycles = p_tools.sort_cycles(names, model)
                 for i, c_names in enumerate(cycle_names):
-                    if not c_names:
-                        raise AttributeError
+                    _check_cycle_names()
                     params = [
                         product,
                         c_names,
@@ -95,10 +115,8 @@ def generate_L3_day_plots(
                         model_name,
                         save_path,
                         image_name,
-                        show,
-                        cycles[i],
-                        title,
                     ]
+                    kwargs = {"show": show, "title": title, "cycle": cycles[i]}
                     if fig_type == "statistic":
                         params = [
                             product,
@@ -109,11 +127,12 @@ def generate_L3_day_plots(
                             stats,
                             save_path,
                             image_name,
-                            show,
-                            cycles[i],
-                            title,
                         ]
-                    getattr(cls, f"get_{fig_type}_plots")(*params)
+                    figs, axes = getattr(cls, f"get_{fig_type}_plots")(
+                        *params, **kwargs
+                    )
+                    for fig, ax in zip(figs, axes, strict=False):
+                        dimensions.append(Dimensions(fig, ax))
             except AttributeError:
                 params = [
                     product,
@@ -123,10 +142,8 @@ def generate_L3_day_plots(
                     model_name,
                     save_path,
                     image_name,
-                    show,
-                    "",
-                    title,
                 ]
+                kwargs = {"show": show, "title": title}
                 if fig_type == "statistic":
                     params = [
                         product,
@@ -137,11 +154,11 @@ def generate_L3_day_plots(
                         stats,
                         save_path,
                         image_name,
-                        show,
-                        "",
-                        title,
                     ]
-                getattr(cls, f"get_{fig_type}_plots")(*params)
+                figs, axes = getattr(cls, f"get_{fig_type}_plots")(*params, **kwargs)
+                for fig, ax in zip(figs, axes, strict=False):
+                    dimensions.append(Dimensions(fig, ax))
+    return dimensions
 
 
 def get_group_plots(
@@ -152,14 +169,17 @@ def get_group_plots(
     model_name: str,
     save_path: str,
     image_name: str,
+    *,
     show: bool,
     cycle: str = "",
     title: bool = True,
-):
+    include_xlimits: bool = False,
+) -> tuple:
     """Group subplot visualization for both standard and advection downsampling.
     Generates group subplot figure for product with model and all different
     downsampling methods. Generates separated figures for standard and advection
     timegrids. All model cycles if any will be generated to their own figures.
+
     Args:
         product (str): Name of the product
         names (list): List of variables to be visualized to same fig
@@ -171,6 +191,7 @@ def get_group_plots(
         show (bool): Show figure before saving if True
         cycle (str): Name of cycle if exists
         title (bool): True or False if wanted to add title to subfig
+        include_xlimits (bool): Show labels at the ends of x-axis
     """
     fig, ax = initialize_figure(len(names))
     model_run = model
@@ -178,22 +199,28 @@ def get_group_plots(
     j = 0
     for j, name in enumerate(names):
         variable_info = ATTRIBUTES[product]
-        cloud_plt.set_ax(ax[j], 12, ylabel=None)
+        set_yax(ax[j], 12, ylabel=None)
+        set_xax(ax[j], include_xlimits=include_xlimits)
         if title:
             _set_title(ax[j], name, product, variable_info)
         if j == 0 and title:
             _set_title(ax[j], model, product, variable_info, model_name)
         data, x, y = p_tools.read_data_characters(nc_file, name, model)
         plot_colormesh(ax[j], data, (x, y), variable_info)
-    casedate = cloud_plt.set_labels(fig, ax[j], nc_file)
+    casedate = set_labels(fig, ax[j], nc_file)
     if "adv" in name:
         product = product + "_adv"
     if len(cycle) > 1:
         fig.text(0.64, 0.885, f"Cycle: {cycle}", fontsize=13)
         model_run = f"{model}_{cycle}"
-    cloud_plt.handle_saving(
-        image_name, save_path, show, casedate, [product, model_run, "group"]
+    handle_saving(
+        image_name,
+        save_path,
+        casedate,
+        [product, model_run, "group"],
+        show=show,
     )
+    return fig, ax
 
 
 def get_pair_plots(
@@ -204,14 +231,17 @@ def get_pair_plots(
     model_name: str,
     save_path: str,
     image_name: str,
-    show: bool,
     cycle: str = "",
+    *,
+    show: bool = False,
     title: bool = True,
-):
+    include_xlimits: bool = False,
+) -> None:
     """Pair subplots of model and product method.
     In upper subplot is model product and lower subplot one of the
     downsampled method of select product. Function generates all product methods
     in a given nc-file in loop.
+
     Args:
         product (str): Name of the product
         names (list): List of variables to be visualized to same fig
@@ -223,6 +253,7 @@ def get_pair_plots(
         show (bool): Show figure before saving if True
         cycle (str): Name of cycle if exists
         title (bool): True or False if wanted add title to subfig
+        include_xlimits (bool): Show labels at the ends of x-axis
     """
     variable_info = ATTRIBUTES[product]
     model_ax = names[0]
@@ -230,8 +261,11 @@ def get_pair_plots(
         if i == 0:
             continue
         fig, ax = initialize_figure(2)
-        cloud_plt.set_ax(ax[0], 12, ylabel=None)
-        cloud_plt.set_ax(ax[-1], 12, ylabel=None)
+        set_yax(ax[0], 12, ylabel=None)
+        set_yax(ax[-1], 12, ylabel=None)
+        set_xax(ax[0], include_xlimits=include_xlimits)
+        set_xax(ax[-1], include_xlimits=include_xlimits)
+
         if title:
             _set_title(ax[0], model, product, variable_info, model_name)
             _set_title(ax[-1], name, product, variable_info)
@@ -239,10 +273,16 @@ def get_pair_plots(
         data, x, y = p_tools.read_data_characters(nc_file, name, model)
         plot_colormesh(ax[0], model_data, (mx, my), variable_info)
         plot_colormesh(ax[-1], data, (x, y), variable_info)
-        casedate = cloud_plt.set_labels(fig, ax[-1], nc_file)
+        casedate = set_labels(fig, ax[-1], nc_file)
         if len(cycle) > 1:
             fig.text(0.64, 0.889, f"Cycle: {cycle}", fontsize=13)
-        cloud_plt.handle_saving(image_name, save_path, show, casedate, [name, "pair"])
+        handle_saving(
+            image_name,
+            save_path,
+            casedate,
+            [name, "pair"],
+            show=show,
+        )
 
 
 def get_single_plots(
@@ -253,11 +293,14 @@ def get_single_plots(
     model_name: str,
     save_path: str,
     image_name: str,
+    *,
     show: bool,
     cycle: str = "",
     title: bool = True,
-):
+    include_xlimits: bool = False,
+) -> tuple[list, list]:
     """Generates figures of each product variable from given file in loop.
+
     Args:
         product (str): Name of the product
         names (list): List of variables to be visualized to same fig
@@ -269,38 +312,55 @@ def get_single_plots(
         show (bool): Show figure before saving if True
         cycle (str): Name of cycle if exists
         title (bool): True or False if wanted to add title to subfig
+        include_xlimits (bool): Show labels at the ends of x-axis
     """
+    figs = []
+    axes = []
     variable_info = ATTRIBUTES[product]
-    for _, name in enumerate(names):
+    for name in names:
         fig, ax = initialize_figure(1)
-        cloud_plt.set_ax(ax[0], 12, ylabel=None)
+        figs.append(fig)
+        axes.append(ax)
+        set_yax(ax[0], 12, ylabel=None)
+        set_xax(ax[0], include_xlimits=include_xlimits)
+
         if title:
             _set_title(ax[0], name, product, variable_info)
         data, x, y = p_tools.read_data_characters(nc_file, name, model)
         plot_colormesh(ax[0], data, (x, y), variable_info)
-        casedate = cloud_plt.set_labels(fig, ax[0], nc_file, sub_title=title)
+        casedate = set_labels(fig, ax[0], nc_file, sub_title=title)
         if title:
             if len(cycle) > 1:
                 fig.text(0.64, 0.9, f"{model_name} cycle: {cycle}", fontsize=13)
             else:
                 fig.text(0.64, 0.9, f"{model_name}", fontsize=13)
-        cloud_plt.handle_saving(image_name, save_path, show, casedate, [name, "single"])
+        handle_saving(
+            image_name,
+            save_path,
+            casedate,
+            [name, "single"],
+            show=show,
+        )
+    return figs, axes
 
 
-def plot_colormesh(ax, data: np.ndarray, axes: tuple, variable_info):
+def plot_colormesh(
+    ax: Axes, data: npt.NDArray, axes: tuple, variable_info: PlotMeta
+) -> None:
     vmin, vmax = variable_info.plot_range
     if variable_info.plot_scale == "logarithmic":
-        data, vmin, vmax = cloud_plt.lin2log(data, vmin, vmax)
+        data, vmin, vmax = lin2log(data, vmin, vmax)
     cmap = plt.get_cmap(variable_info.cbar, 22)
     data[data < vmin] = ma.masked
     pl = ax.pcolormesh(*axes, data, vmin=vmin, vmax=vmax, cmap=cmap)
     colorbar = init_colorbar(pl, ax)
     if variable_info.plot_scale == "logarithmic":
-        tick_labels = cloud_plt.generate_log_cbar_ticklabel_list(vmin, vmax)
-        colorbar.set_ticks(np.arange(vmin, vmax + 1))
+        tick_labels = get_log_cbar_tick_labels(vmin, vmax)
+        colorbar.set_ticks(np.arange(vmin, vmax + 1).tolist())  # type: ignore[arg-type]
         colorbar.ax.set_yticklabels(tick_labels)
     ax.set_facecolor("white")
-    colorbar.set_label(variable_info.clabel, fontsize=13)
+    if variable_info.clabel is not None:
+        colorbar.set_label(variable_info.clabel, fontsize=13)
 
 
 def get_statistic_plots(
@@ -312,10 +372,11 @@ def get_statistic_plots(
     stats: list,
     save_path: str,
     image_name: str,
+    *,
     show: bool,
     cycle: str = "",
     title: bool = True,
-):
+) -> tuple:
     """Statistical subplots for day scale products.
     Statistical analysis can be done by day scale with relative error ('error'),
     total data area analysis ('area'), histogram ('hist') or vertical profiles
@@ -323,6 +384,7 @@ def get_statistic_plots(
     per statistical method for a select product. All different downsampled method
     are in a same fig. Standard and advection timegrids are separated to own figs
     as well as different cycle runs.
+
     Args:
         product (str): Name of the product
         names (list): List of variables to be visualized to same fig
@@ -335,18 +397,34 @@ def get_statistic_plots(
         image_name (str, optional): Saving name of generated fig
         show (bool): Show figure before saving if True
         cycle (str): Name of cycle if exists
+        title (bool): True or False if wanted to add title to subfig
     """
-    # pylint: disable=too-many-branches
-    # pylint: disable=too-many-nested-blocks
     model_run = model
     name = ""
     j = 0
+
+    def _check_data() -> None:
+        if model_missing and obs_missing:
+            _raise()
+
+    def _check_data2() -> None:
+        if "error" in stat and np.all(day_stat.model_stat.mask is True):
+            _raise()
+
+    def _raise() -> None:
+        err_msg = f"No data in {model_name} or observation"
+        raise ValueError(err_msg)
+
+    figs = []
+    axes = []
     for stat in stats:
         try:
             obs_missing = False
             model_missing = False
             variable_info = ATTRIBUTES[product]
             fig, ax = initialize_figure(len(names) - 1, stat)
+            figs.append(fig)
+            axes.append(ax)
             model_data, _, _ = p_tools.read_data_characters(nc_file, names[0], model)
             if np.all(model_data.mask is True):
                 model_missing = True
@@ -354,65 +432,73 @@ def get_statistic_plots(
                 data, x, y = p_tools.read_data_characters(nc_file, name, model)
                 if np.all(data.mask is True):
                     obs_missing = True
-                if model_missing and obs_missing:
-                    raise ValueError("No data in either dataset")
-                if product == "cf" and stat == "error":
-                    stat = "aerror"
+                _check_data()
+                statistics = "aerror" if product == "cf" and stat == "error" else stat
                 if j > 0:
-                    name = ""
-                    name = _get_stat_titles(name, product, variable_info)
+                    name_new = ""
+                    name_new = _get_stat_titles(name_new, product, variable_info)
                     day_stat = DayStatistics(
-                        stat, [product, model_name, name], model_data, data
+                        statistics,
+                        [product, model_name, name_new],
+                        model_data,
+                        data,
                     )
-                    if "error" in stat:
-                        if np.all(day_stat.model_stat.mask is True):
-                            raise ValueError("No data to calculate relative error")
+                    _check_data2()
                     initialize_statistic_plots(
                         j,
                         len(names) - 1,
                         ax[j - 1],
-                        stat,
+                        statistics,
                         day_stat,
                         model_data,
                         data,
                         (x, y),
                         variable_info,
-                        title,
+                        title=title,
                     )
-        except ValueError as e:
-            logging.error(e)
+        except ValueError:
+            logging.exception("Exception occurred")
         if stat not in ("hist", "vertical"):
-            casedate = cloud_plt.set_labels(fig, ax[j - 1], nc_file, sub_title=title)
+            casedate = set_labels(fig, ax[j - 1], nc_file, sub_title=title)
         else:
-            casedate = cloud_plt.read_date(nc_file)
-            _name = cloud_plt.read_location(nc_file)
+            casedate = read_date(nc_file)
+            _name = read_location(nc_file)
             if title:
-                cloud_plt.add_subtitle(fig, casedate, _name.capitalize())
+                add_subtitle(fig, casedate, _name.capitalize())
         if len(cycle) > 1:
             fig.text(0.64, 0.885, f"Cycle: {cycle}", fontsize=13)
             model_run = f"{model}_{cycle}"
-        cloud_plt.handle_saving(
-            image_name, save_path, show, casedate, [name, stat, model_run]
+        handle_saving(
+            image_name,
+            save_path,
+            casedate,
+            [name, stat, model_run],
+            show=show,
         )
+    return figs, axes
 
 
 def initialize_statistic_plots(
     j: int,
     max_len: int,
-    ax,
+    ax: Axes,
     method: str,
     day_stat: DayStatistics,
     model: ma.MaskedArray,
     obs: ma.MaskedArray,
     args: tuple,
-    variable_info,
+    variable_info: PlotMeta,
+    *,
     title: bool = True,
-):
+    include_xlimits: bool = False,
+) -> None:
     if method in ("error", "aerror"):
         plot_relative_error(ax, day_stat.model_stat.T, args, method)
         if title:
             ax.set_title(day_stat.title, fontsize=14)
-        cloud_plt.set_ax(ax, 12, ylabel=None)
+        set_yax(ax, 12, ylabel=None)
+        set_xax(ax, include_xlimits=include_xlimits)
+
     if method == "area":
         plot_data_area(ax, day_stat, model, obs, args, title=title)
         ax.text(
@@ -423,7 +509,9 @@ def initialize_statistic_plots(
             ha="center",
             transform=ax.transAxes,
         )
-        cloud_plt.set_ax(ax, 12, ylabel=None)
+        set_yax(ax, 12, ylabel=None)
+        set_xax(ax, include_xlimits=include_xlimits)
+
     if method == "hist":
         plot_histogram(ax, day_stat, variable_info)
         if j == max_len - 1 and (max_len % 2) == 0:
@@ -459,19 +547,19 @@ def initialize_statistic_plots(
             )
 
 
-def plot_relative_error(ax, error: ma.MaskedArray, axes: tuple, method: str):
+def plot_relative_error(
+    ax: Axes, error: ma.MaskedArray, axes: tuple, method: str
+) -> None:
     pl = ax.pcolormesh(*axes, error[:-1, :-1].T, cmap="RdBu_r", vmin=-50, vmax=50)
     colorbar = init_colorbar(pl, ax)
     colorbar.set_label("%", fontsize=13)
     error[np.isnan(error)] = ma.masked
-    error = ma.round(error, decimals=4)
     median_error = ma.median(error.compressed())
-    median_error = np.round(median_error, 3)
     if method == "aerror":
         ax.text(
             0.9,
             -0.17,
-            f"Median absolute error: {median_error} %",
+            f"Median absolute error: {median_error:.2f} %",
             size=12,
             ha="center",
             transform=ax.transAxes,
@@ -480,7 +568,7 @@ def plot_relative_error(ax, error: ma.MaskedArray, axes: tuple, method: str):
         ax.text(
             0.9,
             -0.17,
-            f"Median relative error: {median_error} %",
+            f"Median relative error: {median_error:.2f} %",
             size=12,
             ha="center",
             transform=ax.transAxes,
@@ -488,24 +576,39 @@ def plot_relative_error(ax, error: ma.MaskedArray, axes: tuple, method: str):
 
 
 def plot_data_area(
-    ax,
+    ax: Axes,
     day_stat: DayStatistics,
     model: ma.MaskedArray,
     obs: ma.MaskedArray,
     axes: tuple,
+    *,
     title: bool = True,
-):
-    data, cmap = p_tools.create_segment_values([model.mask, obs.mask])
-    pl = ax.pcolormesh(*axes, data, cmap=cmap)
+) -> None:
+    data = p_tools.create_segment_values(model, obs)
+
+    colors = mpl.colormaps["YlGnBu"]
+    newcolors = colors(np.linspace(0, 1, 256))
+    c_map = {
+        0: "white",
+        1: "khaki",
+        2: newcolors[90],
+        3: newcolors[140],
+    }
+    unique_values = sorted(np.unique(data))
+    c_list = [c_map[value] for value in unique_values if value in c_map]
+    cmap = ListedColormap(c_list)
+
+    ax.pcolormesh(*axes, data, cmap=cmap)
+
     if title:
-        colorbar = init_colorbar(pl, ax)
-        colorbar.set_ticks(np.arange(1, 1, 3))
         ax.set_title(f"{day_stat.title}", fontsize=14)
+
     ax.set_facecolor("black")
+
     legend_elements = [
         Patch(facecolor="khaki", edgecolor="k", label="Model"),
-        Patch(facecolor=cmap(0.5), edgecolor="k", label="Common"),
-        Patch(facecolor=cmap(1.0), edgecolor="k", label="Observation"),
+        Patch(facecolor=newcolors[90], edgecolor="k", label="Common"),
+        Patch(facecolor=newcolors[140], edgecolor="k", label="Observation"),
     ]
     ax.legend(
         handles=legend_elements,
@@ -516,13 +619,13 @@ def plot_data_area(
     )
 
 
-def plot_histogram(ax, day_stat: DayStatistics, variable_info):
+def plot_histogram(ax: Axes, day_stat: DayStatistics, variable_info: PlotMeta) -> None:
     weights = np.ones_like(day_stat.model_stat) / float(len(day_stat.model_stat))
     hist_bins = np.histogram(day_stat.observation_stat, density=True)[-1]
     ax.hist(
         day_stat.model_stat,
         weights=weights,
-        bins=hist_bins,
+        bins=list(hist_bins),
         alpha=0.7,
         facecolor="khaki",
         edgecolor="k",
@@ -530,32 +633,35 @@ def plot_histogram(ax, day_stat: DayStatistics, variable_info):
     )
 
     weights = np.ones_like(day_stat.observation_stat) / float(
-        len(day_stat.observation_stat)
+        len(day_stat.observation_stat),
     )
     ax.hist(
         day_stat.observation_stat,
         weights=weights,
-        bins=hist_bins,
+        bins=list(hist_bins),
         alpha=0.7,
         facecolor="steelblue",
         edgecolor="k",
         label="Observation",
     )
-    ax.set_xlabel(variable_info.x_title, fontsize=13)
+    if variable_info.x_title is not None:
+        ax.set_xlabel(variable_info.x_title, fontsize=13)
     if variable_info.plot_scale == "logarithmic":
         ax.ticklabel_format(axis="x", style="sci", scilimits=(0, 0))
     ax.set_ylabel("Relative frequency %", fontsize=13)
-    ax.yaxis.grid(True, "major")
+    ax.yaxis.grid(which="major")
     ax.set_title(f"{day_stat.title[-1]}", fontsize=14)
 
 
-def plot_vertical_profile(ax, day_stat: DayStatistics, axes: tuple, variable_info):
+def plot_vertical_profile(
+    ax: Axes,
+    day_stat: DayStatistics,
+    axes: tuple,
+    variable_info: PlotMeta,
+) -> None:
     mrm = p_tools.rolling_mean(day_stat.model_stat)
     orm = p_tools.rolling_mean(day_stat.observation_stat)
-    if len(axes[-1].shape) > 1:
-        axes = axes[-1][0]
-    else:
-        axes = axes[-1]
+    axes = axes[-1][0] if len(axes[-1].shape) > 1 else axes[-1]
     ax.plot(day_stat.model_stat, axes, "o", markersize=5.5, color="k")
     ax.plot(day_stat.observation_stat, axes, "o", markersize=5.5, color="k")
     ax.plot(
@@ -588,15 +694,16 @@ def plot_vertical_profile(ax, day_stat: DayStatistics, axes: tuple, variable_inf
     ax.plot(orm, axes, "-", color="green", lw=2, label="Mean of observation")
 
     ax.set_title(f"{day_stat.title[-1]}", fontsize=14)
-    ax.set_xlabel(variable_info.x_title, fontsize=13)
+    if variable_info.x_title is not None:
+        ax.set_xlabel(variable_info.x_title, fontsize=13)
     if variable_info.plot_scale == "logarithmic":
         ax.ticklabel_format(axis="x", style="sci", scilimits=(0, 0))
-    ax.yaxis.grid(True, "major")
-    ax.xaxis.grid(True, "major")
+    ax.yaxis.grid(which="major")
+    ax.xaxis.grid(which="major")
 
 
-def initialize_figure(n_subplots: int, stat: str = "") -> tuple:
-    """Set up fig and ax object, if subplot"""
+def initialize_figure(n_subplots: int, stat: str = "") -> tuple[Figure, list[Axes]]:
+    """Set up fig and ax object, if subplot."""
     if n_subplots <= 0:
         n_subplots = 1
     fig, axes = plt.subplots(n_subplots, 1, figsize=(16, 4 + (n_subplots - 1) * 4.8))
@@ -622,7 +729,6 @@ def initialize_figure(n_subplots: int, stat: str = "") -> tuple:
                 right=0.75,
                 hspace=0.2,
             )
-            axes = axes.flatten()
     if stat == "vertical" and n_subplots > 1:
         fig, axes = plt.subplots(
             int(n_subplots / 2),
@@ -639,20 +745,24 @@ def initialize_figure(n_subplots: int, stat: str = "") -> tuple:
             right=0.75,
             hspace=0.16,
         )
-        axes = axes.flatten()
-    if n_subplots == 1:
-        axes = [axes]
-    return fig, axes
+    axes_list = [axes] if isinstance(axes, Axes) else axes.flatten().tolist()
+    return fig, axes_list
 
 
-def init_colorbar(plot, axis):
+def init_colorbar(plot: ColorizingArtist, axis: Axes) -> Colorbar:
     divider = make_axes_locatable(axis)
     cax = divider.append_axes("right", size="1%", pad=0.25)
     return plt.colorbar(plot, fraction=1.0, ax=axis, cax=cax)
 
 
-def _set_title(ax, field_name: str, product: str, variable_info, model_name: str = ""):
-    """Generates subtitles for different product types"""
+def _set_title(
+    ax: Axes,
+    field_name: str,
+    product: str,
+    variable_info: PlotMeta,
+    model_name: str = "",
+) -> None:
+    """Generates subtitles for different product types."""
     parts = field_name.split("_")
     if parts[0] == product:
         title = _get_product_title(variable_info)
@@ -672,14 +782,14 @@ def _set_title(ax, field_name: str, product: str, variable_info, model_name: str
             ax.set_title(f"Simulated {name}")
 
 
-def _get_cf_title(field_name: str, variable_info) -> str:
+def _get_cf_title(field_name: str, variable_info: PlotMeta) -> str:
     title = f"{variable_info.name}, Area"
     if "V" in field_name:
         title = f"{variable_info.name}, Volume"
     return title
 
 
-def _get_iwc_title(field_name: str, variable_info) -> str:
+def _get_iwc_title(field_name: str, variable_info: PlotMeta) -> str:
     name = variable_info.name
     if "att" in field_name:
         title = f"{name} with good attenuation"
@@ -690,12 +800,11 @@ def _get_iwc_title(field_name: str, variable_info) -> str:
     return title
 
 
-def _get_product_title(variable_info) -> str:
-    title = f"{variable_info.name}"
-    return title
+def _get_product_title(variable_info: PlotMeta) -> str:
+    return f"{variable_info.name}"
 
 
-def _get_stat_titles(field_name: str, product: str, variable_info) -> str:
+def _get_stat_titles(field_name: str, product: str, variable_info: PlotMeta) -> str:
     title = _get_product_title_stat(variable_info)
     if product == "cf":
         title = _get_cf_title_stat(field_name, variable_info)
@@ -707,7 +816,7 @@ def _get_stat_titles(field_name: str, product: str, variable_info) -> str:
     return title
 
 
-def _get_cf_title_stat(field_name: str, variable_info) -> str:
+def _get_cf_title_stat(field_name: str, variable_info: PlotMeta) -> str:
     name = variable_info.name
     title = f"{name} area"
     if "V" in field_name:
@@ -715,7 +824,7 @@ def _get_cf_title_stat(field_name: str, variable_info) -> str:
     return title
 
 
-def _get_iwc_title_stat(field_name: str, variable_info) -> str:
+def _get_iwc_title_stat(field_name: str, variable_info: PlotMeta) -> str:
     name = variable_info.name
     if "att" in field_name:
         title = f"{name} with good attenuation"
@@ -726,7 +835,108 @@ def _get_iwc_title_stat(field_name: str, variable_info) -> str:
     return title
 
 
-def _get_product_title_stat(variable_info) -> str:
+def _get_product_title_stat(variable_info: PlotMeta) -> str:
     name = variable_info.name
-    title = f"{name}"
-    return title
+    return f"{name}"
+
+
+def set_yax(ax: Axes, max_y: float, ylabel: str | None, min_y: float = 0.0) -> None:
+    """Sets yticks, ylim and ylabel for yaxis of axis."""
+    ax.set_ylim(min_y, max_y)
+    ax.set_ylabel("Height (km)", fontsize=13)
+    if ylabel is not None:
+        ax.set_ylabel(ylabel, fontsize=13)
+
+
+def set_xax(ax: Axes, *, include_xlimits: bool = False) -> None:
+    """Sets xticks and xtick labels for plt.imshow()."""
+    ticks_x_labels = _get_standard_time_ticks(include_xlimits=include_xlimits)
+    ax.set_xticks(np.arange(0, 25, 4, dtype=int))
+    ax.set_xticklabels(ticks_x_labels, fontsize=12)
+    ax.set_xlim(0, 24)
+
+
+def _get_standard_time_ticks(
+    resolution: int = 4,
+    *,
+    include_xlimits: bool = False,
+) -> list:
+    """Returns typical ticks / labels for a time vector between 0-24h."""
+    if include_xlimits:
+        return [
+            f"{int(i):02d}:00" if 24 >= i >= 0 else ""
+            for i in np.arange(0, 24.01, resolution)
+        ]
+    return [
+        f"{int(i):02d}:00" if 24 > i > 0 else ""
+        for i in np.arange(0, 24.01, resolution)
+    ]
+
+
+def set_labels(fig: Figure, ax: Axes, nc_file: str, *, sub_title: bool = True) -> date:
+    ax.set_xlabel("Time (UTC)", fontsize=13)
+    case_date = read_date(nc_file)
+    site_name = read_location(nc_file)
+    if sub_title:
+        add_subtitle(fig, case_date, site_name)
+    return case_date
+
+
+def read_location(nc_file: str) -> str:
+    """Returns site name."""
+    with netCDF4.Dataset(nc_file) as nc:
+        return nc.location
+
+
+def read_date(nc_file: str) -> date:
+    """Returns measurement date."""
+    with netCDF4.Dataset(nc_file) as nc:
+        return date(int(nc.year), int(nc.month), int(nc.day))
+
+
+def add_subtitle(fig: Figure, case_date: date, site_name: str) -> None:
+    """Adds subtitle into figure."""
+    text = _get_subtitle_text(case_date, site_name)
+    fig.suptitle(
+        text,
+        fontsize=13,
+        y=0.885,
+        x=0.07,
+        horizontalalignment="left",
+        verticalalignment="bottom",
+    )
+
+
+def _get_subtitle_text(case_date: date, site_name: str) -> str:
+    site_name = site_name.replace("-", " ")
+    return f"{site_name}, {case_date.strftime('%d %b %Y').lstrip('0')}"
+
+
+def _create_save_name(
+    save_path: str,
+    case_date: date,
+    field_names: list,
+    fix: str = "",
+) -> str:
+    """Creates file name for saved images."""
+    date_string = case_date.strftime("%Y%m%d")
+    return f"{save_path}{date_string}_{'_'.join(field_names)}{fix}.png"
+
+
+def handle_saving(
+    image_name: str | None,
+    save_path: str | None,
+    case_date: date,
+    field_names: list,
+    fix: str = "",
+    *,
+    show: bool = False,
+) -> None:
+    if image_name:
+        plt.savefig(image_name, bbox_inches="tight")
+    elif save_path:
+        file_name = _create_save_name(save_path, case_date, field_names, fix)
+        plt.savefig(file_name, bbox_inches="tight")
+    if show:
+        plt.show()
+    plt.close()