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, combine=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
combine	list[tuple[str, str]]	Groups matching values of a field into non-overlapping combinations.	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,
    combine: Optional[list[tuple[str, 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`.
    combine : list[tuple[str, str]], optional
        Groups matching values of a field into non-overlapping combinations.
    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`.
    """
    data = data.copy()
    ignore = [] if ignore is None else ignore.copy()
    merge = [] if merge is None else merge
    combine = [] if combine is None else combine

    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]]))
            ]
        )

    for independent, to_combine in combine:
        columns = [independent, to_combine]
        missing = [col for col in columns if col not in data]
        if missing:
            raise ValueError(f"Columns not found in data: {missing}")

        group = np.asarray(data[independent])
        values = np.asarray(data[to_combine])
        result = np.asarray(group, dtype=object).copy()
        signatures = {}
        for g in np.unique(group):
            mask = group == g
            signatures[g] = tuple(sorted(map(str, values[mask])))
        signatures = {str(k): np.unique(v) for k, v in signatures.items()}
        unique_values = list({tuple(x) for x in signatures.values()})
        unique_values = [list(x) for x in unique_values]

        max_n = max(len(v) for v in unique_values)

        for i in range(max_n):
            uvs = [u[i] for u in unique_values if len(u) > i]
            label = "+".join(uvs)
            msk = np.isin(values, uvs)
            result[msk] = label

        max_n = np.max([len(l) for l in unique_values])
        for i in range(max_n):
            uvs = [u[i] for u in unique_values if len(u) > i]
            label = "+".join(uvs)
            msk = np.isin(values, uvs)
            result[msk] = label
        data[to_combine] = result

    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])))
        fields = fields[::-1]

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

        if len(fields) > len(can_add):
            for i, cat in enumerate(cats):
                j = len(can_add) + i
                if j >= len(fields):
                    continue
                to_combine = []
                if cat in kwargs and kwargs[cat] is not None:
                    to_combine = [kwargs[cat]]
                to_combine += fields[j:][:: len(cats)]
                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[cat] = 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_debug_map(outdir, msg, fjobstr, imap, ivar, hot, ps2d, lmap, lmin, lmax, wn, cn)

Generate and save a diagnostic plot of a beam map and noise estimates.

The figure contains the input map with hot-pixel regions overlaid, the two-dimensional power spectrum when available, and the inverse- variance map. The figure title includes the ratio of correlated to white noise and the logarithms of the estimated white and correlated noise amplitudes.

Parameters:

Name	Type	Description	Default
outdir	str	Directory in which to save the diagnostic plot.	required
msg	str	Additional message to include in the figure title.	required
fjobstr	str	Identifier used for the plot title and output filename.	required
imap	Float[ndmap, '... ny nx']	Input map to display.	required
ivar	Float[ndmap, '... ny nx']	Inverse-variance map corresponding to imap.	required
hot	Bool[ndarray, 'ny nx']	Boolean mask identifying hot pixels in imap. These regions are overlaid on the map.	required
ps2d	Optional[Float[ndmap, 'ny nx']]	Two-dimensional power spectrum to display. If None, the power spectrum panel is left empty.	required
lmap	Optional[Float[ndmap, 'ny nx']]	Two-dimensional multipole coordinate maps corresponding to ps2d. The first element contains the y-direction multipoles and the second contains the x-direction multipoles. If None, the power spectrum is not plotted.	required
lmin	float	Minimum multipole used to define the white-noise region in the power spectrum.	required
lmax	float	Maximum multipole used to define the correlated-noise region in the power spectrum.	required
wn	float	Estimated white-noise amplitude.	required
cn	float	Estimated correlated-noise amplitude.	required

Source code in lat_beams/plotting.py

def plot_debug_map(
    outdir: str,
    msg: str,
    fjobstr: str,
    imap: Float[enmap.ndmap, "... ny nx"],
    ivar: Float[enmap.ndmap, "... ny nx"],
    hot: Bool[np.ndarray, "ny nx"],
    ps2d: Optional[Float[enmap.ndmap, "ny nx"]],
    lmap: Optional[Float[enmap.ndmap, "2 ny nx"]],
    lmin: float,
    lmax: float,
    wn: float,
    cn: float,
) -> None:
    """
    Generate and save a diagnostic plot of a beam map and noise estimates.

    The figure contains the input map with hot-pixel regions overlaid,
    the two-dimensional power spectrum when available, and the inverse-
    variance map. The figure title includes the ratio of correlated to
    white noise and the logarithms of the estimated white and correlated
    noise amplitudes.

    Parameters
    ----------
    outdir : str
        Directory in which to save the diagnostic plot.
    msg : str
        Additional message to include in the figure title.
    fjobstr : str
        Identifier used for the plot title and output filename.
    imap : Float[enmap.ndmap, "... ny nx"]
        Input map to display.
    ivar : Float[enmap.ndmap, "... ny nx"]
        Inverse-variance map corresponding to `imap`.
    hot : Bool[np.ndarray, "ny nx"]
        Boolean mask identifying hot pixels in `imap`. These regions are
        overlaid on the map.
    ps2d : Optional[Float[enmap.ndmap, "ny nx"]]
        Two-dimensional power spectrum to display. If `None`, the power
        spectrum panel is left empty.
    lmap : Optional[Float[enmap.ndmap, "ny nx"]]
        Two-dimensional multipole coordinate maps corresponding to `ps2d`.
        The first element contains the y-direction multipoles and the
        second contains the x-direction multipoles. If `None`, the power
        spectrum is not plotted.
    lmin : float
        Minimum multipole used to define the white-noise region in the
        power spectrum.
    lmax : float
        Maximum multipole used to define the correlated-noise region in
        the power spectrum.
    wn : float
        Estimated white-noise amplitude.
    cn : float
        Estimated correlated-noise amplitude.
    """
    ratio = cn / wn
    plt.close("all")
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    posmap = np.rad2deg(imap.posmap()) * 3600
    _norm = SymLogNorm(linthresh=1e-4, clip=True, vmin=-1, vmax=1)
    plt_extent = (posmap[1].max(), posmap[1].min(), posmap[0].min(), posmap[0].max())
    im0 = axes[0].imshow(imap[0], origin="lower", extent=plt_extent, norm=_norm)
    axes[0].set_title(f"Map")
    fig.colorbar(im0, ax=axes[0])
    hp = hot.astype(float)
    hp[~hot] = np.nan
    axes[0].imshow(
        hp, origin="lower", extent=plt_extent, cmap="Greys", alpha=0.3, vmin=0, vmax=1
    )
    if ps2d is not None and lmap is not None:
        im1 = axes[1].pcolormesh(
            enmap.fftshift(lmap[1]),
            enmap.fftshift(lmap[0]),
            enmap.fftshift(np.log10(ps2d)),
            shading="auto",
        )
        axes[1].axvline(lmin, color="r", ls="--")
        axes[1].axvline(-lmin, color="r", ls="--")
        axes[1].axhline(lmax, color="b", ls="--")
        axes[1].axhline(-lmax, color="b", ls="--")
        fig.colorbar(im1, ax=axes[1])
    axes[1].set_title("2D Power Spectrum")
    im2 = axes[2].imshow(ivar[0], origin="lower", extent=plt_extent, cmap="Purples")
    axes[2].set_title("Map (ivar)")
    fig.colorbar(im2, ax=axes[2])
    plt.suptitle(
        f"{fjobstr} ({np.round(ratio, 2)}, {np.round(np.log10(wn), 1)}, {np.round(np.log10(cn), 1)})\n{msg}"
    )
    plt.tight_layout()
    plt.savefig(os.path.join(outdir, f"{fjobstr}.png"), dpi=150)
    plt.close(fig)

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_list, posmap, pixsize, extent, cent, plot_dir, title, titles, log=False, log_thresh=0.001, append='', units='"')

Plot multiple beam maps and their radial profiles in a single figure.

Each column corresponds to one map. The top row contains the map image, and the bottom row contains its radial profile.

Parameters:

Name	Type	Description	Default
data_list	list[Float[ndmap, 'nx ny']]	List of beam maps, all with the same geometry.	required
posmap	Float[ndmap, '2 nx ny']	Position map generated by enmap.posmap, with shape (2, nx, ny).	required
pixsize	float	Pixel size in the same units as posmap.	required
extent	float	Radius of the plot in the same units as posmap.	required
cent	tuple[float, float]	Center of the maps, in (Xi, Eta) coordinates.	required
plot_dir	str	Directory in which to save the figure.	required
title	str	Title for the figure. Also used to determint filename.	required
titles	list[str]	Title for each map, one per column.	required
log	bool	If True, use a symmetric logarithmic color normalization and plot the signed log10 of the radial profiles.	False
log_thresh	float	Linear threshold for the symmetric logarithmic color normalization.	1e-3
append	str	Optional string appended to the plot title and filename.	""
units	str	Units used in the position and radius axis labels.	'"'

Source code in lat_beams/plotting.py

def plot_map(
    data_list: list[Float[enmap.ndmap, "nx ny"]],
    posmap: Float[enmap.ndmap, "2 nx ny"],
    pixsize: float,
    extent: float,
    cent: tuple[float, float],
    plot_dir: str,
    title: str,
    titles: list[str],
    log: bool = False,
    log_thresh: float = 1e-3,
    append: str = "",
    units: str = '"',
):
    """
    Plot multiple beam maps and their radial profiles in a single figure.

    Each column corresponds to one map. The top row contains the map image,
    and the bottom row contains its radial profile.

    Parameters
    ----------
    data_list : list[Float[enmap.ndmap, "nx ny"]]
        List of beam maps, all with the same geometry.
    posmap : Float[enmap.ndmap, "2 nx ny"]
        Position map generated by enmap.posmap, with shape (2, nx, ny).
    pixsize : float
        Pixel size in the same units as posmap.
    extent : float
        Radius of the plot in the same units as posmap.
    cent : tuple[float, float]
        Center of the maps, in (Xi, Eta) coordinates.
    plot_dir : str
        Directory in which to save the figure.
    title : str
        Title for the figure. Also used to determint filename.
    titles : list[str]
        Title for each map, one per column.
    log : bool, default=False
        If True, use a symmetric logarithmic color normalization and
        plot the signed log10 of the radial profiles.
    log_thresh : float, default=1e-3
        Linear threshold for the symmetric logarithmic color normalization.
    append : str, default=""
        Optional string appended to the plot title and filename.
    units : str, default='"'
        Units used in the position and radius axis labels.
    """
    if len(data_list) != len(titles):
        raise ValueError("data_list and titles must have the same length.")
    if len(data_list) == 0:
        raise ValueError("data_list must contain at least one map.")
    os.makedirs(plot_dir, exist_ok=True)

    # Restrict all maps to the same local patch.
    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)
    if w[0].size == 0:
        raise ValueError("The requested extent does not overlap the position map.")
    ext_sl = (
        slice(np.min(w[0]), np.max(w[0]) + 1),
        slice(np.min(w[1]), np.max(w[1]) + 1),
    )
    pos_xi = posmap[1][ext_sl]
    pos_eta = posmap[0][ext_sl]
    # posmap_crop = (pos_eta, pos_xi)
    # Coordinates for imshow.
    plt_extent = (
        pos_xi.max(),
        pos_xi.min(),
        pos_eta.min(),
        pos_eta.max(),
    )
    rsq = (pos_eta - cent[1]) ** 2 + (pos_xi - cent[0]) ** 2
    rcent = np.unravel_index(np.argmin(rsq), rsq.shape)

    # Radial profiles and cropped maps.
    maps = [data[ext_sl] for data in data_list]
    profiles = []
    for data in maps:
        rprof = radial_profile(data, (int(rcent[1]), int(rcent[0])))[
            : int(0.5 * min(*data.shape))
        ]
        if log:
            with np.errstate(divide="ignore", invalid="ignore"):
                rprof = np.sign(rprof) * np.log10(np.abs(rprof))
        profiles.append(rprof)
    radii = [
        np.linspace(0, pixsize * len(profile), len(profile)) for profile in profiles
    ]

    n_maps = len(data_list)
    fig = plt.figure(figsize=(5 * n_maps, 8), constrained_layout=True)

    gs = fig.add_gridspec(
        2,
        2 * n_maps,
        height_ratios=[1, 1],
        width_ratios=[width for _ in range(n_maps) for width in (1, 0.05)],
    )

    axes_map = []
    axes_prof = []

    for i in range(n_maps):
        map_col = 2 * i
        cbar_col = map_col + 1

        ax_map = fig.add_subplot(gs[0, map_col])
        ax_cbar = fig.add_subplot(gs[0, cbar_col])
        ax_prof = fig.add_subplot(gs[1, map_col])

        # Reserve the colorbar column in the profile row too.
        fig.add_subplot(gs[1, cbar_col]).set_axis_off()

        axes_map.append((ax_map, ax_cbar))
        axes_prof.append(ax_prof)

    for i, (data, rprof, x, tit) in enumerate(zip(maps, profiles, radii, titles)):
        ax_map, ax_cbar = axes_map[i]
        ax_prof = axes_prof[i]

        if log:
            norm = SymLogNorm(
                linthresh=log_thresh,
                clip=True,
                vmin=-1,
                vmax=1,
            )
            im = ax_map.imshow(
                data,
                origin="lower",
                extent=plt_extent,
                norm=norm,
                cmap="coolwarm",
            )
        else:
            vminmax = np.max(np.abs(data)[rsq < extent**2])
            im = ax_map.imshow(
                data,
                origin="lower",
                extent=plt_extent,
                vmin=-vminmax,
                vmax=vminmax,
                cmap="coolwarm",
            )

        fig.colorbar(im, cax=ax_cbar)
        ax_map.grid()
        ax_map.set_xlabel(f"Xi ({units})")
        ax_map.set_ylabel(f"Eta ({units})")
        ax_map.set_title(f"{tit} map")
        ax_map.set_xlim(cent[0] - extent, cent[0] + extent)
        ax_map.set_ylim(cent[1] - extent, cent[1] + extent)

        if log:
            pos_prof = np.ma.masked_where(rprof <= 0, rprof)
            neg_prof = np.ma.masked_where(rprof > 0, rprof)
            ax_prof.plot(x, -1 * np.abs(rprof), color="grey", linestyle="--", alpha=0.5)
            ax_prof.scatter(x, -1 * pos_prof, color="b", marker=".")
            ax_prof.scatter(x, neg_prof, color="r", marker=".")
        else:
            ax_prof.plot(x, rprof, color="grey", alpha=0.5)
            ax_prof.scatter(x, rprof, color="grey", marker=".")
        ax_prof.set_xlabel(f"Radius ({units})")
        ax_prof.set_ylabel("log10(Profile)" if log else "Profile")
        ax_prof.set_title(f"{tit} profile")
        ax_prof.set_xlim(0, extent)
        ax_prof.grid()
    if append:
        title += f" {append}"
    if log:
        title += " log10"
    fig.suptitle("\n".join(textwrap.wrap(title, 50)))

    filename = title.replace(" ", "_") + ".png"
    fig.savefig(
        os.path.join(plot_dir, filename),
        bbox_inches="tight",
        dpi=150,
    )
    plt.close(fig)

plot_map_tqu(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[ndmap, 'ncomp nx ny']	The beam map to plot.	required
posmap	Float[ndmap, '2 nx ny']	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} {append} {'log10'*log}". The output file will use the full title as the filename but with whitespace replaced by '_'	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_tqu(
    data: Float[enmap.ndmap, "ncomp nx ny"],
    posmap: Float[enmap.ndmap, "2 nx ny"],
    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[enmap.ndmap, "ncomp nx ny"]
        The beam map to plot.
    posmap : Float[enmap.ndmap, "2 nx ny"]
        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} {append} {'log10'*log}".
        The output file will use the full title as the filename but with whitespace replaced by '_'
    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:
        maps = [data]
    else:
        maps = [data[i] for i in range(len(data))]
    Q_r, U_r = None, None
    if qrur:
        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 log in (False, True):
        plot_map(
            maps if not log else [m * lognorm for m in maps],  # type: ignore
            posmap,
            pixsize,
            extent,
            cent,
            plot_dir,
            title,
            [comps[i] for i in range(len(comps))],
            log=log,
            log_thresh=log_thresh,
            append=append,
            units=units,
        )
        if not qrur or Q_r is None or U_r is None:
            continue
        plot_map(
            [Q_r, U_r] if not log else [Q_r * lognorm, U_r * lognorm],
            posmap,
            pixsize,
            extent,
            cent,
            plot_dir,
            title + " QrUr",
            ["Q_r", "U_r"],
            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)