plotting

Plotting utilities. These are all fairly implementation specific right now, they should be refactored to rely on more generic units.

auto_relplot(data, x, y, ignore=None, merge=None, auto=True, errorbar=None, max_errorbars=20, **kwargs)

Make a relplot. This function tries to be semi-clever about automatically selecting which terms to use for hue, col, row, and style.

First it determines the categories of data as any fields other than the ones specified in x and y or in ignore. We then assign the category with the most unique entries to hue, the second most to col, then row, and then style. Any categories beyond that are combined with whatever the hue category is. Note that you can still manually specify things with **kwargs.

Parameters:

Name	Type	Description	Default
data	dict[str, NDArray]	The data to plot. Should be a dict that seaborn understands as a long form dataset.	required
x	str	The field to plot as x.	required
y	str	The field to plot as y.	required
ignore	list[str]	List of fields to ignore.	None
merge	list[list[str]]	Fields to forcibly merge. Each element in this list should be a list of fields. The merged field will have a name with the format {field1}+{field2}+.... The fields specified here cannot overlap with each other, ignore, or fields specified by **kwargs.	None
auto	bool	If False then all automatic features to pick categories for the relplot are turned off and this is simply a wrapper to relplot with the ability to specify merged fields.	True
errorbar	str	The field containing the precomputed 1-sigma errors for y. Requires kind="line". Errorbars inherit the hue and style of the corresponding lines.	None
**kwargs		Keyword arguments to pass to relplot. Note that you can pass categorical variables like hue here if you want to specify them manually.	{}

Returns:

Name	Type	Description
plot	FacetGrid	The plot output by relplot.

Source code in lat_beams/plotting.py

