preprocessing.py

"""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('<U10'))
    cids[m_hetatm] = cids_hetatm
    structure['chain_name'] = np.array(list(cids)).astype(str)

    return structure


def remove_duplicate_tagged_subunits(subunits):
    tagged_cids = [cid for cid in subunits if (len(cid.split(':')) == 3)]
    for i in range(len(tagged_cids)):
        cid_i = tagged_cids[i]
        for j in range(i + 1, len(tagged_cids)):
            cid_j = tagged_cids[j]

            if (cid_i in subunits) and (cid_j in subunits):
                xyz0 = subunits[cid_i]['xyz']
                xyz1 = subunits[cid_j]['xyz']

                if xyz0.shape[0] == xyz1.shape[0]:
                    d_min = np.min(np.linalg.norm(xyz0 - xyz1, axis=1))
                    if d_min < 0.2:
                        subunits.pop(cid_j)

    return subunits


def filter_non_atomic_subunits(subunits):
    """Drop coarse-grained subunits (one atom per residue), e.g. CA-only traces."""
    for sname in list(subunits):
        n_res = np.unique(subunits[sname]['resid']).shape[0]
        n_atm = subunits[sname]['xyz'].shape[0]

        if (n_atm == n_res) & (n_atm > 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