openscm_calibration.emcee_plotting

Support for plotting emcee during run and afterwards

Functions:

Name	Description
get_neg_log_likelihood_ylim_default	Get the y-limits for the negative log likelihood axes
plot_chains	Plot chains in MCMC run
plot_corner	Plot corner plot of distribution from MCMC run
plot_dist	Plot distributions from MCMC run
plot_emcee_progress	Plot MCMC progress
plot_parameter_chains	Plot chains for a single parameter in an MCMC run
plot_tau	Plot the autocorrelation time, tau

Attributes:

Name	Type	Description
DEFAULT_PLOT_CORNER_LABEL_KWARGS	dict[str, Any]	Default value for label_kwargs used by plot_corner
DEFAULT_PLOT_CORNER_TITLE_KWARGS	dict[str, Any]	Default value for title_kwargs used by plot_corner
DEFAULT_PROGRESS_KWARGS		Default arguments to use for displaying progress bars

DEFAULT_PLOT_CORNER_LABEL_KWARGS module-attribute

DEFAULT_PLOT_CORNER_LABEL_KWARGS: dict[str, Any] = {
    "fontsize": "x-small"
}

Default value for label_kwargs used by plot_corner

Provided to give the user an easy to modify these defaults if they wish or a starting point

DEFAULT_PLOT_CORNER_TITLE_KWARGS module-attribute

DEFAULT_PLOT_CORNER_TITLE_KWARGS: dict[str, Any] = {
    "fontsize": 12
}

Default value for title_kwargs used by plot_corner

Provided to give the user an easy to modify these defaults if they wish or a starting point

DEFAULT_PROGRESS_KWARGS module-attribute

DEFAULT_PROGRESS_KWARGS = {'leave': True}

Default arguments to use for displaying progress bars

get_neg_log_likelihood_ylim_default

get_neg_log_likelihood_ylim_default(
    neg_ll_values: NDArray[floating[Any] | integer[Any]],
    median_scaling: float = 1.5,
    max_scaling: float = 2.0,
) -> tuple[float, float]

Get the y-limits for the negative log likelihood axes

This is the default algorithm

Parameters:

Name	Type	Description	Default
neg_ll_values	NDArray[floating[Any] | integer[Any]]	Negative log likelihood values being plotted	required
median_scaling	float	Scaling to apply to the median value	1.5
max_scaling	float	Scaling to apply to the maximum value	2.0

Returns:

Type	Description
y-limits

Notes

The algorithm is

.. math::

\text{median_scaled} = \text{median_scaling} \times \text{med}(\text{neg_ll_values}) \\
\text{max_scaled} = \text{max_scaling} \times \max(\text{neg_ll_values}) \\
\text{ymin} = \min(0, \text{median_scaled}, \text{max_scaled}) \\
\text{ymax} = \max(0, \text{median_scaled}, \text{max_scaled})

Source code in src/openscm_calibration/emcee_plotting.py

def get_neg_log_likelihood_ylim_default(
    neg_ll_values: np.typing.NDArray[np.floating[Any] | np.integer[Any]],
    median_scaling: float = 1.5,
    max_scaling: float = 2.0,
) -> tuple[float, float]:
    r"""
    Get the y-limits for the negative log likelihood axes

    This is the default algorithm

    Parameters
    ----------
    neg_ll_values
        Negative log likelihood values being plotted

    median_scaling
        Scaling to apply to the median value

    max_scaling
        Scaling to apply to the maximum value

    Returns
    -------
        y-limits

    Notes
    -----
    The algorithm is

    .. math::

        \text{median_scaled} = \text{median_scaling} \times \text{med}(\text{neg_ll_values}) \\
        \text{max_scaled} = \text{max_scaling} \times \max(\text{neg_ll_values}) \\
        \text{ymin} = \min(0, \text{median_scaled}, \text{max_scaled}) \\
        \text{ymax} = \max(0, \text{median_scaled}, \text{max_scaled})
    """  # noqa: E501
    median_scaled = float(median_scaling * np.median(neg_ll_values))
    if not np.isfinite(median_scaled):
        median_scaled = 0.0

    max_scaled = float(max_scaling * np.max(neg_ll_values))
    if not np.isfinite(max_scaled):
        max_scaled = 0.0

    ymin = float(
        min(
            0.0,
            median_scaled,
            max_scaled,
        )
    )
    ymax = float(
        max(
            0.0,
            median_scaled,
            max_scaled,
        )
    )

    return (ymin, ymax)

