blpd.plot

  1"""
  2Create plots of LPD model results.
  3
  4The functions here all take `xr.Dataset`s as their first argument.
  5"""
  6
  7from itertools import cycle
  8
  9import matplotlib as mpl
 10import matplotlib.pyplot as plt
 11import numpy as np
 12from matplotlib.collections import LineCollection as _LineCollection
 13from mpl_toolkits.mplot3d import Axes3D as _Axes3D  # noqa: F401 unused import
 14
 15from . import utils
 16from .utils import _add_snippets
 17
 18# ^ could add `as _{}` to all of these
 19# to indicate that these are not intended to be public parts of the module name space
 20# since __all__ is not respected by linters or autocompleters
 21
 22__all__ = (
 23    "conc_2d",
 24    "conc_scatter",
 25    "conc_xline",
 26    "final_pos_hist",
 27    "final_pos_hist2d",
 28    "final_pos_scatter",
 29    "trajectories",
 30    "ws_hist",
 31)
 32
 33# TODO: create some base classes for plots to reduce repeating of code
 34#       - allow passing fig and ax kwargs in __init__
 35#       - stars to mark sources; check_fig_num, labeling, etc.
 36#       - alpha calculation based on number of particles & spread?
 37
 38# TODO: really all plots could have the auto-bounds stuff. and (optionally?) print message about it?
 39
 40
 41_SOURCE_MARKER_PROPS = dict(marker="*", c="gold", ms=11, mec="0.35", mew=1.0)
 42
 43utils._SNIPPETS.update(
 44    spc_params="""
 45spc : str
 46    ASCII identifier of a chemical species. For example, `'apinene'`.
 47    """.strip(),
 48    ds_lpd_params="""
 49ds : xr.Dataset
 50    Dataset of Lagrangian particle positions, e.g., from `blpd.model.Model.to_xr`.
 51    """.strip(),
 52    ds_chem_params="""
 53ds : xr.Dataset
 54    Dataset of Lagrangian particle positions + chemistry, e.g., from
 55    `blpd.chem.calc_relative_levels_fix_oxidants`.
 56    """.strip(),
 57    log_cnorm_params="""
 58log_cnorm : bool
 59    Whether to use a logarithmic colormap normalization.
 60    """.strip(),
 61)
 62
 63
 64@_add_snippets
 65def final_pos_scatter(ds, sdim="xy"):
 66    """Scatter plot of particle end positions.
 67
 68    Parameters
 69    ----------
 70    %(ds_lpd_params)s
 71    %(sdim_params)s
 72    """
 73    p = utils.load_p(ds)
 74
 75    dims = utils.dims_from_sdim(sdim)
 76    if len(dims) == 3:
 77        sdim = "xyz"
 78        x = ds.x.values
 79        y = ds.y.values
 80        z = ds.z.values
 81        subplot_kw = {"projection": "3d"}
 82        coords = (x, y, z)
 83        plot_kw = dict(alpha=0.5, mew=0, ms=7)
 84    elif len(dims) == 2:
 85        x = ds[dims[0]].values
 86        y = ds[dims[1]].values
 87        subplot_kw = {}
 88        coords = (x, y)
 89        plot_kw = dict(alpha=0.5, mfc="none", mew=0.8, ms=5)
 90    else:
 91        raise ValueError("invalid `sdim`")
 92
 93    num = utils.check_fig_num(f"final-pos-scatter-{sdim}")
 94    fig, ax = plt.subplots(num=num, subplot_kw=subplot_kw)
 95
 96    ax.plot(*coords, "o", **plot_kw)
 97
 98    ax.set_xlabel(f"${dims[0]}$")
 99    ax.set_ylabel(f"${dims[1]}$")
