Plotting and Visualization

def cornerplot(results, **kwargs):
    """
    Create corner plot showing posterior distributions and correlations.
    
    Parameters:
    - results (dict): Results from nested sampling run
    - **kwargs: Additional plotting options including:
        - parameters (list): Subset of parameters to plot
        - labels (list): Custom parameter labels
        - show_titles (bool): Display parameter statistics as titles
        - title_quantiles (list): Quantiles for title statistics
        - color (str): Color scheme for plots
        - bins (int): Number of histogram bins
        - smooth (bool): Apply smoothing to contours
        - contour_levels (list): Contour levels to display
        - range (list): Plot ranges for each parameter
        - figsize (tuple): Figure size
        - dpi (int): Resolution for saved plots
    
    Returns:
    matplotlib figure: Corner plot figure
    """

import matplotlib.pyplot as plt
from ultranest.plot import cornerplot

# Create corner plot from results
fig = cornerplot(
    results,
    labels=['Amplitude', 'Mean', 'Sigma'],
    show_titles=True,
    title_quantiles=[0.16, 0.5, 0.84],
    smooth=True,
    bins=30
)

plt.savefig('posterior_corner.png', dpi=300, bbox_inches='tight')
plt.show()

def runplot(results, **kwargs):
    """
    Create run diagnostic plots showing sampling progress and convergence.
    
    Parameters:
    - results (dict): Results from nested sampling run
    - **kwargs: Plotting options including:
        - figsize (tuple): Figure size
        - dpi (int): Plot resolution
        - show_information (bool): Display information content evolution
        - show_evidence (bool): Display evidence estimation
        - show_weights (bool): Display point weights
    
    Returns:
    matplotlib figure: Run diagnostic plots
    """

from ultranest.plot import runplot

# Create convergence diagnostic plots
fig = runplot(
    results,
    figsize=(12, 8),
    show_information=True,
    show_evidence=True
)

plt.savefig('convergence_diagnostics.png', dpi=300, bbox_inches='tight')
plt.show()

def traceplot(results, **kwargs):
    """
    Create trace plots showing parameter evolution during sampling.
    
    Parameters:
    - results (dict): Results from nested sampling run
    - **kwargs: Plotting options including:
        - parameters (list): Parameters to include in trace plot
        - labels (list): Custom parameter labels
        - figsize (tuple): Figure size
        - thin (int): Thinning factor for trace lines
        - alpha (float): Transparency for trace lines
    
    Returns:
    matplotlib figure: Trace plot figure
    """

from ultranest.plot import traceplot

# Create parameter trace plots
fig = traceplot(
    results,
    parameters=['amplitude', 'mean', 'sigma'],
    labels=['Amplitude', 'Mean', 'Sigma'],
    figsize=(10, 6),
    thin=10,
    alpha=0.7
)

plt.savefig('parameter_traces.png', dpi=300, bbox_inches='tight')
plt.show()

def highest_density_interval_from_samples(
    xsamples, 
    xlo=None, 
    xhi=None, 
    probability_level=0.68
):
    """
    Compute highest density interval from samples.
    
    Parameters:
    - xsamples (array): Posterior samples
    - xlo (float, optional): Lower bound for samples
    - xhi (float, optional): Upper bound for samples  
    - probability_level (float): Credible interval level (default: 0.68)
    
    Returns:
    tuple: (lower_bound, upper_bound) of highest density interval
    """

class PredictionBand:
    """
    Class for creating prediction bands from posterior samples.
    
    Provides methods for computing credible intervals and prediction
    bands from posterior predictive distributions.
    """
    
    def __init__(self, x, y_samples):
        """
        Initialize prediction band computation.
        
        Parameters:
        - x (array): Independent variable values
        - y_samples (array): Posterior predictive samples, shape (n_samples, n_points)
        """
    
    def compute_bands(self, alpha_levels=[0.68, 0.95]):
        """
        Compute prediction bands at specified confidence levels.
        
        Parameters:
        - alpha_levels (list): Confidence levels for bands
        
        Returns:
        dict: Band boundaries for each confidence level
        """