plot_chains

plot_chains(
    inp: Backend | EnsembleSampler,
    burnin: int,
    parameter_order: tuple[str, ...],
    neg_log_likelihood_name: str,
    axes_d: dict[str, Axes],
    get_neg_log_likelihood_ylim: Callable[
        [NDArray[floating[Any] | integer[Any]]],
        tuple[float, float],
    ]
    | None = None,
    **kwargs: Any,
) -> None

Plot chains in MCMC run

The y-limits of the log likelihood axis are set dynamically

Parameters:

Name	Type	Description	Default
inp	Backend | EnsembleSampler	Object from which to plot the state	required
burnin	int	Number of iterations to treat as burn in	required
parameter_order	tuple[str, ...]	Order of model parameters. This must match the order used by inp.	required
neg_log_likelihood_name	str	Label when plotting negative log likelihood. Must match the expected name in axes_d	required
axes_d	dict[str, Axes]	Axes on which to plot the chains. Must have a key for each name in parameter_order plus a key equal to the value of neg_log_likelihood_name	required
get_neg_log_likelihood_ylim	Callable[[NDArray[floating[Any] | integer[Any]]], tuple[float, float]] | None	Function which gets the y-limits for the negative log likelihood plot based on the costs. If not provided, :func:get_neg_log_likelihood_ylim_default is used	None
**kwargs	Any	Passed to :func:plot_parameter_chains	{}

Source code in src/openscm_calibration/emcee_plotting.py

def plot_chains(  # noqa: PLR0913
    inp: emcee.backends.Backend | emcee.ensemble.EnsembleSampler,
    burnin: int,
    parameter_order: tuple[str, ...],
    neg_log_likelihood_name: str,
    axes_d: dict[str, matplotlib.axes.Axes],
    get_neg_log_likelihood_ylim: Callable[
        [np.typing.NDArray[np.floating[Any] | np.integer[Any]]], tuple[float, float]
    ]
    | None = None,
    **kwargs: Any,
) -> None:
    """
    Plot chains in MCMC run

    The y-limits of the log likelihood axis are set dynamically

    Parameters
    ----------
    inp
        Object from which to plot the state

    burnin
        Number of iterations to treat as burn in

    parameter_order
        Order of model parameters. This must match the order used by  ``inp``.

    neg_log_likelihood_name
        Label when plotting negative log likelihood. Must match the expected
        name in ``axes_d``

    axes_d
        Axes on which to plot the chains. Must have a key for each name in
        ``parameter_order`` plus a key equal to the value of
        ``neg_log_likelihood_name``

    get_neg_log_likelihood_ylim
        Function which gets the y-limits for the negative log likelihood plot
        based on the costs. If not provided,
        :func:`get_neg_log_likelihood_ylim_default` is used

    **kwargs
        Passed to :func:`plot_parameter_chains`
    """
    labelled_data = get_labelled_chain_data(
        inp,
        parameter_order,
        neg_log_likelihood_name,
        burnin=0,
        thin=1,
    )

    if get_neg_log_likelihood_ylim is None:
        get_neg_log_likelihood_ylim = get_neg_log_likelihood_ylim_default

    for label, to_plot in labelled_data.items():
        ax = axes_d[label]
        plot_parameter_chains(
            ax,
            to_plot,
            burnin=burnin,
            **kwargs,
        )
        ax.set_ylabel(label)

        if label == neg_log_likelihood_name:
            ax.set_ylim(*get_neg_log_likelihood_ylim(to_plot))

plot_corner

plot_corner(
    inp: Backend | EnsembleSampler,
    burnin: int,
    thin: int,
    parameter_order: tuple[str, ...],
    fig: Figure,
    bins: int = 30,
    plot_contours: bool = True,
    smooth: bool = True,
    quantiles: Sequence[float] = (
        0.05,
        0.17,
        0.5,
        0.83,
        0.95,
    ),
    show_titles: bool = True,
    title_quantiles: Sequence[float] = (0.05, 0.5, 0.95),
    title_kwargs: dict[str, Any] | None = None,
    title_fmt: str = ".3f",
    label_kwargs: dict[str, Any] | None = None,
    **kwargs: Any,
) -> None

Plot corner plot of distribution from MCMC run

