failed to save the state of the chain to a file "chain.dat"

General information:

• emcee version: 3.0.2
• platform: Ubuntu, python3
• installation method (pip/conda/source/other?): pip3

Problem description:

I want to save the state of the chain to a file "chain.dat" as the method you introduced in the web
https://emcee.readthedocs.io/en/v2.2.1/user/advanced/
but I failed.

Below is my code:

import numpy as np

# Choose the "true" parameters.
m_true = -0.9594
b_true = 4.294
f_true = 0.534

# Generate some synthetic data from the model.
N = 50
x = np.sort(10*np.random.rand(N))
yerr = 0.1+0.5*np.random.rand(N)
y = m_true*x+b_true
y += np.abs(f_true*y) * np.random.randn(N)
y += yerr * np.random.randn(N)

A = np.vstack((np.ones_like(x), x)).T
C = np.diag(yerr * yerr)
cov = np.linalg.inv(np.dot(A.T, np.linalg.solve(C, A)))
b_ls, m_ls = np.dot(cov, np.dot(A.T, np.linalg.solve(C, y)))


def lnlike(theta, x, y, yerr):
    m, b, lnf = theta
    model = m * x + b
    inv_sigma2 = 1.0/(yerr**2 + model**2*np.exp(2*lnf))
    return -0.5*(np.sum((y-model)**2*inv_sigma2 - np.log(inv_sigma2)))


import scipy.optimize as op
nll = lambda *args: -lnlike(*args)
result = op.minimize(nll, [m_true, b_true, np.log(f_true)], args=(x, y, yerr))
m_ml, b_ml, lnf_ml = result["x"]

def lnprior(theta):
    m, b, lnf = theta
    if -5.0 < m < 0.5 and 0.0 < b < 10.0 and -10.0 < lnf < 1.0:
        return 0.0
    return -np.inf

def lnprob(theta, x, y, yerr):
    lp = lnprior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + lnlike(theta, x, y, yerr)

ndim, nwalkers = 3, 100
pos = [result["x"] + 1e-4*np.random.randn(ndim) for i in range(nwalkers)]

import emcee
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(x, y, yerr))

f = open("chain.dat", "w")
f.close()

for res in sampler.sample(pos, iterations=500, store=False):
    position = res[0]
    f = open("chain.dat", "a")
    for k in range(position.shape[0]):
        f.write("{0:4d} {1:s}\n".format(k, " ".join(position[k])))
    f.close()


# sample code goes here...

When I run the above code, I get the following error message:

TypeError Traceback (most recent call last)
in
1 for res in sampler.sample(pos, iterations=500, store=False):
----> 2 position = res[0]
3 f = open("chain.dat", "a")
4 for k in range(position.shape[0]):
5 f.write("{0:4d} {1:s}\n".format(k, " ".join(position[k])))

TypeError: 'State' object is not subscriptable

Start at the sampler.sample loop in the provided example and inspect the returned State object at the failing res[0] line. Compare the emcee 3.0.2 sampler API with the older documentation, then run the example to verify that chain.dat is written successfully for all iterations.