Add functionality to prune changing atom types from mapping.

python/BioSimSpace/Align/_align.py

@@ -727,6 +727,7 @@ def matchAtoms(
     roi=None,
     prune_perturbed_constraints=False,
     prune_crossing_constraints=False,
+    prune_atom_types=False,
     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'")
 
@@ -1108,6 +1119,10 @@ def _matchAtoms(
             _prune_crossing_constraints(molecule0, molecule1, mapping)
             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:
@@ -1203,6 +1218,7 @@ def _roiMatch(
     roi,
     prune_perturbed_constraints=False,
     prune_crossing_constraints=False,
+    prune_atom_types=False,
     **kwargs,
 ):
     """
@@ -1235,6 +1251,10 @@ def _roiMatch(
         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
         molecules.
@@ -1312,6 +1332,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", {})
@@ -1479,12 +1502,16 @@ def _roiMatch(
     }
 
     # Optionally post-process the MCS for use with AMBER.
+
     if prune_perturbed_constraints:
         full_mapping = _prune_perturbed_constraints(molecule0, molecule1, full_mapping)
 
     if prune_crossing_constraints:
         full_mapping = _prune_crossing_constraints(molecule0, molecule1, full_mapping)
 
+    if prune_atom_types:
+        full_mapping = _prune_atom_types(molecule0, molecule1, full_mapping)
+
     return full_mapping
 
 
@@ -3172,3 +3199,41 @@ def _prune_crossing_constraints(molecule0, molecule1, mapping):
             mapping = new_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 : class:`Molecule <BioSimSpace._SireWrappers.Molecule>
+        The first molecule (used at lambda = 0).
+
+    molecule1 : class:`Molecule <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 mapping.
+    """
+    new_mapping = {}
+
+    for idx0, idx1 in mapping.items():
+        atom0 = molecule0.getAtoms()[idx0]
+        atom1 = molecule1.getAtoms()[idx1]
+        elem0 = atom0._sire_object.property("element")
+        elem1 = atom1._sire_object.property("element")
+
+        if elem0 == elem1:
+            new_mapping[idx0] = idx1
+
+    return new_mapping