def get_posterior(z, t0, lc, seed, temp_dir, interped):
my_model = PerfectRedshift([lc], [z], t0, name="posterior%d" % seed)
sampler = EnsembleSampler(temp_dir=temp_dir, num_burn=500, num_steps=1500)
c = my_model.fit(sampler)
chain = c.chains[-1]
parameters = c.parameters[-1]
chain, parameters = add_mu_to_chain(interped, chain, parameters)
return chain, parameters
def get_posterior(z, t0, lc, seed, temp_dir, interped, special):
my_model = PerfectRedshift([lc], [z], t0, name="posterior%d" % seed)
num_steps = 1500
if special:
num_steps *= 10
sampler = EnsembleSampler(temp_dir=temp_dir, num_burn=500, num_steps=num_steps)
c = my_model.fit(sampler)
chain = c.chains[-1]
parameters = c.parameters[-1]
chain, parameters = add_mu_to_chain(interped, chain, parameters)
# for i in range(chain.shape[1]):
# print(parameters[i], skew(chain[:, i]), skewtest(chain[:, i]))
return chain, parameters
def random_obs(temp_dir, seed):
np.random.seed(seed)
interp = generate_and_return()
x1 = np.random.normal()
# colour = np.random.normal(scale=0.1)
colour = 0
x0 = 1e-5
# t0 = np.random.uniform(low=1000, high=2000)
t0 = 1000
z = np.random.uniform(low=0.1, high=1.0)
# deltat = np.random.uniform(low=-20, high=0)
# num_obs = np.random.randint(low=10, high=40)
num_obs = 20
deltat = -35
filename = temp_dir + "/save_%d.npy" % seed
if not os.path.exists(filename):
ts = np.arange(t0 + deltat, (t0 + deltat) + 5 * num_obs, 5)
times = np.array([[t, t + 0.05, t + 0.1, t + 0.2] for t in ts]).flatten()
bands = [b for t in ts for b in ["desg", "desr", "desi", "desz"]]
gains = np.ones(times.shape)
skynoise = np.random.uniform(low=20, high=800) * np.ones(times.shape)
zp = 30 * np.ones(times.shape)
zpsys = ["ab"] * times.size
obs = Table({"time": times, "band": bands, "gain": gains, "skynoise": skynoise, "zp": zp, "zpsys": zpsys})
model = sncosmo.Model(source="salt2")
p = {"z": z, "t0": t0, "x0": x0, "x1": x1, "c": colour}
model.set(z=z)
print(seed, " Vals are ", p)
lc = sncosmo.realize_lcs(obs, model, [p])[0]
ston = (lc["flux"] / lc["fluxerr"]).max()
model.set(t0=t0, x1=x1, c=colour, x0=x0)
try:
res, fitted_model = sncosmo.fit_lc(
lc, model, ["t0", "x0", "x1", "c"], guess_amplitude=False, guess_t0=False
)
except ValueError:
return np.nan, np.nan, x1, colour, num_obs, ston, deltat, z, 0
fig = sncosmo.plot_lc(lc, model=fitted_model, errors=res.errors)
fig.savefig(temp_dir + os.sep + "lc_%d.png" % seed, bbox_inches="tight", dpi=300)
my_model = PerfectRedshift([lc], [z], t0, name="posterior%d" % seed)
sampler = EnsembleSampler(temp_dir=temp_dir, num_burn=400, num_steps=1500)
c = ChainConsumer()
my_model.fit(sampler, chain_consumer=c)
map = {"x0": "$x_0$", "x1": "$x_1$", "c": "$c$", "t0": "$t_0$"}
parameters = [map[a] for a in res.vparam_names]
mu1 = get_mu_from_chain(interped, c.chains[-1], c.parameters[-1])
c.parameteers[-1].append(r"$\mu$")
c.chains[-1] = np.hstack((c.chains[-1], mu1[:, None]))
chain2 = np.random.multivariate_normal(res.parameters[1:], res.covariance, size=int(1e5))
chain2 = np.hstack((chain2, get_mu_from_chain(interp, chain2, parameters)[:, None]))
c.add_chain(chain2, parameters=parameters, name="Gaussian")
figfilename = filename.replace(".npy", ".png")
c.plot(filename=figfilename, truth={"$t_0$": t0, "$x_0$": x0, "$x_1$": x1, "$c$": colour})
means = []
stds = []
isgood = (
(np.abs(x1 - res.parameters[3]) < 4) & (np.abs(colour - res.parameters[4]) < 2) & (res.parameters[2] > 0.0)
)
isgood *= 1.0
for i in range(len(c.chains)):
a = c.chains[i][:, -1]
means.append(a.mean())
stds.append(np.std(a))
diffmu = np.diff(means)[0]
diffstd = np.diff(stds)[0]
np.save(filename, np.array([diffmu, diffstd, ston, 1.0 * isgood]))
else:
vals = np.load(filename)
diffmu = vals[0]
diffstd = vals[1]
ston = vals[2]
isgood = vals[3]
return diffmu, diffstd, x1, colour, num_obs, ston, deltat, z, isgood
model = sncosmo.Model(source='salt2-extended')
p = {'z': z, 't0': t0, 'x0': x0, 'x1': x1, 'c': colour}
model.set(z=z)
print("Realise LCs")
lcs = sncosmo.realize_lcs(obs, model, [p])
print("Fit LCs")
res, fitted_model = sncosmo.fit_lc(lcs[0], model, ['t0', 'x0', 'x1', 'c'])
dir_name = os.path.dirname(__file__)
temp_dir = dir_name + os.sep + "output"
surface = temp_dir + os.sep + "surfaces_simple.png"
mu_simple = temp_dir + os.sep + "mu_simple.png"
mcmc_chain = temp_dir + os.sep + "mcmc_simple.npy"
c = ChainConsumer()
print("Fit model")
my_model = PerfectRedshift(lcs, [z], t0, name="My posterior")
sampler = EnsembleSampler(temp_dir=temp_dir, num_steps=20000)
my_model.fit(sampler, chain_consumer=c)
c.add_chain(np.random.multivariate_normal(res.parameters[1:], res.covariance, size=int(1e7)),
name="Summary Stats", parameters=["$t_0$", "$x_0$", "$x_1$", "$c$"])
if False:
if not os.path.exists(mcmc_chain):
res2, fitted_model2 = sncosmo.mcmc_lc(lcs[0], model, ['t0', 'x0', 'x1', 'c'], nwalkers=20,
nburn=500, nsamples=4000)
mcchain = res2.samples
np.save(mcmc_chain, mcchain)
else:
mcchain = np.load(mcmc_chain)
c.add_chain(mcchain, name="sncosmo mcmc", parameters=["$t_0$", "$x_0$", "$x_1$", "$c$"])
print("Plot surfaces")
