Added function to group counts columns by aa.

libcodonusage/__init__.py

 __copyright__ = "Copyright (C) 2022 Blaise Li"
 __licence__ = "GNU GPLv3"
-__version__ = 0.4
+__version__ = 0.5
 from .libcodonusage import (
     aa2colour,
     codon2aa,
@@ -11,6 +11,7 @@ from .libcodonusage import (
     make_aa_codon_columns,
     make_counts_only,
     render_md,
+    sort_counts_by_aa,
     violin_usage,
     violin_usage_vertical,
     violin_usage_by_clusters,

libcodonusage/libcodonusage.py

@@ -72,6 +72,9 @@ columns_by_aa = groupby(itemgetter(0), codon_headers)
 ######################################
 # Associating colours to amino-acids #
 ######################################
+# We load amino-acid colouring information to use in graphical representations.
+# We'll use colour schemes provided by the biotite package:
+# https://www.biotite-python.org/examples/gallery/sequence/color_schemes_protein.html
 # We look directly at the json file
 # instead of using the "native" biotite mechanisms
 with Path(bgraphs.colorschemes._scheme_dir).joinpath(
@@ -269,9 +272,36 @@ def make_aa_codon_columns(counts_table):
     counts_table.columns = pd.MultiIndex.from_tuples(
         ((codon2aa[codon], codon) for codon in counts_table.columns),
         names=("aa", "codon"))
+    render_md(
+        "We associated amino-acid information to codons, "
+        "as an extra level in the table columns.")
     return counts_table
 
 
+def sort_counts_by_aa(counts_table):
+    """
+    Sort columns so that codons are grouped by corresponding amino-acid.
+    """
+    aacodons_order = list(aa2colour.keys())
+    # We need to include both aas and codons in the list
+    # because it will be used for both the aa and the codon levels
+    # when sorting the (aa, codon) MultiIndex
+    aacodons_order.extend(codon2aa.keys())
+
+    def aa_sortkey(col):
+        """
+        Key function to use when sorting a MultiIndex consisting in
+        (aa, codon) pairs.
+        """
+        return col.map(aacodons_order.index)
+    sorted_counts = counts_table.sort_index(
+        axis=1, level=(0, 1), ascending=(True, True), key=aa_sortkey)
+    render_md(
+        "The columns were sorted in order to group codons"
+        " by associated amino-acid.")
+    return sorted_counts
+
+
 def load_bias_table(table_path, nb_cluster_series=2):
     """
     Load a table containing by-amino-acid codon usage biases.