41 lines
2.3 KiB
Python
41 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
import argparse
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
import pandas as pd
|
|
from panfamily.io import iter_gff, read_fasta, write_fasta
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Choose one transcript per gene deterministically")
|
|
ap.add_argument("--accession", required=True); ap.add_argument("--gff3", required=True, type=Path)
|
|
ap.add_argument("--pep", required=True, type=Path); ap.add_argument("--cds", required=True, type=Path)
|
|
ap.add_argument("--out-prefix", required=True, type=Path); args = ap.parse_args()
|
|
pep, cds = read_fasta(args.pep), read_fasta(args.cds)
|
|
tx_to_gene: dict[str, str] = {}; gene_features = []
|
|
for seqid, source, feature, start, end, score, strand, phase, attrs in iter_gff(args.gff3):
|
|
if feature == "gene" and "ID" in attrs: gene_features.append((seqid, source, feature, start, end, score, strand, phase, attrs))
|
|
if feature in {"mRNA", "transcript"} and "ID" in attrs and "Parent" in attrs: tx_to_gene[attrs["ID"]] = attrs["Parent"].split(",")[0]
|
|
if not tx_to_gene:
|
|
common = set(pep) & set(cds); tx_to_gene = {item: item for item in common}
|
|
by_gene: dict[str, list[str]] = defaultdict(list)
|
|
for tx, gene in tx_to_gene.items():
|
|
if tx in pep and tx in cds: by_gene[gene].append(tx)
|
|
selected, mapping = {}, []
|
|
for gene, transcripts in sorted(by_gene.items()):
|
|
def rank(tx):
|
|
p, c = pep[tx], cds[tx]
|
|
complete = int(len(c) % 3 == 0 and not ("*" in p.rstrip("*")))
|
|
return (-complete, -len(c), -len(p), tx)
|
|
tx = sorted(transcripts, key=rank)[0]; standard = f"{args.accession}|{gene}"
|
|
selected[standard] = (pep[tx], cds[tx]); mapping.append({"accession": args.accession, "gene_id": gene, "transcript_id": tx, "standard_gene_id": standard, "protein_length": len(pep[tx]), "cds_length": len(cds[tx])})
|
|
args.out_prefix.parent.mkdir(parents=True, exist_ok=True)
|
|
write_fasta({k:v[0] for k,v in selected.items()}, str(args.out_prefix)+".primary.pep.fa")
|
|
write_fasta({k:v[1] for k,v in selected.items()}, str(args.out_prefix)+".primary.cds.fa")
|
|
pd.DataFrame(mapping).to_csv(str(args.out_prefix)+".id_mapping.tsv", sep="\t", index=False)
|
|
|
|
|
|
if __name__ == "__main__": main()
|
|
|