diff --git a/nesso/data/featurizer.py b/nesso/data/featurizer.py index 742304a..5efb2ba 100644 --- a/nesso/data/featurizer.py +++ b/nesso/data/featurizer.py @@ -469,12 +469,20 @@ def process_atom_features( center = center / resolved_mask.sum().clamp(min=1) coords = coords - center[:, None] - # Apply random roto-translation to the input conformers + # Apply random roto-translation to the input conformers. + # The generator is seeded from `random` (a per-record RandomState) rather than + # drawn from the global torch RNG, which is seeded per DataLoader worker and so + # would make `ref_pos` depend on `--num_workers`. + aug_generator = torch.Generator() + aug_generator.manual_seed(int(random.randint(0, 2**31 - 1))) for i in range(torch.max(ref_space_uid) + 1): included = ref_space_uid == i if torch.sum(included) > 0 and torch.any(resolved_mask[included]): ref_pos[included] = center_random_augmentation( - ref_pos[included][None], resolved_mask[included][None], centering=True + ref_pos[included][None], + resolved_mask[included][None], + centering=True, + generator=aug_generator, )[0] num_token_classes = max_tokens if max_tokens is not None else L diff --git a/nesso/data/inference.py b/nesso/data/inference.py index 4a5e2ab..26c7626 100644 --- a/nesso/data/inference.py +++ b/nesso/data/inference.py @@ -3,6 +3,7 @@ Adapted from https://github.com/jwohlwend/boltz, MIT License, Copyright (c) 2024 Jeremy Wohlwend. """ +import hashlib import pickle from dataclasses import replace from pathlib import Path @@ -109,6 +110,21 @@ def collate( INFERENCE_EXCLUDED_KEYS = frozenset({"record", "exception", "record_id"}) +# numpy seeds must fit in uint32. +_MAX_NUMPY_SEED = 2**32 + + +def record_rng_seed(record_id: str) -> int: + """Per-record RNG seed derived from the record id. + + Keyed on identity rather than on the dataset index, so a record's features do + not change when the surrounding batch does. The index shifts whenever inputs + are added, removed, renamed, or reordered, and ``hash()`` is salted per + process, so neither is usable here. + """ + digest = hashlib.sha256(record_id.encode("utf-8")).digest() + return int.from_bytes(digest[:4], "big") % _MAX_NUMPY_SEED + def inference_collate(data: list[dict[str, Tensor]]) -> dict[str, Tensor]: # Must be module-level (picklable) for DataLoader multiprocessing on macOS/Windows. @@ -199,7 +215,7 @@ def __getitem__(self, idx: int) -> dict[str, Any]: structure=structure, record=record, ) - random = RandomState(idx) + random = RandomState(record_rng_seed(str(record_id))) molecules = self._setup_molecules(structure, str(record_id)) features = self.featurizer.process( tokenized, diff --git a/nesso/data/yaml_input.py b/nesso/data/yaml_input.py index de9fe51..57989aa 100644 --- a/nesso/data/yaml_input.py +++ b/nesso/data/yaml_input.py @@ -107,10 +107,42 @@ def _load_mol( return pickle.load(f) # noqa: S301 -def get_conformer(mol: Chem.Mol) -> Chem.Conformer: +DEFAULT_CONFORMER_SEED = 42 + +# RDKit takes a signed 32-bit seed; -1 means "pick a random one". +_MAX_RDKIT_SEED = 2**31 - 1 + + +def conformer_seed(mol: Chem.Mol, base_seed: int) -> int: + """Derive a stable ETKDG seed for ``mol``. + + The seed is a function of ``base_seed`` and the molecule's canonical SMILES, + deliberately **not** of processing order: ``preprocess_yamls`` parses inputs in a + ``ProcessPoolExecutor``, so a counter-based scheme would vary with worker + scheduling and with ``--num_workers``. Keying on the molecule also means the same + ligand embeds identically wherever it appears. + """ + try: + key = Chem.MolToSmiles(mol) + except Exception: # unsanitized/exotic mol: fall back to a structural key + key = f"{mol.GetNumAtoms()}:{mol.GetNumBonds()}" + digest = hashlib.sha256(f"{base_seed}:{key}".encode("utf-8")).digest() + return int.from_bytes(digest[:4], "big") % _MAX_RDKIT_SEED + + +def get_conformer( + mol: Chem.Mol, base_seed: int | None = DEFAULT_CONFORMER_SEED +) -> Chem.Conformer: + """Return conformer 0, embedding one with ETKDG if the molecule has none. + + ``base_seed`` makes embedding reproducible; pass ``None`` for RDKit's default + non-deterministic behaviour. + """ if mol.GetNumConformers() == 0: opts = AllChem.ETKDGv3() opts.clearConfs = False + if base_seed is not None: + opts.randomSeed = conformer_seed(mol, base_seed) cid = AllChem.EmbedMolecule(mol, opts) if cid < 0: opts.useRandomCoords = True @@ -124,6 +156,7 @@ def _standard_residue( res_idx: int, *, ccd_dict: dict[str, Chem.Mol] | None = None, + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> dict[str, Any]: name = "MET" if res_name == "MSE" else res_name if name not in const.ref_atoms or not const.ref_atoms[name]: @@ -131,7 +164,7 @@ def _standard_residue( raise ValueError(msg) mol = _load_mol(mol_dir, name, ccd_dict=ccd_dict) mol = Chem.RemoveHs(mol, sanitize=False) - conf = get_conformer(mol) + conf = get_conformer(mol, base_seed) by_name = {a.GetProp("name"): a for a in mol.GetAtoms()} atoms: list[dict[str, Any]] = [] for atom_name in const.ref_atoms[name]: @@ -162,21 +195,27 @@ def _protein_residues( mol_dir: Path | None, *, ccd_dict: dict[str, Chem.Mol] | None = None, + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[int, list[dict[str, Any]]]: cmap = const.prot_letter_to_token unk = const.unk_token["PROTEIN"] mol_type = const.chain_type_ids["PROTEIN"] residues = [ - _standard_residue(cmap.get(c, unk), mol_dir, j, ccd_dict=ccd_dict) + _standard_residue( + cmap.get(c, unk), mol_dir, j, ccd_dict=ccd_dict, base_seed=base_seed + ) for j, c in enumerate(raw_seq) ] return mol_type, residues def _ligand_residue_from_mol( - mol: Chem.Mol, res_name: str, res_idx: int + mol: Chem.Mol, + res_name: str, + res_idx: int, + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> dict[str, Any]: - conf = get_conformer(mol) + conf = get_conformer(mol, base_seed) idx_map: dict[int, int] = {} atoms: list[dict[str, Any]] = [] for atom in mol.GetAtoms(): @@ -222,6 +261,7 @@ def _ligand_from_conformer_pkl( mol_dir: Path | None = None, rid: str = "", record_id: str = "", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[int, list[dict[str, Any]]]: """Load an RDKit ``Mol`` (heavy-atom, with conformer) from ``path`` as a ligand residue.""" with Path(path).open("rb") as f: @@ -236,7 +276,7 @@ def _ligand_from_conformer_pkl( atom.SetProp("name", f"{sym}{int(rnk) + 1}"[:4]) _dump_ligand_mol(mol, mol_dir, rid, record_id=record_id) return const.chain_type_ids["NONPOLYMER"], [ - _ligand_residue_from_mol(mol, lig_tag, 0) + _ligand_residue_from_mol(mol, lig_tag, 0, base_seed) ] @@ -263,6 +303,7 @@ def _ligand_smiles_residue( rid: str = "", *, record_id: str = "", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[int, list[dict[str, Any]]]: mol = Chem.MolFromSmiles(smi) if mol is None: @@ -270,12 +311,12 @@ def _ligand_smiles_residue( raise ValueError(msg) mol = Chem.AddHs(mol) Chem.AssignStereochemistry(mol, force=True, cleanIt=True) - get_conformer(mol) + get_conformer(mol, base_seed) mol_nh = Chem.RemoveHs(mol, sanitize=False) _assign_ligand_atom_names(mol_nh) _dump_ligand_mol(mol_nh, mol_dir, rid, record_id=record_id) return const.chain_type_ids["NONPOLYMER"], [ - _ligand_residue_from_mol(mol_nh, rid or "LIG", 0) + _ligand_residue_from_mol(mol_nh, rid or "LIG", 0, base_seed) ] @@ -286,14 +327,15 @@ def _ligand_ccd_residue( ccd_dict: dict[str, Chem.Mol] | None = None, rid: str = "", record_id: str = "", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[int, list[dict[str, Any]]]: mol = _load_mol(mol_dir, ccd_code, ccd_dict=ccd_dict) mol = Chem.RemoveHs(mol, sanitize=False) _assign_ligand_atom_names(mol) - get_conformer(mol) + get_conformer(mol, base_seed) _dump_ligand_mol(mol, mol_dir, rid, record_id=record_id) return const.chain_type_ids["NONPOLYMER"], [ - _ligand_residue_from_mol(mol, ccd_code, 0) + _ligand_residue_from_mol(mol, ccd_code, 0, base_seed) ] @@ -303,6 +345,7 @@ def _ligand_sdf_residue( rid: str = "", *, record_id: str = "", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[int, list[dict[str, Any]]]: path = Path(sdf_path) if not path.exists(): @@ -313,10 +356,10 @@ def _ligand_sdf_residue( msg = f"Invalid SDF: {path}" raise ValueError(msg) _assign_ligand_atom_names(mol) - get_conformer(mol) + get_conformer(mol, base_seed) _dump_ligand_mol(mol, mol_dir, rid, record_id=record_id) return const.chain_type_ids["NONPOLYMER"], [ - _ligand_residue_from_mol(mol, rid or "LIG", 0) + _ligand_residue_from_mol(mol, rid or "LIG", 0, base_seed) ] @@ -344,10 +387,11 @@ def _chain_data_for_entity( ccd_dict: dict[str, Chem.Mol] | None = None, rid: str = "", record_id: str = "", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> _ChainData: if isinstance(entity, _EntityProtein): mol_type, residues = _protein_residues( - entity.sequence, mol_dir, ccd_dict=ccd_dict + entity.sequence, mol_dir, ccd_dict=ccd_dict, base_seed=base_seed ) elif isinstance(entity, _EntityLigandSmiles): smiles_ligand_idx[0] += 1 @@ -358,10 +402,15 @@ def _chain_data_for_entity( mol_dir=mol_dir, rid=rid, record_id=record_id, + base_seed=base_seed, ) else: mol_type, residues = _ligand_smiles_residue( - entity.smiles, mol_dir=mol_dir, rid=rid, record_id=record_id + entity.smiles, + mol_dir=mol_dir, + rid=rid, + record_id=record_id, + base_seed=base_seed, ) elif isinstance(entity, _EntityLigandCCD): mol_type, residues = _ligand_ccd_residue( @@ -370,10 +419,15 @@ def _chain_data_for_entity( ccd_dict=ccd_dict, rid=rid, record_id=record_id, + base_seed=base_seed, ) elif isinstance(entity, _EntityLigandSDF): mol_type, residues = _ligand_sdf_residue( - entity.sdf, mol_dir=mol_dir, rid=rid, record_id=record_id + entity.sdf, + mol_dir=mol_dir, + rid=rid, + record_id=record_id, + base_seed=base_seed, ) else: raise TypeError(entity) @@ -592,9 +646,13 @@ def parse_schema( *, ccd_dict: dict[str, Chem.Mol] | None = None, record_id: str = "record", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[Structure, Record, dict[int, str], dict[int, str]]: """Parse a YAML schema dict into ``(Structure, Record, entity_to_seq, entity_to_esm_path)`` + + ``base_seed`` seeds ETKDG conformer embedding so parsing is reproducible; + pass ``None`` for RDKit's default non-deterministic behaviour. """ if not isinstance(schema, dict) or schema.get("version", 1) != 1: raise ValueError("Schema must be a dict with version: 1") @@ -630,6 +688,7 @@ def parse_schema( ccd_dict=ccd_dict, rid=lid, record_id=record_id, + base_seed=base_seed, ) data.entity_id = entity_id if isinstance(entity, _EntityProtein): @@ -671,11 +730,13 @@ def parse_yaml( ccd_pkl: Path | None = None, ccd_dict: dict[str, Chem.Mol] | None = None, record_id: str | None = None, + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[Structure, Record, dict[int, str], dict[int, str]]: """Parse YAML into ``(Structure, Record, entity_to_seq, entity_to_esm_path)``. ``ccd_dict`` (in-memory) takes precedence over ``ccd_pkl`` (disk); ``_load_mol`` raises if a standard residue needs a CCD source that was not provided. + ``base_seed`` seeds ETKDG conformer embedding so parsing is reproducible. """ if ccd_dict is None and ccd_pkl is not None: ccd_dict = load_ccd_mol_dict(ccd_pkl) @@ -684,7 +745,9 @@ def parse_yaml( schema = yaml.safe_load(f) rid = record_id if record_id is not None else path.stem - return parse_schema(schema, mol_dir, ccd_dict=ccd_dict, record_id=rid) + return parse_schema( + schema, mol_dir, ccd_dict=ccd_dict, record_id=rid, base_seed=base_seed + ) def esm_keys(entity_to_seq: dict[int, str]) -> dict[str, str]: diff --git a/nesso/main.py b/nesso/main.py index c29a46e..55af005 100644 --- a/nesso/main.py +++ b/nesso/main.py @@ -24,7 +24,11 @@ from nesso.data.inference import NessoInferenceDataModule from nesso.data.types import Manifest, Record from nesso.data.writer import NessoWriter -from nesso.data.yaml_input import parse_yaml, validate_schema +from nesso.data.yaml_input import ( + DEFAULT_CONFORMER_SEED, + parse_yaml, + validate_schema, +) from nesso.model.models.nesso1 import Nesso1 from nesso.data.esm import ( @@ -164,8 +168,11 @@ def _process_single_yaml( mol_dir: Path, structures_dir: Path, records_dir: Path, + seed: int, ) -> Record: - struct, rec, _, _ = parse_yaml(yp, mol_dir, ccd_dict=_worker_ccd_dict) + struct, rec, _, _ = parse_yaml( + yp, mol_dir, ccd_dict=_worker_ccd_dict, base_seed=seed + ) struct.dump(structures_dir / f"{rec.id}.npz") rec.dump(records_dir / f"{rec.id}.json") return rec @@ -178,6 +185,7 @@ def preprocess_yamls( structures_dir: Path, records_dir: Path, num_workers: int = 2, + seed: int = DEFAULT_CONFORMER_SEED, ) -> tuple[Manifest, list[str]]: """Parse YAMLs into a Manifest, reporting which inputs failed. @@ -188,26 +196,32 @@ def preprocess_yamls( """ structures_dir.mkdir(parents=True, exist_ok=True) records_dir.mkdir(parents=True, exist_ok=True) - records: list[Record] = [] failed: list[str] = [] + # Results are collected by submission index, not completion order: + # `as_completed` yields whichever worker finishes first, so appending here + # would make the manifest order depend on scheduling and hence on + # `--num_workers`. Downstream that order is visible as the dataset index. + by_index: dict[int, Record] = {} + with ProcessPoolExecutor( - max_workers=num_workers, initializer=_init_worker, initargs=(ccd_pkl,) + max_workers=max(1, num_workers), initializer=_init_worker, initargs=(ccd_pkl,) ) as executor: futures = { executor.submit( - _process_single_yaml, yp, mol_dir, structures_dir, records_dir - ): yp - for yp in yaml_paths + _process_single_yaml, yp, mol_dir, structures_dir, records_dir, seed + ): (i, yp) + for i, yp in enumerate(yaml_paths) } for future in tqdm(as_completed(futures), total=len(futures), desc="YAML"): - yp = futures[future] + i, yp = futures[future] try: - records.append(future.result()) + by_index[i] = future.result() except Exception as e: failed.append(yp.stem) print(f"Error processing YAML {yp.name}: {e}", file=sys.stderr) + records = [by_index[i] for i in sorted(by_index)] return Manifest(records), failed @@ -547,6 +561,7 @@ def predict( paths.structures_dir, paths.records_dir, num_workers=num_workers, + seed=seed, ) manifest.dump(paths.manifest_path) if failed_preprocessing: diff --git a/nesso/model/modules/utils.py b/nesso/model/modules/utils.py index 5a9db9a..fd02e09 100644 --- a/nesso/model/modules/utils.py +++ b/nesso/model/modules/utils.py @@ -31,8 +31,10 @@ def forward( return F.silu(gates) * x -def randomly_rotate(coords, return_second_coords=False, second_coords=None): - R = random_rotations(len(coords), coords.dtype, coords.device) +def randomly_rotate( + coords, return_second_coords=False, second_coords=None, generator=None +): + R = random_rotations(len(coords), coords.dtype, coords.device, generator=generator) if return_second_coords: return torch.einsum("bmd,bds->bms", coords, R), torch.einsum( @@ -50,6 +52,7 @@ def center_random_augmentation( centering=True, return_second_coords=False, second_coords=None, + generator=None, ): """Algorithm 19""" if centering: @@ -64,9 +67,15 @@ def center_random_augmentation( if augmentation: atom_coords, second_coords = randomly_rotate( - atom_coords, return_second_coords=True, second_coords=second_coords + atom_coords, + return_second_coords=True, + second_coords=second_coords, + generator=generator, ) - random_trans = torch.randn_like(atom_coords[:, 0:1, :]) * s_trans + # randn_like() takes no generator; empty_like().normal_() is equivalent. + trans_shape = atom_coords[:, 0:1, :] + random_trans = torch.empty_like(trans_shape).normal_(generator=generator) + random_trans = random_trans * s_trans atom_coords = atom_coords + random_trans if second_coords is not None: @@ -132,7 +141,10 @@ def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor: def random_quaternions( - n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None + n: int, + dtype: Optional[torch.dtype] = None, + device: Optional[Device] = None, + generator: Optional[torch.Generator] = None, ) -> torch.Tensor: """ Generate random quaternions representing rotations, @@ -149,14 +161,17 @@ def random_quaternions( """ if isinstance(device, str): device = torch.device(device) - o = torch.randn((n, 4), dtype=dtype, device=device) + o = torch.randn((n, 4), dtype=dtype, device=device, generator=generator) s = (o * o).sum(1) o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None] return o def random_rotations( - n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None + n: int, + dtype: Optional[torch.dtype] = None, + device: Optional[Device] = None, + generator: Optional[torch.Generator] = None, ) -> torch.Tensor: """ Generate random rotations as 3x3 rotation matrices. @@ -170,5 +185,5 @@ def random_rotations( Returns: Rotation matrices as tensor of shape (n, 3, 3). """ - quaternions = random_quaternions(n, dtype=dtype, device=device) + quaternions = random_quaternions(n, dtype=dtype, device=device, generator=generator) return quaternion_to_matrix(quaternions) diff --git a/tests/test_determinism.py b/tests/test_determinism.py new file mode 100644 index 0000000..90e4cda --- /dev/null +++ b/tests/test_determinism.py @@ -0,0 +1,236 @@ +"""Predictions must be reproducible for a fixed seed. + +Two inputs to the model were previously drawn from unseeded RNGs, so the same +YAML gave different atomistic features on every run even though +``docs/prediction.md`` documents ``--seed`` as giving "reproducible predictions": + +* the ligand 3D conformer, embedded by RDKit ETKDG with no ``randomSeed`` (its + default is ``-1``, i.e. random); +* the per-residue roto-translation of ``ref_pos``, drawn from the *global* torch + RNG, which ``seed_everything(..., workers=True)`` seeds per DataLoader worker, + making the result depend on ``--num_workers``. + +A third source sat between them: ``preprocess_yamls`` collected results via +``as_completed``, so the manifest order followed worker scheduling rather than +input order, and the per-record RNG was seeded from that position. Seeding the +augmentation was therefore not enough on its own. + +These tests pin all three. They are ligand-only so they need no CCD asset and run +on plain CI. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch +from rdkit import Chem +from torch.utils.data import DataLoader + +from nesso.data.featurizer import NessoFeaturizer +from nesso.data.inference import ( + InferenceDataset, + inference_collate, + record_rng_seed, +) +from nesso.data.types import Manifest +from nesso.data.yaml_input import ( + DEFAULT_CONFORMER_SEED, + conformer_seed, + get_conformer, + parse_yaml, +) + +# Flexible, drug-like: many rotatable bonds, so ETKDG genuinely varies. +_SMILES = "Cc1ccc(NC(=O)c2ccc(CN3CCN(C)CC3)cc2)cc1Nc1nccc(-c2cccnc2)n1" +_YAML = f'version: 1\nsequences:\n - ligand:\n id: B\n smiles: "{_SMILES}"\n' + + +def _parse_coords(tmp_path: Path, name: str, seed: int | None) -> np.ndarray: + work = tmp_path / name + mol_dir = work / "rdkit_conformers" + mol_dir.mkdir(parents=True, exist_ok=True) + yaml_path = work / "lig.yaml" + yaml_path.write_text(_YAML) + struct, _, _, _ = parse_yaml( + yaml_path, mol_dir, ccd_dict=None, record_id="lig", base_seed=seed + ) + return struct.coords["coords"].copy() + + +def _internal_distances(x: np.ndarray) -> np.ndarray: + """Rotation- and translation-invariant description of a conformer.""" + return np.linalg.norm(x[:, None, :] - x[None, :, :], axis=-1) + + +def test_same_seed_gives_identical_conformer(tmp_path: Path) -> None: + a = _parse_coords(tmp_path, "a", DEFAULT_CONFORMER_SEED) + b = _parse_coords(tmp_path, "b", DEFAULT_CONFORMER_SEED) + assert np.array_equal(a, b) + # Invariant under pose, so this also rules out "same shape, different frame". + assert np.array_equal(_internal_distances(a), _internal_distances(b)) + + +def test_different_seed_gives_different_conformer(tmp_path: Path) -> None: + """Guards against the seed being accepted but ignored.""" + a = _parse_coords(tmp_path, "a", DEFAULT_CONFORMER_SEED) + c = _parse_coords(tmp_path, "c", DEFAULT_CONFORMER_SEED + 1) + assert not np.array_equal(a, c) + + +def test_conformer_seed_depends_on_molecule_not_order() -> None: + """`preprocess_yamls` embeds in a ProcessPoolExecutor, so the seed must not + depend on processing order.""" + mol_a = Chem.AddHs(Chem.MolFromSmiles(_SMILES)) + mol_b = Chem.AddHs(Chem.MolFromSmiles(_SMILES)) + other = Chem.AddHs(Chem.MolFromSmiles("CC(=O)Oc1ccccc1C(=O)O")) + + assert conformer_seed(mol_a, 42) == conformer_seed(mol_b, 42) + assert conformer_seed(mol_a, 42) != conformer_seed(other, 42) + assert conformer_seed(mol_a, 42) != conformer_seed(mol_a, 43) + # RDKit takes a signed 32-bit seed. + assert 0 <= conformer_seed(mol_a, 42) < 2**31 - 1 + + +def test_conformer_seed_can_be_opted_out() -> None: + """`base_seed=None` keeps RDKit's default (non-deterministic) behaviour.""" + mol = Chem.AddHs(Chem.MolFromSmiles(_SMILES)) + get_conformer(mol, None) + assert mol.GetNumConformers() == 1 + + +def _ref_pos_by_record( + tmp_path: Path, num_workers: int, n_records: int = 3, reverse: bool = False +) -> dict[str, np.ndarray]: + """Featurize several records through a real DataLoader at a given worker count.""" + mol_dir = tmp_path / "rdkit_conformers" + struct_dir = tmp_path / "structures" + esm_dir = tmp_path / "esm" + for directory in (mol_dir, struct_dir, esm_dir): + directory.mkdir(parents=True, exist_ok=True) + + records = [] + for i in range(n_records): + yaml_path = tmp_path / f"lig{i}.yaml" + yaml_path.write_text(_YAML) + struct, record, _, _ = parse_yaml( + yaml_path, mol_dir, ccd_dict=None, record_id=f"lig{i}" + ) + struct.dump(struct_dir / f"{record.id}.npz") + records.append(record) + + if reverse: + records = list(reversed(records)) + + dataset = InferenceDataset( + manifest=Manifest(records), + 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, + ) + loader = DataLoader( + dataset, + batch_size=1, + shuffle=False, + num_workers=num_workers, + collate_fn=inference_collate, + ) + return {b["record"][0].id: b["ref_pos"][0].numpy().copy() for b in loader} + + +def test_ref_pos_independent_of_num_workers(tmp_path: Path) -> None: + """ref_pos augmentation must not read the global (per-worker) torch RNG.""" + torch.manual_seed(0) + single = _ref_pos_by_record(tmp_path / "w0", 0) + torch.manual_seed(0) + multi = _ref_pos_by_record(tmp_path / "w2", 2) + + assert set(single) == set(multi) + for key in single: + assert np.array_equal(single[key], multi[key]), key + + +def test_ref_pos_unaffected_by_global_rng_state(tmp_path: Path) -> None: + """Consuming the global RNG beforehand must not change featurization.""" + torch.manual_seed(0) + baseline = _ref_pos_by_record(tmp_path / "base", 0) + torch.manual_seed(0) + _ = torch.randn(17) # perturb global RNG stream position + shifted = _ref_pos_by_record(tmp_path / "shift", 0) + + for key in baseline: + assert np.array_equal(baseline[key], shifted[key]), key + + +_OTHER_SMILES = [ + "CC(=O)Oc1ccccc1C(=O)O", + "CN1C=NC2=C1C(=O)N(C)C(=O)N2C", + "CC(C)Cc1ccc(cc1)C(C)C(=O)O", +] + + +def _screen(tmp_path: Path, name: str, smiles: list[str]) -> dict[str, np.ndarray]: + """Parse a batch of ligands, returning coordinates keyed by SMILES.""" + work = tmp_path / name + mol_dir = work / "rdkit_conformers" + mol_dir.mkdir(parents=True, exist_ok=True) + out: dict[str, np.ndarray] = {} + for i, smi in enumerate(smiles): + yaml_path = work / f"c{i}.yaml" + yaml_path.write_text( + f'version: 1\nsequences:\n - ligand:\n id: B\n smiles: "{smi}"\n' + ) + struct, _, _, _ = parse_yaml( + yaml_path, mol_dir, ccd_dict=None, record_id=f"c{i}" + ) + out[smi] = struct.coords["coords"].copy() + return out + + +def test_conformer_independent_of_batch_composition_and_order(tmp_path: Path) -> None: + """A ligand's geometry must depend only on that ligand. + + Previously RDKit's unseeded RNG advanced with every embedding, so a compound's + conformer depended on what else was in the batch and on the order it was + processed in. That silently changed already-computed results whenever a + screening library was extended or its inputs reordered. + """ + alone = _screen(tmp_path, "alone", [_SMILES]) + with_others = _screen(tmp_path, "with_others", [*_OTHER_SMILES, _SMILES]) + reordered = _screen(tmp_path, "reordered", [_SMILES, *reversed(_OTHER_SMILES)]) + + assert np.array_equal(alone[_SMILES], with_others[_SMILES]) + assert np.array_equal(alone[_SMILES], reordered[_SMILES]) + + +def test_record_rng_seed_depends_on_id_not_position() -> None: + """The per-record RNG must be keyed on identity, not on the dataset index. + + ``as_completed`` in ``preprocess_yamls`` made the manifest order follow worker + scheduling, so seeding from the index made ``ref_pos`` depend on + ``--num_workers``. Adding, removing or renaming inputs shifts it too. + """ + assert record_rng_seed("lig") == record_rng_seed("lig") + assert record_rng_seed("lig") != record_rng_seed("other") + # numpy requires a uint32 seed. + assert 0 <= record_rng_seed("lig") < 2**32 + + +def test_ref_pos_independent_of_record_position(tmp_path: Path) -> None: + """The same record must featurize identically wherever it sits in the batch. + + Seeding the per-record RNG from the dataset index made a record's features + depend on its position, so reordering the manifest (which `as_completed` did + on its own at ``--num_workers > 1``) changed ``ref_pos``. + """ + forward = _ref_pos_by_record(tmp_path / "fwd", 0, n_records=3) + # Identical records, reversed order, so every index changes. + backward = _ref_pos_by_record(tmp_path / "rev", 0, n_records=3, reverse=True) + + assert set(forward) == set(backward) + for key in forward: + assert np.array_equal(forward[key], backward[key]), key