[nomenclature._sort_all_consensus_labels] API CHANGE return also sort…

…ed indices
gph82 · Aug 12, 2024 · 5d76a56 · 5d76a56
1 parent 739074e
commit 5d76a56
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 6 deletions.
diff --git a/mdciao/nomenclature/nomenclature.py b/mdciao/nomenclature/nomenclature.py
@@ -1549,7 +1549,7 @@ def __init__(self, maps, tops=None):
         self._residxs["consensus"] = self._residxs.index.values
         self._residxs=self._residxs[["consensus"]+[key for key in self._residxs.keys() if key !="consensus"]]
 
-        sorted_keys = _sort_all_consensus_labels(self._residxs["consensus"], append_diffset=False)
+        sorted_keys = _sort_all_consensus_labels(self._residxs["consensus"], append_diffset=False)[0]
         assert len(sorted_keys)==len(self._residxs["consensus"]),  (len(sorted_keys), len(self._residxs["consensus"]))
         self._residxs = self._residxs.sort_values("consensus", key=lambda col: col.map(lambda x: sorted_keys.index(x)))
         self._residxs.index = _np.arange(len(self._residxs))
@@ -2428,7 +2428,7 @@ def _conslabel2fraglabel(labelres, defrag="@", prefix_GPCR=True):
         label = _GPCR_num2lett[label]
     return label
 
-def _sort_all_consensus_labels(labels, append_diffset=True, order=["GPCR","CGN","KLIFS"], ):
+def _sort_all_consensus_labels(labels, append_diffset=True, order=["GPCR","CGN","KLIFS"], return_argsort=False):
     r"""
     Sort a mix of consensus labels GPCR, CGN, KLIFS
 
@@ -2442,7 +2442,7 @@ def _sort_all_consensus_labels(labels, append_diffset=True, order=["GPCR","CGN",
         end of `sorted_labels` unless
         explicitly deactivated with
         `append_diffset`.
-        append_diffset : bool, default is True
+    append_diffset : bool, default is True
         Append the non-consensus labels
         at the end of `sorted_labels`
     order : list
@@ -2458,6 +2458,11 @@ def _sort_all_consensus_labels(labels, append_diffset=True, order=["GPCR","CGN",
     -------
     sorted_labels : list
         Sorted consensus labels
+    sorted_indices : 1D _np.ndarray
+        The indices of `labels` that return
+        the sorted `soted_labels`. Depending
+        on `append_diffset` it will contain
+        (or not) all indices of `labels`
     """
 
     lambdas = {"GPCR":  lambda labels: _sort_GPCR_consensus_labels(labels, append_diffset=False),
@@ -2470,7 +2475,10 @@ def _sort_all_consensus_labels(labels, append_diffset=True, order=["GPCR","CGN",
     if append_diffset:
         sorted_labels += [lab for lab in labels if lab not in sorted_labels]
 
-    return sorted_labels
+    sorted_indices = [_np.flatnonzero(lab==_np.array(labels)) for lab in sorted_labels]
+    sorted_indices = _np.hstack([si for si in sorted_indices if len(si)>0]).squeeze()
+
+    return sorted_labels, sorted_indices
 
 _GPCR_num2lett = {
     "1": "TM1 ",

diff --git a/tests/test_nomenclature.py b/tests/test_nomenclature.py
@@ -857,16 +857,22 @@ def test_KLIFS_dont_append(self):
                                        sorted)
 
     def test_sort_all_consensus_labels(self):
-        sorted = nomenclature._sort_all_consensus_labels(self.tosort, append_diffset=False, order = ["CGN","KLIFS", "GPCR"])
+        sorted, sorted_indices = nomenclature._sort_all_consensus_labels(self.tosort, append_diffset=False, order = ["CGN","KLIFS", "GPCR"])
         _np.testing.assert_array_equal(
             [
              "G.H1.1", "G.H1.10", "H.HA.10", "H.HA.20",
              'αC.25', 'αD.55', 'a.l.85', 
              "2.50", "3.50", "H8.1", "H8.10"],
             sorted)
 
+        _np.testing.assert_array_equal(
+            [9, 0, 10, 1,
+             4, 11, 2,
+             8, 7, 6, 3],
+            sorted_indices)
+
     def test_sort_all_consensus_labels_append(self):
-        sorted = nomenclature._sort_all_consensus_labels(self.tosort, append_diffset=True,
+        sorted, sorted_indices = nomenclature._sort_all_consensus_labels(self.tosort, append_diffset=True,
                                                          order=["CGN", "KLIFS"])
         _np.testing.assert_array_equal(
             [
@@ -876,6 +882,12 @@ def test_sort_all_consensus_labels_append(self):
                  ],
             sorted)
 
+        _np.testing.assert_array_equal(
+            [9, 0, 10, 1,
+             4, 11, 2,
+             3, 5, 6, 7, 8],
+            sorted_indices)
+
 class Test_compatible_consensus_fragments(TestClassSetUpTearDown_CGN_local):
 
     def setUp(self):