eryn 1.2.4__py3-none-any.whl → 1.2.6__py3-none-any.whl

eryn/utils/plot.py

@@ -0,0 +1,1393 @@
+# *-- coding: utf-8 --*
+import os
+import numpy as np
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+from matplotlib.patches import Ellipse, Rectangle
+
+from matplotlib.colors import to_rgba
+
+import corner
+import typing
+
+from eryn.utils.updates import UpdateStep
+
+import pandas as pd
+import seaborn as sns
+
+from eryn.utils.utility import stepping_stone_log_evidence, get_integrated_act
+DEFAULT_PALETTE = "icefire"
+
+try:
+    import scienceplots
+    plt.style.use(['science'])
+except (ImportError, ModuleNotFoundError):
+    pass
+
+# increase default font size
+mpl.rcParams.update({'font.size': 16})
+
+class Backend:
+    """A placeholder Backend class for type hinting."""
+    pass
+
+def save_or_show(fig, filename=None):
+    """
+    Save the figure to a file or show it.
+    
+    Args:
+        fig (matplotlib.figure.Figure): Figure to save or show.
+        filename (str, optional): If provided, saves the figure to this filename.
+    """
+    if filename:
+        fig.savefig(filename, dpi=200, bbox_inches='tight')
+        plt.close(fig)
+    else:
+        plt.show()
+
+def cov_ellipse(mean, cov, ax, n_std=1.0, **kwargs):
+    """
+    Plot a covariance ellipse using eigendecomposition.
+
+    The ellipse axes are aligned with the eigenvectors of the covariance matrix,
+    and scaled by sqrt(eigenvalue) * n_std.
+
+    Args:
+        mean (array-like): Center of the ellipse (mean_x, mean_y).
+        cov (np.ndarray): 2x2 covariance matrix.
+        ax (matplotlib.axes.Axes): Axes object on which to plot the ellipse.
+        n_std (float, optional): Number of standard deviations for ellipse radius. Default is 1.0.
+        **kwargs: Additional keyword arguments passed to matplotlib.patches.Ellipse.
+
+    Returns:
+        matplotlib.patches.Ellipse: The covariance ellipse added to the axes.
+    """
+    # Eigendecomposition: eigenvalues are variances along principal axes
+    eigenvalues, eigenvectors = np.linalg.eigh(cov)
+    
+    # Sort by eigenvalue (largest first) for consistent orientation
+    order = eigenvalues.argsort()[::-1]
+    eigenvalues = eigenvalues[order]
+    eigenvectors = eigenvectors[:, order]
+    
+    # Ellipse dimensions: 2 * n_std * sqrt(eigenvalue) for width/height
+    width, height = 2 * n_std * np.sqrt(eigenvalues)
+    
+    # Rotation angle from the first eigenvector (major axis direction)
+    angle = np.degrees(np.arctan2(eigenvectors[1, 0], eigenvectors[0, 0]))
+    
+    ellipse = Ellipse(xy=mean, width=width, height=height, angle=angle, **kwargs)
+    return ax.add_patch(ellipse)
+
+def overlay_fim_covariance(
+    fig,
+    covariance,
+    means=None,
+    nsigmas=[1, 2, 3],
+    plot_1d=False,
+    colors=None,
+    linestyles=None,
+    linewidths=None,
+    alpha=0.7,
+    labels=None,
+):
+    """
+    Overlay Fisher Information Matrix confidence contours on corner plot axes.
+    
+    For 2D subplots, draws elliptical contours at specified confidence levels.
+    For 1D diagonal plots, draws vertical lines at ±nσ from the mean.
+    
+    Args:
+        fig (matplotlib.figure.Figure): Figure object containing corner plot axes.
+        covariance (np.ndarray): Covariance matrix from Fisher analysis, shape (n_params, n_params).
+        means (np.ndarray, optional): Mean values for each parameter. If None, uses origin (0, 0, ...).
+        nsigmas (list, optional): List of sigma levels to plot (e.g., [1, 2, 3]). Default [1, 2, 3].
+        plot_1d (bool, optional): Whether to plot 1D contours on diagonal plots. Default is False.
+        colors (list, optional): Colors for each sigma level. If None, uses default color cycle.
+        linestyles (list, optional): Line styles for each sigma level. If None, uses solid lines.
+        linewidths (list, optional): Line widths for each sigma level. If None, uses default (1.5).
+        alpha (float, optional): Transparency of contours. Default 0.7.
+        labels (list, optional): Labels for legend entries. If None, uses "nσ FIM".
+    
+    Returns:
+        matplotlib.figure.Figure: The figure with overlaid contours.
+    """
+    
+    # Convert axs to numpy array if it's a list
+    axs = np.array(fig.get_axes())
+    
+    # Infer number of parameters from covariance matrix
+    n_params = covariance.shape[0]
+    
+    if covariance.shape != (n_params, n_params):
+        raise ValueError(f"Covariance matrix must be square, got shape {covariance.shape}")
+    
+    # Set default means to origin
+    if means is None:
+        means = np.zeros(n_params)
+    elif len(means) != n_params:
+        raise ValueError(f"means must have length {n_params}, got {len(means)}")
+    
+    # Set default colors
+    if colors is None:
+        colors = [f'C{i}' for i in range(len(nsigmas))]
+    elif isinstance(colors, str):
+        colors = [colors] * len(nsigmas)
+    elif len(colors) != len(nsigmas):
+        raise ValueError(f"colors must have length {len(nsigmas)}, got {len(colors)}")
+    
+    # Set default linestyles
+    if linestyles is None:
+        linestyles = ['-'] * len(nsigmas)
+    elif len(linestyles) != len(nsigmas):
+        raise ValueError(f"linestyles must have length {len(nsigmas)}, got {len(linestyles)}")
+    
+    # Set default linewidths
+    if linewidths is None:
+        linewidths = [1.5] * len(nsigmas)
+    elif len(linewidths) != len(nsigmas):
+        raise ValueError(f"linewidths must have length {len(nsigmas)}, got {len(linewidths)}")
+    
+    # Set default labels
+    if labels is None:
+        labels = [f'{n}$\\sigma$ FIM' for n in nsigmas]
+    elif len(labels) != len(nsigmas):
+        raise ValueError(f"labels must have length {len(nsigmas)}, got {len(labels)}")
+    
+    # Extract standard deviations for 1D plots
+    sigmas = np.sqrt(np.diag(covariance))
+    
+    # Reshape axes into 2D grid if needed
+    if axs.ndim == 1:
+        # Corner plot axes are typically returned as 1D array
+        # Reshape to (n_params, n_params) grid
+        n_axs = int(np.sqrt(len(axs)))
+        axs_grid = np.empty((n_axs, n_axs), dtype=object)
+        idx = 0
+        for i in range(n_axs):
+            for j in range(n_axs):
+                axs_grid[j, i] = axs[idx]
+                idx += 1
+                # if j <= i:
+                #     axs_grid[j, i] = axs[idx]
+                #     idx += 1
+                # else:
+                #     axs_grid[j, i] = None
+    else:
+        axs_grid = axs
+    
+    # Loop over each sigma level
+    for sigma_idx, (n_sigma, color, ls, lw, label) in enumerate(
+        zip(nsigmas, colors, linestyles, linewidths, labels)
+    ):        
+        # Loop over all subplots
+        for i in range(n_params):
+            for j in range(i, n_params):
+                ax = axs_grid[i, j]
+                
+                if ax is None:
+                    continue
+                
+                if i == j:
+                    if plot_1d:
+                        # 1D diagonal plot - draw vertical lines at mean ± n*sigma
+                        mean_val = means[i]
+                        sigma_val = sigmas[i]
+                        
+                        # Get y-limits for vertical lines
+                        ylim = ax.get_ylim()
+                        
+                        # Draw vertical lines
+                        for sign in [-1, 1]:
+                            line = ax.axvline(
+                                mean_val + sign * n_sigma * sigma_val,
+                                color=color,
+                                linestyle=ls,
+                                linewidth=lw,
+                                alpha=alpha,
+                                zorder=10,
+                            )
+                        else:
+                            continue
+                
+                else:
+                    # 2D off-diagonal plot - draw ellipse
+                    # Extract 2x2 subcovariance for parameters j and i
+                    cov = np.array(
+                    (
+                        (covariance[i][i], covariance[i][j]),
+                        (covariance[j][i], covariance[j][j]),
+                    )
+                )
+                # print(cov)
+
+                    mean = np.array((means[i], means[j]))
+                    cov_ellipse(mean, cov, ax, n_std=n_sigma, 
+                                edgecolor=color, facecolor='none', linestyle=ls, linewidth=lw, 
+                                zorder=10, alpha=alpha
+                                )
+    
+    return fig
+
+def cornerplot(data, *args, means=None, overlay_covariance=None, legend_label='Samples', overlay_label='Information Matrix Covariance', filename=None, **kwargs):
+    """
+    Create a corner plot with optional Information Matrix covariance overlay. This is centered around the means if provided.
+    Wrapper around `corner.corner()` that adds Fisher Information Matrix covariance contours.
+
+    Args:
+        data (array-like): Input data for corner plot (e.g., MCMC samples).
+        *args: Positional arguments passed to `corner.corner()`.
+        means (array-like, optional): Mean values for each parameter to center the overlay. If None, uses 'truths' from kwargs or mean of data.
+        overlay_covariance (np.ndarray, optional): Covariance matrix to overlay. If None, no overlay is added.
+        legend_label (str, optional): Label for the sample distribution in the legend. Default is 'Samples'.
+        overlay_label (str, optional): Label for the overlay covariance in the legend. Default is 'Information Matrix Covariance', assuming FIM.
+        filename (str, optional): If provided, saves the figure to this filename.
+        **kwargs: Keyword arguments passed to `corner.corner()`.
+
+    Returns:
+        matplotlib.figure.Figure: The corner plot figure with optional overlays. If `filename` is provided, the figure is saved instead.
+    """
+
+    corner_kwargs = {
+    #'quantiles': [0.16, 0.5, 0.84],
+    'levels': (1 - np.exp(-0.5), 1 - np.exp(-2), 1 - np.exp(-4.5)),
+    'show_titles': True,
+    'title_fmt': '.1e',
+    'title_kwargs': {'fontsize': 12},
+    'label_kwargs': {'fontsize': 14},
+    'hist_kwargs': {'density': True, 'linewidth': 2},
+    'plot_datapoints': False,
+    'fill_contours': True,
+    'color': 'steelblue',
+    'truth_color': 'red',
+    }
+
+    # Update base kwargs with any user-provided kwargs
+    corner_kwargs.update(kwargs)
+
+    # how do we deal with 2d truths for reversible jump?
+    truths = corner_kwargs.get('truths', None)
+    if truths is not None and len(truths.shape) > 1:  
+        corner_kwargs['truths'] = None # add truths later
+
+    fig = corner.corner(data=data, *args, **corner_kwargs)
+
+    # add the other truths as vertical/horizontal lines
+    if truths is not None and len(truths.shape) > 1:  
+        corner_kwargs['truths'] = truths # add back in for legend
+        n_params = data.shape[1]
+        axs = np.array(fig.get_axes()).reshape((n_params, n_params))
+        for truth in truths:
+            for i in range(n_params):
+                ax = axs[i, i]
+                ax.axvline(truth[i], color=corner_kwargs['truth_color'], linestyle='-', linewidth=1)
+
+            for i in range(n_params):
+                for j in range(i):
+                    ax = axs[i, j]
+                    ax.axhline(truth[i], color=corner_kwargs['truth_color'], linestyle='-', linewidth=1)
+                    ax.axvline(truth[j], color=corner_kwargs['truth_color'], linestyle='-', linewidth=1)
+
+    # prepare handles for legend
+    handles = []
+    handles_labels = []
+
+    handles.append(mpl.lines.Line2D([], [], color=corner_kwargs['color'],
+                                    linestyle='-', linewidth=2))
+    handles_labels.append(legend_label)
+    if 'truths' in corner_kwargs and corner_kwargs['truths'] is not None:
+        handles.append(mpl.lines.Line2D([], [], color=corner_kwargs['truth_color'],
+                                        linestyle='-', linewidth=1))
+        handles_labels.append('Truths')
+
+
+    # Overlay covariance contours if provided
+    if overlay_covariance is not None:
+        if means is None:
+            means = kwargs.get('truths', None)
+        
+        if means is None:
+            # If no means or truths provided, default to the mean of the samples
+            means = np.mean(data, axis=0)
+
+        overlay_fim_covariance(
+            fig,
+            overlay_covariance,
+            means=means,
+            plot_1d=False,
+            alpha=0.8,
+            nsigmas=[1, 2 ,3],   
+            colors='k'
+        )
+
+        handles.append(mpl.lines.Line2D([], [], color='k', linestyle='-', linewidth=1.5))
+        handles_labels.append(overlay_label)
+    
+    # Add legend to the second leftmost top subplot
+    ax_legend = fig.get_axes()[1]
+    ax_legend.legend(handles, handles_labels, loc='upper left', fontsize=18)
+    
+    save_or_show(fig, filename)
+
+
+def traceplot(chain, labels=None, truths=None, filename=None):
+    """
+    Create trace plots for MCMC chains.
+
+    Args:
+        chain (np.ndarray): MCMC chain of shape (nsteps, nwalkers, nleaves, ndim).
+        labels (list, optional): List of parameter names for x-axis labels. 
+        truths (array-like, optional): True parameter values to overlay as horizontal lines.
+        filename (str, optional): If provided, saves the figure to this filename.
+
+    Returns:
+        matplotlib.figure.Figure: The trace plot figure. If `filename` is provided, the figure is saved instead.
+    """
+
+    nsteps, nwalkers, nleaves, ndim = chain.shape
+    fig, axs = plt.subplots(ndim, 1, figsize=(10, 2.5 * ndim), sharex=True)
+    
+    for i in range(ndim):
+        for w in range(nwalkers):
+            axs[i].plot(chain[:, w, :, i], alpha=0.5, rasterized=True)
+            if truths is not None:
+                truths = np.atleast_2d(truths)
+                for j in range(truths.shape[0]):
+                    axs[i].axhline(truths[j, i], color='k', linestyle='--')
+            if labels is not None:
+                axs[i].set_ylabel(labels[i])
+
+    axs[-1].set_xlabel('Step')
+    plt.tight_layout()
+
+    save_or_show(fig, filename)
+
+
+def plot_loglikelihood(logl, filename=None):
+    nsteps, nwalkers = logl.shape
+    fig = plt.figure(figsize=(10, 6))
+    for j in range(nwalkers):
+        plt.plot(logl[:, j], color=f"C{j % 10}", alpha=0.8, rasterized=True)
+    plt.xlabel("Sampler Iteration")
+    plt.ylabel("Log-Likelihood")
+    
+    save_or_show(fig, filename)
+
+    # plot a facet grid of loglikelihood evolution  for each walker
+    # Reshape: logl is (nsteps, nwalkers), need to flatten properly
+    max_logl = np.max(logl, axis=(0,1))
+    facet_logl = logl - max_logl
+
+    step = np.tile(range(nsteps), nwalkers)
+    walker = np.repeat(np.arange(nwalkers, dtype=int), nsteps)
+
+    df = pd.DataFrame(np.c_[facet_logl.flat, step, walker],
+                      columns=[r"$\Delta \log\mathcal{L}$", "step", "walker"])
+    
+    # Initialize a grid of plots with an Axes for each walker
+    grid = sns.FacetGrid(df, col="walker", hue="walker", #palette="tab20c",
+                     col_wrap=int(np.floor(np.sqrt(nwalkers))), height=1.5)
+    
+    # Draw a line plot to show the trajectory of each random walk
+    grid.map(plt.plot, "step", r"$\Delta \log\mathcal{L}$", marker=".", rasterized=True)
+
+    grid.refline(y=0, linestyle=":") # Add a horizontal reference line at y=0 ~ average loglikelihood at each step
+
+    # Adjust the arrangement of the plots
+    grid.set_titles(col_template="Walker {col_name:.0f}")
+    grid.set_axis_labels("Step", r"$\Delta \log\mathcal{L}$")
+    # Disable tight_layout to avoid warning
+    grid.tight_layout = lambda *args, **kwargs: None
+    grid.tight_layout()
+
+    # add overall title
+    plt.subplots_adjust(top=0.9)
+    grid.figure.suptitle(r"$\Delta \log\mathcal{L}_w = \log\mathcal{L}_w - \max(\log\mathcal{L})$", fontsize=16)
+
+    save_or_show(grid.figure, filename.replace('.png', '_facet.png') if filename else None)
+
+def tempering_ridgeplot(chain, labels=None, palette=None, 
+                        bw_adjust=0.5, aspect=5, height=0.5, hspace=-0.25,
+                        max_samples=10000, filename=None):
+    """
+    Create ridge plots of tempered distributions using overlapping KDE plots for all parameters.
+
+    This creates a visually appealing ridge plot (also known as joy plot) showing
+    how the posterior distribution broadens at higher temperatures. Each temperature
+    level is shown as a separate row with overlapping density estimates.
+    All parameters are shown as columns in a single FacetGrid figure.
+
+    Args:
+        chain (np.ndarray): MCMC chain of shape (nsteps, ntemps, nwalkers, nleaves, ndim).
+        labels (list, optional): List of parameter names. If provided, uses labels for column titles.
+        palette (str or list, optional): Seaborn color palette name or list of colors. 
+            Default uses cubehelix_palette.
+        bw_adjust (float, optional): Bandwidth adjustment factor for KDE. Default is 0.5.
+        aspect (float, optional): Aspect ratio of each facet. Default is 5.
+        height (float, optional): Height of each temperature row in inches. Default is 0.5.
+        hspace (float, optional): Vertical spacing between temperature rows (negative for overlap). Default is -0.25.
+        max_samples (int, optional): Maximum number of samples per temperature for KDE. Default is 5000.
+            Subsampling speeds up KDE computation for large chains.
+        filename (str, optional): If provided, saves figure to this filename.
+        
+    Returns:
+        matplotlib.figure.Figure: The figure containing all ridge plots.
+            If `filename` is provided, the figure is saved instead.
+    """
+    # Use seaborn context manager to temporarily set theme
+    with sns.axes_style("white", {"axes.facecolor": (0, 0, 0, 0)}):
+
+        nsteps, ntemps, nwalkers, nleaves, ndim = chain.shape
+        
+        # Create color palette (blue=cold/β=1 at top, red=hot/β→0 at bottom)
+        if palette is None:
+            # Use coolwarm reversed: blue for cold (β=1), red for hot (β→0)
+            pal = sns.color_palette(DEFAULT_PALETTE, ntemps)
+        elif isinstance(palette, str):
+            pal = sns.color_palette(palette, ntemps)
+        else:
+            pal = palette
+        
+        # Subsampling RNG
+        rng = np.random.default_rng(42)
+        
+        # Build dataframe with samples from all temperatures and parameters
+        data_list = []
+        for param_idx in range(ndim):
+            param_label = labels[param_idx] if labels is not None else fr'$x_{param_idx}$'
+            for t in range(ntemps):
+                # Flatten samples across steps, walkers, leaves for the selected parameter
+                samples = chain[:, t, :, :, param_idx].reshape(-1)
+                
+                # Remove NaNs
+                samples = samples[~np.isnan(samples)]
+                
+                # Subsample if needed for faster KDE
+                if len(samples) > max_samples:
+                    samples = rng.choice(samples, size=max_samples, replace=False)
+                
+                temp_label = rf"$\beta_{{{t}}}$" #fr"$\beta$={betas[t]:.1e}"
+                data_list.append(pd.DataFrame({
+                    'x': samples,
+                    'temp': temp_label,
+                    'temp_idx': t,
+                    'param': param_label,
+                    'param_idx': param_idx
+                }))
+        
+        df = pd.concat(data_list, ignore_index=True)
+        
+        # Get unique temps in order (beta=1 first at top)
+        temp_order = df.drop_duplicates('temp_idx').sort_values('temp_idx', ascending=True)['temp'].tolist()
+        
+        # Get unique params in order
+        param_order = df.drop_duplicates('param_idx').sort_values('param_idx', ascending=True)['param'].tolist()
+        
+        # Compute x-axis limits from the cold posterior (β=1, temp_idx=0) for each parameter
+        # This ensures the cold posterior is always visible even when hot distributions are much wider
+        xlims = {}
+        cold_df = df[df['temp_idx'] == 0]
+        for param in param_order:
+            param_data = cold_df[cold_df['param'] == param]['x']
+            q_low, q_high = param_data.quantile([0.001, 0.999])
+            margin = (q_high - q_low) * 0.3  # Add 30% margin to show some broadening
+            xlims[param] = (q_low - margin, q_high + margin)
+        
+        # Initialize the FacetGrid with row=temp, col=param
+        # Suppress tight_layout warnings during initialization
+        import warnings
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", message=".*Tight layout.*")
+            g = sns.FacetGrid(df, row="temp", col="param", hue="temp", 
+                            aspect=aspect, height=height, 
+                            palette=pal, row_order=temp_order, col_order=param_order,
+                            sharex=False, sharey=False)
+        
+        # Disable tight_layout to avoid warning with negative hspace
+        g.tight_layout = lambda *args, **kwargs: None
+        
+        # Custom plotting function that clips KDE to the parameter's xlim
+        def plot_kde_clipped(x, color, label, **kwargs):
+            ax = plt.gca()
+            # Get the parameter for this column
+            col_idx = ax.get_subplotspec().colspan.start
+            param = param_order[col_idx]
+            clip = xlims[param]
+            
+            sns.kdeplot(x, ax=ax, bw_adjust=bw_adjust, clip_on=False,
+                        fill=True, alpha=1, linewidth=1.5, color=color, clip=clip)
+            sns.kdeplot(x, ax=ax, clip_on=False, color="w", lw=2, bw_adjust=bw_adjust, clip=clip)
+            ax.set_xlim(clip)
+        
+        g.map(plot_kde_clipped, "x")
+        
+        # Add horizontal reference line at y=0
+        g.refline(y=0, linewidth=2, linestyle="-", color=None, clip_on=False)
+        
+        # Label function for temperature labels (only on first column)
+        def label_temp(x, color, label):
+            ax = plt.gca()
+            # Only add label if this is the first column
+            if ax.get_subplotspec().colspan.start == 0:
+                ax.text(-0.2, 0.2, label, fontweight="bold", color=color,
+                        ha="left", va="center", transform=ax.transAxes, fontsize=14)
+        
+        g.map(label_temp, "x")
+        
+        # Set subplots to overlap vertically
+        g.figure.subplots_adjust(hspace=hspace, wspace=0.1)
+
+        # use parameter labels in the x axes (on the last row)
+        for ax, param in zip(g.axes[-1], param_order):
+            ax.set_xlabel(param, fontsize=14)
+
+        # Remove xticks from all rows except the bottom
+        for row_idx in range(len(temp_order) - 1):
+            for col_idx in range(len(param_order)):
+                g.axes[row_idx, col_idx].set_xticks([])
+
+                # set xlims again to ensure consistency
+                param = param_order[col_idx]
+                g.axes[row_idx, col_idx].set_xlim(xlims[param])
+        
+        # Remove axes details that don't work well with overlap
+        g.set_titles("")
+        g.set(yticks=[], ylabel="")
+        g.despine(bottom=True, left=True)
+        
+        # Add overall title
+        g.figure.suptitle('Tempered Distributions', y=1.02, fontsize=14)
+        
+        save_or_show(g.figure, filename)
+
+    # # Properly restore the original seaborn style
+    # sns.set_theme(style=original_style['style'], rc=original_style)
+    
+    # return g.figure
+
+def plot_swap_acceptance(swap_acceptance_fraction, palette=None, filename=None):
+    """
+    Plot the temperature swap acceptance fraction between adjacent temperature levels.
+
+    Args:
+        swap_acceptance_fraction (np.ndarray): Swap acceptance fraction between adjacent 
+            temperatures, shape (ntemps-1,). Element i corresponds to swaps between 
+            temperature i and i+1.
+        palette (str or list, optional): Seaborn color palette name or list of colors.
+        filename (str, optional): If provided, saves the figure to this filename.
+
+    Returns:
+        matplotlib.figure.Figure: The swap acceptance plot figure.
+    """
+    ntemps = swap_acceptance_fraction.shape[0] + 1
+    
+    fig, ax = plt.subplots(figsize=(8, 5))
+    
+    # X-axis: temperature pair indices (swap between temp i and i+1)
+    x = np.arange(ntemps - 1)
+    
+    # Create labels for each swap pair
+    labels = [fr'{i}$\leftrightarrow${i+1}' for i in range(ntemps - 1)]
+    
+    # Color by temperature (use coolwarm to match ridgeplot)
+    palette = palette if palette is not None else DEFAULT_PALETTE
+    colors = sns.color_palette(palette, ntemps - 1)
+    
+    # Bar plot
+    bars = ax.bar(x, swap_acceptance_fraction, color=colors, edgecolor='white', linewidth=0.5)
+    
+    # Add horizontal reference line at common target acceptance (0.2-0.4 is often good)
+    ax.axhline(y=0.25, color='gray', linestyle='--', linewidth=1, alpha=0.7, label='0.25')
+    
+    # Add beta values as secondary labels
+    # ax2 = ax.twiny()
+    # ax2.set_xlim(ax.get_xlim())
+    # ax2.set_xticks(x)
+    # beta_labels = [fr'$\beta$={betas[i]:.1e}$\leftrightarrow${betas[i+1]:.1e}' for i in range(ntemps - 1)]
+    # ax2.set_xticklabels(beta_labels, fontsize=8, rotation=45, ha='left')
+    
+    # Labels and formatting
+    ax.set_xticks(x)
+    ax.set_xticklabels(labels, fontsize=10, rotation=45, ha='right')
+    ax.set_xlabel('Temperature Pair (index)', fontsize=12)
+    ax.set_ylabel('Swap Acceptance Fraction', fontsize=12)
+    ax.set_title('Temperature Swap Acceptance', fontsize=14)
+    ax.set_ylim(0, 1.1)
+    ax.legend(loc='upper left', fontsize=10)
+    
+    # Add value annotations on bars
+    for bar, val in zip(bars, swap_acceptance_fraction):
+        height = bar.get_height()
+        ax.annotate(f'{val:.2f}',
+                    xy=(bar.get_x() + bar.get_width() / 2, height),
+                    xytext=(0, 3),  # 3 points vertical offset
+                    textcoords="offset points",
+                    ha='center', va='bottom', fontsize=9)
+    
+    plt.tight_layout()
+    
+    save_or_show(fig, filename)
+    
+    return fig
+    
+def plot_logl_betas(betas: np.ndarray, 
+                    logl: np.ndarray, 
+                    palette: str = None,
+                    filename: str = None
+                    ):
+        """
+        Plots the evolution of log-likelihood values for each temperature.
+
+        Args:
+            betas (numpy.ndarray): Array of inverse temperatures.
+            logl (numpy.ndarray): Array of log-likelihood values.
+            palette (str, optional): Seaborn color palette name or list of colors.
+            filename (str, optional): If provided, saves the figure to this filename.
+        Returns:
+            None
+        """
+        fig = plt.figure(figsize=(10, 6))
+        ntemp = betas.shape[1]
+        tempcolors = sns.color_palette(palette if palette is not None else DEFAULT_PALETTE, ntemp)
+        for temp in range(ntemp):
+            plt.semilogx(betas[-1, temp], np.mean(logl[:, temp]), '.', c=tempcolors[temp], label=f'$T_{temp}$')
+
+        logZ, dlogZ = stepping_stone_log_evidence(betas[-1], logl)
+        
+        plt.ylabel(r'$<\log{\mathcal{L}}>_{\beta}$')
+        plt.xlabel(r'$\beta$')
+        plt.title(r'$\log{\mathcal{Z}} = %.2f \pm %.2f$' % (logZ, dlogZ))
+        
+        save_or_show(fig, filename)
+
+def plot_betas_evolution(betas: np.ndarray, palette: str = None, filename: str = None):
+    """
+    Plots the evolution of inverse temperatures (betas) over sampling steps.
+
+    Args:
+        betas (numpy.ndarray): Array of inverse temperatures of shape (nsteps, ntemps).
+        palette (str, optional): Seaborn color palette name or list of colors.
+        filename (str, optional): If provided, saves the figure to this filename.
+    Returns:
+        None
+    """
+    nsteps, ntemps = betas.shape
+    tempcolors = sns.color_palette(palette if palette is not None else DEFAULT_PALETTE, ntemps)
+    fig = plt.figure(figsize=(10, 6))
+    for temp in range(ntemps):
+        plt.plot(range(nsteps), betas[:, temp], color=tempcolors[temp], linewidth=1.5, alpha=0.8, rasterized=True)
+    plt.xlabel('Sampler Iteration')
+    plt.ylabel(r'Inverse Temperature ($\beta$)')
+    plt.title('Evolution of Inverse Temperatures')
+    
+    # Create temperature color gradient legend
+    
+    ax = plt.gca()
+    legend_width = 0.15
+    legend_height = 0.03
+    legend_x = 0.75
+    legend_y = 0.9
+    
+    # Create gradient patches
+    for i, color in enumerate(tempcolors[::-1]):  # Reverse to show cold->hot left to right
+        rect = Rectangle(
+            (legend_x + i * legend_width / ntemps, legend_y),
+            legend_width / ntemps, legend_height,
+            transform=ax.transAxes,
+            facecolor=to_rgba(color, 0.7),
+            edgecolor='none'
+        )
+        ax.add_patch(rect)
+    
+    # Add border and labels
+    border = Rectangle(
+        (legend_x, legend_y), legend_width, legend_height,
+        transform=ax.transAxes,
+        facecolor='none',
+        edgecolor='black',
+        linewidth=0.5
+    )
+    ax.add_patch(border)
+    
+    # Add text labels
+    ax.text(legend_x - 0.01, legend_y + legend_height / 2, r'$T_{\rm max}$',
+            transform=ax.transAxes, ha='right', va='center', fontsize=11, fontweight='normal', antialiased=True)
+    ax.text(legend_x + legend_width + 0.01, legend_y + legend_height / 2, r'$T_0$',
+            transform=ax.transAxes, ha='left', va='center', fontsize=11, fontweight='normal', antialiased=True)
+    
+    save_or_show(fig, filename)
+
+
+# RJ plots
+def plot_leaves(nleaves: np.ndarray,
+                nleaves_min: int,
+                nleaves_max: int,
+                palette: str = None,
+                iteration: int = 0,
+                filename: str = None):
+    """
+    Plot the histogram of the number of leaves for each temperature.
+
+    This method plots a histogram of the number of leaves for each temperature in the `rj_branches` dictionary.
+    It uses the `self.backend` object to get the number of leaves for each temperature.
+    The histogram is plotted using the `plt.hist` function from the `matplotlib.pyplot` module.
+    The plot includes temperature-specific colors and a legend for the colors.
+
+    Returns:
+    None
+    """
+    bns = (np.arange(nleaves_min, nleaves_max + 2) - 0.5)
+    ntemps = nleaves.shape[1]
+    tempcolors = sns.color_palette(palette if palette is not None else DEFAULT_PALETTE, ntemps)
+
+    fig = plt.figure(figsize=(8, 5))
+
+    for temp, tempcolor in enumerate(tempcolors):
+        plt.hist(nleaves[:, temp].flatten(), bins=bns, histtype="stepfilled", edgecolor=tempcolor,
+                    facecolor=to_rgba(tempcolor, 0.2), density=True, ls='-', zorder=100 - temp, rasterized=True)
+
+    plt.xlabel('Number of leaves')
+    plt.ylabel('Density')
+    
+    # Create temperature color gradient legend
+    # Use a horizontal gradient showing cold (blue) to hot (red)
+    
+    # Add a color bar as legend showing temperature progression
+    ax = plt.gca()
+    legend_width = 0.15
+    legend_height = 0.03
+    legend_x = 0.75
+    legend_y = 0.9
+    
+    # Create gradient patches
+    n_gradient = len(tempcolors)
+    for i, color in enumerate(tempcolors[::-1]):  # Reverse to show cold->hot left to right
+        rect = Rectangle(
+            (legend_x + i * legend_width / n_gradient, legend_y),
+            legend_width / n_gradient, legend_height,
+            transform=ax.transAxes,
+            facecolor=to_rgba(color, 0.7),
+            edgecolor='none'
+        )
+        ax.add_patch(rect)
+    
+    # Add border and labels
+    border = Rectangle(
+        (legend_x, legend_y), legend_width, legend_height,
+        transform=ax.transAxes,
+        facecolor='none',
+        edgecolor='black',
+        linewidth=0.5
+    )
+    ax.add_patch(border)
+    
+    # Add text labels
+    ax.text(legend_x - 0.01, legend_y + legend_height / 2, r'$T_{\rm max}$',
+            transform=ax.transAxes, ha='right', va='center', fontsize=11, fontweight='normal', antialiased=True)
+    ax.text(legend_x + legend_width + 0.01, legend_y + legend_height / 2, r'$T_0$',
+            transform=ax.transAxes, ha='left', va='center', fontsize=11, fontweight='normal', antialiased=True)
+
+    fig.text(0.07, 0.08, f"Step: {iteration}", ha='left', va='top', fontfamily='serif', c='k')
+    #plt.title(key)
+    save_or_show(fig, filename)
+
+def plot_leaves_evolution(nleaves: np.ndarray,
+                          filename: str = None):
+    """
+    Plot the evolution of the number of leaves per walker in the cold chain over sampling steps.
+    Args:
+        nleaves (np.ndarray): Array of number of leaves in the cold chain, shape (nsteps, nwalkers).
+        filename (str, optional): If provided, saves the figure to this filename.
+    Returns:
+        None
+    """
+    nsteps, nwalkers = nleaves.shape
+    fig = plt.figure(figsize=(10, 6))
+    for w in range(nwalkers):
+        plt.plot(range(nsteps), nleaves[:, w], color=f"C{w % 10}", linewidth=1.5, alpha=0.8, rasterized=True)
+    plt.xlabel('Sampler Iteration')
+    plt.ylabel('Number of Leaves')
+    plt.title('Evolution of Number of Leaves in Cold Chain')
+    save_or_show(fig, filename)
+
+def plot_acceptance_fraction(steps: typing.Union[np.ndarray, list],
+                            total_acceptance_fraction: np.ndarray,
+                             moves_acceptance_fraction: dict,
+                             filename: str = None):
+    """
+    Plot the acceptance fraction for different moves over sampling steps.
+
+    Args:
+
+    """
+
+    fig = plt.figure(figsize=(10, 6))
+    # cold chain total acceptance fraction
+    plt.plot(steps, total_acceptance_fraction[:, 0].mean(axis=1), label='Total', color='black', linewidth=2)
+    
+    # skip if moves_acceptance_fraction is empty
+    if len(moves_acceptance_fraction) != 0:
+        for move, acc_fraction in moves_acceptance_fraction.items():
+            plt.plot(steps, acc_fraction[:, 0].mean(axis=1), marker='o', label=move)
+
+    plt.axhline(y=0.234, color='gray', linestyle='--', linewidth=1, alpha=0.7, label='0.234')
+    plt.legend()
+    plt.xlabel('Sampler Iteration')
+    plt.ylabel('Acceptance Fraction')  
+    plt.title('Acceptance Fraction Over Time')
+
+    save_or_show(fig, filename)
+
+def plot_tempered_acceptance_fraction(steps: typing.Union[np.ndarray, list],
+                            total_acceptance_fraction: np.ndarray,
+                            palette: str = None,
+                            filename: str = None):
+    """
+    Plot the acceptance fraction for different moves over sampling steps.
+
+    Args:
+        steps (np.ndarray or list): Array of sampling steps.
+        total_acceptance_fraction (np.ndarray): Array of total acceptance fractions, shape (nsteps, ntemps, nwalkers).
+        palette (str or list, optional): Seaborn color palette name or list of colors.
+        filename (str, optional): If provided, saves the figure to this filename.
+    """
+    ntemps = total_acceptance_fraction.shape[1]
+    tempcolors = sns.color_palette(palette if palette is not None else DEFAULT_PALETTE, ntemps)
+
+    fig = plt.figure(figsize=(10, 6))
+    
+    for temp in range(ntemps):
+        plt.plot(steps, total_acceptance_fraction[:, temp].mean(axis=1), color=tempcolors[temp], linewidth=1.5, marker='o', alpha=0.8, rasterized=True)
+
+    ax = plt.gca()
+    legend_width = 0.15
+    legend_height = 0.03
+    legend_x = 0.75
+    legend_y = 0.9
+    
+    # Create gradient patches
+    for i, color in enumerate(tempcolors[::-1]):  # Reverse to show cold->hot left to right
+        rect = Rectangle(
+            (legend_x + i * legend_width / ntemps, legend_y),
+            legend_width / ntemps, legend_height,
+            transform=ax.transAxes,
+            facecolor=to_rgba(color, 0.7),
+            edgecolor='none'
+        )
+        ax.add_patch(rect)
+    
+    # Add border and labels
+    border = Rectangle(
+        (legend_x, legend_y), legend_width, legend_height,
+        transform=ax.transAxes,
+        facecolor='none',
+        edgecolor='black',
+        linewidth=0.5
+    )
+    ax.add_patch(border)
+    
+    # Add text labels
+    ax.text(legend_x - 0.01, legend_y + legend_height / 2, r'$T_{\rm max}$',
+            transform=ax.transAxes, ha='right', va='center', fontsize=11, fontweight='normal', antialiased=True)
+    ax.text(legend_x + legend_width + 0.01, legend_y + legend_height / 2, r'$T_0$',
+            transform=ax.transAxes, ha='left', va='center', fontsize=11, fontweight='normal', antialiased=True)
+
+    plt.legend()
+    plt.xlabel('Sampler Iteration')
+    plt.ylabel('Acceptance Fraction')  
+
+    ymin, ymax = plt.ylim()
+    plt.ylim(ymin, 1.2 * ymax)
+
+    plt.title('Acceptance Fraction Over Time')
+
+    save_or_show(fig, filename)
+
+def plot_act_evolution(chain: dict,
+                       iteration: int = 0,
+                       parent_folder: str = '.'):
+    
+    """
+    Plot the evolution of the autocorrelation time for each branch in the chain. Also plots the ACT values per parameter in each branch.
+
+    Args:
+        chain (Dict): Dictionary of MCMC chains for different branches.
+        iteration (int, optional): Current iteration number for labeling. Default is 0.
+        parent_folder (str, optional): Folder to save the plots. Default is current directory.
+    """
+
+    NPOINTS = 10
+    points = np.exp(np.linspace(np.log(min(100, iteration)), np.log(iteration), NPOINTS)).astype(int)
+    
+    taus = {}
+    for branch, samples in chain.items():
+        # create branch folder
+        branch_folder = os.path.join(parent_folder, branch)
+        os.makedirs(branch_folder, exist_ok=True)
+
+        nsteps, ntemps, nwalkers, nleaves, ndim = samples.shape
+        cold_chain = samples[:, 0, :, :, :].reshape(nsteps, 1, nwalkers, nleaves * ndim)
+        #cold_chain = cold_chain[~np.isnan(cold_chain).any(axis=-1)] # remove NaNs
+        if np.isnan(cold_chain).any(axis=-1).any():
+            print(f"Skipping ACT plot for branch {branch} due to NaNs in the cold chain.")
+            continue
+        tmp = []
+        for i, point in enumerate(points):
+            if point > cold_chain.shape[0]:
+                continue
+            
+            tau = get_integrated_act(cold_chain[:point], average=True)
+            tmp.append(tau.squeeze())
+
+        taus[branch] = np.array(tmp)
+        
+        fig = plt.figure(figsize=(10, 6))
+        for d in range(ndim):
+            plt.loglog(points[:len(taus[branch])], taus[branch][:, d], marker='o', label=fr'$x_{d}$')
+
+        # add tau = iteration / 50 line for reference
+        xaxis = np.logspace(0, np.log10(iteration), 100)
+        #xaxis[0] = 1  # avoid plotting tau=0 at step 0
+
+        plt.loglog(xaxis, xaxis / 50, linestyle='--', color='k', lw=2, label=r'$\tau$ = Nsteps / 50')
+        plt.xlabel('Number of Steps')
+        plt.ylabel('Integrated Autocorrelation Time')
+        plt.title(f'Autocorrelation Time Evolution - Branch: {branch}')
+
+        ymax = np.max(taus[branch]) * 1.2
+        ymin = np.min(taus[branch]) * 0.5
+        if ymin < ymax:
+            plt.ylim(ymin, ymax)
+        else:
+            plt.ylim(0, ymax)
+        plt.xlim(min(points) * 0.9, iteration * 1.1)
+        plt.legend()
+        save_or_show(fig, os.path.join(branch_folder, f'act_evolution.png'))
+    
+        # plot the maximum ACT across all parameters for each branch
+    if len(taus) > 0:
+        fig = plt.figure(figsize=(10, 6))
+        for branch in taus.keys():
+            max_tau = np.max(taus[branch], axis=1)
+            plt.loglog(points[:len(max_tau)], max_tau, marker='o', label=f'Branch: {branch}')
+        
+        plt.loglog(xaxis, xaxis / 50, linestyle='--', color='k', lw=2, label=r'$\tau$ = Nsteps / 50')
+        plt.xlabel('Number of Steps')
+        plt.ylabel('Maximum Integrated Autocorrelation Time')
+        plt.title('Maximum Autocorrelation Time Evolution Across Branches')
+        ymax = max([np.max(np.max(taus[branch], axis=1)) for branch in taus.keys()]) * 1.2
+        ymin = min([np.min(np.max(taus[branch], axis=1)) for branch in taus.keys()]) * 0.5
+        if ymin < ymax:
+            plt.ylim(ymin, ymax)
+        else:
+            plt.ylim(0, ymax)
+        plt.xlim(min(points) * 0.9, iteration * 1.1)
+        plt.legend()
+        save_or_show(fig, os.path.join(parent_folder, f'max_act_evolution.png'))
+            
+
+def produce_base_plots(chain: dict, 
+                       logl: np.ndarray, 
+                       truths: dict = None, 
+                       overlay_covariance: dict = None, 
+                       labels: dict = None, 
+                       iteration: int = 0,
+                       parent_folder: str = '.',
+                       ):
+
+    """
+    Produce a set of standard diagnostic plot. These include:
+    
+    * corner plots for the cold chain per branch,
+    * trace plots for the cold chain per branch,
+    * log-likelihood evolution plots.
+
+    Args:
+        chain (Dict): Dictionary of MCMC chains for different branches.
+        logl (np.ndarray): Log-likelihood array of shape (nsteps, ntemperatures, nwalkers).
+        truths (Dict, optional): Dictionary of true parameter values for different branches.
+        overlay_covariance (Dict, optional): Dictionary of covariance matrices to overlay on corner plots.
+        labels (Dict, optional): Dictionary of parameter labels for different branches.
+        parent_folder (str, optional): Folder to save the plots. Default is current directory.
+    """
+
+    # create dictionary with a color shade for each branch
+    branches = list(chain.keys())
+    legend_label = 'Samples at step %d' % iteration
+
+    palette = 'Blues'
+    colors = sns.color_palette(palette, n_colors=len(branches))
+    #colors = sns.color_palette("mako", n_colors=len(branches)+2)[2:]  # avoid very light colors
+    branch_colors = {branch: colors[i] for i, branch in enumerate(branches)}
+
+    for branch, samples in chain.items():
+        branch_labels = labels.get(branch, None) if labels else None
+        branch_truths = truths.get(branch, None) if truths else None
+        branch_cov = overlay_covariance.get(branch, None) if overlay_covariance else None
+
+        # create branch folder
+        branch_folder = os.path.join(parent_folder, branch)
+        os.makedirs(branch_folder, exist_ok=True)
+
+        nsteps, ntemps, nwalkers, nleaves, ndim = samples.shape
+        cold_chain = samples[:, 0, :, :, :].reshape(-1, ndim)
+        cold_chain = cold_chain[~np.isnan(cold_chain).any(axis=1)] # remove NaNs
+
+        cornerplot(
+            cold_chain,
+            means=branch_truths,
+            overlay_covariance=branch_cov,
+            legend_label=legend_label,
+            truths=branch_truths,
+            overlay_label='Information Matrix Covariance' if branch_cov is not None else None,
+            labels=branch_labels,
+            color=branch_colors[branch],
+            filename=os.path.join(branch_folder, f'cornerplot.png')
+        )
+
+        traceplot(
+            samples[:, 0, :, :, :],
+            labels=branch_labels,
+            truths=branch_truths,
+            filename=os.path.join(branch_folder, f'traceplot.png')
+        )
+
+        plot_loglikelihood(
+            logl[:, 0, :],
+            filename=os.path.join(parent_folder, f'loglikelihood.png')
+        )
+
+def produce_tempering_plots(chain: dict, 
+                            betas: np.ndarray,
+                            logl: np.ndarray,
+                            swap_acceptance_fraction: np.ndarray,
+                            labels: dict = None, 
+                            parent_folder: str = '.',
+                            palette: str = None
+                            ):
+    """
+    Produce tempering ridge plots for each branch in the chain. These include:
+
+    * ridge plots of the tempered distributions per parameter per branch,
+    * the swap acceptance fraction between adjacent temperatures,
+    * the averaged log-likelihood vs. betas plot,
+    * the evolution of betas over sampling steps.
+
+    Args:
+        chain (Dict): Dictionary of MCMC chains for different branches.
+        betas (np.ndarray): Inverse temperatures of shape (nsteps, ntemps,).
+        swap_acceptance_fraction (np.ndarray): Swap acceptance fraction between adjacent temperatures.
+        labels (Dict, optional): Dictionary of parameter labels for different branches.
+        logl (np.ndarray): Log-likelihood array of shape (nsteps, ntemperatures, nwalkers).
+        parent_folder (str, optional): Folder to save the plots. Default is current directory.
+        palette (str or list, optional): Seaborn color palette name or list of colors.
+    """
+
+    for branch, samples in chain.items():
+        branch_labels = labels.get(branch, None) if labels else None
+
+        # create branch folder
+        branch_folder = os.path.join(parent_folder, branch)
+        os.makedirs(branch_folder, exist_ok=True)
+
+        tempering_ridgeplot(
+            samples,
+            labels=branch_labels,
+            palette=palette,
+            filename=os.path.join(branch_folder, f'tempering_ridgeplot.png')
+        )
+
+    plot_swap_acceptance(
+        swap_acceptance_fraction,
+        palette=palette,
+        filename=os.path.join(parent_folder, f'swap_acceptance.png')
+    )
+
+    plot_logl_betas(
+        betas,
+        logl,
+        palette=palette,
+        filename=os.path.join(parent_folder, f'logl_betas.png')
+    )
+
+    plot_betas_evolution(
+        betas,
+        palette=palette,
+        filename=os.path.join(parent_folder, f'betas_evolution.png')
+    )
+
+def produce_advanced_plots(steps: typing.Union[np.ndarray, list],
+                           total_acceptance_fraction: np.ndarray,
+                           moves_acceptance_fraction: dict,
+                           palette: str = None,
+                           iteration: int = 0,
+                           chain: dict = None,
+                           parent_folder: str = '.'):
+    """
+    Produce advanced diagnostic plots. These include:
+        
+    * autocorrelation time evolution per parameter per branch in the cold chain, 
+    * the comparison of the maximum autocorrelation  time in each branch against the number of steps, 
+    * the acceptance fraction evolution over steps in the cold chain (both overall and per move), 
+    * the overall acceptance fraction evolution over steps per temperature.
+    
+    Args:
+        steps (Union[np.ndarray, list]): Array or list of sampling steps.
+        total_acceptance_fraction (np.ndarray): Total acceptance fraction array of shape (nsteps, ntemps, nwalkers).
+        moves_acceptance_fraction (Dict): Dictionary of acceptance fractions for different moves.
+        parent_folder (str, optional): Folder to save the plots. Default is current directory.
+    """
+
+    plot_acceptance_fraction(
+        steps,
+        total_acceptance_fraction,
+        moves_acceptance_fraction,
+        filename=os.path.join(parent_folder, f'acceptance_fraction.png')
+    )
+
+    plot_tempered_acceptance_fraction(
+        steps,
+        total_acceptance_fraction,
+        palette=palette,
+        filename=os.path.join(parent_folder, f'tempered_acceptance_fraction.png')
+    )
+
+    plot_act_evolution(
+        chain,
+        iteration=iteration,
+        parent_folder=parent_folder
+    )
+
+def produce_rj_plots(nleaves: dict,
+                     nleaves_min: dict,
+                     nleaves_max: dict,
+                     palette: str = None,
+                     parent_folder: str = '.',
+                     iteration: int = 0):
+    
+    """
+    Produce RJ diagnostic plots for each branch in the chain. At present, only plots the histogram of the number of leaves across temperatures.
+    
+    Args:
+        nleaves (Dict): Dictionary of number of leaves arrays for different branches.
+        nleaves_min (Dict): Dictionary of minimum number of leaves for different branches.
+        nleaves_max (Dict): Dictionary of maximum number of leaves for different branches.
+        palette (str or list, optional): Seaborn color palette name or list of colors.
+        parent_folder (str, optional): Folder to save the plots. Default is current directory.
+        iteration (int, optional): Current iteration number for labeling plots.
+    """
+
+    for branch, leaves in nleaves.items():
+        branch_nleaves_min = nleaves_min.get(branch, 1)
+        branch_nleaves_max = nleaves_max.get(branch, 1)
+
+        # skip if there is no reversible-jump sampling for this branch
+        if branch_nleaves_max <= branch_nleaves_min:
+            continue
+        # create branch folder
+        branch_folder = os.path.join(parent_folder, branch)
+        os.makedirs(branch_folder, exist_ok=True)
+        
+        plot_leaves(
+            leaves,
+            branch_nleaves_min,
+            branch_nleaves_max,
+            palette=palette,
+            iteration=iteration,
+            filename=os.path.join(branch_folder, f'rj_nleaves.png')
+        )
+
+        # plot_leaves_evolution(
+        #     leaves[:, 0],
+        #     filename=os.path.join(branch_folder, f'rj_nleaves_evolution.png')
+        # )
+    
+    
+
+
+class PlotContainer:
+    """
+    An Update that generates diagnostic plots at specified intervals
+    
+    Args:
+        plots (list or str): List of plot types to generate. Options are 'base', 'tempering', 'advanced', 'rj', or 'all'.  If multiple plot types are desired, provide a list of strings.
+        branches (list, optional): List of branch names to generate plots for. If None, all branches are used.
+        truths (dict, optional): Dictionary of true parameter values for different branches.
+        overlay_covariance (dict, optional): Dictionary of covariance matrices to overlay on corner plots.
+        tempering_palette (str or list, optional): Seaborn color palette name or list of colors for tempering plots. If None, it defaults to `icefire`.
+        parent_folder (str, optional): Folder to save the plots. Default is current directory.
+        discard (float, optional): Number of initial samples to discard from the chain before plotting. If between 0 and 1, it is treated as fraction of total samples. Default is 0.
+        stop (int, optional): Maximum number of steps to generate plots for. Default is 10000.
+    """
+    
+    def __init__(self, 
+                 backend: Backend = None,
+                 plots : typing.Union[list, str] = 'base',
+                 branches : list = None,
+                 truths: dict = None,
+                 overlay_covariance: dict = None,
+                 tempering_palette: str = None,
+                 parent_folder: str = '.',
+                 discard: float = 0,
+                 stop: int = int(1e4), 
+                 ):
+        """
+        Initialize the PlotContainer.
+        """
+
+        self.backend = backend
+
+        self.parent_folder = parent_folder
+        os.makedirs(self.parent_folder, exist_ok=True)
+
+        allowable_plots = ['base', 'tempering', 'advanced', 'rj']
+        self.branches = branches
+
+        if isinstance(plots, str):
+            if plots == 'all':
+                plots = allowable_plots
+            else:
+                plots = [plots]
+
+        for plot in plots:
+            if plot not in allowable_plots:
+                raise ValueError(f"Plot type '{plot}' not recognized. Allowable types: {allowable_plots}")
+        self.plots = plots
+        
+        #self.labels = labels
+        self.truths = truths
+        self.overlay_covariance = overlay_covariance
+        self.discard = discard
+        self.tempering_palette = tempering_palette
+
+        self.steps = []
+        self.total_acceptance_fraction = None
+        self.move_acceptance_fractions = {}
+
+        self.stop = stop
+
+    @property
+    def backend(self):
+        return self._backend
+    @backend.setter
+    def backend(self, value):
+        self._backend = value
+
+    @property
+    def truths(self):
+        return self._truths
+    @truths.setter
+    def truths(self, value):
+        self._truths = value
+
+    @property
+    def overlay_covariance(self):
+        return self._overlay_covariance
+    @overlay_covariance.setter
+    def overlay_covariance(self, value):
+        self._overlay_covariance = value
+
+    def produce_plots(self, sampler=None) -> None:
+        """
+        Generate diagnostic plots at specified intervals.
+
+        Args:
+            sampler: The sampler object. If not provided, uses self.backend. In this cases not all the plots could be available.
+        Returns:
+            None
+        """
+
+        if self.backend.iteration > self.stop:
+            return
+
+        labels = self.backend.key_order
+    
+        discard = int(self.discard) if self.discard >= 1 else int(self.discard * self.backend.iteration) 
+        chain = self.backend.get_chain(discard=discard)
+        logl = self.backend.get_log_like(discard=discard)
+        betas = self.backend.get_betas(discard=discard)
+
+        if self.branches is not None:
+            chain = {branch: chain[branch] for branch in self.branches if branch in chain}
+            logl = {branch: logl[branch] for branch in self.branches if branch in logl}
+            betas = {branch: betas[branch] for branch in self.branches if branch in betas}  
+
+        for plot in self.plots:
+            base_folder = os.path.join(self.parent_folder, plot)
+            os.makedirs(base_folder, exist_ok=True)    
+
+            if plot == 'base':           
+                produce_base_plots(
+                    chain=chain,
+                    logl=logl,
+                    truths=self.truths,
+                    overlay_covariance=self.overlay_covariance,
+                    iteration=self.backend.iteration,
+                    labels=labels,
+                    parent_folder=base_folder
+                )
+            elif plot == 'tempering':      
+                swap_acceptance_fraction = self.backend.swaps_accepted / float(self.backend.iteration * self.backend.nwalkers)         
+                produce_tempering_plots(
+                    chain=chain,
+                    betas=betas,
+                    logl=logl,
+                    swap_acceptance_fraction=swap_acceptance_fraction,
+                    labels=labels,
+                    parent_folder=base_folder,
+                    palette=self.tempering_palette
+                )
+
+            elif plot == 'advanced':
+                self.steps.append(self.backend.iteration)
+                if self.total_acceptance_fraction is None:
+                    self.total_acceptance_fraction = (self.backend.accepted / float(self.backend.iteration))[np.newaxis, ...]
+                else:
+                    self.total_acceptance_fraction = np.vstack((self.total_acceptance_fraction, (self.backend.accepted / float(self.backend.iteration))[np.newaxis, ...])) # shape (niterations, ntemps, nwalkers)
+                
+                if sampler is not None:
+                    moves = sampler.moves
+                elif hasattr(self.backend, moves):
+                    moves = self.backend.moves
+                else:
+                    moves = None
+
+                if moves is not None:
+                    for move in moves:
+                        name = move.__class__.__name__
+                        if name not in self.move_acceptance_fractions:
+                            self.move_acceptance_fractions[name] = move.acceptance_fraction[np.newaxis, ...]
+                        else:
+                            self.move_acceptance_fractions[name] = np.vstack((self.move_acceptance_fractions[name], move.acceptance_fraction[np.newaxis, ...])) # shape (niterations, ntemps, nwalkers)
+
+                full_chain = self.backend.get_chain(discard=0) if discard > 0 else chain
+                
+                produce_advanced_plots(steps=self.steps,
+                                        total_acceptance_fraction=self.total_acceptance_fraction,   
+                                        moves_acceptance_fraction=self.move_acceptance_fractions,
+                                        palette=self.tempering_palette,
+                                        iteration=self.backend.iteration,
+                                        chain=full_chain,
+                                        parent_folder=base_folder
+                                    )
+
+            elif plot == 'rj':
+                if self.backend.rj is False:
+                    continue
+    
+                nleaves = self.backend.get_nleaves(discard=discard)
+
+                nleaves_min = sampler.nleaves_min if sampler is not None else dict(zip(self.backend.rj_branches, [0]*len(self.backend.rj_branches)))
+                nleaves_max = self.backend.nleaves_max
+
+                produce_rj_plots(
+                    nleaves=nleaves,
+                    nleaves_min=nleaves_min,
+                    nleaves_max=nleaves_max,
+                    palette=self.tempering_palette,
+                    parent_folder=base_folder,
+                    iteration=self.backend.iteration
+                )
+            
+            
+