fix: correct bond reading mechanism from input coordinates

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"] = {}

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)

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:

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")