def join_with_broadcast(mod1, mod2, rv, gamma, new_x):
broadcast_result = inherent_alpha_model(mod1.xaxis,
mod1.flux,
mod2.flux,
rvs=rv,
gammas=gamma)
broadcast_values = broadcast_result(new_x)
return broadcast_values.squeeze()
def join_with_broadcast(mod1, mod2, rv, gamma):
broadcast_result = inherent_alpha_model(mod1.xaxis,
mod1.flux,
mod2.flux,
rvs=rv,
gammas=gamma)
broadcast_values = broadcast_result(mod1.xaxis)
return Spectrum(flux=broadcast_values.squeeze(), xaxis=mod1.xaxis)
def fake_obs(simnum, snr=200, suffix=None, plot=False, mode="iam"):
snr = 200
params1 = [5300, 4.5, 0.0]
params2 = [2500, 4.5, 0.0]
gamma = 5
rv = -3
normalization_limits = [2100, 2180]
mod1_spec = load_starfish_spectrum(params1,
limits=normalization_limits,
hdr=True,
normalize=False,
area_scale=True,
flux_rescale=True)
mod2_spec = load_starfish_spectrum(params2,
limits=normalization_limits,
hdr=True,
normalize=False,
area_scale=True,
flux_rescale=True)
if broadcast:
broadcast_result = inherent_alpha_model(mod1_spec.xaxis,
mod1_spec.flux,
mod2_spec.flux,
rvs=rv,
gammas=gamma)
broadcast_values = broadcast_result(obs_spec.xaxis)
result_spectrum = Spectrum(flux=broadcast_values,
xaxis=mod1_spec.xaxis)
else:
# Manually redo the join
mod2_spec.doppler_shift(rv)
mod_combine = mod1_spec.copy()
mod_combine += mod2_spec
mod_combine.doppler_shift(gamma)
mod_combine.normalize(method="exponential")
mod_combine.interpolate(obs_spec.xaxis)
result_spectrum = mod_combine
# result_spectrum = Spectrum(flux=broadcast_values, xaxis=obs_spec.xaxis)
result_spectrum.add_noise(snr)
# Save as
# Detector limits
dect_limits = [(2112, 2123), (2127, 2137), (2141, 2151), (2155, 2165)]
for ii, dect in enumerate(dect_limits):
spec = result_spectrum.copy()
spec.wav_select(*dect)
spec.resample(1024)
name = "HDsim-{0}_{1}_snr_{2}".format(sim_num, ii, snr)
def test_inherinent_model_same_as_alpha_0(host, comp):
tcm = two_comp_model(host.xaxis,
host.flux,
comp.flux,
alphas=0,
rvs=[0, 2, 4],
gammas=[1, 2, 3])
iam = inherent_alpha_model(host.xaxis,
host.flux,
comp.flux,
rvs=[0, 2, 4],
gammas=[1, 2, 3])
host.wav_select(2100.5, 2104.5) # cut to avoid Nans from doppler shifts
tcm_value = tcm(host.xaxis)
iam_value = iam(host.xaxis)
# print(tcm_value)
# print(iam_value)
assert tcm_value.squeeze().shape == iam_value.shape
def iam_wrapper(num, params1, model2_pars, rvs, gammas, obs_spec, norm=False,
verbose=False, save_only=True, chip=None, prefix=None, errors=None,
area_scale=True, wav_scale=True, grid_slices=False, norm_method="scalar",
fudge=None):
"""Wrapper for iteration loop of iam. params1 fixed, model2_pars are many.
fudge is multiplicative on companion spectrum.
"""
if prefix is None:
sf = os.path.join(
simulators.paths["output_dir"], obs_spec.header["OBJECT"].upper(),
"iam_{0}_{1}-{2}_part{6}_host_pars_[{3}_{4}_{5}].csv".format(
obs_spec.header["OBJECT"].upper(), int(obs_spec.header["MJD-OBS"]), chip,
params1[0], params1[1], params1[2], num))
prefix = os.path.join(
simulators.paths["output_dir"], obs_spec.header["OBJECT"].upper()) # for fudge
else:
sf = "{0}_part{4}_host_pars_[{1}_{2}_{3}].csv".format(
prefix, params1[0], params1[1], params1[2], num)
save_filename = sf
if os.path.exists(save_filename) and save_only:
print("'{0}' exists, so not repeating calculation.".format(save_filename))
return None
else:
if not save_only:
iam_grid_chisqr_vals = np.empty(len(model2_pars))
for jj, params2 in enumerate(model2_pars):
if verbose:
print(("Starting iteration with parameters: "
"{0}={1},{2}={3}").format(num, params1, jj, params2))
# Main Part
rv_limits = observation_rv_limits(obs_spec, rvs, gammas)
obs_spec = obs_spec.remove_nans()
assert ~np.any(np.isnan(obs_spec.flux)), "Observation has nan"
# Load phoenix models and scale by area and wavelength limit
mod1_spec, mod2_spec = \
prepare_iam_model_spectra(params1, params2, limits=rv_limits,
area_scale=area_scale, wav_scale=wav_scale)
# Estimated flux ratio from models
inherent_alpha = continuum_alpha(mod1_spec, mod2_spec, chip)
# Combine model spectra with iam model
mod1_spec.plot(label=params1)
mod2_spec.plot(label=params2)
plt.close()
if fudge or (fudge is not None):
fudge_factor = float(fudge)
mod2_spec.flux *= fudge_factor # fudge factor multiplication
mod2_spec.plot(label="fudged {0}".format(params2))
plt.title("fudges models")
plt.legend()
fudge_prefix = os.path.basename(os.path.normpath(prefix))
fname = os.path.join(simulators.paths["output_dir"],
obs_spec.header["OBJECT"].upper(), "iam", "fudgeplots",
"{1}_fudged_model_spectra_factor={0}_num={2}_iter_{3}.png".format(fudge_factor,
fudge_prefix,
num, jj))
plt.savefig(fname)
plt.close()
warnings.warn("Using a fudge factor = {0}".format(fudge_factor))
iam_grid_func = inherent_alpha_model(mod1_spec.xaxis, mod1_spec.flux, mod2_spec.flux,
rvs=rvs, gammas=gammas)
iam_grid_models = iam_grid_func(obs_spec.xaxis)
# Continuum normalize all iam_gird_models
def axis_continuum(flux):
"""Continuum to apply along axis with predefined variables parameters."""
return continuum(obs_spec.xaxis, flux, splits=20, method="exponential", top=20)
iam_grid_continuum = np.apply_along_axis(axis_continuum, 0, iam_grid_models)
iam_grid_models = iam_grid_models / iam_grid_continuum
# RE-NORMALIZATION
if chip == 4:
# Quadratically renormalize anyway
obs_spec = renormalization(obs_spec, iam_grid_models, normalize=True, method="quadratic")
obs_flux = renormalization(obs_spec, iam_grid_models, normalize=norm, method=norm_method)
if grid_slices:
# Long execution plotting.
plot_iam_grid_slices(obs_spec.xaxis, rvs, gammas, iam_grid_models,
star=obs_spec.header["OBJECT"].upper(),
xlabel="wavelength", ylabel="rv", zlabel="gamma",
suffix="iam_grid_models", chip=chip)
old_shape = iam_grid_models.shape
# Arbitrary_normalization of observation
iam_grid_models, arb_norm = arbitrary_rescale(iam_grid_models,
*simulators.sim_grid["arb_norm"])
# print("Arbitrary Normalized iam_grid_model shape.", iam_grid_models.shape)
assert iam_grid_models.shape == (*old_shape, len(arb_norm))
# Calculate Chi-squared
obs_flux = np.expand_dims(obs_flux, -1) # expand on last axis to match rescale
iam_norm_grid_chisquare = chi_squared(obs_flux, iam_grid_models, error=errors)
# Take minimum chi-squared value along Arbitrary normalization axis
iam_grid_chisquare, arbitrary_norms = arbitrary_minimums(iam_norm_grid_chisquare, arb_norm)
npix = obs_flux.shape[0] # Number of pixels used
if grid_slices:
# Long execution plotting.
plot_iam_grid_slices(rvs, gammas, arb_norm, iam_norm_grid_chisquare,
star=obs_spec.header["OBJECT"].upper(),
xlabel="rv", ylabel="gamma", zlabel="Arbitrary Normalization",
suffix="iam_grid_chisquare", chip=chip)
if not save_only:
iam_grid_chisqr_vals[jj] = iam_grid_chisquare.ravel()[np.argmin(iam_grid_chisquare)]
save_full_iam_chisqr(save_filename, params1, params2,
inherent_alpha, rvs, gammas,
iam_grid_chisquare, arbitrary_norms, npix, verbose=verbose)
if save_only:
return None
else:
return iam_grid_chisqr_vals
def fake_iam_simulation(wav,
params1,
params2,
gamma,
rv,
limits=(2070, 2180),
noise=None,
header=False,
fudge=None,
area_scale=True):
"""Make a fake spectrum with binary params and radial velocities."""
mod1_spec, mod2_spec = prepare_iam_model_spectra(params1,
params2,
limits,
area_scale=area_scale)
if fudge is not None:
mod2_spec.flux = mod2_spec.flux * fudge
warnings.warn("Fudging fake companion by '*{0}'".format(fudge))
# Combine model spectra with iam model
iam_grid_func = inherent_alpha_model(mod1_spec.xaxis,
mod1_spec.flux,
mod2_spec.flux,
rvs=rv,
gammas=gamma)
if wav is None:
delta = spec_max_delta(mod1_spec, rv, gamma)
assert np.all(np.isfinite(mod1_spec.xaxis))
mask = (mod1_spec.xaxis > mod1_spec.xaxis[0] + 2 * delta) * (
mod1_spec.xaxis < mod1_spec.xaxis[-1] - 2 * delta)
wav = mod1_spec.xaxis[mask]
iam_grid_models = iam_grid_func(wav).squeeze()
logging.debug(__("iam_grid_func(wav).squeeze() = {}", iam_grid_models))
logging.debug(
__("number of nans {}", np.sum(~np.isfinite(iam_grid_models))))
logging.debug(__("iam_grid_models = {}", iam_grid_models))
logging.debug("Continuum normalizing")
# Continuum normalize all iam_gird_models
def axis_continuum(flux):
"""Continuum to apply along axis with predefined variables parameters."""
return continuum(wav, flux, splits=20, method="exponential", top=20)
iam_grid_continuum = np.apply_along_axis(axis_continuum, 0,
iam_grid_models)
iam_grid_models = iam_grid_models / iam_grid_continuum
# This noise is added after continuum normalization.
if noise is not None or noise is not 0:
# Add 1 / snr noise to continuum normalized spectra
iam_grid_models = add_noise(iam_grid_models, noise, use_mu=False)
else:
logging.warning(
"\n!!!\n\nNot adding any noise to fake simulation!!\n\n!!!!!\n")
print("\n!!!\n\nNot adding any noise to fake simulation!!\n\n!!!!!\n")
if header:
return wav, iam_grid_models.squeeze(), mod1_spec.header
else:
return wav, iam_grid_models.squeeze()
def compare_spectra(table, params):
"""Plot the min chi2 result against the observations."""
extractor = DBExtractor(table)
for ii, chi2_val in enumerate(chi2_names[0:-2]):
df = extractor.minimum_value_of(chi2_val)
df = df[["teff_1", "logg_1", "feh_1", "gamma",
"teff_2", "logg_2", "feh_2", "rv", chi2_val]]
params1 = [df["teff_1"].values[0], df["logg_1"].values[0], df["feh_1"].values[0]]
params2 = [df["teff_2"].values[0], df["logg_2"].values[0], df["feh_2"].values[0]]
params1 = [float(param1) for param1 in params1]
params2 = [float(param2) for param2 in params2]
gamma = df["gamma"].values
rv = df["rv"].values
from simulators.iam_module import iam_helper_function
obs_name, obs_params, output_prefix = iam_helper_function(params["star"], params["obsnum"], ii + 1)
obs_spec = load_spectrum(obs_name)
# Mask out bad portion of observed spectra
# obs_spec = spectrum_masking(obs_spec, params["star"], params["obsnum"], ii + 1)
# Barycentric correct spectrum
_obs_spec = barycorr_crires_spectrum(obs_spec, extra_offset=None)
normalization_limits = [obs_spec.xaxis[0] - 5, obs_spec.xaxis[-1] + 5]
# models
# print("params for models", params1, params2)
mod1 = load_starfish_spectrum(params1, limits=normalization_limits,
hdr=True, normalize=False, area_scale=True,
flux_rescale=True)
mod2 = load_starfish_spectrum(params2, limits=normalization_limits,
hdr=True, normalize=False, area_scale=True,
flux_rescale=True)
iam_grid_func = inherent_alpha_model(mod1.xaxis, mod1.flux, mod2.flux,
rvs=rv, gammas=gamma)
iam_grid_model = iam_grid_func(obs_spec.xaxis).squeeze()
iam_grid_model_full = iam_grid_func(mod1.xaxis).squeeze()
model_spec_full = Spectrum(flux=iam_grid_model_full, xaxis=mod1.xaxis)
model_spec = Spectrum(flux=iam_grid_model, xaxis=obs_spec.xaxis)
model_spec = model_spec.remove_nans()
model_spec = model_spec.normalize(method="exponential")
model_spec_full = model_spec_full.remove_nans()
model_spec_full = model_spec_full.normalize(method="exponential")
fig, ax = plt.subplots(1, 1)
plt.plot(obs_spec.xaxis, obs_spec.flux, label="Observation")
plt.plot(model_spec.xaxis, model_spec.flux, label="Minimum \chi^2 model")
plt.plot(model_spec_full.xaxis, model_spec_full.flux, "--", label="Model_full_res")
plt.legend()
fig.tight_layout()
name = "{0}-{1}_{2}_{3}_min_chi2_spectrum_comparison_{4}.png".format(
params["star"], params["obsnum"], params["chip"], chi2_val, params["suffix"])
plt.savefig(os.path.join(params["path"], "plots", name))
plt.close()
plt.plot(obs_spec.xaxis, obs_spec.flux, label="Observation")
plt.show()
def main(star,
obsnum,
teff_1,
logg_1,
feh_1,
teff_2,
logg_2,
feh_2,
gamma,
rv,
plot_name=None):
fig, axis = plt.subplots(2, 2, figsize=(15, 8), squeeze=False)
for chip, ax in zip(range(1, 5), axis.flatten()):
# Get observation data
obs_name, params, output_prefix = iam_helper_function(
star, obsnum, chip)
# Load observed spectrum
obs_spec = load_spectrum(obs_name)
# Mask out bad portion of observed spectra
obs_spec = spectrum_masking(obs_spec, star, obsnum, chip)
# Barycentric correct spectrum
_obs_spec = barycorr_crires_spectrum(obs_spec, extra_offset=None)
# Determine Spectrum Errors
# error_off = False
# errors = spectrum_error(star, obsnum, chip, error_off=error_off)
# Create model with given parameters
host = load_starfish_spectrum([teff_1, logg_1, feh_1],
limits=[2110, 2165],
area_scale=True,
hdr=True)
if teff_2 is None:
assert (logg_2 is None) and (feh_2 is None) and (
rv == 0), "All must be None for bhm case."
companion = Spectrum(xaxis=host.xaxis,
flux=np.zeros_like(host.flux))
else:
companion = load_starfish_spectrum([teff_2, logg_2, feh_2],
limits=[2110, 2165],
area_scale=True,
hdr=True)
joint_model = inherent_alpha_model(host.xaxis,
host.flux,
companion.flux,
gammas=gamma,
rvs=rv)
model_spec = Spectrum(xaxis=host.xaxis,
flux=joint_model(host.xaxis).squeeze())
model_spec = model_spec.remove_nans()
model_spec = model_spec.normalize("exponential")
# plot
obs_spec.plot(axis=ax, label="{}-{}".format(star, obsnum))
model_spec.plot(axis=ax, linestyle="--", label="Chi-squared model")
ax.set_xlim([obs_spec.xmin() - 0.5, obs_spec.xmax() + 0.5])
ax.set_title("{} obs {} chip {}".format(star, obsnum, chip))
ax.legend()
plt.tight_layout()
if plot_name is None:
fig.show()
else:
if not (plot_name.endswith(".png") or plot_name.endswith(".pdf")):
raise ValueError("plot_name does not end with .pdf or .png")
fig.savefig(plot_name)
return 0
def fake_simulation(wav,
params1,
params2,
gamma,
rv,
chip=None,
limits=(2070, 2180),
noise=None):
"""Make a fake spectrum with binary params and radial velocities."""
mod1_spec, mod2_spec = prepare_iam_model_spectra(params1, params2, limits)
mod1_spec.plot()
mod2_spec.plot()
plt.show()
# Estimated flux ratio from models
if chip is not None:
inherent_alpha = continuum_alpha(mod1_spec, mod2_spec, chip)
print("inherent flux ratio = {}, chip={}".format(inherent_alpha, chip))
# Combine model spectra with iam model
iam_grid_func = inherent_alpha_model(mod1_spec.xaxis,
mod1_spec.flux,
mod2_spec.flux,
rvs=rv,
gammas=gamma)
if wav is None:
delta = spec_max_delta(mod1_spec, rv, gamma)
assert np.all(np.isfinite(mod1_spec.xaxis))
mask = (mod1_spec.xaxis > mod1_spec.xaxis[0] + delta) * (
mod1_spec.xaxis < mod1_spec.xaxis[-1] - delta)
wav = mod1_spec.xaxis[mask]
print("wav masked", wav)
iam_grid_models = iam_grid_func(wav).squeeze()
print(iam_grid_models)
assert np.all(np.isfinite(iam_grid_models))
if isinstance(noise, (int, float)):
snr = noise
else:
snr = None
# Continuum normalize all iam_gird_models
def axis_continuum(flux):
"""Continuum to apply along axis with predefined variables parameters."""
return continuum(wav, flux, splits=50, method="exponential", top=5)
iam_grid_continuum = np.apply_along_axis(axis_continuum, 0,
iam_grid_models)
iam_grid_models = iam_grid_models / iam_grid_continuum
# grid_spectrum = Spectrum(xaxis=wav, flux=iam_grid_models)
# iam_grid_models = grid_spectrum.normalize("exponential")
# Add the noise
from mingle.utilities.simulation_utilities import add_noise
if snr is not None:
iam_grid_models = add_noise(iam_grid_models, snr)
if np.any(np.isnan(iam_grid_models)):
print("there was some nans")
pass
return wav, iam_grid_models
