PySDM_examples.Szumowski_et_al_1998.gui_viewer

  1import warnings
  2
  3import numpy as np
  4from matplotlib import pyplot, rcParams
  5from open_atmos_jupyter_utils import save_and_make_link
  6from PySDM_examples.Szumowski_et_al_1998.plots import (
  7    _ImagePlot,
  8    _SpectrumPlot,
  9    _TemperaturePlot,
 10    _TerminalVelocityPlot,
 11    _TimeseriesPlot,
 12)
 13from PySDM_examples.utils.widgets import (
 14    Box,
 15    Button,
 16    Dropdown,
 17    HBox,
 18    IntRangeSlider,
 19    IntSlider,
 20    Layout,
 21    Output,
 22    Play,
 23    VBox,
 24    clear_output,
 25    display,
 26    jslink,
 27)
 28
 29from PySDM.physics import constants as const
 30
 31
 32class GUIViewer:
 33    def __init__(self, storage, settings):
 34        self.storage = storage
 35        self.settings = settings
 36
 37        self.play = Play(interval=1000)
 38        self.step_slider = IntSlider(continuous_update=False, description="t/dt_out:")
 39        self.product_select = Dropdown()
 40        self.spectrum_select = Dropdown()
 41        self.plots_box = Box()
 42
 43        self.timeseriesOutput = None
 44        self.timeseriesPlot = None
 45        self.plot_box = None
 46        self.spectrum_box = None
 47        self.outputs = None
 48        self.plots = None
 49        self.spectrumOutputs = None
 50        self.spectrumPlots = None
 51        self.products = None
 52
 53        self.slider = {}
 54        self.lines = {"X": [{}, {}], "Z": [{}, {}]}
 55        for xz in ("X", "Z"):
 56            self.slider[xz] = IntRangeSlider(
 57                min=0,
 58                max=1,
 59                description=f"{xz}",
 60                continuous_update=False,
 61                orientation="horizontal" if xz == "X" else "vertical",
 62            )
 63
 64        self.reinit({})
 65
 66    def clear(self):
 67        self.plots_box.children = ()
 68
 69    def reinit(self, products):
 70        self.products = products
 71        self.product_select.options = tuple(
 72            (f"{val.name} [{val.unit}]", key)
 73            for key, val in sorted(self.products.items(), key=lambda item: item[1].name)
 74            if len(val.shape) == 2
 75        )
 76        opts = [
 77            ("particle size spectra", "size"),
 78            ("terminal velocity", "terminal velocity"),
 79        ]
 80        if "freezable specific concentration" in products:
 81            opts.append(("freezing temperature spectra", "temperature"))
 82        if "immersed surface area" in products:
 83            pass  # TODO #599
 84        self.spectrum_select.options = tuple(opts)
 85
 86        r_bins = self.settings.r_bins_edges.copy()
 87        const.convert_to(r_bins, const.si.micrometres)
 88        self.spectrumOutputs = {}
 89        self.spectrumPlots = {}
 90        for key in ("size", "terminal velocity", "temperature"):
 91            self.spectrumOutputs[key] = Output()
 92            with self.spectrumOutputs[key]:
 93                self.spectrumPlots[key] = (
 94                    _SpectrumPlot(r_bins, self.settings.spectrum_per_mass_of_dry_air)
 95                    if key == "size"
 96                    else (
 97                        _TemperaturePlot(
 98                            self.settings.T_bins_edges, self.settings.formulae
 99                        )
100                        if key == "temperature"
101                        else _TerminalVelocityPlot(
102                            self.settings.terminal_velocity_radius_bin_edges,
103                            self.settings.formulae,
104                        )
105                    )
106                )
107                clear_output()
108
109        self.timeseriesOutput = Output()
110        with self.timeseriesOutput:
111            default_figsize = rcParams["figure.figsize"]
112            fig_kw = {"figsize": (2.5 * default_figsize[0], default_figsize[1] / 2)}
113            fig, ax = pyplot.subplots(1, 1, **fig_kw)
114            self.timeseriesPlot = _TimeseriesPlot(
115                fig, ax, self.settings.output_steps * self.settings.dt
116            )
117            clear_output()
118
119        self.plots = {}
120        self.outputs = {}
121        for key, product in products.items():
122            if len(product.shape) == 2:
123                self.outputs[key] = Output()
124                with self.outputs[key]:
125                    fig, ax = pyplot.subplots(1, 1)
126                    self.plots[key] = _ImagePlot(
127                        fig,
128                        ax,
129                        self.settings.grid,
130                        self.settings.size,
131                        product,
132                        show=True,
133                        lines=True,
134                    )
135                    clear_output()
136
137        self.plot_box = Box()
138        self.spectrum_box = Box()
139        if len(products.keys()) > 0:
140            layout_flex_end = Layout(display="flex", justify_content="flex-end")
141            save_map = Button(icon="save")
142            save_map.on_click(self.handle_save_map)
143            save_spe = Button(icon="save")
144            save_spe.on_click(self.handle_save_spe)
145            self.plots_box.children = (
146                VBox(
147                    children=(
148                        HBox(
149                            children=(
150                                VBox(
151                                    children=(
152                                        Box(
153                                            layout=layout_flex_end,
154                                            children=(save_map, self.product_select),
155                                        ),
156                                        HBox((self.slider["Z"], self.plot_box)),
157                                        HBox(
158                                            (self.slider["X"],), layout=layout_flex_end
159                                        ),
160                                    )
161                                ),
162                                VBox(
163                                    layout=Layout(),
164                                    children=(
165                                        Box(
166                                            children=(save_spe, self.spectrum_select),
167                                            layout=layout_flex_end,
168                                        ),
169                                        self.spectrum_box,
170                                    ),
171                                ),
172                            )
173                        ),
174                        HBox((self.timeseriesOutput,)),
175                    )
176                ),
177            )
178
179        for widget in (self.step_slider, self.play):
180            widget.value = 0
181            widget.max = len(self.settings.output_steps) - 1
182
183        for j, xz in enumerate(("X", "Z")):
184            slider = self.slider[xz]
185            mx = self.settings.grid[j]
186            slider.max = mx
187            slider.value = (0, mx)
188
189        self.replot()
190
191    def handle_save_map(self, _):
192        display(save_and_make_link(self.plots[self.product_select.value].fig))
193
194    def handle_save_spe(self, _):
195        display(save_and_make_link(self.spectrumPlots[self.spectrum_select.value].fig))
196
197    def replot(self, *_):
198        selectedImage = self.product_select.value
199        if not (selectedImage is None or selectedImage not in self.plots):
200            self.replot_image()
201            self.outputs[selectedImage].clear_output(wait=True)
202            with self.outputs[selectedImage]:
203                display(self.plots[selectedImage].fig)
204
205        selectedSpectrum = self.spectrum_select.value
206        if not (selectedSpectrum is None or selectedSpectrum not in self.spectrumPlots):
207            self.replot_spectra()
208            self.spectrumOutputs[selectedSpectrum].clear_output(wait=True)
209            with self.spectrumOutputs[selectedSpectrum]:
210                display(self.spectrumPlots[selectedSpectrum].fig)
211
212        self.update_timeseries()
213        self.timeseriesOutput.clear_output(wait=True)
214        with self.timeseriesOutput:
215            display(self.timeseriesPlot.fig)
216
217    def update_spectra(self):
218        selected = self.spectrum_select.value
219        self.spectrum_box.children = [self.spectrumOutputs[selected]]
220        plot = self.spectrumPlots[selected]
221
222        step = self.step_slider.value
223
224        xrange = slice(*self.slider["X"].value)
225        yrange = slice(*self.slider["Z"].value)
226
227        if selected == "size":
228            for key in ("Particles Wet Size Spectrum", "Particles Dry Size Spectrum"):
229                if xrange.start == xrange.stop or yrange.start == yrange.stop:
230                    continue
231                try:
232                    data = self.storage.load(key, self.settings.output_steps[step])
233                    data = data[xrange, yrange, :]
234                    data = np.mean(np.mean(data, axis=0), axis=0)
235                    data = np.concatenate(((0,), data))
236                    if key == "Particles Wet Size Spectrum":
237                        plot.update_wet(data, step)
238                    if key == "Particles Dry Size Spectrum":
239                        plot.update_dry(data)
240                except self.storage.Exception:
241                    pass
242        elif selected == "terminal velocity":
243            try:
244                data = self.storage.load(
245                    "radius binned number averaged terminal velocity",
246                    self.settings.output_steps[step],
247                )
248
249                data = data[xrange, yrange, :]
250                data = np.where(data != 0, data, np.nan)
251                try:
252                    with warnings.catch_warnings(record=True) as _:
253                        warnings.simplefilter("ignore")
254                        data_min = np.nanmin(np.nanmin(data, axis=0), axis=0)
255                        data_max = np.nanmax(np.nanmax(data, axis=0), axis=0)
256                except RuntimeWarning:
257                    pass
258                plot.update(data_min, data_max, step)
259            except self.storage.Exception:
260                pass
261        elif selected == "temperature":
262            try:
263                dT = abs(self.settings.T_bins_edges[1] - self.settings.T_bins_edges[0])
264                np.testing.assert_allclose(np.diff(self.settings.T_bins_edges), dT)
265
266                conc = self.storage.load(
267                    "particle specific concentration", self.settings.output_steps[step]
268                )
269                # TODO #705: assert unit == mg^-1
270                conc = conc[xrange, yrange]
271
272                data = self.storage.load(
273                    "freezable specific concentration", self.settings.output_steps[step]
274                )
275                data = data[xrange, yrange, :]
276
277                data = np.sum(np.sum(data, axis=0), axis=0) / np.sum(
278                    np.sum(conc, axis=0), axis=0
279                )
280                data = np.concatenate(((0,), dT * np.cumsum(data[::-1])))[::-1]
281
282                plot.update(data, step)
283            except self.storage.Exception:
284                pass
285        else:
286            raise NotImplementedError()
287
288    def replot_spectra(self, *_):
289        self.update_spectra()
290
291        selected = self.product_select.value
292        if selected is None or selected not in self.plots:
293            return
294        self.plots[selected].update_lines(
295            self.slider["X"].value, self.slider["Z"].value
296        )
297
298        self.outputs[selected].clear_output(wait=True)
299
300        key = self.spectrum_select.value
301        self.spectrumOutputs[key].clear_output(wait=True)
302        with self.outputs[selected]:
303            display(self.plots[selected].fig)
304        with self.spectrumOutputs[key]:
305            display(self.spectrumPlots[key].fig)
306
307    def update_image(self):
308        selected = self.product_select.value
309
310        if selected in self.outputs:
311            self.plot_box.children = [self.outputs[selected]]
312
313        step = self.step_slider.value
314        try:
315            data = self.storage.load(selected, self.settings.output_steps[step])
316        except self.storage.Exception:
317            data = None
318
319        self.plots[selected].update(
320            data, step, self.storage.data_range(selected) if data is not None else None
321        )
322
323    def replot_image(self, *_):
324        selected = self.product_select.value
325        if selected is None or selected not in self.plots:
326            return
327
328        self.update_image()
329        self.outputs[selected].clear_output(wait=True)
330        with self.outputs[selected]:
331            display(self.plots[selected].fig)
332
333    def update_timeseries(self):
334        try:
335            data = self.storage.load("surf_precip")
336        except self.storage.Exception:
337            data = None
338        self.timeseriesPlot.update(
339            data, self.storage.data_range("surf_precip") if data is not None else None
340        )
341
342    def replot_timeseries(self):
343        self.update_timeseries()
344        self.timeseriesOutput.clear_output(wait=True)
345        with self.timeseriesOutput:
346            display(self.timeseriesPlot.fig)
347
348    def box(self):
349        jslink((self.play, "value"), (self.step_slider, "value"))
350        self.step_slider.observe(self.replot, "value")
351        self.product_select.observe(self.replot_image, "value")
352        self.spectrum_select.observe(self.replot_spectra, "value")
353        for xz in ("X", "Z"):
354            self.slider[xz].observe(self.replot_spectra, "value")
355        return VBox([Box([self.play, self.step_slider]), self.plots_box])