Power Spectrum Plotting module

pk

Power spectrum and posterior visualization functions.

This module provides utilities for plotting primordial power spectrum data, including posteriors, priors, and residuals from MCMC chains.

posteriors_delta(k, samples, colors=None, Nbins=20, kmin=0.001, kmax=0.34, figsize=(6, 8), chain_entries=None, **kwargs)

Plot posteriors for deviations only: \(1 + \delta(k)\).

PARAMETER	DESCRIPTION
k	Array of \(k\) values TYPE: ndarray
samples	Dict mapping labels to lists of samples (each sample is \(1 + \delta(k)\)) TYPE: dict
colors	List of colors for each dataset (optional, required if chain_entries not provided) TYPE: list DEFAULT: None
Nbins	Number of bins for the node visualization (default: 20) TYPE: int DEFAULT: 20
kmin	Minimum k value (default: 1e-3) TYPE: float DEFAULT: 0.001
kmax	Maximum k value (default: 0.34) TYPE: float DEFAULT: 0.34
chain_entries	Optional ChainsCollection or list of ChainEntry objects for auto color extraction TYPE: Optional[List] DEFAULT: None
kwargs	Additional keyword arguments for add_nodes DEFAULT: {}

RETURNS	DESCRIPTION
Figure	Matplotlib Figure object.

Source code in src/primefeat/plots/pk.py

def posteriors_delta(
    k: np.ndarray,
    samples: dict,
    colors: list = None,
    Nbins: int = 20,
    kmin: float = 1e-3,
    kmax: float = 0.34,
    figsize: tuple = (6, 8),
    chain_entries: Optional[List] = None,
    **kwargs,
) -> plt.Figure:
    """
    Plot posteriors for deviations only: $1 + \\delta(k)$.

    Args:
        k: Array of $k$ values
        samples: Dict mapping labels to lists of samples (each sample is $1 + \\delta(k)$)
        colors: List of colors for each dataset (optional, required if chain_entries not provided)
        Nbins: Number of bins for the node visualization (default: 20)
        kmin: Minimum k value (default: 1e-3)
        kmax: Maximum k value (default: 0.34)
        chain_entries: Optional ChainsCollection or list of ChainEntry objects for auto color extraction
        kwargs: Additional keyword arguments for `add_nodes`

    Returns:
        Matplotlib Figure object.
    """
    # Extract colors from chain_entries if provided
    if chain_entries is not None:
        colors = [entry.color for entry in chain_entries]

    if colors is None:
        raise ValueError("Either colors or chain_entries must be provided")

    fig, axs = plt.subplots(
        len(samples),
        sharex=True,
        figsize=figsize,
        tight_layout=True,
        gridspec_kw={"hspace": 0.05},
    )
    # Ensure axs is always iterable (even for single subplot)
    if len(samples) == 1:
        axs = [axs]
    for ax, color, (lbl, Pk) in zip(axs, colors, samples.items()):
        ax.semilogx([], [], color=color, label=lbl)
        ax = plot_fill_between(k, Pk, ax=ax, color=color, alpha_contour=0.5)
        ax.axhline(1.0, color="k", linestyle="--", alpha=0.7)
        # add_nodes(get_bin_centers(kmin, kmax, Nbins), chain, ax=ax, **kwargs)
        ax.legend(frameon=False)
    return fig

posteriors_PPS(k, samples=None, colors=None, mode='delta', ax=None, show_binning_range=False, k_start=None, k_end=None, alpha_contour=0.5, add_inset=False, inset_klim=(0.01, 0.2), inset_bbox=(0.55, 0.55, 0.4, 0.4), inset_ylim=None, sigma_levels=2, show_mean=False, normalize=1000000000.0, add_2nd_xaxis=True, fig_kw=None, chain_entries=None)

Plot posteriors for primordial power spectrum.

PARAMETER	DESCRIPTION
k	Array of \(k\) values TYPE: ndarray
samples	Dict mapping labels to lists of samples. If None and chain_entries is provided, samples will be extracted from the entries. TYPE: dict DEFAULT: None
colors	List of colors for each dataset (optional). If not provided and chain_entries is None, will use registered colors from get_chains(), or auto-generated defaults. If chain_entries is provided, colors will be extracted from the entries. Supported formats: hex codes ('#2E86AB'), named colors ('red', 'blue'), or RGB tuples. TYPE: Optional[List[str]] DEFAULT: None
mode	"delta" for \(1+\delta(k)\) or "full" for full \(P(k) = A_s \cdot k^{(n_s-1)} \cdot [1+\delta(k)]\) TYPE: str DEFAULT: 'delta'
ax	Existing matplotlib axes (optional) TYPE: Optional[Axes] DEFAULT: None
show_binning_range	If True, show gray shaded region for binning range TYPE: bool DEFAULT: False
k_start	Start of binning range (only used if show_binning_range=True) TYPE: float DEFAULT: None
k_end	End of binning range (only used if show_binning_range=True) TYPE: float DEFAULT: None
alpha_contour	Transparency for confidence bands (default: 0.5) TYPE: float DEFAULT: 0.5
add_inset	If True, add an inset plot with zoomed view TYPE: bool DEFAULT: False
inset_klim	Tuple of (kmin, kmax) for inset zoom range TYPE: tuple DEFAULT: (0.01, 0.2)
inset_bbox	Tuple of (x, y, width, height) for inset position in axes coordinates TYPE: tuple DEFAULT: (0.55, 0.55, 0.4, 0.4)
inset_ylim	Tuple of (ymin, ymax) for inset y-axis limits (optional) TYPE: tuple DEFAULT: None
sigma_levels	Number of \(\sigma\) levels to show in contours (default: 2) TYPE: int DEFAULT: 2
show_mean	If True, show mean instead of median (default: False) TYPE: bool DEFAULT: False
normalize	Normalization factor (default: 1e9) TYPE: float DEFAULT: 1000000000.0
fig_kw	Additional keyword arguments for plt.subplots (e.g., figsize) TYPE: Optional[dict] DEFAULT: None
chain_entries	Optional ChainsCollection or list of ChainEntry objects. When provided, colors and labels are extracted automatically. For backward compatibility, this can be omitted and samples + colors passed explicitly. TYPE: Optional[List] DEFAULT: None