100    if len(dims) == 3:
101        ax.set_zlabel(f"${dims[2]}$")
102    ax.set_title(utils.s_t_info(p), loc="left")
103    ax.set_title(f"$N_p = {p['Np_tot']}$", loc="right")
104
105    for xs, ys in p["source_positions"]:
106        sp = dict(x=xs, y=ys, z=p["release_height"])
107        if len(dims) == 3:
108            ax.plot([sp[dims[0]]], [sp[dims[1]]], [sp[dims[2]]], "*", c="gold", ms=10)
109        else:
110            ax.plot(sp[dims[0]], sp[dims[1]], "*", c="gold", ms=10)
111
112    fig.set_tight_layout(True)
113
114
115# TODO: add option to do trajectories for continuous release runs, colored by time out
116# TODO: trajectories for hist run in 3-D?
117
118
119@_add_snippets
120def trajectories(ds, *, smooth=False, smooth_window_size=None, color_sources=False):
121    """Plot particle trajectories.
122
123    Parameters
124    ----------
125    %(ds_lpd_params)s
126    smooth : bool
127        Whether to smooth the trajectories.
128    smooth_window_size : int, optional
129        How many points to include in the moving average window.
130    color_sources : bool
131        Whether to color particle trajectories by their sources
132        to help in differentiating them.
133    """
134    if "t" not in ds.dims:
135        raise ValueError("time 't' must be a dimension")
136
137    p = utils.load_p(ds)
138    pos = np.stack((ds.x.values, ds.y.values, ds.z.values), axis=-1)  # same shape as hist["pos"]
139
140    t_tot = p["t_tot"]
141    dt = p["dt_out"]  # note output timestep not that of the integration
142    # N_t = p["N_t"]
143    Np = p["Np_tot"]
144    # N = Np * N_t
145    assert pos.shape[0] == Np
146
147    ltitle = utils.s_t_info(p)
148    rtitle = utils.s_sample_size(p)
149
150    # allow specifying smooth_window_size only
151    if smooth_window_size is not None and not smooth:
152        smooth = True
153    if smooth:
154        if smooth_window_size is None:
155            n = 100
156        else:
157            n = int(smooth_window_size)
158
159        if n * dt > 0.5 * t_tot:
160            raise ValueError(
161                "can't do any smoothing with the requested window size (not enough points)"
162            )
163        pos0 = pos[:, 0, :][:, np.newaxis, :]
164        pos = np.swapaxes(utils.moving_average(np.swapaxes(pos, 0, 1), n=n, axis=0), 0, 1)
165        # ^ `np.swapaxes` should return views, not create new arrays (ver >= 1.10.0)
166
167        # preserve starting point
168        pos = np.concatenate((pos0, pos), axis=1)  # axis=1 => basically `hstack`
169
170        ltitle = f"$N_{{smooth}} = {n}$ ({n*dt:} s)\n{ltitle}"  # add smoothing info to left title
171
172    num = utils.check_fig_num("trajectories")
173    fig, ax = plt.subplots(num=num)
174
175    if color_sources:
176        if isinstance(color_sources, (list, np.ndarray)):
177            colors = color_sources  # assume is a list of colors (TODO: should check)
178        else:
179            colors = plt.get_cmap("Dark2").colors
180            # Dark2 is a ListedColormap with 8 colors; `plt.cm.Dark2` same but pylint complains plt.cm `no-member` Dark2
181
182        N_sources = p["N_sources"]
183        for i_source, color in zip(range(N_sources), cycle(colors)):
184            segs = [pos[i, :, :2] for i in range(i_source, Np, N_sources)]
185
186            lc = _LineCollection(
187                segs,
188                linewidths=0.5,
189                colors=color,
190                linestyles="solid",
191                alpha=0.3,
192            )
193            ax.add_collection(lc)
194
195    else:
196        segs = [pos[i, :, :2] for i in range(Np)]
197
198        lc = _LineCollection(
199            segs,
200            linewidths=0.5,
201            colors="0.6",
202            linestyles="solid",
203            alpha=0.5,
204        )
205        ax.add_collection(lc)
206
207    for x, y in p["source_positions"]:
208        ax.plot(x, y, **_SOURCE_MARKER_PROPS)
209
210    ax.autoscale()  # `ax.add_collection` won't do this automatically
211    ax.set_xlabel("$x$")
212    ax.set_ylabel("$y$")
213    ax.set_title(ltitle, loc="left")
214    ax.set_title(rtitle, loc="right")
215
216    fig.set_tight_layout(True)
217
218
219@_add_snippets
220def conc_scatter(
221    ds,
222    spc="apinene",
223    sdim="xy",
224    *,
225    cmap="gnuplot",
226    log_cnorm=False,
227    vmax=100,
228    vmin=None,
229    ax=None,
230):
231    """Plot species relative level in particle as a scatter (2- or 3-d).
232
233    Parameters
234    ----------
235    %(ds_chem_params)s
236    %(spc_params)s
237    %(sdim_params)s
238    %(log_cnorm_params)s
239    """
240    p = utils.load_p(ds)
241
242    dims = utils.dims_from_sdim(sdim)
243    if len(dims) == 3:
244        sdim = "xyz"
245        x = ds.x.values
246        y = ds.y.values
247        z = ds.z.values
248        subplot_kw = {"projection": "3d"}
249        coords = (x, y, z)
250        # plot_kw = dict(alpha=0.5, mew=0, ms=7)
251    elif len(dims) == 2:
252        x = ds[dims[0]].values
253        y = ds[dims[1]].values
254        subplot_kw = {}
255        coords = (x, y)
256        # plot_kw = dict(alpha=0.5, mfc="none", mew=0.8, ms=5)
257    else:
258        raise ValueError("invalid `sdim`")
259
260    conc = ds.f_r.sel(spc=spc).values
261    spc_display_name = str(ds.display_name.sel(spc=spc).values)
262
263    num = utils.check_fig_num(f"horizontal-end-positions-with-conc_{spc}_scatter_{sdim}")
264    fig, ax = utils.maybe_new_figure(num, ax=ax, subplot_kw=subplot_kw)
265
266    norm, _ = utils.maybe_log_cnorm(log_cnorm=log_cnorm, levels=None, vmin=vmin, vmax=vmax)
267
268    im = ax.scatter(
269        *coords,
270        c=conc,
271        s=7,
272        marker="o",
273        alpha=0.4,
274        linewidths=0,
275        cmap=cmap,
276        norm=norm,
277    )
278    cb = fig.colorbar(im, ax=ax, drawedges=False)
279    cb.set_label(f"{spc_display_name} relative conc. (%)")
280
281    for xs, ys in p["source_positions"]:
282        sp = dict(x=xs, y=ys, z=p["release_height"])
283        ax.plot(*tuple(sp[dim] for dim in dims), **_SOURCE_MARKER_PROPS)
284
285    ax.set_title(utils.s_t_info(p), loc="left")
286    ax.set(
287        xlabel=f"{ds[dims[0]].attrs['long_name']} [{ds[dims[0]].attrs['units']}]",
288        ylabel=f"{ds[dims[1]].attrs['long_name']} [{ds[dims[1]].attrs['units']}]",
289    )
290    if len(dims) == 3:
291        ax.set_zlabel(f"{ds[dims[2]].attrs['long_name']} [{ds[dims[2]].attrs['units']}]")
292    fig.set_tight_layout(True)
293
294
295@_add_snippets
296def conc_2d(
297    ds,
298    spc="apinene",
299    *,
300    bins=(20, 10),
301    plot_type="pcolor",
302    levels=30,  # for contourf
303    cmap="gnuplot",
304    log_cnorm=False,
305    vmax=100,
306    vmin=None,
307    ax=None,
308):
309    """Plot species 2-d (*xy*) binned average species relative level.
310
311    Parameters
312    ----------
313    %(ds_chem_params)s
314    %(spc_params)s
315    %(bins_params)s
316    plot_type : {'pcolor', 'contourf'}
317    levels : int
318        Only used for plot type `'contourf'`.
319    %(log_cnorm_params)s
320    """
321    p = utils.load_p(ds)
322    X = ds.x.values
323    Y = ds.y.values
324    conc = ds.f_r.sel(spc=spc).values
325    spc_display_name = str(ds.display_name.sel(spc=spc).values)
326
327    if "x" not in ds.dims:
328        # We were passed the particles dataset, need to do the binning
329        binned = utils.bin_values_xy(X, Y, conc, bins=bins)
330    else:
331        raise NotImplementedError("already binned?")
332
333    fig, ax = utils.maybe_new_figure(f"horizontal-end-positions-with-conc_{spc}_{plot_type}", ax=ax)
334
335    norm, locator = utils.maybe_log_cnorm(
336        log_cnorm=log_cnorm,
337        levels=levels if plot_type == "contourf" else None,
338        vmin=vmin,
339        vmax=vmax,
340    )
341
342    if plot_type == "pcolor":
343        im = ax.pcolormesh(binned.xe, binned.ye, binned.v, cmap=cmap, norm=norm)
344    elif plot_type == "contourf":
345        im = ax.contourf(
346            binned.x, binned.y, binned.v, levels, cmap=cmap, norm=norm, locator=locator
347        )
348    else:
349        raise ValueError("`plot_type` should be 'pcolor' or 'contourf'")
350
351    cb = fig.colorbar(im, ax=ax, drawedges=False)
352    cb.set_label(f"{spc_display_name} relative conc. (%)")
353
354    for x, y in p["source_positions"]:
355        ax.plot(x, y, **_SOURCE_MARKER_PROPS)
356
357    ax.autoscale(enable=True, axis="both", tight=True)
358    ax.set_title(utils.s_t_info(p), loc="left")
359    ax.set(
360        xlabel=f"{ds.x.attrs['long_name']} [{ds.x.attrs['units']}]",
361        ylabel=f"{ds.y.attrs['long_name']} [{ds.y.attrs['units']}]",
362    )
363    fig.set_tight_layout(True)
364
365
366@_add_snippets
367def conc_xline(
368    ds,
369    spc="apinene",
370    y=0.0,
371    *,
372    dy=1.0,
373    ax=None,  # TODO: select z as well?
374    legend=True,
375    label=None,
376    legend_title="Chemical species",
377):
378    r"""Plot species average relative level in the *x*-direction at a certain approximate y value.
379    `spc` can be ``'all'`` (all will fit in one plot since we are plotting lines).
380
381    Parameters
382    ----------
383    %(ds_chem_params)s
384    %(spc_params)s
385    y : float
386        *y* value of our "line".
387    dy : float
388        We will include Lagrangian particles within $[-\delta y / 2, \delta y / 2]$.
389    """
390    p = utils.load_p(ds)
391    X = ds.x.values
392    Y = ds.y.values
393
394    fig, ax = utils.maybe_new_figure(f"horizontal-end-positions-with-conc_{spc}_line", ax=ax)
395
396    spc_to_plot = ds.spc.values if spc == "all" else [spc]
397
398    # Define bins (edges)
399    xe = utils._auto_bins_1d(X, nx_max=100, std_mult=2.0, method="sqrt")
400    ye = np.r_[y - 0.5 * dy, y + 0.5 * dy]
401    # ze = np.r_[z - 0.5*dz, z + 0.5*dz]
402    bins = [xe, ye]
403
404    for spc in spc_to_plot:
405        spc_display_name = str(ds.display_name.sel(spc=spc).values)
406        conc = ds.f_r.sel(spc=spc).values
407        binned = utils.bin_values_xy(X, Y, conc, bins=bins)
408        ax.plot(binned.x, binned.v.squeeze(), label=spc_display_name if label is None else label)
409
410    for xs, ys in p["source_positions"]:
411        if abs(ys - y) <= 5:
412            ax.plot(xs, np.nanmin(binned.v), **_SOURCE_MARKER_PROPS)
413
414    if legend:
415        fig.legend(
416            ncol=2,
417            fontsize="small",
418            title=legend_title,
419        )
420    ax.set_title(utils.s_t_info(p), loc="left")
421    ax.set(
422        xlabel=f"{ds.x.attrs['long_name']} [{ds.x.attrs['units']}]",
423        ylabel="Relative concentration",
424    )
425    fig.set_tight_layout(True)
426
427
428# TODO: x-y (or u-v) hists for different height (z) bins
429
430
431@_add_snippets
432def ws_hist(
433    ds,
434    *,
435    bounds=None,
436):
437    """Plot histograms of particle wind speed components (one subplot for each).
438    For a single- or continuous-release run.
439
440    Parameters
441    ----------
442    %(ds_lpd_params)s
443    bounds : 2-tuple(float), optional
444        If provided, used as wind speed limits for all histograms.
445    """
446    p = utils.load_p(ds)
447    u_all = np.ravel(ds.u.values)  # or can do `.ravel()`
448    v_all = np.ravel(ds.v.values)
449    w_all = np.ravel(ds.w.values)
450
451    num = utils.check_fig_num("ws-hist-all")
452    fig, axs = plt.subplots(3, 1, num=num, sharex=True)
453
454    if bounds is None:
455        bins = 100
456    else:
457        bins = np.linspace(bounds[0], bounds[1], 100)
458
459    labels = ["$u$", "$v$", "$w$"]
460    for i, (ui, ax) in enumerate(zip([u_all, v_all, w_all], axs.flat)):
461        ax.hist(ui, bins)
462        ax.text(0.01, 0.98, labels[i], va="top", ha="left", fontsize=13, transform=ax.transAxes)
463
464    if bounds:
465        axs[0].set_xlim(bounds)
466
467    axs[0].set_title(utils.s_t_info(p), loc="left")
468    axs[0].set_title(utils.s_sample_size(p), loc="right")
469    axs[0].set_xlabel(f"[{ds.u.attrs['units']}]")
470    fig.set_tight_layout(True)
471
472
473@_add_snippets
474def final_pos_hist(
475    ds,
476    *,
477    bounds=None,
478):
479    """Plot histograms of final position components (*x*, *y*, and *z*).
480
481    Parameters
482    ----------
483    %(ds_lpd_params)s
484    bounds : 2-tuple(float), optional
485        If provided, used as position component limits for all histograms.
486    """
487    p = utils.load_p(ds)
488    if "t" in ds.dims:
489        xf = ds.x.isel(t=-1).values
490        yf = ds.y.isel(t=-1).values
491        zf = ds.z.isel(t=-1).values
492    else:
493        xf = ds.x.values
494        yf = ds.y.values
495        zf = ds.z.values
496
497    num = utils.check_fig_num("final-pos-hist")
498    fig, axs = plt.subplots(3, 1, num=num, sharex=True)
499
500    if bounds is None:
501        bins = 100
502    else:
503        bins = np.linspace(bounds[0], bounds[1], 100)
504
505    labels = ["$x$", "$y$", "$z$"]
506    for i, (xi, ax) in enumerate(zip([xf, yf, zf], axs.flat)):
507        ax.hist(xi, bins)
508        ax.text(0.01, 0.98, labels[i], va="top", ha="left", fontsize=13, transform=ax.transAxes)
509
510    if bounds:
511        axs[0].set_xlim(bounds)
512
513    axs[0].set_title(utils.s_t_info(p), loc="left")
514    axs[0].set_title(utils.s_sample_size(p, N_p_only=True), loc="right")
515    axs[-1].set_xlabel(f"[{ds.x.attrs['units']}]")
516    fig.set_tight_layout(True)
517
518
519@_add_snippets
520def final_pos_hist2d(
521    ds,
522    *,
523    sdim="xy",
524    bins=50,
525    plot_type="pcolor",
526    log_cnorm=False,
527    vmax=None,
528):
529    """Plot 2-d histogram of selected final position components.
530
531    Parameters
532    ----------
533    %(ds_lpd_params)s
534    %(sdim_params)s
535    %(bins_params)s
536    %(log_cnorm_params)s
537    """
538    p = utils.load_p(ds)
539
540    dims = utils.dims_from_sdim(sdim)
541    if "t" in ds.dims:
542        x = ds[dims[0]].isel(t=-1).values
543        y = ds[dims[1]].isel(t=-1).values
544    else:
545        x = ds[dims[0]].values
546        y = ds[dims[1]].values
547
548    num = utils.check_fig_num(f"final-pos-hist-{sdim}")
549    fig, ax = plt.subplots(num=num)
550
551    if bins == "auto":
552        bins = utils.auto_bins_xy([x, y])
553
554    if log_cnorm:
555        norm = mpl.colors.LogNorm(vmin=1.0, vmax=vmax)
556    else:
557        norm = mpl.colors.Normalize(vmin=1.0, vmax=vmax)
558
559    _, _, _, im = ax.hist2d(x, y, bins=bins, norm=norm)
560    # ^ returns h (nx, ny), xedges, yedges, image
561
562    cb = fig.colorbar(im, ax=ax)
563    cb.set_label("Particle count")
564
565    for x, y in p["source_positions"]:
566        ax.plot(x, y, "*", c="gold", ms=11, mec="0.35", mew=1.0)
567
568    ax.set_xlabel(f"${dims[0]}$ [{ds.x.attrs['units']}]")
569    ax.set_ylabel(f"${dims[1]}$ [{ds.x.attrs['units']}]")
570    ax.autoscale(enable=True, axis="both", tight=True)
571    ax.set_title(utils.s_t_info(p), loc="left")
572    ax.set_title(utils.s_sample_size(p, N_p_only=True), loc="right")
573    fig.set_tight_layout(True)