Temperatures & Adaptation¶

For setting the temperatures we will expand on the [Quickstart] example:

In [13]:

from IPython.display import Image, display

from matplotlib import colormaps
import matplotlib.pyplot as pl
from matplotlib.ticker import MaxNLocator

import numpy as np
import reddemcee

np.random.seed(1234)

from IPython.display import Image, display from matplotlib import colormaps import matplotlib.pyplot as pl from matplotlib.ticker import MaxNLocator import numpy as np import reddemcee np.random.seed(1234)

Definitions¶

We will use the same definitions as in Quickstart:

In [2]:

ndim_ = 2  # n dimensions
r_ = 2.  # radius
w_ = 0.1  # width
hard_limit = 6  # hard search boundary

limits_ = [-hard_limit,  hard_limit]
c1_ = np.zeros(ndim_)
c1_[0] = -3.5
c2_ = np.zeros(ndim_)
c2_[0] = 3.5
const_ = np.log(1. / np.sqrt(2. * np.pi * w_**2))

ndim_ = 2 # n dimensions r_ = 2. # radius w_ = 0.1 # width hard_limit = 6 # hard search boundary limits_ = [-hard_limit, hard_limit] c1_ = np.zeros(ndim_) c1_[0] = -3.5 c2_ = np.zeros(ndim_) c2_[0] = 3.5 const_ = np.log(1. / np.sqrt(2. * np.pi * w_**2))

In [3]:

def logcirc(theta, c):
    # log-likelihood of a single shell
    d = np.sqrt(np.sum((theta - c)**2, axis=-1))  # |theta - c|
    return const_ - (d - r_)**2 / (2. * w_**2)


def loglike(theta):
    # log-likelihood of two shells
    return np.logaddexp(logcirc(theta, c1_), logcirc(theta, c2_))


def logprior(theta):
    # prior for our parameters
    lp = 0.
    for i in range(ndim_):
        if  theta[i] <= limits_[0] or limits_[1] <= theta[i]:
            return -np.inf
    return lp

def logcirc(theta, c): # log-likelihood of a single shell d = np.sqrt(np.sum((theta - c)**2, axis=-1)) # |theta - c| return const_ - (d - r_)**2 / (2. * w_**2) def loglike(theta): # log-likelihood of two shells return np.logaddexp(logcirc(theta, c1_), logcirc(theta, c2_)) def logprior(theta): # prior for our parameters lp = 0. for i in range(ndim_): if theta[i] <= limits_[0] or limits_[1] <= theta[i]: return -np.inf return lp

In [4]:

def display_samples(sampler, temp=0):
    nd = sampler.ndim
    fig, axes = pl.subplots(1, nd, figsize=(8, 2*nd))
    
    samples = sampler.get_chain(flat=True)
    
    for i in range(len(axes)):
        axes[i].hist(samples[temp][:, i], 100, histtype="step", lw=1)
        axes[i].set_xlabel(fr"$\theta_{i}$")
        axes[i].set_ylabel(fr"$p(\theta_{i})$")
    pl.gca().set_yticks([])
    fig.suptitle('Samples')
    
def display_chains(sampler, dens=False, temp=0):
    nd = sampler.ndim
    fig, axes = pl.subplots(nd, 1, sharex=True, figsize=(8, nd*3))
    samples = sampler.get_chain(flat=False)
    for i in range(len(axes)):
        if dens:
             axes[i].plot(samples[temp][:, :, i], marker='o', alpha=0.75, lw=0)
        else:
             axes[i].plot(samples[temp][:, :, i], alpha=0.75, lw=1)
        
        axes[i].yaxis.set_major_locator(MaxNLocator(5))
        axes[i].set_ylabel(fr"$p(\theta_{i})$")
        
    fig.suptitle('Chains')
    fig.supxlabel('Step')

def display_samples(sampler, temp=0): nd = sampler.ndim fig, axes = pl.subplots(1, nd, figsize=(8, 2*nd)) samples = sampler.get_chain(flat=True) for i in range(len(axes)): axes[i].hist(samples[temp][:, i], 100, histtype="step", lw=1) axes[i].set_xlabel(fr"$\theta_{i}$") axes[i].set_ylabel(fr"$p(\theta_{i})$") pl.gca().set_yticks([]) fig.suptitle('Samples') def display_chains(sampler, dens=False, temp=0): nd = sampler.ndim fig, axes = pl.subplots(nd, 1, sharex=True, figsize=(8, nd*3)) samples = sampler.get_chain(flat=False) for i in range(len(axes)): if dens: axes[i].plot(samples[temp][:, :, i], marker='o', alpha=0.75, lw=0) else: axes[i].plot(samples[temp][:, :, i], alpha=0.75, lw=1) axes[i].yaxis.set_major_locator(MaxNLocator(5)) axes[i].set_ylabel(fr"$p(\theta_{i})$") fig.suptitle('Chains') fig.supxlabel('Step')

Setup¶

We will set 10 temperatures and a bit of a longer chain for this example:

In [5]:

ntemps, nwalkers, nsweeps, nsteps = [10, 100, 300, 2]
p0 = np.random.uniform(limits_[0], limits_[1], [ntemps, nwalkers, ndim_])

ntemps, nwalkers, nsweeps, nsteps = [10, 100, 300, 2] p0 = np.random.uniform(limits_[0], limits_[1], [ntemps, nwalkers, ndim_])

Temperatures¶

You can:

• set ntemps=int, to the number of temperatures you want to use.
• set betas=array, to the be the initial temperature ladder.

