PySDM_examples.Jensen_and_Nugent_2017.plotting

  1from matplotlib import pyplot
  2import numpy as np
  3from PySDM.physics import si, in_unit
  4from PySDM.physics.constants import PER_CENT
  5
  6CLOUD_BASE = 300 * si.m
  7
  8
  9def find_drop_ids_by_dry_size(plot_drops_with_dry_radii_um, simulation_r_dry):
 10    drop_ids = []
 11    for drop_size_um in plot_drops_with_dry_radii_um:
 12        drop_id = (np.abs(simulation_r_dry - drop_size_um * si.um)).argmin()
 13        drop_ids.append(drop_id)
 14    return drop_ids
 15
 16
 17def compute_table_values(height_above_cb_m, height_cb_m, products, ascent_mask):
 18    """constructing data for Table 4"""
 19    data = {"ascent": [], "descent": []}
 20    for level_idx, height_m in enumerate(products["z"]):
 21        rounded_height_above_cb_m = np.round(height_m - height_cb_m)
 22        if rounded_height_above_cb_m in height_above_cb_m:
 23            r_mean = products["r_mean_act"][level_idx]
 24            r_standard_deviation = products["r_std_act"][level_idx]
 25            r_relative_dispersion = r_standard_deviation / r_mean
 26
 27            data["ascent" if ascent_mask[level_idx] else "descent"].append(
 28                (
 29                    f"{rounded_height_above_cb_m:.0f}",
 30                    f"{in_unit(r_mean, si.um):.2f}",
 31                    f"{in_unit(r_standard_deviation, si.um):.2f}",
 32                    f"{r_relative_dispersion:.3f}",
 33                )
 34            )
 35    data["ascent"].append(data["descent"][0])
 36    return data
 37
 38
 39def figure(
 40    *,
 41    output,
 42    settings,
 43    simulation,
 44    plot_drops_with_dry_radii_um,
 45    xlim_r_um: tuple,
 46    xlim_S_percent: tuple,
 47    return_masks: bool = False,
 48):
 49    y_axis = np.asarray(output["products"]["z"]) - settings.z0 - CLOUD_BASE
 50
 51    masks = {}
 52    if settings.t_end_of_ascent is None:
 53        masks["ascent"] = np.full_like(output["products"]["t"], True, dtype=bool)
 54    else:
 55        masks["ascent"] = np.asarray(output["products"]["t"]) < settings.t_end_of_ascent
 56        masks["descent"] = np.logical_not(masks["ascent"])
 57
 58    colors = {"ascent": "r", "descent": "b"}
 59
 60    _, axs = pyplot.subplot_mosaic(
 61        mosaic=[["r", "S"]], width_ratios=[3, 1], sharey=True, tight_layout=True
 62    )
 63
 64    for label, mask in masks.items():
 65        axs["S"].plot(
 66            in_unit(np.asarray(output["products"]["S_max"]), PER_CENT)[mask],
 67            y_axis[mask],
 68            label=label,
 69            color=colors[label],
 70        )
 71    axs["S"].set_xlim(*xlim_S_percent)
 72    axs["S"].set_xlabel("S (%)", fontsize="15")
 73    axs["S"].legend(fontsize="10")
 74
 75    drop_ids = find_drop_ids_by_dry_size(
 76        plot_drops_with_dry_radii_um=plot_drops_with_dry_radii_um,
 77        simulation_r_dry=simulation.r_dry,
 78    )
 79
 80    for (
 81        drop_id
 82    ) in (
 83        drop_ids
 84    ):  # TODO #1266: bug! why rightmost drop is not 500 nm if size range is set to end at 500 nm???
 85        for label, mask in masks.items():
 86            axs["r"].plot(
 87                in_unit(np.asarray(output["attributes"]["radius"][drop_id]), si.um)[
 88                    mask
 89                ],
 90                y_axis[mask],
 91                label=(
 92                    f"{in_unit(simulation.r_dry[drop_id], si.um):.2} µm"
 93                    if label == "ascent"
 94                    else ""
 95                ),
 96                color=colors[label],
 97            )
 98    axs["r"].legend()
 99    axs["r"].set_xlim(*xlim_r_um)
100    axs["r"].set_xlabel("r$_c$ (µm)", fontsize="15")
101    axs["r"].set_ylabel("Height above cloud base (m)", fontsize="15")
102    # pyplot.xticks(fontsize=12)
103    # pyplot.yticks(fontsize=12)
104
105    for ax in axs.values():
106        ax.grid()
107
108    if return_masks:
109        return masks
110    return None