import matplotlib.pyplot as plt
from ultranest import ReactiveNestedSampler
from ultranest.plot import cornerplot, runplot, traceplot

# Run nested sampling
sampler = ReactiveNestedSampler(
    param_names=['amplitude', 'mean', 'sigma'],
    loglike=loglike,
    transform=prior_transform
)

results = sampler.run()

# Create comprehensive visualization suite
plt.style.use('seaborn-v0_8')  # Set plotting style

# 1. Corner plot for posterior analysis
fig1 = cornerplot(
    results,
    labels=['A', r'$\mu$', r'$\sigma$'],
    show_titles=True,
    title_quantiles=[0.16, 0.5, 0.84],
    smooth=True,
    bins=30,
    figsize=(10, 10)
)
fig1.savefig('posterior_corner.png', dpi=300, bbox_inches='tight')

# 2. Run diagnostics for convergence assessment
fig2 = runplot(
    results,
    figsize=(12, 8),
    show_information=True,
    show_evidence=True
)
fig2.savefig('run_diagnostics.png', dpi=300, bbox_inches='tight')

# 3. Parameter traces for sampling quality
fig3 = traceplot(
    results,
    labels=['Amplitude', 'Mean', 'Sigma'],
    figsize=(12, 8),
    thin=5
)
fig3.savefig('parameter_traces.png', dpi=300, bbox_inches='tight')

plt.show()

import numpy as np
import matplotlib.pyplot as plt
from ultranest.plot import highest_density_interval_from_samples

# Extract posterior samples
samples = results['samples']
weights = results['weights']

# Weighted resampling for equal-weight samples
from ultranest.utils import resample_equal
equal_weight_samples = resample_equal(samples, weights)

# Custom analysis
for i, param_name in enumerate(['amplitude', 'mean', 'sigma']):
    param_samples = equal_weight_samples[:, i]
    
    # Compute statistics
    median = np.median(param_samples)
    hdi_lower, hdi_upper = highest_density_interval_from_samples(
        param_samples, probability_level=0.68
    )
    
    print(f"{param_name}: {median:.3f} [{hdi_lower:.3f}, {hdi_upper:.3f}]")
    
    # Custom histogram
    plt.figure(figsize=(8, 5))
    plt.hist(param_samples, bins=50, alpha=0.7, density=True)
    plt.axvline(median, color='red', linestyle='--', label='Median')
    plt.axvline(hdi_lower, color='orange', linestyle=':', label='68% HDI')
    plt.axvline(hdi_upper, color='orange', linestyle=':')
    plt.xlabel(param_name)
    plt.ylabel('Posterior Density')
    plt.legend()
    plt.title(f'Posterior Distribution: {param_name}')
    plt.savefig(f'{param_name}_posterior.png', dpi=300, bbox_inches='tight')
    plt.show()

# Compare multiple models
models = ['model_A', 'model_B', 'model_C']
evidences = [results_A['logz'], results_B['logz'], results_C['logz']]
evidence_errors = [results_A['logzerr'], results_B['logzerr'], results_C['logzerr']]

# Evidence comparison plot
plt.figure(figsize=(10, 6))
plt.errorbar(models, evidences, yerr=evidence_errors, 
             fmt='o', capsize=5, capthick=2, markersize=8)
plt.ylabel('Log Evidence')
plt.title('Model Comparison via Bayesian Evidence')
plt.grid(True, alpha=0.3)

# Add Bayes factors
best_evidence = max(evidences)
for i, (model, evidence) in enumerate(zip(models, evidences)):
    bayes_factor = np.exp(evidence - best_evidence)
    plt.text(i, evidence + 0.5, f'BF: {bayes_factor:.2f}', 
             ha='center', fontsize=10)

plt.tight_layout()
plt.savefig('model_comparison.png', dpi=300, bbox_inches='tight')
plt.show()