If betas is not specified, it default to a geometrically spaced ladder, defined in (Vousden et al).

First we will take a look at a 'standard' Parallel tempering, without adaptation:

In [6]:

sampler = reddemcee.PTSampler(nwalkers, ndim_, loglike, logprior,
                              ntemps=ntemps)
    
silent = sampler.run_mcmc(p0, nsweeps, nsteps, progress=True)

sampler = reddemcee.PTSampler(nwalkers, ndim_, loglike, logprior, ntemps=ntemps) silent = sampler.run_mcmc(p0, nsweeps, nsteps, progress=True)

100%|██████████| 6000/6000 [00:07<00:00, 772.50it/s]

We take a quick look at what the samples look like for the cold chain:

In [7]:

display_samples(sampler)

display_samples(sampler)

We would expect to see flatter distributions in hotter chains:

In [8]:

display_samples(sampler, temp=2)

display_samples(sampler, temp=2)

Adaptation¶

We take a look at the temperature adaptation using different configurations, using a linearly spaced initial ladder, to better visualise temperature changes.

When initialising the sampler object PTSampler, you can set:

adapt_tau=float, adaptation decay halflife

adapt_nu=float, adaptation rate (more like inverse coefficient)

adapt_mode=str['NONE', 'SAR','SMD','SGG','GAO','ETL'], for different adaptation schemes:

• 'SAR': Uniform swap acceptance rate
• 'SMD': Swap mean distance
• 'SGG': Small Gaussian gap approximation'
• 'GAO': Gaussian area overlap
• 'ETL': Equalised thermodynamic length
• 'NONE': Skip adaptation

To visualize changes we define a plot function:

In [9]:

cmap_plasma = colormaps['plasma']
def running(arr, window_size=10):
    """Calculate running average of the last n values
    for better plot visualisation."""
    averages = np.zeros_like(arr)
    for i in range(arr.shape[0]):
        start_idx = max(0, i - window_size + 1)  # Start of the window (avoid negative index)
        averages[i] = np.mean(arr[start_idx:i + 1], axis=0)  # Average of the last `window_size` points
    return averages

def plot_betas_ratios(sampler, setup):
    bh = sampler.get_betas()
    rh = sampler.get_tsw()

    colors = cmap_plasma(np.linspace(0, 0.95, setup[0]))

    fig, axes = pl.subplots(2, 1, figsize=(9, 5), sharex=True)

    for t in range(setup[0]-2):
        # PLOT BETAS
        y_bet = bh[t]
        axes[0].plot(1/y_bet, c=colors[t])

        # PLOT TS_ACCEPTANCE
        x0 = np.arange(setup[2]) * setup[3]
        y_tsw = running(rh[:, t])
        axes[1].plot(x0, y_tsw, alpha=0.75, color=colors[t])

    # plot options
    axes[0].set_ylabel(r"$T$")
    axes[1].set_ylabel(r"$a_{frac}$")

    axes[1].set_xlabel("N Step")

    axes[0].set_xscale('log')
    axes[0].set_yscale('log')

    pl.tight_layout()

cmap_plasma = colormaps['plasma'] def running(arr, window_size=10): """Calculate running average of the last n values for better plot visualisation.""" averages = np.zeros_like(arr) for i in range(arr.shape[0]): start_idx = max(0, i - window_size + 1) # Start of the window (avoid negative index) averages[i] = np.mean(arr[start_idx:i + 1], axis=0) # Average of the last `window_size` points return averages def plot_betas_ratios(sampler, setup): bh = sampler.get_betas() rh = sampler.get_tsw() colors = cmap_plasma(np.linspace(0, 0.95, setup[0])) fig, axes = pl.subplots(2, 1, figsize=(9, 5), sharex=True) for t in range(setup[0]-2): # PLOT BETAS y_bet = bh[t] axes[0].plot(1/y_bet, c=colors[t]) # PLOT TS_ACCEPTANCE x0 = np.arange(setup[2]) * setup[3] y_tsw = running(rh[:, t]) axes[1].plot(x0, y_tsw, alpha=0.75, color=colors[t]) # plot options axes[0].set_ylabel(r"$T$") axes[1].set_ylabel(r"$a_{frac}$") axes[1].set_xlabel("N Step") axes[0].set_xscale('log') axes[0].set_yscale('log') pl.tight_layout()

In [10]:

my_betas = np.linspace(1, 0, ntemps)
setup = ntemps, nwalkers, nsweeps, nsteps

my_betas = np.linspace(1, 0, ntemps) setup = ntemps, nwalkers, nsweeps, nsteps

In [11]:

sampler = reddemcee.PTSampler(nwalkers, ndim_, loglike, logprior,
                              betas=my_betas,
                              adapt_tau=600,
                              adapt_nu=0.3,
                              adapt_mode='SAR')

silent = sampler.run_mcmc(p0, nsweeps, nsteps, progress=True)

sampler = reddemcee.PTSampler(nwalkers, ndim_, loglike, logprior, betas=my_betas, adapt_tau=600, adapt_nu=0.3, adapt_mode='SAR') silent = sampler.run_mcmc(p0, nsweeps, nsteps, progress=True)

  0%|          | 0/6000 [00:00<?, ?it/s]

100%|██████████| 6000/6000 [00:07<00:00, 774.59it/s]

We can see how the temperatures vary over time, until their acceptance fractions converge to the same value.

In [12]:

plot_betas_ratios(sampler, setup)

plot_betas_ratios(sampler, setup)

I highly encourage playing with the three adaptive parameters in this example to check how they behave.