Visualization Tools

The visualization module provides tools for plotting MCMC chains and fit results.

Visualization tools module for MCMC chains and fit results.

This module provides specialized plotting functions for visualizing MCMC chains, posterior distributions, and convergence diagnostics. These functions are designed to work seamlessly with ezfit’s Model class and can also be used standalone.

Features

  • Corner plots for posterior distributions and parameter correlations

  • Trace plots for chain convergence visualization

  • Posterior distribution histograms with percentile markers

  • Arviz integration for advanced MCMC diagnostics

  • Support for both 2D and 3D chain formats

ezfit.visualization.plot_arviz_summary(chain: ndarray, param_names: list[str] | None = None, **kwargs: dict[str, Any]) Figure | None[source]

Create arviz summary plots for MCMC chain.

This function uses arviz for advanced MCMC diagnostics and visualization.

Parameters:
  • chain (np.ndarray) – MCMC chain array of shape (n_walkers, n_steps, n_params) or (n_samples, n_params).

  • param_names (list[str] | None, optional) – List of parameter names for labels. If None, uses “param0”, “param1”, etc., by default None.

  • **kwargs (dict[str, Any]) – Additional keyword arguments passed to arviz plotting functions.

Returns:

matplotlib Figure if arviz is available, None otherwise.

Return type:

Figure | None

ezfit.visualization.plot_corner(chain: ndarray, param_names: list[str] | None = None, fig: Figure | None = None, **kwargs: dict[str, Any]) tuple[Figure, ndarray][source]

Create a corner plot from MCMC chain.

Parameters:
  • chain (np.ndarray) – MCMC chain array of shape (n_samples, n_params) or (n_walkers, n_steps, n_params).

  • param_names (list[str] | None, optional) – List of parameter names for labels. If None, uses “param0”, “param1”, etc., by default None.

  • fig (Figure | None, optional) – Existing matplotlib Figure to plot on. If None, creates new figure, by default None.

  • **kwargs (dict[str, Any]) – Additional keyword arguments passed to corner.corner().

Returns:

Tuple of (figure, axes_array).

Return type:

tuple[Figure, np.ndarray]

Raises:
  • ImportError – If corner package is not installed.

  • ValueError – If chain shape is invalid or param_names length doesn’t match chain dimensions.

ezfit.visualization.plot_posterior(chain: ndarray, param_names: list[str] | None = None, fig: Figure | None = None, axes: np.ndarray | None = None, bins: int = 50) tuple[Figure, ndarray][source]

Create posterior distribution plots for MCMC chain.

Parameters:
  • chain (np.ndarray) – MCMC chain array of shape (n_samples, n_params) or (n_walkers, n_steps, n_params).

  • param_names (list[str] | None, optional) – List of parameter names for labels. If None, uses “param0”, “param1”, etc., by default None.

  • fig (Figure | None, optional) – Existing matplotlib Figure to plot on. If None, creates new figure, by default None.

  • axes (np.ndarray | None, optional) – Existing axes array to plot on. If None, creates new axes, by default None.

  • bins (int, optional) – Number of bins for histograms, by default 50.

Returns:

Tuple of (figure, axes_array).

Return type:

tuple[Figure, np.ndarray]

Raises:

ValueError – If chain shape is invalid or param_names length doesn’t match chain dimensions.

ezfit.visualization.plot_trace(chain: ndarray, param_names: list[str] | None = None, fig: Figure | None = None, axes: np.ndarray | None = None) tuple[Figure, ndarray][source]

Create trace plots for MCMC chain.

Parameters:
  • chain (np.ndarray) – MCMC chain array of shape (n_walkers, n_steps, n_params) or (n_samples, n_params).

  • param_names (list[str] | None, optional) – List of parameter names for labels. If None, uses “param0”, “param1”, etc., by default None.

  • fig (Figure | None, optional) – Existing matplotlib Figure to plot on. If None, creates new figure, by default None.

  • axes (np.ndarray | None, optional) – Existing axes array to plot on. If None, creates new axes, by default None.

Returns:

Tuple of (figure, axes_array).

Return type:

tuple[Figure, np.ndarray]

Raises:

ValueError – If chain shape is invalid or param_names length doesn’t match chain dimensions.

