File size: 9,207 Bytes
f2ad668 | 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 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 | """Taxonomy-paired MSA construction for ESMFold2 inference.
Taxonomy IDs are read from FASTA headers as ``key=N`` tokens. Rows
where any chain has ``key=-1`` (or no ``key=`` at all) are treated as
unpaired and assigned to that chain's block-diagonal section after
the paired rows.
"""
import re
import numpy as np
from .esmfold2_constants import (
MSA_GAP_TOKEN_ID,
PROTEIN_3TO1,
PROTEIN_RESIDUE_TO_RES_TYPE,
PROTEIN_UNK_RES_TYPE,
)
from .esmfold2_msa import MSA
_KEY_RE = re.compile(r"key=(-?\d+)")
def protein_letter_to_res_type() -> dict[str, int]:
"""Return the protein 1-letter → res_type mapping used by the MSA encoder."""
mapping: dict[str, int] = {}
for three, one in PROTEIN_3TO1.items():
if three in PROTEIN_RESIDUE_TO_RES_TYPE:
mapping[one] = PROTEIN_RESIDUE_TO_RES_TYPE[three]
mapping["-"] = MSA_GAP_TOKEN_ID
mapping["X"] = PROTEIN_UNK_RES_TYPE
return mapping
def _taxonomy_from_header(header: str) -> int:
if not header:
return -1
m = _KEY_RE.search(header)
return int(m.group(1)) if m else -1
def msa_to_res_type_and_deletions(
msa: MSA, letter_to_res_type: dict[str, int]
) -> tuple[np.ndarray, np.ndarray]:
"""Convert an :class:`MSA` to ``(res_type[M, L], deletion_count[M, L])``.
Handles a3m insertion convention: lowercase letters and ``.`` are
insertions and are not emitted; their count is accumulated into the
next non-insertion position's deletion value. ``L`` is the query
length after stripping insertions from row 0.
"""
query = msa.entries[0].sequence
L = sum(1 for ch in query if not (ch.islower() or ch == "."))
M = msa.depth
res_type = np.full((M, L), MSA_GAP_TOKEN_ID, dtype=np.int64)
deletions = np.zeros((M, L), dtype=np.float32)
for r, entry in enumerate(msa.entries):
col = 0
ins = 0
for ch in entry.sequence:
if ch == "." or (ch.islower() and ch != "-"):
ins += 1
continue
if col >= L:
break
if ch == "-":
res_type[r, col] = MSA_GAP_TOKEN_ID
else:
res_type[r, col] = letter_to_res_type.get(
ch.upper(), PROTEIN_UNK_RES_TYPE
)
if ins > 0:
deletions[r, col] = float(ins)
ins = 0
col += 1
return res_type, deletions
def _dummy_msa_residues(query_res_types: np.ndarray) -> np.ndarray:
"""Single-row 'MSA' for chains without one — just the query."""
return query_res_types[None, :] # [1, L]
def construct_paired_msa(
chain_msas: dict[int, MSA | None],
chain_query_res_types: dict[int, np.ndarray],
token_asym_ids: np.ndarray,
token_res_ids: np.ndarray,
letter_to_res_type: dict[str, int] | None = None,
*,
max_pairs: int = 8192,
max_total: int = 16384,
max_seqs: int = 16384,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Build paired MSA features.
Parameters
----------
chain_msas
``asym_id -> MSA`` (or ``None`` for chains without an MSA).
chain_query_res_types
``asym_id -> np.ndarray[L_c]`` of res-type ids for the chain's
query. Used to build dummy MSAs when a chain has no MSA.
token_asym_ids
Per-token asym_id, length ``T``. Must be non-decreasing.
token_res_ids
Per-token residue index within chain, length ``T``.
letter_to_res_type
1-letter → res-type mapping. Defaults to
:func:`protein_letter_to_res_type`.
Returns
-------
msa_residues : ``np.ndarray[M, T]`` int64
deletion_value : ``np.ndarray[M, T]`` float32 (raw deletion counts; the
``arctan(/3) * pi/2`` transform is applied by the caller)
is_paired : ``np.ndarray[M, T]`` float32 broadcast of per-row,
per-chain paired flags.
"""
if letter_to_res_type is None:
letter_to_res_type = protein_letter_to_res_type()
chain_ids: list[int] = sorted(chain_msas.keys())
# Build per-chain (res_type, deletions, taxonomy) tables.
chain_res_type: dict[int, np.ndarray] = {}
chain_deletions: dict[int, np.ndarray] = {}
chain_taxonomies: dict[int, list[int]] = {}
for c in chain_ids:
m = chain_msas.get(c)
if m is None or m.depth == 0:
qres = chain_query_res_types[c]
chain_res_type[c] = _dummy_msa_residues(qres)
chain_deletions[c] = np.zeros((1, qres.shape[0]), dtype=np.float32)
chain_taxonomies[c] = [-1]
continue
rt, dl = msa_to_res_type_and_deletions(m, letter_to_res_type)
chain_res_type[c] = rt
chain_deletions[c] = dl
chain_taxonomies[c] = [_taxonomy_from_header(e.header) for e in m.entries]
# Group by taxonomy, skip query row and unpaired (-1) entries.
taxonomy_map: dict[int, list[tuple[int, int]]] = {}
for c in chain_ids:
for seq_idx, taxon in enumerate(chain_taxonomies[c]):
if seq_idx == 0 or taxon == -1:
continue
taxonomy_map.setdefault(taxon, []).append((c, seq_idx))
taxonomy_map = {k: v for k, v in taxonomy_map.items() if len(v) > 1}
# Order taxonomies by number of distinct chains, descending.
sorted_taxa = sorted(
taxonomy_map.items(), key=lambda kv: len({c for c, _ in kv[1]}), reverse=True
)
visited = {s for _, items in taxonomy_map.items() for s in items}
available: dict[int, list[int]] = {
c: [i for i in range(1, len(chain_taxonomies[c])) if (c, i) not in visited]
for c in chain_ids
}
pairing: list[dict[int, int]] = [{c: 0 for c in chain_ids}]
is_paired: list[dict[int, int]] = [{c: 1 for c in chain_ids}]
for _, pairs in sorted_taxa:
per_chain: dict[int, list[int]] = {}
for c, seq_idx in pairs:
per_chain.setdefault(c, []).append(seq_idx)
max_occ = max(len(v) for v in per_chain.values())
for i in range(max_occ):
row_pairing: dict[int, int] = {}
row_is_paired: dict[int, int] = {}
for c, seq_idxs in per_chain.items():
row_pairing[c] = seq_idxs[i % len(seq_idxs)]
row_is_paired[c] = 1
for c in chain_ids:
if c in row_pairing:
continue
row_is_paired[c] = 0
if available[c]:
row_pairing[c] = available[c].pop(0)
else:
row_pairing[c] = -1
pairing.append(row_pairing)
is_paired.append(row_is_paired)
if len(pairing) >= max_pairs:
break
if len(pairing) >= max_pairs:
break
max_left = max((len(v) for v in available.values()), default=0)
for _ in range(min(max_total - len(pairing), max_left)):
row_pairing = {}
row_is_paired = {}
for c in chain_ids:
row_is_paired[c] = 0
if available[c]:
row_pairing[c] = available[c].pop(0)
else:
row_pairing[c] = -1
pairing.append(row_pairing)
is_paired.append(row_is_paired)
if len(pairing) >= max_total:
break
pairing = pairing[:max_seqs]
is_paired = is_paired[:max_seqs]
M = len(pairing)
T = len(token_asym_ids)
msa_residues = np.full((M, T), MSA_GAP_TOKEN_ID, dtype=np.int64)
deletion_value = np.zeros((M, T), dtype=np.float32)
paired_mask = np.zeros((M, T), dtype=np.float32)
# Vectorize per chain: gather chain rows according to pairing[c], then
# index into them by the chain's token residue ids.
for c in chain_ids:
rt = chain_res_type[c]
dl = chain_deletions[c]
Lc = rt.shape[1]
chain_pairing = np.array([row[c] for row in pairing], dtype=np.int64)
chain_paired = np.array([row[c] for row in is_paired], dtype=np.float32)
token_mask = token_asym_ids == c
if not token_mask.any():
continue
token_res_in_chain = token_res_ids[token_mask]
# Clamp residue indices to the MSA's column range. Modified-residue
# tokens that exceed the query length fall back to the last column.
cols = np.minimum(token_res_in_chain, Lc - 1)
# Rows where pairing == -1 fall back to gap (already initialized).
valid_rows = chain_pairing >= 0
if valid_rows.any():
gathered_rt = rt[chain_pairing[valid_rows]][:, cols]
gathered_dl = dl[chain_pairing[valid_rows]][:, cols]
valid_idx = np.where(valid_rows)[0]
token_idx = np.where(token_mask)[0]
msa_residues[np.ix_(valid_idx, token_idx)] = gathered_rt
deletion_value[np.ix_(valid_idx, token_idx)] = gathered_dl
paired_mask[:, token_mask] = chain_paired[:, None]
return msa_residues, deletion_value, paired_mask
|