diff --git a/modules/mol/alg/pymod/chain_mapping.py b/modules/mol/alg/pymod/chain_mapping.py index 865e3b2ab447559c8f8b7ada35dd37b1feb51270..4b90652f2fee00f5eb32d58462e51ffa29f0aaa6 100644 --- a/modules/mol/alg/pymod/chain_mapping.py +++ b/modules/mol/alg/pymod/chain_mapping.py @@ -89,6 +89,25 @@ class MappingResult: """ return self._alns + def GetFlatMapping(self, mdl_as_key=False): + """ Returns flat mapping as :class:`dict` for all mapable chains + + :param mdl_as_key: Default is target chain name as key and model chain + name as value. This can be reversed with this flag. + :returns: :class:`dict` with :class:`str` as key/value that describe + one-to-one mapping + """ + flat_mapping = dict() + for trg_chem_group, mdl_chem_group in zip(self.chem_groups, + self.mapping): + for a,b in zip(trg_chem_group, mdl_chem_group): + if a is not None and b is not None: + if mdl_as_key: + flat_mapping[b] = a + else: + flat_mapping[a] = b + return flat_mapping + def JSONSummary(self): """ Returns JSON serializable summary of results """ diff --git a/modules/mol/alg/tests/test_chain_mapping.py b/modules/mol/alg/tests/test_chain_mapping.py index d9a9362eba4bc5fea3488c77a62ac36c884a9af5..0a663081ce80dbafb9b9978e13450138baf83bec 100644 --- a/modules/mol/alg/tests/test_chain_mapping.py +++ b/modules/mol/alg/tests/test_chain_mapping.py @@ -284,6 +284,19 @@ class TestChainMapper(unittest.TestCase): greedy_rigid_res = mapper.GetRigidMapping(mdl, strategy="greedy_iterative_rmsd") self.assertEqual(greedy_rigid_res.mapping, [['X', 'Y'],[None],['Z']]) + # test flat mapping functionality of MappingResult + flat_map = greedy_rigid_res.GetFlatMapping() + self.assertEqual(len(flat_map), 3) + self.assertEqual(flat_map[greedy_rigid_res.chem_groups[0][0]], 'X') + self.assertEqual(flat_map[greedy_rigid_res.chem_groups[0][1]], 'Y') + self.assertEqual(flat_map[greedy_rigid_res.chem_groups[2][0]], 'Z') + flat_map = greedy_rigid_res.GetFlatMapping(mdl_as_key=True) + self.assertEqual(len(flat_map), 3) + self.assertEqual(greedy_rigid_res.chem_groups[0][0], flat_map['X']) + self.assertEqual(greedy_rigid_res.chem_groups[0][1], flat_map['Y']) + self.assertEqual(greedy_rigid_res.chem_groups[2][0], flat_map['Z']) + + if __name__ == "__main__": from ost import testutils if testutils.SetDefaultCompoundLib():