File size: 2,603 Bytes
d149fb1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | from argparse import ArgumentParser
from pathlib import Path
import json
def main(args):
# get the super index
with open(args.alignment_db_super_index_path, "r") as fp:
super_index = json.load(fp)
# get all chains and sequences
chains_to_seqs = {}
with open(args.all_chains_fasta, "r") as fp:
lines = fp.readlines()
# iterate through chain-sequence pairs
for chain_idx in range(0, len(lines), 2):
chain = lines[chain_idx][1:].strip()
seq = lines[chain_idx + 1].strip()
chains_to_seqs[chain] = seq
chains_w_alignments = set(super_index.keys())
chains_wo_alignments = set(chains_to_seqs.keys()) - chains_w_alignments
seq_to_chain_w_alignment = {
chains_to_seqs[chain]: chain for chain in chains_w_alignments
}
print("Unique sequences with alignments:", len(seq_to_chain_w_alignment))
# map chain without alignment to alignment entry of another chain with the
# same sequence
remaining_unaligned_chains = []
for chain in chains_wo_alignments:
seq = chains_to_seqs[chain]
try:
corresponding_alignment = super_index[seq_to_chain_w_alignment[seq]]
# no corresponding chain with alignment found
except KeyError:
remaining_unaligned_chains.append(chain)
continue
super_index[chain] = corresponding_alignment
with open(args.output_path, "w") as fp:
json.dump(super_index, fp)
print(
f"No corresponding alignment found for the following {len(remaining_unaligned_chains)} chains:",
remaining_unaligned_chains,
)
if __name__ == "__main__":
parser = ArgumentParser(
description="""
If the alignment-db index was created on unique-chain alignments only,
this will add the missing chain entries to the super-index file based on
a .fasta file that contains sequences for all chains.
Note that this only modifies the index and not the database itself, as
the duplicate sequences will just point to the same alignments.
"""
)
parser.add_argument(
"alignment_db_super_index_path",
type=Path,
help="Path to alignment-db super index file.",
)
parser.add_argument(
"output_path", type=Path, help="Write the output super index to this path."
)
parser.add_argument(
"all_chains_fasta",
type=Path,
help="Path to the fasta file containing sequences for all chains.",
)
args = parser.parse_args()
main(args)
|