File size: 3,369 Bytes
6aab6b3 | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | import gzip
from pathlib import Path
from typing import Optional, TextIO
import numpy as np
from redis import Redis
from boltzgen.data import const
from boltzgen.data.data import MSA, MSADeletion, MSAResidue, MSASequence
def _process_a3m(
lines: TextIO,
taxonomy: Optional[Redis] = None,
max_seqs: Optional[int] = None,
) -> MSA:
"""Process an MSA file.
Parameters
----------
lines : TextIO
The lines of the MsSA file.
taxonomy : Redis
The taxonomy database.
max_seqs : int, optional
The maximum number of sequences.
Returns
-------
MSA
The MSA object.
"""
visited = set()
sequences = []
deletions = []
residues = []
seq_idx = 0
for line in lines:
line: str
line = line.strip() # noqa: PLW2901
if not line or line.startswith("#"):
continue
# Get taxonomy, if annotated
if line.startswith(">"):
header = line.split()[0]
if taxonomy is None:
taxonomy_id = -1
elif header.startswith(">UniRef100"):
uniref_id = header.split("_")[1]
taxonomy_id = taxonomy.get(uniref_id)
if taxonomy_id is None:
taxonomy_id = -1
else:
taxonomy_id = -1
continue
# Skip if duplicate sequence
str_seq = line.replace("-", "").upper()
if str_seq not in visited:
visited.add(str_seq)
else:
continue
# Process sequence
residue = []
deletion = []
count = 0
res_idx = 0
for c in line:
if c != "-" and c.islower():
count += 1
continue
token = const.prot_letter_to_token[c]
token = const.token_ids[token]
residue.append(token)
if count > 0:
deletion.append((res_idx, count))
count = 0
res_idx += 1
res_start = len(residues)
res_end = res_start + len(residue)
del_start = len(deletions)
del_end = del_start + len(deletion)
sequences.append((seq_idx, taxonomy_id, res_start, res_end, del_start, del_end))
residues.extend(residue)
deletions.extend(deletion)
seq_idx += 1
if (max_seqs is not None) and (seq_idx >= max_seqs):
break
# Create MSA object
msa = MSA(
residues=np.array(residues, dtype=MSAResidue),
deletions=np.array(deletions, dtype=MSADeletion),
sequences=np.array(sequences, dtype=MSASequence),
)
return msa
def process_a3m(
path: Path,
taxonomy: Optional[Redis] = None,
max_seqs: Optional[int] = None,
) -> MSA:
"""Process an A3M file.
Parameters
----------
path : Path
The path to the a3m(.gz) file.
taxonomy : Redis
The taxonomy database.
max_seqs : int, optional
The maximum number of sequences.
Returns
-------
MSA
The MSA object.
"""
# Read the file
if path.suffix == ".gz":
with gzip.open(str(path), "rt") as f:
msa = _process_a3m(f, taxonomy, max_seqs)
else:
with path.open("r") as f:
msa = _process_a3m(f, taxonomy, max_seqs)
return msa
|