def test_nearby_peaks():
assert nearby_peaks(w_test, 34, delta = 9) == (2, 6)
assert nearby_peaks(w_test, 34, delta = 5) == (2, 4)
assert nearby_peaks(w_test, 34, delta = 3) == (4, 4)
assert nearby_peaks(w_test, 25, delta = 7) == (1, 4)
assert nearby_peaks(w_test, 45, delta = 7) == (4, 7)
assert nearby_peaks(w_test, 55, delta = 7) == (6, 8)
assert nearby_peaks(w_test, 35, delta = 100) == (0, 8)
assert nearby_peaks(w_test, 5, delta = 100) == (0, 8)
assert nearby_peaks(w_test, 5, delta = 3) == (0, 0)
assert nearby_peaks(w_test, 65, delta = 3) == (8, 8)
def filter_elements(
observed_peak_wavelengths,
known_emission_lines,
drop_if_below_fraction,
delta,
group_by):
peaks_grouped = defaultdict(set)
for row_index, row in known_emission_lines.iterrows():
key = tuple([row[var] for var in group_by])
peaks_grouped[key].add(row["wav_mars"])
overall_mask_to_drop = np.zeros(len(known_emission_lines), dtype='bool')
for key, emission_wavs in peaks_grouped.items():
denom = 0
numer = 0
for w in emission_wavs:
denom += 1
i_lo, i_hi = nearby_peaks(
observed_peak_wavelengths, w, delta = delta)
nearby = observed_peak_wavelengths[i_lo : i_hi]
if len(nearby) > 0:
numer += 1
fraction = numer / float(denom)
if fraction < drop_if_below_fraction:
mask = np.ones(len(known_emission_lines), dtype='bool')
for var_index, var in enumerate(group_by):
mask &= (known_emission_lines[var] == key[var_index])
overall_mask_to_drop |= mask
omitted = known_emission_lines[overall_mask_to_drop]
filtered = known_emission_lines[~overall_mask_to_drop]
assert len(omitted) + len(filtered) == len(known_emission_lines)
return filtered