priorities.py

"""
calculate priorties from index and proximities
"""
importargparse
fromrandomimportshuffle
fromcollectionsimportdefaultdict
importnumpyasnp
importpandasaspd

if__name__=='__main__':
parser=argparse.ArgumentParser(
description="generate priorities files based on genetic proximity to focal sample",
formatter_class=argparse.ArgumentDefaultsHelpFormatter
 )
parser.add_argument("--sequence-index", type=str, required=True, help="sequence index file")
parser.add_argument("--proximities", type=str, required=True, help="tsv file with proximities")
parser.add_argument("--Nweight", type=float, default=0.003, required=False, help="parameterizes de-prioritization of incomplete sequences")
parser.add_argument("--crowding-penalty", type=float, default=0.1, required=False, help="parameterizes how priorities decrease when there is many very similar sequences")
parser.add_argument("--output", type=str, required=True, help="tsv file with the priorities")
args=parser.parse_args()

proximities=pd.read_csv(args.proximities, sep='\t', index_col=0)
index=pd.read_csv(args.sequence_index, sep='\t', index_col=0)
combined=pd.concat([proximities, index], axis=1)

closest_matches=combined.groupby('closest strain')
candidates= {}
forfocal_seq, seqsinclosest_matches.groups.items():
tmp=combined.loc[seqs, ["distance", "N"]]
# penalize larger distances and more undetermined sites. 1/args.Nweight are 'as bad' as one extra mutation
tmp["priority"] =-tmp.distance-tmp.N*args.Nweight
name_prior= [(name, d.priority) forname, dintmp.iterrows()]
shuffle(name_prior)
candidates[focal_seq] =sorted(name_prior, key=lambdax:x[1], reverse=True)

# export priorities
crowding=args.crowding_penalty
withopen(args.output, 'w') asfh:
# loop over lists of sequences that are closest to particular focal sequences
forcsincandidates.values():
# these sets have been sorted by priorities after shuffling -- reduce priorities in this shuffled/sorted order
fori, (name, pr) inenumerate(cs):
fh.write(f"{name}\t{pr-i*crowding:1.2f}\n")