When psf based method does not work fall back to restart file only parsing

Author: yingyue2030699

ionerdss/model/pdb/structure_validation.py

@@ -435,11 +435,43 @@ def _parse_xyz_coordinates(xyz_file: Union[str, Path]) -> np.ndarray:
     return np.asarray(coords, dtype=float)
 
 
+def _parse_restart_template_names_by_iface_count(lines: list[str]) -> Dict[int, str]:
+    """Infer molecule names from the MolTemplates section keyed by unique interface counts."""
+    molecule_section_idx = None
+    for idx, line in enumerate(lines):
+        stripped = line.strip()
+        if stripped == "#All Molecules and coordinates":
+            molecule_section_idx = idx
+            break
+
+    if molecule_section_idx is None:
+        return {}
+
+    names_by_iface_count: Dict[int, set[str]] = defaultdict(set)
+    for line in lines[:molecule_section_idx]:
+        parts = line.split()
+        if (
+            len(parts) != 3
+            or not parts[0].isdigit()
+            or not parts[2].isdigit()
+            or not parts[1].isalpha()
+        ):
+            continue
+        names_by_iface_count[int(parts[2])].add(parts[1])
+
+    return {
+        iface_count: next(iter(names))
+        for iface_count, names in names_by_iface_count.items()
+        if len(names) == 1
+    }
+
+
 def _parse_restart_snapshot(
     restart_file: Union[str, Path],
-) -> Tuple[Dict[int, set[int]], Dict[int, Tuple[float, float, float]]]:
-    """Parse molecule connectivity and COM coordinates from a NERDSS restart file."""
+) -> Tuple[Dict[int, set[int]], Dict[int, Tuple[float, float, float]], Dict[int, str]]:
+    """Parse molecule connectivity, COM coordinates, and inferred molecule names from a NERDSS restart file."""
     lines = Path(restart_file).read_text(encoding="utf-8", errors="replace").splitlines()
+    template_names_by_iface_count = _parse_restart_template_names_by_iface_count(lines)
 
     start_idx = None
     for idx, line in enumerate(lines):
