gh-140677 Improve heatmap colors (#142241)

Co-authored-by: Pablo Galindo Salgado <pablogsal@gmail.com>

Lib/profiling/sampling/_heatmap_assets/heatmap.css

@@ -1094,18 +1094,34 @@ #scroll_marker .marker {
 }
 
 #scroll_marker .marker.cold {
+  background: var(--heat-1);
+}
+
+#scroll_marker .marker.cool {
   background: var(--heat-2);
 }
 
+#scroll_marker .marker.mild {
+  background: var(--heat-3);
+}
+
 #scroll_marker .marker.warm {
-  background: var(--heat-5);
+  background: var(--heat-4);
 }
 
 #scroll_marker .marker.hot {
+  background: var(--heat-5);
+}
+
+#scroll_marker .marker.very-hot {
+  background: var(--heat-6);
+}
+
+#scroll_marker .marker.intense {
   background: var(--heat-7);
 }
 
-#scroll_marker .marker.vhot {
+#scroll_marker .marker.extreme {
   background: var(--heat-8);
 }
 

Lib/profiling/sampling/_heatmap_assets/heatmap.js

@@ -26,6 +26,7 @@ function toggleTheme() {
     if (btn) {
         btn.innerHTML = next === 'dark' ? '☼' : '☾';  // sun or moon
     }
+    applyLineColors();
 
     // Rebuild scroll marker with new theme colors
     buildScrollMarker();
@@ -160,13 +161,6 @@ function getSampleCount(line) {
     return parseInt(text) || 0;
 }
 
-function getIntensityClass(ratio) {
-    if (ratio > 0.75) return 'vhot';
-    if (ratio > 0.5) return 'hot';
-    if (ratio > 0.25) return 'warm';
-    return 'cold';
-}
-
 // ============================================================================
 // Scroll Minimap
 // ============================================================================
@@ -194,7 +188,7 @@ function buildScrollMarker() {
 
         const lineTop = Math.floor(line.offsetTop * markerScale);
         const lineNumber = index + 1;
-        const intensityClass = maxSamples > 0 ? getIntensityClass(samples / maxSamples) : 'cold';
+        const intensityClass = maxSamples > 0 ? (intensityToClass(samples / maxSamples) || 'cold') : 'cold';
 
         if (lineNumber === prevLine + 1 && lastMark?.classList.contains(intensityClass)) {
             lastMark.style.height = `${lineTop + lineHeight - lastTop}px`;
@@ -212,6 +206,21 @@ function buildScrollMarker() {
     document.body.appendChild(scrollMarker);
 }
 
+function applyLineColors() {
+    const lines = document.querySelectorAll('.code-line');
+    lines.forEach(line => {
+        let intensity;
+        if (colorMode === 'self') {
+            intensity = parseFloat(line.getAttribute('data-self-intensity')) || 0;
+        } else {
+            intensity = parseFloat(line.getAttribute('data-cumulative-intensity')) || 0;
+        }
+
+        const color = intensityToColor(intensity);
+        line.style.background = color;
+    });
+}
+
 // ============================================================================
 // Toggle Controls
 // ============================================================================
@@ -264,20 +273,7 @@ function applyHotFilter() {
 
 function toggleColorMode() {
     colorMode = colorMode === 'self' ? 'cumulative' : 'self';
-    const lines = document.querySelectorAll('.code-line');
+    applyLineColors();
-
-    lines.forEach(line => {
-        let bgColor;
-        if (colorMode === 'self') {
-            bgColor = line.getAttribute('data-self-color');
-        } else {
-            bgColor = line.getAttribute('data-cumulative-color');
-        }
-
-        if (bgColor) {
-            line.style.background = bgColor;
-        }
-    });
 
     updateToggleUI('toggle-color-mode', colorMode === 'cumulative');
 
@@ -295,14 +291,7 @@ function toggleColorMode() {
 document.addEventListener('DOMContentLoaded', function() {
     // Restore UI state (theme, etc.)
     restoreUIState();
-
+    applyLineColors();
-    // Apply background colors
-    document.querySelectorAll('.code-line[data-bg-color]').forEach(line => {
-        const bgColor = line.getAttribute('data-bg-color');
-        if (bgColor) {
-            line.style.background = bgColor;
-        }
-    });
 
     // Initialize navigation buttons
     document.querySelectorAll('.nav-btn').forEach(button => {

Lib/profiling/sampling/_heatmap_assets/heatmap_index.js

@@ -1,6 +1,19 @@
 // Tachyon Profiler - Heatmap Index JavaScript
 // Index page specific functionality
 
+// ============================================================================
+// Heatmap Bar Coloring
+// ============================================================================
+
+function applyHeatmapBarColors() {
+    const bars = document.querySelectorAll('.heatmap-bar[data-intensity]');
+    bars.forEach(bar => {
+        const intensity = parseFloat(bar.getAttribute('data-intensity')) || 0;
+        const color = intensityToColor(intensity);
+        bar.style.backgroundColor = color;
+    });
+}
+
 // ============================================================================
 // Theme Support
 // ============================================================================
@@ -17,6 +30,8 @@ function toggleTheme() {
     if (btn) {
         btn.innerHTML = next === 'dark' ? '☼' : '☾';  // sun or moon
     }
+
+    applyHeatmapBarColors();
 }
 
 function restoreUIState() {
@@ -108,4 +123,5 @@ function collapseAll() {
 
 document.addEventListener('DOMContentLoaded', function() {
     restoreUIState();
+    applyHeatmapBarColors();
 });

Lib/profiling/sampling/_heatmap_assets/heatmap_shared.js

@@ -0,0 +1,40 @@
+// Tachyon Profiler - Shared Heatmap JavaScript
+// Common utilities shared between index and file views
+
+// ============================================================================
+// Heat Level Mapping (Single source of truth for intensity thresholds)
+// ============================================================================
+
+// Maps intensity (0-1) to heat level (0-8). Level 0 = no heat, 1-8 = heat levels.
+function intensityToHeatLevel(intensity) {
+    if (intensity <= 0) return 0;
+    if (intensity <= 0.125) return 1;
+    if (intensity <= 0.25) return 2;
+    if (intensity <= 0.375) return 3;
+    if (intensity <= 0.5) return 4;
+    if (intensity <= 0.625) return 5;
+    if (intensity <= 0.75) return 6;
+    if (intensity <= 0.875) return 7;
+    return 8;
+}
+
+// Class names corresponding to heat levels 1-8 (used by scroll marker)
+const HEAT_CLASS_NAMES = ['cold', 'cool', 'mild', 'warm', 'hot', 'very-hot', 'intense', 'extreme'];
+
+function intensityToClass(intensity) {
+    const level = intensityToHeatLevel(intensity);
+    return level === 0 ? null : HEAT_CLASS_NAMES[level - 1];
+}
+
+// ============================================================================
+// Color Mapping (Intensity to Heat Color)
+// ============================================================================
+
+function intensityToColor(intensity) {
+    const level = intensityToHeatLevel(intensity);
+    if (level === 0) {
+        return 'transparent';
+    }
+    const rootStyle = getComputedStyle(document.documentElement);
+    return rootStyle.getPropertyValue(`--heat-${level}`).trim();
+}

Lib/profiling/sampling/heatmap_collector.py

@@ -5,6 +5,7 @@
 import html
 import importlib.resources
 import json
+import math
 import os
 import platform
 import site
@@ -44,31 +45,6 @@ class TreeNode:
     children: Dict[str, 'TreeNode'] = field(default_factory=dict)
 
 
-@dataclass
-class ColorGradient:
-    """Configuration for heatmap color gradient calculations."""
-    # Color stops thresholds
-    stop_1: float = 0.2  # Blue to cyan transition
-    stop_2: float = 0.4  # Cyan to green transition
-    stop_3: float = 0.6  # Green to yellow transition
-    stop_4: float = 0.8  # Yellow to orange transition
-    stop_5: float = 1.0  # Orange to red transition
-
-    # Alpha (opacity) values
-    alpha_very_cold: float = 0.3
-    alpha_cold: float = 0.4
-    alpha_medium: float = 0.5
-    alpha_warm: float = 0.6
-    alpha_hot_base: float = 0.7
-    alpha_hot_range: float = 0.15
-
-    # Gradient multiplier
-    multiplier: int = 5
-
-    # Cache for calculated colors
-    cache: Dict[float, Tuple[int, int, int, float]] = field(default_factory=dict)
-
-
 # ============================================================================
 # Module Path Analysis
 # ============================================================================
@@ -224,8 +200,9 @@ def _load_templates(self):
             self.file_css = css_content
 
             # Load JS
-            self.index_js = (assets_dir / "heatmap_index.js").read_text(encoding="utf-8")
+            shared_js = (assets_dir / "heatmap_shared.js").read_text(encoding="utf-8")
-            self.file_js = (assets_dir / "heatmap.js").read_text(encoding="utf-8")
+            self.index_js = f"{shared_js}\n{(assets_dir / 'heatmap_index.js').read_text(encoding='utf-8')}"
+            self.file_js = f"{shared_js}\n{(assets_dir / 'heatmap.js').read_text(encoding='utf-8')}"
 
             # Load Python logo
             logo_dir = template_dir / "_assets"
@@ -321,18 +298,13 @@ def _calculate_node_stats(node: TreeNode) -> Tuple[int, int]:
 class _HtmlRenderer:
     """Renders hierarchical tree structures as HTML."""
 
-    def __init__(self, file_index: Dict[str, str], color_gradient: ColorGradient,
+    def __init__(self, file_index: Dict[str, str]):
-                 calculate_intensity_color_func):
+        """Initialize renderer with file index.
-        """Initialize renderer with file index and color calculation function.
 
         Args:
             file_index: Mapping from filenames to HTML file names
-            color_gradient: ColorGradient configuration
-            calculate_intensity_color_func: Function to calculate colors
         """
         self.file_index = file_index
-        self.color_gradient = color_gradient
-        self.calculate_intensity_color = calculate_intensity_color_func
         self.heatmap_bar_height = 16
 
     def render_hierarchical_html(self, trees: Dict[str, TreeNode]) -> str:
@@ -450,8 +422,6 @@ def _render_file_item(self, stat: FileStats, indent: str = '') -> str:
         module_name = html.escape(stat.module_name)
 
         intensity = stat.percentage / 100.0
-        r, g, b, alpha = self.calculate_intensity_color(intensity)
-        bg_color = f"rgba({r}, {g}, {b}, {alpha})"
         bar_width = min(stat.percentage, 100)
 
         html_file = self.file_index[stat.filename]
@@ -459,7 +429,7 @@ def _render_file_item(self, stat: FileStats, indent: str = '') -> str:
         return (f'{indent}<div class="file-item">\n'
                 f'{indent}  <a href="{html_file}" class="file-link" title="{full_path}">📄 {module_name}</a>\n'
                 f'{indent}  <span class="file-samples">{stat.total_samples:,} samples</span>\n'
-                f'{indent}  <div class="heatmap-bar-container"><div class="heatmap-bar" style="width: {bar_width}px; background-color: {bg_color}; height: {self.heatmap_bar_height}px;"></div></div>\n'
+                f'{indent}  <div class="heatmap-bar-container"><div class="heatmap-bar" style="width: {bar_width}px; height: {self.heatmap_bar_height}px;" data-intensity="{intensity:.3f}"></div></div>\n'
                 f'{indent}</div>\n')
 
 
@@ -501,20 +471,12 @@ def __init__(self, *args, **kwargs):
         self._path_info = get_python_path_info()
         self.stats = {}
 
-        # Color gradient configuration
-        self._color_gradient = ColorGradient()
-
         # Template loader (loads all templates once)
         self._template_loader = _TemplateLoader()
 
         # File index (populated during export)
         self.file_index = {}
 
-    @property
-    def _color_cache(self):
-        """Compatibility property for accessing color cache."""
-        return self._color_gradient.cache
-
     def set_stats(self, sample_interval_usec, duration_sec, sample_rate, error_rate=None, missed_samples=None, **kwargs):
         """Set profiling statistics to include in heatmap output.
 
@@ -746,8 +708,7 @@ def _generate_index_html(self, index_path: Path, file_stats: List[FileStats]):
         tree = _TreeBuilder.build_file_tree(file_stats)
 
         # Render tree as HTML
-        renderer = _HtmlRenderer(self.file_index, self._color_gradient,
+        renderer = _HtmlRenderer(self.file_index)
-                                self._calculate_intensity_color)
         sections_html = renderer.render_hierarchical_html(tree)
 
         # Format error rate and missed samples with bar classes
@@ -809,56 +770,6 @@ def _generate_index_html(self, index_path: Path, file_stats: List[FileStats]):
         except (IOError, OSError) as e:
             raise RuntimeError(f"Failed to write index file {index_path}: {e}") from e
 
-    def _calculate_intensity_color(self, intensity: float) -> Tuple[int, int, int, float]:
-        """Calculate RGB color and alpha for given intensity (0-1 range).
-
-        Returns (r, g, b, alpha) tuple representing the heatmap color gradient:
-        blue -> green -> yellow -> orange -> red
-
-        Results are cached to improve performance.
-        """
-        # Round to 3 decimal places for cache key
-        cache_key = round(intensity, 3)
-        if cache_key in self._color_gradient.cache:
-            return self._color_gradient.cache[cache_key]
-
-        gradient = self._color_gradient
-        m = gradient.multiplier
-
-        # Color stops with (threshold, rgb_func, alpha_func)
-        stops = [
-            (gradient.stop_1,
-             lambda i: (0, int(150 * i * m), 255),
-             lambda i: gradient.alpha_very_cold),
-            (gradient.stop_2,
-             lambda i: (0, 255, int(255 * (1 - (i - gradient.stop_1) * m))),
-             lambda i: gradient.alpha_cold),
-            (gradient.stop_3,
-             lambda i: (int(255 * (i - gradient.stop_2) * m), 255, 0),
-             lambda i: gradient.alpha_medium),
-            (gradient.stop_4,
-             lambda i: (255, int(200 - 100 * (i - gradient.stop_3) * m), 0),
-             lambda i: gradient.alpha_warm),
-            (gradient.stop_5,
-             lambda i: (255, int(100 * (1 - (i - gradient.stop_4) * m)), 0),
-             lambda i: gradient.alpha_hot_base + gradient.alpha_hot_range * (i - gradient.stop_4) * m),
-        ]
-
-        result = None
-        for threshold, rgb_func, alpha_func in stops:
-            if intensity < threshold or threshold == gradient.stop_5:
-                r, g, b = rgb_func(intensity)
-                result = (r, g, b, alpha_func(intensity))
-                break
-
-        # Fallback
-        if result is None:
-            result = (255, 0, 0, 0.75)
-
-        # Cache the result
-        self._color_gradient.cache[cache_key] = result
-        return result
-
     def _generate_file_html(self, output_path: Path, filename: str,
                           line_counts: Dict[int, int], self_counts: Dict[int, int],
                           file_stat: FileStats):
@@ -913,25 +824,23 @@ def _build_line_html(self, line_num: int, line_content: str,
 
         # Calculate colors for both self and cumulative modes
         if cumulative_samples > 0:
-            cumulative_intensity = cumulative_samples / max_samples if max_samples > 0 else 0
+            log_cumulative = math.log(cumulative_samples + 1)
-            self_intensity = self_samples / max_self_samples if max_self_samples > 0 and self_samples > 0 else 0
+            log_max = math.log(max_samples + 1)
+            cumulative_intensity = log_cumulative / log_max if log_max > 0 else 0
 
-            # Default to self-based coloring
+            if self_samples > 0 and max_self_samples > 0:
-            intensity = self_intensity if self_samples > 0 else cumulative_intensity
+                log_self = math.log(self_samples + 1)
-            r, g, b, alpha = self._calculate_intensity_color(intensity)
+                log_max_self = math.log(max_self_samples + 1)
-            bg_color = f"rgba({r}, {g}, {b}, {alpha})"
+                self_intensity = log_self / log_max_self if log_max_self > 0 else 0
-
+            else:
-            # Pre-calculate colors for both modes (for JS toggle)
+                self_intensity = 0
-            self_bg_color = self._format_color_for_intensity(self_intensity) if self_samples > 0 else "transparent"
-            cumulative_bg_color = self._format_color_for_intensity(cumulative_intensity)
 
             self_display = f"{self_samples:,}" if self_samples > 0 else ""
             cumulative_display = f"{cumulative_samples:,}"
             tooltip = f"Self: {self_samples:,}, Total: {cumulative_samples:,}"
         else:
-            bg_color = "transparent"
+            cumulative_intensity = 0
-            self_bg_color = "transparent"
+            self_intensity = 0
-            cumulative_bg_color = "transparent"
             self_display = ""
             cumulative_display = ""
             tooltip = ""
@@ -939,13 +848,14 @@ def _build_line_html(self, line_num: int, line_content: str,
         # Get navigation buttons
         nav_buttons_html = self._build_navigation_buttons(filename, line_num)
 
-        # Build line HTML
+        # Build line HTML with intensity data attributes
         line_html = html.escape(line_content.rstrip('\n'))
         title_attr = f' title="{html.escape(tooltip)}"' if tooltip else ""
 
         return (
-            f'        <div class="code-line" data-bg-color="{bg_color}" '
+            f'        <div class="code-line" '
-            f'data-self-color="{self_bg_color}" data-cumulative-color="{cumulative_bg_color}" '
+            f'data-self-intensity="{self_intensity:.3f}" '
+            f'data-cumulative-intensity="{cumulative_intensity:.3f}" '
             f'id="line-{line_num}"{title_attr}>\n'
             f'            <div class="line-number">{line_num}</div>\n'
             f'            <div class="line-samples-self">{self_display}</div>\n'
@@ -955,11 +865,6 @@ def _build_line_html(self, line_num: int, line_content: str,
             f'        </div>\n'
         )
 
-    def _format_color_for_intensity(self, intensity: float) -> str:
-        """Format color as rgba() string for given intensity."""
-        r, g, b, alpha = self._calculate_intensity_color(intensity)
-        return f"rgba({r}, {g}, {b}, {alpha})"
-
     def _build_navigation_buttons(self, filename: str, line_num: int) -> str:
         """Build navigation buttons for callers/callees."""
         line_key = (filename, line_num)

Lib/test/test_profiling/test_heatmap.py

@@ -147,12 +147,6 @@ def test_init_sets_total_samples_to_zero(self):
         collector = HeatmapCollector(sample_interval_usec=100)
         self.assertEqual(collector._total_samples, 0)
 
-    def test_init_creates_color_cache(self):
-        """Test that color cache is initialized."""
-        collector = HeatmapCollector(sample_interval_usec=100)
-        self.assertIsInstance(collector._color_cache, dict)
-        self.assertEqual(len(collector._color_cache), 0)
-
     def test_init_gets_path_info(self):
         """Test that path info is retrieved during init."""
         collector = HeatmapCollector(sample_interval_usec=100)