RETURNS	DESCRIPTION
Figure	Matplotlib Figure object.

Source code in src/primefeat/plots/pk.py

def posteriors_PPS(
    k: np.ndarray,
    samples: dict = None,
    colors: Optional[List[str]] = None,
    mode: str = "delta",
    ax: Optional[plt.Axes] = None,
    show_binning_range: bool = False,
    k_start: float = None,
    k_end: float = None,
    alpha_contour: float = 0.5,
    add_inset: bool = False,
    inset_klim: tuple = (1e-2, 0.2),
    inset_bbox: tuple = (0.55, 0.55, 0.4, 0.4),
    inset_ylim: tuple = None,
    sigma_levels: int = 2,
    show_mean: bool = False,
    normalize: float = 1e9,
    add_2nd_xaxis: bool = True,
    fig_kw: Optional[dict] = None,
    chain_entries: Optional[List] = None,
) -> plt.Figure:
    """
    Plot posteriors for primordial power spectrum.

    Args:
        k: Array of $k$ values
        samples: Dict mapping labels to lists of samples. If None and chain_entries
                is provided, samples will be extracted from the entries.
        colors: List of colors for each dataset (optional). If not provided and
                chain_entries is None, will use registered colors from get_chains(),
                or auto-generated defaults. If chain_entries is provided, colors
                will be extracted from the entries. Supported formats: hex codes
                ('#2E86AB'), named colors ('red', 'blue'), or RGB tuples.
        mode: "delta" for $1+\\delta(k)$ or "full" for full $P(k) = A_s \\cdot k^{(n_s-1)} \\cdot [1+\\delta(k)]$
        ax: Existing matplotlib axes (optional)
        show_binning_range: If True, show gray shaded region for binning range
        k_start: Start of binning range (only used if show_binning_range=True)
        k_end: End of binning range (only used if show_binning_range=True)
        alpha_contour: Transparency for confidence bands (default: 0.5)
        add_inset: If True, add an inset plot with zoomed view
        inset_klim: Tuple of (kmin, kmax) for inset zoom range
        inset_bbox: Tuple of (x, y, width, height) for inset position in axes coordinates
        inset_ylim: Tuple of (ymin, ymax) for inset y-axis limits (optional)
        sigma_levels: Number of $\\sigma$ levels to show in contours (default: 2)
        show_mean: If True, show mean instead of median (default: False)
        normalize: Normalization factor (default: 1e9)
        fig_kw: Additional keyword arguments for plt.subplots (e.g., figsize)
        chain_entries: Optional ChainsCollection or list of ChainEntry objects.
                      When provided, colors and labels are extracted automatically.
                      For backward compatibility, this can be omitted and samples
                      + colors passed explicitly.

    Returns:
        Matplotlib Figure object.
    """
    # Handle chain_entries if provided
    if chain_entries is not None:
        # Extract colors and labels from chain_entries
        colors = [entry.color for entry in chain_entries]
        # Note: samples dict should already be provided by the caller
        # (typically via collection.compute_pk_samples(k))
        if samples is None:
            raise ValueError(
                "chain_entries provided but samples dict is None. "
                "Please pass the samples dict computed from the collection."
            )
    else:
        # Resolve colors: explicit arg > registered > defaults
        if colors is None:
            # Try to get registered colors from get_chains()
            colors = [get_chain_color(label) for label in samples.keys()]

            # Fill in missing colors with defaults
            if None in colors:
                n_chains = len(samples)
                default_palette = get_default_color_palette(n_chains)
                colors = [
                    c if c else default
                    for c, default in zip(colors, default_palette)
                ]

    # Set y-label based on mode
    if mode == "full":
        ylabel = r"$10^9~\mathcal{P}_\zeta(k)$"
    else:
        ylabel = r"$1+\delta\mathcal{P}(k)$"

    # Create canvas with dual k/ell axes
    fig, ax, _ = create_pk_canvas(
        ax=ax,
        fig_kw=fig_kw,
        primary_axis="k",
        ylabel=ylabel,
        xlim=(k.min(), k.max()),
        add_2nd_xaxis=add_2nd_xaxis,
    )

    # Show binning range if requested (inverted shading: gray outside, white inside)
    if show_binning_range and k_start is not None and k_end is not None:
        # Shade region before k_start
        ax.axvspan(k.min(), k_start, alpha=0.15, color="gray", zorder=0)
        # Shade region after k_end
        ax.axvspan(
            k_end,
            k.max(),
            alpha=0.15,
            color="gray",
            zorder=0,
            label="Extrapolation region",
        )

    # Plot each dataset using plot_fill_between
    for i, (label, sample_list) in enumerate(samples.items()):
        sample_array = np.array(sample_list)
        color = colors[i]

        # Use the existing plot_fill_between function
        ax.plot([], [], color=color, label=label)  # For legend
        ax = plot_fill_between(
            k,
            normalize * sample_array,
            ax=ax,
            color=color,
            alpha_contour=alpha_contour,
            quantiles=_SIGMA_LEVELS_TO_QUANTILES[sigma_levels],
        )

    ax.legend(fontsize=10, loc="best")

    # Add reference line for delta mode
    if mode != "full":
        ax.axhline(1.0, color="k", linestyle="--", alpha=0.7, lw=1)

    # Add inset plot if requested
    if add_inset:
        # Create inset axes
        axins = ax.inset_axes(inset_bbox)

        # Filter k values within inset range
        k_mask = (k >= inset_klim[0]) & (k <= inset_klim[1])
        k_inset = k[k_mask]

        # Plot each dataset in the inset
        for i, (label, sample_list) in enumerate(samples.items()):
            sample_array = np.array(sample_list)
            # Filter samples to inset range
            sample_inset = sample_array[:, k_mask]
            color = colors[i]

            # Use plot_fill_between for inset
            axins = plot_fill_between(
                k_inset,
                sample_inset,
                ax=axins,
                color=color,
                alpha_contour=alpha_contour,
            )

        # Configure inset axes
        axins.set_xscale("log")
        axins.set_xlim(inset_klim[0], inset_klim[1])
        axins.axhline(0.0, color="k", linestyle="--", alpha=0.7, lw=1)

        if inset_ylim is not None:
            axins.set_ylim(inset_ylim[0], inset_ylim[1])

        axins.grid(True, alpha=0.3)
        axins.tick_params(labelsize=8)

        # Store inset axes as a figure attribute for easy access
        fig.inset_axes = axins

    plt.tight_layout()
    return fig

priors_PPS(k, delta_prior, As_prior=None, ns_prior=None, n_samples=10000, k_pivot=0.05, sigma_levels=2, alpha_contour=0.5, normalize=1000000000.0, color='steelblue', label='Prior', ax=None, fig_kw=None, **line_kwargs)

Plot the prior distribution of the primordial power spectrum P(k).

Randomly draws n_samples parameter sets from the user-specified priors, evaluates P(k) for each draw, then displays the resulting distribution as sigma-level bands (identical style to :func:posteriors_PPS).

Two modes are supported automatically:

• Delta mode (default) — plots :math:1 + \delta(k) when neither As_prior nor ns_prior is supplied.
• Full mode — plots :math:10^9 A_s (k/k_{\mathrm{pivot}})^{n_s-1} [1+\delta(k)] when at least one of As_prior / ns_prior is provided. Missing parameters are held fixed at their fiducial values (:math:A_s = 2.1 \times 10^{-9}, :math:n_s = 0.965).

Prior specifications

A prior spec is either:

• (low, high) — draws from :math:\mathcal{U}(\text{low}, \text{high}).
• ('normal', mean, std) — draws from :math:\mathcal{N}(\text{mean}, \text{std}^2).

For delta_prior you may supply a single spec applied to all k-bins, or a list of specs with one entry per bin.

PARAMETER	DESCRIPTION
k	Array of :math:k values (Mpc :sup:-1). TYPE: ndarray
delta_prior	Prior specification for the :math:\delta(k) bins. Either a single spec (low, high) / ('normal', mean, std) applied to every bin, or a list of such specs (one per bin). TYPE: Union[Tuple, List[Tuple]]
As_prior	Optional prior for :math:A_s in units of :math:10^{-9}. E.g. (1.5, 2.5) for a uniform prior. If None and full mode is triggered via ns_prior, the fiducial value :math:A_s = 2.1 \times 10^{-9} is used. TYPE: Optional[Tuple] DEFAULT: None
ns_prior	Optional prior for :math:n_s. E.g. (0.9, 1.05) or ('normal', 0.965, 0.01). If None and full mode is triggered via As_prior, the fiducial value :math:n_s = 0.965 is used. TYPE: Optional[Tuple] DEFAULT: None
n_samples	Number of random draws from the prior (default: 10 000). TYPE: int DEFAULT: 10000
k_pivot	Pivot scale in Mpc :sup:-1 (default: 0.05). TYPE: float DEFAULT: 0.05
sigma_levels	Number of :math:\sigma contours to display — 1, 2, or 3 (default: 2). TYPE: int DEFAULT: 2
alpha_contour	Transparency of the confidence bands (default: 0.5). TYPE: float DEFAULT: 0.5
normalize	Overall normalization factor (default: :math:10^9). Applied as a multiplicative prefactor to the plotted quantity. TYPE: float DEFAULT: 1000000000.0
color	Color for the prior bands and median line (default: 'steelblue'). TYPE: str DEFAULT: 'steelblue'
label	Legend label (default: 'Prior'). TYPE: Optional[str] DEFAULT: 'Prior'
ax	Existing matplotlib axes. When provided the function plots directly onto this axes without setting up the dual-axis canvas (no secondary :math:\ell axis, no axis-label overrides). If None, a new figure with the standard P(k) canvas is created. TYPE: Optional[Axes] DEFAULT: None
fig_kw	Keyword arguments forwarded to plt.subplots() when ax is None. TYPE: Optional[dict] DEFAULT: None
**line_kwargs	Additional keyword arguments forwarded to the median line (e.g. ls='--', lw=1.5). The color kwarg here takes precedence over the color argument. DEFAULT: {}

RETURNS	DESCRIPTION
Matplotlib	class:~matplotlib.figure.Figure object. TYPE: Figure

Examples:

>>> import numpy as np
>>> import primefeat.plots as pf_plot
>>> k = np.logspace(-3, 0, 200)
>>> # Prior bands on 1 + delta(k) with uniform prior on each bin
>>> fig = pf_plot.priors_PPS(k, delta_prior=(-0.5, 0.5))

>>> # Full P(k) prior including A_s and n_s
>>> fig = pf_plot.priors_PPS(
...     k,
...     delta_prior=(-0.3, 0.3),
...     As_prior=(1.5, 2.5),
...     ns_prior=('normal', 0.965, 0.02),
...     sigma_levels=2,
...     color='darkorange',
... )

Source code in src/primefeat/plots/pk.py

def priors_PPS(
    k: np.ndarray,
    delta_prior: Union[Tuple, List[Tuple]],
    As_prior: Optional[Tuple] = None,
    ns_prior: Optional[Tuple] = None,
    n_samples: int = 10_000,
    k_pivot: float = 0.05,
    sigma_levels: int = 2,
    alpha_contour: float = 0.5,
    normalize: float = 1e9,
    color: str = "steelblue",
    label: Optional[str] = "Prior",
    ax: Optional[plt.Axes] = None,
    fig_kw: Optional[dict] = None,
    **line_kwargs,
) -> plt.Figure:
    """
    Plot the prior distribution of the primordial power spectrum P(k).

    Randomly draws ``n_samples`` parameter sets from the user-specified
    priors, evaluates P(k) for each draw, then displays the resulting
    distribution as sigma-level bands (identical style to
    :func:`posteriors_PPS`).

    Two modes are supported automatically:

    * **Delta mode** (default) — plots :math:`1 + \\delta(k)` when neither
      ``As_prior`` nor ``ns_prior`` is supplied.
    * **Full mode** — plots
      :math:`10^9 A_s (k/k_{\\mathrm{pivot}})^{n_s-1} [1+\\delta(k)]`
      when at least one of ``As_prior`` / ``ns_prior`` is provided. Missing
      parameters are held fixed at their fiducial values
      (:math:`A_s = 2.1 \\times 10^{-9}`, :math:`n_s = 0.965`).

    Prior specifications
    --------------------
    A *prior spec* is either:

    * ``(low, high)`` — draws from :math:`\\mathcal{U}(\\text{low}, \\text{high})`.
    * ``('normal', mean, std)`` — draws from :math:`\\mathcal{N}(\\text{mean}, \\text{std}^2)`.

    For ``delta_prior`` you may supply a single spec applied to **all**
    k-bins, or a list of specs with one entry per bin.

    Args:
        k: Array of :math:`k` values (Mpc :sup:`-1`).
        delta_prior: Prior specification for the :math:`\\delta(k)` bins.
            Either a single spec ``(low, high)`` / ``('normal', mean, std)``
            applied to every bin, or a list of such specs (one per bin).
        As_prior: Optional prior for :math:`A_s` in units of
            :math:`10^{-9}`. E.g. ``(1.5, 2.5)`` for a uniform prior.
            If ``None`` and full mode is triggered via ``ns_prior``, the
            fiducial value :math:`A_s = 2.1 \\times 10^{-9}` is used.
        ns_prior: Optional prior for :math:`n_s`.
            E.g. ``(0.9, 1.05)`` or ``('normal', 0.965, 0.01)``.
            If ``None`` and full mode is triggered via ``As_prior``, the
            fiducial value :math:`n_s = 0.965` is used.
        n_samples: Number of random draws from the prior (default: 10 000).
        k_pivot: Pivot scale in Mpc :sup:`-1` (default: 0.05).
        sigma_levels: Number of :math:`\\sigma` contours to display —
            1, 2, or 3 (default: 2).
        alpha_contour: Transparency of the confidence bands (default: 0.5).
        normalize: Overall normalization factor (default: :math:`10^9`).
            Applied as a multiplicative prefactor to the plotted quantity.
        color: Color for the prior bands and median line (default:
            ``'steelblue'``).
        label: Legend label (default: ``'Prior'``).
        ax: Existing matplotlib axes. When provided the function plots
            directly onto this axes **without** setting up the dual-axis
            canvas (no secondary :math:`\\ell` axis, no axis-label
            overrides). If ``None``, a new figure with the standard P(k)
            canvas is created.
        fig_kw: Keyword arguments forwarded to ``plt.subplots()`` when
            ``ax`` is ``None``.
        **line_kwargs: Additional keyword arguments forwarded to the median
            line (e.g. ``ls='--'``, ``lw=1.5``). The ``color`` kwarg here
            takes precedence over the ``color`` argument.

    Returns:
        Matplotlib :class:`~matplotlib.figure.Figure` object.

    Examples:
        >>> import numpy as np
        >>> import primefeat.plots as pf_plot
        >>> k = np.logspace(-3, 0, 200)
        >>> # Prior bands on 1 + delta(k) with uniform prior on each bin
        >>> fig = pf_plot.priors_PPS(k, delta_prior=(-0.5, 0.5))

        >>> # Full P(k) prior including A_s and n_s
        >>> fig = pf_plot.priors_PPS(
        ...     k,
        ...     delta_prior=(-0.3, 0.3),
        ...     As_prior=(1.5, 2.5),
        ...     ns_prior=('normal', 0.965, 0.02),
        ...     sigma_levels=2,
        ...     color='darkorange',
        ... )
    """
    k = np.asarray(k)
    n_bins = len(k)
    full_mode = (As_prior is not None) or (ns_prior is not None)

    # --- Sample deltas ---------------------------------------------------
    if isinstance(delta_prior, list) and not isinstance(
        delta_prior[0], (int, float, str)
    ):
        # Per-bin prior: list of specs, one per bin
        if len(delta_prior) != n_bins:
            raise ValueError(
                f"delta_prior has {len(delta_prior)} entries but k has {n_bins} bins."
            )
        delta_samples = np.column_stack(
            [_sample_from_prior(spec, n_samples) for spec in delta_prior]
        )
    else:
        # Single spec applied to all bins
        delta_samples = _sample_from_prior(delta_prior, n_samples, n_bins=n_bins)

    # --- Build P(k) samples ----------------------------------------------
    if full_mode:
        # Sample A_s (units: 1e-9)
        if As_prior is not None:
            As_samples = _sample_from_prior(As_prior, n_samples)  # shape (n,)
        else:
            As_samples = np.full(n_samples, 2.1)  # fiducial

        # Sample n_s
        if ns_prior is not None:
            ns_samples = _sample_from_prior(ns_prior, n_samples)  # shape (n,)
        else:
            ns_samples = np.full(n_samples, 0.965)  # fiducial

        # P(k) = normalize * As * (k/kpivot)^(ns-1) * (1+delta)
        # As_samples is in units of 1e-9, normalize=1e9 → net factor ~1
        powerlaw = (k[np.newaxis, :] / k_pivot) ** (ns_samples[:, np.newaxis] - 1)
        Pk_samples = (
            normalize
            * As_samples[:, np.newaxis]
            * 1e-9
            * powerlaw
            * (1.0 + delta_samples)
        )
        ylabel = r"$10^9~\mathcal{P}_\zeta(k)$"
    else:
        Pk_samples = 1.0 + delta_samples
        ylabel = r"$1+\delta\mathcal{P}(k)$"

    # --- Canvas ----------------------------------------------------------
    if ax is None:
        fig, ax, _ = create_pk_canvas(
            fig_kw=fig_kw,
            primary_axis="k",
            ylabel=ylabel,
            xlim=(k.min(), k.max()),
        )
    else:
        fig = ax.figure

    # Merge color: line_kwargs override the dedicated color argument
    plot_color = line_kwargs.pop("color", color)

    # Legend proxy
    ax.plot([], [], color=plot_color, label=label, **line_kwargs)

    ax = plot_fill_between(
        k,
        Pk_samples,
        ax=ax,
        color=plot_color,
        alpha_contour=alpha_contour,
        quantiles=_SIGMA_LEVELS_TO_QUANTILES[sigma_levels],
        **line_kwargs,
    )

    # Reference line
    if not full_mode:
        ax.axhline(1.0, color="k", linestyle="--", alpha=0.7, lw=1)

    ax.legend(fontsize=10, loc="best")
    plt.tight_layout()
    return fig

