In [1]:

from IPython.display import Image, display, Math

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, Math import matplotlib.pyplot as pl from matplotlib.ticker import MaxNLocator import numpy as np import reddemcee np.random.seed(1234)

Fitting a model to data¶

We are going to fit a simple sinusoid model, with some white noise. This model has 4 parameters, 3 for the sinusoid plus the jitter related to the white noise:

In [2]:

# True parameters
K_ = 55
P_ = 4.1

j_ = 5  # white noise
ndim_ = 3

theta_ = [K_, P_, j_]

# We generate the observation time stamps
N_ = 100
baseline = 6
x_ = np.sort(np.random.uniform(0, baseline, N_))

# And their values
# Sinusoid
def sinusoid(theta):
    K, P = theta[:2]
    freq = 2. * np.pi / P
    y = K * np.cos(freq*x_)
    return y
    
y_ = sinusoid(theta_)

# True parameters K_ = 55 P_ = 4.1 j_ = 5 # white noise ndim_ = 3 theta_ = [K_, P_, j_] # We generate the observation time stamps N_ = 100 baseline = 6 x_ = np.sort(np.random.uniform(0, baseline, N_)) # And their values # Sinusoid def sinusoid(theta): K, P = theta[:2] freq = 2. * np.pi / P y = K * np.cos(freq*x_) return y y_ = sinusoid(theta_)

Now we need to add some error to our observations:

In [3]:

# we assume a measured error of N(1, 7)
yerr_ = abs(1 + 7 * np.random.rand(N_))

# we add it to our observations
y_ += yerr_ * np.random.randn(N_)

# and we add an error not accounted for in the measurements
y_ += j_ * np.random.randn(N_)

# we assume a measured error of N(1, 7) yerr_ = abs(1 + 7 * np.random.rand(N_)) # we add it to our observations y_ += yerr_ * np.random.randn(N_) # and we add an error not accounted for in the measurements y_ += j_ * np.random.randn(N_)

In [4]:

# Visualise data
pl.errorbar(x_, y_, yerr=yerr_, fmt=".k", capsize=0)

# true model
x0 = np.linspace(0, baseline, 500)
y0 = K_ * np.cos((2. * np.pi / P_)*x0)

pl.plot(x0, y0, "k", alpha=0.3, lw=3)
pl.xlabel("x")
pl.ylabel("y");

# Visualise data pl.errorbar(x_, y_, yerr=yerr_, fmt=".k", capsize=0) # true model x0 = np.linspace(0, baseline, 500) y0 = K_ * np.cos((2. * np.pi / P_)*x0) pl.plot(x0, y0, "k", alpha=0.3, lw=3) pl.xlabel("x") pl.ylabel("y");

Using the sampler¶

The sampler requires a likelihood and a prior function.

The likelihood is modeled with the residuals from the observations and the model, as well as the jitter parameter.

We will use uniform priors for every parameter.

We also need to define some limits for the prior volume (our search space):

In [5]:

limits = [[10, 100],
          [0.1, 8],
          [np.log(1), np.log(20)],]


def loglike(theta):
    model = sinusoid(theta)
    jitter = theta[-1]
    sigma2 = yerr_**2 + np.exp(2 * jitter)
    return -0.5 * np.sum((y_ - model) ** 2 / sigma2 + np.log(sigma2)) - 0.5 * np.log(2*np.pi) * N_


def logprior(theta):
    lp = 0.
    for p in range(ndim_):
        if limits[p][0] <= theta[p] <= limits[p][1]:
            lp += np.log(1 / (limits[p][1] - limits[p][0]))
        else:
            return -np.inf

    return lp

limits = [[10, 100], [0.1, 8], [np.log(1), np.log(20)],] def loglike(theta): model = sinusoid(theta) jitter = theta[-1] sigma2 = yerr_**2 + np.exp(2 * jitter) return -0.5 * np.sum((y_ - model) ** 2 / sigma2 + np.log(sigma2)) - 0.5 * np.log(2*np.pi) * N_ def logprior(theta): lp = 0. for p in range(ndim_): if limits[p][0] <= theta[p] <= limits[p][1]: lp += np.log(1 / (limits[p][1] - limits[p][0])) else: return -np.inf return lp