This is a thin wrapper around :func:corner.corner which uses some sensible defaults.

Parameters:

Name	Type	Description	Default
inp	Backend | EnsembleSampler	Object from which to plot the state	required
burnin	int	Number of iterations to treat as burn in	required
thin	int	Thinning to use when sampling the chains	required
parameter_order	tuple[str, ...]	Order of model parameters	required
fig	Figure	Figure to use for plotting, should be empty i.e. have been cleared before being passed here	required
bins	int	Number of bins to use in histograms. See docstring of :func:corner.corner for full details.	30
plot_contours	bool	Whether to plot contours on the 2D distribution plots. See docstring of :func:corner.corner for full details.	True
smooth	bool	Whether to smooth the contours on the 2D distribution plots. See docstring of :func:corner.corner for full details.	True
quantiles	Sequence[float]	Quantiles at which to draw vertical lines in the histogram plots. See docstring of :func:corner.corner for full details.	(0.05, 0.17, 0.5, 0.83, 0.95)
show_titles	bool	Whether to show titles on the histogram plots. See docstring of :func:corner.corner for full details.	True
title_quantiles	Sequence[float]	Quantiles to put in the titles of the histogram plots. See docstring of :func:corner.corner for full details.	(0.05, 0.5, 0.95)
title_kwargs	dict[str, Any] | None	Keyword arguments to use when making the titles on the histogram plots. If not supplied, we use DEFAULT_PLOT_CORNER_TITLE_KWARGS.	None
title_fmt	str	Format string to use when creating the titles on the histogram plots. If not supplied, our own internal defaults are used (see source code for values). See docstring of :func:corner.corner for full details.	'.3f'
label_kwargs	dict[str, Any] | None	Keyword arguments to use when creating the labels for the plot. If not supplied, we use DEFAULT_PLOT_CORNER_LABEL_KWARGS.	None
**kwargs	Any	Passed to :func:corner.corner	{}

Source code in src/openscm_calibration/emcee_plotting.py

def plot_corner(  # noqa: PLR0913,too-many-locals
    inp: emcee.backends.Backend | emcee.ensemble.EnsembleSampler,
    burnin: int,
    thin: int,
    parameter_order: tuple[str, ...],
    fig: matplotlib.figure.Figure,
    bins: int = 30,
    plot_contours: bool = True,
    smooth: bool = True,
    quantiles: Sequence[float] = (0.05, 0.17, 0.5, 0.83, 0.95),
    show_titles: bool = True,
    title_quantiles: Sequence[float] = (0.05, 0.5, 0.95),
    title_kwargs: dict[str, Any] | None = None,
    title_fmt: str = ".3f",
    label_kwargs: dict[str, Any] | None = None,
    **kwargs: Any,
) -> None:
    """
    Plot corner plot of distribution from MCMC run

    This is a thin wrapper around :func:`corner.corner` which uses some
    sensible defaults.

    Parameters
    ----------
    inp
        Object from which to plot the state

    burnin
        Number of iterations to treat as burn in

    thin
        Thinning to use when sampling the chains

    parameter_order
        Order of model parameters

    fig
        Figure to use for plotting, should be empty i.e. have been cleared
        before being passed here

    bins
        Number of bins to use in histograms. See docstring of
        :func:`corner.corner` for full details.

    plot_contours
        Whether to plot contours on the 2D distribution plots. See docstring
        of :func:`corner.corner` for full details.

    smooth
        Whether to smooth the contours on the 2D distribution plots. See
        docstring of :func:`corner.corner` for full details.

    quantiles
        Quantiles at which to draw vertical lines in the histogram plots. See
        docstring of :func:`corner.corner` for full details.

    show_titles
        Whether to show titles on the histogram plots. See docstring of
        :func:`corner.corner` for full details.

    title_quantiles
        Quantiles to put in the titles of the histogram plots. See docstring
        of :func:`corner.corner` for full details.

    title_kwargs
        Keyword arguments to use when making the titles on the histogram
        plots. If not supplied, we use
        ``DEFAULT_PLOT_CORNER_TITLE_KWARGS``.

    title_fmt
        Format string to use when creating the titles on the histogram plots.
        If not supplied, our own internal defaults are used (see source code
        for values). See docstring of :func:`corner.corner` for full details.

    label_kwargs
        Keyword arguments to use when creating the labels for the plot. If not
        supplied, we use ``DEFAULT_PLOT_CORNER_LABEL_KWARGS``.

    **kwargs
        Passed to :func:`corner.corner`
    """
    try:
        import corner
    except ImportError as exc:
        raise MissingOptionalDependencyError(
            "plot_corner", requirement="corner"
        ) from exc

    if title_kwargs is None:
        title_kwargs = DEFAULT_PLOT_CORNER_TITLE_KWARGS

    if label_kwargs is None:
        label_kwargs = DEFAULT_PLOT_CORNER_LABEL_KWARGS

    burnt_in_samples_flat = inp.get_chain(discard=burnin, thin=thin, flat=True)

    corner.corner(
        burnt_in_samples_flat,
        labels=parameter_order,
        fig=fig,
        bins=bins,
        plot_contours=plot_contours,
        smooth=smooth,
        quantiles=quantiles,
        show_titles=show_titles,
        title_quantiles=title_quantiles,
        title_kwargs=title_kwargs,
        title_fmt=title_fmt,
        label_kwargs=label_kwargs,
        **kwargs,
    )

