In [ ]:
from astropy.io import fits
import numpy as np
import scipy as sp
from scipy.interpolate import interp1d
from scipy.stats import chisquare
from PyAstronomy.pyasl import dopplerShift
import matplotlib.pyplot as plt
% matplotlib
In [ ]:
def two_comp_model(wav, model1, model2, alphas, rvs, gammas):
# Make 2 component simulations, broadcasting over alpha, rv, gamma values.
# Enable single scalar inputs (turn to 1d np.array)
if not hasattr(alphas, "__len__"):
alphas = np.asarray(alphas)[np.newaxis]
if not hasattr(rvs, "__len__"):
rvs = np.asarray(rvs)[np.newaxis]
if not hasattr(gammas, "__len__"):
gammas = np.asarray(gammas)[np.newaxis]
print(len(gammas))
am2 = model2[:,np.newaxis] * alphas # alpha * Model2 (am2)
# print(am2.shape)
am2rv = np.empty(am2.shape + (len(rvs),)) # am2rv = am2 with rv doppler-shift
print(am2rv.shape)
for i, rv in enumerate(rvs):
#nflux, wlprime = dopplerShift(wav, am2, rv)
#am2rv[:, :, i] = nflux
wav_i = (1 + rv / c) * wav
am2rv[:, :, i] = interp1d(wav_i, am2, axis=0, bounds_error=False)(wav)
# Normalize by (1 / 1 + alpha)
am2rv = am2rv / (1 + alphas)[np.newaxis, :, np.newaxis]
am2rvm1 = h[:, np.newaxis, np.newaxis] + am2rv # am2rvm1 = am2rv + model_1
print(am2rvm1.shape)
am2rvm1g = np.empty(am2rvm1.shape + (len(gammas),)) # am2rvm1g = am2rvm1 with gamma doppler-shift
for j, gamma in enumerate(gammas):
wav_j = (1 + gamma / 299792.458) * wav
am2rvm1g[:, :, :, j] = interp1d(wav_j, am2rvm1, axis=0, bounds_error=False)(wav)
return interp1d(w, am2rvm1g, axis=0) # pass it the wavelength values to return
In [ ]:
wav = "/home/jneal/Phd/data/phoenixmodels/WAVE_PHOENIX-ACES-AGSS-COND-2011.fits"
host = "/home/jneal/Phd/data/phoenixmodels/HD30501-lte05200-4.50-0.0.PHOENIX-ACES-AGSS-COND-2011-HiRes.fits"
comp = "/home/jneal/Phd/data/phoenixmodels/HD30501b-lte02500-5.00-0.0.PHOENIX-ACES-AGSS-COND-2011-HiRes.fits"
w = fits.getdata(wav) / 10
h = fits.getdata(host)
c = fits.getdata(comp)
In [ ]:
mask = (2111 < w) & (w < 2117)
w = w[mask]
h = h[mask]
c = c[mask]
# crude normalization
h = h/np.max(h)
c = c/np.max(c)
In [ ]:
# Create a simulated spectrum
# Parameters
c_kms = 299792.458 # km/s
s_alpha = np.array([0.1])
s_rv = np.array([1.5])
s_gamma = np.array([0.5])
answers = (s_alpha, s_rv, s_gamma)
# COMPACT SIMULATION
comp = interp1d((1 + s_rv / c_kms) * w, s_alpha * c, bounds_error=False)(w)
Sim_func = interp1d((1 + s_gamma / c_kms) * w, (h + comp) / (1 + s_alpha), bounds_error=False, axis=0)
sim_f_orgw = Sim_func(w)
sim_w = np.linspace(2114, 2115, 1024)
sim_f = Sim_func(sim_w)
In [ ]:
# Compare output to tcm
tcm_sim_f = two_comp_model(w, h, c, s_alpha, s_rv, s_gamma)(sim_w)
ocm_sim_f = one_comp_model(w, h, s_gamma)(sim_w)
In [ ]:
plt.close()
plt.plot(w, sim_f_orgw, label="org_w")
plt.plot(sim_w, sim_f, label="sim")
plt.plot(sim_w, np.squeeze(tcm_sim_f), label="tcm sim")
plt.plot(sim_w, np.squeeze(ocm_sim_f), label="ocm sim")
plt.legend()
plt.show()
sim_f.shape
# sim_w, sim_f are the observations to perform chisquared against!
In [ ]:
alphas = np.linspace(0.1, 0.3, 40)
rvs = np.arange(1.1, 2, 0.05)
gammas = np.arange(-0.9, 1, 0.015)
print(len(alphas), len(rvs), len(gammas))
In [ ]:
tcm = two_comp_model(w, h, c, alphas=alphas, rvs=rvs, gammas=gammas)
In [ ]:
# Two component model
tcm_obs = tcm(sim_w)
tcm_obs.shape
In [ ]:
chi2 = chisquare(sim_f[:, np.newaxis, np.newaxis, np.newaxis], tcm_obs).statistic
print(chi2.shape)
min_indx = np.unravel_index(chi2.argmin(), chi2.shape)
print("sim results", alphas[min_indx[0]], min_rvs[indx[1]], gammas[min_indx[2]])
print("answer", answers)
In [ ]:
# Putting resulted sim min values back into tcm model
res = two_comp_model(w, h, c, alphas[min_indx[0]], rvs[min_indx[1]], gammas[min_indx[2]])
res_f = res(sim_w) # Flux at the min min chisquare model evaulated at obs points.
In [ ]:
# Compare to tcm generated simulation
chi2_tcm = chisquare(tcm_sim_f, tcm_obs).statistic
min_indx_tcm = np.unravel_index(chi2.argmin(), chi2.shape)
print("tcm results", alphas[min_indx_tcm[0]], rvs[min_indx_tcm[1]], gammas[min_indx_tcm[2]])
print("answer", answers)
In [ ]:
# Putting resulted tcm sim min values back into tcm model
res_tcm = two_comp_model(w, h, c, alphas[min_indx[0]], rvs[min_indx[1]], gammas[min_indx[2]])
res_tcm_f = res_tcm(sim_w) # Flux at the min min chisquare model evaulated at obs points.
In [ ]:
plt.plot(sim_w, sim_f, "--", label="org")
plt.plot(sim_w, np.squeeze(res_f), label= "2 comp")
plt.plot(sim_w, np.squeeze(res_tcm_f), label="fit to tcm sim")
plt.title("Comparison to Simulation")
plt.legend()
plt.show()
In [ ]:
plt.close()
plt.figure()
In [ ]:
plt.figure()
plt.contourf(chi2[:,:,0])
plt.figure()
plt.contourf(chi2[0,:,:])
In [ ]:
plt.figure()
plt.contourf(chi2[:,1,:])
plt.figure()
In [ ]:
# Slice arrays to make contour maps
xslice = np.arange(0, chi2.shape[0], 5)
yslice = np.arange(0, chi2.shape[1], 5)
zslice = np.arange(0, chi2.shape[2], 5)
for xs in xslice:
plt.figure()
plt.contourf(chi2[xs, :, :])
plt.colorbar()
plt.title("x alpha = {}".format(alphas[xs]))
plt.show()
In [ ]:
for ys in yslice:
plt.figure()
plt.contourf(chi2[:, ys, :])
plt.colorbar()
plt.title("y rvs = {}".format(rvs[ys]))
plt.show()
In [ ]:
for zs in zslice:
plt.figure()
plt.contourf(chi2[:, :, zs])
plt.colorbar()
plt.title("z gammas = {}".format(gammas[zs]))
plt.show()
In [ ]:
for xs in np.concatenate([xslice, yslice, zslice]):
plt.close()
In [ ]:
In [ ]: