def main():
"""Visually test the addition of Noise using add_noise function."""
org_flux = np.ones(500)
for i, snr in enumerate([50, 100, 200, 300, 50000]):
# Test that the standard deviation of the noise is close to the snr level
print("Applying a snr of {}".format(snr))
noisey_flux = add_noise(org_flux, snr)
spec = Spectrum(
flux=copy.copy(org_flux)) # Copy becasue org_flux is mutable.
spec.add_noise(snr)
std = np.std(noisey_flux)
print("Standard deviation of signal = {}".format(std))
snr_est = 1 / std
snr_spectrum = 1. / np.std(spec.flux)
print(
"Estimated SNR from stddev = {}".format(snr_est))
print("Estimated SNR from stddev of Spectrum class = {}".format(
snr_spectrum))
plt.plot(noisey_flux + 0.1 * i, label="snr={}".format(snr))
plt.plot(spec.flux + 0.1 * i,
"--",
label="Spectrum snr={}".format(snr))
# Calculate chisqr from one
chi2 = np.sum((noisey_flux - 1)**2 / 1**2)
print("Chisqr for SNR {} = \t\t{}".format(snr, chi2))
chi2_spectrum = np.sum((spec.flux - 1)**2 / 1**2)
print("Chisqr for snr_spectrum {} = \t\t{}".format(snr, chi2_spectrum))
plt.legend()
plt.show()
def fake_bhm_simulation(wav,
params,
gamma,
limits=(2070, 2180),
noise=None,
header=False):
"""Make a fake spectrum with binary params and radial velocities."""
mod_spec = load_starfish_spectrum(params,
limits=limits,
hdr=True,
normalize=True,
wav_scale=True)
bhm_grid_func = one_comp_model(mod_spec.xaxis, mod_spec.flux, gammas=gamma)
if wav is None:
delta = spec_max_delta(mod_spec, 0, gamma)
assert np.all(np.isfinite(mod_spec.xaxis))
mask = (mod_spec.xaxis > mod_spec.xaxis[0] + 2 * delta) * (
mod_spec.xaxis < mod_spec.xaxis[-1] - 2 * delta)
wav = mod_spec.xaxis[mask]
bhm_grid_values = bhm_grid_func(wav).squeeze()
logging.debug(
__("number of bhm nans = {}", np.sum(~np.isfinite(bhm_grid_values))))
if noise is not None or noise is not 0:
bhm_grid_values = add_noise(bhm_grid_values, noise)
else:
logging.warning(
"\n!!!\n\nNot adding any noise to bhm fake simulation!!!!!\n\n!!!!!\n"
)
print("\n!!!!\n\nNot adding any noise to fake simulation!!!!\n\n!\n")
if header:
return wav, bhm_grid_values.squeeze(), mod_spec.header
else:
return wav, bhm_grid_values.squeeze()
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 test_add_noise_with_mu(mu, noise):
x = mu * np.ones(10000)
assert np.allclose(np.std(add_noise(x, noise, use_mu=True)), mu / noise,
1e-2)
def test_add_noise(noise):
x = np.ones(10000)
assert np.allclose(np.std(add_noise(x, noise)), 1. / noise, 1e-2)
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