plot_dist

plot_dist(
    inp: Backend | EnsembleSampler,
    burnin: int,
    thin: int,
    parameter_order: tuple[str, ...],
    axes_d: dict[str, Axes],
    common_norm: bool = False,
    fill: bool = True,
    legend: bool = False,
    **kwargs: Any,
) -> None

Plot distributions from MCMC run

This is a thin wrapper around :func:sns.kdeplot that sets helpful defaults.

Parameters:

Name	Type	Description	Default
inp	Backend | EnsembleSampler	Object from which to plot the state	required
burnin	int	Number of iterations to treat as burn in	required
thin	int	Thinning to use when sampling the chains	required
parameter_order	tuple[str, ...]	Order of model parameters	required
axes_d	dict[str, Axes]	Axes on which to plot the distributions	required
common_norm	bool	Should all the distributions use the same normalisation? We generally set this to False as we want to see the chains independently. See docstring of :func:sns.kdeplot for full details.	False
fill	bool	Should the KDE plots be filled? See docstring of :func:sns.kdeplot for full details.	True
legend	bool	Should a legend be added to the plots? The legend is pretty meaningless (it is just the chain numbers) so we generally set this to False. See docstring of :func:sns.kdeplot for full details.	False
**kwargs	Any	Passed to :func:sns.kdeplot.	{}

Source code in src/openscm_calibration/emcee_plotting.py

def plot_dist(  # noqa: PLR0913
    inp: emcee.backends.Backend | emcee.ensemble.EnsembleSampler,
    burnin: int,
    thin: int,
    parameter_order: tuple[str, ...],
    axes_d: dict[str, matplotlib.axes.Axes],
    common_norm: bool = False,
    fill: bool = True,
    legend: bool = False,
    **kwargs: Any,
) -> None:
    """
    Plot distributions from MCMC run

    This is a thin wrapper around :func:`sns.kdeplot` that sets helpful
    defaults.

    Parameters
    ----------
    inp
        Object from which to plot the state

    burnin
        Number of iterations to treat as burn in

    thin
        Thinning to use when sampling the chains

    parameter_order
        Order of model parameters

    axes_d
        Axes on which to plot the distributions

    common_norm
        Should all the distributions use the same normalisation? We generally
        set this to ``False`` as we want to see the chains independently. See
        docstring of :func:`sns.kdeplot` for full details.

    fill
        Should the KDE plots be filled? See docstring of :func:`sns.kdeplot`
        for full details.

    legend
        Should a legend be added to the plots? The legend is pretty
        meaningless (it is just the chain numbers) so we generally set this to
        ``False``. See docstring of :func:`sns.kdeplot` for full details.

    **kwargs
        Passed to :func:`sns.kdeplot`.
    """
    try:
        import seaborn as sns
    except ImportError as exc:
        raise MissingOptionalDependencyError(
            "plot_dist", requirement="seaborn"
        ) from exc

    burnt_in_samples_labelled = get_labelled_chain_data(
        inp,
        parameter_order,
        burnin=burnin,
        thin=thin,
    )

    for label, to_plot in burnt_in_samples_labelled.items():
        ax = axes_d[label]

        sns.kdeplot(
            data=to_plot,
            ax=ax,
            common_norm=common_norm,
            fill=fill,
            legend=legend,
            **kwargs,
        )
        ax.set_xlabel(label)

