Improve Heatmap

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import pandas as pd
 import numpy as np
 from bokeh.plotting import figure
-from bokeh.models import ColumnDataSource, DataTable, TableColumn, CustomJS, Select, Button, HoverTool, LinearColorMapper, ColorBar
+from bokeh.models import ColumnDataSource, DataTable, TableColumn, CustomJS, Select, Button, HoverTool, LinearColorMapper, ColorBar, FuncTickFormatter
 from bokeh.layouts import column
 from bokeh.palettes import Reds9, Blues9, Oranges9, Purples9, Greys9, BuGn9, Greens9
 from sklearn.decomposition import PCA
@@ -12,6 +12,7 @@ import io
 import ot
 from sklearn.linear_model import LinearRegression
 from scipy.stats import binned_statistic_2d
+import json
 
 
 N_COMPONENTS = 2
@@ -1059,11 +1060,10 @@ def run_model(model_name):
         # --- INICIO DEL BLOQUE: Heatmap de características ---
         st.markdown("## Heatmap de Características")
 
-        # Cargar el CSV que contiene las características para el heatmap.
-        # Se asume que el archivo se encuentra en "data/heatmaps.csv"
         try:
-            df_heat = pd.read_csv("data/heatmaps.csv", sep=";")
-            df_heat.columns = [col.strip("'\"") for col in df_heat.columns]
+            df_heat = pd.read_csv("data/heatmaps.csv")
+            # Si fuera necesario, se pueden limpiar los nombres de las columnas:
+            # df_heat.columns = [col.strip("'\"") for col in df_heat.columns]
         except Exception as e:
             st.error(f"Error al cargar heatmaps.csv: {e}")
             df_heat = None
@@ -1073,19 +1073,18 @@ def run_model(model_name):
             if 'img' not in df_all["real"].columns:
                 st.error("La columna 'img' no se encuentra en las muestras reales para hacer el merge con heatmaps.csv.")
             else:
-                # Crear la columna 'name' extrayendo el nombre sin la extensión '.png'
+                # Crear la columna 'name' extrayendo el nombre final de la URL y removiendo ".png"
                 df_all["real"]["name"] = df_all["real"]["img"].apply(
                     lambda x: x.split("/")[-1].replace(".png", "") if isinstance(x, str) else x
                 )
-                print(df_all["real"]["name"])
-                print(df_heat)
                 
                 # Hacemos merge de las posiciones reales con el CSV de heatmaps usando la columna 'name'
                 df_heatmap = pd.merge(df_all["real"], df_heat, on="name", how="inner")
                 
                 # Extraemos las características disponibles (excluyendo 'name')
                 feature_options = [col for col in df_heat.columns if col != "name"]
-                selected_feature = st.selectbox("Seleccione la característica para el heatmap:", options=feature_options)
+                selected_feature = st.selectbox("Seleccione la característica para el heatmap:", 
+                                                options=feature_options, key=f"heatmap_{model_name}")
                 
                 # Determinar el rango de las posiciones (x, y) de las muestras reales
                 x_min, x_max = df_heatmap['x'].min(), df_heatmap['x'].max()
@@ -1096,31 +1095,62 @@ def run_model(model_name):
                 x_bins = np.linspace(x_min, x_max, grid_size + 1)
                 y_bins = np.linspace(y_min, y_max, grid_size + 1)
                 
-                # Utilizamos binned_statistic_2d para calcular la media de la característica seleccionada en cada celda
-                heat_stat, x_edges, y_edges, binnumber = binned_statistic_2d(
-                    df_heatmap['x'], df_heatmap['y'], df_heatmap[selected_feature],
-                    statistic='mean', bins=[x_bins, y_bins]
-                )
+                # Si la variable seleccionada no es numérica, la convertimos a códigos numéricos
+                # y guardamos la correspondencia para la leyenda.
+                cat_mapping = None
+                if not pd.api.types.is_numeric_dtype(df_heatmap[selected_feature]):
+                    cat = df_heatmap[selected_feature].astype('category')
+                    cat_mapping = list(cat.cat.categories)
+                    df_heatmap[selected_feature] = cat.cat.codes
+                
+                # Intentamos calcular el heatmap; si falla, aplicamos la conversión a categoría
+                try:
+                    heat_stat, x_edges, y_edges, binnumber = binned_statistic_2d(
+                        df_heatmap['x'], df_heatmap['y'], df_heatmap[selected_feature],
+                        statistic='mean', bins=[x_bins, y_bins]
+                    )
+                except:
+                    cat = df_heatmap[selected_feature].astype('category')
+                    cat_mapping = list(cat.cat.categories)
+                    df_heatmap[selected_feature] = cat.cat.codes
+                    heat_stat, x_edges, y_edges, binnumber = binned_statistic_2d(
+                        df_heatmap['x'], df_heatmap['y'], df_heatmap[selected_feature],
+                        statistic='mean', bins=[x_bins, y_bins]
+                    )
                 
-                # La función image de Bokeh espera una lista de arrays y el arreglo debe estar orientado correctamente
-                heatmap_data = heat_stat.T  # Transponemos para alinear los ejes
+                # La función image de Bokeh espera una lista de arrays; se transpone para alinear los ejes.
+                heatmap_data = heat_stat.T
                 
                 # Crear el mapa de color
                 color_mapper = LinearColorMapper(palette="Viridis256", low=np.nanmin(heatmap_data), high=np.nanmax(heatmap_data))
                 
-                # Crear la figura para el heatmap
+                # Crear la figura para el heatmap con fondo blanco
                 heatmap_fig = figure(title=f"Heatmap de '{selected_feature}'",
                                     x_range=(x_min, x_max), y_range=(y_min, y_max),
                                     width=600, height=600,
                                     tools="pan,wheel_zoom,reset,save", active_scroll="wheel_zoom")
+                heatmap_fig.background_fill_color = "white"
                 
-                # Dibujar el heatmap usando la imagen (se pasa la lista con el array de datos)
+                # Dibujar el heatmap usando la imagen
                 heatmap_fig.image(image=[heatmap_data], x=x_min, y=y_min,
                                 dw=x_max - x_min, dh=y_max - y_min,
                                 color_mapper=color_mapper)
                 
-                # Agregar barra de colores
+                # Crear la barra de colores
                 color_bar = ColorBar(color_mapper=color_mapper, location=(0, 0))
+                # Si se usó conversión a categoría, formateamos la barra para mostrar las etiquetas originales
+                if cat_mapping is not None:
+
+                    categories_json = json.dumps(cat_mapping)
+                    color_bar.formatter = FuncTickFormatter(code=f"""
+                        var categories = {categories_json};
+                        var index = Math.round(tick);
+                        if(index >= 0 && index < categories.length) {{
+                            return categories[index];
+                        }} else {{
+                            return "";
+                        }}
+                    """)
                 heatmap_fig.add_layout(color_bar, 'right')
                 
                 st.bokeh_chart(heatmap_fig)