"""Structure parsing, cleaning and feature encoding for the BindCraft-targetability GNN. Vendored from the PeSTo repo (`src/structure_io.py`, `src/structure.py`, `src/data_encoding.py`, `src/dataset.py`) so the Space is self-contained. The encoding must stay bit-identical to what the model was trained on -- do not "improve" the element/residue/atom vocabularies or the neighbour count. Two additions over the original: * `extract_topology_knn` streams the pairwise distance matrix in row blocks instead of materialising an [N, N, 3] displacement tensor. Same k-nearest neighbour indices, but memory grows as O(block * N) rather than O(N^2 * 3), which is what lets a 1000+ residue target run inside a Space container. * `map_scores_to_original` writes the scores back onto the *input* file's numbering. `clean_structure` renumbers residues consecutively from 1 across all chains, so a score at PeSTo residue 412 is generally not residue 412 of the deposited structure. Matching on CA coordinates avoids that trap. """ import os import numpy as np import torch as pt import gemmi # --------------------------------------------------------------------------- # encoding vocabularies (frozen -- part of the trained model's input contract) # --------------------------------------------------------------------------- # standard elements (sorted by abundance) (29 + unknown = 30 input channels) std_elements = np.array([ 'C', 'O', 'N', 'S', 'P', 'Se', 'Mg', 'Cl', 'Zn', 'Fe', 'Ca', 'Na', 'F', 'Mn', 'I', 'K', 'Br', 'Cu', 'Cd', 'Ni', 'Co', 'Sr', 'Hg', 'W', 'As', 'B', 'Mo', 'Ba', 'Pt' ]) # standard residue names: AA/RNA/DNA (sorted by abundance) (28) std_resnames = np.array([ 'LEU', 'GLU', 'ARG', 'LYS', 'VAL', 'ILE', 'PHE', 'ASP', 'TYR', 'ALA', 'THR', 'SER', 'GLN', 'ASN', 'PRO', 'GLY', 'HIS', 'TRP', 'MET', 'CYS', 'G', 'A', 'C', 'U', 'DG', 'DA', 'DT', 'DC' ]) # standard atom names contained in standard residues (sorted by abundance) (63) std_names = np.array([ 'CA', 'N', 'C', 'O', 'CB', 'CG', 'CD2', 'CD1', 'CG1', 'CG2', 'CD', 'OE1', 'OE2', 'OG', 'OG1', 'OD1', 'OD2', 'CE', 'NZ', 'NE', 'CZ', 'NH2', 'NH1', 'ND2', 'CE2', 'CE1', 'NE2', 'OH', 'ND1', 'SD', 'SG', 'NE1', 'CE3', 'CZ3', 'CZ2', 'CH2', 'P', "C3'", "C4'", "O3'", "C5'", "O5'", "O4'", "C1'", "C2'", "O2'", 'OP1', 'OP2', 'N9', 'N2', 'O6', 'N7', 'C8', 'N1', 'N3', 'C2', 'C4', 'C6', 'C5', 'N6', 'N4', 'O2', 'O4' ]) res3to1 = { 'CYS': 'C', 'ASP': 'D', 'SER': 'S', 'GLN': 'Q', 'LYS': 'K', 'ILE': 'I', 'PRO': 'P', 'THR': 'T', 'PHE': 'F', 'ASN': 'N', 'GLY': 'G', 'HIS': 'H', 'LEU': 'L', 'ARG': 'R', 'TRP': 'W', 'ALA': 'A', 'VAL': 'V', 'GLU': 'E', 'TYR': 'Y', 'MET': 'M', } # residues the model was trained to score; everything else (ligands, ions, # nucleic acids) still enters the graph as context but is not reported protein_resnames = set(res3to1) # --------------------------------------------------------------------------- # structure I/O # --------------------------------------------------------------------------- def read_pdb(pdb_filepath): """Parse a PDB/mmCIF file into the flat per-atom dict PeSTo operates on. PDB files go through `gemmi.read_pdb(..., max_line_length=80)`, exactly as PeSTo does -- that cap is what keeps the element column at 77-78 for files that append extra columns. Nothing here may drift from the reference reader: atom order and het_flag both feed the renumbering in `clean_structure`. """ if os.path.splitext(pdb_filepath)[1].lower() in ('.cif', '.mmcif'): st = gemmi.read_structure(pdb_filepath) st.setup_entities() else: st = gemmi.read_pdb(pdb_filepath, max_line_length=80) # altloc memory altloc_l = [] # data storage atom_element, atom_name, atom_xyz = [], [], [] residue_name, seq_id, het_flag, chain_name, icodes = [], [], [], [], [] # parse structure for mid, model in enumerate(st): for a in model.all(): # altloc check (keep first encountered) if a.atom.has_altloc(): key = f"{a.chain.name}_{a.residue.seqid.num}_{a.atom.name}" if key in altloc_l: continue else: altloc_l.append(key) # insertion code (kept for the renumbering logic, dropped afterwards) icodes.append(a.residue.seqid.icode.strip()) # store data atom_element.append(a.atom.element.name) atom_name.append(a.atom.name) atom_xyz.append([a.atom.pos.x, a.atom.pos.y, a.atom.pos.z]) residue_name.append(a.residue.name) seq_id.append(a.residue.seqid.num) het_flag.append(a.residue.het_flag) chain_name.append(f"{a.chain.name}:{mid}") if len(atom_xyz) == 0: raise ValueError("no atoms found in the input structure") # pack data return { 'xyz': np.array(atom_xyz, dtype=np.float32), 'name': np.array(atom_name), 'element': np.array(atom_element), 'resname': np.array(residue_name), 'resid': np.array(seq_id, dtype=np.int32), 'het_flag': np.array(het_flag), 'chain_name': np.array(chain_name), 'icode': np.array(icodes), } def save_pdb(subunits, filepath): """Write chains back out in PeSTo's own (renumbered) convention.""" with open(filepath, 'w') as fs: for cn in subunits: N = subunits[cn]['xyz'].shape[0] for i in range(N): h = "ATOM" if subunits[cn]['het_flag'][i] == 'A' else "HETATM" n = subunits[cn]['name'][i] rn = subunits[cn]['resname'][i] e = subunits[cn]['element'][i] ri = subunits[cn]['resid'][i] xyz = subunits[cn]['xyz'][i] bf = subunits[cn]['bfactor'][i] if "bfactor" in subunits[cn] else 0.0 # extract single character chain name c = cn.split(':')[0][0] pdb_line = ( "{:<6s}{:>5d} {:<4s} {:>3s} {:1s}{:>4d} " "{:8.3f}{:8.3f}{:8.3f}{:6.2f}{:6.2f} {:<2s} " ).format(h, i + 1, n, rn, c, ri, xyz[0], xyz[1], xyz[2], bf, bf, e) fs.write(pdb_line + '\n') fs.write("TER\n") fs.write("END") # --------------------------------------------------------------------------- # structure cleaning / splitting # --------------------------------------------------------------------------- def clean_structure(structure, rm_wat=True): """Drop water and hydrogens, then renumber residues 1..N across all chains.""" m_wat = (structure["resname"] == "HOH") m_h = (structure["element"] == "H") m_d = (structure["element"] == "D") m_hwat = (structure["resname"] == "DOD") if rm_wat: mask = ((~m_wat) & (~m_h) & (~m_d) & (~m_hwat)) else: mask = ((~m_h) & (~m_d) & (~m_hwat)) structure["resid"][m_wat] = -999 # filter structure atoms structure = {key: structure[key][mask] for key in structure} # find changes due to chain chains = structure["chain_name"] ids_chains = np.where(np.array(chains).reshape(-1, 1) == np.unique(chains).reshape(1, -1))[1] delta_chains = np.abs(np.sign(np.concatenate([[0], np.diff(ids_chains)]))) # find changes due to insertion code icodes = structure["icode"] ids_icodes = np.where(np.array(icodes).reshape(-1, 1) == np.unique(icodes).reshape(1, -1))[1] delta_icodes = np.abs(np.sign(np.concatenate([[0], np.diff(ids_icodes)]))) # find changes due to resids resids = structure["resid"] delta_resids = np.abs(np.sign(np.concatenate([[0], np.diff(resids)]))) # renumber resids structure['resid'] = np.cumsum(np.sign(delta_chains + delta_resids + delta_icodes)) + 1 # remove unnecessary icodes structure.pop("icode") return structure def atom_select(structure, sel): return {key: structure[key][sel] for key in structure} def split_by_chain(structure): chains = {} cnames = structure["chain_name"] ucnames = np.unique(cnames) m_chains = (cnames.reshape(-1, 1) == ucnames.reshape(1, -1)) for i in range(len(ucnames)): chain = atom_select(structure, m_chains[:, i]) chain.pop("chain_name") chains[ucnames[i]] = chain return chains def concatenate_chains(chains): keys = set.intersection(*[set(chains[cid]) for cid in chains]) structure = {key: np.concatenate([chains[cid][key] for cid in chains]) for key in keys} structure['chain_name'] = np.concatenate( [np.array([cid] * chains[cid]['xyz'].shape[0]) for cid in chains] ) return structure def tag_hetatm_chains(structure): """Give every HETATM group its own pseudo-chain so ligands split cleanly.""" m_hetatm = (structure['het_flag'] == "H") resids_hetatm = structure['resid'][m_hetatm] delta_hetatm = np.cumsum(np.abs(np.sign(np.concatenate([[0], np.diff(resids_hetatm)])))) cids_hetatm = np.array([f"{cid}:{hid}" for cid, hid in zip(structure['chain_name'][m_hetatm], delta_hetatm)]) cids = structure['chain_name'].copy().astype(np.dtype(' 1): subunits.pop(sname) return subunits def preprocess_structure(structure): """`StructuresDataset(with_preprocessing=True)` in one call.""" structure = clean_structure(structure) structure = tag_hetatm_chains(structure) subunits = split_by_chain(structure) subunits = filter_non_atomic_subunits(subunits) subunits = remove_duplicate_tagged_subunits(subunits) return subunits # --------------------------------------------------------------------------- # feature encoding # --------------------------------------------------------------------------- def onehot(x, v): m = (x.reshape(-1, 1) == np.array(v).reshape(1, -1)) return np.concatenate([m, ~np.any(m, axis=1).reshape(-1, 1)], axis=1) def encode_structure(structure, device=pt.device("cpu")): if isinstance(structure['xyz'], pt.Tensor): X = structure['xyz'].to(device) else: X = pt.from_numpy(structure['xyz'].astype(np.float32)).to(device) if isinstance(structure['resid'], pt.Tensor): resids = structure['resid'].to(device) else: resids = pt.from_numpy(structure['resid']).to(device) M = (resids.unsqueeze(1) == pt.unique(resids).unsqueeze(0)) return X, M def encode_features(structure, device=pt.device("cpu")): qe = pt.from_numpy(onehot(structure['element'], std_elements).astype(np.float32)).to(device) qr = pt.from_numpy(onehot(structure['resname'], std_resnames).astype(np.float32)).to(device) qn = pt.from_numpy(onehot(structure['name'], std_names).astype(np.float32)).to(device) return qe, qr, qn def extract_topology_knn(X, num_nn=64, block=None): """k-nearest-neighbour indices, computed in row blocks. Equivalent to PeSTo's `extract_topology(X, num_nn)[0]`: distances below 1e-2 (self, and any coincident atoms) are pushed past the global maximum so they never win the topk. Reproducing that needs the global max, hence two passes over the distance matrix; both are cheap next to the 32 message-passing layers. """ N = X.shape[0] if block is None: # ~4M floats of scratch per block, and never more than the whole matrix block = max(1, min(N, int(4e6) // max(N, 1))) # pass 1: global maximum distance d_max = pt.zeros((), dtype=X.dtype) for i in range(0, N, block): Dc = pt.norm(X.unsqueeze(0) - X[i:i + block].unsqueeze(1), dim=2) d_max = pt.maximum(d_max, pt.max(Dc)) # pass 2: masked topk per row block knn = min(num_nn, N) ids_topk = pt.zeros((N, knn), dtype=pt.long) for i in range(0, N, block): Dc = pt.norm(X.unsqueeze(0) - X[i:i + block].unsqueeze(1), dim=2) Dc = Dc + d_max * (Dc < 1e-2).float() ids_topk[i:i + block] = pt.topk(Dc, knn, dim=1, largest=False)[1] return ids_topk def collate_batch_features(batch_data, max_num_nn=64): """Pack one or more structures into a single graph, reserving index 0 as the sink.""" X = pt.cat([data[0] for data in batch_data], dim=0) q = pt.cat([data[2] for data in batch_data], dim=0) sizes = pt.tensor([data[3].shape for data in batch_data]) ids_topk = pt.zeros((X.shape[0], max_num_nn), dtype=pt.long, device=X.device) M = pt.zeros(pt.Size(pt.sum(sizes, dim=0)), dtype=pt.float, device=X.device) for size, data in zip(pt.cumsum(sizes, dim=0), batch_data): ix1 = size[0] ix0 = ix1 - data[3].shape[0] iy1 = size[1] iy0 = iy1 - data[3].shape[1] ids_topk[ix0:ix1, :data[1].shape[1]] = data[1] + ix0 + 1 M[ix0:ix1, iy0:iy1] = data[3] return X, ids_topk, q, M def encode_bfactor(structure, p): """Broadcast a per-residue score onto every atom of its residue.""" names = structure["name"] elements = structure["element"] het_flags = structure["het_flag"] m_ca = (names == "CA") & (elements == "C") & (het_flags == "A") resids = structure["resid"] if p.shape[0] == m_ca.shape[0]: structure['bfactor'] = p elif p.shape[0] == np.sum(m_ca): bf = np.zeros(len(resids), dtype=np.float32) for i in np.unique(resids): m_ri = (resids == i) i_rca = np.where(m_ri[m_ca])[0] if len(i_rca) > 0: bf[m_ri] = float(np.max(p[i_rca])) structure['bfactor'] = bf elif p.shape[0] == np.unique(resids).shape[0]: uresids = np.unique(resids) bf = np.zeros(len(resids), dtype=np.float32) for i in uresids: m_ri = (resids == i) m_uri = (uresids == i) bf[m_ri] = float(np.max(p[m_uri])) structure['bfactor'] = bf else: raise ValueError( f"cannot map {p.shape[0]} scores onto {len(resids)} atoms / " f"{np.unique(resids).shape[0]} residues" ) return structure # --------------------------------------------------------------------------- # mapping scores back onto the input file's numbering # --------------------------------------------------------------------------- def _coord_key(x, y, z): # PDB coordinates carry exactly three decimals, so this is an exact match, # not a tolerance -- cleaning only ever removes atoms, never moves them return (round(float(x), 3), round(float(y), 3), round(float(z), 3)) def map_scores_to_original(pdb_filepath, xyz, bfactors, out_filepath, unscored=-1.0): """Write per-atom scores into a copy of the input file, numbering untouched. `xyz`/`bfactors` come from the cleaned structure, whose residue ids PeSTo has renumbered. Atoms are matched back by coordinate, then each original residue takes the score of its matched atoms. Atoms PeSTo dropped -- water, hydrogens, alternate locations -- inherit their residue's score; residues with no match at all (a fully stripped water, say) get `unscored` so they are obviously not a prediction of 0. Returns (gemmi structure with B-factors set, list of per-residue records). """ score_by_coord = {_coord_key(*c): float(b) for c, b in zip(xyz, bfactors)} st = gemmi.read_structure(pdb_filepath) st.setup_entities() records = [] for model in st: for chain in model: for residue in chain: # collect the scores of this residue's matched atoms matched = [ score_by_coord[k] for k in (_coord_key(a.pos.x, a.pos.y, a.pos.z) for a in residue) if k in score_by_coord ] score = max(matched) if matched else unscored for atom in residue: atom.b_iso = score if matched: records.append({ 'chain': chain.name, 'resid': residue.seqid.num, 'icode': residue.seqid.icode.strip(), 'resname': residue.name, 'aa': res3to1.get(residue.name, 'X'), 'score': score, 'is_protein': residue.name in protein_resnames, }) if out_filepath is not None: st.write_pdb(out_filepath) return st, records