plot_emcee_progress

plot_emcee_progress(
    sampler: EnsembleSampler,
    iterations: int,
    burnin: int,
    thin: int,
    plot_every: int,
    parameter_order: tuple[str, ...],
    neg_log_likelihood_name: str,
    holder_chain: DisplayHandle,
    figure_chain: Figure,
    axes_chain: dict[str, Axes],
    holder_dist: DisplayHandle,
    figure_dist: Figure,
    axes_dist: dict[str, Axes],
    holder_corner: DisplayHandle,
    figure_corner: Figure,
    holder_tau: DisplayHandle,
    figure_tau: Figure,
    ax_tau: Axes,
    start: NDArray[float64] | None = None,
    sample_kwargs: dict[str, Any] | None = None,
    progress: str | bool = "notebook",
    progress_kwargs: dict[str, Any] | None = None,
    min_samples_before_plot: int = 2,
    corner_kwargs: dict[str, Any] | None = None,
    convergence_ratio: float = 50,
) -> Iterator[ChainProgressInfo]

Plot MCMC progress

Parameters:

Name	Type	Description	Default
sampler	EnsembleSampler	Sampler used for sampling the posterior distribution	required
iterations	int	Number of iterations to perform	required
burnin	int	Burn-in to apply to the chains	required
thin	int	Thinning to apply to the chains	required
plot_every	int	How many steps to wait before updating the plots	required
parameter_order	tuple[str, ...]	Order of the parameters. This is used for ensuring that the plots are labelled correctly.	required
neg_log_likelihood_name	str	Name to use for the negative log likelihood in plots	required
holder_chain	DisplayHandle	Holder of the figure that displays the chains	required
figure_chain	Figure	Figure that displays the chains	required
axes_chain	dict[str, Axes]	Axes on which to plot the chains. Each parameter in parameter_order should be a key in axes_chain.	required
holder_dist	DisplayHandle	Holder of the figure that displays the parameter distributions	required
figure_dist	Figure	Figure that displays the parameter distributions	required
axes_dist	dict[str, Axes]	Axes on which to plot the parameter distributions. Each parameter in parameter_order should be a key in axes_dist.	required
holder_corner	DisplayHandle	Holder of the figure that displays the corner plot	required
figure_corner	Figure	Figure that displays the corner plot	required
holder_tau	DisplayHandle	Holder of the figure that displays the autocorrelation	required
figure_tau	Figure	Figure that displays the autocorrelation	required
ax_tau	Axes	Axes on which to plot the autocorrelation	required
start	NDArray[float64] | None	Starting point for the sampling. Only required if the sampler has not performed iterations already.	None
sample_kwargs	dict[str, Any] | None	Arguments to pass to sampler.sample.	None
progress	str | bool	Whether to show a progress bar or not. If this is a string, it must match the values supported by sampler.sample.	'notebook'
progress_kwargs	dict[str, Any] | None	Arguments to pass to the progress bar, if used.	None
min_samples_before_plot	int	Minimum number of samples that must be taken before any plot can be made.	2
corner_kwargs	dict[str, Any] | None	Passed to plot_corner.	None
convergence_ratio	float	Ratio to use to check whether the chains have converged or not. Passed to get_autocorrelation_info.	50

Yields:

Type	Description
ChainProgressInfo	Information about the chain's progress. The yield occurs each time the plots are updated, so can also be used as a way to perform other actions.

Source code in src/openscm_calibration/emcee_plotting.py

