PySDM_examples.Shima_et_al_2009.tutorial_plotter

  1import numpy as np
  2from matplotlib import pyplot
  3from open_atmos_jupyter_utils import show_plot
  4
  5from PySDM.physics.constants import si
  6
  7
  8class SpectrumColors:
  9    def __init__(self, begining="#2cbdfe", end="#b317b1"):
 10        self.b = begining
 11        self.e = end
 12
 13    def __call__(self, value: float):
 14        bR, bG, bB = int(self.b[1:3], 16), int(self.b[3:5], 16), int(self.b[5:7], 16)
 15        eR, eG, eB = int(self.e[1:3], 16), int(self.e[3:5], 16), int(self.e[5:7], 16)
 16        R = bR + int((eR - bR) * value)
 17        G = bG + int((eG - bG) * value)
 18        B = bB + int((eB - bB) * value)
 19        result = f"#{hex(R)[2:4]}{hex(G)[2:4]}{hex(B)[2:4]}"
 20        return result
 21
 22
 23class SpectrumPlotter:
 24    def __init__(self, settings, grid=True, legend=True, log_base=10):
 25        self.settings = settings
 26        self.format = "pdf"
 27        self.colors = SpectrumColors()
 28        self.smooth = False
 29        self.smooth_scope = 2
 30        self.legend = legend
 31        self.grid = grid
 32        self.xlabel = "particle radius [µm]"
 33        self.ylabel = "dm/dlnr [g/m$^3$]"
 34        self.log_base = log_base
 35        self.ax = pyplot
 36        self.fig = pyplot
 37        self.finished = False
 38
 39    def finish(self):
 40        if self.finished:
 41            return
 42        self.finished = True
 43        if self.grid:
 44            self.ax.grid()
 45
 46        self.ax.xscale("log")
 47        self.ax.xlabel(self.xlabel)
 48        self.ax.ylabel(self.ylabel)
 49        if self.legend:
 50            self.ax.legend()
 51
 52    def show(self):
 53        self.finish()
 54        pyplot.tight_layout()
 55        show_plot()
 56
 57    def save(self, file):
 58        self.finish()
 59        pyplot.savefig(file, format=self.format)
 60
 61    def plot(self, spectrum, t):
 62        self.plot_analytic_solution(self.settings, t)
 63        self.plot_data(self.settings, t, spectrum)
 64
 65    def plot_analytic_solution(self, settings, t):
 66        def analytic_solution(x):
 67            return settings.norm_factor * settings.kernel.analytic_solution(
 68                x=x, t=t, x_0=settings.X0, N_0=settings.n_part
 69            )
 70
 71        if t == 0:
 72            analytic_solution = settings.spectrum.size_distribution
 73
 74        volume_bins_edges = self.settings.formulae.trivia.volume(
 75            settings.radius_bins_edges
 76        )
 77        dm = np.diff(volume_bins_edges)
 78        dr = np.diff(settings.radius_bins_edges)
 79
 80        pdf_m_x = volume_bins_edges[:-1] + dm / 2
 81        pdf_m_y = analytic_solution(pdf_m_x)
 82
 83        pdf_r_x = settings.radius_bins_edges[:-1] + dr / 2
 84        pdf_r_y = pdf_m_y * dm / dr * pdf_r_x
 85
 86        x = pdf_r_x * si.metres / si.micrometres
 87        y_true = (
 88            pdf_r_y
 89            * self.settings.formulae.trivia.volume(radius=pdf_r_x)
 90            * settings.rho
 91            / settings.dv
 92            * si.kilograms
 93            / si.grams
 94        )
 95
 96        self.ax.plot(x, y_true, color="black")
 97
 98    def plot_data(self, settings, t, spectrum):
 99        if self.smooth:
100            scope = self.smooth_scope
101            if t != 0:
102                new = np.copy(spectrum)
103                for _ in range(2):
104                    for i in range(scope, len(spectrum) - scope):
105                        new[i] = np.mean(spectrum[i - scope : i + scope + 1])
106                    scope = 1
107                    for i in range(scope, len(spectrum) - scope):
108                        spectrum[i] = np.mean(new[i - scope : i + scope + 1])
109
110            x = settings.radius_bins_edges[:-scope]
111            dx = np.diff(x)
112            self.ax.plot(
113                (x[:-1] + dx / 2) * si.metres / si.micrometres,
114                spectrum[:-scope] * si.kilograms / si.grams,
115                label=f"t = {t}s",
116                color=self.colors(
117                    t / (self.settings.output_steps[-1] * self.settings.dt)
118                ),
119            )
120        else:
121            self.ax.step(
122                settings.radius_bins_edges[:-1] * si.metres / si.micrometres,
123                spectrum * si.kilograms / si.grams,
124                where="post",
125                label=f"t = {t}s",
126                color=self.colors(
127                    t / (self.settings.output_steps[-1] * self.settings.dt)
128                ),
129            )