diff --git a/.gitignore b/.gitignore index f8042c90..e2a72483 100644 --- a/.gitignore +++ b/.gitignore @@ -115,3 +115,5 @@ venv.bak/ .DS_Store /beamphysics/_version.py + +/tests/artifacts diff --git a/beamphysics/__init__.py b/beamphysics/__init__.py index c569eeec..e15dbe36 100644 --- a/beamphysics/__init__.py +++ b/beamphysics/__init__.py @@ -25,6 +25,8 @@ "Wavefront": ".wavefront", "WavefrontK": ".wavefront", "pmd_init": ".writers", + "set_default_backend": ".plot_dispatch", + "get_default_backend": ".plot_dispatch", } @@ -35,6 +37,8 @@ "particle_paths", "pmd_init", "single_particle", + "set_default_backend", + "get_default_backend", "Wavefront", "WavefrontK", ] diff --git a/beamphysics/particles.py b/beamphysics/particles.py index 3f631206..bb101bbc 100644 --- a/beamphysics/particles.py +++ b/beamphysics/particles.py @@ -21,7 +21,7 @@ from .interfaces.lucretia import write_lucretia from .interfaces.opal import write_opal from .interfaces.simion import write_simion -from .plot import density_plot, marginal_plot, slice_plot, wakefield_plot +from .plot_dispatch import get_backend from .readers import particle_array, particle_paths from .species import charge_of, mass_of from .statistics import ( @@ -1381,6 +1381,7 @@ def plot( tex=True, nice=True, ellipse=False, + backend=None, **kwargs, ): """ @@ -1397,7 +1398,7 @@ def plot( Parameters ---------- - key1 : str, default = 't' + key1 : str, default = 'x' Key to bin on the x-axis key2 : str, default = None @@ -1423,25 +1424,28 @@ def plot( 2x2 sigma matrix return_figure : bool, default = False - If true, return a matplotlib.figure.Figure object + If true, return the figure/layout object - **kwargs - Any additional kwargs to send to the the plot in: plt.subplots(**kwargs) + backend : str, optional + Plot backend: ``'mpl'`` or ``'bokeh'``. + Defaults to the module-level setting (see ``set_default_backend``). + **kwargs + Additional keyword arguments passed to the backend plot function. Returns ------- - None or fig: matplotlib.figure.Figure - This only returns a figure object if return_figure=T, otherwise returns None - + None or figure object + Returns the figure only if ``return_figure=True``. """ + be = get_backend(backend) if not key2: - fig = density_plot( + fig = be.density_plot( self, key=key1, bins=bins, xlim=xlim, tex=tex, nice=nice, **kwargs ) else: - fig = marginal_plot( + fig = be.marginal_plot( self, key1=key1, key2=key2, @@ -1514,6 +1518,7 @@ def slice_plot( return_figure=False, xlim=None, ylim=None, + backend=None, **kwargs, ): """ @@ -1533,20 +1538,23 @@ def slice_plot( nice : bool, optional Scale to nice units. Default is True. return_figure : bool, optional - If True, return the matplotlib Figure. Default is False. + If True, return the figure/layout object. Default is False. xlim : tuple of float, optional Manual x-axis limits in raw units. ylim : tuple of float, optional Manual y-axis limits in raw units. + backend : str, optional + Plot backend: ``'mpl'`` or ``'bokeh'``. **kwargs - Additional keyword arguments passed to ``plt.subplots``. + Additional keyword arguments passed to the backend plot function. Returns ------- - None or matplotlib.figure.Figure - Returns a Figure only if ``return_figure=True``. + None or figure object + Returns the figure only if ``return_figure=True``. """ - fig = slice_plot( + be = get_backend(backend) + fig = be.slice_plot( self, *keys, n_slice=n_slice, @@ -1619,11 +1627,11 @@ def wakefield_plot( wake: WakefieldBase, key=None, nice=True, - ax=None, xlim=None, ylim=None, tex=True, bins=None, + backend=None, **kwargs, ): """ @@ -1647,9 +1655,6 @@ def wakefield_plot( nice : bool, default=True If True, applies unit-aware scaling using SI prefixes (e.g., mm, ns). - ax : matplotlib.axes.Axes, optional - An existing Axes to plot into. If None, a new figure and axes are created. - xlim : tuple of float, optional Limits to apply to the x-axis, in native units. @@ -1662,21 +1667,25 @@ def wakefield_plot( bins : int or str, optional Number of bins to use for the density histogram. - kwargs : dict - Additional keyword arguments passed to `plt.subplots()` if a new axis is created. + backend : str, optional + Plot backend: ``'mpl'`` or ``'bokeh'``. + + ax : matplotlib.axes.Axes, optional + Matplotlib-only: An existing Axes to plot into. (part of **kwargs) + + **kwargs + Additional keyword arguments passed to the backend plot function. Returns ------- - fig : matplotlib.figure.Figure - The matplotlib figure containing the plot. + figure object """ - - wakefield_plot( + be = get_backend(backend) + return be.wakefield_plot( self, wake, key=key, nice=nice, - ax=ax, xlim=xlim, ylim=ylim, tex=tex, diff --git a/beamphysics/plot.py b/beamphysics/plot.py index 379e6a4c..4bd9b5f2 100644 --- a/beamphysics/plot.py +++ b/beamphysics/plot.py @@ -1,7 +1,9 @@ """ """ +from __future__ import annotations + from copy import copy -from typing import Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING import matplotlib import matplotlib.pyplot as plt @@ -15,16 +17,25 @@ from mpl_toolkits.axes_grid1 import make_axes_locatable from .labels import mathlabel -from .statistics import slice_statistics, twiss_ellipse_points +from .plot_base import ( + Limit, + PlotPreparationError, + prepare_density_and_slice_plot, + prepare_density_plot, + prepare_marginal_plot, + prepare_slice_plot, + prepare_wakefield_plot, +) from .units import ( - c_light, nice_array, - nice_scale_prefix, pg_units, plottable_array, pmd_unit, ) +if TYPE_CHECKING: + from .particles import ParticleGroup + CMAP0 = copy(plt.get_cmap("viridis")) CMAP0.set_under("white") @@ -59,124 +70,64 @@ def slice_plot( Parameters ---------- - particle_group: ParticleGroup + particle_group : ParticleGroup The object to plot - keys: iterable of str + keys : iterable of str Keys to calculate the statistics, e.g. `sigma_x`. - n_slice: int, default = 40 + n_slice : int, default = 40 Number of slices - slice_key: str, default = None + slice_key : str, default = None The dimension to slice in. This is typically `t` or `z`. `delta_t`, etc. are also allowed. If None, `t` or `z` will automatically be determined. - ylim: tuple, default = None + ylim : tuple, default = None Manual setting of the y-axis limits. - tex: bool, default = True + tex : bool, default = True Use TEX for labels Returns ------- - fig: matplotlib.figure.Figure + fig : matplotlib.figure.Figure """ - - # Allow a single key - # if isinstance(keys, str): - # - # keys = (keys, ) - - if slice_key is None: - if particle_group.in_t_coordinates: - slice_key = "z" - else: - slice_key = "t" - - # Special case for delta_ - if slice_key.startswith("delta_"): - slice_key = slice_key[6:] - has_delta_prefix = True - else: - has_delta_prefix = False - - # Get all data - x_key = "mean_" + slice_key - slice_dat = particle_group.slice_statistics( - *keys, n_slice=n_slice, slice_key=slice_key + pdata = prepare_slice_plot( + particle_group, + *keys, + n_slice=n_slice, + slice_key=slice_key, + xlim=xlim, + ylim=ylim, + nice=nice, + tex=tex, ) - slice_dat["density"] = slice_dat["charge"] / slice_dat["ptp_" + slice_key] - y2_key = "density" - - # X-axis - x = slice_dat["mean_" + slice_key] - if has_delta_prefix: - x -= particle_group["mean_" + slice_key] - slice_key = "delta_" + slice_key # restore - - x, f1, p1, xmin, xmax = plottable_array(x, nice=nice, lim=xlim) - ux = p1 + str(particle_group.units(slice_key)) - # Y-axis - - # Units check - ulist = [particle_group.units(k).unitSymbol for k in keys] - uy = ulist[0] - if not all([u == uy for u in ulist]): - raise ValueError(f"Incompatible units: {ulist}") - - ymin = max([slice_dat[k].min() for k in keys]) - ymax = max([slice_dat[k].max() for k in keys]) - - _, f2, p2, ymin, ymax = plottable_array(np.array([ymin, ymax]), nice=nice, lim=ylim) - uy = p2 + uy - - # Form Figure fig, ax = plt.subplots(**kwargs) # Main curves - if len(keys) == 1: - color = "black" - else: - color = None - - for k in keys: - label = mathlabel(k, units=uy, tex=tex) - ax.plot(x, slice_dat[k] / f2, label=label, color=color) - if len(keys) > 1: + color = "black" if len(pdata.curves) == 1 else None + for curve in pdata.curves: + ax.plot(pdata.x, curve.values, label=curve.label, color=color) + if len(pdata.curves) > 1: ax.legend() - # Density on r.h.s - y2, _, prey2, _, _ = plottable_array(slice_dat[y2_key], nice=nice, lim=None) + ax.set_xlabel(pdata.x_label) + ax.set_ylabel(pdata.y_label) - # Convert to Amps if possible - y2_units = f"C/{particle_group.units(x_key)}" - if y2_units == "C/s": - y2_units = "A" - y2_units = prey2 + y2_units - - # Labels - labelx = mathlabel(slice_key, units=ux, tex=tex) - labely = mathlabel(*keys, units=uy, tex=tex) - labely2 = mathlabel(y2_key, units=y2_units, tex=tex) - - ax.set_xlabel(labelx) - ax.set_ylabel(labely) - - # rhs plot + # Density on r.h.s ax2 = ax.twinx() - ax2.set_ylabel(labely2) - ax2.fill_between(x, 0, y2, color="black", alpha=0.2) + ax2.set_ylabel(pdata.density_label) + ax2.fill_between(pdata.x, 0, pdata.density_values, color="black", alpha=0.2) ax2.set_ylim(0, None) - # Actual plot limits, considering scaling - if xlim: - ax.set_xlim(xmin / f1, xmax / f1) - if ylim: - ax.set_ylim(ymin / f2, ymax / f2) + if pdata.xlim: + ax.set_xlim(*pdata.xlim) + if pdata.ylim: + ax.set_ylim(*pdata.ylim) return fig @@ -184,12 +135,12 @@ def slice_plot( def density_plot( particle_group, key: str = "x", - bins: Optional[Union[int, str]] = None, + bins: int | str | None = None, *, - xlim: Optional[Tuple[float, float]] = None, + xlim: Limit | None = None, tex: bool = True, nice: bool = True, - ax: Optional[Axes] = None, + # ax: Axes | None = None, # <-- handled in kwargs to maintain protocol color="grey", alpha=1, **kwargs, @@ -231,214 +182,142 @@ def density_plot( fig : matplotlib.figure.Figure The created or parent figure. """ - if bins is None: - n = len(particle_group) - bins = int(n / 100) - - x, f1, p1, xmin, xmax = plottable_array(particle_group[key], nice=nice, lim=xlim) - w = particle_group["weight"] - u1 = particle_group.units(key).unitSymbol - ux = p1 + u1 - - labelx = mathlabel(key, units=ux, tex=tex) + ax: Axes | None = kwargs.pop("ax", None) + pdata = prepare_density_plot( + particle_group, key=key, bins=bins, xlim=xlim, nice=nice, tex=tex + ) if ax is None: fig, ax = plt.subplots(**kwargs) else: fig = ax.get_figure() - hist, bin_edges = np.histogram(x, bins=bins, weights=w) - hist_x = bin_edges[:-1] + np.diff(bin_edges) / 2 - hist_width = np.diff(bin_edges) - hist_y, hist_f, hist_prefix, _hist_xmin, _hist_xmax = plottable_array( - hist / hist_width, nice=nice + ax.bar( + pdata.hist_centers, + pdata.hist_values, + pdata.hist_width, + color=color, + alpha=alpha, ) - - ax.bar(hist_x, hist_y, hist_width, color=color, alpha=alpha) - - if u1 == "s": - _, hist_prefix = nice_scale_prefix(hist_f / f1) - ax.set_ylabel(f" density ({hist_prefix}A)") - else: - ax.set_ylabel(f"{hist_prefix}C/{ux}") + ax.set_ylabel(pdata.y_label) if hasattr(ax, "get_shared_x_axes") and ax.get_shared_x_axes().joined( ax, ax.figure.axes[0] ): - ax.figure.axes[0].set_xlabel(labelx) + ax.figure.axes[0].set_xlabel(pdata.x_label) else: - ax.set_xlabel(labelx) + ax.set_xlabel(pdata.x_label) - if xlim: - ax.set_xlim(xmin / f1, xmax / f1) + if pdata.xlim: + ax.set_xlim(*pdata.xlim) return fig def marginal_plot( - particle_group, - key1="t", - key2="p", - bins=None, + particle_group: ParticleGroup, + key1: str = "t", + key2: str = "p", + bins: int | None = None, *, - xlim=None, - ylim=None, - tex=True, - nice=True, - ellipse=False, + xlim: Limit | None = None, + ylim: Limit | None = None, + tex: bool = True, + nice: bool = True, + ellipse: bool = False, **kwargs, ): """ - Density plot and projections - - Example: - - marginal_plot(P, 't', 'energy', bins=200) - + Density plot and projections with matplotlib. Parameters ---------- - particle_group: ParticleGroup + particle_group : ParticleGroup The object to plot - - key1: str, default = 't' + key1 : str, default = 't' Key to bin on the x-axis - - key2: str, default = 'p' + key2 : str, default = 'p' Key to bin on the y-axis - - bins: int, default = None - Number of bins. If None, this will use a heuristic: bins = sqrt(n_particle/4) - - xlim: tuple, default = None + bins : int, default = None + Number of bins. If None, this will use a heuristic: + `bins = sqrt(n_particle/4)` + xlim : tuple, default = None Manual setting of the x-axis limits. - - ylim: tuple, default = None + ylim : tuple, default = None Manual setting of the y-axis limits. - - tex: bool, default = True + tex : bool, default = True Use TEX for labels + nice : bool, default = True - nice: bool, default = True - - ellipse: bool, default = True + ellipse : bool, default = True If True, plot an ellipse representing the 2x2 sigma matrix + **kwargs : + Passed to `plt.figure`. Returns ------- - fig: matplotlib.figure.Figure + matplotlib.figure.Figure + Examples + -------- + >>> P = ParticleGroup("particles.h5") + >>> marginal_plot(P, 't', 'energy', bins=200) """ - if not bins: - n = len(particle_group) - bins = int(np.sqrt(n / 4)) - - # Scale to nice units and get the factor, unit prefix - x = particle_group[key1] - y = particle_group[key2] - - if len(x) == 1: - bins = 100 - - if xlim is None: - (x0,) = x - if np.isclose(x0, 0.0): - xlim = (-1, 1) - else: - xlim = tuple(sorted((0.9 * x0, 1.1 * x0))) - if ylim is None: - (y0,) = y - if np.isclose(y0, 0.0): - ylim = (-1, 1) - else: - ylim = tuple(sorted((0.9 * y0, 1.1 * y0))) - - # Form nice arrays - x, f1, p1, xmin, xmax = plottable_array(x, nice=nice, lim=xlim) - y, f2, p2, ymin, ymax = plottable_array(y, nice=nice, lim=ylim) - - w = particle_group["weight"] - - u1 = particle_group.units(key1).unitSymbol - u2 = particle_group.units(key2).unitSymbol - ux = p1 + u1 - uy = p2 + u2 - - # Handle labels. - labelx = mathlabel(key1, units=ux, tex=tex) - labely = mathlabel(key2, units=uy, tex=tex) fig = plt.figure(**kwargs) - if np.all(np.isnan(x)): - fig.text(0.5, 0.5, f"{key1} is all NaN", ha="center", va="center") - return fig - if np.all(np.isnan(y)): - fig.text(0.5, 0.5, f"{key2} is all NaN", ha="center", va="center") + try: + pdata = prepare_marginal_plot( + particle_group, + key1=key1, + key2=key2, + bins=bins, + xlim=xlim, + ylim=ylim, + nice=nice, + ellipse=ellipse, + ) + except PlotPreparationError as ex: + fig.text(0.5, 0.5, str(ex), ha="center", va="center") return fig gs = GridSpec(4, 4) - ax_joint = fig.add_subplot(gs[1:4, 0:3]) ax_marg_x = fig.add_subplot(gs[0, 0:3]) ax_marg_y = fig.add_subplot(gs[1:4, 3]) - # ax_info = fig.add_subplot(gs[0, 3:4]) - # ax_info.table(cellText=['a']) - # Main plot - # Proper weighting - if len(x) == 1: - ax_joint.scatter(x, y) + if len(pdata.x.data) == 1: + ax_joint.scatter(pdata.x.data, pdata.y.data) else: ax_joint.hexbin( - x, - y, - C=w, + pdata.x.data, + pdata.y.data, + C=pdata.weights, reduce_C_function=np.sum, - gridsize=bins, + gridsize=pdata.bins, cmap=CMAP0, vmin=1e-20, ) - if ellipse: - sigma_mat2 = particle_group.cov(key1, key2) - x_ellipse, y_ellipse = twiss_ellipse_points(sigma_mat2) - x_ellipse += particle_group.avg(key1) - y_ellipse += particle_group.avg(key2) - ax_joint.plot(x_ellipse / f1, y_ellipse / f2, color="red") - - # Manual histogramming version - # H, xedges, yedges = np.histogram2d(x, y, weights=w, bins=bins) - # extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]] - # ax_joint.imshow(H.T, cmap=cmap, vmin=1e-16, origin='lower', extent=extent, aspect='auto') + if pdata.ellipse_x is not None and pdata.ellipse_y is not None: + ax_joint.plot(pdata.ellipse_x, pdata.ellipse_y, color="red") # Top histogram - # Old method: - # dx = x.ptp()/bins - # ax_marg_x.hist(x, weights=w/dx/f1, bins=bins, color='gray') - hist, bin_edges = np.histogram(x, bins=bins, weights=w) - hist_x = bin_edges[:-1] + np.diff(bin_edges) / 2 - hist_width = np.diff(bin_edges) - hist_y, hist_f, hist_prefix = nice_array(hist / hist_width) - ax_marg_x.bar(hist_x, hist_y, hist_width, color="gray") - # Special label for C/s = A - if u1 == "s": - _, hist_prefix = nice_scale_prefix(hist_f / f1) - ax_marg_x.set_ylabel(f"{hist_prefix}A") - else: - ax_marg_x.set_ylabel(f"{hist_prefix}" + mathlabel(f"C/{ux}")) # Always use tex - - # Side histogram - # Old method: - # dy = y.ptp()/bins - # ax_marg_y.hist(y, orientation="horizontal", weights=w/dy, bins=bins, color='gray') - hist, bin_edges = np.histogram(y, bins=bins, weights=w) - hist_x = bin_edges[:-1] + np.diff(bin_edges) / 2 - hist_width = np.diff(bin_edges) - hist_y, hist_f, hist_prefix = nice_array(hist / hist_width) - ax_marg_y.barh(hist_x, hist_y, hist_width, color="gray") - ax_marg_y.set_xlabel(f"{hist_prefix}" + mathlabel(f"C/{uy}")) # Always use tex + ax_marg_x.bar( + pdata.x.hist_centers, pdata.x.hist_values, pdata.x.hist_width, color="gray" + ) + + # Right histogram + ax_marg_y.barh( + pdata.y.hist_centers, pdata.y.hist_values, pdata.y.hist_width, color="gray" + ) + + labelx = mathlabel(key1, units=pdata.x.full_unit, tex=tex) + labely = mathlabel(key2, units=pdata.y.full_unit, tex=tex) + + ax_marg_x.set_ylabel(pdata.x.axis_label) + ax_marg_y.set_xlabel(pdata.y.axis_label) # Turn off tick labels on marginals plt.setp(ax_marg_x.get_xticklabels(), visible=False) @@ -449,13 +328,13 @@ def marginal_plot( ax_joint.set_ylabel(labely) # Actual plot limits, considering scaling - if xlim: - ax_joint.set_xlim(xmin / f1, xmax / f1) - ax_marg_x.set_xlim(xmin / f1, xmax / f1) + if xlim is not None: + ax_joint.set_xlim(pdata.x.lim) + ax_marg_x.set_xlim(pdata.x.lim) - if ylim: - ax_joint.set_ylim(ymin / f2, ymax / f2) - ax_marg_y.set_ylim(ymin / f2, ymax / f2) + if ylim is not None: + ax_joint.set_ylim(pdata.y.lim) + ax_marg_y.set_ylim(pdata.y.lim) return fig @@ -464,81 +343,72 @@ def density_and_slice_plot( particle_group, key1="t", key2="p", - stat_keys=["norm_emit_x", "norm_emit_y"], + stat_keys=None, bins=100, n_slice=30, tex=True, + **kwargs, ): """ - Density plot and projections - - Example: + 2D density plot with overlaid slice statistics. - marginal_plot(P, 't', 'energy', bins=200) + Parameters + ---------- + particle_group : ParticleGroup + The object to plot. + key1 : str, default = 't' + Key for x-axis (also used as slice key). + key2 : str, default = 'p' + Key for y-axis (density). + stat_keys : list of str, optional + Slice statistics to overlay. Default: ``['norm_emit_x', 'norm_emit_y']``. + bins : int, default = 100 + Number of bins for the 2D histogram. + n_slice : int, default = 30 + Number of slices. + tex : bool, default = True + Use TeX for labels. + Returns + ------- + matplotlib.figure.Figure """ - - # Scale to nice units and get the factor, unit prefix - x, f1, p1, xmin, xmax = plottable_array(particle_group[key1]) - y, f2, p2, ymin, ymax = plottable_array(particle_group[key2]) - w = particle_group["weight"] - - u1 = particle_group.units(key1).unitSymbol - u2 = particle_group.units(key2).unitSymbol - ux = p1 + u1 - uy = p2 + u2 - - labelx = mathlabel(key1, units=ux, tex=tex) - labely = mathlabel(key2, units=uy, tex=tex) - - fig, ax = plt.subplots() - - ax.set_xlabel(labelx) - ax.set_ylabel(labely) - - # Proper weighting - # ax_joint.hexbin(x, y, C=w, reduce_C_function=np.sum, gridsize=bins, cmap=cmap, vmin=1e-15) - - # Manual histogramming version - H, xedges, yedges = np.histogram2d(x, y, weights=w, bins=bins) - extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]] - ax.imshow(H.T, cmap=CMAP0, vmin=1e-16, origin="lower", extent=extent, aspect="auto") - - # Slice data - slice_dat = slice_statistics( + pdata = prepare_density_and_slice_plot( particle_group, + key1=key1, + key2=key2, + stat_keys=stat_keys, + bins=bins, n_slice=n_slice, - slice_key=key1, - keys=stat_keys + ["ptp_" + key1, "mean_" + key1, "charge"], + tex=tex, ) - slice_dat["density"] = slice_dat["charge"] / slice_dat["ptp_" + key1] - - # - ax2 = ax.twinx() - # ax2.set_ylim(0, 1e-6) - x2 = slice_dat["mean_" + key1] / f1 - ulist = [particle_group.units(k).unitSymbol for k in stat_keys] + fig, ax = plt.subplots(**kwargs) - max2 = max([np.ptp(slice_dat[k]) for k in stat_keys]) + ax.set_xlabel(pdata.x_label) + ax.set_ylabel(pdata.y_label) - f3, p3 = nice_scale_prefix(max2) + ax.imshow( + pdata.hist2d.T, + cmap=CMAP0, + vmin=1e-16, + origin="lower", + extent=pdata.extent, + aspect="auto", + ) - u2 = ulist[0] - assert all([u == u2 for u in ulist]) - u2 = p3 + u2 - labely2 = mathlabel(*stat_keys, units=u2, tex=tex) - for k in stat_keys: - label = mathlabel(k, units=u2, tex=tex) - ax2.plot(x2, slice_dat[k] / f3, label=label) + # Slice statistics on secondary y-axis + ax2 = ax.twinx() + for curve in pdata.slice_curves: + ax2.plot(pdata.slice_x, curve.values, label=curve.label) ax2.legend() - ax2.set_ylabel(labely2) + ax2.set_ylabel(pdata.slice_y_label) ax2.set_ylim(bottom=0) - # Add density - y2 = slice_dat["density"] - y2 = y2 * max2 / y2.max() / f3 / 2 - ax2.fill_between(x2, 0, y2, color="black", alpha=0.1) + # Density overlay + ax2.fill_between(pdata.slice_x, 0, pdata.slice_density, color="black", alpha=0.1) + + return fig # ------------------------------------- @@ -1000,28 +870,30 @@ def plot_fieldmesh_rectangular_2d( def plot_1d_density( - x: Union[str, np.ndarray], - y: Union[str, np.ndarray], + x: str | np.ndarray, + y: str | np.ndarray, x_name: str = "", - y_name: Optional[str] = None, - x_units: Optional[str] = None, - y_units: Optional[str] = None, - figsize: Tuple[float, float] = (6, 4), + y_name: str | None = None, + x_units: str | None = None, + y_units: str | None = None, + figsize: Limit = (6, 4), log_scale_y: bool = False, show_cdf: bool = False, cdf_label: str = "CDF", - cdf_style: Optional[Dict[str, Union[str, float]]] = None, + cdf_style: dict[str, str | float] | None = None, kind: str = "bar", - plot_style: Optional[Dict[str, Union[str, float]]] = None, - xlim: Optional[Tuple[float, float]] = None, - ylim: Optional[Tuple[float, float]] = (0, None), - ax: Optional[plt.Axes] = None, + plot_style: dict[str, str | float] | None = None, + xlim: Limit | None = None, + ylim: Limit | None = (0, None), + # ax: Axes | None = None, # <-- handled in kwargs to maintain protocol nice: bool = True, auto_label: bool = False, tex: bool = True, - data: Optional[Dict[str, np.ndarray]] = None, + data: dict[str, np.ndarray] | None = None, + return_figure: bool = False, return_axes: bool = False, -) -> Optional[Tuple[plt.Figure, Dict[str, plt.Axes]]]: + **kwargs, +) -> tuple[plt.Figure, dict[str, plt.Axes]] | plt.Figure | None: """ Plot a 1D density distribution with optional cumulative distribution function (CDF). @@ -1104,7 +976,7 @@ def plot_1d_density( >>> plot_1d_density("t", "norm_emit_x", data=data, auto_label=True) # Will automatically use TeX labels and proper units """ - # Handle data dict indexing (matplotlib pattern) + ax: Axes | None = kwargs.pop("ax", None) # Handle data dict indexing (matplotlib pattern) x_key = None y_key = None @@ -1290,42 +1162,43 @@ def plot_1d_density( axes["cdf"] = ax_cdf - # Return axes if requested if return_axes: return fig, axes + if return_figure: + return fig def plot_2d_density_with_marginals( data: np.ndarray, - dx: Optional[float] = 1, - dy: Optional[float] = 1, - xmin: Optional[float] = None, - ymin: Optional[float] = None, + dx: float | None = 1, + dy: float | None = 1, + xmin: float | None = None, + ymin: float | None = None, x_name: str = "", y_name: str = "", z_name: str = "", - x_units: Optional[str] = None, - y_units: Optional[str] = None, - z_units: Optional[str] = None, + x_units: str | None = None, + y_units: str | None = None, + z_units: str | None = None, cmap: str = "inferno", - figsize: Tuple[float, float] = (5, 5), + figsize: Limit = (5, 5), log_scale_z: bool = False, log_scale_marginals: bool = False, - marginal_titles: Tuple[Optional[str], Optional[str]] = (None, None), - highlight_regions: Optional[ - List[Dict[str, Union[float, Tuple[float, float]]]] - ] = None, - marginal_style: Optional[Dict[str, Union[str, float]]] = None, + marginal_titles: tuple[str | None, str | None] = (None, None), + highlight_regions: None | (list[dict[str, float | Limit]]) = None, + marginal_style: dict[str, str | float] | None = None, show_stats: bool = False, show_colorbar: bool = True, - xlim: Tuple[float, float] = None, - ylim: Tuple[float, float] = None, - vmin: Optional[float] = None, - vcenter: Optional[float] = None, - vmax: Optional[float] = None, - aspect: Optional[str] = "auto", + xlim: Limit | None = None, + ylim: Limit | None = None, + vmin: float | None = None, + vcenter: float | None = None, + vmax: float | None = None, + aspect: str | None = "auto", + return_figure: bool = False, return_axes: bool = False, -) -> Optional[Tuple[plt.Figure, Dict[str, plt.Axes]]]: + **kwargs, +) -> tuple[plt.Figure, dict[str, plt.Axes]] | plt.Figure | None: """ Basic plot for a 2D density map with marginal histograms. @@ -1471,21 +1344,22 @@ def plot_2d_density_with_marginals( ax_main.set_ylim(ylim) ax_right.set_ylim(ylim) - # Return axes if requested if return_axes: return fig, axes + if return_figure: + return fig def wakefield_plot( particle_group, wake, - key: Optional[str] = None, + key: str | None = None, nice: bool = True, - ax: Optional[Axes] = None, - xlim: Optional[Tuple[float, float]] = None, - ylim: Optional[Tuple[float, float]] = None, + # ax: Axes | None = None, # <-- handled in kwargs to maintain protocol + xlim: Limit | None = None, + ylim: Limit | None = None, tex: bool = True, - bins: Optional[Union[int, str]] = None, + bins: int | str | None = None, **kwargs, ) -> Figure: """ @@ -1536,48 +1410,42 @@ def wakefield_plot( fig : matplotlib.figure.Figure The matplotlib figure containing the plot. """ - if key is None: - if particle_group.in_t_coordinates: - key = "delta_z/c" - else: - key = "delta_t" + ax: Axes | None = kwargs.pop("ax", None) + pdata = prepare_wakefield_plot( + particle_group, + wake, + key=key, + nice=nice, + tex=tex, + xlim=xlim, + ylim=ylim, + bins=bins, + ) if ax is None: fig, ax = plt.subplots(**kwargs) else: fig = ax.get_figure() - # Plot density on twin axis + # Density on twin axis ax2 = ax.twinx() - density_plot(particle_group, key=key, ax=ax2, nice=nice, alpha=0.5, bins=bins) - - # Wake kicks - x_raw = particle_group[key] - - if particle_group.in_t_coordinates: - z = np.asarray(particle_group.z) - else: - z = -c_light * np.asarray(particle_group.t) - kicks = wake.particle_kicks(z=z, weight=particle_group.weight) - - x, f1, p1, xmin, xmax = plottable_array(x_raw, nice=nice, lim=xlim) - y, f2, p2, ymin, ymax = plottable_array(kicks, nice=nice, lim=ylim) - - ax.scatter(x, y, marker=".", color="black", s=0.5) - - # Labels - ux = p1 + particle_group.units(key).unitSymbol - labelx = mathlabel(key, units=ux, tex=tex) - ax.set_xlabel(labelx) + ax2.bar( + pdata.density.hist_centers, + pdata.density.hist_values, + pdata.density.hist_width, + color="grey", + alpha=0.5, + ) + ax2.set_ylabel(pdata.density.y_label) - uy = p2 + "eV/m" - labely = mathlabel("W_z", units=uy, tex=tex) - ax.set_ylabel(labely) + # Wake kicks scatter + ax.scatter(pdata.scatter_x, pdata.scatter_y, marker=".", color="black", s=0.5) + ax.set_xlabel(pdata.x_label) + ax.set_ylabel(pdata.y_label) - # Limits - if xlim: - ax.set_xlim(xmin / f1, xmax / f1) - if ylim: - ax.set_ylim(ymin / f2, ymax / f2) + if pdata.xlim: + ax.set_xlim(*pdata.xlim) + if pdata.ylim: + ax.set_ylim(*pdata.ylim) return fig diff --git a/beamphysics/plot_base.py b/beamphysics/plot_base.py new file mode 100644 index 00000000..b31a669b --- /dev/null +++ b/beamphysics/plot_base.py @@ -0,0 +1,537 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, cast + +import numpy as np + +from .labels import mathlabel +from .statistics import twiss_ellipse_points +from .units import c_light, nice_array, nice_scale_prefix, plottable_array + +if TYPE_CHECKING: + from .particles import ParticleGroup + +Limit = tuple[float | None, float | None] + + +class PlotPreparationError(Exception): ... + + +class NanDataError(PlotPreparationError): ... + + +@dataclass +class MarginalAxisData: + """ + Single axis of a marginal plot. + """ + + key: str + data: np.ndarray + lim: tuple[float, float] + unit_factor: float + unit_prefix: str + unit_symbol: str + + # Histogram/Profile data + hist_centers: np.ndarray + hist_values: np.ndarray + hist_width: np.ndarray + hist_unit_factor: float + hist_label_prefix: str + + @property + def full_unit(self) -> str: + """Returns the combined scale prefix and root symbol, e.g. 'mm'""" + return f"{self.unit_prefix}{self.unit_symbol}" + + @property + def axis_label(self) -> str: + if self.unit_symbol == "s": + _, hist_prefix = nice_scale_prefix(self.hist_unit_factor / self.unit_factor) + return f"{hist_prefix}A" + return self.hist_label_prefix + mathlabel(f"C/{self.full_unit}") + + +@dataclass +class MarginalPlotData: + """ + Marginal plot data. + """ + + x: MarginalAxisData + y: MarginalAxisData + weights: np.ndarray + bins: int + ellipse_x: np.ndarray | None = None # Scaled ellipse coords + ellipse_y: np.ndarray | None = None # Scaled ellipse coords + + +@dataclass +class DensityPlotData: + """Prepared data for a 1D density histogram.""" + + key: str + hist_centers: np.ndarray + hist_values: np.ndarray + hist_width: np.ndarray + x_label: str + y_label: str + xlim: tuple[float, float] | None + x_factor: float + + +def prepare_density_plot( + particle_group: ParticleGroup, + key: str = "x", + bins: int | str | None = None, + *, + xlim: tuple[float, float] | None = None, + nice: bool = True, + tex: bool = True, +) -> DensityPlotData: + """ + Prepare data for a 1D density plot. + + Computes the histogram, scales units, and formats labels. + """ + if bins is None: + n = len(particle_group) + bins = max(1, int(n / 100)) + + x, f1, p1, xmin, xmax = plottable_array(particle_group[key], nice=nice, lim=xlim) + w = particle_group["weight"] + u1 = particle_group.units(key).unitSymbol + ux = p1 + u1 + + x_label = mathlabel(key, units=ux, tex=tex) + + hist, bin_edges = np.histogram(x, bins=bins, weights=w) + hist_x = bin_edges[:-1] + np.diff(bin_edges) / 2 + hist_width = np.diff(bin_edges) + hist_y, hist_f, hist_prefix, _hist_xmin, _hist_xmax = plottable_array( + hist / hist_width, nice=nice + ) + + if u1 == "s": + _, hist_prefix = nice_scale_prefix(hist_f / f1) + y_label = f"density ({hist_prefix}A)" + else: + y_label = f"{hist_prefix}C/{ux}" + + scaled_xlim = (xmin / f1, xmax / f1) if xlim else None + + return DensityPlotData( + key=key, + hist_centers=hist_x, + hist_values=hist_y, + hist_width=hist_width, + x_label=x_label, + y_label=y_label, + xlim=scaled_xlim, + x_factor=f1, + ) + + +@dataclass +class SliceCurve: + """A single curve in a slice plot.""" + + key: str + label: str + values: np.ndarray # scaled + + +@dataclass +class SlicePlotData: + """Prepared data for a slice statistics plot.""" + + x: np.ndarray # scaled slice positions + x_label: str + curves: list[SliceCurve] + y_label: str + density_values: np.ndarray # scaled density for rhs overlay + density_label: str + xlim: tuple[float, float] | None # scaled + ylim: tuple[float, float] | None # scaled + x_factor: float + y_factor: float + + +def prepare_slice_plot( + particle_group: ParticleGroup, + *keys: str, + n_slice: int = 40, + slice_key: str | None = None, + xlim: tuple[float, float] | None = None, + ylim: tuple[float, float] | None = None, + nice: bool = True, + tex: bool = True, +) -> SlicePlotData: + """ + Prepare data for a slice statistics plot. + + Computes slice statistics, scales units, and formats labels. + """ + if slice_key is None: + if particle_group.in_t_coordinates: + slice_key = "z" + else: + slice_key = "t" + + # Special case for delta_ + if slice_key.startswith("delta_"): + slice_key = slice_key[6:] + has_delta_prefix = True + else: + has_delta_prefix = False + + # Get all data + x_key = "mean_" + slice_key + slice_dat = particle_group.slice_statistics( + *keys, n_slice=n_slice, slice_key=slice_key + ) + slice_dat["density"] = slice_dat["charge"] / slice_dat["ptp_" + slice_key] + + # X-axis + x = slice_dat["mean_" + slice_key] + if has_delta_prefix: + x -= particle_group["mean_" + slice_key] + slice_key = "delta_" + slice_key # restore + + x, f1, p1, xmin, xmax = plottable_array(x, nice=nice, lim=xlim) + ux = p1 + str(particle_group.units(slice_key)) + + # Y-axis - units check + ulist = [particle_group.units(k).unitSymbol for k in keys] + uy = ulist[0] + if not all(u == uy for u in ulist): + raise ValueError(f"Incompatible units: {ulist}") + + ymin = max(slice_dat[k].min() for k in keys) + ymax = max(slice_dat[k].max() for k in keys) + + _, f2, p2, ymin, ymax = plottable_array(np.array([ymin, ymax]), nice=nice, lim=ylim) + uy = p2 + uy + + # Curves + curves = [] + for k in keys: + label = mathlabel(k, units=uy, tex=tex) + curves.append(SliceCurve(key=k, label=label, values=slice_dat[k] / f2)) + + # Density on r.h.s + y2, _, prey2, _, _ = plottable_array(slice_dat["density"], nice=nice, lim=None) + + y2_units = f"C/{particle_group.units(x_key)}" + if y2_units == "C/s": + y2_units = "A" + y2_units = prey2 + y2_units + + x_label = mathlabel(slice_key, units=ux, tex=tex) + y_label = mathlabel(*keys, units=uy, tex=tex) + density_label = mathlabel("density", units=y2_units, tex=tex) + + scaled_xlim = (xmin / f1, xmax / f1) if xlim else None + scaled_ylim = (ymin / f2, ymax / f2) if ylim else None + + return SlicePlotData( + x=x, + x_label=x_label, + curves=curves, + y_label=y_label, + density_values=y2, + density_label=density_label, + xlim=scaled_xlim, + ylim=scaled_ylim, + x_factor=f1, + y_factor=f2, + ) + + +@dataclass +class WakefieldPlotData: + """Prepared data for a wakefield plot.""" + + scatter_x: np.ndarray # scaled x positions + scatter_y: np.ndarray # scaled kick values + x_label: str + y_label: str + density: DensityPlotData # for overlay + xlim: tuple[float, float] | None + ylim: tuple[float, float] | None + + +def prepare_wakefield_plot( + particle_group: ParticleGroup, + wake, + key: str | None = None, + nice: bool = True, + tex: bool = True, + xlim: tuple[float, float] | None = None, + ylim: tuple[float, float] | None = None, + bins: int | str | None = None, +) -> WakefieldPlotData: + """ + Prepare data for a wakefield plot. + + Computes wake kicks, density histogram, and formats labels. + """ + if key is None: + if particle_group.in_t_coordinates: + key = "delta_z/c" + else: + key = "delta_t" + + # Density data for overlay + density_data = prepare_density_plot( + particle_group, key=key, bins=bins, xlim=xlim, nice=nice, tex=tex + ) + + # Wake kicks + if particle_group.in_t_coordinates: + z = np.asarray(particle_group.z) + else: + z = -c_light * np.asarray(particle_group.t) + kicks = wake.particle_kicks(z=z, weight=particle_group.weight) + + x_raw = particle_group[key] + x, f1, p1, xmin, xmax = plottable_array(x_raw, nice=nice, lim=xlim) + y, f2, p2, ymin, ymax = plottable_array(kicks, nice=nice, lim=ylim) + + ux = p1 + particle_group.units(key).unitSymbol + x_label = mathlabel(key, units=ux, tex=tex) + + uy = p2 + "eV/m" + y_label = mathlabel("W_z", units=uy, tex=tex) + + scaled_xlim = (xmin / f1, xmax / f1) if xlim else None + scaled_ylim = (ymin / f2, ymax / f2) if ylim else None + + return WakefieldPlotData( + scatter_x=x, + scatter_y=y, + x_label=x_label, + y_label=y_label, + density=density_data, + xlim=scaled_xlim, + ylim=scaled_ylim, + ) + + +@dataclass +class DensityAndSlicePlotData: + """Prepared data for a combined 2D density + slice statistics plot.""" + + # 2D histogram + hist2d: np.ndarray # shape (nbins_x, nbins_y), transposed for imshow + extent: list[float] # [xmin, xmax, ymin, ymax] in scaled units + x_label: str + y_label: str + + # Slice statistics curves (on secondary y-axis) + slice_x: np.ndarray # scaled slice positions + slice_curves: list[SliceCurve] + slice_y_label: str + slice_y_factor: float + + # Slice density overlay (normalized to fit on the stat axis) + slice_density: np.ndarray # scaled to overlay on stat axis + + +def prepare_density_and_slice_plot( + particle_group: ParticleGroup, + key1: str = "t", + key2: str = "p", + stat_keys: list[str] | None = None, + bins: int = 100, + n_slice: int = 30, + tex: bool = True, +) -> DensityAndSlicePlotData: + """ + Prepare data for a combined 2D density + slice statistics plot. + """ + if stat_keys is None: + stat_keys = ["norm_emit_x", "norm_emit_y"] + + from .statistics import slice_statistics as _slice_statistics + + # Scale to nice units + x, f1, p1, xmin, xmax = plottable_array(particle_group[key1]) + y, f2, p2, ymin, ymax = plottable_array(particle_group[key2]) + w = particle_group["weight"] + + u1 = particle_group.units(key1).unitSymbol + u2 = particle_group.units(key2).unitSymbol + ux = p1 + u1 + uy = p2 + u2 + + x_label = mathlabel(key1, units=ux, tex=tex) + y_label = mathlabel(key2, units=uy, tex=tex) + + # 2D histogram + H, xedges, yedges = np.histogram2d(x, y, weights=w, bins=bins) + extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]] + + # Slice data + slice_dat = _slice_statistics( + particle_group, + n_slice=n_slice, + slice_key=key1, + keys=stat_keys + ["ptp_" + key1, "mean_" + key1, "charge"], + ) + slice_dat["density"] = slice_dat["charge"] / slice_dat["ptp_" + key1] + + slice_x = slice_dat["mean_" + key1] / f1 + + # Stat curves scaling + ulist = [particle_group.units(k).unitSymbol for k in stat_keys] + u2_stat = ulist[0] + max2 = max(np.ptp(slice_dat[k]) for k in stat_keys) + f3, p3 = nice_scale_prefix(max2) + u2_stat = p3 + u2_stat + slice_y_label = mathlabel(*stat_keys, units=u2_stat, tex=tex) + + curves = [] + for k in stat_keys: + label = mathlabel(k, units=u2_stat, tex=tex) + curves.append(SliceCurve(key=k, label=label, values=slice_dat[k] / f3)) + + # Density overlay normalized to fit on the stat axis + density_raw = slice_dat["density"] + density_scaled = density_raw * max2 / density_raw.max() / f3 / 2 + + return DensityAndSlicePlotData( + hist2d=H, + extent=extent, + x_label=x_label, + y_label=y_label, + slice_x=slice_x, + slice_curves=curves, + slice_y_label=slice_y_label, + slice_y_factor=f3, + slice_density=density_scaled, + ) + + +def calculate_marginal( + data_scaled: np.ndarray, + weights: np.ndarray, + bins_count: int, +): + hist, bin_edges = np.histogram(data_scaled, bins=bins_count, weights=weights) + h_width = np.diff(bin_edges) + h_centers = bin_edges[:-1] + h_width / 2 + + profile = hist / h_width + + profile_scaled, profile_f, profile_prefix = nice_array(profile) + return profile_scaled, h_centers, h_width, profile_f, profile_prefix + + +def prepare_marginal_plot( + particle_group: ParticleGroup, + key1: str = "t", + key2: str = "p", + bins: int | None = None, + *, + xlim: Limit | None = None, + ylim: Limit | None = None, + nice: bool = True, + ellipse: bool = False, +) -> MarginalPlotData: + """ + Prepares the data structures required for marginal plotting. + + Calculates units, scaling, limits, histograms, and ellipses. + """ + if not bins: + n = len(particle_group) + bins = int(np.sqrt(n / 4)) + + x_raw = cast(np.ndarray, particle_group[key1]) + y_raw = cast(np.ndarray, particle_group[key2]) + + if np.all(np.isnan(x_raw)): + raise NanDataError(f"{key1} is all NaN") + + if np.all(np.isnan(y_raw)): + raise NanDataError(f"{key2} is all NaN") + + if len(x_raw) == 1: + bins = 100 + if xlim is None: + (x0,) = x_raw + if np.isclose(x0, 0.0): + xlim = (-1.0, 1.0) + else: + params = sorted((0.9 * x0, 1.1 * x0)) + xlim = (params[0], params[1]) + if ylim is None: + (y0,) = y_raw + if np.isclose(y0, 0.0): + ylim = (-1.0, 1.0) + else: + params = sorted((0.9 * y0, 1.1 * y0)) + ylim = (params[0], params[1]) + + x_data, f1, p1, xmin_raw, xmax_raw = plottable_array(x_raw, nice=nice, lim=xlim) + y_data, f2, p2, ymin_raw, ymax_raw = plottable_array(y_raw, nice=nice, lim=ylim) + + u1 = particle_group.units(key1).unitSymbol + u2 = particle_group.units(key2).unitSymbol + + weights = cast(np.ndarray, particle_group["weight"]) + + # X marginal (top) + x_prof, x_cents, x_width, x_prof_factor, x_prof_prefix = calculate_marginal( + x_data, weights, bins + ) + + # Y marginal (right) + y_prof, y_cents, y_width, y_prof_factor, y_prof_prefix = calculate_marginal( + y_data, weights, bins + ) + + ell_x, ell_y = None, None + if ellipse and len(x_data) > 1: + sigma_mat2 = particle_group.cov(key1, key2) + x_ellipse, y_ellipse = twiss_ellipse_points(sigma_mat2) + x_ellipse += particle_group.avg(key1) + y_ellipse += particle_group.avg(key2) + ell_x = x_ellipse / f1 + ell_y = y_ellipse / f2 + + return MarginalPlotData( + x=MarginalAxisData( + key=key1, + data=x_data, + lim=(xmin_raw / f1, xmax_raw / f1), + unit_factor=f1, + unit_prefix=p1, + unit_symbol=u1, + hist_centers=x_cents, + hist_values=x_prof, + hist_width=x_width, + hist_unit_factor=x_prof_factor, + hist_label_prefix=x_prof_prefix, + ), + y=MarginalAxisData( + key=key2, + data=y_data, + lim=(ymin_raw / f2, ymax_raw / f2), + unit_factor=f2, + unit_prefix=p2, + unit_symbol=u2, + hist_centers=y_cents, + hist_values=y_prof, + hist_width=y_width, + hist_unit_factor=y_prof_factor, + hist_label_prefix=y_prof_prefix, + ), + weights=weights, + bins=bins, + ellipse_x=ell_x, + ellipse_y=ell_y, + ) diff --git a/beamphysics/plot_bokeh.py b/beamphysics/plot_bokeh.py new file mode 100644 index 00000000..e036bedf --- /dev/null +++ b/beamphysics/plot_bokeh.py @@ -0,0 +1,1388 @@ +from __future__ import annotations + +import dataclasses +import logging +from typing import Literal + +import numpy as np +from bokeh.core.enums import SizingModeType +from bokeh.layouts import column, gridplot, row +from bokeh.models import ( + ColorBar, # pyright: ignore[reportPrivateImportUsage] + ColumnDataSource, # pyright: ignore[reportPrivateImportUsage] + Div, # pyright: ignore[reportPrivateImportUsage] + HoverTool, # pyright: ignore[reportPrivateImportUsage] + LayoutDOM, # pyright: ignore[reportPrivateImportUsage] + LinearAxis, # pyright: ignore[reportPrivateImportUsage] + LinearColorMapper, # pyright: ignore[reportPrivateImportUsage] + Range1d, # pyright: ignore[reportPrivateImportUsage] + Spacer, # pyright: ignore[reportPrivateImportUsage] +) +from bokeh.io import show as _bokeh_show +from bokeh.palettes import Palette, Viridis256 +from bokeh.plotting import figure + +from .labels import mathlabel +from .plot_base import ( + Limit, + prepare_density_and_slice_plot, + prepare_density_plot, + prepare_marginal_plot, + prepare_slice_plot, + prepare_wakefield_plot, +) +from .units import c_light + +logger = logging.getLogger(__name__) + + +def initialize_jupyter(): + # Is this public bokeh API? An attempt at forward-compatibility + try: + from bokeh.io.state import curstate + except ImportError: + pass + else: + state = curstate() + if getattr(state, "notebook", False): + # Jupyter already initialized + logger.debug("Bokeh output_notebook already called; not re-initializing") + return + + from bokeh.plotting import output_notebook + + output_notebook() + + +def _maybe_show(layout: LayoutDOM, show: bool = True) -> LayoutDOM: + """Call ``bokeh.io.show`` on *layout* if *show* is truthy.""" + if show: + _bokeh_show(layout) + return layout + + +@dataclasses.dataclass +class FontPlotSettings: + """Font settings for a single Bokeh figure's axes.""" + + axis_label_text_font_size: str = "14px" + axis_label_text_font_style: str = "italic" + major_label_text_font_size: str = "12px" + major_label_text_font_style: str = "normal" + + def apply(self, *axes) -> None: + """Apply these font settings to one or more Bokeh axis objects.""" + for axis in axes: + axis.axis_label_text_font_size = self.axis_label_text_font_size + axis.axis_label_text_font_style = self.axis_label_text_font_style + axis.major_label_text_font_size = self.major_label_text_font_size + axis.major_label_text_font_style = self.major_label_text_font_style + + +@dataclasses.dataclass +class MarginalFontSettings: + """Font settings for Bokeh marginal plots.""" + + text_font: str | None = None + annotation_text_font_size: str | None = None + main: FontPlotSettings = dataclasses.field(default_factory=FontPlotSettings) + top: FontPlotSettings = dataclasses.field( + default_factory=lambda: FontPlotSettings( + axis_label_text_font_size="10px", + major_label_text_font_size="8px", + ) + ) + right: FontPlotSettings = dataclasses.field( + default_factory=lambda: FontPlotSettings( + axis_label_text_font_size="10px", + major_label_text_font_size="8px", + ) + ) + + +def mathjax_fix(label: str) -> str: + """ + Adjust the Matplotlib-style LaTeX label for bokeh/MathJax. + """ + label = label.replace("µ", r" \mu ") + label = label.replace("$", "$$") + return label + + +@dataclasses.dataclass +class StatsAnnotation: + """A single beam statistic annotation row.""" + + label: str + sub_label: str + value: str + units: str + + +def get_annotations(particle_group, key1: str, key2: str) -> list[StatsAnnotation]: + """ + Return beam-statistic annotations for a given key combination. + + Parameters + ---------- + particle_group : ParticleGroup + The particle group to compute statistics from. + key1 : str + The x-axis key. + key2 : str + The y-axis key. + + Returns + ------- + list[StatsAnnotation] + """ + + # Longitudinal phase space: delta_z/c or z/c vs energy + if key1 in ("delta_z/c", "z/c") and key2 == "energy": + sigma_z = particle_group["sigma_z"] + sigma_p = particle_group["sigma_p"] + p0 = particle_group["mean_p"] + return [ + StatsAnnotation("σ", "z", f"{sigma_z / c_light * 1e15:.0f}", "fs"), + StatsAnnotation("σ", "δ", f"{sigma_p / p0 * 1e4:.1f} × 10⁻⁴", ""), + StatsAnnotation( + "⟨E⟩", "", f"{particle_group['mean_energy'] / 1e6:.1f}", "MeV" + ), + ] + + # Transverse spot: x vs y + if key1 == "x" and key2 == "y": + return [ + StatsAnnotation("⟨x⟩", "", f"{particle_group['mean_x'] * 1e6:.1f}", "µm"), + StatsAnnotation("⟨y⟩", "", f"{particle_group['mean_y'] * 1e6:.1f}", "µm"), + StatsAnnotation("σ", "x", f"{particle_group['sigma_x'] * 1e6:.1f}", "µm"), + StatsAnnotation("σ", "y", f"{particle_group['sigma_y'] * 1e6:.1f}", "µm"), + ] + + # Horizontal phase space: x vs xp or px + if key1 == "x" and key2 in ("xp", "px"): + return [ + StatsAnnotation( + "ε", "n,x", f"{particle_group['norm_emit_x'] * 1e6:.2f}", "mm-mrad" + ), + ] + + # Vertical phase space: y vs yp or py + if key1 == "y" and key2 in ("yp", "py"): + return [ + StatsAnnotation( + "ε", "n,y", f"{particle_group['norm_emit_y'] * 1e6:.2f}", "mm-mrad" + ), + ] + + return [] + + +def _annotations_to_html( + annotations: list[StatsAnnotation], + horizontal: bool = False, +) -> str | None: + """Convert a list of Annotation objects to an HTML table. + + Parameters + ---------- + annotations : list[StatsAnnotation] + horizontal : bool + If True, render items in a single row separated by spacing. + """ + if not annotations: + return None + + if horizontal: + items = [] + for a in annotations: + label = f"{a.label}{a.sub_label}" if a.sub_label else a.label + items.append(f"{label} = {a.value} {a.units}") + sep = "  ·  " + return f"{sep.join(items)}" + + rows = [] + for a in annotations: + label = f"{a.label}{a.sub_label}" if a.sub_label else a.label + rows.append( + f"{label}" + f"{a.value} {a.units}" + ) + return "" + "".join(rows) + "
" + + +def density_plot( + particle_group, + key: str = "x", + bins: int | str | None = None, + *, + xlim: Limit | None = None, + tex: bool = False, + nice: bool = True, + width: int = 600, + height: int = 400, + color: str = "gray", + alpha: float = 0.7, + sizing_mode: SizingModeType | None = None, + title: str | None = None, + show: bool = True, + **kwargs, +) -> LayoutDOM: + """ + 1D density histogram with Bokeh. + + Parameters + ---------- + particle_group : ParticleGroup + The object to plot. + key : str, default = 'x' + Which quantity to plot. + bins : int or str, optional + Number of bins. + xlim : tuple of float, optional + Manual x-axis limits. + nice : bool, default = True + Use nice unit scaling. + width, height : int + Figure dimensions in pixels. + show : bool, default = True + Display the plot. + + Returns + ------- + LayoutDOM + """ + pdata = prepare_density_plot( + particle_group, key=key, bins=bins, xlim=xlim, nice=nice, tex=tex + ) + + fig = figure( + width=width, + height=height, + x_axis_label=mathjax_fix(pdata.x_label), + y_axis_label=mathjax_fix(pdata.y_label), + tools="pan,wheel_zoom,box_zoom,save,reset", + toolbar_location="right", + ) + + fig.vbar( + x=pdata.hist_centers, + top=pdata.hist_values, + width=pdata.hist_width, + bottom=0, + fill_color=color, + line_color=color, + fill_alpha=alpha, + ) + + if pdata.xlim: + fig.x_range.start, fig.x_range.end = pdata.xlim + + if title: + fig.title.text = title + + if sizing_mode is not None: + fig.sizing_mode = sizing_mode + + fig.toolbar.logo = None + + return _maybe_show(fig, show) + + +def marginal_plot( + particle_group, + key1: str = "t", + key2: str = "p", + bins: int | None = None, + *, + xlim: Limit | None = None, + ylim: Limit | None = None, + nice: bool = True, + ellipse: bool = False, + width: int = 600, + height: int = 600, + colorbar: bool = False, + sizing_mode: SizingModeType | None = None, + x_label_orientation: float | None = np.pi / 4, + marginal_fraction: float = 0.33, + palette: Palette = Viridis256, + low_color: str = "#ffffff00", + text: str | None = None, + title: str | None = None, + font_settings: MarginalFontSettings | None = None, + stats_location: Literal["bottom", "top-right"] = "top-right", + show: bool = True, + **kwargs, +) -> LayoutDOM: + """ + Density plot and projections with bokeh. + + Parameters + ---------- + particle_group: ParticleGroup + The object to plot + key1: str, default = 't' + Key to bin on the x-axis + key2: str, default = 'p' + Key to bin on the y-axis + bins: int, default = None + Number of bins. If None, this will use a heuristic: + `bins = sqrt(n_particle/4)` + xlim: tuple, default = None + Manual setting of the x-axis limits. + ylim: tuple, default = None + Manual setting of the y-axis limits. + tex: bool, default = True + Use TEX for labels + nice: bool, default = True + Use "nice" prefixes. + ellipse: bool, default = True + If True, plot an ellipse representing the 2x2 sigma matrix. + sizing_mode: str, default = None + Bokeh sizing mode for responsive layout. When set (e.g. + ``"stretch_width"``), the layout uses ``row``/``column`` instead of + ``gridplot`` so that the marginal histograms scale correctly with the + main figure. + By default (None), a fixed-size ``GridPlot`` is returned. + marginal_fraction : float, default = 0.2 + Fraction of the plot to use for the marginal plots. + palette : bokeh.palettes.Palette, default=Viridis256 + Color map. + text : str or None, optional + Custom HTML text to display in the top-right corner. If ``None`` + (the default), automatic beam-statistics text is generated for + recognized key combinations (e.g. ``x``/``y``, ``x``/``px``, + ``delta_z/c``/``energy``). Pass an empty string ``""`` to suppress + automatic text. + title : str or None, optional + Title to set on the main density plot. + + Returns + ------- + LayoutDOM + + Examples + -------- + + >>> P = ParticleGroup("particles.h5") + >>> obj = marginal_plot(P, 't', 'energy', bins=200) + >>> bokeh.io.save(obj, "t_vs_energy.html") + """ + if kwargs: + logger.debug("Unused kwargs (may be for another backend): %s", kwargs) + if font_settings is None: + font_settings = MarginalFontSettings() + + pdata = prepare_marginal_plot( + particle_group, + key1=key1, + key2=key2, + bins=bins, + xlim=xlim, + ylim=ylim, + nice=nice, + ellipse=ellipse, + ) + + labelx = mathlabel(key1, units=pdata.x.full_unit, tex=False) + labely = mathlabel(key2, units=pdata.y.full_unit, tex=False) + + # Layout Sizes + main_w = int(width * (1.0 - marginal_fraction)) + main_h = int(height * (1.0 - marginal_fraction)) + marg_w = int(width * marginal_fraction) + marg_h = int(height * marginal_fraction) + + # Main Joint Figure + fig_joint = figure( + width=main_w, + height=main_h, + x_axis_label=labelx, + y_axis_label=labely, + x_range=pdata.x.lim, + y_range=pdata.y.lim, + tools="pan,wheel_zoom,box_zoom,save,reset", + toolbar_location="left", + ) + + font_settings.main.apply(fig_joint.xaxis, fig_joint.yaxis) + + if len(pdata.x.data) == 1: + fig_joint.scatter(pdata.x.data, pdata.y.data, size=10, color="navy") + else: + H, xedges, yedges = np.histogram2d( + pdata.x.data, pdata.y.data, bins=pdata.bins, weights=pdata.weights + ) + H = H.T + + h_min = np.min(H[H > 0]) if np.any(H > 0) else 0 + h_max = np.max(H) if np.any(H) else 1 + + mapper = LinearColorMapper( + palette=palette, + low=h_min, + high=h_max, + low_color=low_color, + ) + + if colorbar: + color_bar = ColorBar(color_mapper=mapper, location=(0, 0)) + fig_joint.add_layout(color_bar, "left") + + source_img = ColumnDataSource( + { + "image": [H], + "x": [xedges[0]], + "y": [yedges[0]], + "dw": [xedges[-1] - xedges[0]], + "dh": [yedges[-1] - yedges[0]], + } + ) + + image_renderer = fig_joint.image( + image="image", + x="x", + y="y", + dw="dw", + dh="dh", + source=source_img, + color_mapper=mapper, + ) + + hover = HoverTool( + renderers=[image_renderer], + tooltips=[ + (labelx, "$x"), + (labely, "$y"), + ("density", "@image"), + ], + ) + fig_joint.add_tools(hover) + + if pdata.ellipse_x is not None and pdata.ellipse_y is not None: + fig_joint.line( + pdata.ellipse_x, + pdata.ellipse_y, + color="red", + line_width=2, + alpha=0.8, + ) + + # Marginal Plots + + # Top (X projection) + p_top = figure( + width=main_w, + height=marg_h, + x_range=fig_joint.x_range, + y_axis_location="left", + min_border=0, + outline_line_color=None, + tools="", + ) + p_top.vbar( + x=pdata.x.hist_centers, + top=pdata.x.hist_values, + width=pdata.x.hist_width, + bottom=0, + fill_color="gray", + line_color="gray", + ) + + p_top.yaxis.axis_label = mathjax_fix(pdata.x.axis_label) + # p_top.yaxis.axis_label_orientation = ... + p_top.xaxis.visible = False + + # Right (Y projection) + p_right = figure( + width=marg_w, + height=main_h, + y_range=fig_joint.y_range, + x_axis_location="below", + min_border=0, + outline_line_color=None, + tools="", + ) + p_right.hbar( + y=pdata.y.hist_centers, + right=pdata.y.hist_values, + height=pdata.y.hist_width, + left=0, + fill_color="gray", + line_color="gray", + ) + p_right.xaxis.axis_label = mathjax_fix(pdata.y.axis_label) + # p_right.xaxis.axis_label_orientation = ... + p_right.yaxis.visible = False + + plots = [p_right, p_top, fig_joint] + if x_label_orientation is not None: + for plot in plots: + plot.xaxis.major_label_orientation = x_label_orientation + for plot in plots: + plot.toolbar.logo = None + + font_settings.top.apply(p_top.yaxis) + font_settings.right.apply(p_right.xaxis) + + if font_settings.text_font is not None: + for plot in plots: + for axis in (plot.xaxis, plot.yaxis): + axis.axis_label_text_font = font_settings.text_font + axis.major_label_text_font = font_settings.text_font + + annotations = get_annotations(particle_group, key1, key2) if text is None else [] + custom_text = text.replace("\n", "
") if text else None + + if title: + fig_joint.title.text = title + + # Build the stats Div (if any) with style depending on location + stats_div: Div | None = None + popup_font_size = font_settings.annotation_text_font_size or "12px" + popup_css = "" + if font_settings.text_font is not None: + popup_css += f"font-family: {font_settings.text_font}; " + + if custom_text or annotations: + if stats_location == "top-right": + content = custom_text or _annotations_to_html(annotations) + stats_div = Div( + text=f""" +
+ {content} +
+ """, + width=marg_w, + height=marg_h, + ) + else: + # "bottom" - horizontal stats bar below the plot + content = custom_text or _annotations_to_html(annotations, horizontal=True) + stats_div = Div( + text=f""" +
+ {content} +
+ """, + ) + + # Assemble layout + top_right: LayoutDOM = ( + stats_div + if stats_div is not None and stats_location == "top-right" + else Spacer(width=marg_w, height=marg_h) + ) + + if sizing_mode is not None: + fig_joint.sizing_mode = "scale_both" + fig_joint.aspect_ratio = main_h / main_w + p_top.sizing_mode = "stretch_width" + p_right.sizing_mode = "stretch_height" + + left_col = column(p_top, fig_joint, sizing_mode=sizing_mode) + right_col = column( + top_right, + p_right, + sizing_mode="stretch_height", + width=marg_w, + ) + plot_layout = row(left_col, right_col, sizing_mode=sizing_mode) + else: + plot_layout = gridplot( + [ + [p_top, top_right], + [fig_joint, p_right], + ], + merge_tools=True, + toolbar_location="left", + ) + + if stats_div is not None and stats_location == "bottom": + layout = column(plot_layout, stats_div) + else: + layout = plot_layout + + return _maybe_show(layout, show) + + +# Default Bokeh color cycle for multi-curve plots +_BOKEH_COLORS = [ + "#1f77b4", + "#ff7f0e", + "#2ca02c", + "#d62728", + "#9467bd", + "#8c564b", + "#e377c2", + "#7f7f7f", + "#bcbd22", + "#17becf", +] + + +def slice_plot( + particle_group, + *keys: str, + n_slice: int = 40, + slice_key: str | None = None, + xlim: Limit | None = None, + ylim: Limit | None = None, + nice: bool = True, + tex: bool = False, + width: int = 700, + height: int = 400, + sizing_mode: SizingModeType | None = None, + title: str | None = None, + density_alpha: float = 0.2, + show: bool = True, + **kwargs, +) -> LayoutDOM: + """ + Slice statistics plot with Bokeh. + + Plots slice statistics as lines on the primary y-axis and + the bunch density as a filled area on a secondary y-axis. + + Parameters + ---------- + particle_group : ParticleGroup + The object to plot. + keys : str + Statistical quantities to plot (e.g. ``'sigma_x'``, ``'norm_emit_x'``). + n_slice : int, default = 40 + Number of slices. + slice_key : str, optional + Dimension to slice in (``'t'``, ``'z'``, ``'delta_t'``, etc.). + xlim, ylim : tuple of float, optional + Manual axis limits. + nice : bool, default = True + Use nice unit scaling. + width, height : int + Figure dimensions in pixels. + + Returns + ------- + LayoutDOM + """ + if kwargs: + logger.debug("Unused kwargs (may be for another backend): %s", kwargs) + + pdata = prepare_slice_plot( + particle_group, + *keys, + n_slice=n_slice, + slice_key=slice_key, + xlim=xlim, + ylim=ylim, + nice=nice, + tex=tex, + ) + + fig = figure( + width=width, + height=height, + x_axis_label=mathjax_fix(pdata.x_label), + y_axis_label=mathjax_fix(pdata.y_label), + tools="pan,wheel_zoom,box_zoom,save,reset", + toolbar_location="right", + ) + + # Main curves + for i, curve in enumerate(pdata.curves): + color = ( + "black" if len(pdata.curves) == 1 else _BOKEH_COLORS[i % len(_BOKEH_COLORS)] + ) + fig.line( + pdata.x, + curve.values, + legend_label=mathjax_fix(curve.label), + color=color, + line_width=2, + ) + + if len(pdata.curves) > 1: + fig.legend.click_policy = "hide" + + # Density on secondary y-axis + density_max = ( + float(np.max(pdata.density_values)) if len(pdata.density_values) > 0 else 1.0 + ) + fig.extra_y_ranges["density"] = Range1d(start=0, end=density_max * 1.1) + fig.add_layout( + LinearAxis( + y_range_name="density", + axis_label=mathjax_fix(pdata.density_label), + ), + "right", + ) + + fig.varea( + x=pdata.x, + y1=0, + y2=pdata.density_values, + y_range_name="density", + fill_color="black", + fill_alpha=density_alpha, + ) + + if pdata.xlim: + fig.x_range.start, fig.x_range.end = pdata.xlim + if pdata.ylim: + fig.y_range.start, fig.y_range.end = pdata.ylim + + if title: + fig.title.text = title + + if sizing_mode is not None: + fig.sizing_mode = sizing_mode + + fig.toolbar.logo = None + + return _maybe_show(fig, show) + + +def wakefield_plot( + particle_group, + wake, + key: str | None = None, + nice: bool = True, + xlim: Limit | None = None, + ylim: Limit | None = None, + tex: bool = False, + bins: int | str | None = None, + width: int = 700, + height: int = 400, + sizing_mode: SizingModeType | None = None, + title: str | None = None, + density_alpha: float = 0.3, + scatter_size: float = 2, + show: bool = True, + **kwargs, +) -> LayoutDOM: + """ + Wakefield kicks scatter plot with density overlay using Bokeh. + + Parameters + ---------- + particle_group : ParticleGroup + The particle distribution. + wake : WakefieldBase + Wakefield object providing ``particle_kicks(z, weight)``. + key : str, optional + Independent variable key. Auto-detected if None. + nice : bool, default = True + Use nice unit scaling. + width, height : int + Figure dimensions in pixels. + density_alpha : float + Alpha for the density overlay bars. + scatter_size : float + Size of scatter markers. + + Returns + ------- + LayoutDOM + """ + + if kwargs: + logger.debug("Unused kwargs (may be for another backend): %s", kwargs) + pdata = prepare_wakefield_plot( + particle_group, + wake, + key=key, + nice=nice, + tex=tex, + xlim=xlim, + ylim=ylim, + bins=bins, + ) + + fig = figure( + width=width, + height=height, + x_axis_label=mathjax_fix(pdata.x_label), + y_axis_label=mathjax_fix(pdata.y_label), + tools="pan,wheel_zoom,box_zoom,save,reset", + toolbar_location="right", + ) + + # Density overlay on secondary y-axis + density_max = ( + float(np.max(pdata.density.hist_values)) + if len(pdata.density.hist_values) > 0 + else 1.0 + ) + fig.extra_y_ranges["density"] = Range1d(start=0, end=density_max * 1.1) + fig.add_layout( + LinearAxis( + y_range_name="density", + axis_label=mathjax_fix(pdata.density.y_label), + ), + "right", + ) + + fig.vbar( + x=pdata.density.hist_centers, + top=pdata.density.hist_values, + width=pdata.density.hist_width, + bottom=0, + y_range_name="density", + fill_color="gray", + line_color="gray", + fill_alpha=density_alpha, + ) + + # Wake kicks scatter + fig.scatter( + pdata.scatter_x, + pdata.scatter_y, + size=scatter_size, + color="black", + ) + + if pdata.xlim: + fig.x_range.start, fig.x_range.end = pdata.xlim + if pdata.ylim: + fig.y_range.start, fig.y_range.end = pdata.ylim + + if title: + fig.title.text = title + + if sizing_mode is not None: + fig.sizing_mode = sizing_mode + + fig.toolbar.logo = None + + return _maybe_show(fig, show) + + +def density_and_slice_plot( + particle_group, + key1: str = "t", + key2: str = "p", + stat_keys: list[str] | None = None, + bins: int = 100, + n_slice: int = 30, + tex: bool = False, + width: int = 700, + height: int = 450, + sizing_mode: SizingModeType | None = None, + title: str | None = None, + density_alpha: float = 0.1, + palette: Palette = Viridis256, + show: bool = True, + **kwargs, +) -> LayoutDOM: + """ + 2D density plot with overlaid slice statistics using Bokeh. + + Parameters + ---------- + particle_group : ParticleGroup + The object to plot. + key1 : str, default = 't' + Key for x-axis (also used as slice key). + key2 : str, default = 'p' + Key for y-axis (density). + stat_keys : list of str, optional + Slice statistics to overlay. + bins : int, default = 100 + Number of bins for the 2D histogram. + n_slice : int, default = 30 + Number of slices. + width, height : int + Figure dimensions in pixels. + + Returns + ------- + LayoutDOM + """ + pdata = prepare_density_and_slice_plot( + particle_group, + key1=key1, + key2=key2, + stat_keys=stat_keys, + bins=bins, + n_slice=n_slice, + tex=tex, + ) + + ext = pdata.extent # [xmin, xmax, ymin, ymax] + + # Color mapper for the 2D histogram + H = pdata.hist2d + h_min = float(np.min(H[H > 0])) if np.any(H > 0) else 0 + h_max = float(np.max(H)) if np.any(H) else 1 + + mapper = LinearColorMapper( + palette=palette, + low=h_min, + high=h_max, + low_color="#ffffff00", + ) + + fig = figure( + width=width, + height=height, + x_axis_label=mathjax_fix(pdata.x_label), + y_axis_label=mathjax_fix(pdata.y_label), + x_range=(ext[0], ext[1]), + y_range=(ext[2], ext[3]), + tools="pan,wheel_zoom,box_zoom,save,reset", + toolbar_location="right", + ) + + fig.image( + image=[H.T], + x=ext[0], + y=ext[2], + dw=ext[1] - ext[0], + dh=ext[3] - ext[2], + color_mapper=mapper, + ) + + # Slice statistics on secondary y-axis + stat_max = ( + max(float(np.max(c.values)) for c in pdata.slice_curves) + if pdata.slice_curves + else 1.0 + ) + fig.extra_y_ranges["stats"] = Range1d(start=0, end=stat_max * 1.1) + fig.add_layout( + LinearAxis( + y_range_name="stats", + axis_label=mathjax_fix(pdata.slice_y_label), + ), + "right", + ) + + for i, curve in enumerate(pdata.slice_curves): + color = _BOKEH_COLORS[i % len(_BOKEH_COLORS)] + fig.line( + pdata.slice_x, + curve.values, + y_range_name="stats", + legend_label=mathjax_fix(curve.label), + color=color, + line_width=2, + ) + + if len(pdata.slice_curves) > 1: + fig.legend.click_policy = "hide" + + # Density overlay + fig.varea( + x=pdata.slice_x, + y1=0, + y2=pdata.slice_density, + y_range_name="stats", + fill_color="black", + fill_alpha=density_alpha, + ) + + if title: + fig.title.text = title + if sizing_mode is not None: + fig.sizing_mode = sizing_mode + fig.toolbar.logo = None + + return _maybe_show(fig, show) + + +# --------------------------------------------------------------------------- +# Generic plotting functions (used by Wavefront, etc.) +# --------------------------------------------------------------------------- + + +def plot_1d_density( + x, + y, + x_name: str = "", + y_name: str | None = None, + x_units: str | None = None, + y_units: str | None = None, + log_scale_y: bool = False, + show_cdf: bool = False, + cdf_label: str = "CDF", + kind: str = "bar", + plot_style: dict | None = None, + xlim: Limit | None = None, + ylim: Limit | None = (0, None), + nice: bool = True, + auto_label: bool = False, + tex: bool = False, + data: dict | None = None, + width: int = 600, + height: int = 400, + sizing_mode: SizingModeType | None = None, + title: str | None = None, + show: bool = True, + **kwargs, +) -> LayoutDOM: + """ + Generic 1D density distribution plot with Bokeh. + + Mirrors the API of the matplotlib ``plot_1d_density``. + + Parameters + ---------- + x, y : array or str + Data arrays or string keys into *data* dict. + x_name, y_name : str + Axis labels. + x_units, y_units : str, optional + Units appended to labels. + log_scale_y : bool + Log scale on y-axis. + show_cdf : bool + Show cumulative distribution on secondary y-axis. + kind : str + ``'bar'`` or ``'line'``. + nice : bool + Use nice unit scaling. + data : dict, optional + Dict mapping string keys to arrays. + width, height : int + Figure dimensions. + + Returns + ------- + LayoutDOM + """ + + if kwargs: + logger.debug("Unused kwargs (may be for another backend): %s", kwargs) + from .units import pg_units, plottable_array + + # Resolve data dict + x_key = None + y_key = None + + if isinstance(x, str): + if data is None: + raise ValueError("If `x` is a string, `data` dict must be provided") + x_key = x + x = np.asarray(data[x_key]) + else: + x = np.asarray(x) + + if isinstance(y, str): + if data is None: + raise ValueError("If `y` is a string, `data` dict must be provided") + y_key = y + y = np.asarray(data[y_key]) + else: + y = np.asarray(y) + + if x_key is not None and x_name == "": + x_name = x_key + if y_key is not None and y_name is None: + y_name = y_key + if y_name is None: + y_name = "Density" + + # Auto-label + if auto_label: + if x_key and x_units is None: + try: + x_units = pg_units(x_key).unitSymbol + except (ValueError, KeyError): + pass + if y_key and y_units is None: + try: + y_units = pg_units(y_key).unitSymbol + except (ValueError, KeyError): + pass + + # Nice scaling + x, f1, p1, x_min, x_max = plottable_array(x, nice=nice, lim=xlim) + y, f2, p2, y_min, y_max = plottable_array(y, nice=nice, lim=ylim) + + if x_units: + x_units = p1 + str(x_units) + elif p1: + x_units = p1 + + if y_units: + y_units = p2 + str(y_units) + elif p2: + y_units = p2 + + # Labels + if auto_label and x_key: + x_label = mathjax_fix(mathlabel(x_key, units=x_units, tex=tex)) + else: + x_label = f"{x_name} ({x_units})" if x_units else x_name + + if auto_label and y_key: + y_label = mathjax_fix(mathlabel(y_key, units=y_units, tex=tex)) + else: + y_label = f"{y_name} ({y_units})" if y_units else y_name + + # Bar widths + if len(x) > 1: + widths = np.diff(x) + widths = np.append(widths, widths[-1]) + else: + widths = np.ones_like(x) + + fig = figure( + width=width, + height=height, + x_axis_label=x_label, + y_axis_label=y_label, + tools="pan,wheel_zoom,box_zoom,save,reset", + toolbar_location="right", + y_axis_type="log" if log_scale_y else "auto", + ) + + if plot_style is None: + plot_style = {} + + if kind == "bar": + color = plot_style.get("color", "gray") + alpha = plot_style.get("alpha", 0.7) + fig.vbar( + x=x, + top=y, + width=widths, + bottom=0, + fill_color=color, + line_color=color, + fill_alpha=alpha, + ) + elif kind == "line": + color = plot_style.get("color", "blue") + line_width = plot_style.get("linewidth", plot_style.get("line_width", 2)) + fig.line(x, y, color=color, line_width=line_width) + else: + raise ValueError(f"kind must be 'bar' or 'line', got '{kind}'") + + if xlim is not None: + fig.x_range.start, fig.x_range.end = x_min / f1, x_max / f1 + if ylim is not None: + if ylim[0] is not None: + fig.y_range.start = y_min / f2 + if ylim[1] is not None: + fig.y_range.end = y_max / f2 + + # CDF on secondary y-axis + if show_cdf: + cdf = np.cumsum(y * widths) * f1 * f2 + cdf_scaled, _, cdf_prefix, _, _ = plottable_array(cdf, nice=nice) + + cdf_max = float(np.max(cdf_scaled)) if len(cdf_scaled) > 0 else 1.0 + fig.extra_y_ranges["cdf"] = Range1d(start=0, end=cdf_max) + cdf_axis_label = f"{cdf_label} ({cdf_prefix})" if cdf_prefix else cdf_label + fig.add_layout( + LinearAxis(y_range_name="cdf", axis_label=cdf_axis_label), + "right", + ) + fig.line(x, cdf_scaled, y_range_name="cdf", color="blue", line_width=2) + + if title: + fig.title.text = title + if sizing_mode is not None: + fig.sizing_mode = sizing_mode + fig.toolbar.logo = None + + return _maybe_show(fig, show) + + +def plot_2d_density_with_marginals( + data: np.ndarray, + dx: float = 1, + dy: float = 1, + xmin: float | None = None, + ymin: float | None = None, + x_name: str = "", + y_name: str = "", + z_name: str = "", + x_units: str | None = None, + y_units: str | None = None, + z_units: str | None = None, + log_scale_z: bool = False, + log_scale_marginals: bool = False, + show_colorbar: bool = True, + xlim: Limit | None = None, + ylim: Limit | None = None, + vmin: float | None = None, + vmax: float | None = None, + width: int = 600, + height: int = 600, + marginal_fraction: float = 0.25, + palette: Palette = Viridis256, + sizing_mode: SizingModeType | None = None, + title: str | None = None, + show: bool = True, + **kwargs, +) -> LayoutDOM: + """ + 2D density map with marginal histograms using Bokeh. + + Mirrors the API of the matplotlib ``plot_2d_density_with_marginals``. + + Parameters + ---------- + data : np.ndarray + 2D array of density values, shape ``(nx, ny)``. + dx, dy : float + Grid spacing. + xmin, ymin : float, optional + Origin of the grid. Default centers at 0. + x_name, y_name, z_name : str + Axis labels. + x_units, y_units, z_units : str, optional + Units appended to labels. + log_scale_z : bool + Log color mapping. + palette : Palette + Bokeh color palette. + width, height : int + Figure dimensions. + + Returns + ------- + LayoutDOM + """ + if kwargs: + logger.debug("Unused kwargs (may be for another backend): %s", kwargs) + nx, ny = data.shape + + if xmin is None: + xmin = -((nx - 1) * dx) / 2 + if ymin is None: + ymin = -((ny - 1) * dy) / 2 + + xmax = xmin + (nx - 1) * dx + ymax = ymin + (ny - 1) * dy + + xvec = np.linspace(xmin, xmax, nx) + yvec = np.linspace(ymin, ymax, ny) + + x_marginal = np.sum(data, axis=1) * dy + y_marginal = np.sum(data, axis=0) * dx + + vmin = vmin if vmin is not None else float(np.min(data)) + vmax = vmax if vmax is not None else float(np.max(data)) + + x_label = mathjax_fix(f"{x_name} ({x_units})" if x_units else x_name) + y_label = mathjax_fix(f"{y_name} ({y_units})" if y_units else y_name) + + # Layout sizes + main_w = int(width * (1.0 - marginal_fraction)) + main_h = int(height * (1.0 - marginal_fraction)) + marg_w = int(width * marginal_fraction) + marg_h = int(height * marginal_fraction) + + # Color mapper + if log_scale_z: + low = max(vmin, vmax * 1e-6) + mapper = LinearColorMapper(palette=palette, low=low, high=vmax) + else: + mapper = LinearColorMapper(palette=palette, low=vmin, high=vmax) + + # Main density figure + x_range = xlim or (xmin - dx / 2, xmax + dx / 2) + y_range = ylim or (ymin - dy / 2, ymax + dy / 2) + + fig_main = figure( + width=main_w, + height=main_h, + x_axis_label=x_label, + y_axis_label=y_label, + x_range=x_range, + y_range=y_range, + tools="pan,wheel_zoom,box_zoom,save,reset", + toolbar_location="left", + ) + + fig_main.image( + image=[data.T], + x=xmin - dx / 2, + y=ymin - dy / 2, + dw=xmax - xmin + dx, + dh=ymax - ymin + dy, + color_mapper=mapper, + ) + + if show_colorbar: + cbar_label = mathjax_fix(f"{z_name} ({z_units})" if z_units else z_name) + color_bar = ColorBar(color_mapper=mapper, title=cbar_label, location=(0, 0)) + fig_main.add_layout(color_bar, "left") + + if title: + fig_main.title.text = title + + # Top marginal (X projection) + p_top = figure( + width=main_w, + height=marg_h, + x_range=fig_main.x_range, + y_axis_type="log" if log_scale_marginals else "auto", + min_border=0, + outline_line_color=None, + tools="", + ) + p_top.vbar( + x=xvec, + top=x_marginal, + width=dx, + bottom=0, + fill_color="gray", + line_color="gray", + ) + if z_units and y_units: + p_top.yaxis.axis_label = mathjax_fix(f"{z_units} {y_units}") + p_top.xaxis.visible = False + + # Right marginal (Y projection) + p_right = figure( + width=marg_w, + height=main_h, + y_range=fig_main.y_range, + x_axis_type="log" if log_scale_marginals else "auto", + min_border=0, + outline_line_color=None, + tools="", + ) + p_right.hbar( + y=yvec, + right=y_marginal, + height=dy, + left=0, + fill_color="gray", + line_color="gray", + ) + if z_units and x_units: + p_right.xaxis.axis_label = mathjax_fix(f"{z_units} {x_units}") + p_right.yaxis.visible = False + + for p in (fig_main, p_top, p_right): + p.toolbar.logo = None + + top_right = Spacer(width=marg_w, height=marg_h) + + if sizing_mode is not None: + fig_main.sizing_mode = "scale_both" + # fig_main.aspect_ratio = main_h / main_w + p_top.sizing_mode = "stretch_width" + p_right.sizing_mode = "stretch_height" + left_col = column(p_top, fig_main, sizing_mode=sizing_mode) + right_col = column( + top_right, p_right, sizing_mode="stretch_height", width=marg_w + ) + layout = row(left_col, right_col, sizing_mode=sizing_mode) + else: + layout = gridplot( + [ + [p_top, top_right], + [fig_main, p_right], + ], + merge_tools=True, + toolbar_location="left", + ) + + return _maybe_show(layout, show) diff --git a/beamphysics/plot_dispatch.py b/beamphysics/plot_dispatch.py new file mode 100644 index 00000000..a03cd22c --- /dev/null +++ b/beamphysics/plot_dispatch.py @@ -0,0 +1,292 @@ +from __future__ import annotations + +import os + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol + +import numpy as np + +from .plot_base import Limit + +import functools +import logging +import sys + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from .particles import ParticleGroup + + +# --------------------------------------------------------------------------- +# Protocols – common parameter signatures for each plot function +# --------------------------------------------------------------------------- + + +class DensityPlotFn(Protocol): + """1D density histogram of a single particle key.""" + + def __call__( + self, + particle_group: ParticleGroup, + key: str = ..., + bins: int | str | None = ..., + *, + xlim: Limit | None = ..., + nice: bool = ..., + **kwargs: Any, + ) -> Any: ... + + +class MarginalPlotFn(Protocol): + """2D density with marginal histograms.""" + + def __call__( + self, + particle_group: ParticleGroup, + key1: str = ..., + key2: str = ..., + bins: int | None = ..., + *, + xlim: Limit | None = ..., + ylim: Limit | None = ..., + nice: bool = ..., + ellipse: bool = ..., + **kwargs: Any, + ) -> Any: ... + + +class SlicePlotFn(Protocol): + """Slice statistics with density overlay.""" + + def __call__( + self, + particle_group: ParticleGroup, + *keys: str, + n_slice: int = ..., + slice_key: str | None = ..., + xlim: Limit | None = ..., + ylim: Limit | None = ..., + nice: bool = ..., + **kwargs: Any, + ) -> Any: ... + + +class WakefieldPlotFn(Protocol): + """Wakefield kicks scatter with density overlay.""" + + def __call__( + self, + particle_group: ParticleGroup, + wake: Any, + key: str | None = ..., + nice: bool = ..., + xlim: Limit | None = ..., + ylim: Limit | None = ..., + **kwargs: Any, + ) -> Any: ... + + +class Plot1dDensityFn(Protocol): + """Generic 1D density distribution plot.""" + + def __call__( + self, + x: str | np.ndarray, + y: str | np.ndarray, + *, + nice: bool = ..., + xlim: Limit | None = ..., + ylim: Limit | None = ..., + **kwargs: Any, + ) -> Any: ... + + +class Plot2dDensityWithMarginalsFn(Protocol): + """Generic 2D density map with marginal histograms.""" + + def __call__( + self, + data: np.ndarray, + dx: float = ..., + dy: float = ..., + **kwargs: Any, + ) -> Any: ... + + +class DensityAndSlicePlotFn(Protocol): + """2D density with overlaid slice statistics.""" + + def __call__( + self, + particle_group: ParticleGroup, + key1: str = ..., + key2: str = ..., + stat_keys: list[str] = ..., + bins: int = ..., + n_slice: int = ..., + **kwargs: Any, + ) -> Any: ... + + +# --------------------------------------------------------------------------- +# PlotBackend – holds one callable per plot type +# --------------------------------------------------------------------------- + + +@dataclass +class PlotBackend: + """ + Container for a set of plot functions from a single backend. + + Instantiated lazily on first access via `get_backend`. + """ + + name: str + density_plot: DensityPlotFn + marginal_plot: MarginalPlotFn + slice_plot: SlicePlotFn + wakefield_plot: WakefieldPlotFn + density_and_slice_plot: DensityAndSlicePlotFn + plot_1d_density: Plot1dDensityFn + plot_2d_density_with_marginals: Plot2dDensityWithMarginalsFn + + +# --------------------------------------------------------------------------- +# Module-level default and resolution +# --------------------------------------------------------------------------- + +_default_backend: str = os.environ.get("BEAMPHYSICS_PLOT", "mpl") +_backend_cache: dict[str, PlotBackend] = {} + + +def set_default_backend(name: str) -> None: + """ + Set the module-level default plot backend. + + Parameters + ---------- + name : str + ``"mpl"`` for Matplotlib or ``"bokeh"`` for Bokeh. + """ + global _default_backend + if name not in ("mpl", "bokeh"): + raise ValueError(f"Unknown backend {name!r}. Choose 'mpl' or 'bokeh'.") + _default_backend = name + + +def get_default_backend() -> str: + """Return the current module-level default backend name.""" + return _default_backend + + +def resolve_backend(backend: str | None = None, obj: Any = None) -> str: + """ + Determine which backend to use. + + Priority: explicit *backend* argument > ``obj.plot_backend`` > module default. + """ + if backend is not None: + return backend + if obj is not None: + obj_backend = getattr(obj, "plot_backend", None) + if obj_backend is not None: + return obj_backend + return _default_backend + + +@functools.cache +def is_jupyter() -> bool: + """ + Determine if we're in a Jupyter notebook session. + + This works by way of interacting with IPython display and seeing what + choice it makes regarding reprs. + + Returns + ------- + bool + """ + if "IPython" not in sys.modules or "IPython.display" not in sys.modules: + return False + + from IPython.display import display + + class ReprCheck: + def _repr_html_(self) -> str: + self.mode = "jupyter" + logger.info("Detected Jupyter. Using the notebook graph backend.") + return "" + + def __repr__(self) -> str: + self.mode = "console" + return "" + + check = ReprCheck() + display(check) + return check.mode == "jupyter" + + +def _load_mpl_backend() -> PlotBackend: + from . import plot as mod + + return PlotBackend( + name="mpl", + density_plot=mod.density_plot, + marginal_plot=mod.marginal_plot, + slice_plot=mod.slice_plot, + wakefield_plot=mod.wakefield_plot, + density_and_slice_plot=mod.density_and_slice_plot, + plot_1d_density=mod.plot_1d_density, + plot_2d_density_with_marginals=mod.plot_2d_density_with_marginals, + ) + + +def _load_bokeh_backend() -> PlotBackend: + from . import plot_bokeh as mod + + backend = PlotBackend( + name="bokeh", + density_plot=mod.density_plot, + marginal_plot=mod.marginal_plot, + slice_plot=mod.slice_plot, + wakefield_plot=mod.wakefield_plot, + density_and_slice_plot=mod.density_and_slice_plot, + plot_1d_density=mod.plot_1d_density, + plot_2d_density_with_marginals=mod.plot_2d_density_with_marginals, + ) + + if is_jupyter(): + mod.initialize_jupyter() + return backend + + +_backend_loaders = { + "mpl": _load_mpl_backend, + "bokeh": _load_bokeh_backend, +} + + +def get_backend(backend: str | None = None, obj: Any = None) -> PlotBackend: + """ + Resolve the backend name and return a :class:`PlotBackend`. + + The backend module is lazy-loaded on first access and cached. + + Parameters + ---------- + backend : str or None + Explicit backend name (``"mpl"`` or ``"bokeh"``). + If ``None``, falls through to *obj* and then the module default. + obj : object, optional + An object with an optional ``plot_backend`` attribute + (e.g. a :class:`ParticleGroup`). + """ + name = resolve_backend(backend, obj) + if name not in _backend_cache: + loader = _backend_loaders.get(name) + if loader is None: + raise ValueError(f"Unknown backend {name!r}. Choose 'mpl' or 'bokeh'.") + _backend_cache[name] = loader() + return _backend_cache[name] diff --git a/beamphysics/wavefront/wavefront.py b/beamphysics/wavefront/wavefront.py index fec146eb..3b8ded3b 100644 --- a/beamphysics/wavefront/wavefront.py +++ b/beamphysics/wavefront/wavefront.py @@ -11,7 +11,6 @@ import h5py import matplotlib.pyplot as plt import numpy as np -from matplotlib.colors import LogNorm from numpy.fft import fftfreq, fftshift, ifftn, ifftshift from scipy.constants import c, e, epsilon_0, hbar @@ -19,7 +18,7 @@ load_genesis4_fields, wavefront_write_genesis4, ) -from ..plot import plot_1d_density, plot_2d_density_with_marginals +from ..plot_dispatch import get_backend from ..statistics import mean_calc, mean_variance_calc from ..units import Z0, c_light from ..wavefront.propagators import drift_wavefront @@ -923,63 +922,107 @@ def intensity(self) -> np.ndarray | float: """ return self.intensity_x + self.intensity_y - def plot_spectral_intensity(self, cmap="inferno", logscale=False): - """ - Simple projected intensity plot - + def plot_spectral_intensity( + self, + cmap="inferno", + logscale=False, + backend=None, + return_figure=False, + **kwargs, + ): """ + Projected spectral intensity plot with marginals. - xlabel = r"$\theta_x$ (µrad)" - ylabel = r"$\theta_y$ (µrad)" - xfactor = 1e6 + Parameters + ---------- + cmap : str, default = 'inferno' + Colormap name. + logscale : bool, default = False + Use log scale for color and marginals. + backend : str, optional + Plot backend: ``'mpl'`` or ``'bokeh'``. + return_figure : bool, default = False + If True, return the figure/layout object. + """ + xfactor = 1e6 # rad -> µrad yfactor = 1e6 zfactor = self.k0**2 / (1e6 * 1e6) - label = r"Spectral $F$ (J/µrad$^2$)" - extent = ( - self.thetaxmin * xfactor, - self.thetaxmax * xfactor, - self.thetaymin * yfactor, - self.thetaymax * yfactor, - ) + F = zfactor * self.spectral_fluence # shape (nx, ny) - # Alternatively: - # extent = (self.kxmin, self.kxmax, self.kymin, self.kymax) - # xlabel = r'$k_x$ (rad/m)' - # ylabel = r'$k_y$ (rad/m)' - # zfactor = 1 - # label = r"Spectral $F$ (J$\cdot$m$^2$)" - - F = zfactor * self.spectral_fluence - Fmax = np.max(F) + # Grid spacing in µrad + dthetax = ( + (self.thetaxmax - self.thetaxmin) / (self.nx - 1) * xfactor + if self.nx > 1 + else 1.0 + ) + dthetay = ( + (self.thetaymax - self.thetaymin) / (self.ny - 1) * yfactor + if self.ny > 1 + else 1.0 + ) - fig, ax = plt.subplots() - im = ax.imshow( - F.T, + be = get_backend(backend) + fig = be.plot_2d_density_with_marginals( + F, + dx=dthetax, + dy=dthetay, + xmin=self.thetaxmin * xfactor, + ymin=self.thetaymin * yfactor, + x_name=r"$\theta_x$", + x_units="µrad", + y_name=r"$\theta_y$", + y_units="µrad", + z_name=r"Spectral $F$", + z_units="J/µrad$^2$", cmap=cmap, - extent=extent, - origin="lower", - ) # Note data.T and origin='lower' are required - if logscale: - norm = LogNorm(vmin=Fmax / 1e6, vmax=Fmax) - im.set_norm(norm) - - fig.colorbar(im, ax=ax, label=label) + log_scale_marginals=logscale, + log_scale_z=logscale, + return_figure=True, + **kwargs, + ) + if return_figure: + return fig - ax.set_xlabel(xlabel) - ax.set_ylabel(ylabel) + def plot_photon_energy_spectrum( + self, xlim=None, ax=None, backend=None, return_figure=False, **kwargs + ): + """ + Photon energy spectrum plot. - def plot_photon_energy_spectrum(self, xlim=None, ax=None): + Parameters + ---------- + xlim : tuple of float, optional + X-axis limits. + ax : matplotlib.axes.Axes, optional + Existing axes. Only used by the ``'mpl'`` backend. + backend : str, optional + Plot backend: ``'mpl'`` or ``'bokeh'``. + return_figure : bool, default = False + If True, return the figure/layout object. + """ x = self.photon_energy_vec # eV - y = self.photon_energy_spectrum # J/eV - - if ax is None: - _, ax = plt.subplots() - ax.plot(x, y * 1e6, color="purple") - ax.set_xlabel("photon energy (eV)") - ax.set_ylabel("photon spectral energy density (µJ/eV)") - ax.set_ylim(0, None) - ax.set_xlim(xlim) + y = self.photon_energy_spectrum * 1e6 # µJ/eV + + be = get_backend(backend) + fig = be.plot_1d_density( + x, + y, + x_name="photon energy", + y_name="photon spectral energy density", + x_units="eV", + y_units="µJ/eV", + kind="line", + plot_style={"color": "purple"}, + ylim=(0, None), + xlim=xlim, + ax=ax, + nice=False, + return_figure=True, + **kwargs, + ) + if return_figure: + return fig # Statistics @@ -1333,13 +1376,17 @@ def plot_power( nice=True, log_scale_y=False, show_cdf=False, + backend=None, + return_figure=False, + **kwargs, ): x = self.zvec / c y = self.power data = {"z/c": x, "power": y} - return plot_1d_density( + be = get_backend(backend) + fig = be.plot_1d_density( "z/c", "power", data=data, @@ -1352,47 +1399,67 @@ def plot_power( plot_style={"color": "purple"}, kind="bar", nice=nice, + return_figure=True, + **kwargs, ) + if return_figure: + return fig - def plot_fluence(self, cmap="inferno", logscale=False): + def plot_fluence( + self, + cmap="inferno", + logscale=False, + backend=None, + return_figure=False, + **kwargs, + ): """ - Simple fluence plot + Fluence plot with marginal projections. + Parameters + ---------- + cmap : str, default = 'inferno' + Colormap name. + logscale : bool, default = False + Use log scale for color and marginals. + backend : str, optional + Plot backend: ``'mpl'`` or ``'bokeh'``. """ - - xlabel = r"$x$ (cm)" - ylabel = r"$y$ (cm)" xfactor = 100 yfactor = 100 zfactor = 1 / (100 * 100) # 1/m^2 -> 1/cm^2 - label = r"$F$ (J/cm$^2$)" - extent = ( - self.xmin * xfactor, - self.xmax * xfactor, - self.ymin * yfactor, - self.ymax * yfactor, - ) - - F = self.fluence * zfactor - Fmax = np.max(F) + F = self.fluence - fig, ax = plt.subplots() - im = ax.imshow( - F.T, + be = get_backend(backend) + fig = be.plot_2d_density_with_marginals( + F * zfactor, + dx=self.dx * xfactor, + dy=self.dy * yfactor, + xmin=self.xmin * xfactor, + ymin=self.ymin * yfactor, + x_name=r"$x$", + x_units="cm", + y_name=r"$y$", + y_units="cm", + z_name=r"$F$", + z_units="J/cm$^2$", cmap=cmap, - extent=extent, - origin="lower", - ) # Note data.T and origin='lower' are required - if logscale: - norm = LogNorm(vmin=Fmax / 1e6, vmax=Fmax) - im.set_norm(norm) - - fig.colorbar(im, ax=ax, label=label) - - ax.set_xlabel(xlabel) - ax.set_ylabel(ylabel) + log_scale_marginals=logscale, + log_scale_z=logscale, + return_figure=True, + **kwargs, + ) + if return_figure: + return fig - def plot2(self, cmap="inferno", logscale=False): + def plot2( + self, + cmap="inferno", + logscale=False, + backend=None, + return_figure=False, + **kwargs, + ): """ Simple fluence plot @@ -1401,15 +1468,13 @@ def plot2(self, cmap="inferno", logscale=False): This is experimental. """ - # xlabel = r"$x$ (cm)" - # ylabel = r"$y$ (cm)" xfactor = 100 yfactor = 100 zfactor = 1 / (100 * 100) # 1/m^2 -> 1/cm^2 - # label = r"$F$ (J/cm$^2$)" F = self.fluence - plot_2d_density_with_marginals( + be = get_backend(backend) + fig = be.plot_2d_density_with_marginals( F * zfactor, dx=self.dx * xfactor, dy=self.dy * yfactor, @@ -1424,12 +1489,11 @@ def plot2(self, cmap="inferno", logscale=False): cmap=cmap, log_scale_marginals=logscale, log_scale_z=logscale, + return_figure=True, + **kwargs, ) - - # if logscale: - # Fmax = np.max(F) - # norm = LogNorm(vmin=Fmax / 1e6, vmax=Fmax) - # im.set_norm(norm) + if return_figure: + return fig @property def dkx(self) -> float: diff --git a/docs/examples/plot_backend_comparison.ipynb b/docs/examples/plot_backend_comparison.ipynb new file mode 100644 index 00000000..0dbf70c1 --- /dev/null +++ b/docs/examples/plot_backend_comparison.ipynb @@ -0,0 +1,411 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "14b974e8", + "metadata": {}, + "source": [ + "# Plot Backend Comparison: Matplotlib vs Bokeh\n", + "\n", + "Side-by-side comparison of every plot type using both backends." + ] + }, + { + "cell_type": "markdown", + "id": "5891afab", + "metadata": {}, + "source": [ + "### Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fda61932", + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import base64\n", + "import io\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from bokeh.embed import file_html\n", + "from bokeh.resources import Resources\n", + "from IPython.display import HTML, display\n", + "\n", + "from beamphysics import ParticleGroup\n", + "from beamphysics.plot_dispatch import get_backend\n", + "\n", + "mpl_be = get_backend(\"mpl\")\n", + "bokeh_be = get_backend(\"bokeh\")\n", + "\n", + "ROW_HEIGHT = 800\n", + "\n", + "\n", + "def mpl_to_base64(fig: plt.Figure) -> str:\n", + " \"\"\"Render a matplotlib figure to a base64-encoded PNG.\"\"\"\n", + " buf = io.BytesIO()\n", + " fig.savefig(buf, format=\"png\", dpi=120, bbox_inches=\"tight\")\n", + " plt.close(fig)\n", + " buf.seek(0)\n", + " return base64.b64encode(buf.read()).decode()\n", + "\n", + "\n", + "def bokeh_to_srcdoc(layout) -> str:\n", + " \"\"\"Render a bokeh layout to an HTML string suitable for iframe srcdoc.\"\"\"\n", + " html = file_html(layout, resources=Resources(mode=\"inline\"), title=\"\")\n", + " return html.replace(\"&\", \"&\").replace('\"', \""\")\n", + "\n", + "\n", + "def compare(title: str, mpl_fig, bokeh_layout):\n", + " \"\"\"Display a side-by-side comparison with fixed row height.\"\"\"\n", + " img_b64 = mpl_to_base64(mpl_fig)\n", + " srcdoc = bokeh_to_srcdoc(bokeh_layout)\n", + " display(\n", + " HTML(f\"\"\"\n", + "

{title}

\n", + "
\n", + "
\n", + "

Matplotlib

\n", + " \n", + "
\n", + "
\n", + "

Bokeh

\n", + " \n", + "
\n", + "
\n", + " \"\"\")\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e529ede3", + "metadata": {}, + "outputs": [], + "source": [ + "P = ParticleGroup(\"data/bmad_particles2.h5\")\n", + "P.t = P.t - P[\"mean_t\"]" + ] + }, + { + "cell_type": "markdown", + "id": "b9b52f36", + "metadata": {}, + "source": [ + "## Density Plot (1D)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7122d682", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"P.plot('t')\",\n", + " mpl_be.density_plot(P, key=\"t\"),\n", + " bokeh_be.density_plot(P, key=\"t\", show=False),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "8386dde8", + "metadata": {}, + "source": [ + "## Marginal Plot — x vs y" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7a17a7b3", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"P.plot('x', 'y')\",\n", + " mpl_be.marginal_plot(P, key1=\"x\", key2=\"y\"),\n", + " bokeh_be.marginal_plot(P, key1=\"x\", key2=\"y\", show=False),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "b9c97155", + "metadata": {}, + "source": [ + "## Marginal Plot — x vs px with ellipse" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "74da6e0d", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"P.plot('x', 'px', ellipse=True)\",\n", + " mpl_be.marginal_plot(P, key1=\"x\", key2=\"px\", ellipse=True),\n", + " bokeh_be.marginal_plot(P, key1=\"x\", key2=\"px\", ellipse=True, show=False),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "ecfae361", + "metadata": {}, + "source": [ + "## Marginal Plot — t vs energy" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "70559c73", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"P.plot('t', 'energy')\",\n", + " mpl_be.marginal_plot(P, key1=\"t\", key2=\"energy\"),\n", + " bokeh_be.marginal_plot(P, key1=\"t\", key2=\"energy\", show=False),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "823e48c6", + "metadata": {}, + "source": [ + "## Marginal Plot — x vs y (spot)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4cd837c4", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"P.plot('x', 'y')\",\n", + " mpl_be.marginal_plot(P, key1=\"x\", key2=\"y\"),\n", + " bokeh_be.marginal_plot(P, key1=\"x\", key2=\"y\", show=False),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "6945099f", + "metadata": {}, + "source": [ + "## Slice Plot" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a6d4f9f9", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"P.slice_plot('sigma_x', 'sigma_y')\",\n", + " mpl_be.slice_plot(P, \"sigma_x\", \"sigma_y\"),\n", + " bokeh_be.slice_plot(P, \"sigma_x\", \"sigma_y\", show=False),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "f30151f5", + "metadata": {}, + "source": [ + "## Density + Slice Plot" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e37bcd60", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"density_and_slice_plot(P, 't', 'energy', ...)\",\n", + " mpl_be.density_and_slice_plot(\n", + " P, \"t\", \"energy\", stat_keys=[\"sigma_x\", \"sigma_y\"], n_slice=200, bins=200\n", + " ),\n", + " bokeh_be.density_and_slice_plot(\n", + " P,\n", + " \"t\",\n", + " \"energy\",\n", + " stat_keys=[\"sigma_x\", \"sigma_y\"],\n", + " n_slice=200,\n", + " bins=200,\n", + " show=False,\n", + " ),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "da989f45", + "metadata": {}, + "source": [ + "# Wavefront Plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "36677637", + "metadata": {}, + "outputs": [], + "source": [ + "from beamphysics import Wavefront\n", + "\n", + "W = Wavefront.from_gaussian(\n", + " shape=(51, 51, 21),\n", + " dx=10e-6,\n", + " dy=10e-6,\n", + " dz=10e-6,\n", + " wavelength=1e-9,\n", + " sigma0=50e-6,\n", + " energy=1.0,\n", + ")\n", + "Wk = W.to_kspace()" + ] + }, + { + "cell_type": "markdown", + "id": "69ee4e3e", + "metadata": {}, + "source": [ + "## Wavefront Power" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aef05f3e", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"W.plot_power()\",\n", + " W.plot_power(backend=\"mpl\", return_figure=True),\n", + " W.plot_power(backend=\"bokeh\", return_figure=True, show=False),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "3cbfeec3", + "metadata": {}, + "source": [ + "## Wavefront Fluence" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11a4f4b6", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"W.plot_fluence()\",\n", + " W.plot_fluence(backend=\"mpl\", return_figure=True),\n", + " W.plot_fluence(backend=\"bokeh\", return_figure=True, show=False),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "30254b07", + "metadata": {}, + "source": [ + "## Wavefront Fluence with Marginals (plot2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d6f69944", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"W.plot2()\",\n", + " W.plot2(backend=\"mpl\", return_figure=True),\n", + " W.plot2(backend=\"bokeh\", return_figure=True, show=False),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "c1b0cac9", + "metadata": {}, + "source": [ + "## Spectral Intensity (k-space)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "35e46e18", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"Wk.plot_spectral_intensity()\",\n", + " Wk.plot_spectral_intensity(backend=\"mpl\", return_figure=True),\n", + " Wk.plot_spectral_intensity(backend=\"bokeh\", return_figure=True, show=False),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "aad3631c", + "metadata": {}, + "source": [ + "## Photon Energy Spectrum (k-space)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aea47298", + "metadata": {}, + "outputs": [], + "source": [ + "compare(\n", + " \"Wk.plot_photon_energy_spectrum()\",\n", + " Wk.plot_photon_energy_spectrum(backend=\"mpl\", return_figure=True),\n", + " Wk.plot_photon_energy_spectrum(backend=\"bokeh\", return_figure=True, show=False),\n", + ")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/examples/plot_examples.ipynb b/docs/examples/plot_examples.ipynb index a60e3311..06905aea 100644 --- a/docs/examples/plot_examples.ipynb +++ b/docs/examples/plot_examples.ipynb @@ -4,7 +4,13 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Plot examples" + "# Plot examples\n", + "\n", + "openPMD-beamphysics supports Matplotlib and Bokeh for plots. By default, matplotlib will be used.\n", + "\n", + "You may set the environment variable BEAMPHYSICS_PLOT to either `\"mpl\"` or `\"bokeh\"` to select the backend prior to launching Python.\n", + "\n", + "In code, you can use `beamphysics.set_default_backend`." ] }, { @@ -42,7 +48,26 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Density plots" + "# Matplotlib" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from beamphysics import set_default_backend\n", + "\n", + "# Set matplotlib as the default for this section\n", + "set_default_backend(\"mpl\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Density plots" ] }, { @@ -58,7 +83,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Slice statistics Plots" + "## Slice statistics Plots" ] }, { @@ -75,7 +100,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Marginal plots" + "## Marginal plots" ] }, { @@ -91,7 +116,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Combined density and slice plot" + "## Combined density and slice plot" ] }, { @@ -115,20 +140,119 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "tags": [] - }, + "metadata": {}, "outputs": [], "source": [ "density_and_slice_plot(\n", " P, \"t\", \"energy\", stat_keys=[\"sigma_x\", \"sigma_y\"], n_slice=200, bins=200\n", ")" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Bokeh" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "set_default_backend(\"bokeh\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Density plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "P.plot(\"t\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Slice statistics Plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "P.t = P.t - P[\"mean_t\"]\n", + "P.slice_plot(\"sigma_x\", \"sigma_y\", xlim=(None, 80e-15), ylim=(0, 30e-6))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Marginal plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "P.plot(\"y_bar\", \"py_bar\", ellipse=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Combined density and slice plot" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from beamphysics.plot_bokeh import density_and_slice_plot" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "P.species" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "density_and_slice_plot(\n", + " P, \"t\", \"energy\", stat_keys=[\"sigma_x\", \"sigma_y\"], n_slice=200, bins=200\n", + ");" + ] } ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "kyber", "language": "python", "name": "python3" }, diff --git a/environment.yml b/environment.yml index 66d20c82..1b29540f 100644 --- a/environment.yml +++ b/environment.yml @@ -7,6 +7,7 @@ dependencies: - numpy - scipy>=1.0.0 - matplotlib + - bokeh # optional - h5py - python-dateutil # Developer diff --git a/tests/test_plot_bokeh.py b/tests/test_plot_bokeh.py new file mode 100644 index 00000000..9b8e3619 --- /dev/null +++ b/tests/test_plot_bokeh.py @@ -0,0 +1,339 @@ +"""Tests for the Bokeh plotting backend.""" + +from __future__ import annotations + + +import numpy as np +import pytest + +try: + from bokeh.io import save + from bokeh.models.layouts import LayoutDOM + from bokeh.resources import Resources +except ImportError: + # <-- I like sorted imports without 'noqa' everywhere; repeat import then + # skip here + bokeh = pytest.importorskip("bokeh") + raise + + +from beamphysics import ParticleGroup, set_default_backend +from beamphysics.particles import single_particle +from beamphysics.plot_dispatch import get_backend, get_default_backend, _default_backend +from beamphysics.wavefront.wavefront import Wavefront + +from conftest import test_artifacts + +import beamphysics.plot_bokeh as _plot_bokeh_mod + +P = ParticleGroup("docs/examples/data/bmad_particles.h5") + +_bokeh_artifacts = test_artifacts / "bokeh" +_bokeh_artifacts.mkdir(exist_ok=True) +_resources = Resources() + + +@pytest.fixture(autouse=True) +def _bokeh_show_to_save(request, monkeypatch): + """Intercept bokeh show() calls and save to HTML artifacts instead.""" + index = 0 + node_name = request.node.name.replace("/", "_") + + def _save_instead(layout): + nonlocal index + filename = _bokeh_artifacts / f"{node_name}_{index}.html" + save(layout, filename=str(filename), resources=_resources, title=node_name) + print(f"Saved bokeh artifact to {filename}") + index += 1 + + monkeypatch.setattr(_plot_bokeh_mod, "_bokeh_show", _save_instead) + + +# --------------------------------------------------------------------------- +# Dispatch tests +# --------------------------------------------------------------------------- + + +def test_dispatch_default(): + assert get_default_backend() == _default_backend + be = get_backend() + assert be.name == _default_backend + + +def test_dispatch_explicit_bokeh(): + be = get_backend("bokeh") + assert be.name == "bokeh" + + +def test_dispatch_set_default(): + old = get_default_backend() + try: + set_default_backend("bokeh") + assert get_default_backend() == "bokeh" + be = get_backend() + assert be.name == "bokeh" + finally: + set_default_backend(old) + + +def test_dispatch_invalid(): + with pytest.raises(ValueError, match="Unknown backend"): + set_default_backend("plotly") + with pytest.raises(ValueError, match="Unknown backend"): + get_backend("plotly") + + +# --------------------------------------------------------------------------- +# ParticleGroup density_plot (1D) +# --------------------------------------------------------------------------- + + +def test_density_plot(): + result = P.plot("x", backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"density_plot_x") + + +def test_density_plot_with_options(): + result = P.plot("t", backend="bokeh", return_figure=True, bins=50) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"density_plot_t_bins50") + + +# --------------------------------------------------------------------------- +# ParticleGroup marginal_plot (2D) +# --------------------------------------------------------------------------- + +MARGINAL_PAIRS = [ + ("x", "px"), + ("y", "py"), + ("t", "energy"), + ("x", "y"), +] + + +@pytest.mark.parametrize("key1,key2", MARGINAL_PAIRS, ids=lambda p: str(p)) +def test_marginal_plot(key1, key2): + result = P.plot(key1, key2, backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,f"marginal_plot_{key1}_{key2}") + + +def test_marginal_plot_with_ellipse(): + result = P.plot("x", "px", backend="bokeh", return_figure=True, ellipse=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"marginal_plot_x_px_ellipse") + + +# --------------------------------------------------------------------------- +# ParticleGroup slice_plot +# --------------------------------------------------------------------------- + +SLICE_KEYS = ["sigma_x", "norm_emit_x"] + + +@pytest.mark.parametrize("stat_key", SLICE_KEYS) +def test_slice_plot(stat_key): + result = P.slice_plot(stat_key, backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,f"slice_plot_{stat_key}") + + +def test_slice_plot_multi_keys(): + result = P.slice_plot("sigma_x", "sigma_y", backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + + +# --------------------------------------------------------------------------- +# ParticleGroup density_and_slice_plot +# --------------------------------------------------------------------------- + + +def test_density_and_slice_plot(): + be = get_backend("bokeh") + result = be.density_and_slice_plot(P, key1="t", key2="p", return_figure=True) + assert isinstance(result, LayoutDOM) + + +def test_density_and_slice_plot_custom_keys(): + be = get_backend("bokeh") + result = be.density_and_slice_plot( + P, key1="t", key2="energy", stat_keys=["sigma_x", "sigma_y"], return_figure=True + ) + assert isinstance(result, LayoutDOM) + + +# --------------------------------------------------------------------------- +# Single-particle edge case +# --------------------------------------------------------------------------- + + +@pytest.mark.filterwarnings("ignore:.*invalid value encountered in.*") +@pytest.mark.filterwarnings("ignore:.*divide by zero.*") +@pytest.mark.filterwarnings("ignore:.*Degrees of freedom.*") +@pytest.mark.filterwarnings("ignore:.*The fit may be poorly conditioned.*") +def test_single_particle_density(): + Ps = single_particle(pz=10e6) + result = Ps.plot("x", backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"single_particle_density") + + +@pytest.mark.filterwarnings("ignore:.*invalid value encountered in.*") +@pytest.mark.filterwarnings("ignore:.*divide by zero.*") +@pytest.mark.filterwarnings("ignore:.*Degrees of freedom.*") +@pytest.mark.filterwarnings("ignore:.*The fit may be poorly conditioned.*") +def test_single_particle_marginal(): + Ps = single_particle(pz=10e6) + result = Ps.plot("x", "px", backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"single_particle_marginal") + + +# --------------------------------------------------------------------------- +# Generic plot_1d_density (used by Wavefront) +# --------------------------------------------------------------------------- + + +def test_plot_1d_density_bar(): + be = get_backend("bokeh") + x = np.linspace(0, 10, 50) + y = np.exp(-x) + result = be.plot_1d_density( + x, y, x_name="x", y_name="f(x)", kind="bar", return_figure=True + ) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"plot_1d_density_bar") + + +def test_plot_1d_density_line(): + be = get_backend("bokeh") + x = np.linspace(0, 10, 50) + y = np.exp(-x) + result = be.plot_1d_density( + x, y, x_name="x", y_name="f(x)", kind="line", return_figure=True + ) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"plot_1d_density_line") + + +def test_plot_1d_density_with_data_dict(): + be = get_backend("bokeh") + data = {"time": np.linspace(0, 1, 100), "signal": np.random.randn(100)} + result = be.plot_1d_density( + "time", "signal", data=data, kind="line", return_figure=True + ) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"plot_1d_density_data_dict") + + +def test_plot_1d_density_with_cdf(): + be = get_backend("bokeh") + x = np.linspace(0, 10, 50) + y = np.exp(-x) + result = be.plot_1d_density(x, y, show_cdf=True, return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"plot_1d_density_cdf") + + +# --------------------------------------------------------------------------- +# Generic plot_2d_density_with_marginals (used by Wavefront) +# --------------------------------------------------------------------------- + + +def test_plot_2d_density_with_marginals(): + be = get_backend("bokeh") + data = np.random.rand(50, 50) + result = be.plot_2d_density_with_marginals( + data, dx=0.1, dy=0.1, x_name="x", y_name="y", return_figure=True + ) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"plot_2d_density_with_marginals") + + +def test_plot_2d_density_with_log_scale(): + be = get_backend("bokeh") + data = np.random.rand(50, 50) + 0.01 + result = be.plot_2d_density_with_marginals( + data, + dx=0.1, + dy=0.1, + log_scale_z=True, + log_scale_marginals=True, + return_figure=True, + ) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"plot_2d_density_log_scale") + + +# --------------------------------------------------------------------------- +# Wavefront plots +# --------------------------------------------------------------------------- + + +def _make_wavefront(): + return Wavefront.from_gaussian( + shape=(51, 51, 21), + dx=10e-6, + dy=10e-6, + dz=10e-6, + wavelength=1e-9, + sigma0=50e-6, + energy=1.0, + ) + + +@pytest.mark.filterwarnings("ignore:.*identical low and high.*:UserWarning") +def test_wavefront_plot_power(): + W = _make_wavefront() + result = W.plot_power(backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"wavefront_plot_power") + + +def test_wavefront_plot_fluence(): + W = _make_wavefront() + result = W.plot_fluence(backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"wavefront_plot_fluence") + + +def test_wavefront_plot2(): + W = _make_wavefront() + result = W.plot2(backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"wavefront_plot2") + + +def test_wavefront_plot_spectral_intensity(): + W = _make_wavefront() + Wk = W.to_kspace() + result = Wk.plot_spectral_intensity(backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"wavefront_plot_spectral_intensity") + + +def test_wavefront_plot_photon_energy_spectrum(): + W = _make_wavefront() + Wk = W.to_kspace() + result = Wk.plot_photon_energy_spectrum(backend="bokeh", return_figure=True) + assert isinstance(result, LayoutDOM) + # Artifact saved automatically by _bokeh_show_to_save fixture + # _save_bokeh(result,"wavefront_plot_photon_energy_spectrum")