powerlaw_residuals(k, samples_Pk, colors=None, ax=None, alpha_contour=0.5, sigma_levels=2, show_binning_range=False, k_start=None, k_end=None, add_inset=False, inset_klim=(0.01, 0.2), inset_bbox=(0.55, 0.55, 0.4, 0.4), inset_ylim=None, show_nodes=False, nbins=20, node_marker='o', node_size=40, use_median=False, fig_kw=None, chain_entries=None)

Plot residuals with respect to a fiducial power-law power spectrum.

This function plots \(\delta(k)\), the fractional deviations from the best-fit power-law model \(P_{\mathrm{PL}}(k) = A_s \cdot (k/k_{\mathrm{pivot}})^{(n_s-1)}\). The residuals represent the primordial features in the power spectrum.

PARAMETER	DESCRIPTION
k	Array of \(k\) values (wavenumbers) in Mpc\(^{-1}\). TYPE: ndarray
samples_Pk	Dict mapping labels to arrays of \(\delta(k)\) samples. Each array should have shape (n_samples, n_k_bins). \(\delta(k)\) values represent fractional deviations from power-law. TYPE: dict
colors	List of colors for each dataset (optional, required if chain_entries not provided). TYPE: list DEFAULT: None
ax	Existing matplotlib axes. If None, creates new figure. TYPE: Optional[Axes] DEFAULT: None
alpha_contour	Transparency for confidence bands (default: 0.5). TYPE: float DEFAULT: 0.5
sigma_levels	Number of \(\sigma\) levels for contours (1, 2, or 3). TYPE: int DEFAULT: 2
show_binning_range	If True, shade extrapolation regions. TYPE: bool DEFAULT: False
k_start	Start of binning range (for extrapolation shading). TYPE: float DEFAULT: None
k_end	End of binning range (for extrapolation shading). TYPE: float DEFAULT: None
add_inset	If True, add an inset plot with zoomed view. TYPE: bool DEFAULT: False
inset_klim	Tuple (kmin, kmax) for inset zoom range. TYPE: tuple DEFAULT: (0.01, 0.2)
inset_bbox	Tuple (x, y, width, height) for inset position. TYPE: tuple DEFAULT: (0.55, 0.55, 0.4, 0.4)
inset_ylim	Tuple (ymin, ymax) for inset y-axis limits. TYPE: tuple DEFAULT: None
show_nodes	If True, plot node values at bin centers as scatter points. TYPE: bool DEFAULT: False
nbins	Number of bins (used if show_nodes=True). TYPE: int DEFAULT: 20
node_marker	Marker style for node points (default: 'o'). TYPE: str DEFAULT: 'o'
node_size	Size of node markers (default: 40). TYPE: float DEFAULT: 40
use_median	If True, use median instead of mean for node values. TYPE: bool DEFAULT: False
fig_kw	Keyword arguments passed to plt.subplots(). TYPE: Optional[dict] DEFAULT: None

RETURNS	DESCRIPTION
Figure	Matplotlib Figure object.

Examples:

>>> import primefeat.plots as plot
>>> fig = plot.powerlaw_residuals(
...     k=k_centers,
...     samples={"Planck": delta_samples_planck, "ACT": delta_samples_act},
...     colors=["C0", "C1"],
...     sigma_levels=2,
... )

Source code in src/primefeat/plots/pk.py