@@ -456,6 +488,7 @@ def _parse_restart_snapshot(
 
     adjacency: Dict[int, set[int]] = defaultdict(set)
     restart_coords: Dict[int, Tuple[float, float, float]] = {}
+    restart_mol_names: Dict[int, str] = {}
     idx = start_idx + 2
     for _ in range(molecule_count):
         header = lines[idx].split()
@@ -502,6 +535,7 @@ def _parse_restart_snapshot(
         if not iface_count_line:
             raise ValueError(f"Malformed interface count in restart file near line {idx + 1}")
         iface_count = int(iface_count_line[0])
+        restart_mol_names[mol_id] = template_names_by_iface_count.get(iface_count, str(iface_count))
         idx += 1
 
         if bound_list_size != partner_count:
@@ -542,15 +576,45 @@ def _parse_restart_snapshot(
                 raise ValueError(f"Truncated {list_name} list in restart file near line {idx + 1}")
             idx += 1
 
-    return adjacency, restart_coords
+    return adjacency, restart_coords, restart_mol_names
 
 
 def _parse_restart_molecule_partners(restart_file: Union[str, Path]) -> Dict[int, set[int]]:
     """Parse final-frame molecule connectivity from a NERDSS restart file."""
-    adjacency, _ = _parse_restart_snapshot(restart_file)
+    adjacency, _, _ = _parse_restart_snapshot(restart_file)
     return adjacency
 
 
+def _find_matching_components(
+    adjacency: Mapping[int, set[int]],
+    mol_id_to_name: Mapping[int, str],
+    target_counts: Mapping[str, int],
+) -> list[list[int]]:
+    """Return connected components whose composition matches the target counts."""
+    matching_components: list[list[int]] = []
+    visited: set[int] = set()
+    for mol_id in sorted(mol_id_to_name):
+        if mol_id in visited:
+            continue
+
+        component: list[int] = []
+        stack = [mol_id]
+        visited.add(mol_id)
+        while stack:
+            current = stack.pop()
+            component.append(current)
+            for neighbor in sorted(adjacency.get(current, ())):
+                if neighbor in mol_id_to_name and neighbor not in visited:
+                    visited.add(neighbor)
+                    stack.append(neighbor)
+
+        component_counts = Counter(mol_id_to_name[node_id] for node_id in component)
+        if dict(component_counts) == dict(target_counts):
+            matching_components.append(sorted(component))
+
+    return matching_components
+
+
 def _restart_snapshot_sort_key(restart_path: Path) -> Tuple[Tuple[int, ...], str]:
     """Sort restart snapshot files from earliest to latest by embedded numeric suffixes."""
     numeric_parts = tuple(int(part) for part in re.findall(r"\d+", restart_path.stem))
@@ -589,12 +653,12 @@ def _extract_observed_com_coordinates(
     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 = _parse_restart_snapshot(restart_file)
+    adjacency, restart_coords, restart_mol_names = _parse_restart_snapshot(restart_file)
 
-    mol_id_to_name: Dict[int, str] = {}
     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:
-        mol_id_to_name[mol_id] = mol_name
+        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:
@@ -603,26 +667,17 @@ def _extract_observed_com_coordinates(
             raise ValueError(f"Missing coordinates for molecule id {mol_id} in restart snapshot {restart_file}")
         adjacency.setdefault(mol_id, set())
 
-    matching_components: list[list[int]] = []
-    visited: set[int] = set()
-    for mol_id in sorted(mol_id_to_name):
-        if mol_id in visited:
-            continue
-
-        component: list[int] = []
-        stack = [mol_id]
-        visited.add(mol_id)
-        while stack:
-            current = stack.pop()
-            component.append(current)
-            for neighbor in sorted(adjacency.get(current, ())):
-                if neighbor in mol_id_to_name and neighbor not in visited:
-                    visited.add(neighbor)
-                    stack.append(neighbor)
+    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)
 
-        component_counts = Counter(mol_id_to_name[node_id] for node_id in component)
-        if dict(component_counts) == dict(target_counts):
-            matching_components.append(sorted(component))
+    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
 
     if not matching_components:
         raise ValueError(
@@ -641,7 +696,7 @@ def _extract_observed_com_coordinates(
     observed = {}
     type_counts: Dict[str, int] = {}
     for mol_id in selected_component:
-        mol_name = mol_id_to_name[mol_id]
+        mol_name = selected_mol_id_to_name[mol_id]
         copy_idx = type_counts.get(mol_name, 0)
         type_counts[mol_name] = copy_idx + 1
 

ionerdss/tests/test_structure_validation.py

@@ -442,3 +442,72 @@ def _restart_block(mol_id: int, partners: list[int], coord) -> list[str]:
         "H": (1.0, 0.0, 0.0),
         "L": (0.0, 1.0, 0.0),
     }
+
+
+def test_observed_structure_extraction_falls_back_to_restart_native_molecule_types():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        psf_path = tmp_path / "system.psf"
+        restart_path = tmp_path / "restart.dat"
+
+        # Deliberately mismap the COM entries so PSF-based composition matching fails.
+        psf_path.write_text(
+            "\n".join(
+                [
+                    "PSF",
+                    "",
+                    "    3 !NATOM",
+                    "       1    A    0  COM    O    0          0    1.0         0",
+                    "       2    A    1  COM    O    0          0    1.0         0",
+                    "       3    A    2  COM    O    0          0    1.0         0",
+                ]
+            ),
+            encoding="utf-8",
+        )
+
+        def _restart_block(mol_id: int, coord, iface_count: int, partners: list[int]) -> list[str]:
+            lines = [
+                f"{mol_id} 0 0 0 0",
+                "1.0 0 0 0 0 0",
+                f"{coord[0]} {coord[1]} {coord[2]}",
+                " ".join([str(iface_count)] + [str(i) for i in range(iface_count)]),
+                " ".join([str(len(partners))] + [str(i) for i in range(len(partners))]),
+                " ".join([str(len(partners))] + [str(pid) for pid in partners]),
+                str(iface_count),
+            ]
+
+            bound_iface_indexes = set(range(max(iface_count - len(partners), 0), iface_count))
+            for iface_idx in range(iface_count):
+                is_bound = 1 if iface_idx in bound_iface_indexes else 0
+                lines.append(f"{iface_idx} {iface_idx} 0 0 0 {is_bound}")
+                lines.append("0.0 0.0 0.0")
+                if is_bound:
+                    lines.append("0 0 0")
+
+            lines.extend(["0", "0", "0", "0", "0", "0"])
+            return lines
+
+        restart_lines = [
+            "#MolTemplates",
+            "0 A 4",
+            "1 B 2",
+            "#All Molecules and coordinates",
+            "3 3",
+        ]
+        restart_lines.extend(_restart_block(0, (0.0, 0.0, 0.0), 4, [2]))
+        restart_lines.extend(_restart_block(1, (1.0, 0.0, 0.0), 4, [2]))
+        restart_lines.extend(_restart_block(2, (0.0, 1.0, 0.0), 2, [0, 1]))
+        restart_path.write_text("\n".join(restart_lines), encoding="utf-8")
+
+        observed = _extract_observed_com_coordinates(
+            system_psf_file=psf_path,
+            final_coords_file=None,
+            restart_file=restart_path,
+            target_counts={"A": 2, "B": 1},
+        )
+
+    assert observed == {
+        "A_0": (0.0, 0.0, 0.0),
+        "A_1": (1.0, 0.0, 0.0),
+        "B": (0.0, 1.0, 0.0),
+    }