We need to choose our sampler setup and a starting position:

In [6]:

setup = [4, 100, 1000, 2]
ntemps, nwalkers, nsweeps, nsteps = setup

setup = [4, 100, 1000, 2] ntemps, nwalkers, nsweeps, nsteps = setup

In [7]:

p0 = np.array([theta_[:ndim_-1]+[np.log(theta_[-1])] + 0.1 * np.random.randn(nwalkers, ndim_) for _ in range(ntemps)])

p0 = np.array([theta_[:ndim_-1]+[np.log(theta_[-1])] + 0.1 * np.random.randn(nwalkers, ndim_) for _ in range(ntemps)])

And we run the sampler:

In [8]:

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

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

100%|██████████| 8000/8000 [00:17<00:00, 444.48it/s]

We can take a look at the full chains:

In [9]:

# Retrieve chains
ch = sampler.get_chain(flat=True, discard=0)
ch0 = ch[0]


# Visualise chains
fig, axes = pl.subplots(ndim_, figsize=(10, 7), sharex=True)

labels = ['K', 'P', 'Jitter']

for i in range(ndim_):
    ax = axes[i]
    chain = ch[0][:, i]
    ax.plot(chain, 'k', alpha=0.3)
    ax.set_ylabel(labels[i])
    #ax.yaxis.set_label_coords(-0.1, 0.5)
    ax.hlines(np.mean(chain), 0, len(chain), label=f'Mean = {np.mean(chain):.2f}', ls='--')
    ax.legend(loc=1, fontsize=8)
axes[-1].set_xlabel('N step');

# Retrieve chains ch = sampler.get_chain(flat=True, discard=0) ch0 = ch[0] # Visualise chains fig, axes = pl.subplots(ndim_, figsize=(10, 7), sharex=True) labels = ['K', 'P', 'Jitter'] for i in range(ndim_): ax = axes[i] chain = ch[0][:, i] ax.plot(chain, 'k', alpha=0.3) ax.set_ylabel(labels[i]) #ax.yaxis.set_label_coords(-0.1, 0.5) ax.hlines(np.mean(chain), 0, len(chain), label=f'Mean = {np.mean(chain):.2f}', ls='--') ax.legend(loc=1, fontsize=8) axes[-1].set_xlabel('N step');

We discard the initial samples as burn-in, select our cold chain and we corner plot:

In [10]:

import corner
flat_samples = sampler.get_chain(flat=True, discard=500)[0]

import corner flat_samples = sampler.get_chain(flat=True, discard=500)[0]

In [11]:

fig = corner.corner(flat_samples, labels=labels, truths=theta_[:ndim_-1]+[np.log(theta_[-1])])

fig = corner.corner(flat_samples, labels=labels, truths=theta_[:ndim_-1]+[np.log(theta_[-1])])

We can use as point estimate the median and the 16 and 84 percentiles:

In [12]:

for i in range(ndim_):
    res = np.percentile(flat_samples[:, i], [16, 50, 84])
    q = np.diff(res)
    display(Math(rf'{labels[i]} = {res[0]:.3f}^{{ + {q[0]:.3f} }}_{{- {q[1]:.3f} }}'))

for i in range(ndim_): res = np.percentile(flat_samples[:, i], [16, 50, 84]) q = np.diff(res) display(Math(rf'{labels[i]} = {res[0]:.3f}^{{ + {q[0]:.3f} }}_{{- {q[1]:.3f} }}'))

$\displaystyle K = 54.109^{ + 1.004 }_{- 1.018 }$

$\displaystyle P = 4.087^{ + 0.013 }_{- 0.013 }$

$\displaystyle Jitter = 1.517^{ + 0.117 }_{- 0.114 }$