ganesh 0.26.0

Minimization and sampling in Rust, simplified
Documentation
#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "corner",
#     "matplotlib",
#     "numpy",
# ]
# ///
import pickle
from pathlib import Path

import corner
import matplotlib.pyplot as plt
import numpy as np

if __name__ == '__main__':
    with Path('data.pkl').open('rb') as f:
        (chains, taus) = pickle.load(f)
    chains = np.array(
        [[[step[0][i] for i in range(5)] for step in chain] for chain in chains]
    )
    plt.plot(np.arange(len(taus)) * 50, taus)
    plt.title('Integrated Autocorrelation Time')
    plt.xlabel('Step')
    plt.ylabel(r'Mean $\tau$')
    plt.savefig('iat.svg')
    plt.close()
    _, ax = plt.subplots(nrows=5, sharex=True, figsize=(10, 50))
    steps = np.arange(chains.shape[1])
    burn_in = int(taus[-1] * 50)
    for i in range(chains.shape[2]):
        for j in range(chains.shape[0]):
            ax[i].plot(steps[burn_in:], chains[j, burn_in:, i], color='k', alpha=0.1)
            ax[i].plot(
                steps[:burn_in], chains[j, :burn_in, i], color='k', ls='--', alpha=0.1
            )
        ax[i].plot(steps[burn_in:], chains[0, burn_in:, i], color='r', label='Walker 0')
        ax[i].plot(
            steps[:burn_in],
            chains[0, :burn_in, i],
            color='r',
            ls='--',
            label='Walker 0 (burn-in)',
        )
        ax[i].set_xlabel('Step')
        ax[i].set_ylabel(f'Parameter {i}')
        ax[i].legend()
    plt.savefig('traces.svg')
    plt.close()
    flat_chain = chains[:, burn_in:, :].reshape(-1, chains.shape[2])
    fig = corner.corner(
        flat_chain,
        labels=[f'Parameter {i}' for i in range(5)],
        quantiles=[0.16, 0.5, 0.84],
        show_titles=True,
    )
    plt.savefig('corner_plot.svg')
    plt.close()