perf(inference): drop training-only targets the inference forward never reads

nesso/data/featurizer.py

@@ -70,21 +70,6 @@ def extract_esm_features(
     return s_chain
 
 
-def _compute_disto_target(
-    disto_coords: Tensor,
-    min_dist: float = 2.0,
-    max_dist: float = 22.0,
-    num_bins: int = 64,
-) -> Tensor:
-    """Compute one-hot distogram from token disto coordinates."""
-    t_dists = torch.cdist(disto_coords.float(), disto_coords.float())
-    boundaries = torch.linspace(
-        min_dist, max_dist, num_bins - 1, device=disto_coords.device
-    )
-    distogram = (t_dists.unsqueeze(-1) > boundaries).sum(dim=-1).long().contiguous()
-    return one_hot(distogram, num_classes=num_bins).float()
-
-
 def select_subset_from_mask(mask: np.ndarray, p: float) -> np.ndarray:
     """Subsample True entries in mask using a geometric draw."""
     num_true = int(np.sum(mask))
@@ -164,12 +149,6 @@ def process_token_features(  # noqa: C901, PLR0915, PLR0912
     pad_mask = torch.ones(len(token_data), dtype=torch.float)
     disto_mask = from_numpy(token_data["disto_mask"].copy()).float()
 
-    # Distogram target from token disto_coords
-    disto_coords = from_numpy(token_data["disto_coords"].copy())
-    disto_target = _compute_disto_target(
-        disto_coords, min_dist=min_dist, max_dist=max_dist, num_bins=num_dist_bins
-    )
-
     # Token bond features
     if max_tokens is not None:
         pad_len = max_tokens - len(token_data)
@@ -275,7 +254,6 @@ def process_token_features(  # noqa: C901, PLR0915, PLR0912
             res_type = pad_dim(res_type, 0, pad_len)
             pad_mask = pad_dim(pad_mask, 0, pad_len)
             disto_mask = pad_dim(disto_mask, 0, pad_len)
-            disto_target = pad_dim(pad_dim(disto_target, 0, pad_len), 1, pad_len)
             unspecified_oh = torch.zeros(
                 pad_len, len(const.pocket_contact_info), dtype=torch.float32
             )
@@ -294,7 +272,6 @@ def process_token_features(  # noqa: C901, PLR0915, PLR0912
         "type_bonds": bonds_type,
         "token_pad_mask": pad_mask,
         "token_disto_mask": disto_mask,
-        "disto_target": disto_target,
         "pocket_feature": pocket_feature,
     }
 
@@ -353,7 +330,6 @@ def process_atom_features(
     ref_space_uid_list = []
     coord_data_list = []
     atom_to_token_list = []
-    token_to_rep_atom_list = []
     resolved_mask_list = []
 
     chain_res_ids = {}
@@ -428,11 +404,6 @@ def process_atom_features(
             conformer_pos = random.randn(n_atoms, 3).astype(np.float32)
         atom_conformer_list.append(conformer_pos)
 
-        disto_in_valid = next(
-            i for i, v in enumerate(valid_indices) if v == token["disto_idx"]
-        )
-        token_to_rep_atom_list.append(atom_idx + disto_in_valid)
-
         token_coords = structure.coords[offset + np.array(valid_indices)]["coords"]
         coord_data_list.append(token_coords[np.newaxis, ...])
         resolved_mask_list.append(token_atoms["is_present"].copy())
@@ -482,14 +453,6 @@ def process_atom_features(
         torch.tensor(atom_to_token_list, dtype=torch.long),
         num_classes=num_token_classes,
     )
-    token_to_rep_atom = one_hot(
-        torch.tensor(token_to_rep_atom_list, dtype=torch.long).clamp(
-            0, max(0, num_atoms - 1)
-        ),
-        num_classes=num_atoms,
-    )
-    if max_tokens is not None and L < max_tokens:
-        token_to_rep_atom = pad_dim(token_to_rep_atom, 0, max_tokens - L)
 
     pad_len = (
         (num_atoms - 1) // atoms_per_window_queries + 1
@@ -507,9 +470,8 @@ def process_atom_features(
         ref_space_uid = pad_dim(ref_space_uid, 0, pad_len)
         coords = pad_dim(coords, 1, pad_len)
         atom_to_token = pad_dim(atom_to_token, 0, pad_len)
-        token_to_rep_atom = pad_dim(token_to_rep_atom, 1, pad_len)
 
-    return {
+    atom_features = {
         "ref_pos": ref_pos,
         "atom_resolved_mask": resolved_mask,
         "ref_atom_name_chars": ref_atom_name_chars,
@@ -521,9 +483,9 @@ def process_atom_features(
         "coords": coords,
         "atom_pad_mask": pad_mask,
         "atom_to_token": atom_to_token,
-        "token_to_rep_atom": token_to_rep_atom,
         "res_index_to_conf_id": res_index_to_conf_id,
     }
+    return atom_features
 
 
 def process_esm_features(

tests/test_inference_features.py

@@ -0,0 +1,116 @@
+"""Inference featurization emits exactly the expected set of features.
+
+The forward pass never reads ``disto_target`` or ``token_to_rep_atom`` (the
+ground-truth distogram label and the token-to-representative-atom gather
+inherited from the Boltz training featurizer). Since the codebase is
+inference-only, those tensors are not built at all.
+
+Asserting the exact key set (rather than only the absence of those two) makes
+this a contract on the featurizer output: adding or dropping any feature without
+updating ``EXPECTED_FEATURE_KEYS`` fails here. Ligand-only input keeps this
+CCD-free so it runs on plain CI.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from numpy.random import RandomState
+
+from nesso.data.featurizer import NessoFeaturizer
+from nesso.data.inference import InferenceDataset
+from nesso.data.tokenize import tokenize_structure
+from nesso.data.types import Manifest, Structure, Tokenized
+from nesso.data.yaml_input import parse_yaml
+
+# The complete set of tensors `NessoFeaturizer.process` is expected to return.
+EXPECTED_FEATURE_KEYS = frozenset(
+    {
+        # token-level
+        "asym_id",
+        "entity_id",
+        "mol_type",
+        "pocket_feature",
+        "res_type",
+        "residue_index",
+        "sym_id",
+        "token_bonds",
+        "token_disto_mask",
+        "token_index",
+        "token_pad_mask",
+        "type_bonds",
+        # atom-level
+        "atom_pad_mask",
+        "atom_resolved_mask",
+        "atom_to_token",
+        "coords",
+        "ref_atom_name_chars",
+        "ref_charge",
+        "ref_chirality",
+        "ref_element",
+        "ref_hybridization",
+        "ref_pos",
+        "ref_space_uid",
+        # sequence embedding
+        "s_esm",
+    }
+)
+
+_LIGAND_SMILES = "Cc1ccc(NC(=O)c2ccc(CN3CCN(C)CC3)cc2)cc1Nc1nccc(-c2cccnc2)n1"
+_YAML = (
+    "version: 1\n"
+    "sequences:\n"
+    "  - ligand:\n"
+    "      id: B\n"
+    f'      smiles: "{_LIGAND_SMILES}"\n'
+)
+
+
+def _raw_features(tmp_path: Path) -> dict:
+    """Exactly what ``NessoFeaturizer.process`` returns."""
+    mol_dir = tmp_path / "rdkit_conformers"
+    structures_dir = tmp_path / "structures"
+    esm_dir = tmp_path / "esm"
+    for d in (mol_dir, structures_dir, esm_dir):
+        d.mkdir(parents=True, exist_ok=True)
+
+    yaml_path = tmp_path / "lig.yaml"
+    yaml_path.write_text(_YAML)
+    struct, record, _, _ = parse_yaml(
+        yaml_path, mol_dir, ccd_dict=None, record_id="lig"
+    )
+    struct.dump(structures_dir / f"{record.id}.npz")
+
+    struct = Structure.load(structures_dir / f"{record.id}.npz")
+    struct = struct.remove_invalid_chains(struct.mask.copy())
+    tokens, bonds = tokenize_structure(struct)
+    tokenized = Tokenized(tokens=tokens, bonds=bonds, structure=struct, record=record)
+
+    ds = InferenceDataset(
+        manifest=Manifest([record]),
+        target_dir=tmp_path,
+        featurizer=NessoFeaturizer(
+            esm_emb_dir=esm_dir, esm_emb_dim=1280, esm_num_layers=33
+        ),
+        ligand_dir=mol_dir,
+        ccd_pkl=None,
+    )
+    return ds.featurizer.process(
+        tokenized,
+        record=record,
+        molecules=ds._setup_molecules(struct, str(record.id)),
+        random=RandomState(0),
+        atoms_per_window_queries=32,
+        binder_pocket_conditioned_prop=0.0,
+        max_tokens=None,
+    )
+
+
+def test_featurizer_emits_expected_keys(tmp_path: Path) -> None:
+    """The featurizer output must match the expected feature set exactly."""
+    keys = set(_raw_features(tmp_path))
+
+    missing = EXPECTED_FEATURE_KEYS - keys
+    unexpected = keys - EXPECTED_FEATURE_KEYS
+    assert not missing, f"missing expected features: {sorted(missing)}"
+    assert not unexpected, f"unexpected features: {sorted(unexpected)}"