Functions

ezfit.visualization.plot_corner(chain: ndarray, param_names: list[str] | None = None, fig: Figure | None = None, **kwargs: dict[str, Any]) tuple[Figure, ndarray][source]

Create a corner plot from MCMC chain.

Parameters:
  • chain (np.ndarray) – MCMC chain array of shape (n_samples, n_params) or (n_walkers, n_steps, n_params).

  • param_names (list[str] | None, optional) – List of parameter names for labels. If None, uses “param0”, “param1”, etc., by default None.

  • fig (Figure | None, optional) – Existing matplotlib Figure to plot on. If None, creates new figure, by default None.

  • **kwargs (dict[str, Any]) – Additional keyword arguments passed to corner.corner().

Returns:

Tuple of (figure, axes_array).

Return type:

tuple[Figure, np.ndarray]

Raises:
  • ImportError – If corner package is not installed.

  • ValueError – If chain shape is invalid or param_names length doesn’t match chain dimensions.

ezfit.visualization.plot_trace(chain: ndarray, param_names: list[str] | None = None, fig: Figure | None = None, axes: np.ndarray | None = None) tuple[Figure, ndarray][source]

Create trace plots for MCMC chain.

Parameters:
  • chain (np.ndarray) – MCMC chain array of shape (n_walkers, n_steps, n_params) or (n_samples, n_params).

  • param_names (list[str] | None, optional) – List of parameter names for labels. If None, uses “param0”, “param1”, etc., by default None.

  • fig (Figure | None, optional) – Existing matplotlib Figure to plot on. If None, creates new figure, by default None.

  • axes (np.ndarray | None, optional) – Existing axes array to plot on. If None, creates new axes, by default None.

Returns:

Tuple of (figure, axes_array).

Return type:

tuple[Figure, np.ndarray]

Raises:

ValueError – If chain shape is invalid or param_names length doesn’t match chain dimensions.

ezfit.visualization.plot_posterior(chain: ndarray, param_names: list[str] | None = None, fig: Figure | None = None, axes: np.ndarray | None = None, bins: int = 50) tuple[Figure, ndarray][source]

Create posterior distribution plots for MCMC chain.

Parameters:
  • chain (np.ndarray) – MCMC chain array of shape (n_samples, n_params) or (n_walkers, n_steps, n_params).

  • param_names (list[str] | None, optional) – List of parameter names for labels. If None, uses “param0”, “param1”, etc., by default None.

  • fig (Figure | None, optional) – Existing matplotlib Figure to plot on. If None, creates new figure, by default None.

  • axes (np.ndarray | None, optional) – Existing axes array to plot on. If None, creates new axes, by default None.

  • bins (int, optional) – Number of bins for histograms, by default 50.

Returns:

Tuple of (figure, axes_array).

Return type:

tuple[Figure, np.ndarray]

Raises:

ValueError – If chain shape is invalid or param_names length doesn’t match chain dimensions.

ezfit.visualization.plot_arviz_summary(chain: ndarray, param_names: list[str] | None = None, **kwargs: dict[str, Any]) Figure | None[source]

Create arviz summary plots for MCMC chain.

This function uses arviz for advanced MCMC diagnostics and visualization.

Parameters:
  • chain (np.ndarray) – MCMC chain array of shape (n_walkers, n_steps, n_params) or (n_samples, n_params).

  • param_names (list[str] | None, optional) – List of parameter names for labels. If None, uses “param0”, “param1”, etc., by default None.

  • **kwargs (dict[str, Any]) – Additional keyword arguments passed to arviz plotting functions.

Returns:

matplotlib Figure if arviz is available, None otherwise.

Return type:

Figure | None

Examples

Corner Plot

from ezfit.visualization import plot_corner

# After MCMC fit
fig, axes = plot_corner(
    model.sampler_chain,
    param_names=["A", "B", "C"]
)
plt.show()

Trace Plot

from ezfit.visualization import plot_trace

fig, axes = plot_trace(
    model.sampler_chain,
    param_names=["A", "B", "C"]
)
plt.show()

Using Model Methods

# Convenience methods on Model object
model.plot_corner()  # Corner plot
model.plot_trace()   # Trace plots