PyMPDATA_examples.Olesik_et_al_2022.plotter

  1import matplotlib
  2import numpy as np
  3from matplotlib import pyplot
  4
  5
  6def to_n_n(y, _, __):
  7    return y
  8
  9
 10def to_n_s(y, r1, r2):
 11    return y * (r2**2 + r1**2 + r1 * r2) * 4 / 3 * np.pi
 12
 13
 14def to_n_v(y, r1, r2):
 15    return y * (r2 + r1) * (r2**2 + r1**2) / 4 * 4 / 3 * np.pi
 16
 17
 18def from_n_n(n_n, _):
 19    return 1 * n_n
 20
 21
 22class Plotter:
 23    def __init__(self, settings, plots=("n", "m")):
 24        self.settings = settings
 25        self.plots = plots
 26        self.cdfarg, self.dcdfarg = np.linspace(
 27            settings.r_min.magnitude, settings.r_max.magnitude, 512, retstep=True
 28        )
 29        self.cdfarg *= settings.r_min.units
 30        self.dcdfarg *= settings.r_max.units
 31
 32        matplotlib.rcParams.update({"font.size": 16})
 33        if len(plots) == 1:
 34            self.figsize = (10, 6)
 35        else:
 36            self.figsize = (10, 11)
 37        self.fig, self.axs = pyplot.subplots(len(plots), 1, figsize=self.figsize)
 38        if len(plots) == 1:
 39            self.axs = (self.axs,)
 40        self.fig.tight_layout(pad=5.0)
 41        self.style_dict = {}
 42        self.style_palette = ["-", "-", "-", "-", "-."]
 43
 44        self.settings.si.setup_matplotlib()
 45
 46        if "n" in plots:
 47            self.axs[plots.index("n")].yaxis.set_units(
 48                1 / self.settings.si.micrometre / self.settings.si.centimetre**3
 49            )
 50        if "m" in plots:
 51            self.axs[plots.index("m")].yaxis.set_units(1 / self.settings.si.micrometre)
 52
 53        for i in range(len(plots)):
 54            self.axs[i].xaxis.set_units(self.settings.si.micrometre)
 55            self.axs[i].grid()
 56            self.axs[i].set_xlabel("particle radius [$μm$]")
 57
 58        if "m" in plots:
 59            self.axs[plots.index("m")].set_ylabel("(dM/dr)/M [$μm^{-1}$]")
 60        if "n" in plots:
 61            self.axs[plots.index("n")].set_ylabel("dN/dr [$cm^{-3} μm^{-1}$]")
 62
 63    def pdf_curve(self, pdf, mnorm, color="black"):
 64        x = self.cdfarg
 65        y = pdf(x)
 66
 67        # normalised mass distribution
 68        if "m" in self.plots:
 69            y_mass = (
 70                y
 71                * x**3
 72                * 4
 73                / 3
 74                * np.pi
 75                * self.settings.rho_w
 76                / self.settings.rho_a
 77                / mnorm
 78            )
 79            self.axs[self.plots.index("m")].plot(x, y_mass, color=color, linestyle=":")
 80
 81        # number distribution
 82        if "n" in self.plots:
 83            self.axs[self.plots.index("n")].plot(x, y, color=color, linestyle=":")
 84
 85    def _plot(self, index, x, y, fill, lbl, color, linewidth):
 86        if fill:
 87            self.axs[self.plots.index(index)].fill_between(
 88                x, y, step="mid", label=lbl, color="gainsboro", alpha=1
 89            )
 90        else:
 91            self.axs[self.plots.index(index)].step(
 92                x, y, where="mid", label=lbl, color=color, linewidth=linewidth
 93            )
 94        self.xlim(index)
 95
 96    def pdf_histogram(
 97        self, x, y, bin_boundaries, label, mnorm, color="black", linewidth=1, fill=True
 98    ):
 99        r1 = bin_boundaries[:-1]
100        r2 = bin_boundaries[1:]
101
102        lbl = label
103        if label not in self.style_dict:
104            self.style_dict[label] = self.style_palette[len(self.style_dict)]
105        else:
106            lbl = ""
107
108        # number distribution
109        if "n" in self.plots:
110            self._plot(
111                "n",
112                x,
113                to_n_n(y, r1, r2),
114                fill=fill,
115                lbl=lbl,
116                color=color,
117                linewidth=linewidth,
118            )
119
120        # normalised mass distribution
121        if "m" in self.plots:
122            self._plot(
123                "m",
124                x,
125                to_n_v(y, r1, r2) * self.settings.rho_w / self.settings.rho_a / mnorm,
126                fill=fill,
127                color=color,
128                linewidth=linewidth,
129                lbl=lbl,
130            )
131
132    def xlim(self, plot):
133        self.axs[self.plots.index(plot)].set_xlim((0, self.settings.r_max.magnitude))