feat: add merged_lipid_classes

Author: Selok

lta/helpers/pipeline.py

@@ -100,9 +100,7 @@ def __post_init__(self) -> None:
             logger.exception(f"{self.file} is a directory. A full traceback follows...")
             raise
         except pd.errors.EmptyDataError:
-            logger.exception(
-                f"{self.file} contains no data. A full traceback follows..."
-            )
+            logger.exception(f"{self.file} contains no data. A full traceback follows...")
             raise
         logger.debug("Binarizing data...")
         self.binary = {
@@ -117,19 +115,13 @@ def __post_init__(self) -> None:
         }
         logger.debug("Filtering data...")
         self.filtered = {
-            group: df.loc[self.binary[group].index, :]
-            for group, df in data.groupby(axis="columns", level=self.mode)
+            group: df.loc[self.binary[group].index, :] for group, df in data.groupby(axis="columns", level=self.mode)
         }
         Path(self.output, "enfc").mkdir(exist_ok=True, parents=True)
         Path(self.output, "jaccard").mkdir(exist_ok=True, parents=True)
 
-        conditions = [
-            df.columns.get_level_values(self.level).unique()
-            for df in self.filtered.values()
-        ]
-        self.conditions = [
-            val for mode in conditions for val in mode if val != self.control
-        ]
+        conditions = [df.columns.get_level_values(self.level).unique() for df in self.filtered.values()]
+        self.conditions = [val for mode in conditions for val in mode if val != self.control]
 
     def _calculate_enfc(self) -> Dict[str, Dict[str, pd.DataFrame]]:
         """Calculate error-normalised fold change.
@@ -178,9 +170,7 @@ def _get_a_lipids(self) -> Dict[str, pd.DataFrame]:
         logger.info("Calculating A-lipids...")
         results = {
             f"a_{mode}": (
-                df.groupby(axis="columns", level=self.level)
-                .all()
-                .pipe(lambda x: x.loc[x.any(axis="columns"), :])
+                df.groupby(axis="columns", level=self.level).all().pipe(lambda x: x.loc[x.any(axis="columns"), :])
             )
             for mode, df in self.binary.items()
         }
@@ -231,16 +221,10 @@ def _get_b_lipids(self, picky: bool = True) -> Dict[str, pd.DataFrame]:
             # This assumes that self.binary and a_lip will have the same keys
             # Which is definitely True
             if picky:
-                data = {
-                    mode: df.drop(index=a_lip[f"a_{mode}"])
-                    for mode, df in self.binary.items()
-                }
+                data = {mode: df.drop(index=a_lip[f"a_{mode}"]) for mode, df in self.binary.items()}
                 subtype = "p"
             else:
-                data = {
-                    mode: df.loc[a_lip[f"a_{mode}"], :]
-                    for mode, df in self.binary.items()
-                }
+                data = {mode: df.loc[a_lip[f"a_{mode}"], :] for mode, df in self.binary.items()}
                 subtype = "c"
 
         logger.info(f"Calculating B{subtype}-lipids...")
@@ -296,47 +280,32 @@ def _get_n_lipids(self, n: int) -> Dict[str, pd.DataFrame]:
             compartments = df.columns.get_level_values(self.compartment)
             # Mask required to prevent dropping levels
             # The initial check must be done with all compartments (only n)...
-            mask = (
-                df.groupby(axis="columns", level=self.compartment)
-                .any()
-                .sum(axis="columns")
-                == n
-            )
+            mask = df.groupby(axis="columns", level=self.compartment).any().sum(axis="columns") == n
 
             data = [
                 (group, df.loc[mask, compartments.isin(group)])
                 for group in itertools.combinations(compartments.unique(), n)
             ]
-            logger.debug(
-                f"N{n} compartment groups before filtering: {[group for group, _ in data]}"
-            )
+            logger.debug(f"N{n} compartment groups before filtering: {[group for group, _ in data]}")
             # ...which necessitates a second check to drop those that are not
             # Again, mask necessary for keeping info
             # Also, we only care for groups with lipids
             masks = [
-                df.groupby(axis="columns", level=self.compartment)
-                .any()
-                .sum(axis="columns")
-                == n
-                for _, df in data
+                df.groupby(axis="columns", level=self.compartment).any().sum(axis="columns") == n for _, df in data
             ]
             data = [
                 (group, df.loc[mask, :].groupby(axis="columns", level=self.level).all())
                 for (group, df), mask in zip(data, masks)
                 if mask.sum() != 0
             ]
-            logger.debug(
-                f"N{n} compartment groups after filtering: {[group for group, _ in data]}"
-            )
+            logger.debug(f"N{n} compartment groups after filtering: {[group for group, _ in data]}")
             for compartments, df in data:
                 n_type = "u" if n == 1 else f"n{n}"
                 group = "_".join([x.upper() for x in compartments])
                 results[f"{n_type}_{group}_{mode}"] = df
         return results
 
-    def _jaccard(
-        self, data: Dict[str, pd.DataFrame], group: str
-    ) -> Dict[str, Dict[str, pd.DataFrame]]:
+    def _jaccard(self, data: Dict[str, pd.DataFrame], group: str) -> Dict[str, Dict[str, pd.DataFrame]]:
         """Calculate jaccard similarity and p-values.
 
         This takes a dictionary of data.
