Source code for arpes.plotting.tof

"""Plotting routines for count-like data from a ToF or delay line.

This module is a bit of a misnomer, in that it also applies perfectly well to data collected by a delay
line on a hemisphere, the important point is that the data in any given channel should correspond to the true number of
electrons that arrived in that channel.

Plotting routines here are ones that include statistical errorbars. Generally for datasets in PyARPES, an xr.Dataset
will hold the standard deviation data for a given variable on `{var_name}_std`.
"""
import matplotlib.pyplot as plt
import numpy as np

import xarray as xr
from arpes.plotting.utils import path_for_plot
from arpes.provenance import save_plot_provenance
from arpes.typing import DataType

__all__ = (
    "plot_with_std",
    "scatter_with_std",
)


[docs]@save_plot_provenance def plot_with_std(data: DataType, name_to_plot=None, ax=None, out=None, **kwargs): """Makes a fill-between line plot with error bars from associated statistical errors.""" if name_to_plot is None: var_names = [k for k in data.data_vars.keys() if "_std" not in k] assert len(var_names) == 1 name_to_plot = var_names[0] assert (name_to_plot + "_std") in data.data_vars.keys() fig = None if ax is None: fig, ax = plt.subplots( figsize=kwargs.pop( "figsize", ( 7, 5, ), ) ) data.data_vars[name_to_plot].plot(ax=ax, **kwargs) x, y = data.data_vars[name_to_plot].G.to_arrays() std = data.data_vars[name_to_plot + "_std"].values ax.fill_between(x, y - std, y + std, alpha=0.3, **kwargs) if out is not None: plt.savefig(path_for_plot(out), dpi=400) return path_for_plot(out) ax.set_xlim([np.min(x), np.max(x)]) return fig, ax
[docs]@save_plot_provenance def scatter_with_std(data: DataType, name_to_plot=None, ax=None, fmt="o", out=None, **kwargs): """Makes a scatter plot of data with error bars generated from associated statistical errors.""" if name_to_plot is None: var_names = [k for k in data.data_vars.keys() if "_std" not in k] assert len(var_names) == 1 name_to_plot = var_names[0] assert (name_to_plot + "_std") in data.data_vars.keys() fig = None if ax is None: fig, ax = plt.subplots( figsize=kwargs.pop( "figsize", ( 7, 5, ), ) ) x, y = data.data_vars[name_to_plot].G.to_arrays() std = data.data_vars[name_to_plot + "_std"].values ax.errorbar(x, y, yerr=std, fmt=fmt, markeredgecolor="black", **kwargs) if out is not None: plt.savefig(path_for_plot(out), dpi=400) return path_for_plot(out) ax.set_xlim([np.min(x), np.max(x)]) return fig, ax