fix(mpnn): honor parser overrides for AtomArray inputs

models/mpnn/src/mpnn/utils/inference.py

@@ -58,6 +58,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,
@@ -240,6 +241,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(
@@ -629,6 +642,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,
@@ -708,6 +722,10 @@ 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()
@@ -721,7 +739,12 @@ def from_atom_array_and_dict(
             # AtomWorks 3 requires charges before template/bond preparation.
             # Derive absent charges from chemistry, preserving supplied charges.
             add_annotations_from_ccd(atom_array, annotations=["charge"])
-        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
@@ -976,7 +999,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")
@@ -994,6 +1017,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:
         """
@@ -1560,10 +1596,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 with the standard processing options and user overrides.
         data = parse(

models/mpnn/tests/test_inference_utils.py

@@ -16,8 +16,11 @@
 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 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,
@@ -272,6 +275,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 +331,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 +497,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({"GOL": "OCC(O)CO", "HOH": "O"})
+
+    result = MPNNInferenceInput.from_atom_array_and_dict(
+        atom_array=atom_array, input_dict={}
+    ).atom_array
+
+    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.
+    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
 ###############################################################################