From f4ea4faa5aa3cf86009925a0524ef305fd23913a Mon Sep 17 00:00:00 2001 From: rish-16 Date: Sun, 25 Oct 2026 19:57:10 +0000 Subject: [PATCH] Fix bond reading mechanism from input coordinates --- src/nucleus/inference_input_parsing.py | 11 +++++ src/nucleus/utils/bond_sidecar.py | 57 ++++++++++++++++++++++++++ src/nucleus/utils/datasets.py | 48 ++++------------------ tests/test_bond_sidecar.py | 45 ++++++++++++++++++++ 4 files changed, 121 insertions(+), 40 deletions(-) diff --git a/src/nucleus/inference_input_parsing.py b/src/nucleus/inference_input_parsing.py index a5a8d07..9a3d832 100644 --- a/src/nucleus/inference_input_parsing.py +++ b/src/nucleus/inference_input_parsing.py @@ -23,6 +23,7 @@ fetch_mask_from_idx, ) +from nucleus.utils.bond_sidecar import perceive_ligand_bonds_from_native_coordinates from nucleus.utils.materials import ( annotate_material_properties, cartesian_to_fractional, @@ -1013,6 +1014,16 @@ def design_input_to_pipeline_input_preserving_hydrogens(self, example_id): data = prepare_pipeline_input_preserving_hydrogens(atom_array) data["example_id"] = example_id + # ... Perceive molecule bonds from native coordinates; build() has zeroed them. + # A bond sidecar, when configured, replaces these bonds downstream. + has_disconnected_fragments = perceive_ligand_bonds_from_native_coordinates( + data["atom_array"], + getattr(self, "_nucleus_native_coordinates", None), + example_id=str(getattr(self, "example_id", None) or example_id), + ) + if has_disconnected_fragments is not None: + data["has_disconnected_fragments"] = has_disconnected_fragments + # ... Wrap up with additional features if "extra" not in spec_dict: spec_dict["extra"] = {} diff --git a/src/nucleus/utils/bond_sidecar.py b/src/nucleus/utils/bond_sidecar.py index c5ee388..b9fe655 100644 --- a/src/nucleus/utils/bond_sidecar.py +++ b/src/nucleus/utils/bond_sidecar.py @@ -320,6 +320,63 @@ def perceive_and_attach_bonds(atom_array: Any, *, example_id: str) -> bool: return bool(record["has_disconnected_fragments"]) +def _can_perceive_ligand_bonds(atom_array, native_coordinates) -> bool: + """Return whether native coordinates define a complete non-material ligand.""" + if native_coordinates is None: + return False + + native_coordinates = np.asarray(native_coordinates) + if native_coordinates.shape != atom_array.coord.shape: + return False + if not np.isfinite(native_coordinates).all(): + return False + + categories = set(atom_array.get_annotation_categories()) + if "is_ligand" in categories: + if not np.asarray(atom_array.is_ligand, dtype=bool).all(): + return False + elif not np.asarray(atom_array.res_name == "LIG", dtype=bool).all(): + # The generic inference parser assigns ``is_ligand`` later in the + # transform pipeline. At this boundary, molecule-only inputs are the + # non-polymer LIG residue produced by the contig parser. + return False + if ( + "is_polymer" in categories + and np.asarray(atom_array.is_polymer, dtype=bool).any() + ): + return False + if ( + "is_material" in categories + and np.asarray(atom_array.is_material, dtype=bool).any() + ): + return False + return True + + +def perceive_ligand_bonds_from_native_coordinates( + atom_array: Any, native_coordinates: Any, *, example_id: str +) -> bool | None: + """Attach bonds perceived from native coordinates to a molecule-only input. + + Inference zeroes generated coordinates before the pipeline runs, so bonds are + perceived on a copy carrying the preserved native coordinates. + + Returns: + Whether the perceived graph contains disconnected fragments, or ``None`` + when the input is not a fully resolved non-material ligand. + """ + if not _can_perceive_ligand_bonds(atom_array, native_coordinates): + return None + perception_atom_array = atom_array.copy() + perception_atom_array.coord = np.asarray(native_coordinates).copy() + has_disconnected_fragments = perceive_and_attach_bonds( + perception_atom_array, + example_id=example_id, + ) + atom_array.bonds = perception_atom_array.bonds + return has_disconnected_fragments + + def encode_bond_record(record: dict[str, Any]) -> bytes: """Serialize one sidecar record.""" return pickle.dumps(record, protocol=pickle.HIGHEST_PROTOCOL) diff --git a/src/nucleus/utils/datasets.py b/src/nucleus/utils/datasets.py index e466cc6..1d9b990 100644 --- a/src/nucleus/utils/datasets.py +++ b/src/nucleus/utils/datasets.py @@ -39,7 +39,7 @@ ) from nucleus.utils.bond_sidecar import ( BondSidecarASEAtomsLoader, - perceive_and_attach_bonds, + perceive_ligand_bonds_from_native_coordinates, ) ranked_logger = RankedLogger(__name__, rank_zero_only=True) @@ -71,39 +71,6 @@ def _coerce_ligand_design_specification(spec): return specification_cls.safe_init(**spec) -def _can_perceive_ligand_bonds(atom_array, native_coordinates) -> bool: - """Return whether native coordinates define a complete non-material ligand.""" - if native_coordinates is None: - return False - - native_coordinates = np.asarray(native_coordinates) - if native_coordinates.shape != atom_array.coord.shape: - return False - if not np.isfinite(native_coordinates).all(): - return False - - categories = set(atom_array.get_annotation_categories()) - if "is_ligand" in categories: - if not np.asarray(atom_array.is_ligand, dtype=bool).all(): - return False - elif not np.asarray(atom_array.res_name == "LIG", dtype=bool).all(): - # The generic inference parser assigns ``is_ligand`` later in the - # transform pipeline. At this boundary, molecule-only inputs are the - # non-polymer LIG residue produced by the contig parser. - return False - if ( - "is_polymer" in categories - and np.asarray(atom_array.is_polymer, dtype=bool).any() - ): - return False - if ( - "is_material" in categories - and np.asarray(atom_array.is_material, dtype=bool).any() - ): - return False - return True - - class ZatomContigJsonDataset: """Contig JSON dataset with Nucleus typed generated-contig support. @@ -172,14 +139,15 @@ def __getitem__(self, idx: int) -> Any: data["atom_array"], example_id=str(source_example_id), ) - elif _can_perceive_ligand_bonds(data["atom_array"], native_coordinates): - perception_atom_array = data["atom_array"].copy() - perception_atom_array.coord = np.asarray(native_coordinates).copy() - data["has_disconnected_fragments"] = perceive_and_attach_bonds( - perception_atom_array, + elif "has_disconnected_fragments" not in data: + # The rfd3 backend already perceives bonds in to_pipeline_input. + has_disconnected_fragments = perceive_ligand_bonds_from_native_coordinates( + data["atom_array"], + native_coordinates, example_id=str(source_example_id), ) - data["atom_array"].bonds = perception_atom_array.bonds + if has_disconnected_fragments is not None: + data["has_disconnected_fragments"] = has_disconnected_fragments data = self._dataset.transform(data) if native_coordinates is not None and "ground_truth" in data: diff --git a/tests/test_bond_sidecar.py b/tests/test_bond_sidecar.py index 80cf8d4..621ef65 100644 --- a/tests/test_bond_sidecar.py +++ b/tests/test_bond_sidecar.py @@ -339,6 +339,51 @@ def test_contig_validation_perceives_bonds_without_sidecar() -> None: assert data["feats"]["token_bonds"].sum() == 2 * len(expected_record["bond_edges"]) +@pytest.mark.skipif(openbabel is None, reason="Open Babel is not installed") +def test_contig_inference_perceives_bonds_without_sidecar() -> None: + # `zatom2 design` loads specs through RFD3's ContigJsonDataset rather than + # ZatomContigJsonDataset, so perception must happen in to_pipeline_input. + from nucleus.inference_input_parsing import ( + ZatomDesignInputSpecification, + install_molecule_contig_support, + ) + from rfd3.inference.datasets import ContigJsonDataset + + install_molecule_contig_support() + native_atom_array = _ligand_atom_array() + expected_record = openbabel_bond_record( + native_atom_array, + example_id="on-the-fly-example", + ) + specification = ZatomDesignInputSpecification.safe_init( + input=None, + atom_array_input=native_atom_array, + contig="A1", + select_fixed_atoms=False, + select_unfixed_sequence=False, + ) + dataset = ContigJsonDataset( + data={"on-the-fly-example": specification}, + cif_parser_args={}, + transform=_add_structure_prediction_token_bonds, + name="on-the-fly-bond-inference-test", + subset_to_keys=None, + eval_every_n=1, + ) + + data = dataset[0] + + np.testing.assert_allclose(data["atom_array"].coord, 0.0) + np.testing.assert_array_equal( + data["atom_array"].bonds.as_array()[:, :2], + expected_record["bond_edges"], + ) + assert data["has_disconnected_fragments"] is bool( + expected_record["has_disconnected_fragments"] + ) + assert data["feats"]["token_bonds"].sum() == 2 * len(expected_record["bond_edges"]) + + @pytest.mark.skipif(openbabel is None, reason="Open Babel is not installed") def test_openbabel_bond_record_stores_edges_without_bond_types() -> None: record = openbabel_bond_record(_water_atom_array(), example_id="water")