def plot_emcee_progress(  # noqa: PLR0913
    sampler: emcee.ensemble.EnsembleSampler,
    iterations: int,
    burnin: int,
    thin: int,
    plot_every: int,
    parameter_order: tuple[str, ...],
    neg_log_likelihood_name: str,
    holder_chain: IPython.core.display_functions.DisplayHandle,
    figure_chain: matplotlib.figure.Figure,
    axes_chain: dict[str, matplotlib.axes.Axes],
    holder_dist: IPython.core.display_functions.DisplayHandle,
    figure_dist: matplotlib.figure.Figure,
    axes_dist: dict[str, matplotlib.axes.Axes],
    holder_corner: IPython.core.display_functions.DisplayHandle,
    figure_corner: matplotlib.figure.Figure,
    holder_tau: IPython.core.display_functions.DisplayHandle,
    figure_tau: matplotlib.figure.Figure,
    ax_tau: matplotlib.axes.Axes,
    start: nptype.NDArray[np.float64] | None = None,
    sample_kwargs: dict[str, Any] | None = None,
    progress: str | bool = "notebook",
    progress_kwargs: dict[str, Any] | None = None,
    min_samples_before_plot: int = 2,
    corner_kwargs: dict[str, Any] | None = None,
    convergence_ratio: float = 50,
) -> Iterator[ChainProgressInfo]:
    """
    Plot MCMC progress

    Parameters
    ----------
    sampler
        Sampler used for sampling the posterior distribution

    iterations
        Number of iterations to perform

    burnin
        Burn-in to apply to the chains

    thin
        Thinning to apply to the chains

    plot_every
        How many steps to wait before updating the plots

    parameter_order
        Order of the parameters.

        This is used for ensuring that the plots are labelled correctly.

    neg_log_likelihood_name
        Name to use for the negative log likelihood in plots

    holder_chain
        Holder of the figure that displays the chains

    figure_chain
        Figure that displays the chains

    axes_chain
        Axes on which to plot the chains.

        Each parameter in `parameter_order` should be a key in `axes_chain`.

    holder_dist
        Holder of the figure that displays the parameter distributions

    figure_dist
        Figure that displays the parameter distributions

    axes_dist
        Axes on which to plot the parameter distributions.

        Each parameter in `parameter_order` should be a key in `axes_dist`.

    holder_corner
        Holder of the figure that displays the corner plot

    figure_corner
        Figure that displays the corner plot

    holder_tau
        Holder of the figure that displays the autocorrelation

    figure_tau
        Figure that displays the autocorrelation

    ax_tau
        Axes on which to plot the autocorrelation

    start
        Starting point for the sampling.

        Only required if the sampler has not performed iterations already.

    sample_kwargs
        Arguments to pass to `sampler.sample`.

    progress
        Whether to show a progress bar or not.

        If this is a string, it must match the values supported by `sampler.sample`.

    progress_kwargs
        Arguments to pass to the progress bar, if used.

    min_samples_before_plot
        Minimum number of samples that must be taken before any plot can be made.

    corner_kwargs
        Passed to [`plot_corner`][openscm_calibration.emcee_plotting.plot_corner].

    convergence_ratio
        Ratio to use to check whether the chains have converged or not.

        Passed to
        [`get_autocorrelation_info`][openscm_calibration.emcee_utils.get_autocorrelation_info].

    Yields
    ------
    :
        Information about the chain's progress.

        The yield occurs each time the plots are updated,
        so can also be used as a way to perform other actions.
    """
    if sample_kwargs is None:
        sample_kwargs = {}

    if progress_kwargs is None:
        progress_kwargs = DEFAULT_PROGRESS_KWARGS

    if corner_kwargs is None:
        corner_kwargs = {}

    start_use = get_start(
        sampler=sampler,
        start=start,
    )

    # Stores for autocorrelation values
    # (as a function of the number of steps performed).
    # This reserves more memory than is needed,
    # but it is so cheap for a plain array that we don't worry about it.
    autocorr = np.zeros((iterations, sampler.ndim))
    autocorr_steps = np.zeros(iterations, dtype=np.int64)
    autocorr_index = 0

    # Helper function
    def enough_autocorr_values_to_plot(autocorr: nptype.NDArray[np.float64]) -> bool:
        """Check whether there is any auto-correlation information to plot"""
        values_to_plot_present = np.logical_and(
            autocorr > 0.0,
            np.logical_not(np.isnan(autocorr)),
        )

        return bool(np.any(np.sum(values_to_plot_present, axis=0) > 1))

    for sample in sampler.sample(
        start_use,
        iterations=iterations,
        progress=progress,
        progress_kwargs={"leave": True},
        **sample_kwargs,
    ):
        if (
            sampler.iteration % plot_every
            or sampler.iteration < min_samples_before_plot
        ):
            continue

        for ax in axes_chain.values():
            ax.clear()

        plot_chains(
            inp=sampler,
            burnin=burnin,
            parameter_order=parameter_order,
            axes_d=axes_chain,
            neg_log_likelihood_name=neg_log_likelihood_name,
        )
        figure_chain.tight_layout()
        holder_chain.update(figure_chain)  # type: ignore

        steps_post_burnin = max(0, sampler.iteration - burnin)
        if steps_post_burnin < 1:
            # In burn in
            acceptance_fraction = np.nan

        else:
            chain_post_burnin = sampler.get_chain(discard=burnin)
            acceptance_fraction = float(
                np.mean(get_acceptance_fractions(chain_post_burnin))
            )

            for ax in axes_dist.values():
                ax.clear()

                plot_dist(
                    inp=sampler,
                    burnin=burnin,
                    thin=thin,
                    parameter_order=parameter_order,
                    axes_d=axes_dist,
                    warn_singular=False,
                )

            figure_dist.tight_layout()
            holder_dist.update(figure_dist)  # type: ignore

            # Not sure why this was wrapped in try-except in the notebooks.
            # Might want to re-instate in future.
            # try:
            figure_corner.clear()
            plot_corner(
                inp=sampler,
                burnin=burnin,
                thin=thin,
                parameter_order=parameter_order,
                fig=figure_corner,
                **corner_kwargs,
            )
            figure_corner.tight_layout()
            holder_corner.update(figure_corner)  # type: ignore
            # except AssertionError:
            #     pass

            autocorr_bits = get_autocorrelation_info(
                sampler,
                burnin=burnin,
                thin=thin,
                convergence_ratio=convergence_ratio,
            )

            if autocorr_bits.any_non_nan_tau():
                autocorr[autocorr_index, :] = autocorr_bits.tau
                autocorr_steps[autocorr_index] = autocorr_bits.steps_post_burnin
                autocorr_index += 1

            if enough_autocorr_values_to_plot(autocorr):
                ax_tau.clear()
                plot_tau(
                    ax=ax_tau,
                    autocorr=autocorr[:autocorr_index, :],
                    steps=autocorr_steps[:autocorr_index],
                    parameter_order=parameter_order,
                    convergence_ratio=convergence_ratio,
                    convergence_ratio_line_kwargs=dict(color="k", linestyle="--"),
                )

                figure_tau.tight_layout()
                holder_tau.update(figure_tau)  # type: ignore

        yield ChainProgressInfo(
            steps=sampler.iteration,
            steps_post_burnin=steps_post_burnin,
            acceptance_fraction=acceptance_fraction,
        )

