From ebff913766029a0b843f5e84d30eee19f2a0f608 Mon Sep 17 00:00:00 2001
From: Blaise Li <blaise.li__git@nsup.org>
Date: Thu, 14 Sep 2023 15:54:41 +0200
Subject: [PATCH] Add function to split table with info in index.

---
 libcodonusage/__init__.py      |  1 +
 libcodonusage/libcodonusage.py | 24 ++++++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/libcodonusage/__init__.py b/libcodonusage/__init__.py
index 266ff34..fcfc26c 100644
--- a/libcodonusage/__init__.py
+++ b/libcodonusage/__init__.py
@@ -30,6 +30,7 @@ from .libcodonusage import (
     render_md,
     save_counts_table,
     sort_counts_by_aa,
+    split_info_index,
     star2stop,
     to_long_form,
     violin_usage,
diff --git a/libcodonusage/libcodonusage.py b/libcodonusage/libcodonusage.py
index dbd1848..cf50fbf 100644
--- a/libcodonusage/libcodonusage.py
+++ b/libcodonusage/libcodonusage.py
@@ -398,6 +398,30 @@ def load_table_with_info_index(table_path, nb_info_cols, nb_cluster_series=0):
         header=[0, 1])
 
 
+def split_info_index(table, keep_index_cols):
+    """
+    Split table *table* into info and data.
+
+    Info is supposed to be contained in the index of DataFrame *table*,
+    data in its columns.
+
+    Return a pair of tables, one for the data, one for the info,
+    where the index contains only the levels listed in *keep_index_cols*
+    """
+    drop_info_cols = [
+        colname for colname in table.index.names
+        if colname not in keep_index_cols]
+    info_table = table.reset_index()[
+        [*keep_index_cols, *drop_info_cols]].set_index(keep_index_cols)
+    # To avoid loss of index levels in input by side effect:
+    out_table = table.copy()
+    out_table.index = out_table.index.droplevel(drop_info_cols)
+    # To ensure indices have their levels in the same order
+    # in out_table and in infoinfo_table:
+    out_table = out_table.reset_index().set_index(keep_index_cols)
+    return (out_table, info_table)
+
+
 def filter_on_idx_levels(counts_table, filter_dict):
     """
     Filter a table *counts_table* based on values of certain index levels.
-- 
GitLab