def auto_relplot(
    data: dict[str, NDArray],
    x: str,
    y: str,
    ignore: Optional[list[str]] = None,
    merge: Optional[list[list[str]]] = None,
    auto: bool = True,
    errorbar: Optional[str] = None,
    max_errorbars: Optional[int] = 20,
    **kwargs,
) -> sns.FacetGrid:
    """
    Make a relplot.
    This function tries to be semi-clever about automatically selecting
    which terms to use for `hue`, `col`, `row`, and `style`.

    First it determines the categories of `data` as any fields other than
    the ones specified in `x` and `y` or in `ignore`. We then assign the
    category with the most unique entries to `hue`, the second most to
    `col`, then `row`, and then `style`. Any categories beyond that are
    combined with whatever the `hue` category is. Note that you can still
    manually specify things with `**kwargs`.

    Parameters
    ----------
    data : dict[str, NDArray]
        The data to plot. Should be a dict that `seaborn` understands as a
        long form dataset.
    x : str
        The field to plot as x.
    y : str
        The field to plot as y.
    ignore : list[str], optional
        List of fields to ignore.
    merge : list[list[str]], optional
        Fields to forcibly merge. Each element in this list should be a list
        of fields. The merged field will have a name with the format
        `{field1}+{field2}+...`. The fields specified here cannot overlap
        with each other, `ignore`, or fields specified by `**kwargs`.
    auto : bool, default=True
        If False then all automatic features to pick categories for the
        relplot are turned off and this is simply a wrapper to `relplot`
        with the ability to specify merged fields.
    errorbar : str, optional
        The field containing the precomputed 1-sigma errors for `y`.
        Requires `kind="line"`. Errorbars inherit the `hue` and `style`
        of the corresponding lines.
    **kwargs
        Keyword arguments to pass to `relplot`. Note that you can pass
        categorical variables like `hue` here if you want to specify them
        manually.

    Returns
    -------
    plot : sns.FacetGrid
        The plot output by `relplot`.
    """
    ignore = [] if ignore is None else ignore.copy()
    merge = [] if merge is None else merge

    if errorbar is not None:
        if errorbar not in data:
            raise ValueError(f"Errorbar field {errorbar!r} not found in data.")
        if kwargs.get("kind", "scatter") != "line":
            raise ValueError("Explicit errorbars require kind='line'.")
        if errorbar not in ignore:
            ignore.append(errorbar)

    cats = ["hue", "row", "col", "style"]
    in_kwargs = [kwargs[c] for c in cats if kwargs.get(c) is not None]
    can_add = [c for c in cats if kwargs.get(c) is None]

    for mlist in merge:
        if not mlist:
            continue

        for m in mlist:
            if m in ignore or m in in_kwargs:
                raise ValueError(
                    f"Can't perform merge with {m}, it already has a role specified!"
                )
            ignore.append(m)

        name = "+".join(mlist)
        data[name] = np.array(
            [
                "+".join(str(data[c][i]) for c in mlist)
                for i in range(len(data[mlist[0]]))
            ]
        )

    if auto:
        fields = [
            k
            for k in data
            if k not in (x, y) and k not in in_kwargs and k not in ignore
        ]
        fields.sort(key=lambda c: len(set(data[c])))

        for field, cat in zip(fields, can_add):
            kwargs[cat] = field

        if len(fields) > len(can_add):
            to_combine = [kwargs["hue"]] + fields[len(can_add) :]
            name = "+".join(to_combine)
            data[name] = np.array(
                [
                    "+".join(str(data[c][i]) for c in to_combine)
                    for i in range(len(data[x]))
                ]
            )
            kwargs["hue"] = name

    if np.issubdtype(data[x].dtype, np.str_):
        nx = len(np.unique(data[x]))
        kwargs.setdefault("height", 5)
        kwargs.setdefault("aspect", max(6, nx * 0.4) / kwargs["height"])
        order = np.argsort(data[x])
        data = {k: np.asarray(v)[order] for k, v in data.items()}

    plot = sns.relplot(data=data, x=x, y=y, **kwargs)

    if errorbar is not None:
        hue = kwargs.get("hue")
        style = kwargs.get("style")

        for (row_i, col_i, _), facet_data in plot.facet_data():
            ax = plot.facet_axis(row_i, col_i)

            groups = [(None, None, facet_data)]
            if hue is not None or style is not None:
                keys = [c for c in (hue, style) if c is not None]
                groups = [
                    (
                        key[0] if hue is not None else None,
                        key[-1] if style is not None else None,
                        group,
                    )
                    for key, group in facet_data.groupby(keys, sort=False)
                ]

            for hue_value, style_value, group in groups:
                # Subsample errorbars while preserving the full curve.
                if max_errorbars is not None:
                    if max_errorbars < 1:
                        raise ValueError("max_errorbars must be >= 1.")

                    if len(group) > max_errorbars:
                        order = np.argsort(np.asarray(group[x]))

                        if max_errorbars == 1:
                            take = np.array([len(order) // 2])
                        elif max_errorbars == 2:
                            take = np.array([0, len(order) - 1])
                        else:
                            take = np.linspace(
                                0,
                                len(order) - 1,
                                max_errorbars,
                            ).astype(int)

                        group = group.iloc[order[take]]

                # Get the color from the line corresponding to this group.
                color = None

                if hue is not None:
                    for line in ax.lines:
                        if line.get_label() == str(hue_value):
                            color = line.get_color()
                            break

                if color is None:
                    color = ax.lines[0].get_color() if ax.lines else "C0"

                ax.errorbar(
                    np.asarray(group[x]),
                    np.asarray(group[y]),
                    yerr=np.asarray(group[errorbar]),
                    fmt="none",
                    color=color,
                    capsize=2,
                    capthick=1,
                    elinewidth=1,
                    alpha=0.7,
                )

    for ax in plot.axes.flat:
        ax.label_outer()

    return plot

plot_focal_plane(focal_plane, fit_plot_dir, ufm, obs_id)

Plot the results of a the focal plane fit. We assume here that all angular values are in radians and the amplitude is in pW.

The following plots will be made:

• Xi vs Eta scatter, file = "{ufm}_fp.png"
• Az vs El scatter, file = "{ufm}_enc.png"
• amplitude histogram, "{ufm}_fp_amp.png"
• FWHM histogram, "{ufm}_fp_fwhm.png"
• Hits histogram, "{ufm}_fp_hits.png"
• Reduced Chi Squared histogram, "{ufm}_fp_red_chisq.png"
• R^2 histogram, "{ufm}_fp_r2.png"

Parameters:

Name	Type	Description	Default
focal_plane	AxisManager	The fit focal plane. Should contain the following fit parameters (all with the same number of elements): xi, eta, az, el, amp, fwhm, hits, reduced_chisq, R2.	required
fit_plot_dir	str	The directory to save plots to.	required
ufm	str	The UFM that was fit. Used to determine filenames (see above).	required

Source code in lat_beams/plotting.py

def plot_focal_plane(
    focal_plane: AxisManager, fit_plot_dir: str, ufm: str, obs_id: str
):
    """
    Plot the results of a the focal plane fit.
    We assume here that all angular values are in radians
    and the amplitude is in pW.

    The following plots will be made:

    * Xi vs Eta scatter, file = "{ufm}_fp.png"
    * Az vs El scatter, file = "{ufm}_enc.png"
    * amplitude histogram, "{ufm}_fp_amp.png"
    * FWHM histogram, "{ufm}_fp_fwhm.png"
    * Hits histogram, "{ufm}_fp_hits.png"
    * Reduced Chi Squared histogram, "{ufm}_fp_red_chisq.png"
    * R^2 histogram, "{ufm}_fp_r2.png"

    Parameters
    ----------
    focal_plane : AxisManager
        The fit focal plane.
        Should contain the following fit parameters (all with the same number of elements):
        `xi`, `eta`, `az`, `el`, `amp`, `fwhm`, `hits`, `reduced_chisq`, `R2`.
    fit_plot_dir : str
        The directory to save plots to.
    ufm : str
        The UFM that was fit.
        Used to determine filenames (see above).
    """
    plt.close()

    # Convert to np arrays to make pyright happy
    fp_data = {
        attr: np.array(getattr(focal_plane, attr))
        for attr in [
            "xi",
            "eta",
            "az",
            "el",
            "amp",
            "fwhm",
            "hits",
            "reduced_chisq",
            "R2",
        ]
    }

    # Define plots with tuples (data, xlabel, ylabel, filename, plot_type)
    plots = [
        (
            (fp_data["xi"], fp_data["eta"]),
            "Xi (rad)",
            "Eta (rad)",
            f"{ufm}_fp.png",
            "scatter",
        ),
        (
            (fp_data["az"], fp_data["el"]),
            "Az (rad)",
            "El (rad)",
            f"{ufm}_enc.png",
            "scatter",
        ),
        (fp_data["amp"], "Amp (pW)", "Dets (#)", f"{ufm}_fp_amp.png", "hist"),
        (fp_data["fwhm"], "FWHM (rad)", "Dets (#)", f"{ufm}_fp_fwhm.png", "hist"),
        (fp_data["hits"], "Hits (#)", "Dets (#)", f"{ufm}_fp_hits.png", "hist"),
        (
            fp_data["reduced_chisq"],
            "Reduced Chi Squared",
            "Dets (#)",
            f"{ufm}_fp_red_chisq.png",
            "hist",
        ),
        (fp_data["R2"], "R2", "Dets (#)", f"{ufm}_fp_r2.png", "hist"),
    ]

    for data, xlabel, ylabel, filename, plot_type in plots:
        plt.clf()
        if plot_type == "scatter":
            plt.scatter(*data, alpha=0.25)
        elif plot_type == "hist":
            plt.hist(data, bins=30, alpha=0.25)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.title(f"{obs_id} {ufm}")
        plt.savefig(os.path.join(fit_plot_dir, filename), bbox_inches="tight")
    plt.close()

plot_map(data, posmap, pixsize, extent, cent, plot_dir, title, comp='T', log=False, log_thresh=0.001, append='', units='"')

Plot a beam map using imshow using the coolwarm colormap. This also plots the profile of the map.

Parameters:

Name	Type	Description	Default
data	Float[ndarray, 'nx ny'] | ndmap	The beam map to plot.	required
posmap	Float[ndarray, '2 nx ny'] | ndmap	The position map for the beam map. Should be generated by enmap.posmap.	required
extent	float	The radius of the plot to make in the same units as posmap.	required
cent	tuple[float, float]	The center of the beam map in the same units as posmap.	required
plot_dir	str	The directory to save plots to.	required
title	str	The base title of the plot. The full title will have the format "{title} {comp} {append} {'log10'*log}". The output file will use the full title as the filename but with whitespace replaced by '_' and 'map' or 'prof' inserted between "{title}" and "{comp}" depending on which plot it is.	required
comp	str	The comp of the map we are plotting. Should be "T", "Q", "U", "Qr", or "Ur".	"T"
log	bool	If True then plot the log10 of the beam instead. This uses a symetric log norm, or signs are preserved.	True
log_thresh	float	The threshold at which the log plot transitions to linear to avoid issues near values of 0.	1e-3
append	str	String to append ti title and filename. See title for details.	""
units	str	The units of posmap to appear in axes labels.	'"'

Source code in lat_beams/plotting.py

def plot_map(
    data: Float[np.ndarray, "nx ny"] | enmap.ndmap,
    posmap: Float[np.ndarray, "2 nx ny"] | enmap.ndmap,
    pixsize: float,
    extent: float,
    cent: tuple[float, float],
    plot_dir: str,
    title: str,
    comp: str = "T",
    log: bool = False,
    log_thresh: float = 1e-3,
    append: str = "",
    units: str = '"',
):
    """
    Plot a beam map using `imshow` using the `coolwarm` colormap.
    This also plots the profile of the map.

    Parameters
    ----------
    data : Float[np.ndarray, "nx ny"] | enmap.ndmap
        The beam map to plot.
    posmap : Float[np.ndarray, "2 nx ny"] | enmap.ndmap
        The position map for the beam map.
        Should be generated by `enmap.posmap`.
    extent : float
        The radius of the plot to make in the same units as `posmap`.
    cent : tuple[float, float]
        The center of the beam map in the same units as `posmap`.
    plot_dir : str
        The directory to save plots to.
    title : str
        The base title of the plot.
        The full title will have the format "{title} {comp} {append} {'log10'*log}".
        The output file will use the full title as the filename but with whitespace replaced by '_'
        and 'map' or 'prof' inserted between "{title}" and "{comp}" depending on which plot it is.
    comp : str, default: "T"
        The comp of the map we are plotting.
        Should be "T", "Q", "U", "Qr", or "Ur".
    log : bool, default: True
        If True then plot the log10 of the beam instead.
        This uses a symetric log norm, or signs are preserved.
    log_thresh : float, default: 1e-3
        The threshold at which the log plot transitions to linear to avoid issues near values of 0.
    append : str, default: ""
        String to append ti title and filename.
        See `title` for details.
    units : str, default: '"'
        The units of `posmap` to appear in axes labels.
    """
    plt.close()
    # Restrict to our local patch (oversample by 2)
    r = np.sqrt((posmap[0] - cent[1]) ** 2 + (posmap[1] - cent[0]) ** 2)
    ext_msk = r <= extent * np.sqrt(2)
    w = np.where(ext_msk)
    ext_sl = (
        slice(np.min(w[0]), np.max(w[0]) + 1),
        slice(np.min(w[1]), np.max(w[1]) + 1),
    )
    data = data[ext_sl]
    posmap = (posmap[0][ext_sl], posmap[1][ext_sl])

    # Use posmap to setup coordinates
    # posmap, pixsize, extent, and cent must all be in the same units
    # We are assuming square pixels here
    plt_extent = (posmap[1].max(), posmap[1].min(), posmap[0].min(), posmap[0].max())

    # Get radial profile and plot map with appropriate norm
    _norm = None
    label = f"_{comp}"
    if append != "":
        label += f"_{append}"
    rsq = (posmap[0] - cent[1]) ** 2 + (posmap[1] - cent[0]) ** 2
    rcent = np.unravel_index(np.argmin(rsq), data.shape)[::-1]
    rprof = radial_profile(
        data,
        (int(rcent[0]), int(rcent[1])),
    )[: int(0.5 * min(*data.shape))]
    if log:
        _norm = SymLogNorm(linthresh=log_thresh, clip=True, vmin=-1, vmax=1)
        label += f"_log10"
        with np.errstate(divide="ignore", invalid="ignore"):
            rprof = np.sign(rprof) * np.log10(np.abs(rprof))
        plt.imshow(data, origin="lower", extent=plt_extent, norm=_norm)
    else:
        vminmax = np.max(np.abs(data)[rsq < extent**2])
        plt.imshow(
            data, origin="lower", extent=plt_extent, vmin=-1 * vminmax, vmax=vminmax
        )
    plt.colorbar()
    plt.grid()
    plt.xlabel(f"Xi ({units})")
    plt.ylabel(f"Eta ({units})")
    plt.title(f"{title}{label.replace('_', ' ')}")

    plt.xlim((cent[0] - extent, cent[0] + extent))
    plt.ylim((cent[1] - extent, cent[1] + extent))
    plt.savefig(os.path.join(plot_dir, f"{title.replace(' ', '_')}_map{label}.png"))

    # Profile
    plt.close()
    x = np.linspace(0, pixsize * len(rprof), len(rprof))
    plt.plot(x, rprof)
    plt.xlabel(f"Radius ({units})")
    plt.title("\n".join(textwrap.wrap(f"{title}{label.replace('_', ' ')}")))
    plt.xlim((0, extent))
    plt.savefig(
        os.path.join(plot_dir, f"{title.replace(' ', '_')}_prof{label}.png"),
        bbox_inches="tight",
    )

plot_map_complete(data, posmap, pixsize, extent, cent, plot_dir, title, comps='TQU', log_thresh=0.001, append='', units='"', lognorm=1.0, qrur=False)

Wrapper for plot_map that plots all comps in both linear and log scales.

Parameters:

Name	Type	Description	Default
data	Float[ndarray, 'nx ny'] | ndmap	The beam map to plot.	required
posmap	Float[ndarray, '2 nx ny'] | ndmap	The position map for the beam map. Should be generated by enmap.posmap.	required
extent	float	The radius of the plot to make in the same units as posmap.	required
cent	tuple[float, float]	The center of the beam map in the same units as posmap.	required
plot_dir	str	The directory to save plots to.	required
title	str	The base title of the plot. The full title will have the format "{title} {comp} {append} {'log10'*log}". The output file will use the full title as the filename but with whitespace replaced by '_' and 'map' or 'prof' inserted between "{title}" and "{comp}" depending on which plot it is.	required
comps	str	The comps of the map we are plotting. Should be "T" or "TQU".	"TQU"
log_thresh	float	The threshold at which the log plot transitions to linear to avoid issues near values of 0.	1e-3
append	str	String to append ti title and filename. See title for details.	""
units	str	The units of posmap to appear in axes labels.	'"'
lognorm	float	Value to normalize map by before plotting in log scale. All comps are normalized by this factor. This is useful for seeing the polarization of the beam relative to the peak of the temperature beam.	1
qrur	bool	If True then also plot the radial QU maps. This will raise an error if Q and U are not in comps.	False

Source code in lat_beams/plotting.py

def plot_map_complete(
    data: Float[np.ndarray, "nx ny"] | enmap.ndmap,
    posmap: Float[np.ndarray, "2 nx ny"] | enmap.ndmap,
    pixsize: float,
    extent: float,
    cent: tuple[float, float],
    plot_dir: str,
    title: str,
    comps: str = "TQU",
    log_thresh: float = 1e-3,
    append: str = "",
    units: str = '"',
    lognorm: float = 1.0,
    qrur: bool = False,
):
    """
    Wrapper for `plot_map` that plots all comps in both linear and log scales.

    Parameters
    ----------
    data : Float[np.ndarray, "nx ny"] | enmap.ndmap
        The beam map to plot.
    posmap : Float[np.ndarray, "2 nx ny"] | enmap.ndmap
        The position map for the beam map.
        Should be generated by `enmap.posmap`.
    extent : float
        The radius of the plot to make in the same units as `posmap`.
    cent : tuple[float, float]
        The center of the beam map in the same units as `posmap`.
    plot_dir : str
        The directory to save plots to.
    title : str
        The base title of the plot.
        The full title will have the format "{title} {comp} {append} {'log10'*log}".
        The output file will use the full title as the filename but with whitespace replaced by '_'
        and 'map' or 'prof' inserted between "{title}" and "{comp}" depending on which plot it is.
    comps : str, default: "TQU"
        The comps of the map we are plotting.
        Should be "T" or "TQU".
    log_thresh : float, default: 1e-3
        The threshold at which the log plot transitions to linear to avoid issues near values of 0.
    append : str, default: ""
        String to append ti title and filename.
        See `title` for details.
    units : str, default: '"'
        The units of `posmap` to appear in axes labels.
    lognorm : float, default: 1
        Value to normalize map by before plotting in log scale.
        All comps are normalized by this factor.
        This is useful for seeing the polarization of the beam relative
        to the peak of the temperature beam.
    qrur : bool, default: False
        If True then also plot the radial QU maps.
        This will raise an error if Q and U are not in `comps`.
    """
    if len(data.shape) == 2:
        data = data[None, ...]
    for i, comp in enumerate(comps):
        for log in (False, True):
            plot_map(
                data[i] * (lognorm * log + (not log)),
                posmap,
                pixsize,
                extent,
                cent,
                plot_dir,
                title,
                comp=comp,
                log=log,
                log_thresh=log_thresh,
                append=append,
                units=units,
            )
    if not qrur:
        return
    if "Q" not in comps or "U" not in comps:
        raise ValueError("Cannot plot Qr and Ur without Q and U")
    Q = data[comps.find("Q")]
    U = data[comps.find("U")]
    eta, xi = posmap
    theta = np.arctan2(eta, xi)
    Q_r = Q * np.cos(2 * theta) + U * np.sin(2 * theta)
    U_r = U * np.cos(2 * theta) - Q * np.sin(2 * theta)
    for dat, comp in [(Q_r, "Qr"), (U_r, "Ur")]:
        for log in (False, True):
            plot_map(
                dat * (lognorm * log + (not log)),
                posmap,
                pixsize,
                extent,
                cent,
                plot_dir,
                title,
                comp=comp,
                log=log,
                log_thresh=log_thresh,
                append=append,
                units=units,
            )

plot_tod(aman, sig_filt, tod_plot_dir, file_label, max_dets)

Plot a TOD and its filtered signal.

Parameters:

Name	Type	Description	Default
aman	AxisManager	A TOD containing the unfiltered signal.	required
sig_filt	Float[ndarray, 'ndet nsamp']	The filtered signal.	required
tod_plot_dir	str	The directory to save plots to.	required
file_label	str	String to determine the file name. For the unfiltered signal the file is "{file_label}_tod.png". For the filtered signal the file is "{file_label}_tod_filt.png"	required
max_dets	str	The maximum number of detectors to plot. If ndets is greated than max_dets then the first max_dets detectors are plotted.	required

Source code in lat_beams/plotting.py

def plot_tod(
    aman: AxisManager,
    sig_filt: Float[np.ndarray, "ndet nsamp"],
    tod_plot_dir: str,
    file_label: str,
    max_dets: int,
):
    """
    Plot a TOD and its filtered signal.

    Parameters
    ----------
    aman : AxisManager
        A TOD containing the unfiltered signal.
    sig_filt : Float[np.ndarray, "ndet nsamp"]
        The filtered signal.
    tod_plot_dir : str
        The directory to save plots to.
    file_label : str
        String to determine the file name.
        For the unfiltered signal the file is "{file_label}_tod.png".
        For the filtered signal the file is "{file_label}_tod_filt.png"
    max_dets : str
        The maximum number of detectors to plot.
        If `ndets` is greated than `max_dets` then the first `max_dets`
        detectors are plotted.
    """
    plt.close()
    for data, append, ylabel in [
        (np.array(aman.signal)[:max_dets], "tod", "Signal (pW)"),
        (sig_filt[:max_dets], "tod_filt", "Filtered Signal (pW)"),
    ]:
        fig, ax = plt.subplots()
        nsamp = data.shape[1]
        x = np.arange(nsamp)
        segments = [np.column_stack([x, y]) for y in data]

        lc = LineCollection(segments, alpha=0.3)
        ax.add_collection(lc)
        ax.autoscale(enable=True, axis="both")
        ax.set_xlabel("Samples")
        ax.set_ylabel(ylabel)
        ax.set_title(file_label.replace("_", " "))

        fig.savefig(os.path.join(tod_plot_dir, f"{file_label}_{append}.png"))
        plt.close(fig)