plot_parameter_chains

plot_parameter_chains(
    ax: Axes,
    chain_values: NDArray[number[Any]],
    burnin: int,
    alpha_chain: float = 0.3,
    linewidth: float = 0.5,
    color: str = "0.2",
    alpha_vspan: float = 0.3,
    kwargs_chain: dict[str, Any] | None = None,
    kwargs_vspan: dict[str, Any] | None = None,
) -> Axes

Plot chains for a single parameter in an MCMC run

Parameters:

Name	Type	Description	Default
ax	Axes	Axes on which to plot	required
chain_values	NDArray[number[Any]]	Chain values to plot (should be 2D)	required
burnin	int	Number of iterations to treat as burn in	required
alpha_chain	float	Alpha to use for the chains	0.3
linewidth	float	Linewidth to use for the chains	0.5
color	str	Colour to use for the chains	'0.2'
alpha_vspan	float	Alpha to use for the vertical span (which shows the burnin period)	0.3
kwargs_chain	dict[str, Any] | None	Pass to :meth:ax.plot when plotting the chains	None
kwargs_vspan	dict[str, Any] | None	Pass to :meth:ax.axvspan when plotting the vertical span	None

Source code in src/openscm_calibration/emcee_plotting.py

def plot_parameter_chains(  # noqa: PLR0913
    ax: matplotlib.axes.Axes,
    chain_values: np.typing.NDArray[np.number[Any]],
    burnin: int,
    alpha_chain: float = 0.3,
    linewidth: float = 0.5,
    color: str = "0.2",
    alpha_vspan: float = 0.3,
    kwargs_chain: dict[str, Any] | None = None,
    kwargs_vspan: dict[str, Any] | None = None,
) -> matplotlib.axes.Axes:
    """
    Plot chains for a single parameter in an MCMC run

    Parameters
    ----------
    ax
        Axes on which to plot

    chain_values
        Chain values to plot (should be 2D)

    burnin
        Number of iterations to treat as burn in

    alpha_chain
        Alpha to use for the chains

    linewidth
        Linewidth to use for the chains

    color
        Colour to use for the chains

    alpha_vspan
        Alpha to use for the vertical span (which shows the burnin period)

    kwargs_chain
        Pass to :meth:`ax.plot` when plotting the chains

    kwargs_vspan
        Pass to :meth:`ax.axvspan` when plotting the vertical span
    """
    if not kwargs_chain:
        kwargs_chain = {}

    if not kwargs_vspan:
        kwargs_vspan = {}

    ax.plot(
        chain_values,
        lw=linewidth,
        alpha=alpha_chain,
        color=color,
        **kwargs_chain,
    )
    # If burnin is 2, then we burn the first two steps
    # i.e. 0 and 1 so the span should cover 0 and 1, but
    # not 2 i.e. the span covers burnin -1, but not burnin.
    # To make this visually easier to see, we do a span that
    # ends halfway between the last step in burnin and the first
    # in the chain i.e. we end at burnin - 0.5
    ax.axvspan(0, burnin - 0.5, alpha=alpha_vspan, **kwargs_vspan)
    ax.set_xlim(0, chain_values.shape[0])

    return ax

