pxmeter 0.1.2__tar.gz → 0.1.4__tar.gz

{pxmeter-0.1.2/pxmeter.egg-info → pxmeter-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pxmeter
-Version: 0.1.2
+Version: 0.1.4
 Summary: PXMeter is a comprehensive toolkit for evaluating the quality of         structures generated by biomolecular structure prediction models.
 Author: Bytedance Inc.
 Author-email: ai4s-bio@bytedance.com

{pxmeter-0.1.2 → pxmeter-0.1.4}/pxmeter/calc_metric.py

@@ -34,6 +34,7 @@ from pxmeter.configs.run_config import RUN_CONFIG
 from pxmeter.constants import IONS, LIGAND
 from pxmeter.data.ccd import get_ccd_mol_from_chain_atom_array
 from pxmeter.data.struct import Structure
+from pxmeter.metrics.clashes import check_clashes_by_vdw
 from pxmeter.metrics.lddt_metrics import LDDT
 from pxmeter.metrics.rmsd_metrics import RMSDMetrics
 
@@ -554,6 +555,16 @@ class MetricResult:
         meta_info_dict["ref_to_model_chain_mapping"] = chain_map
         meta_info_dict["ref_chain_info"] = cls._get_chain_info(ref_struct)
 
+        # Calculate clashes
+        if metric_config.calc_clashes:
+            clashes = check_clashes_by_vdw(
+                model_struct.atom_array,
+                vdw_scale_factor=metric_config.clashes.vdw_scale_factor,
+            )
+            complex_result_dict["clashes"] = len(
+                {x for a, b in clashes for x in (a, b)}
+            )
+
         # Calculate RMSD (if ligand and pocket specified in ref_features)
         if metric_config.calc_rmsd and interested_lig_label_asym_id:
             rmsd_metrics = RMSDMetrics(

{pxmeter-0.1.2 → pxmeter-0.1.4}/pxmeter/configs/run_config.py

@@ -23,6 +23,7 @@ RUN_CONFIG = ConfigDict(
             "enumerate_all_anchors": True,
         },
         "metric": {
+            "calc_clashes": True,
             "calc_lddt": True,
             "calc_dockq": True,
             "calc_rmsd": True,
@@ -32,6 +33,9 @@ RUN_CONFIG = ConfigDict(
                 "nucleotide_threshold": 30.0,
                 "non_nucleotide_threshold": 15.0,
             },
+            "clashes": {
+                "vdw_scale_factor": 0.5,
+            },
         },
     }
 )

{pxmeter-0.1.2 → pxmeter-0.1.4}/pxmeter/data/struct.py

@@ -432,6 +432,38 @@ class Structure:
             chain_id_to_entity_id[chain_id] = entity_id
         return chain_id_to_entity_id
 
