def test_improved_gradient_reduces_precision(test_spec):
"""Check that the gradient produces larger RV error."""
wav = test_spec[0]
flux = test_spec[1]
transmission = test_spec[2]
a = Q.rv_precision(wav, flux, mask=transmission, grad=False).value
b = Q.rv_precision(wav, flux, mask=transmission, grad=True).value
assert a <= b
def test_relation_of_rv_to_sqrtsumwis(test_spec, wav_unit, flux_unit, trans_unit):
"""Test relation of sqrtsumwis to rv_precision."""
wav = test_spec[0] * wav_unit
flux = test_spec[1] * flux_unit
mask = test_spec[2]
if test_spec[2] is not None:
mask *= trans_unit
mask = mask ** 2
assert np.all(
Q.rv_precision(wav, flux, mask=mask) == c / Q.sqrt_sum_wis(wav, flux, mask=mask)
)
def test_sqrt_sum_wis_with_mask_with_unit_fails(
test_spec, wav_unit, flux_unit, trans_unit2
):
"""Assert a transmission with a unit fails with type error."""
wav = test_spec[0] * wav_unit
flux = test_spec[1] * flux_unit
transmission = np.random.rand(len(wav)) * trans_unit2
with pytest.raises(TypeError):
Q.sqrt_sum_wis(wav, flux, mask=transmission ** 2)
with pytest.raises(TypeError):
Q.rv_precision(wav, flux, mask=transmission ** 2)
def test_transmission_reduces_precision(test_spec):
"""Check that a transmission vector reduces precision calculation."""
wav = test_spec[0]
flux = test_spec[1]
transmission = test_spec[2]
# Value should be less then normal if trans <=1
if transmission is not None:
assert Q.rv_precision(wav, flux, mask=None) < Q.rv_precision(
wav, flux, mask=transmission
)
# mask=None is the same as mask of all 1.
assert Q.rv_precision(wav, flux, mask=None) == Q.rv_precision(
wav, flux, mask=np.ones_like(wav)
)
def test_sqrt_sum_wis(test_spec, wav_unit, flux_unit, trans_unit):
"""Test that sqrt_sum_wis can handle inputs as Quantities or unitless.
Returns a dimensionless unscaled Quantity.
"""
wav = test_spec[0] * wav_unit
flux = test_spec[1] * flux_unit
mask = test_spec[2]
if test_spec[2] is not None:
mask *= trans_unit
sqrtsumwis = Q.sqrt_sum_wis(wav, flux, mask)
print("wav", wav, type(wav))
print("flux", flux, type(flux))
print("mask", mask, type(mask))
print("sqrtsumwis", sqrtsumwis, type(sqrtsumwis))
if (
(isinstance(wav, Quantity))
or (isinstance(flux, Quantity))
or (isinstance(mask, Quantity) and (test_spec[2] is not None))
):
assert isinstance(sqrtsumwis, u.Quantity)
# Check is unscaled and dimensionless Quantity
assert sqrtsumwis.unit == u.dimensionless_unscaled
assert not hasattr(sqrtsumwis.value, "__len__") # assert value is a scalar
else:
assert not hasattr(sqrtsumwis, "__len__") # assert value is a scalar
def test_increments_rv_accumulate_same_as_full(real_spec, increment_percent, no_mask):
"""Assuming that the weighted rv from the steps should equal the rv from the band."""
wav, flux, mask = real_spec[0], real_spec[1], real_spec[2]
if no_mask:
# Try with mask= None also.
mask = None
rv_full = Q.rv_precision(wav, flux, mask=mask).value
x, incremented_rv = Q.incremental_rv(
wav, flux, mask=mask, percent=increment_percent
)
incremented_weighted = weighted_error(incremented_rv)
assert np.round(rv_full, 2) == np.round(incremented_weighted, 2)
assert x[0] > wav[0]
assert [-1] < wav[-1]
def test_sqrt_sum_wis_with_no_units(test_spec):
"""Test that sqrt_sum_wis can handle inputs as Quantities or unitless.
Returns a dimensionless unscaled Quantity.
"""
sqrtsumwis = Q.sqrt_sum_wis(test_spec[0], test_spec[1], test_spec[2])
# Doesn't turn into quantity if does not have to.
assert not isinstance(sqrtsumwis, u.Quantity)
assert not hasattr(sqrtsumwis, "__len__") # assert value is a scalar
def test_sqrtsumwis_warns_nonfinite(grad_flag):
"""Some warning tests."""
with pytest.warns(UserWarning, match="This will cause infinite errors."):
Q.sqrt_sum_wis(
np.array([1, 2, 3]),
np.array([1, 2, 3]),
np.array([0, 0, 0]),
grad=grad_flag,
) # All masked
with pytest.warns(UserWarning, match="Weight sum is not finite"):
Q.sqrt_sum_wis(
np.array([2, 2, 2]),
np.array([1, 2, 3]),
np.array([1, 1, 1]),
grad=grad_flag,
) # infinite gradient
def test_rvprev_calc_with_lists(test_spec):
"""Test that it can handle list input also."""
wav = list(test_spec[0])
flux = list(test_spec[1])
mask = test_spec[2]
rv = Q.rv_precision(wav, flux, mask)
assert not hasattr(rv.value, "__len__") # assert value is a scalar
assert isinstance(rv, u.Quantity)
assert rv.unit == m_per_s
def test_quality_independent_of_units(test_spec, wav_unit, flux_unit):
"""Quality should be returned as unitless (not a quantity)."""
wave = test_spec[0] * wav_unit
flux = test_spec[1] * flux_unit
q = Q.quality(wave, flux)
assert not isinstance(q, Quantity)
assert isinstance(q, float)
assert not hasattr(q, "__len__") # assert value is a scalar
def test_normalize_flux_new_verse_old(resampled_data):
"""Test only small differences due to new normalization."""
id_string, wav, flux = resampled_data
print("wav in max =", wav[0], wav[-1])
new_norm = eniric.obsolete.snr_norm.normalize_flux(flux,
id_string,
new=True)
old_norm = eniric.obsolete.snr_norm.normalize_flux(flux,
id_string,
new=False)
print("new norm", new_norm)
print("old_norm", old_norm)
rvprec_new = Q.rv_precision(wav, new_norm)
rvprec_old = Q.rv_precision(wav, old_norm)
print("new rv=", rvprec_new, "old rv=", rvprec_old)
assert np.abs(rvprec_new.value - rvprec_old.value) < 0.4
def test_rvprev_calc(test_spec, wav_unit, flux_unit, trans_unit):
"""Test that rv_precision can handle inputs as Quantities or unitless and returns a scalar Quantity."""
wav = test_spec[0] * wav_unit
flux = test_spec[1] * flux_unit
mask = test_spec[2]
if test_spec[2] is not None:
mask *= trans_unit
rv = Q.rv_precision(wav, flux, mask)
assert rv.unit == m_per_s
assert not hasattr(rv.value, "__len__") # assert value is a scalar
assert isinstance(rv, u.Quantity)
def test_increments_rv_gives_reasonable_length(real_spec, increment_percent):
"""The expected number of steps would be between the
wavelength difference divided by the
first and last point * the percent increment.
"""
wav, flux, mask = real_spec[0], real_spec[1], real_spec[2]
x, rv = Q.incremental_rv(wav, flux, mask=mask, percent=increment_percent)
d1 = wav[0] * increment_percent / 100
d2 = wav[-1] * increment_percent / 100
dlambda = wav[-1] - wav[0]
len_rv = len(rv)
assert len_rv >= np.floor(dlambda / d1)
assert len_rv <= np.ceil(dlambda / d2 + 1)
assert len(x) == len_rv
def test_increment_quality_gives_reasonable_length(real_spec, increment_percent):
"""The expected number of steps would be between the
wavelength difference divided by the
first and last point * the percent increment.
"""
print(real_spec)
wav, flux = real_spec[0], real_spec[1]
x, q = Q.incremental_quality(wav, flux, percent=increment_percent)
d1 = wav[0] * increment_percent / 100
d2 = wav[-1] * increment_percent / 100
dlambda = wav[-1] - wav[0]
len_q = len(q)
assert len_q >= np.floor(dlambda / d1)
assert len_q <= np.ceil(dlambda / d2 + 1)
assert len(x) == len_q
def test_sqrt_sum_wis_transmission_outofbounds(test_spec, wav_unit, flux_unit):
"""Transmission must be within 0-1."""
wav = test_spec[0] * wav_unit
flux = test_spec[1] * flux_unit
mask_1 = np.random.randn(len(wav))
mask_2 = np.random.rand(len(wav))
mask_1[0] = 5 # Outside 0-1
mask_2[-1] = -2 # Outside 0-1
# Higher value
with pytest.raises(ValueError):
Q.rv_precision(wav, flux, mask=mask_1)
with pytest.raises(ValueError):
Q.sqrt_sum_wis(wav, flux, mask=mask_1)
# Lower value
with pytest.raises(ValueError):
Q.sqrt_sum_wis(wav, flux, mask=mask_2)
with pytest.raises(ValueError):
Q.sqrt_sum_wis(wav, flux, mask=mask_2)
def calculate_prec(
spectral_types,
bands,
vsini,
resolution,
sampling,
plot_atm=False,
plot_ste=False,
plot_flux=True,
paper_plots=True,
rv_offset=0.0,
use_unshifted=False,
snr=100,
ref_band="J",
new=True,
grad=True,
):
"""Calculate precisions for given combinations.
grad: bool
Use more precise gradient function.
"""
# TODO: iterate over band last so that the J band normalization value can be
# obtained first and applied to each band.
print("using new config.yaml file here!!!!!!!!!!!!!!")
results = {} # creating empty dictionary for the results
wav_plot_m0 = [] # creating empty lists for the plots
flux_plot_m0 = []
wav_plot_m3 = []
flux_plot_m3 = []
wav_plot_m6 = []
flux_plot_m6 = []
wav_plot_m9 = []
flux_plot_m9 = []
for band in bands:
if use_unshifted:
atmmodel = os.path.join(
eniric.paths["atmmodel"],
"{0}_{1}.dat".format(eniric.atmmodel["base"], band),
)
print("Reading atmospheric model...")
atm = Atmosphere.from_file(atmmodel)
wav_atm, flux_atm, std_flux_atm, mask_atm = (
atm.wl,
atm.transmission,
atm.std,
atm.mask,
)
print(
(
"There were {0:d} unmasked pixels out of {1:d}., or {2:.1%}." ""
).format(
np.sum(mask_atm), len(mask_atm), np.sum(mask_atm) / len(mask_atm)
)
)
print(
"The model ranges from {0:4.2f} to {1:4.2f} micron.".format(
wav_atm[0], wav_atm[-1]
)
)
print("Done.")
print("Calculating impact of Barycentric movement on mask...")
# mask_atm = atm.old_barycenter_shift(wav_atm, mask_atm, rv_offset=rv_offset)
mask_atm = barycenter_shift(wav_atm, mask_atm, rv_offset=rv_offset)
else:
shifted_atmmodel = os.path.join(
eniric.paths["atmmodel"],
"{0}_{1}_bary.dat".format(eniric.atmmodel["base"], band),
)
print("Reading pre-doppler-shifted atmospheric model...")
atm = Atmosphere.from_file(shifted_atmmodel)
wav_atm, flux_atm, std_flux_atm, mask_atm = (
atm.wl,
atm.transmission,
atm.std,
atm.mask,
)
print("Done.")
print(
("There were {0:d} unmasked pixels out of {1:d}, or {2:.1%}." "").format(
np.sum(mask_atm), len(mask_atm), np.sum(mask_atm) / len(mask_atm)
)
)
if plot_atm:
# moved plotting code to separate code, eniric.obsolete.plotting_functions.py
plt_functions.plot_atmosphere_model(wav_atm, flux_atm, mask_atm)
# theoretical ratios calculation
# wav_m0, flux_m0, wav_m3, flux_m3, wav_m6, flux_m6, wav_m9, flux_m9 = read_nIRspectra()
iterations = itertools.product(spectral_types, vsini, resolution, sampling)
# for star in spectral_types:
# for vel in vsini:
# for res in resolution:
# for smpl in sampling:
for (star, vel, res, smpl) in iterations:
file_to_read = (
"Spectrum_{0}-PHOENIX-ACES_{1}band_vsini{2}_R{3}" "_res{4:2.01f}.dat"
).format(star, band, vel, res, float(smpl))
# print("Working on "+file_to_read+".")
try:
wav_stellar, flux_stellar = io.pdread_2col(
os.path.join(eniric.paths["resampled"], file_to_read)
)
except FileNotFoundError:
# Turn list of strings into strings without symbols ["J", "K"] -> J K
spectral_str = re.sub(r"[\[\]\"\',]", "", str(spectral_types))
band_str = re.sub(r"[\[\]\"\',]", "", str(bands))
vsini_str = re.sub(r"[\[\]\"\',]", "", str(vsini))
res_str = re.sub(r"[\[\]\"\',]", "", str(resolution))
sampling_str = re.sub(r"[\[\]\"\',]", "", str(sampling))
print(
(
"\nFor just this file I suggest you run\n\tpython nIR_run.py -s {0} -b {1} -v {2} -R {3} "
"--sample_rate {4}\nOr for all the combinations you ran here\n\tpython nIR_run.py -s {5}"
" -b {6} -v {7} -R {8} --sample_rate {9}"
""
).format(
star,
band,
vel,
res,
smpl,
spectral_str,
band_str,
vsini_str,
res_str,
sampling_str,
)
)
raise
if len(wav_stellar) == 0 or len(flux_stellar) == 0:
raise Exception("The file {0} is empty".format(file_to_read))
# Removing boundary effects
wav_stellar = wav_stellar[2:-2]
flux_stellar = flux_stellar[2:-2]
# sample was left aside because only one value existed
# TODO: Add metallicity and logg into id string
id_string = "{0:s}-{1:s}-{2:.1f}-{3:s}".format(star, band, float(vel), res)
# Getting the wav, flux and mask values from the atm model
# that are the closest to the stellar wav values, see
# https://stackoverflow.com/questions/2566412/find-nearest-value-in-numpy-array
index_atm = np.searchsorted(wav_atm, wav_stellar)
# replace indexes outside the array, at the very end, by the value at the very end
# index_atm = [index if(index < len(wav_atm)) else len(wav_atm)-1 for index in index_atm]
indx_mask = index_atm >= len(wav_atm) # find broken indexs
index_atm[indx_mask] = len(wav_atm) - 1 # replace with index of end.
wav_atm_selected = wav_atm[index_atm]
flux_atm_selected = flux_atm[index_atm]
mask_atm_selected = mask_atm[index_atm]
# Check mask masks out deep atmosphere absorption
if np.any(flux_atm_selected[mask_atm_selected] < 0.98):
print(
"####WARNGING####\nThis absorption mask does not mask out deep atmosphere transmission!"
)
print(
"Min flux_atm_selected[mask_atm_selected] = {} < 0.98\n####".format(
np.min(flux_atm_selected[mask_atm_selected])
)
)
# Normalize to SNR 100 in middle of J band 1.25 micron!
# flux_stellar = normalize_flux(flux_stellar, id_string)
# flux_stellar = snrnorm.normalize_flux(flux_stellar, id_string, new=True) # snr=100, ref_band="J"
flux_stellar = eniric.obsolete.snr_norm.normalize_flux(
flux_stellar,
id_string,
new=new,
snr=snr,
ref_band=ref_band,
sampling=smpl,
)
if id_string in [
"M0-J-1.0-100k",
"M3-J-1.0-100k",
"M6-J-1.0-100k",
"M9-J-1.0-100k",
]:
index_ref = np.searchsorted(
wav_stellar, 1.25
) # searching for the index closer to 1.25 micron
snr_estimate = np.sqrt(
np.sum(flux_stellar[index_ref - 1 : index_ref + 2])
)
print(
"\tSanity Check: The S/N for the {0:s} reference model was of {1:4.2f}.".format(
id_string, snr_estimate
)
)
elif "J" in id_string:
index_ref = np.searchsorted(
wav_stellar, 1.25
) # searching for the index closer to 1.25 micron
snr_estimate = np.sqrt(
np.sum(flux_stellar[index_ref - 1 : index_ref + 2])
)
print(
"\tSanity Check: The S/N for the {0:s} non-reference model was of {1:4.2f}.".format(
id_string, snr_estimate
)
)
# Precision given by the first method:
print("Performing analysis for: ", id_string)
prec_1 = Qcalculator.rv_precision(wav_stellar, flux_stellar, grad=grad)
# Precision as given by the second_method
wav_stellar_chunks, flux_stellar_chunks = eniric.legacy.mask_clumping(
wav_stellar, flux_stellar, mask_atm_selected
)
prec_2_old = eniric.legacy.RVprec_calc_masked(
wav_stellar_chunks, flux_stellar_chunks, grad=grad
)
prec_2 = eniric.legacy.RVprec_calc_masked(
wav_stellar, flux_stellar, mask_atm_selected, grad=grad
)
assert np.all(prec_2_old == prec_2)
"""
# histogram checking
lengths = [len(chunk) for chunk in flux_stellar_chunks_unformatted]
n, bins, patches = plt.hist(lengths, 500, range=[0.5, 500.5], histtype='stepfilled')
plt.title(id_string)
plt.show()
"""
# Precision as given by the third_method
prec_3 = Qcalculator.rv_precision(
wav_stellar, flux_stellar, mask=flux_atm_selected ** 2, grad=grad
)
# Adding Precision results to the dictionary
results[id_string] = [prec_1, prec_2, prec_3]
# Prepare/Do for the plotting.
if plot_ste or plot_ste == id_string:
plt_functions.plot_stellar_spectum(
wav_stellar, flux_stellar, wav_atm_selected, mask_atm_selected
)
plot_ids = [
"M3-Z-1.0-100k",
"M3-Y-1.0-100k",
"M3-J-1.0-100k",
"M3-H-1.0-100k",
"M3-K-1.0-100k",
]
if plot_flux and id_string in plot_ids:
wav_plot_m0.append(wav_stellar)
flux_plot_m0.append(flux_stellar)
if plot_flux and id_string in plot_ids:
wav_plot_m3.append(wav_stellar)
flux_plot_m3.append(flux_stellar)
if plot_flux and id_string in plot_ids:
wav_plot_m6.append(wav_stellar)
flux_plot_m6.append(flux_stellar)
if plot_flux and id_string in plot_ids:
wav_plot_m9.append(wav_stellar)
flux_plot_m9.append(flux_stellar)
if plot_flux:
plt_functions.plot_nIR_flux()
if paper_plots:
plt_functions.plot_paper_plots()
else:
return results
def test_quality_independent_of_flux_level(scale):
"""Q of a spectrum is independent of flux level."""
wavelength = np.arange(100)
flux = np.random.random(100)
assert np.allclose(Q.quality(wavelength, flux), Q.quality(wavelength, flux * scale))