plot_tau

plot_tau(
    ax: Axes,
    autocorr: NDArray[float64],
    steps: NDArray[int64],
    parameter_order: tuple[str, ...],
    convergence_ratio: float,
    convergence_ratio_line_kwargs: dict[str, Any]
    | None = None,
    legend_loc: str | None = "lower right",
    xlabel: str = "Number steps (post burnin)",
    ylabel: str = "Autocorrelation time, tau",
) -> None

Plot the autocorrelation time, tau

Parameters:

Name	Type	Description	Default
ax	Axes	Axes on which to plot	required
autocorr	NDArray[float64]	Autocorrelation information. This should be multi-dimensional, with each column being the autocorrelation time for a different parameter.	required
steps	NDArray[int64]	The number of steps taken before each row in autocorr was calculated. In other words, the x-axis for the plot.	required
parameter_order	tuple[str, ...]	The order of the parameters in autocorr.	required
convergence_ratio	float	Convergence ratio (used to show the convergence line on the plot)	required
convergence_ratio_line_kwargs	dict[str, Any] | None	Keyword arguments to pass to axline when plotting the convergence ratio line. If not supplied, we use {"color": "k", "linestyle": "--"}.	None
legend_loc	str | None	Location of the legend. If None, we don't explicitly set the legend location.	'lower right'
xlabel	str	x-label to apply to the plot.	'Number steps (post burnin)'
ylabel	str	y-label to apply to the plot.	'Autocorrelation time, tau'

Source code in src/openscm_calibration/emcee_plotting.py

def plot_tau(  # noqa: PLR0913
    ax: matplotlib.axes.Axes,
    autocorr: nptype.NDArray[np.float64],
    steps: nptype.NDArray[np.int64],
    parameter_order: tuple[str, ...],
    convergence_ratio: float,
    convergence_ratio_line_kwargs: dict[str, Any] | None = None,
    legend_loc: str | None = "lower right",
    xlabel: str = "Number steps (post burnin)",
    ylabel: str = "Autocorrelation time, tau",
) -> None:
    """
    Plot the autocorrelation time, tau

    Parameters
    ----------
    ax
        Axes on which to plot

    autocorr
        Autocorrelation information.

        This should be multi-dimensional,
        with each column being the autocorrelation time for a different parameter.

    steps
        The number of steps taken before each row in `autocorr` was calculated.

        In other words, the x-axis for the plot.

    parameter_order
        The order of the parameters in `autocorr`.

    convergence_ratio
        Convergence ratio (used to show the convergence line on the plot)

    convergence_ratio_line_kwargs
        Keyword arguments to pass to [`axline`][matplotlib.axes.Axes.axline]
        when plotting the convergence ratio line.

        If not supplied, we use `{"color": "k", "linestyle": "--"}`.

    legend_loc
        Location of the legend.

        If `None`, we don't explicitly set the legend location.

    xlabel
        x-label to apply to the plot.

    ylabel
        y-label to apply to the plot.
    """
    if convergence_ratio_line_kwargs is None:
        convergence_ratio_line_kwargs = {"color": "k", "linestyle": "--"}

    for i, parameter in enumerate(parameter_order):
        ax.scatter(
            steps,
            autocorr[:, i],
            label=parameter,
        )

    ax.axline((0, 0), slope=1 / convergence_ratio, **convergence_ratio_line_kwargs)

    if legend_loc is not None:
        ax.legend(loc=legend_loc)

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)