@@ -367,84 +336,69 @@ def _jaccard(
                 mode: lipids.loc[:, [group, self.control]]
                 .pipe(lambda df: df.loc[df.sum(axis=1) != 0, :])
                 .groupby(axis="index", level="Category")
-                .apply(
-                    lambda x: jac.bootstrap(
-                        x.loc[:, group], x.loc[:, self.control], n=self.n
-                    )
-                )
+                .apply(lambda x: jac.bootstrap(x.loc[:, group], x.loc[:, self.control], n=self.n))
                 for mode, lipids in data.items()
             }
             for group in self.conditions
         }
         return jaccard
 
-    def _generate_enfc_summary(self) -> None:
+    def _generate_enfc_summary(self) -> pd.DataFrame:
         logger.debug("Generating ENFC summary files...")
         enfcs = self._calculate_enfc()
         frames = []
         levels = set()
-        for group, data in enfcs.items():
+        for phenotype, data in enfcs.items():
             df = pd.concat(data, axis="columns")
-            df.to_csv(
-                self.output
-                / "enfc"
-                / f"{group}_by_{self.control}_individual_lipids.csv"
-            )
-            df.columns = utils.add_level_to_index(index=df.columns, new_level=group, new_level_name="Group")
+            df.to_csv(self.output / "enfc" / f"{phenotype}_by_{self.control}_individual_lipids.csv")
+            df.columns = utils.add_level_to_index(index=df.columns, new_level=phenotype, new_level_name="Phenotype")
             frames.append(df)
             levels.update(df.index.names)
-        utils.merge_dataframe_by_level(datas=frames, levels=levels).to_csv(
-            self.output / "enfc" / f"individual_lipids.csv"
-        )
+        summary = utils.merge_dataframe_by_level(datas=frames, levels=levels)
+        summary.to_csv(self.output / "enfc" / f"individual_lipids.csv")
+        return summary
 
-    def _generate_enfc_class_summary(self) -> None:
+    def _generate_enfc_class_summary(self) -> pd.DataFrame:
         logger.debug("Generating class ENFC summary files...")
-        self.filtered = {
-            mode: df.groupby(axis="index", level="Category").sum()
-            for mode, df in self.filtered.items()
-        }
+        self.filtered = {mode: df.groupby(axis="index", level="Category").sum() for mode, df in self.filtered.items()}
         self.enfcs = self._calculate_enfc()
         frames = []
         levels = set()
-        for group, data in self.enfcs.items():
+        for phenotype, data in self.enfcs.items():
             df = pd.concat(data, axis="columns")
-            df.to_csv(
-                self.output / "enfc" / f"{group}_by_{self.control}_lipid_classes.csv"
-            )
-            df.columns = utils.add_level_to_index(index=df.columns, new_level=group, new_level_name="Group")
+            df.to_csv(self.output / "enfc" / f"{phenotype}_by_{self.control}_lipid_classes.csv")
+            df.columns = utils.add_level_to_index(index=df.columns, new_level=phenotype, new_level_name="Phenotype")
             frames.append(df)
             levels.update(df.index.names)
-        utils.merge_dataframe_by_level(datas=frames, levels=levels).to_csv(
-            self.output / "enfc" / f"lipid_classes.csv"
-        )
+        summary = utils.merge_dataframe_by_level(datas=frames, levels=levels)
+        summary.to_csv(self.output / "enfc" / f"lipid_classes.csv")
+        return summary
 
-    def _generate_jaccard_distance_summary(self) -> None:
+    def _generate_jaccard_distance_summary(self) -> pd.DataFrame:
         logger.debug("Generating Jaccard distance summary files...")
         frames = []
         levels = set()
