Completely disable PSF file based parsing

Author: yingyue2030699

ionerdss/model/pdb/structure_validation.py

@@ -648,36 +648,14 @@ def _extract_observed_com_coordinates(
     restart_file: Union[str, Path],
     target_counts: Mapping[str, int],
 ) -> Dict[str, Tuple[float, float, float]]:
-    """Extract COM coordinates from one connected final-frame assembly."""
-    com_records = _parse_psf_com_records(system_psf_file)
-    xyz_coords = None
-    if final_coords_file is not None and Path(final_coords_file).exists():
-        xyz_coords = _parse_xyz_coordinates(final_coords_file)
-    adjacency, restart_coords, restart_mol_names = _parse_restart_snapshot(restart_file)
+    """Extract COM coordinates from one connected final-frame assembly using restart-native typing only."""
+    del system_psf_file, final_coords_file
 
-    mol_id_to_coord: Dict[int, Tuple[float, float, float]] = {}
-    psf_mol_id_to_name: Dict[int, str] = {}
-    for atom_index, mol_id, mol_name in com_records:
-        psf_mol_id_to_name[mol_id] = mol_name
-        if mol_id in restart_coords:
-            mol_id_to_coord[mol_id] = restart_coords[mol_id]
-        elif xyz_coords is not None:
-            mol_id_to_coord[mol_id] = tuple(xyz_coords[atom_index].tolist())
-        else:
-            raise ValueError(f"Missing coordinates for molecule id {mol_id} in restart snapshot {restart_file}")
+    adjacency, restart_coords, restart_mol_names = _parse_restart_snapshot(restart_file)
+    for mol_id in restart_coords:
         adjacency.setdefault(mol_id, set())
 