def powerlaw_residuals(
    k: np.ndarray,
    samples_Pk: dict,
    colors: list = None,
    ax: Optional[plt.Axes] = None,
    alpha_contour: float = 0.5,
    sigma_levels: int = 2,
    show_binning_range: bool = False,
    k_start: float = None,
    k_end: float = None,
    add_inset: bool = False,
    inset_klim: tuple = (1e-2, 0.2),
    inset_bbox: tuple = (0.55, 0.55, 0.4, 0.4),
    inset_ylim: tuple = None,
    show_nodes: bool = False,
    nbins: int = 20,
    node_marker: str = "o",
    node_size: float = 40,
    use_median: bool = False,
    fig_kw: Optional[dict] = None,
    chain_entries: Optional[List] = None,
) -> plt.Figure:
    """
    Plot residuals with respect to a fiducial power-law power spectrum.

    This function plots $\\delta(k)$, the fractional deviations from the best-fit
    power-law model $P_{\\mathrm{PL}}(k) = A_s \\cdot (k/k_{\\mathrm{pivot}})^{(n_s-1)}$.
    The residuals represent the primordial features in the power spectrum.

    Args:
        k: Array of $k$ values (wavenumbers) in Mpc$^{-1}$.
        samples_Pk: Dict mapping labels to arrays of $\\delta(k)$ samples.
            Each array should have shape (n_samples, n_k_bins).
            $\\delta(k)$ values represent fractional deviations from power-law.
        colors: List of colors for each dataset (optional, required if chain_entries not provided).
        ax: Existing matplotlib axes. If None, creates new figure.
        alpha_contour: Transparency for confidence bands (default: 0.5).
        sigma_levels: Number of $\\sigma$ levels for contours (1, 2, or 3).
        show_binning_range: If True, shade extrapolation regions.
        k_start: Start of binning range (for extrapolation shading).
        k_end: End of binning range (for extrapolation shading).
        add_inset: If True, add an inset plot with zoomed view.
        inset_klim: Tuple (kmin, kmax) for inset zoom range.
        inset_bbox: Tuple (x, y, width, height) for inset position.
        inset_ylim: Tuple (ymin, ymax) for inset y-axis limits.
        show_nodes: If True, plot node values at bin centers as scatter points.
        nbins: Number of bins (used if show_nodes=True).
        node_marker: Marker style for node points (default: 'o').
        node_size: Size of node markers (default: 40).
        use_median: If True, use median instead of mean for node values.
        fig_kw: Keyword arguments passed to plt.subplots().

    Returns:
        Matplotlib Figure object.

    Examples:
        >>> import primefeat.plots as plot
        >>> fig = plot.powerlaw_residuals(
        ...     k=k_centers,
        ...     samples={"Planck": delta_samples_planck, "ACT": delta_samples_act},
        ...     colors=["C0", "C1"],
        ...     sigma_levels=2,
        ... )
    """
    # Extract colors from chain_entries if provided
    if chain_entries is not None:
        colors = [entry.color for entry in chain_entries]

    if colors is None:
        raise ValueError("Either colors or chain_entries must be provided")

    # Create canvas with dual k/ell axes
    fig, ax, _ = create_pk_canvas(
        ax=ax,
        fig_kw=fig_kw,
        primary_axis="k",
        ylabel=r"$\delta\mathcal{P}(k) [\%]$",
        xlim=(k.min(), k.max()),
    )

    # Show binning range if requested (gray shading outside binning region)
    if show_binning_range and k_start is not None and k_end is not None:
        ax.axvspan(k.min(), k_start, alpha=0.15, color="gray", zorder=0)
        ax.axvspan(
            k_end,
            k.max(),
            alpha=0.15,
            color="gray",
            zorder=0,
            label="Extrapolation region",
        )

    # Plot each dataset using plot_fill_between
    for i, (label, sample_array) in enumerate(samples_Pk.items()):
        sample_array = np.array(sample_array)
        color = colors[i]

        ax.plot([], [], color=color, label=label)  # For legend
        ax = plot_fill_between(
            k,
            sample_array,
            ax=ax,
            color=color,
            alpha_contour=alpha_contour,
            quantiles=_SIGMA_LEVELS_TO_QUANTILES[sigma_levels],
        )

    # Reference line at δ = 0 (no deviation from power-law)
    ax.axhline(0.0, color="k", linestyle="--", alpha=0.7, lw=1)

    # Add node values at bin centers if requested
    if show_nodes and k_start is not None and k_end is not None:
        bin_centers = get_bin_centers(k_start, k_end, nbins)

        for i, (label, sample_array) in enumerate(samples_Pk.items()):
            sample_array = np.array(sample_array)
            color = colors[i]

            # Compute mean or median across samples
            if use_median:
                node_values = np.median(sample_array, axis=0)
            else:
                node_values = np.mean(sample_array, axis=0)

            # Find indices in k array closest to bin centers
            bin_indices = [np.argmin(np.abs(k - bc)) for bc in bin_centers]
            node_y = node_values[bin_indices]

            # Plot nodes
            ax.scatter(
                bin_centers,
                node_y,
                s=node_size,
                marker=node_marker,
                color=color,
                edgecolors="white",
                linewidth=1,
                zorder=10,
            )

    ax.legend(fontsize=10, loc="best")

    # Add inset plot if requested
    if add_inset:
        axins = ax.inset_axes(inset_bbox)

        # Filter k values within inset range
        k_mask = (k >= inset_klim[0]) & (k <= inset_klim[1])
        k_inset = k[k_mask]

        # Plot each dataset in the inset
        for i, (label, sample_array) in enumerate(samples_Pk.items()):
            sample_array = np.array(sample_array)
            sample_inset = sample_array[:, k_mask]
            color = colors[i]

            axins = plot_fill_between(
                k_inset,
                sample_inset,
                ax=axins,
                color=color,
                alpha_contour=alpha_contour,
                quantiles=_SIGMA_LEVELS_TO_QUANTILES[sigma_levels],
            )

        axins.set_xscale("log")
        axins.set_xlim(inset_klim[0], inset_klim[1])
        axins.axhline(0.0, color="k", linestyle="--", alpha=0.7, lw=1)

        if inset_ylim is not None:
            axins.set_ylim(inset_ylim[0], inset_ylim[1])

        axins.grid(True, alpha=0.3)
        axins.tick_params(labelsize=8)

        fig.inset_axes = axins

    plt.tight_layout()
    return fig

posteriors_ns_eff(k, ns_eff_samples, colors=None, chain_entries=None, ax=None, show_binning_range=False, k_start=None, k_end=None, alpha_contour=0.5, add_inset=False, inset_klim=(0.01, 0.2), inset_bbox=(0.55, 0.55, 0.4, 0.4), inset_ylim=None, fig_kw=None)

Plot posteriors for effective spectral index \(n_s(k) - 1 = d \ln P(k) / d \ln k\).

PARAMETER	DESCRIPTION
k	Array of \(k\) values TYPE: ndarray
ns_eff_samples	Dict mapping labels to arrays of \(n_s(k) - 1\) samples TYPE: dict
colors	List of colors for each dataset (optional if chain_entries provided) TYPE: Optional[list] DEFAULT: None
chain_entries	ChainsCollection or list of ChainEntry objects; colors auto-extracted if provided TYPE: Optional[List] DEFAULT: None
ax	Existing matplotlib axes (optional) TYPE: Optional[Axes] DEFAULT: None
show_binning_range	If True, show gray shaded region for binning range TYPE: bool DEFAULT: False
k_start	Start of binning range (only used if show_binning_range=True) TYPE: float DEFAULT: None
k_end	End of binning range (only used if show_binning_range=True) TYPE: float DEFAULT: None
alpha_contour	Transparency for confidence bands (default: 0.5) TYPE: float DEFAULT: 0.5
add_inset	If True, add an inset plot with zoomed view TYPE: bool DEFAULT: False
inset_klim	Tuple of (kmin, kmax) for inset zoom range TYPE: tuple DEFAULT: (0.01, 0.2)
inset_bbox	Tuple of (x, y, width, height) for inset position in axes coordinates TYPE: tuple DEFAULT: (0.55, 0.55, 0.4, 0.4)
inset_ylim	Tuple of (ymin, ymax) for inset y-axis limits (optional) TYPE: tuple DEFAULT: None
fig_kw	Additional keyword arguments for plt.subplots (e.g., figsize) TYPE: Optional[dict] DEFAULT: None