-        for group in set(self.conditions):
+        for phenotype in set(self.conditions):
             jaccard = pd.concat(
                 {
-                    **self.a_jaccard[group],
-                    **self.bc_jaccard[group],
-                    **self.bp_jaccard[group],
-                    **self.n2_jaccard[group],
-                    **self.u_jaccard[group],
+                    **self.a_jaccard[phenotype],
+                    **self.bc_jaccard[phenotype],
+                    **self.bp_jaccard[phenotype],
+                    **self.n2_jaccard[phenotype],
+                    **self.u_jaccard[phenotype],
                 },
                 axis="columns",
             )
             jaccard.columns.names = ["type_compartment_mode", "Metrics"]
-            jaccard.to_csv(
-                self.output
-                / "jaccard"
-                / f"{group}_to_{self.control}_jaccard_similarity.csv"
+            jaccard.to_csv(self.output / "jaccard" / f"{phenotype}_to_{self.control}_jaccard_similarity.csv")
+            jaccard.columns = utils.add_level_to_index(
+                index=jaccard.columns, new_level=phenotype, new_level_name="Phenotype"
             )
-            jaccard.columns = utils.add_level_to_index(index=jaccard.columns, new_level=group, new_level_name="Group")
             frames.append(jaccard)
             levels.update(jaccard.index.names)
-        utils.merge_dataframe_by_level(datas=frames, levels=levels).to_csv(
-            self.output / "jaccard" / f"jaccard_similarity.csv"
-        )
+        summary = utils.merge_dataframe_by_level(datas=frames, levels=levels)
+        summary.to_csv(self.output / "jaccard" / f"jaccard_similarity.csv")
+        return summary
 
     def run(self) -> None:
         """Run the full LTA pipeline.
@@ -490,8 +444,30 @@ def run(self) -> None:
         ).fillna(False)
         summary.columns.names = ["type_compartment_mode", "Phenotype"]
         summary.to_csv(self.output / "switch_individual_lipids.csv")
-        summary.groupby(axis="index", level="Category").sum().to_csv(
-            self.output / "switch_lipid_classes.csv"
-        )
+        lipid_classes = summary.groupby(axis="index", level="Category").sum()
+        lipid_classes.to_csv(self.output / "switch_lipid_classes.csv")
+
+        jaccard_similarity = self._generate_jaccard_distance_summary()
 
-        self._generate_jaccard_distance_summary()
+        lipid_classes.columns = utils.add_level_to_index(
+            index=lipid_classes.columns, new_level="-", new_level_name="Metrics"
+        )
+        jaccard_similarity.columns = utils.reorder_index(
+            index=jaccard_similarity.columns, orders=lipid_classes.columns.names
+        )
+        merged_lipid_classes = utils.merge_dataframe_by_level(
+            datas=[
+                lipid_classes,
+                jaccard_similarity,
+            ],
+            levels=lipid_classes.index.names,
+        )
+        merged_lipid_classes = utils.sort_columns(
+            data=merged_lipid_classes, level="type_compartment_mode", 
+            pressing=lipid_classes.columns.get_level_values("type_compartment_mode").unique().to_list()
+        )
+        merged_lipid_classes = utils.sort_columns(data=merged_lipid_classes, level="Phenotype", pressing=[self.control])
+        merged_lipid_classes = utils.sort_columns(
+            data=merged_lipid_classes, level="Metrics", pressing=["-"]
+        )
+        merged_lipid_classes.to_csv(self.output / "merged_lipid_classes.csv")

lta/helpers/utils.py

@@ -1,3 +1,5 @@
+import itertools
+import operator
 import typing
 import pandas as pd
 
@@ -11,14 +13,49 @@ def merge_dataframe_by_level(
 
 
 def add_level_to_index(
-    *,
-    index: typing.Union[pd.MultiIndex, pd.Index],
-    new_level: typing.Any,
-    new_level_name: str
+    *, index: typing.Union[pd.MultiIndex, pd.Index], new_level: typing.Any, new_level_name: str
 ) -> pd.MultiIndex:
     names = index.names + [new_level_name]
     if isinstance(index, pd.MultiIndex):
         index = index.to_flat_index()
-    return pd.MultiIndex.from_tuples(
-        list(val + (new_level,) for val in index.array), names=names
+    return pd.MultiIndex.from_tuples(list(val + (new_level,) for val in index.array), names=names)
+
+
+def reorder_index(
+    *,
+    index: typing.Union[pd.MultiIndex, pd.Index],
+    orders: typing.List[typing.Union[str, int]],
+) -> pd.MultiIndex:
+    if isinstance(index, pd.MultiIndex):
+        return index.reorder_levels(orders)
+    return index
+
+
+def sort_columns(
+    *,
+    data: pd.DataFrame,
+    level: typing.Union[str, int],
+    pressing: typing.List[typing.Union[str, int]],
+) -> pd.DataFrame:
+    idx = data.columns
+    pressed = set(pressing)
+    # get correct ordering of the index
+    level_ordering = pressing + list(
+        # filter out the pressing columns
+        map(
+            operator.itemgetter(0),
+            itertools.groupby(sorted(v for v in idx.get_level_values(level).values if v not in pressed)),
+        )
+    )
+
+    if not isinstance(idx, pd.MultiIndex):
+        return data.reindex(columns=level_ordering, level=level)
+
+    # recreate Multiindex with mapping and reindex all columns
+    new_level_index = data.reindex(level=level, columns=level_ordering).columns
+    new_column_order_map = dict(zip(idx[idx.isin(new_level_index)], new_level_index))
+    new_multi_index = pd.MultiIndex.from_tuples(
+        [new_column_order_map.get(x, x) for x in idx],
+        names=idx.names,
     )
+    return data.reindex(new_multi_index, axis=1)

tests/unit/helpers/__init__.py

@@ -0,0 +1,53 @@
+import pandas as pd
+import numpy as np
+
+from lta.helpers import utils
+
+
+def test_merge_dataframe_by_level() -> None:
+    idx = pd.MultiIndex.from_product([["A", "B", "C"], ["1", "2", "3"]], names=["alphabet", "numeric"])
+    col = ["c1", "c2", "c3", "c4"]
+
+    df1_data, df2_data = (np.array([[f"d{i+1}_{i1}{i2}_{c}" for c in col] for i1, i2 in idx]) for i in range(2))
+    assert (df1_data != df2_data).all()
+    frames = [
+        pd.DataFrame(df1_data, idx, col),
+        pd.DataFrame(df2_data, idx, col),
+    ]
+    res = utils.merge_dataframe_by_level(datas=frames, levels=idx.levels)
+    assert res.index.identical(idx)
+    assert res.columns.identical(pd.Index(col * 2))
+    assert (res.to_numpy() == np.append(df1_data, df2_data, axis=1)).all()
+
+
+def test_add_level_to_multi_index() -> None:
+    idx = pd.MultiIndex.from_product([["A", "B", "C"], ["1", "2", "3"]], names=["alphabet", "numeric"])
+    expected = pd.MultiIndex.from_product(
+        [["A", "B", "C"], ["1", "2", "3"], ["ADDED"]],
+        names=["alphabet", "numeric", "newly added"],
+    )
+    res = utils.add_level_to_index(index=idx, new_level="ADDED", new_level_name="newly added")
+    assert res.identical(expected)
+
+
+def test_reorder_column_index() -> None:
+    idx = pd.MultiIndex.from_product([["A", "B", "C"], ["1", "2", "3"]], names=["alphabet", "numeric"])
+    new_idx = utils.reorder_index(index=idx, orders=["numeric", "alphabet"])
+    col = ["c1", "c2", "c3", "c4"]
+    df_data = np.array([[f"d{i1}{i2}_{c}" for c in col] for i1, i2 in idx])
+    expected = pd.DataFrame(df_data, new_idx, col)
+    res = pd.DataFrame(df_data, idx, col)
+    res.index = new_idx
+    assert expected.equals(res)
+
+
+def test_sort_columns() -> None:
+    idx = ["c1", "c2", "c3", "c4"]
+    expected_col = pd.MultiIndex.from_product([["A", "B", "C"], ["2", "3", "1"]], names=["alphabet", "numeric"])
+    expected_data = np.array([[f"d{i1}{i2}_{c}" for c in idx] for i1, i2 in expected_col]).T
+    expected = pd.DataFrame(expected_data, idx, expected_col)
+    col = pd.MultiIndex.from_product([["A", "B", "C"], ["1", "2", "3"]], names=["alphabet", "numeric"])
+    df_data = np.array([[f"d{i1}{i2}_{c}" for c in idx] for i1, i2 in col]).T
+    df = pd.DataFrame(df_data, idx, col)
+    res = utils.sort_columns(data=df, level="numeric", pressing=["2", "3"])
+    assert expected.equals(res)