-    name_maps_to_try: list[Dict[int, str]] = [psf_mol_id_to_name]
-    if restart_mol_names and restart_mol_names != psf_mol_id_to_name:
-        name_maps_to_try.append(restart_mol_names)
-
-    matching_components: list[list[int]] = []
-    selected_mol_id_to_name = psf_mol_id_to_name
-    for candidate_name_map in name_maps_to_try:
-        matching_components = _find_matching_components(adjacency, candidate_name_map, target_counts)
-        if matching_components:
-            selected_mol_id_to_name = candidate_name_map
-            break
+    matching_components = _find_matching_components(adjacency, restart_mol_names, target_counts)
 
     if not matching_components:
         raise ValueError(
@@ -696,7 +674,7 @@ def _extract_observed_com_coordinates(
     observed = {}
     type_counts: Dict[str, int] = {}
     for mol_id in selected_component:
-        mol_name = selected_mol_id_to_name[mol_id]
+        mol_name = restart_mol_names[mol_id]
         copy_idx = type_counts.get(mol_name, 0)
         type_counts[mol_name] = copy_idx + 1
 
@@ -705,7 +683,7 @@ def _extract_observed_com_coordinates(
         else:
             key = mol_name
 
-        observed[key] = mol_id_to_coord[mol_id]
+        observed[key] = restart_coords[mol_id]
 
     missing = sorted(set(target_counts) - set(type_counts))
     if missing:

ionerdss/tests/test_structure_validation.py

@@ -293,24 +293,24 @@ def test_observed_structure_extraction_uses_one_connected_component():
             encoding="utf-8",
         )
 
-        def _restart_block(mol_id: int, partners: list[int], coord) -> list[str]:
+        def _restart_block(mol_id: int, iface_count: int, partners: list[int], coord) -> list[str]:
             bound_iface_line = " ".join([str(len(partners))] + [str(i) for i in range(len(partners))])
             return [
                 f"{mol_id} 0 0 0 0",
                 "1.0 0 0 0 0 0",
                 f"{coord[0]} {coord[1]} {coord[2]}",
-                "4 0 1 2 3",
+                " ".join([str(iface_count)] + [str(i) for i in range(iface_count)]),
                 bound_iface_line,
                 " ".join([str(len(partners))] + [str(pid) for pid in partners]),
-                "4",
-                "0 0 0 0 \0 0",
-                "0.0 0.0 0.0",
-                "1 1 0 0 \0 0",
-                "0.0 0.0 0.0",
-                "2 2 0 0 \0 0",
-                "0.0 0.0 0.0",
-                "3 3 0 0 \0 0",
-                "0.0 0.0 0.0",
+                str(iface_count),
+                *[
+                    token
+                    for iface_idx in range(iface_count)
+                    for token in (
+                        f"{iface_idx} {iface_idx} 0 0 \0 0",
+                        "0.0 0.0 0.0",
+                    )
+                ],
                 "0",
                 "0",
                 "0",
@@ -320,15 +320,19 @@ def _restart_block(mol_id: int, partners: list[int], coord) -> list[str]:
             ]
 
         restart_lines = [
+            "#MolTemplates",
+            "0 A 4",
+            "1 H 3",
+            "2 L 2",
             "#All Molecules and coordinates",
             "6 6",
         ]
-        restart_lines.extend(_restart_block(0, [1, 2], (0.0, 0.0, 0.0)))
-        restart_lines.extend(_restart_block(1, [0, 2], (1.0, 0.0, 0.0)))
-        restart_lines.extend(_restart_block(2, [0, 1], (0.0, 1.0, 0.0)))
-        restart_lines.extend(_restart_block(3, [4, 5], (10.0, 10.0, 0.0)))
-        restart_lines.extend(_restart_block(4, [3, 5], (10.0, 0.0, 0.0)))
-        restart_lines.extend(_restart_block(5, [3, 4], (10.0, 11.0, 0.0)))
+        restart_lines.extend(_restart_block(0, 4, [1, 2], (0.0, 0.0, 0.0)))
+        restart_lines.extend(_restart_block(1, 3, [0, 2], (1.0, 0.0, 0.0)))
+        restart_lines.extend(_restart_block(2, 2, [0, 1], (0.0, 1.0, 0.0)))
+        restart_lines.extend(_restart_block(3, 4, [4, 5], (10.0, 10.0, 0.0)))
+        restart_lines.extend(_restart_block(4, 3, [3, 5], (10.0, 0.0, 0.0)))
+        restart_lines.extend(_restart_block(5, 2, [3, 4], (10.0, 11.0, 0.0)))
         restart_path.write_text("\n".join(restart_lines), encoding="utf-8")
 
         observed = _extract_observed_com_coordinates(
@@ -385,24 +389,24 @@ def test_restart_snapshot_search_falls_back_to_restart_directory():
             encoding="utf-8",
         )
 
-        def _restart_block(mol_id: int, partners: list[int], coord) -> list[str]:
+        def _restart_block(mol_id: int, iface_count: int, partners: list[int], coord) -> list[str]:
             bound_iface_line = " ".join([str(len(partners))] + [str(i) for i in range(len(partners))])
             return [
                 f"{mol_id} 0 0 0 0",
                 "1.0 0 0 0 0 0",
                 f"{coord[0]} {coord[1]} {coord[2]}",
-                "4 0 1 2 3",
+                " ".join([str(iface_count)] + [str(i) for i in range(iface_count)]),
                 bound_iface_line,
                 " ".join([str(len(partners))] + [str(pid) for pid in partners]),
-                "4",
-                "0 0 0 0 \0 0",
-                "0.0 0.0 0.0",
-                "1 1 0 0 \0 0",
-                "0.0 0.0 0.0",
-                "2 2 0 0 \0 0",
-                "0.0 0.0 0.0",
-                "3 3 0 0 \0 0",
-                "0.0 0.0 0.0",
+                str(iface_count),
+                *[
+                    token
+                    for iface_idx in range(iface_count)
+                    for token in (
+                        f"{iface_idx} {iface_idx} 0 0 \0 0",
+                        "0.0 0.0 0.0",
+                    )
+                ],
                 "0",
                 "0",
                 "0",
@@ -412,21 +416,29 @@ def _restart_block(mol_id: int, partners: list[int], coord) -> list[str]:
             ]
 
         primary_lines = [
+            "#MolTemplates",
+            "0 A 4",
+            "1 H 3",
+            "2 L 2",
             "#All Molecules and coordinates",
             "3 3",
         ]
-        primary_lines.extend(_restart_block(0, [], (10.0, 0.0, 0.0)))
-        primary_lines.extend(_restart_block(1, [], (20.0, 0.0, 0.0)))
-        primary_lines.extend(_restart_block(2, [], (30.0, 0.0, 0.0)))
+        primary_lines.extend(_restart_block(0, 4, [], (10.0, 0.0, 0.0)))
+        primary_lines.extend(_restart_block(1, 3, [], (20.0, 0.0, 0.0)))
+        primary_lines.extend(_restart_block(2, 2, [], (30.0, 0.0, 0.0)))
         primary_restart_path.write_text("\n".join(primary_lines), encoding="utf-8")
 
         older_lines = [
+            "#MolTemplates",
+            "0 A 4",
+            "1 H 3",
+            "2 L 2",
             "#All Molecules and coordinates",
             "3 3",
         ]
-        older_lines.extend(_restart_block(0, [1, 2], (0.0, 0.0, 0.0)))
-        older_lines.extend(_restart_block(1, [0, 2], (1.0, 0.0, 0.0)))
-        older_lines.extend(_restart_block(2, [0, 1], (0.0, 1.0, 0.0)))
+        older_lines.extend(_restart_block(0, 4, [1, 2], (0.0, 0.0, 0.0)))
+        older_lines.extend(_restart_block(1, 3, [0, 2], (1.0, 0.0, 0.0)))
+        older_lines.extend(_restart_block(2, 2, [0, 1], (0.0, 1.0, 0.0)))
         older_restart_path.write_text("\n".join(older_lines), encoding="utf-8")
 
         observed, used_restart = _find_observed_com_coordinates_in_restart_snapshots(