PySDM_examples.featured.Shima_et_al_2009.spectrum_plotter

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