|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +Copied from https://github.com/nextstrain/ncov-ingest/blob/master/bin/join-metadata-and-clades |
| 4 | +""" |
| 5 | +import argparse |
| 6 | +import sys |
| 7 | +from datetime import datetime |
| 8 | +import pandas as pd |
| 9 | +import numpy as np |
| 10 | + |
| 11 | +INSERT_BEFORE_THIS_COLUMN = "pango_lineage" |
| 12 | +METADATA_JOIN_COLUMN_NAME = 'strain' |
| 13 | +NEXTCLADE_JOIN_COLUMN_NAME = 'seqName' |
| 14 | +VALUE_MISSING_DATA = '?' |
| 15 | + |
| 16 | +rate_per_day = 0.0007 * 29903 / 365 |
| 17 | +reference_day = datetime(2020,1,1).toordinal() |
| 18 | + |
| 19 | +column_map = { |
| 20 | + "clade": "Nextstrain_clade", |
| 21 | + "totalMissing": "missing_data", |
| 22 | + "totalSubstitutions": "divergence", |
| 23 | + "totalNonACGTNs": "nonACGTN", |
| 24 | + "privateNucMutations.totalUnlabeledSubstitutions": "rare_mutations", |
| 25 | + "privateNucMutations.totalReversionSubstitutions": "reversion_mutations", |
| 26 | + "privateNucMutations.totalLabeledSubstitutions": "potential_contaminants", |
| 27 | + "qc.missingData.status": "QC_missing_data", |
| 28 | + "qc.mixedSites.status": "QC_mixed_sites", |
| 29 | + "qc.privateMutations.status": "QC_rare_mutations", |
| 30 | + "qc.snpClusters.status": "QC_snp_clusters", |
| 31 | + "qc.frameShifts.status": "QC_frame_shifts", |
| 32 | + "qc.stopCodons.status": "QC_stop_codons", |
| 33 | + "frameShifts": "frame_shifts", |
| 34 | + "deletions": "deletions", |
| 35 | + "insertions": "insertions", |
| 36 | + "substitutions": "substitutions", |
| 37 | + "aaSubstitutions": "aaSubstitutions" |
| 38 | +} |
| 39 | + |
| 40 | +preferred_types = { |
| 41 | + "divergence": "int32", |
| 42 | + "nonACGTN": "int32", |
| 43 | + "missing_data": "int32", |
| 44 | + "snp_clusters": "int32", |
| 45 | + "rare_mutations": "int32" |
| 46 | +} |
| 47 | + |
| 48 | +def reorder_columns(result: pd.DataFrame): |
| 49 | + """ |
| 50 | + Moves the new clade column after a specified column |
| 51 | + """ |
| 52 | + columns = list(result.columns) |
| 53 | + columns.remove(column_map['clade']) |
| 54 | + insert_at = columns.index(INSERT_BEFORE_THIS_COLUMN) |
| 55 | + columns.insert(insert_at, column_map['clade']) |
| 56 | + return result[columns] |
| 57 | + |
| 58 | + |
| 59 | +def parse_args(): |
| 60 | + parser = argparse.ArgumentParser( |
| 61 | + description="Joins metadata file with Nextclade clade output", |
| 62 | + ) |
| 63 | + parser.add_argument("first_file") |
| 64 | + parser.add_argument("second_file") |
| 65 | + parser.add_argument("-o", default=sys.stdout) |
| 66 | + return parser.parse_args() |
| 67 | + |
| 68 | +def datestr_to_ordinal(x): |
| 69 | + try: |
| 70 | + return datetime.strptime(x,"%Y-%m-%d").toordinal() |
| 71 | + except: |
| 72 | + return np.nan |
| 73 | + |
| 74 | +def isfloat(value): |
| 75 | + try: |
| 76 | + float(value) |
| 77 | + return True |
| 78 | + except ValueError: |
| 79 | + return False |
| 80 | + |
| 81 | +def main(): |
| 82 | + args = parse_args() |
| 83 | + |
| 84 | + metadata = pd.read_csv(args.first_file, index_col=METADATA_JOIN_COLUMN_NAME, |
| 85 | + sep='\t', low_memory=False) |
| 86 | + |
| 87 | + # Check for existing annotations in the given metadata. Skip join with |
| 88 | + # Nextclade QC file, if those annotations already exist and none of the |
| 89 | + # columns have empty values. In the case where metadata were combined from |
| 90 | + # different sources with and without annotations, the "clock_deviation" |
| 91 | + # column will exist but some values will be missing. We handle this case as |
| 92 | + # if the annotations do not exist at all and reannotate all columns. We |
| 93 | + # cannot look for missing values across all expected columns as evidence of |
| 94 | + # incomplete annotations, since a complete annotation by Nextclade will |
| 95 | + # include missing values for some columns by design. |
| 96 | + expected_columns = list(column_map.values()) + ["clock_deviation"] |
| 97 | + existing_annotation_columns = metadata.columns.intersection(expected_columns) |
| 98 | + if len(existing_annotation_columns) == len(expected_columns): |
| 99 | + if metadata["clock_deviation"].isnull().sum() == 0: |
| 100 | + print(f"Metadata file '{args.first_file}' has already been annotated with Nextclade QC columns. Skipping re-annotation.") |
| 101 | + metadata.to_csv(args.o, sep="\t") |
| 102 | + return |
| 103 | + |
| 104 | + # Read and rename clade column to be more descriptive |
| 105 | + clades = pd.read_csv(args.second_file, index_col=NEXTCLADE_JOIN_COLUMN_NAME, |
| 106 | + sep='\t', low_memory=False, na_filter = False) \ |
| 107 | + .rename(columns=column_map) |
| 108 | + |
| 109 | + clade_columns = clades.columns.intersection(list(column_map.values())) |
| 110 | + clades = clades[clade_columns] |
| 111 | + |
| 112 | + # Concatenate on columns |
| 113 | + result = pd.merge( |
| 114 | + metadata, clades, |
| 115 | + left_index=True, |
| 116 | + right_index=True, |
| 117 | + how='left', |
| 118 | + suffixes=["_original", ""], |
| 119 | + ) |
| 120 | + all_clades = result.Nextstrain_clade.unique() |
| 121 | + t = result["date"].apply(datestr_to_ordinal) |
| 122 | + div_array = np.array([float(x) if isfloat(x) else np.nan for x in result.divergence]) |
| 123 | + offset_by_clade = {} |
| 124 | + for clade in all_clades: |
| 125 | + ind = result.Nextstrain_clade==clade |
| 126 | + if ind.sum()>100: |
| 127 | + deviation = div_array[ind] - (t[ind] - reference_day)*rate_per_day |
| 128 | + offset_by_clade[clade] = np.mean(deviation[~np.isnan(deviation)]) |
| 129 | + |
| 130 | + # extract divergence, time and offset information into vectors or series |
| 131 | + offset = result["Nextstrain_clade"].apply(lambda x: offset_by_clade.get(x, 2.0)) |
| 132 | + # calculate divergence |
| 133 | + result["clock_deviation"] = np.array(div_array - ((t-reference_day)*rate_per_day + offset), dtype=int) |
| 134 | + result.loc[np.isnan(div_array)|np.isnan(t), "clock_deviation"] = np.nan |
| 135 | + |
| 136 | + for col in list(column_map.values()) + ["clock_deviation"]: |
| 137 | + result[col] = result[col].fillna(VALUE_MISSING_DATA) |
| 138 | + |
| 139 | + # Move the new column so that it's next to other clade columns |
| 140 | + if INSERT_BEFORE_THIS_COLUMN in result.columns: |
| 141 | + result = reorder_columns(result) #.astype(preferred_types) |
| 142 | + |
| 143 | + result.to_csv(args.o, index_label=METADATA_JOIN_COLUMN_NAME, sep='\t') |
| 144 | + |
| 145 | + |
| 146 | +if __name__ == '__main__': |
| 147 | + main() |
0 commit comments