diff --git a/workflow/scripts/weighted_distances.py b/workflow/scripts/weighted_distances.py index 48a229f..e7ea083 100644 --- a/workflow/scripts/weighted_distances.py +++ b/workflow/scripts/weighted_distances.py @@ -57,8 +57,12 @@ def get_frequencies_in_position(site_df: pd.DataFrame, pos: int, reference: Seq) for alt in site_df["ALT"]: frequencies[alt] = float(site_df.loc[site_df["ALT"] == alt, "ALT_FREQ"].iloc[0]) # Obtain frequency for reference allele - ref = 1 - sum(frequencies.values()) - frequencies[reference[pos-1]] += ref + reference_frequency = 1 - sum(frequencies.values()) + reference_allele = reference[pos-1] + if reference_allele in frequencies.keys(): + frequencies[reference_allele] += reference_frequency + else: + frequencies[reference_allele] = reference_frequency return tuple(frequencies[alt] for alt in site_df["ALT"].unique())