geoffryan / afterglowpy Goto Github PK

When I try to fit the afterglow data with afterglowpy and emcee, it always raises AutocorrError. I've tried 10000 and 40000 walkers, only setting 18 chain, but the chain is always shorter than 50 times the integrated autocorrelation time. I have no idea whats wrong with my code. Moreover, The code runs for 10 hours with 40000 walkers, even if I use pool for parallelization. Is it normal for this code to run such a long time? My code is as follows, looking forward to any helpful advice.
`import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize
import emcee
import corner
import afterglowpy as grb
from multiprocessing import Pool
#Jet Parameters
Z = {'jetType': grb.jet.Gaussian, # Gaussian jet
'specType': 0, # Basic Synchrotron Emission Spectrum
'thetaObs': 0.4, # Viewing angle in radians
'E0': 1.0e53, # Isotropic-equivalent energy in erg
'thetaCore': 0.066, # Half-opening angle in radians
'thetaWing': 0.47, # Outer truncation angle
'n0': 1.0e-3, # circumburst density in cm^{-3}
'p': 2.168, # electron energy distribution index
'epsilon_e': 0.04, # epsilon_e
'epsilon_B': 1e-4, # epsilon_Bt
'xi_N': 1.0, # Fraction of electrons accelerated
'd_L': 1.23e26, # Luminosity distance in cm
'z': 0.01} # redshift
log_f = 0.5

global t
global f_obs
global sigma_obs
global nuX
global tp

#import data
tday = np.array([9, 16, 18, 1.1e2, 1.4e2, 1.6e2, 2.5e2, 3.6e2])
f_obs = np.array([1.3e-7, 1.8e-7, 1.3e-7, 7e-7, 6e-7, 5e-7, 4.5e-7, 3.4e-7])
sigma_obs = np.array([1e-7, 0.6e-7, 0.7e-7, 2e-7, 2e-7, 3e-7, 2e-7, 0.5e-7])
t = tday * grb.day2sec
nuX = 1.2e18

plt.errorbar(tday, f_obs, yerr=sigma_obs, fmt='.',capsize=0)
ta = 1e-1 * grb.day2sec
tb = 1e3 * grb.day2sec
t0 = np.geomspace(ta, tb, 100)
t0day = t0 * grb.sec2day
#get the position of f_cal
f_cal = np.zeros(8)
tp = np.array([48, 55, 56, 75, 78, 79, 84, 88])

plt.plot(t0day, grb.fluxDensity(t0, nuX, **Z))
plt.plot(tday, f_cal, '.r')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('t(d)')
plt.ylabel(r'$F_\nu$(mJy)')

def log_likelihood(theta):
Z['thetaObs'], log_E0, Z['thetaCore'], Z['thetaWing'], log_n0, Z['p'], log_epsilon_e, log_epsilon_B, log_f = theta
if 0.01 < Z['thetaObs'] < 0.8 and 50 < log_E0 < 56 and 0.01 < Z['thetaCore'] < 0.5 and 0.01 < Z['thetaWing'] < min(12Z['thetaCore'],np.pi/2) and -4 < log_n0 < 0 and 2 < Z['p'] < 2.5 and -5 < log_epsilon_e < 0 and -5 < log_epsilon_B < 0 and -4 < log_f < 1:
Z['E0'] = 10**(log_E0)
Z['n0'] = 10**(log_n0)
Z['epsilon_e']=10**(log_epsilon_e)
Z['epsilon_B']=10**(log_epsilon_B)
fnu = grb.fluxDensity(t0, nuX, Z)
f_cal = fnu[tp]
model = f_cal - f_obs
sigma2 = sigma_obs2 + np.exp(2log_f)(f_cal)**2
return -0.5np.sum(model**2/sigma2 + np.log(sigma2))
else:
return -np.inf

def log_prior(theta):
Z['thetaObs'], log_E0, Z['thetaCore'], Z['thetaWing'], log_n0, Z['p'], log_epsilon_e, log_epsilon_B, log_f = theta
if 0.01 < Z['thetaObs'] < 0.8 and 50 < log_E0 < 56 and 0.01 < Z['thetaCore'] < 0.5 and 0.01 < Z['thetaWing'] < min(12Z['thetaCore'],np.pi/2) and -4 < log_n0 < 0 and 2 < Z['p'] < 2.5 and -5 < log_epsilon_e < 0 and -5 < log_epsilon_B < 0 and -4 < log_f < 1:
return np.log(np.sin(Z['thetaObs'])) - 0.5((np.cos(Z['thetaObs']) - 0.985)/0.07)**2
else:
return -np.inf

def log_probability(theta):
lp = log_prior(theta)
if not np.isfinite(lp):
return -np.inf
else:
return lp + log_likelihood(theta)
if name == 'main':
#fit with scipy.minimize
np.random.seed(42)
initial = np.array([Z['thetaObs'], np.log10(Z['E0']), Z['thetaCore'], Z['thetaWing'], np.log10(Z['n0']), Z['p'], np.log10(Z['epsilon_e']),np.log10(Z['epsilon_B']), log_f ]) + 0.1*np.random.randn(9)
nll = lambda args: -log_likelihood(args)
soln = minimize(nll, initial)
Z['thetaObs'], log_E0, Z['thetaCore'], Z['thetaWing'], log_n0, Z['p'], log_epsilon_e, log_epsilon_B, log_f = soln.x
Z['E0'] = 10(log_E0)
Z['n0'] = 10**(log_n0)
Z['epsilon_e']=10**(log_epsilon_e)
Z['epsilon_B']=10**(log_epsilon_B)
plt.plot(t0day, grb.fluxDensity(t0, nuX, **Z), '--')
print('Maximum likelihood estimate:')
print(Z)
print(log_f)

#fit with MCMC
pos = soln.x + 1e-4 * np.random.randn(32,9)
nwalkers, ndim = pos.shape
with Pool() as pool:
    sampler = emcee.EnsembleSampler(nwalkers, ndim, log_probability,
    moves=[(emcee.moves.DEMove(), 0.8), (emcee.moves.DESnookerMove(), 0.2)],
    pool=pool)
    sampler.run_mcmc(pos, 40000, progress=True)
    
    labels = ['thetaObs','log_E0','thetaCore','thetaWing','log_n0','p','log_epsilin_e','log_epsilon_B','log_f']
    tau = sampler.get_autocorr_time()
    print(tau)
    thin = 0.5 * np.average(tau)
    discard = 3 * np.average(tau) 
    flat_samples = sampler.get_chain(discard=discard, thin=thin, flat=True)  
    fig = corner.corner(flat_samples, labels=labels)
    plt.savefig('1.png')
    #print the result of parameters
    for i in range(ndim):
        mcmc = np.percentile(flat_samples[:,i], [16, 50, 84])
        q=np.diff(mcmc)
        print(labels[i])
        print('median=', mcmc[i])
        print('diff=', q)    

`
The error is as follows

hello i have a problem using afterglowpy with emcee, i don't know how to define a certain type of jet model，If the model cannot be defined, there is no way to set the prior, likelihood function, etc. Please help me, thank you!

afterglowpy is used broadly enough where installation from conda-forge would be beneficial. The dependencies are simple, so a simple conda feedstock is likely to do it.

We are looking forward to an new version that supports the CSM scenarios, ie. n ∝ r^(-k).

between line 360 and line 361 in flux.py file, there should be q=args[6]