The first toy model

Contents

The first toy model#

from pathlib import Path
import numpy as np
from mpi4py import MPI
import tjsampler as ts
from tjsampler.constants import TWOPI

TMP_DIR = Path("../.tmp")

In this tutotial, we will demonstrate the basic usage of Sampler by a toy model: generating multivariate normal posterior with regard to \(x,\,y\). The covariance is

\[\begin{split} \begin{bmatrix}1 & 1/2 \\ 1/2 & 1/2\end{bmatrix}. \end{split}\]

We wiil adopt a uniform prior and hence the likelihood takes the form of posterior.

# ============== MPI ===============
nwalker = 4
ntemp = 2
comm = MPI.COMM_WORLD

iproc = comm.Get_rank()
ts.random.seed(iproc)

# ============== Distribution ================
cov = np.array([[1, .5], [.5, .5]])
log_cov_det = np.log(np.linalg.det(cov))
cov_inv = np.linalg.inv(cov)

The code snippet implies \(2\) temperatue chains and \(4\) walkers are involved. To ensure independent sampling between processes, we must use different seeds to initialize the ramdom variable generator, which are the process IDs in this simple case. (A better method is numpy SeedSequence.)

# ============== MCMC ================
priors = ts.prior.Prior({"x": ts.prior.Uniform(-3, 3), "y": ts.prior.Uniform(-3, 3)})

class ToyLikelihood:
    def calc_log_likelihood(self, state):
        params = state.params
        ans = -.5 * params @ cov_inv @ params - .5 * log_cov_det - np.log(TWOPI)
        return ans
    
toylikelihood = ToyLikelihood()

move = ts.move.StretchMove(2, nwalker, comm=comm.Dup())
temp = ts.temp.Temperature(ntemp)

The Likelihood should implement calc_log_likelihood method. The type of parameter passed in is defined in state module. We use the affine invariant stretch move is this demo. This proposal need to transfer message between processes and hence we copy the communicator to avoid interference with temperature swap.

nstep = 1000
nburn = 500
thin = 20
outfile = TMP_DIR / "toy_model.h5"
params0 = priors.draw_sample()
sampler = ts.sampler.Sampler(priors, toylikelihood, move, nwalker, temp, 
                             outfile=outfile, comm=comm)
sampler.run(params0, nstep, nburn, thin)

The result will save to a .hdf5 file. The user may manipulate it with h5py. Alternatively, Result module encapsulate some functions. For example,

filename = "../examples/toy_model.h5"
result = tj.result.Result(filename)

samples = result.samples_params()
np.cov(samples[:, 0], samples[:, 1]) #[[1.04255178, 0.51451805], [0.51451805, 0.53013195]])

_ = result.plot_dist_params()
_images/toy_model.pdf

Notes about CLI#

This demo is simple enough to run with

mpiexec -n 8 python toy_model.py

directly. While cluster is often eqiupped with a job scheduling system, we take Slurm as an example. The program can be run by submitting a script like

#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks=8
#SBATCH --cpus-per-task=1

mpiexec -n 8 python -u toy_model.py

Each MCMC chain is a mpi process by design of the package. It is common for the number of processes to exceed the usable cores of CPU. If this happens, ignore the ntask and OVERSUBSCRIB your hosts:

#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4

mpiexec --mca mpi_yield_when_idle 1 --map-by :OVERSUBSCRIB 8 python -u toy_model.py
Contents