| """Protein-only ESMFold2 featurization without the Biohub runtime package. |
| |
| Input is one amino-acid sequence. The transformation expands each residue into |
| the checkpoint atom schema, pads atoms to a multiple of 32, and emits batched |
| token, atom, and single-sequence MSA tensors. Reference coordinates are loaded |
| lazily from a provenance-bearing declarative package asset. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from functools import cache |
| from importlib.resources import files |
| from typing import Any |
|
|
| import torch |
| from torch import Tensor |
|
|
| from .esmfold2_constants import ( |
| CHARGED_ATOMS, |
| ELEMENT_TO_ATOMIC_NUM, |
| ESM_PROTEIN_VOCAB, |
| MOL_TYPE_PROTEIN, |
| PROTEIN_1TO3, |
| PROTEIN_HEAVY_ATOMS, |
| PROTEIN_RESIDUE_TO_RES_TYPE, |
| PROTEIN_UNK_RES_TYPE, |
| ) |
|
|
| _GEOMETRY_ASSET = "protein_reference_geometry.json" |
| _GEOMETRY_SCHEMA = "fastplms.esmfold2.reference_geometry.v1" |
|
|
|
|
| @cache |
| def _reference_geometry() -> dict[str, dict[str, tuple[float, float, float]]]: |
| resource = files(__package__).joinpath(_GEOMETRY_ASSET) |
| with resource.open(mode="r", encoding="utf-8") as handle: |
| payload = json.load(handle) |
| if ( |
| payload.get("schema") != _GEOMETRY_SCHEMA |
| or payload.get("dtype") != "float32" |
| or payload.get("provenance", {}).get("manifest_family") != "esmfold2" |
| ): |
| raise RuntimeError("The ESMFold2 reference-geometry asset has invalid provenance.") |
|
|
| raw_residues = payload.get("residues") |
| if not isinstance(raw_residues, dict): |
| raise RuntimeError("The ESMFold2 reference-geometry asset has no residue table.") |
| geometry: dict[str, dict[str, tuple[float, float, float]]] = {} |
| for residue, atom_positions in raw_residues.items(): |
| if not isinstance(residue, str) or not isinstance(atom_positions, dict): |
| raise RuntimeError("The ESMFold2 reference-geometry residue table is malformed.") |
| geometry[residue] = {} |
| for atom_name, position in atom_positions.items(): |
| if ( |
| not isinstance(atom_name, str) |
| or not isinstance(position, list) |
| or len(position) != 3 |
| ): |
| raise RuntimeError("The ESMFold2 reference-geometry atom table is malformed.") |
| geometry[residue][atom_name] = tuple(float(value) for value in position) |
|
|
| expected_residues = set(PROTEIN_HEAVY_ATOMS) - {"MSE"} |
| if set(geometry) != expected_residues: |
| raise RuntimeError("The ESMFold2 reference-geometry residue set is incomplete.") |
| for residue, atom_names in PROTEIN_HEAVY_ATOMS.items(): |
| if residue == "MSE": |
| continue |
| if set(geometry[residue]) != set(atom_names): |
| raise RuntimeError(f"Reference geometry differs from the atom schema for {residue}.") |
| return geometry |
|
|
|
|
| def _encode_atom_name(atom_name: str) -> tuple[int, int, int, int]: |
| padded = atom_name.ljust(4)[:4] |
| return tuple(ord(character) - 32 if character != " " else 0 for character in padded) |
|
|
|
|
| def _padded_atom_count(actual_count: int) -> int: |
| return max(32, ((actual_count + 31) // 32) * 32) |
|
|
|
|
| def _residue_records(sequence: str) -> tuple[list[dict[str, Any]], list[int], list[int], list[int]]: |
| geometry = _reference_geometry() |
| atoms: list[dict[str, Any]] = [] |
| residue_types: list[int] = [] |
| input_ids: list[int] = [] |
| representative_atoms: list[int] = [] |
|
|
| for token_index, residue_letter in enumerate(sequence): |
| residue_name = PROTEIN_1TO3.get(residue_letter, "UNK") |
| atom_names = PROTEIN_HEAVY_ATOMS[residue_name] |
| atom_start = len(atoms) |
| for atom_name in atom_names: |
| atoms.append( |
| { |
| "token_index": token_index, |
| "name": atom_name, |
| "element": atom_name[0], |
| "charge": CHARGED_ATOMS.get((residue_name, atom_name), 0), |
| "position": geometry[residue_name][atom_name], |
| } |
| ) |
|
|
| representative_name = "CB" if "CB" in atom_names else "CA" |
| representative_atoms.append(atom_start + atom_names.index(representative_name)) |
| residue_types.append(PROTEIN_RESIDUE_TO_RES_TYPE.get(residue_name, PROTEIN_UNK_RES_TYPE)) |
| input_ids.append(ESM_PROTEIN_VOCAB.get(residue_letter, ESM_PROTEIN_VOCAB["X"])) |
|
|
| return atoms, residue_types, input_ids, representative_atoms |
|
|
|
|
| def prepare_protein_features(sequence: str) -> dict[str, Tensor]: |
| """Build the protein-only feature mapping consumed by ESMFold2. |
| |
| Every tensor includes a leading batch dimension. Biological tokens have |
| length ``l``; atom tensors have length ``n_atoms``, where ``n_atoms`` is the |
| smallest multiple of 32 covering all heavy atoms. |
| """ |
|
|
| if not sequence: |
| raise ValueError("sequence must be non-empty") |
|
|
| atoms, residue_types, input_ids, representative_atoms = _residue_records(sequence) |
| sequence_length = len(sequence) |
| n_atoms = _padded_atom_count(len(atoms)) |
|
|
| ref_pos = torch.zeros((n_atoms, 3), dtype=torch.float32) |
| ref_element = torch.zeros(n_atoms, dtype=torch.int64) |
| ref_charge = torch.zeros(n_atoms, dtype=torch.int8) |
| ref_atom_name_chars = torch.zeros((n_atoms, 4), dtype=torch.int64) |
| ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64) |
| atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool) |
| atom_to_token = torch.zeros(n_atoms, dtype=torch.int64) |
|
|
| for atom_index, atom in enumerate(atoms): |
| token_index = atom["token_index"] |
| ref_pos[atom_index] = torch.tensor(atom["position"], dtype=torch.float32) |
| ref_element[atom_index] = ELEMENT_TO_ATOMIC_NUM[atom["element"]] |
| ref_charge[atom_index] = atom["charge"] |
| ref_atom_name_chars[atom_index] = torch.tensor( |
| _encode_atom_name(atom["name"]), dtype=torch.int64 |
| ) |
| ref_space_uid[atom_index] = token_index |
| atom_attention_mask[atom_index] = True |
| atom_to_token[atom_index] = token_index |
|
|
| residue_type_tensor = torch.tensor(residue_types, dtype=torch.int64) |
| msa = residue_type_tensor.unsqueeze(0) |
| features = { |
| "token_index": torch.arange(sequence_length, dtype=torch.int64), |
| "residue_index": torch.arange(sequence_length, dtype=torch.int64), |
| "asym_id": torch.zeros(sequence_length, dtype=torch.int64), |
| "sym_id": torch.zeros(sequence_length, dtype=torch.int64), |
| "entity_id": torch.ones(sequence_length, dtype=torch.int64), |
| "mol_type": torch.full((sequence_length,), MOL_TYPE_PROTEIN, dtype=torch.int64), |
| "res_type": residue_type_tensor, |
| "input_ids": torch.tensor(input_ids, dtype=torch.int64), |
| "token_bonds": torch.zeros((sequence_length, sequence_length, 1), dtype=torch.float32), |
| "token_attention_mask": torch.ones(sequence_length, dtype=torch.bool), |
| "ref_pos": ref_pos, |
| "ref_element": ref_element, |
| "ref_charge": ref_charge, |
| "ref_atom_name_chars": ref_atom_name_chars, |
| "ref_space_uid": ref_space_uid, |
| "atom_attention_mask": atom_attention_mask, |
| "atom_to_token": atom_to_token, |
| "distogram_atom_idx": torch.tensor(representative_atoms, dtype=torch.int64), |
| "msa": msa, |
| "msa_attention_mask": torch.ones_like(msa, dtype=torch.bool), |
| "has_deletion": torch.zeros_like(msa, dtype=torch.bool), |
| "deletion_value": torch.zeros_like(msa, dtype=torch.float32), |
| "deletion_mean": torch.zeros(sequence_length, dtype=torch.float32), |
| } |
| return {name: tensor.unsqueeze(0) for name, tensor in features.items()} |
|
|
|
|
| __all__ = ["prepare_protein_features"] |
|
|