RETURNS	DESCRIPTION
Figure	Matplotlib Figure object.

RAISES	DESCRIPTION
ValueError	If neither colors nor chain_entries is provided.

Source code in src/primefeat/plots/pk.py

def posteriors_ns_eff(
    k: np.ndarray,
    ns_eff_samples: dict,
    colors: Optional[list] = None,
    chain_entries: Optional[List] = None,
    ax: Optional[plt.Axes] = None,
    show_binning_range: bool = False,
    k_start: float = None,
    k_end: float = None,
    alpha_contour: float = 0.5,
    add_inset: bool = False,
    inset_klim: tuple = (1e-2, 0.2),
    inset_bbox: tuple = (0.55, 0.55, 0.4, 0.4),
    inset_ylim: tuple = None,
    fig_kw: Optional[dict] = None,
) -> plt.Figure:
    """
    Plot posteriors for effective spectral index $n_s(k) - 1 = d \\ln P(k) / d \\ln k$.

    Args:
        k: Array of $k$ values
        ns_eff_samples: Dict mapping labels to arrays of $n_s(k) - 1$ samples
        colors: List of colors for each dataset (optional if chain_entries provided)
        chain_entries: ChainsCollection or list of ChainEntry objects; colors auto-extracted if provided
        ax: Existing matplotlib axes (optional)
        show_binning_range: If True, show gray shaded region for binning range
        k_start: Start of binning range (only used if show_binning_range=True)
        k_end: End of binning range (only used if show_binning_range=True)
        alpha_contour: Transparency for confidence bands (default: 0.5)
        add_inset: If True, add an inset plot with zoomed view
        inset_klim: Tuple of (kmin, kmax) for inset zoom range
        inset_bbox: Tuple of (x, y, width, height) for inset position in axes coordinates
        inset_ylim: Tuple of (ymin, ymax) for inset y-axis limits (optional)
        fig_kw: Additional keyword arguments for plt.subplots (e.g., figsize)

    Returns:
        Matplotlib Figure object.

    Raises:
        ValueError: If neither colors nor chain_entries is provided.
    """
    # Handle color extraction from chain_entries
    if chain_entries is not None:
        colors = [e.color for e in chain_entries]
    elif colors is None:
        raise ValueError(
            "Either colors or chain_entries must be provided to posteriors_ns_eff()"
        )
    # Create canvas with dual k/ell axes
    fig, ax, _ = create_pk_canvas(
        ax=ax,
        fig_kw=fig_kw,
        primary_axis="k",
        ylabel=r"$n_s(k) - 1$",
        xlim=(k.min(), k.max()),
    )

    # Show binning range if requested (inverted shading: gray outside, white inside)
    if show_binning_range and k_start is not None and k_end is not None:
        # Shade region before k_start
        ax.axvspan(k.min(), k_start, alpha=0.15, color="gray", zorder=0)
        # Shade region after k_end
        ax.axvspan(
            k_end,
            k.max(),
            alpha=0.15,
            color="gray",
            zorder=0,
            label="Extrapolation region",
        )

    # Plot each dataset using plot_fill_between
    for i, (label, sample_array) in enumerate(ns_eff_samples.items()):
        color = colors[i]

        # Use the existing plot_fill_between function
        ax.plot([], [], color=color, label=label)  # For legend
        ax = plot_fill_between(
            k, sample_array, ax=ax, color=color, alpha_contour=alpha_contour
        )

    # Reference line at n_s - 1 = 0 (scale-invariant Harrison-Zel'dovich spectrum)
    ax.axhline(0.0, color="k", linestyle="--", alpha=0.7, lw=1, label=r"$n_s = 1$ (HZ)")

    ax.legend(fontsize=10, loc="best")

    # Add inset plot if requested
    if add_inset:
        # Create inset axes
        axins = ax.inset_axes(inset_bbox)

        # Filter k values within inset range
        k_mask = (k >= inset_klim[0]) & (k <= inset_klim[1])
        k_inset = k[k_mask]

        # Plot each dataset in the inset
        for i, (label, sample_array) in enumerate(ns_eff_samples.items()):
            # Filter samples to inset range
            sample_inset = sample_array[:, k_mask]
            color = colors[i]

            # Use plot_fill_between for inset
            axins = plot_fill_between(
                k_inset,
                sample_inset,
                ax=axins,
                color=color,
                alpha_contour=alpha_contour,
            )

        # Configure inset axes
        axins.set_xscale("log")
        axins.set_xlim(inset_klim[0], inset_klim[1])
        axins.axhline(0.0, color="k", linestyle="--", alpha=0.7, lw=1)

        if inset_ylim is not None:
            axins.set_ylim(inset_ylim[0], inset_ylim[1])

        axins.grid(True, alpha=0.3)
        axins.tick_params(labelsize=8)

        # Store inset axes as a figure attribute for easy access
        fig.inset_axes = axins

    plt.tight_layout()
    return fig

posteriors_delta_pcolormesh(chain_dict, nbins=20, n_grid_points=100, delta_range=(-0.5, 0.5), cmap='viridis', figsize=(10, 6), log_density=True, title=None)

Plot posterior distributions for all delta_i bins as a 2D pcolormesh density map.

Creates a 2D histogram/density visualization where: - X-axis: bin index (delta_1, delta_2, ..., delta_nbins) - Y-axis: parameter value (delta) - Color: probability density from posterior samples

