added generate_tree function to Kinome object (#51)

* added generate_tree function to Kinome object

* moved generate_tree function to utils.py; created generate_tree wrappers in enrichment results

* improved type checking for NaN values in MEA

* change named path to output_path

Author: harinlee83

CHANGELOG.md

@@ -1,5 +1,48 @@
 # Changelog
 
+## Unreleased
+
+Added
+
+- New feature for generating kinome tree using same logic as website. Original kinome tree SVG sourced from [CORAL](http://phanstiel-lab.med.unc.edu/CORAL/)
+
+## [1.5.0] - 2025-06-27
+
+Added
+
+- New feature for scoring all the possible phosphosites in an entire protein
+
+## [1.4.1] - 2025-06-19
+
+Added
+
+- Function to retrieve label mapping
+
+Changed
+
+- Label retrieval method in enrichment modules
+
+## [1.4.0] - 2025-06-18
+
+Added
+
+- New feature for displaying different label types in volcanoes: gene name, protein name, matrix name, or curated display name
+
+## [1.3.0] - 2025-05-13
+
+Added
+
+- New feature for running MEA with custom kinase-substrate sets
+- The kinase-substrate sets used in the MEA to the results
+
+Changed
+
+- Kinase name in custom sets to uppercase for consistency
+
+Fixed
+
+- Description mispellings
+
 ## [1.2.0] - 2025-04-15
 
 Added

src/kinase_library/databases/kinase_data/base_tree.svg

@@ -615,4 +615,40 @@ def plot_volcano(self, sig_lff=0, sig_pval=0.1, fg_percent_thresh=0, fg_percent_
                                            symmetric_xaxis=symmetric_xaxis, grid=grid, max_window=max_window,
                                            title=title, xlabel=xlabel, ylabel=ylabel,
                                             plot=plot, save_fig=save_fig, return_fig=return_fig,
-                                            ax=ax, **plot_kwargs)
+                                            ax=ax, **plot_kwargs)
+
+    def generate_tree(self, output_path, sort_by: str ='fisher_pval', sort_direction: str = 'ascending', filter_top: int = None, **kwargs):
+        """
+        Generate a colored kinome tree from the enrichment results.
+
+        Parameters
+        ----------
+        output_path : str
+            Destination path for the generated kinome tree image.
+        sort_by : str, optional
+            Column name to sort the DataFrame by before generating the tree. Default is 'fisher_pval'.
+        sort_direction : str, optional
+            Direction to sort the DataFrame, either 'ascending' or 'descending'. Default is 'ascending'.
+        filter_top : int, optional
+            If provided, only the top N rows will be included in the tree. Default is None (no filtering).
+        **kwargs : optional
+            Additional keyword arguments to be passed to the `utils.generate_tree` function.
+        """
+
+        # Check if the sort_by column exists in the enrichment results
+        if sort_by not in self.enrichment_results.columns:
+            raise ValueError(f"Column '{sort_by}' not found in enrichment results. Available columns: {self.enrichment_results.columns.tolist()}")
+
+        # Check if the sort_direction is valid
+        if sort_direction not in ['ascending', 'descending']:
+            raise ValueError("sort_direction must be either 'ascending' or 'descending'.")
+        
+    
+        # Sort the DataFrame based on the specified column and direction
+        df = self.enrichment_results.sort_values(by=sort_by, ascending=(sort_direction == 'ascending'))
+
+        if filter_top is not None:
+            df = df.head(filter_top)
+
+        # This kinome tree coloring will always be based on 'log2_freq_factor'
+        return utils.generate_tree(df, output_path, "log2_freq_factor", { "high": 3.0, "middle": 0.0, "low": -3.0 }, **kwargs)

src/kinase_library/enrichment/binary_enrichment.py

@@ -9,7 +9,7 @@
 import matplotlib.pyplot as plt
 import warnings
 
-from ..utils import _global_vars, exceptions
+from ..utils import _global_vars, exceptions, utils
 from ..modules import data, enrichment
 from ..objects import phosphoproteomics as pps
 from ..enrichment import binary_enrichment as be
@@ -899,4 +899,40 @@ def plot_down_up_comb_volcanos(self, sig_lff=0, sig_pval=0.1, adj_pval=True, kin
             plt.close(fig)
 
         if return_fig:
-            return fig
+            return fig
+
+    def generate_tree(self, output_path, sort_by: str ='most_sig_log2_freq_factor', sort_direction: str = 'ascending', filter_top: int = None, **kwargs):
+        """
+        Generate a colored kinome tree from the combined enrichment results.
+
+        Parameters
+        ----------
+        output_path : str
+            Destination path for the generated kinome tree image.
+        sort_by : str, optional
+            Column name to sort the DataFrame by before generating the tree. Default is 'most_sig_log2_freq_factor'.
+        sort_direction : str, optional
+            Direction to sort the DataFrame, either 'ascending' or 'descending'. Default is 'ascending'.
+        filter_top : int, optional
+            If provided, only the top N rows will be included in the tree. Default is None (no filtering).
+        **kwargs : optional
+            Additional keyword arguments to be passed to the `utils.generate_tree` function.
+        """
+
+        # Check if the sort_by column exists in the combined enrichment results
+        if sort_by not in self.combined_enrichment_results.columns:
+            raise ValueError(f"Column '{sort_by}' not found in combined enrichment results. Available columns: {self.combined_enrichment_results.columns.tolist()}")
+
+        # Check if the sort_direction is valid
+        if sort_direction not in ['ascending', 'descending']:
+            raise ValueError("sort_direction must be either 'ascending' or 'descending'.")
+        
+    
+        # Sort the DataFrame based on the specified column and direction
+        df = self.combined_enrichment_results.sort_values(by=sort_by, ascending=(sort_direction == 'ascending'))
+
+        if filter_top is not None:
+            df = df.head(filter_top)
+
+        # This kinome tree coloring will always be based on 'most_sig_log2_freq_factor'
+        return utils.generate_tree(df, output_path, "most_sig_log2_freq_factor", { "high": 3.0, "middle": 0.0, "low": -3.0 }, **kwargs)

src/kinase_library/enrichment/differential_phosphorylation.py

@@ -8,7 +8,7 @@
 import pandas as pd
 import gseapy as gp
 
-from ..utils import _global_vars, exceptions
+from ..utils import _global_vars, exceptions, utils
 from ..modules import data, enrichment
 from ..objects import phosphoproteomics as pps
 from ..logger import logger
@@ -578,4 +578,46 @@ def plot_volcano(self, sig_es=0, sig_pval=0.1, adj_pval=True, kinases=None,
                                            symmetric_xaxis=symmetric_xaxis, grid=grid, max_window=max_window,
                                            title=title, xlabel=xlabel, ylabel=ylabel,
                                            plot=plot, save_fig=save_fig, return_fig=return_fig,
-                                           ax=ax, **plot_kwargs)
+                                           ax=ax, **plot_kwargs)
+
+    def generate_tree(self, output_path, sort_by: str ='p-value', sort_direction: str = 'ascending', filter_top: int = None, **kwargs):
+        """
+        Generate a colored kinome tree from the enrichment results.
+
+        Parameters
+        ----------
+        output_path : str
+            Destination path for the generated kinome tree image.
+        sort_by : str, optional
+            Column name to sort the DataFrame by before generating the tree. Default is 'p-value'.
+        sort_direction : str, optional
+            Direction to sort the DataFrame, either 'ascending' or 'descending'. Default is 'ascending'.
+        filter_top : int, optional
+            If provided, only the top N rows will be included in the tree. Default is None (no filtering).
+        **kwargs : optional
+            Additional keyword arguments to be passed to the `utils.generate_tree` function.
+        """
+
+        # Check if the sort_by column exists in the enrichment results
+        if sort_by not in self.enrichment_results.columns:
+            raise ValueError(f"Column '{sort_by}' not found in enrichment results. Available columns: {self.enrichment_results.columns.tolist()}")
+
+        # Check if the sort_direction is valid
+        if sort_direction not in ['ascending', 'descending']:
+            raise ValueError("sort_direction must be either 'ascending' or 'descending'.")
+        
+    
+        # Sort the DataFrame based on the specified column and direction
+        df = self.enrichment_results.sort_values(by=sort_by, ascending=(sort_direction == 'ascending'))
+
+        if filter_top is not None:
+            df = df.head(filter_top)
+        
+        # Compute max and min NES values, ignoring NaN or non-numeric values
+        nes_values = df['NES']
+        nes_values = nes_values[pd.to_numeric(nes_values, errors='coerce').notnull()]
+        maxNes = nes_values.max()
+        minNes = nes_values.min()
+
+        # This kinome tree coloring will always be based on 'NES'
+        return utils.generate_tree(df, output_path, "NES", { "high": maxNes, "middle": 0.0, "low": minNes }, **kwargs)

src/kinase_library/enrichment/mea.py

@@ -4,6 +4,7 @@
 ##################################
 """
 import re
+import os
 import string
 import numpy as np
 import pandas as pd
@@ -607,4 +608,142 @@ def list_series_to_df(subs_list, col_name=None):
     if isinstance(subs_list, pd.Series):
             return(subs_list.to_frame(name=col_name))
     if isinstance(subs_list, list):
-            return(pd.Series(subs_list).to_frame(name=col_name))
+            return(pd.Series(subs_list).to_frame(name=col_name))
+
+def generate_tree(
+        kinase_matrix: pd.DataFrame,
+        output_path: str,
+        color_column: str,
+        color_thresholds: dict,
+        node_size: int = 5,
+        branch_color: str = "#663636",
+        low_color: str = "#999acf",
+        mid_color: str = "#c8c8c8",
+        high_color: str = "#fa6464",
+    ):
+        """
+        Generic function to generate a colored kinome tree. See DiffPhosEnrichmentResults.generate_tree(), MeaEnrichmentResults.generate_tree(), and EnrichmentResults.generate_tree() for specific implementations.
+
+        Parameters
+        ----------
+        kinase_matrix : pd.DataFrame
+            DataFrame containing kinases as indices and numerical columns to color the nodes. e.g. the output of kl.Substrate('PSVEPPLsQETFSDL').predict()
+        output_path : str
+            Path to save the tree image.
+        color_column : str
+            Column name in the kinase matrix to use for coloring the nodes.
+        color_thresholds : dict
+            Dictionary containing the color thresholds for low, middle, and high values. e.g. { "high": 3.0, "middle": 0.0, "low": -3.0 }.
+        node_size : int
+            Size of the nodes (SVG circles). Default is 5.
+        branch_color : str
+            Hex color for the tree branches. Default is "#663636".
+        low_color : str
+            Hex color for the low end of the heatmap. Default is "#999acf".
+        mid_color : str
+            Hex color for the midpoint of the heatmap. Default is "#c8c8c8".
+        high_color : str
+            Hex color for the high end of the heatmap. Default is "#fa6464".
+        """
+
+        def hex_to_rgb(hex_color):
+            """Convert hex string like '#FF0000' to (255, 0, 0)"""
+            hex_color = hex_color.lstrip('#')
+            return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
+
+
+        def rgb_to_hex(rgb):
+            """Convert (255, 0, 0) to '#FF0000'"""
+            return '#{:02X}{:02X}{:02X}'.format(*rgb)
+
+
+        def calculate_color(value, high_color_hex, low_color_hex, high_value, low_value):
+            high_color = hex_to_rgb(high_color_hex)
+            low_color = hex_to_rgb(low_color_hex)
+
+            new_value = max(min(value, high_value), low_value)
+            new_value -= low_value
+
+            percentage = new_value / (high_value - low_value) if high_value != low_value else 0
+
+            red = high_color[0] * percentage + low_color[0] * (1.0 - percentage)
+            green = high_color[1] * percentage + low_color[1] * (1.0 - percentage)
+            blue = high_color[2] * percentage + low_color[2] * (1.0 - percentage)
+
+            return rgb_to_hex((round(red), round(green), round(blue)))
+
+
+        def calculate_heatmap_color_midpoint(value, high_color, mid_color, low_color, high_value, mid_point_value, low_value):
+            if value < mid_point_value:
+                return calculate_color(value, mid_color, low_color, mid_point_value, low_value)
+            return calculate_color(value, high_color, mid_color, high_value, mid_point_value)
+
+
+        def calculate_heatmap_color(value, high_value, mid_point_value, low_value, high_color, mid_color, low_color):
+            return calculate_heatmap_color_midpoint(
+                value,
+                high_color,
+                mid_color,
+                low_color,
+                high_value,
+                mid_point_value,
+                low_value
+            )
+
+        # Check if output_path is valid
+        if not isinstance(output_path, str) or not output_path.endswith('.svg'):
+            raise ValueError("Output path must be a valid string ending with '.svg'.")
+
+        if kinase_matrix.get(color_column, None) is None:
+            raise ValueError(f"Column '{color_column}' not found in the kinase matrix. Please provide a valid column name.")
+
+        # Check if color thresholds are valid
+        if not all(key in color_thresholds for key in ["high", "middle", "low"]):
+            raise ValueError("Color thresholds must contain 'high', 'middle', and 'low' keys.")
+
+        kinases = kinase_matrix.index
+        
+        # Create map for quick access to kinase names
+        kinase_uniprot_mapping = {
+            row['UNIPROT_ID']: row['MATRIX_NAME']
+            for _, row in data.get_kinome_info().iterrows()
+        }
+
+        # Load the base SVG
+        import xml.etree.ElementTree as ET
+        current_dir = os.path.dirname(__file__)
+        svg_path = os.path.abspath(os.path.join(current_dir, "../databases/kinase_data/base_tree.svg"))
+        tree = ET.parse(svg_path)
+        root = tree.getroot()
+
+        # Set the branch color if not default
+        if branch_color != "#663636":
+            for line in root.findall('.//svg:path', namespaces={'svg': 'http://www.w3.org/2000/svg'}):
+                line.set('fill', branch_color)
+
+        ns = {'svg': 'http://www.w3.org/2000/svg'}
+
+        # For the top X kinases in kinase matrix, set the opacity to 1
+        for circle in root.findall('.//svg:circle', namespaces=ns):
+            uniprot_id = circle.get('class').split("_")[-1]
+            if kinase_uniprot_mapping.get(uniprot_id, None) in kinases:
+                circle.set('opacity', '1')
+
+        for circle in root.findall('.//svg:circle', namespaces=ns):
+            uniprot_id = circle.get('class').split("_")[-1]
+            if kinase_uniprot_mapping.get(uniprot_id, None) in kinases:
+                val = kinase_matrix.at[kinase_uniprot_mapping.get(uniprot_id, None), color_column]
+                if pd.notna(val) and isinstance(val, (int, float, np.integer, np.floating)):
+                    color = calculate_heatmap_color(val, color_thresholds["high"], color_thresholds["middle"], color_thresholds["low"],
+                                                    high_color, mid_color, low_color)
+                    circle.set('fill', color)
+                    circle.set('stroke', "gray")
+                    circle.set('stroke-width', "0.5px")
+                    circle.set('opacity', '1')
+                    circle.set('r', str(node_size))
+
+        try:
+            os.makedirs(os.path.dirname(output_path), exist_ok=True)
+            tree.write(output_path)
+        except Exception as e:
+            raise Exception(f"Error saving SVG file: {e}")