From 1e797a69a2d3e7f59813df5bf1baa225a6e8e7b9 Mon Sep 17 00:00:00 2001 From: daylight-00 Date: Fri, 2 Oct 2026 14:21:09 +0900 Subject: [PATCH 1/2] fix(mpnn): honor parser overrides for AtomArray inputs `MPNNInferenceInput.from_atom_array_and_dict()` parsed an AtomArray with the parser defaults and ignored `remove_ccds` / `remove_waters`, so a ligand whose label is the CCD code of another molecule (phenytoin labelled `PHT`) was rebuilt from that template. - Forward the parser fields the caller set to `parse_atom_array()` through a helper that `build_atom_array()` now shares; unset fields keep the parser defaults. - Add `add_missing_atoms` (JSON field and `--add_missing_atoms`, unset by default) so the atoms of an input can be kept as given. - Test that `add_missing_atoms=False` keeps all 19 phenytoin atoms, that the AtomArray defaults are unchanged, that both routes forward the overrides, and that an input without an assembly still needs `structure_path`. --- models/mpnn/src/mpnn/utils/inference.py | 45 +++++- models/mpnn/tests/test_inference_utils.py | 170 ++++++++++++++++++++++ 2 files changed, 209 insertions(+), 6 deletions(-) diff --git a/models/mpnn/src/mpnn/utils/inference.py b/models/mpnn/src/mpnn/utils/inference.py index 19092aa5f..fb6acebb0 100644 --- a/models/mpnn/src/mpnn/utils/inference.py +++ b/models/mpnn/src/mpnn/utils/inference.py @@ -57,6 +57,7 @@ # Parser Overrides "remove_ccds": [], "remove_waters": None, + "add_missing_atoms": None, # Pipeline Setup Overrides "occupancy_threshold_sidechain": 0.0, "occupancy_threshold_backbone": 0.0, @@ -239,6 +240,18 @@ def build_arg_parser() -> argparse.ArgumentParser: ), default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["remove_waters"], ) + parser.add_argument( + "--add_missing_atoms", + type=lambda v: none_or_type(v, str2bool), + choices=[True, False, None], + help=( + "If set, override the parser default for adding missing atoms from " + "the CCD (overrides STANDARD_PARSER_ARGS). False keeps only the atoms " + "present in the input. 'None' has special behavior: use the parser " + "default behavior." + ), + default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["add_missing_atoms"], + ) # ---------------- Pipeline Setup Overrides ---------------- # parser.add_argument( @@ -628,6 +641,7 @@ def cli_to_json(args: argparse.Namespace) -> dict[str, Any]: # Parser Overrides "remove_ccds": parse_list_like(args.remove_ccds), "remove_waters": args.remove_waters, + "add_missing_atoms": args.add_missing_atoms, # Pipeline Setup Overrides "occupancy_threshold_sidechain": args.occupancy_threshold_sidechain, "occupancy_threshold_backbone": args.occupancy_threshold_backbone, @@ -707,13 +721,22 @@ def from_atom_array_and_dict( fields from the input dictionary will be ignored. If you would like to override those annotations, please either do so in the atom array or delete the annotations from the atom array before passing it in. + + The parser override fields (remove_ccds, remove_waters, + add_missing_atoms) apply to the atom array as they do to a structure + file; fields left unset keep the parser defaults. """ # Copy input dictionary. input_dict = copy.deepcopy(input_dict) if input_dict is not None else dict() # Copy atom array. atom_array = atom_array.copy() if atom_array is not None else None - parser_output = parse_atom_array(atom_array) if atom_array is not None else {} + # The overrides feed the parser, so validate them first. + MPNNInferenceInput._validate_parser_overrides(input_dict) + overrides = MPNNInferenceInput._parser_overrides(input_dict) + parser_output = ( + parse_atom_array(atom_array, **overrides) if atom_array is not None else {} + ) atom_array = ( parser_output["assemblies"]["1"][0] if len(parser_output.get("assemblies", {})) > 0 @@ -968,7 +991,7 @@ def _validate_sampling_parameters(input_dict: dict[str, Any]) -> None: @staticmethod def _validate_parser_overrides(input_dict: dict[str, Any]) -> None: - """Validate parser override fields: remove_ccds, remove_waters.""" + """Validate override fields: remove_ccds, remove_waters, add_missing_atoms.""" # Check that remove_ccds is a list of strings if provided. remove_ccds = input_dict.get("remove_ccds") @@ -986,6 +1009,19 @@ def _validate_parser_overrides(input_dict: dict[str, Any]) -> None: if remove_waters is not None and not isinstance(remove_waters, bool): raise TypeError("remove_waters must be a bool when provided.") + add_missing_atoms = input_dict.get("add_missing_atoms") + if add_missing_atoms is not None and not isinstance(add_missing_atoms, bool): + raise TypeError("add_missing_atoms must be a bool when provided.") + + @staticmethod + def _parser_overrides(input_dict: dict[str, Any]) -> dict[str, Any]: + """Parser keywords the caller set; unset or None keeps the parser default.""" + return { + key: input_dict[key] + for key in ("remove_ccds", "remove_waters", "add_missing_atoms") + if input_dict.get(key) is not None + } + @staticmethod def _validate_pipeline_override_fields(input_dict: dict[str, Any]) -> None: """ @@ -1552,10 +1588,7 @@ def build_atom_array(input_dict: dict[str, Any]) -> AtomArray: """Build AtomArray from structure_path and parser overrides.""" # Override parser args if specified. parser_args = dict(STANDARD_PARSER_ARGS) - if input_dict["remove_ccds"] is not None: - parser_args["remove_ccds"] = input_dict["remove_ccds"] - if input_dict["remove_waters"] is not None: - parser_args["remove_waters"] = input_dict["remove_waters"] + parser_args.update(MPNNInferenceInput._parser_overrides(input_dict)) # Parse structure file. parser_args is a heterogeneous dict[str, object] # (STANDARD_PARSER_ARGS + overrides); atomworks parse() types each kwarg diff --git a/models/mpnn/tests/test_inference_utils.py b/models/mpnn/tests/test_inference_utils.py index 0c08ef4e2..4f20f2389 100644 --- a/models/mpnn/tests/test_inference_utils.py +++ b/models/mpnn/tests/test_inference_utils.py @@ -16,8 +16,10 @@ from pathlib import Path from typing import Any +import biotite.structure as struc import numpy as np import pytest +from atomworks.io.parser import parse_atom_array from biotite.structure import AtomArray from mpnn.transforms.feature_aggregation.token_encodings import ( MPNN_TOKEN_ENCODING, @@ -272,6 +274,8 @@ def test_cli_to_json_builds_single_input_and_parses_fields() -> None: "A,B", "--remove_waters", "True", + "--add_missing_atoms", + "False", "--occupancy_threshold_sidechain", "0.5", "--occupancy_threshold_backbone", @@ -326,6 +330,7 @@ def test_cli_to_json_builds_single_input_and_parses_fields() -> None: assert inp["number_of_batches"] == 2 assert inp["remove_ccds"] == ["A", "B"] assert inp["remove_waters"] is True + assert inp["add_missing_atoms"] is False assert inp["occupancy_threshold_sidechain"] == 0.5 assert inp["occupancy_threshold_backbone"] == 0.6 assert inp["undesired_res_names"] == ["UNK", "HOH"] @@ -491,6 +496,171 @@ def test_from_atom_array_and_dict_invalid_scalar_settings_raise() -> None: ) +############################################################################### +# MPNNInferenceInput: parser overrides for an AtomArray input +############################################################################### + + +def _atom_array_with_ligands(ligands: dict[str, str]) -> AtomArray: + """CA-only ALA/GLY/SER on chain A, plus one chain per ``{res_name: smiles}``.""" + from rdkit import Chem + from rdkit.Chem import AllChem + + atoms = [ + struc.Atom( + [3.8 * i, 0.0, 0.0], + chain_id="A", + res_id=i + 1, + ins_code="", + res_name=res_name, + atom_name="CA", + element="C", + ) + for i, res_name in enumerate(["ALA", "GLY", "SER"]) + ] + for i, (res_name, smiles) in enumerate(ligands.items()): + mol = Chem.AddHs(Chem.MolFromSmiles(smiles)) + AllChem.EmbedMolecule(mol, randomSeed=7) + mol = Chem.RemoveHs(mol) + conf = mol.GetConformer() + counts: dict[str, int] = {} + for atom in mol.GetAtoms(): + element = atom.GetSymbol() + counts[element] = counts.get(element, 0) + 1 + xyz = conf.GetAtomPosition(atom.GetIdx()) + atoms.append( + struc.Atom( + [xyz.x, xyz.y, xyz.z], + chain_id=chr(ord("B") + i), + res_id=1, + ins_code="", + res_name=res_name, + atom_name=f"{element}{counts[element]}", + element=element, + hetero=True, + ) + ) + atom_array = struc.array(atoms) + atom_array.set_annotation("occupancy", np.ones(len(atom_array))) + return atom_array + + +# PHT is the CCD code of phthalate, not phenytoin. +PHENYTOIN = "O=C1NC(=O)C(N1)(c1ccccc1)c1ccccc1" + + +def test_add_missing_atoms_false_keeps_ligand_atoms_as_given() -> None: + atom_array = _atom_array_with_ligands({"PHT": PHENYTOIN}) + + kept = MPNNInferenceInput.from_atom_array_and_dict( + atom_array=atom_array, input_dict={"add_missing_atoms": False} + ).atom_array + + ligand = kept[kept.res_name == "PHT"] + assert len(ligand) == 19 + assert set(ligand.element) == {"C", "N", "O"} + assert not np.isnan(ligand.coord).any() + + +def test_atom_array_defaults_are_unchanged() -> None: + atom_array = _atom_array_with_ligands( + {"PHT": PHENYTOIN, "GOL": "OCC(O)CO", "HOH": "O"} + ) + + result = MPNNInferenceInput.from_atom_array_and_dict( + atom_array=atom_array, input_dict={} + ).atom_array + + expected = parse_atom_array(atom_array.copy())["assemblies"]["1"][0] + assert list(result.res_name) == list(expected.res_name) + assert list(result.atom_name) == list(expected.atom_name) + # The parser defaults still remove crystallization aids and waters. + assert {"GOL", "HOH"}.isdisjoint(result.res_name) + + +def test_atom_array_route_applies_remove_ccds_and_remove_waters() -> None: + atom_array = _atom_array_with_ligands({"GOL": "OCC(O)CO", "HOH": "O"}) + + kept = MPNNInferenceInput.from_atom_array_and_dict( + atom_array=atom_array, + input_dict={"remove_ccds": [], "remove_waters": False}, + ).atom_array + + assert {"GOL", "HOH"} <= set(kept.res_name) + + +def test_atom_array_route_forwards_only_the_fields_the_caller_set( + monkeypatch, +) -> None: + seen: dict[str, Any] = {} + + def fake_parse_atom_array(atom_array, **kwargs): + seen.clear() + seen.update(kwargs) + return {"assemblies": {"1": [atom_array]}} + + monkeypatch.setattr("mpnn.utils.inference.parse_atom_array", fake_parse_atom_array) + + MPNNInferenceInput.from_atom_array_and_dict( + atom_array=_make_simple_atom_array(), input_dict={} + ) + assert seen == {} + + overrides = { + "remove_ccds": ["GOL"], + "remove_waters": False, + "add_missing_atoms": False, + } + MPNNInferenceInput.from_atom_array_and_dict( + atom_array=_make_simple_atom_array(), input_dict=overrides + ) + assert seen == overrides + + +def test_atom_array_without_assembly_still_needs_structure_path(monkeypatch) -> None: + monkeypatch.setattr( + "mpnn.utils.inference.parse_atom_array", lambda atom_array, **kwargs: {} + ) + + with pytest.raises(ValueError, match="structure_path"): + MPNNInferenceInput.from_atom_array_and_dict( + atom_array=_make_simple_atom_array(), input_dict={} + ) + + +def test_structure_path_route_forwards_parser_overrides(monkeypatch) -> None: + seen: dict[str, Any] = {} + + def fake_parse(filename, **kwargs): + seen.clear() + seen.update(kwargs) + return {"asym_unit": [_make_simple_atom_array()]} + + monkeypatch.setattr("mpnn.utils.inference.parse", fake_parse) + + input_dict: dict[str, Any] = {"structure_path": "structure.cif"} + MPNNInferenceInput.apply_defaults(input_dict) + MPNNInferenceInput.build_atom_array(input_dict) + assert seen["add_missing_atoms"] is True + assert seen["remove_ccds"] == [] + + input_dict.update(add_missing_atoms=False, remove_waters=False) + MPNNInferenceInput.build_atom_array(input_dict) + assert seen["add_missing_atoms"] is False + assert seen["remove_waters"] is False + + +@pytest.mark.parametrize( + "field, value", + [("add_missing_atoms", "no"), ("remove_waters", "yes"), ("remove_ccds", "GOL")], +) +def test_parser_overrides_are_validated(field: str, value: Any) -> None: + with pytest.raises(TypeError, match=field): + MPNNInferenceInput.from_atom_array_and_dict( + atom_array=_make_simple_atom_array(), input_dict={field: value} + ) + + ############################################################################### # MPNNInferenceInput: annotation behaviour ############################################################################### From e12b7f00032060bc5ac749d15af4cfce7160a2cf Mon Sep 17 00:00:00 2001 From: daylight-00 Date: Fri, 16 Oct 2026 14:10:44 +0900 Subject: [PATCH 2/2] test(mpnn): keep the AtomArray defaults check valid on AtomWorks 3 AtomWorks 3 rejects a ligand whose atom names do not match its CCD template, so the defaults check can no longer carry the mislabelled phenytoin. - Check the defaults on glycerol and water only; the phenytoin case stays in the `add_missing_atoms=False` test. - Build the expected array with the CCD charges that `from_atom_array_and_dict()` now adds before parsing. --- models/mpnn/tests/test_inference_utils.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/models/mpnn/tests/test_inference_utils.py b/models/mpnn/tests/test_inference_utils.py index 4f20f2389..e75e7c87a 100644 --- a/models/mpnn/tests/test_inference_utils.py +++ b/models/mpnn/tests/test_inference_utils.py @@ -20,6 +20,7 @@ import numpy as np import pytest from atomworks.io.parser import parse_atom_array +from atomworks.io.utils.ccd import add_annotations_from_ccd from biotite.structure import AtomArray from mpnn.transforms.feature_aggregation.token_encodings import ( MPNN_TOKEN_ENCODING, @@ -563,15 +564,15 @@ def test_add_missing_atoms_false_keeps_ligand_atoms_as_given() -> None: def test_atom_array_defaults_are_unchanged() -> None: - atom_array = _atom_array_with_ligands( - {"PHT": PHENYTOIN, "GOL": "OCC(O)CO", "HOH": "O"} - ) + atom_array = _atom_array_with_ligands({"GOL": "OCC(O)CO", "HOH": "O"}) result = MPNNInferenceInput.from_atom_array_and_dict( atom_array=atom_array, input_dict={} ).atom_array - expected = parse_atom_array(atom_array.copy())["assemblies"]["1"][0] + expected = atom_array.copy() + add_annotations_from_ccd(expected, annotations=["charge"]) + expected = parse_atom_array(expected)["assemblies"]["1"][0] assert list(result.res_name) == list(expected.res_name) assert list(result.atom_name) == list(expected.atom_name) # The parser defaults still remove crystallization aids and waters.