join-metadata-and-clades.py

#!/usr/bin/env python3
"""
Copied from https://github.com/nextstrain/ncov-ingest/blob/master/bin/join-metadata-and-clades
"""
importargparse
importsys
fromdatetimeimportdatetime
importpandasaspd
importnumpyasnp

INSERT_BEFORE_THIS_COLUMN="pango_lineage"
METADATA_JOIN_COLUMN_NAME='strain'
NEXTCLADE_JOIN_COLUMN_NAME='seqName'
VALUE_MISSING_DATA='?'

rate_per_day=0.0007*29903/365
reference_day=datetime(2020,1,1).toordinal()

column_map= {
"clade": "Nextstrain_clade",
"Nextclade_pango": "Nextclade_pango",
"totalMissing": "missing_data",
"totalSubstitutions": "divergence",
"totalNonACGTNs": "nonACGTN",
"privateNucMutations.totalUnlabeledSubstitutions": "rare_mutations",
"privateNucMutations.totalReversionSubstitutions": "reversion_mutations",
"privateNucMutations.totalLabeledSubstitutions": "potential_contaminants",
"qc.overallScore": "QC_overall_score",
"qc.overallStatus": "QC_overall_status",
"qc.missingData.status": "QC_missing_data",
"qc.mixedSites.status": "QC_mixed_sites",
"qc.privateMutations.status": "QC_rare_mutations",
"qc.snpClusters.status": "QC_snp_clusters",
"qc.frameShifts.status": "QC_frame_shifts",
"qc.stopCodons.status": "QC_stop_codons",
"frameShifts": "frame_shifts",
"deletions": "deletions",
"insertions": "insertions",
"substitutions": "substitutions",
"aaSubstitutions": "aaSubstitutions",
"immune_escape": "immune_escape",
"ace2_binding": "ace2_binding",
}

preferred_types= {
"divergence": "int32",
"nonACGTN": "int32",
"missing_data": "int32",
"snp_clusters": "int32",
"rare_mutations": "int32"
}

defreorder_columns(result: pd.DataFrame):
"""
 Moves the new clade column after a specified column
 """
columns=list(result.columns)
columns.remove(column_map['clade'])
insert_at=columns.index(INSERT_BEFORE_THIS_COLUMN)
columns.insert(insert_at, column_map['clade'])
returnresult[columns]


defparse_args():
parser=argparse.ArgumentParser(
description="Joins metadata file with Nextclade clade output",
 )
parser.add_argument("first_file")
parser.add_argument("second_file")
parser.add_argument("-o", default=sys.stdout)
returnparser.parse_args()

defdatestr_to_ordinal(x):
try:
returndatetime.strptime(x,"%Y-%m-%d").toordinal()
except:
returnnp.nan

defisfloat(value):
try:
float(value)
returnTrue
exceptValueError:
returnFalse

defmain():
args=parse_args()

metadata=pd.read_csv(args.first_file, index_col=METADATA_JOIN_COLUMN_NAME,
sep='\t', low_memory=False)

# Check for existing annotations in the given metadata. Skip join with
# Nextclade QC file, if those annotations already exist and none of the
# columns have empty values. In the case where metadata were combined from
# different sources with and without annotations, the "clock_deviation"
# column will exist but some values will be missing. We handle this case as
# if the annotations do not exist at all and reannotate all columns. We
# cannot look for missing values across all expected columns as evidence of
# incomplete annotations, since a complete annotation by Nextclade will
# include missing values for some columns by design.
expected_columns=list(column_map.values()) + ["clock_deviation"]
existing_annotation_columns=metadata.columns.intersection(expected_columns)
iflen(existing_annotation_columns) ==len(expected_columns):
ifmetadata["clock_deviation"].isnull().sum() ==0:
print(f"Metadata file '{args.first_file}' has already been annotated with Nextclade QC columns. Skipping re-annotation.")
metadata.to_csv(args.o, sep="\t")
return

# Read and rename clade column to be more descriptive
clades=pd.read_csv(args.second_file, index_col=NEXTCLADE_JOIN_COLUMN_NAME,
sep='\t', low_memory=False, na_filter=False) \
 .rename(columns=column_map)

clade_columns=clades.columns.intersection(list(column_map.values()))
clades=clades[clade_columns]

# Concatenate on columns
result=pd.merge(
metadata, clades,
left_index=True,
right_index=True,
how='left',
suffixes=["_original", ""],
 )
all_clades=result.Nextstrain_clade.unique()
t=result["date"].apply(datestr_to_ordinal)
div_array=np.array([float(x) ifisfloat(x) elsenp.nanforxinresult.divergence])
offset_by_clade= {}
forcladeinall_clades:
ind=result.Nextstrain_clade==clade
ifind.sum()>100:
deviation=div_array[ind] - (t[ind] -reference_day)*rate_per_day
offset_by_clade[clade] =np.mean(deviation[~np.isnan(deviation)])

# extract divergence, time and offset information into vectors or series
offset=result["Nextstrain_clade"].apply(lambdax: offset_by_clade.get(x, 2.0))
# calculate divergence
result["clock_deviation"] =np.array(div_array- ((t-reference_day)*rate_per_day+offset), dtype=int)
result.loc[np.isnan(div_array)|np.isnan(t), "clock_deviation"] =np.nan

forcolinlist(column_map.values()) + ["clock_deviation"]:
ifcolinresult:
result[col] =result[col].fillna(VALUE_MISSING_DATA)

# Move the new column so that it's next to other clade columns
ifINSERT_BEFORE_THIS_COLUMNinresult.columns:
result=reorder_columns(result) #.astype(preferred_types)

result.to_csv(args.o, index_label=METADATA_JOIN_COLUMN_NAME, sep='\t')


if__name__=='__main__':
main()