+    def get_ligand_polymer_bonds(self) -> np.ndarray:
+        """
+        Get bonds between the bonded ligand and its parent chain.
+
+        Returns:
+            np.ndarray: bond records between the bonded ligand and its parent chain.
+                        e.g. np.array([[atom1, atom2, bond_order]...])
+        """
+        atom_array = self.atom_array
+        bond_array = atom_array.bonds.as_array()
+
+        polymer_mask = np.isin(
+            atom_array.label_entity_id, list(self.entity_poly_type.keys())
+        )
+
+        lig_mask = ~polymer_mask
+
+        idx_i = bond_array[:, 0]
+        idx_j = bond_array[:, 1]
+
+        lig_polymer_bond_indices = np.where(
+            (lig_mask[idx_i] & polymer_mask[idx_j])
+            | (lig_mask[idx_j] & polymer_mask[idx_i])
+        )[0]
+        if lig_polymer_bond_indices.size == 0:
+            # no ligand-polymer bonds
+            lig_polymer_bonds = np.empty((0, 3)).astype(int)
+        else:
+            # np.array([[atom1, atom2, bond_order], ...])
+            lig_polymer_bonds = bond_array[lig_polymer_bond_indices]
+        return lig_polymer_bonds
+
     def clean_structure(
         self,
         mse_to_met=True,

{pxmeter-0.1.2 → pxmeter-0.1.4}/pxmeter/mapping.py

@@ -924,15 +924,28 @@ class MappingCIF:
         # This step will be change the res_id, res_name, atom_name in self.model_struct.atom_array
         self._align_model_lig_atom_to_ref(self.model_struct, model_to_ref_atom_mapping)
 
+        # Re-order model struct by res_id for each chain
+        order = []
+        for chain_id in np.unique(self.model_struct.uni_chain_id):
+            chain_mask = np.where(self.model_struct.uni_chain_id == chain_id)[0]
+            # Remove unmapped ligands
+            valid_chain_mask = chain_mask[
+                ~(
+                    (
+                        (self.model_struct.atom_array.res_name[chain_mask] == ".")
+                        & (self.model_struct.atom_array.atom_name[chain_mask] == ".")
+                    )
+                    | (
+                        self.model_struct.atom_array.res_id[chain_mask] < 0
+                    )  # -1 for unmapped residues
+                )
+            ]
+            res_ids = self.model_struct.atom_array.res_id[valid_chain_mask]
+            order.extend(valid_chain_mask[np.argsort(res_ids)])
+
         # Remove unmapped ligand of model and reset unique atom id
         self.model_struct = self.model_struct.select_substructure(
-            ~(
-                (
-                    (self.model_struct.atom_array.res_name == ".")
-                    & (self.model_struct.atom_array.atom_name == ".")
-                )
-                | (self.model_struct.atom_array.res_id < 0)  # -1 for unmapped residues
-            ),
+            order,
             reset_uni_id=True,
         )
 

pxmeter-0.1.4/pxmeter/metrics/clashes.py

@@ -0,0 +1,88 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Sequence
+
+import numpy as np
+from biotite.structure import AtomArray
+from biotite.structure.info.radii import vdw_radius_single
+from scipy.spatial import KDTree
+
+
+def check_clashes_by_vdw(
+    atom_array: AtomArray,
+    query_mask: Sequence[bool] = None,
+    vdw_scale_factor: float = 0.5,
+) -> list[tuple[int, int]]:
+    """
+    Check clashes between atoms in the given atom array.
+
+    Args:
+        atom_array (AtomArray): The atom array to check for clashes.
+        query_mask (bool, optional): A boolean mask to select atoms to check for clashes.
+                   If None, all atoms are checked.
+        vdw_scale_factor (float, optional): The scale factor to apply to the Van der Waals radii.
+                         Defaults to 0.5.
+
+    Returns:
+        list[tuple[int, int]]: A list of tuples representing the indices of atoms that are in clash.
+    """
+    if query_mask is None:
+        # query all atoms
+        query_mask = np.ones(len(atom_array), dtype=bool)
+    elif not np.any(query_mask):
+        # no query atoms, return empty list
+        return []
+
+    if query_mask is None:
+        query_mask = np.ones(len(atom_array), dtype=bool)
+
+    query_idx_in_ref = np.where(query_mask)[0]
+
+    vdw_radii = np.array([vdw_radius_single(e) for e in atom_array.element])
+    query_vdw_radii = vdw_radii[query_mask]
+
+    clashes = []
+    query_tree = KDTree(atom_array.coord)
+    for query_idx, nbs_idx in enumerate(
+        query_tree.query_ball_point(atom_array.coord[query_mask], r=3.0)
+    ):
+        query_bonds, _query_bond_types = atom_array.bonds.get_bonds(
+            query_idx_in_ref[query_idx]
+        )
+        query_vdw = query_vdw_radii[query_idx]
+        if query_vdw is None:
+            # undefined vdw for elem, use 1.7 as "C"
+            query_vdw = vdw_radius_single("C")
+
+        for nb_idx in nbs_idx:
+            if query_idx_in_ref[query_idx] == nb_idx:
+                # clash with self
+                continue
+
+            if nb_idx in query_bonds:
+                # clash with bonded atoms
+                continue
+
+            nb_vdw = vdw_radii[nb_idx]
+            if nb_vdw is None:
+                # undefined vdw for elem, use 1.7 as "C"
+                nb_vdw = vdw_radius_single("C")
+
+            dist = np.linalg.norm(
+                atom_array.coord[query_mask][query_idx] - atom_array.coord[nb_idx]
+            )
+            if dist < vdw_scale_factor * (query_vdw + nb_vdw):
+                clashes.append((query_idx_in_ref[query_idx], nb_idx))
+    return clashes

{pxmeter-0.1.2 → pxmeter-0.1.4}/pxmeter/permutation/chain.py

@@ -59,6 +59,71 @@ class ChainPermutation:
         self.ref_chain_id_to_entity_id = self.ref_struct.get_chain_id_to_entity_id()
         self.model_chain_id_to_entity_id = self.model_struct.get_chain_id_to_entity_id()
 
+        self.ref_and_model_mapping_ban_set = self._get_ban_set_by_lig_bonded_position()
+
+    @staticmethod
+    def find_bonded_position_for_lig_chains(
+        struct: Structure,
+    ) -> dict[str, tuple[str, int]]:
+        """
+        Find the bonded entity ID and residue ID for ligand chains.
+
+        Args:
+            struct (Structure): Structure object.
+
+        Returns:
+            dict[str, tuple[str, int]]: Mapping of ligand chain ID to bonded
+                                        entity ID and residue ID.
+        """
+        ligand_polymer_bonds = struct.get_ligand_polymer_bonds()
+
+        chain_id_to_bonded_position = {}
+        for bond in ligand_polymer_bonds:
+            atom1, atom2, _ = bond
+            if struct.atom_array.label_entity_id[atom1] not in struct.entity_poly_type:
+                # atom1 is ligand
+                lig_chain_id = struct.uni_chain_id[atom1]
+                entity_id = struct.atom_array.label_entity_id[atom2]
+                res_id = struct.atom_array.res_id[atom2]
+            else:
+                lig_chain_id = struct.uni_chain_id[atom2]
+                entity_id = struct.atom_array.label_entity_id[atom1]
+                res_id = struct.atom_array.res_id[atom1]
+            chain_id_to_bonded_position[lig_chain_id] = (entity_id, res_id)
+        return chain_id_to_bonded_position
+
+    def _get_ban_set_by_lig_bonded_position(
+        self,
+    ) -> set[tuple[str, str]]:
+        ref_chain_id_to_bond_position = (
+            ChainPermutation.find_bonded_position_for_lig_chains(self.ref_struct)
+        )
+        model_chain_id_to_bond_position = (
+            ChainPermutation.find_bonded_position_for_lig_chains(self.model_struct)
+        )
+
+        ban_set = set()
+        for ref_chain_id in np.unique(self.ref_struct.uni_chain_id):
+            ref_bonded_entity, ref_bonded_res_id = ref_chain_id_to_bond_position.get(
+                ref_chain_id, ["-1", -1]
+            )
+            mapped_model_bonded_entity = self.ref_to_model_entity_id.get(
+                ref_bonded_entity, "-1"
+            )
+
+            for model_chain_id in np.unique(self.model_struct.uni_chain_id):
+                (
+                    model_bonded_entity,
+                    model_bonded_res_id,
+                ) = model_chain_id_to_bond_position.get(model_chain_id, ["-1", -1])
+
+                if mapped_model_bonded_entity != "-1" and model_bonded_entity != "-1":
+                    if (mapped_model_bonded_entity != model_bonded_entity) or (
+                        ref_bonded_res_id != model_bonded_res_id
+                    ):
+                        ban_set.add((ref_chain_id, model_chain_id))
+        return ban_set
+
     def find_model_anchor_chains(self) -> str:
         """
         Ref: AlphaFold3 SI Chapter 4.2. -> AlphaFold Multimer Chapter 7.3.1
@@ -180,15 +245,22 @@ class ChainPermutation:
         row_indices = []
         col_indices = []
 
+        dist_matrix_copy = dist_matrix.copy()
         for _ in range(num_cols):
-            min_pos = np.unravel_index(np.argmin(dist_matrix), dist_matrix.shape)
+            min_pos = np.unravel_index(
+                np.argmin(dist_matrix_copy), dist_matrix_copy.shape
+            )
+
+            if dist_matrix_copy[min_pos[0], min_pos[1]] == np.inf:
+                # No more valid pairs
+                break
 
             row_indices.append(min_pos[0])
             col_indices.append(min_pos[1])
 
             # Set the found row and column to np.inf to ignore it
-            dist_matrix[min_pos[0], :] = np.inf
-            dist_matrix[:, min_pos[1]] = np.inf
+            dist_matrix_copy[min_pos[0], :] = np.inf
+            dist_matrix_copy[:, min_pos[1]] = np.inf
         return row_indices, col_indices
 
     @staticmethod
@@ -199,6 +271,7 @@ class ChainPermutation:
         struct2: Structure,
         coords1: np.ndarray,
         coords2: np.ndarray,
+        banned_chain_pairs: set[tuple[str, str]],
     ) -> dict[str, str]:
         """
         Chain mapping between two structures within the same entity using
@@ -212,12 +285,13 @@ class ChainPermutation:
             - Selecting the pair with minimal centroid distance
 
         Args:
-            chain_ids1 (list[str]): Chain IDs from first structure (may become struct2 after swap)
-            chain_ids2 (list[str]): Chain IDs from second structure (may become struct1 after swap)
+            chain_ids1 (list[str]): Chain IDs from first structure
+            chain_ids2 (list[str]): Chain IDs from second structure
             struct1 (Structure): First structure containing chain metadata
             struct2 (Structure): Second structure containing chain metadata
             coords1 (np.ndarray): Atom coordinates for struct1 (shape: [N, 3])
-            coords2 (np.ndarray): Atom coordinates for struct2 (shape: [M, 3])
+            coords2 (np.ndarray): Atom coordinates for struct2 (shape: [N, 3])
+            banned_chain_pairs (set[tuple[str, str]]): Pairs of chain IDs to be banned.
 
         Returns:
             dict[str, str]: Mapping of chain IDs from struct1 to struct2.
@@ -227,6 +301,8 @@ class ChainPermutation:
             chain_ids1, chain_ids2 = chain_ids2, chain_ids1
             struct1, struct2 = struct2, struct1
             coords1, coords2 = coords2, coords1
+            banned_chain_pairs = {(cid2, cid1) for cid1, cid2 in banned_chain_pairs}
+
             swapped = True
         else:
             swapped = False
@@ -241,6 +317,10 @@ class ChainPermutation:
             atoms1 = struct1.uni_atom_id[mask1]
 
             for cid2 in chain_ids2:
+                if (cid1, cid2) in banned_chain_pairs:
+                    # "inf" distance if the pair is banned
+                    continue
+
                 mask2 = struct2.uni_chain_id == cid2
                 atoms2 = struct2.uni_atom_id[mask2]
 
@@ -258,6 +338,9 @@ class ChainPermutation:
         matched_chains = {}
         row_indices, col_indices = ChainPermutation._find_min_indices(dist_mat)
         for row, col in zip(row_indices, col_indices):
+            if np.isinf(dist_mat[row, col]):
+                # "inf" distance if the pair is banned
+                continue
             matched_chains[chain_ids1[row]] = chain_ids2[col]
 
         return (
@@ -309,6 +392,7 @@ class ChainPermutation:
                 struct2=self.model_struct,
                 coords1=aligned_ref_coord,
                 coords2=self.model_struct.atom_array.coord,
+                banned_chain_pairs=self.ref_and_model_mapping_ban_set,
             )
             matched_chains.update(matched_chains_in_curr_entity)
         return matched_chains
@@ -443,6 +527,12 @@ class ChainPermutation:
         anchors = {}
 
         for ref_anchor_chain_id in ref_anchor_candidates:
+            if (
+                ref_anchor_chain_id,
+                model_anchor_chain_id,
+            ) in self.ref_and_model_mapping_ban_set:
+                continue
+
             # Find atoms in ref chain to match atoms in model chain
             ref_chain_mask = self.ref_struct.uni_chain_id == ref_anchor_chain_id
             ref_anchor_coord = self.ref_struct.atom_array.coord[ref_chain_mask]
@@ -555,9 +645,6 @@ class ChainPermutation:
                 - Anchors used for alignment.
         """
         model_anchor_chain_ids = self.find_model_anchor_chains()
-        if not self.enumerate_all_anchors:
-            # Only use the first anchor chain
-            model_anchor_chain_ids = [model_anchor_chain_ids[0]]
 
         ref_to_model_optimal_mapping = None
         best_rmsd = float("inf")
@@ -579,8 +666,17 @@ class ChainPermutation:
             ) = self.find_ref_to_model_optimal_chain_mapping(
                 model_anchor_chain_id, ref_anchor_candidates
             )
+
+            if mapping_i is None:
+                continue
+
             if best_rmsd_i < best_rmsd:
                 best_rmsd = best_rmsd_i
                 ref_to_model_optimal_mapping = mapping_i
                 best_anchors = anchors
+
+                if not self.enumerate_all_anchors:
+                    # Only use the first valid anchor chain
+                    break
+
         return ref_to_model_optimal_mapping, best_anchors

{pxmeter-0.1.2 → pxmeter-0.1.4/pxmeter.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pxmeter
-Version: 0.1.2
+Version: 0.1.4
 Summary: PXMeter is a comprehensive toolkit for evaluating the quality of         structures generated by biomolecular structure prediction models.
 Author: Bytedance Inc.
 Author-email: ai4s-bio@bytedance.com

{pxmeter-0.1.2 → pxmeter-0.1.4}/pxmeter.egg-info/SOURCES.txt

@@ -26,6 +26,7 @@ pxmeter/data/struct.py
 pxmeter/data/utils.py
 pxmeter/data/writer.py
 pxmeter/metrics/__init__.py
+pxmeter/metrics/clashes.py
 pxmeter/metrics/lddt_metrics.py
 pxmeter/metrics/rmsd.py
 pxmeter/metrics/rmsd_metrics.py

{pxmeter-0.1.2 → pxmeter-0.1.4}/setup.py

@@ -20,7 +20,7 @@ with open("requirements.txt") as f:
 setup(
     name="pxmeter",
     python_requires=">=3.11",
-    version="0.1.2",
+    version="0.1.4",
     description="PXMeter is a comprehensive toolkit for evaluating the quality of \
         structures generated by biomolecular structure prediction models.",
     author="Bytedance Inc.",