PARAMETER	DESCRIPTION
chain_dict	Dictionary mapping bin indices to samples, or a GetDist MCSamples object with delta_i parameters TYPE: dict
nbins	Number of delta bins (default: 20) TYPE: int DEFAULT: 20
n_grid_points	Resolution for density grid (default: 100) TYPE: int DEFAULT: 100
delta_range	Tuple of (min, max) for y-axis delta range. Default: (-0.5, 0.5) based on flat prior. Set to None to infer from data with 10% padding. TYPE: Optional[Tuple[float, float]] DEFAULT: (-0.5, 0.5)
cmap	Colormap name (default: 'viridis') TYPE: str DEFAULT: 'viridis'
figsize	Figure size (default: (10, 6)) TYPE: Tuple DEFAULT: (10, 6)
log_density	If True, show log10(density) for better visibility (default: True) TYPE: bool DEFAULT: True
title	Figure title (optional) TYPE: str DEFAULT: None

RETURNS	DESCRIPTION
Tuple[Figure, Axes]	Tuple of (fig, ax) matplotlib objects

Examples:

>>> import primefeat as pf
>>> chains = pf.get_chains(1e-3)
>>> fig, ax = pf.plot.posteriors_delta_pcolormesh(chains['PR3'], nbins=20)
>>> plt.show()

>>> # Multiple datasets in subplots
>>> fig, axs = plt.subplots(1, len(chains), figsize=(12, 5))
>>> for ax, (label, chain) in zip(axs, chains.items()):
...     pf.plot.posteriors_delta_pcolormesh(chain, nbins=20, figsize=(4, 5))

Source code in src/primefeat/plots/pk.py

def posteriors_delta_pcolormesh(
    chain_dict: dict,
    nbins: int = 20,
    n_grid_points: int = 100,
    delta_range: Optional[Tuple[float, float]] = (-0.5, 0.5),
    cmap: str = "viridis",
    figsize: Tuple = (10, 6),
    log_density: bool = True,
    title: str = None,
) -> Tuple[plt.Figure, plt.Axes]:
    r"""
    Plot posterior distributions for all delta_i bins as a 2D pcolormesh density map.

    Creates a 2D histogram/density visualization where:
    - X-axis: bin index (delta_1, delta_2, ..., delta_nbins)
    - Y-axis: parameter value (delta)
    - Color: probability density from posterior samples

    Args:
        chain_dict: Dictionary mapping bin indices to samples, or
                    a GetDist MCSamples object with delta_i parameters
        nbins: Number of delta bins (default: 20)
        n_grid_points: Resolution for density grid (default: 100)
        delta_range: Tuple of (min, max) for y-axis delta range.
                     Default: (-0.5, 0.5) based on flat prior.
                     Set to None to infer from data with 10% padding.
        cmap: Colormap name (default: 'viridis')
        figsize: Figure size (default: (10, 6))
        log_density: If True, show log10(density) for better visibility (default: True)
        title: Figure title (optional)

    Returns:
        Tuple of (fig, ax) matplotlib objects

    Examples:
        >>> import primefeat as pf
        >>> chains = pf.get_chains(1e-3)
        >>> fig, ax = pf.plot.posteriors_delta_pcolormesh(chains['PR3'], nbins=20)
        >>> plt.show()

        >>> # Multiple datasets in subplots
        >>> fig, axs = plt.subplots(1, len(chains), figsize=(12, 5))
        >>> for ax, (label, chain) in zip(axs, chains.items()):
        ...     pf.plot.posteriors_delta_pcolormesh(chain, nbins=20, figsize=(4, 5))
    """
    from scipy.stats import gaussian_kde

    # Extract delta samples
    delta_samples = {}

    # Handle both GetDist MCSamples and dict formats
    if hasattr(chain_dict, "samples"):
        # GetDist MCSamples object
        for i in range(1, nbins + 1):
            param_name = f"delta_{i}"
            if hasattr(chain_dict, "getParam"):
                delta_samples[i] = chain_dict.getParam(param_name).samples
            else:
                # Fallback: try accessing by index
                delta_samples[i] = chain_dict.samples[:, i - 1]
    else:
        # Dictionary of arrays
        delta_samples = chain_dict

    # Determine delta range
    all_samples = np.concatenate([np.atleast_1d(v) for v in delta_samples.values()])
    if delta_range is None:
        min_val, max_val = np.percentile(all_samples, [0.5, 99.5])
        padding = (max_val - min_val) * 0.1
        delta_range = (min_val - padding, max_val + padding)

    # Create 2D grid
    delta_grid = np.linspace(delta_range[0], delta_range[1], n_grid_points)
    bin_indices = np.arange(1, nbins + 1)

    # Compute density for each bin
    density_grid = np.zeros((n_grid_points, nbins))

    for bin_idx, samples in delta_samples.items():
        try:
            # Use KDE to estimate density
            kde = gaussian_kde(samples)
            density_grid[:, bin_idx - 1] = kde(delta_grid)
        except Exception:
            # If KDE fails, use histogram
            hist, _ = np.histogram(
                samples, bins=n_grid_points, range=delta_range, density=True
            )
            density_grid[:, bin_idx - 1] = hist

    # Apply log if requested (avoid log(0))
    if log_density:
        density_grid = np.log10(density_grid + 1e-12)
        cbar_label = r"$\log_{10}$ Probability Density"
    else:
        cbar_label = "Probability Density"

    # Create figure
    fig, ax = plt.subplots(figsize=figsize)

    # Create pcolormesh
    # Transpose to get delta on y-axis, bins on x-axis
    X, Y = np.meshgrid(bin_indices - 0.5, delta_grid)
    im = ax.pcolormesh(X, Y, density_grid, cmap=cmap, shading="auto")

    # Labels and formatting
    ax.set_xlabel("Bin Index", fontsize=12)
    ax.set_ylabel(r"$\delta_i$ value", fontsize=12)
    ax.set_xticks(bin_indices)
    ax.set_xticklabels([f"{i}" for i in bin_indices], fontsize=9)

    # Add colorbar
    cbar = plt.colorbar(im, ax=ax)
    cbar.set_label(cbar_label, fontsize=11)

    # Add reference line at delta=0
    ax.axhline(
        0,
        color="red",
        linestyle="--",
        alpha=0.5,
        linewidth=1.5,
        label=r"$\delta_i = 0$",
    )
    ax.legend(fontsize=10)

    if title is not None:
        ax.set_title(title, fontsize=13)

    fig.tight_layout()
    return fig, ax