hanhou/patchseq

1"""2Color mapping utilities for the scatter plot.3"""4 5from typing import Any, Dict, Union6 7import numpy as np8import pandas as pd9from bokeh.models import CategoricalColorMapper, ColorBar, LinearColorMapper10from bokeh.palettes import all_palettes11from bokeh.plotting import figure12 13from LCNE_patchseq_analysis import REGION_COLOR_MAPPER14 15 16class ColorMapping:17    """Handles color mapping for scatter plots."""18 19    def __init__(self, df_meta: pd.DataFrame, font_size: int = 14):20        """Initialize with metadata dataframe."""21        self.df_meta = df_meta22        self.font_size = font_size23 24    def add_color_bar(25        self,26        color_mapper: Union[CategoricalColorMapper, LinearColorMapper],27        title: str,28        p: figure,29        font_size: int = 14,30    ) -> ColorBar:31        """Add a color bar to the plot with consistent styling."""32        color_bar = ColorBar(33            color_mapper=color_mapper,34            label_standoff=12,35            border_line_color=None,36            location=(0, 0),37            title=title,38            title_text_font_size=f"{font_size*0.8}pt",39            major_label_text_font_size=f"{font_size*0.8}pt",40        )41        p.add_layout(color_bar, "right")42        return color_bar43 44    def determine_color_mapping(  # noqa: C90145        self, color_mapping: str, color_palette: Any, p: figure, font_size: int = 14, if_add_color_bar: bool = True46    ) -> Dict[str, Any]:47        """48        Determine the color mapping for the scatter plot.49 50        Args:51            color_mapping: Column name to use for color mapping52            color_palette: Color palette to use53            p: Bokeh figure to add color bar to54 55        Returns:56            Dictionary with field and transform for scatter plot57        """58        if color_mapping == "injection region":59            color_mapper = {60                key: value61                for key, value in REGION_COLOR_MAPPER.items()62                if key in self.df_meta["injection region"].unique()63            }64            color_mapper = CategoricalColorMapper(65                factors=list(color_mapper.keys()), palette=list(color_mapper.values())66            )67 68            # Add a color bar for categorical data69            if if_add_color_bar:70                self.add_color_bar(color_mapper, color_mapping, p, font_size)71 72            return {"field": color_mapping, "transform": color_mapper}73 74        # If categorical (nunique <= 50), use categorical color mapper75        if self.df_meta[color_mapping].nunique() <= 50:76            n_categories = self.df_meta[color_mapping].nunique()77 78            # Check if the provided color_palette is a string (name in all_palettes)79            if isinstance(color_palette, str) and color_palette in all_palettes:80                # Use the named palette from all_palettes81                # Check if the number of categories is supported by the palette82                max_colors = max(all_palettes[color_palette].keys())83                n_colors = min(n_categories, max_colors)84                categorical_palette = all_palettes[color_palette][n_colors]85                # Make the palette circular by cycling through the colors86                if n_categories > n_colors:87                    categorical_palette = list(categorical_palette)88                    categorical_palette = [89                        categorical_palette[i % n_colors] for i in range(n_categories)90                    ]91            else:92                # For continuous palette lists or palette names not in all_palettes93                # First check if it's a named continuous palette like 'Viridis256'94                named_palette = None95                for name, palettes in all_palettes.items():96                    if isinstance(color_palette, str) and color_palette.startswith(name):97                        # If we found the palette, use it98                        if 256 in palettes:99                            named_palette = palettes[256]100                        elif len(palettes) > 0:101                            # Get the largest available palette102                            max_key = max(palettes.keys())103                            named_palette = palettes[max_key]104                        break105 106                # Use the named continuous palette or the provided palette107                if named_palette is not None:108                    continuous_palette = named_palette109                else:110                    continuous_palette = color_palette111 112                # Uniformly sample from the continuous colormap113                if isinstance(continuous_palette, (list, tuple)):114                    indices = np.linspace(0, len(continuous_palette) - 1, n_categories).astype(int)115                    categorical_palette = [continuous_palette[i] for i in indices]116                else:117                    # If we can't determine the palette type, use a default118                    categorical_palette = all_palettes["Category10"][min(n_categories, 10)]119 120            # Map "None" or NaN factors to "gray"121            # Sort factors alphabetically for consistent color assignment122            factors = sorted(list(self.df_meta[color_mapping].dropna().unique()), reverse=True)123            categorical_palette = list(categorical_palette)124            125            # Ensure we have enough colors for all factors126            if len(categorical_palette) < len(factors):127                # Extend palette by cycling through colors if needed128                categorical_palette = categorical_palette * ((len(factors) // len(categorical_palette)) + 1)129            categorical_palette = categorical_palette[:len(factors)]130            131            for missing in ["None", "unknown", "seq_data_not_available"]:132                if missing in factors:133                    categorical_palette[factors.index(missing)] = "gray"134 135            color_mapper = CategoricalColorMapper(136                factors=factors,137                palette=categorical_palette,138            )139            if if_add_color_bar:140                # Add a color bar for categorical data141                self.add_color_bar(color_mapper, color_mapping, p)142            return {"field": color_mapping, "transform": color_mapper}143 144        # Try to convert the column to numeric145        numeric_data = pd.Series(pd.to_numeric(self.df_meta[color_mapping], errors="coerce"))146        if not numeric_data.isna().all():147            # If conversion is successful, use linear color mapper148            low = numeric_data.quantile(0.01)149            high = numeric_data.quantile(0.99)150            color_mapper = LinearColorMapper(palette=color_palette, low=low, high=high)151            color = {"field": color_mapping, "transform": color_mapper}152 153            # Add a color bar154            if if_add_color_bar:155                self.add_color_bar(color_mapper, color_mapping, p)156            return color157 158        return "black"159