|
| 1 | +from typing import Optional |
| 2 | +from packaging import version |
| 3 | + |
| 4 | +import numpy as np |
| 5 | +import matplotlib |
| 6 | +from matplotlib import pyplot |
| 7 | +from open_atmos_jupyter_utils import show_plot |
| 8 | + |
| 9 | +from PySDM.builder import Builder |
| 10 | +from PySDM.dynamics import Coalescence |
| 11 | +from PySDM.environments import Box |
| 12 | +from PySDM.products import ( |
| 13 | + ParticleVolumeVersusRadiusLogarithmSpectrum, |
| 14 | + WallTime, |
| 15 | + CollisionRateDeficitPerGridbox, |
| 16 | +) |
| 17 | + |
| 18 | +from PySDM import Formulae |
| 19 | +from PySDM.dynamics.collisions.collision_kernels import Golovin |
| 20 | +from PySDM.initialisation import spectra |
| 21 | +from PySDM.physics import si |
| 22 | + |
| 23 | +_matplotlib_version_3_3_3 = version.parse("3.3.0") |
| 24 | +_matplotlib_version_actual = version.parse(matplotlib.__version__) |
| 25 | + |
| 26 | + |
| 27 | +def error_measure(y, y_true, x): |
| 28 | + # The length of each bin on a logarithmic scale. |
| 29 | + dlnr = np.gradient(np.log(x)) |
| 30 | + |
| 31 | + F = np.cumsum(y * dlnr) |
| 32 | + CDF_Golovin = np.cumsum(y_true * dlnr) |
| 33 | + |
| 34 | + return np.trapz(np.abs(CDF_Golovin - F), np.log(x)) |
| 35 | + |
| 36 | + |
| 37 | +class Settings: |
| 38 | + def __init__(self: int, steps: Optional[list] = None): |
| 39 | + steps = steps or [0, 1200, 2400, 3600] |
| 40 | + self.formulae = Formulae() |
| 41 | + self.n_sd = 2**13 |
| 42 | + self.n_part = 2**23 / si.metre**3 |
| 43 | + self.X0 = self.formulae.trivia.volume(radius=30.531 * si.micrometres) |
| 44 | + self.dv = 1e6 * si.metres**3 |
| 45 | + self.norm_factor = self.n_part * self.dv |
| 46 | + self.rho = 1000 * si.kilogram / si.metre**3 |
| 47 | + self.dt = 1 * si.seconds |
| 48 | + self.adaptive = False |
| 49 | + self.steps = steps |
| 50 | + self.kernel = Golovin(b=1.5e3 / si.second) |
| 51 | + self.spectrum = spectra.Exponential(norm_factor=self.norm_factor, scale=self.X0) |
| 52 | + self.radius_bins_edges = np.logspace( |
| 53 | + np.log10(10 * si.um), np.log10(5e3 * si.um), num=128, endpoint=True |
| 54 | + ) |
| 55 | + |
| 56 | + @property |
| 57 | + def output_steps(self): |
| 58 | + return [int(step / self.dt) for step in self.steps] |
| 59 | + |
| 60 | + |
| 61 | +class SpectrumColors: |
| 62 | + def __init__(self, begining="#2cbdfe", end="#b317b1"): |
| 63 | + self.b = begining |
| 64 | + self.e = end |
| 65 | + |
| 66 | + def __call__(self, value: float): |
| 67 | + bR, bG, bB = int(self.b[1:3], 16), int(self.b[3:5], 16), int(self.b[5:7], 16) |
| 68 | + eR, eG, eB = int(self.e[1:3], 16), int(self.e[3:5], 16), int(self.e[5:7], 16) |
| 69 | + R = bR + int((eR - bR) * value) |
| 70 | + G = bG + int((eG - bG) * value) |
| 71 | + B = bB + int((eB - bB) * value) |
| 72 | + result = f"#{hex(R)[2:4]}{hex(G)[2:4]}{hex(B)[2:4]}" |
| 73 | + return result |
| 74 | + |
| 75 | + |
| 76 | +class SpectrumPlotter: |
| 77 | + def __init__(self, settings, title=None, grid=True, legend=True, log_base=10): |
| 78 | + self.settings = settings |
| 79 | + self.format = "pdf" |
| 80 | + self.colors = SpectrumColors() |
| 81 | + self.smooth = False |
| 82 | + self.smooth_scope = 2 |
| 83 | + self.legend = legend |
| 84 | + self.grid = grid |
| 85 | + self.title = title |
| 86 | + self.xlabel = "particle radius [µm]" |
| 87 | + self.ylabel = "dm/dlnr [g/m^3/(unit dr/r)]" |
| 88 | + self.log_base = log_base |
| 89 | + self._ax = None |
| 90 | + self.finished = False |
| 91 | + |
| 92 | + @property |
| 93 | + def ax(self): |
| 94 | + return self._ax or pyplot.gca() |
| 95 | + |
| 96 | + @ax.setter |
| 97 | + def ax(self, value): |
| 98 | + self._ax = value |
| 99 | + |
| 100 | + def finish(self): |
| 101 | + if self.finished: |
| 102 | + return |
| 103 | + self.finished = True |
| 104 | + if self.grid: |
| 105 | + self.ax.grid() |
| 106 | + |
| 107 | + base_arg = { |
| 108 | + "base" |
| 109 | + + ( |
| 110 | + "x" if _matplotlib_version_actual < _matplotlib_version_3_3_3 else "" |
| 111 | + ): self.log_base |
| 112 | + } |
| 113 | + if self.title is not None: |
| 114 | + self.ax.set_title(self.title) |
| 115 | + self.ax.set_xscale("log", **base_arg) |
| 116 | + self.ax.set_xlabel(self.xlabel) |
| 117 | + self.ax.set_ylabel(self.ylabel) |
| 118 | + if self.legend: |
| 119 | + self.ax.legend() |
| 120 | + |
| 121 | + def show(self): |
| 122 | + self.finish() |
| 123 | + pyplot.tight_layout() |
| 124 | + show_plot() |
| 125 | + |
| 126 | + def save(self, file): |
| 127 | + self.finish() |
| 128 | + pyplot.savefig(file, format=self.format) |
| 129 | + |
| 130 | + def plot( |
| 131 | + self, spectrum, t, label=None, color=None, title=None, add_error_to_label=False |
| 132 | + ): |
| 133 | + error = self.plot_analytic_solution(self.settings, t, spectrum, title) |
| 134 | + if label is not None and add_error_to_label: |
| 135 | + label += f" error={error:.4g}" |
| 136 | + self.plot_data(self.settings, t, spectrum, label, color) |
| 137 | + return error |
| 138 | + |
| 139 | + def plot_analytic_solution(self, settings, t, spectrum, title): |
| 140 | + if t == 0: |
| 141 | + analytic_solution = settings.spectrum.size_distribution |
| 142 | + else: |
| 143 | + |
| 144 | + def analytic_solution(x): |
| 145 | + return settings.norm_factor * settings.kernel.analytic_solution( |
| 146 | + x=x, t=t, x_0=settings.X0, N_0=settings.n_part |
| 147 | + ) |
| 148 | + |
| 149 | + volume_bins_edges = self.settings.formulae.trivia.volume( |
| 150 | + settings.radius_bins_edges |
| 151 | + ) |
| 152 | + dm = np.diff(volume_bins_edges) |
| 153 | + dr = np.diff(settings.radius_bins_edges) |
| 154 | + |
| 155 | + pdf_m_x = volume_bins_edges[:-1] + dm / 2 |
| 156 | + pdf_m_y = analytic_solution(pdf_m_x) |
| 157 | + |
| 158 | + pdf_r_x = settings.radius_bins_edges[:-1] + dr / 2 |
| 159 | + pdf_r_y = pdf_m_y * dm / dr * pdf_r_x |
| 160 | + |
| 161 | + x = pdf_r_x * si.metres / si.micrometres |
| 162 | + y_true = ( |
| 163 | + pdf_r_y |
| 164 | + * self.settings.formulae.trivia.volume(radius=pdf_r_x) |
| 165 | + * settings.rho |
| 166 | + / settings.dv |
| 167 | + * si.kilograms |
| 168 | + / si.grams |
| 169 | + ) |
| 170 | + |
| 171 | + self.ax.plot(x, y_true, color="black") |
| 172 | + |
| 173 | + if spectrum is not None: |
| 174 | + y = spectrum * si.kilograms / si.grams |
| 175 | + error = error_measure(y, y_true, x) |
| 176 | + self.title = ( |
| 177 | + title or f"error measure: {error:.2f}" |
| 178 | + ) # TODO #327 relative error |
| 179 | + return error |
| 180 | + return None |
| 181 | + |
| 182 | + def plot_data(self, settings, t, spectrum, label, color): |
| 183 | + if self.smooth: |
| 184 | + scope = self.smooth_scope |
| 185 | + if t != 0: |
| 186 | + new = np.copy(spectrum) |
| 187 | + for _ in range(2): |
| 188 | + for i in range(scope, len(spectrum) - scope): |
| 189 | + new[i] = np.mean(spectrum[i - scope : i + scope + 1]) |
| 190 | + scope = 1 |
| 191 | + for i in range(scope, len(spectrum) - scope): |
| 192 | + spectrum[i] = np.mean(new[i - scope : i + scope + 1]) |
| 193 | + |
| 194 | + x = settings.radius_bins_edges[:-scope] |
| 195 | + dx = np.diff(x) |
| 196 | + self.ax.plot( |
| 197 | + (x[:-1] + dx / 2) * si.metres / si.micrometres, |
| 198 | + spectrum[:-scope] * si.kilograms / si.grams, |
| 199 | + label=label or f"t = {t}s", |
| 200 | + color=color |
| 201 | + or self.colors(t / (self.settings.output_steps[-1] * self.settings.dt)), |
| 202 | + ) |
| 203 | + else: |
| 204 | + self.ax.step( |
| 205 | + settings.radius_bins_edges[:-1] * si.metres / si.micrometres, |
| 206 | + spectrum * si.kilograms / si.grams, |
| 207 | + where="post", |
| 208 | + label=label or f"t = {t}s", |
| 209 | + color=color |
| 210 | + or self.colors(t / (self.settings.output_steps[-1] * self.settings.dt)), |
| 211 | + ) |
| 212 | + |
| 213 | + |
| 214 | +def run(settings, backend, observers=(), sampling_method="deterministic"): |
| 215 | + builder = Builder( |
| 216 | + n_sd=settings.n_sd, |
| 217 | + backend=backend, |
| 218 | + environment=Box(dv=settings.dv, dt=settings.dt), |
| 219 | + dynamics=( |
| 220 | + Coalescence(collision_kernel=settings.kernel, adaptive=settings.adaptive), |
| 221 | + ), |
| 222 | + ) |
| 223 | + builder.particulator.environment["rhod"] = 1.0 |
| 224 | + attributes = {} |
| 225 | + sampling = settings.sampling |
| 226 | + attributes["volume"], attributes["multiplicity"] = getattr( |
| 227 | + sampling, f"sample_{sampling_method}" |
| 228 | + )(settings.n_sd, backend=backend) |
| 229 | + products = ( |
| 230 | + ParticleVolumeVersusRadiusLogarithmSpectrum( |
| 231 | + settings.radius_bins_edges, name="dv/dlnr" |
| 232 | + ), |
| 233 | + WallTime(), |
| 234 | + CollisionRateDeficitPerGridbox(name="deficit"), |
| 235 | + ) |
| 236 | + particulator = builder.build(attributes, products) |
| 237 | + |
| 238 | + for observer in observers: |
| 239 | + particulator.observers.append(observer) |
| 240 | + |
| 241 | + vals = {} |
| 242 | + deficit = 0 |
| 243 | + particulator.products["wall time"].reset() |
| 244 | + for step in settings.output_steps: |
| 245 | + particulator.run(step - particulator.n_steps) |
| 246 | + vals[step] = particulator.products["dv/dlnr"].get()[0] |
| 247 | + vals[step][:] *= settings.rho |
| 248 | + deficit += particulator.products["deficit"].get() |
| 249 | + deficit = deficit / len(settings.output_steps) |
| 250 | + |
| 251 | + exec_time = particulator.products["wall time"].get() |
| 252 | + return vals, exec_time, deficit |
0 commit comments