prune_atom_types added as function

python/BioSimSpace/Align/_align.py

@@ -727,6 +727,7 @@ def matchAtoms(
     roi=None,
     prune_perturbed_constraints=False,
     prune_crossing_constraints=False,
+    prune_atom_types=None,
     property_map0={},
     property_map1={},
 ):
@@ -799,6 +800,10 @@ def matchAtoms(
         constraints between dummy and non-dummy atoms. This option should
         be True when creating mappings to use with the AMBER engine.
 
+    prune_atom_types : bool
+        Whether to remove atoms from the mapping such that there are no
+        changing atom types.
+
     property_map0 : dict
         A dictionary that maps "properties" in molecule0 to their user
         defined values. This allows the user to refer to properties
@@ -873,6 +878,7 @@ def matchAtoms(
             max_scoring_matches=max_scoring_matches,
             prune_perturbed_constraints=prune_perturbed_constraints,
             prune_crossing_constraints=prune_crossing_constraints,
+            prune_atom_types=prune_atom_types,
             property_map0=property_map0,
             property_map1=property_map1,
         )
@@ -883,6 +889,7 @@ def matchAtoms(
             roi=roi,
             prune_perturbed_constraints=prune_perturbed_constraints,
             prune_crossing_constraints=prune_crossing_constraints,
+            prune_atom_types=prune_atom_types,
             use_kartograf=False,
             kartograf_kwargs={},
         )
@@ -901,6 +908,7 @@ def _matchAtoms(
     roi=None,
     prune_perturbed_constraints=False,
     prune_crossing_constraints=False,
+    prune_atom_types=False,
     property_map0={},
     property_map1={},
 ):
@@ -969,6 +977,9 @@ def _matchAtoms(
     if not isinstance(prune_crossing_constraints, bool):
         raise TypeError("'prune_crossing_constraints' must be of type 'bool'")
 
+    if not isinstance(prune_atom_types, bool):
+        raise TypeError("'prune_atom_types' must be of type 'bool'")
+
     if not isinstance(property_map0, dict):
         raise TypeError("'property_map0' must be of type 'dict'")
 
@@ -1109,6 +1120,12 @@ def _matchAtoms(
             for mapping in mappings
         ]
 
+    if prune_atom_types:
+        mappings = [
+            _prune_atom_types(molecule0, molecule1, mapping)
+            for mapping in mappings
+        ]
+
     if matches == 1:
         if return_scores:
             return (mappings[0], scores[0])
@@ -1203,6 +1220,7 @@ def _roiMatch(
     roi,
     prune_perturbed_constraints=False,
     prune_crossing_constraints=False,
+    prune_atom_types=False,
     **kwargs,
 ):
     """
@@ -1234,6 +1252,10 @@ def _roiMatch(
         Whether to remove atoms from the mapping such that there are no
         constraints between dummy and non-dummy atoms. This option should
         be True when creating mappings to use with the AMBER engine.
+        
+    prune_atom_types : bool
+        Whether to remove atoms from the mapping such that there are no
+        changing atom types.
 
     use_kartograf : bool, optional, default=False
         If set to True, will use the kartograf algorithm to match the
@@ -1312,6 +1334,9 @@ def _roiMatch(
     if not isinstance(prune_crossing_constraints, bool):
         raise TypeError("'prune_crossing_constraints' must be of type 'bool'")
 
+    if not isinstance(prune_atom_types, bool):
+        raise TypeError("'prune_atom_types' must be of type 'bool'")
+
     # Check kwargs.
     use_kartograf = kwargs.get("use_kartograf", False)
     kartograf_kwargs = kwargs.get("kartograf_kwargs", {})
@@ -1485,6 +1510,12 @@ def _roiMatch(
     if prune_crossing_constraints:
         full_mapping = _prune_crossing_constraints(molecule0, molecule1, full_mapping)
 
+    if prune_atom_types:
+        mappings = [
+            _prune_atom_types(molecule0, molecule1, mapping)
+            for mapping in mappings
+        ]
+
     return full_mapping
 
 
@@ -3048,6 +3079,37 @@ def _from_sire_mapping(sire_mapping):
 
     return mapping
 
+def _prune_atom_types(molecule0, molecule1, mapping):
+    """
+        Prunes the maximum common substructure (MCS) so that there are no atoms changing type.
+
+       Parameters
+       ----------
+       molecule0 : BioSimSpace._SireWrappers.Molecule
+           The first molecule (used at lambda = 0).
+       molecule1 : BioSimSpace._SireWrappers.Molecule
+           The second molecule (used at lambda = 1).
+       mapping : dict(int, int)
+           A maximum common substructure mapping between both molecules, as generated by e.g.
+           BioSimSpace.Align.matchAtoms().
+
+       Returns
+       -------
+       new_mapping : dict(int, int)
+           The pruned MCS.
+    """
+    new_mapping = {}
+
+    for idx0, idx1 in mapping.items():
+        atom0 = molecule0.getAtoms()[idx0]
+        atom1 = molecule1.getAtoms()[idx1]
+        elem0 = atom0._sire_object.property("element").symbol()
+        elem1 = atom1._sire_object.property("element").symbol()
+
+        if elem0 == elem1:
+            new_mapping[idx0] = idx1
+
+    return new_mapping
 
 def _prune_perturbed_constraints(molecule0, molecule1, mapping):
     """