Add functions to group codons by sum given classes

libcodonusage/__init__.py

@@ -19,6 +19,7 @@ from .libcodonusage import (
     find_valley,
     format_codon_labels,
     gene_wide_codon_usage,
+    group_codons_by_sum,
     load_bias_table,
     load_counts_table,
     load_table_with_info_index,
@@ -27,11 +28,13 @@ from .libcodonusage import (
     make_centroids_cluster_finder,
     make_counts_only,
     plot_codon_usage_for_gene_list,
+    remove_codons,
     render_md,
     save_counts_table,
     sort_counts_by_aa,
     split_info_index,
     star2stop,
+    sum_codon_counts,
     to_long_form,
     violin_usage,
     violin_usage_vertical,

libcodonusage/libcodonusage.py

@@ -141,7 +141,6 @@ def load_counts_table(
     display(codon_counts.head(3))
     return codon_counts
 
-#def filter_codon_counts_table(codon_counts, )
 
 def compute_criteria(codon_counts):
     """
@@ -446,6 +445,37 @@ SUZUKI_DOI = "10.1016/j.febslet.2005.10.032"
 SUZUKI_LINK = f"[Suzuki et al (2005)](https://doi.org/{SUZUKI_DOI})"
 
 
+def remove_codons(codon_counts, codon_list):
+    """
+    Filter out codons in a table *codon_counts* based on codons present in the list *codon_list* (like stop codons).
+    """
+    codon_counts.drop(columns=codon_list, inplace=True)
+    return codon_counts
+
+
+def sum_codon_counts(row, codons):
+    """
+    Perform the row-wise sum of codon counts for the codons present in *codons* list given the row *row*.
+    """
+    sum = 0
+    for cod in codons:
+        sum += row[cod]
+    return sum
+
+
+def group_codons_by_sum(codon_counts, class_name, dict_classes, filter):
+    
+    list_classes = list(dict_classes.items())
+    list_classes_names = []
+    for key, value in dict_classes.items():
+        codon_counts[class_name, key] = codon_counts.apply(lambda row: sum_codon_counts(row, value), axis=1)
+        list_classes_names.append(key)
+    if filter:
+        return codon_counts.loc[:, ([class_name], list_classes_names)]
+    else:
+        return codon_counts
+
+
 def gene_wide_codon_usage(
         codon_counts,
         verbose=False, return_more=False, ref_filter_dict=None):
@@ -482,7 +512,9 @@ using the "l1" norm (which, for positive-only values amounts to the sum).
     colsums = codon_proportions.sum(axis=1).values
     # Due to imprecision in float arithmetics,
     # we can only check that the sums are close to 1
-    assert np.allclose(colsums, np.full(len(colsums), 1))
+    #assert np.allclose(colsums, np.full(len(colsums), 1))
+    print("mean", np.mean(colsums))
+    assert np.isclose(np.mean(colsums), 1)
     if ref_filter_dict is None:
         counts_for_global = codon_counts
     else: