Add array heatmap logging to ViewerGL

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Added
 
+- Add heatmap rendering for scalar arrays logged through `ViewerGL.log_array()`
 - Add `SolverXPBD.update_contacts()` to populate `contacts.force` with per-contact spatial forces (linear force and torque) derived from XPBD constraint impulses
 - Raise process priority automatically in `--benchmark` mode for more stable measurements; add `--realtime` for maximum priority.
 - Import per-shape authored color from USD stages into `ModelBuilder.shape_color`

newton/_src/viewer/viewer_gl.py

@@ -220,6 +220,12 @@ def __init__(
         self._scalar_arrays: dict[str, np.ndarray | None] = {}
         self._scalar_accumulators: dict[str, list[float]] = {}
         self._scalar_smoothing: dict[str, int] = {}
+        self._array_buffers: dict[str, np.ndarray] = {}
+        self._array_dirty: set[str] = set()
+        self._array_textures: dict[str, dict[str, Any]] = {}
+        self._heatmap_min_cell_pixels = 3.0
+        self._heatmap_nan_rgba = np.array([51, 51, 51, 255], dtype=np.uint8)
+        self._heatmap_color_lut = self._build_heatmap_color_lut()
         self._plot_history_size = plot_history_size
 
         super().__init__()
@@ -313,6 +319,30 @@ def _invalidate_pbo(self):
             gl.glDeleteBuffers(1, pbo_id)
             self._pbo = None
 
+    def _delete_array_texture(self, name: str):
+        texture_state = self._array_textures.pop(name, None)
+        if texture_state is None:
+            return
+        gl = getattr(RendererGL, "gl", None)
+        texture_id = texture_state.get("texture_id")
+        if gl is None or texture_id is None:
+            return
+        texture_ids = (gl.GLuint * 1)(texture_id)
+        gl.glDeleteTextures(1, texture_ids)
+
+    def _clear_array_textures(self):
+        if not self._array_textures:
+            return
+        gl = getattr(RendererGL, "gl", None)
+        if gl is None:
+            self._array_textures.clear()
+            return
+        texture_ids = [state["texture_id"] for state in self._array_textures.values() if state.get("texture_id")]
+        if texture_ids:
+            gl_ids = (gl.GLuint * len(texture_ids))(*texture_ids)
+            gl.glDeleteTextures(len(texture_ids), gl_ids)
+        self._array_textures.clear()
+
     def register_ui_callback(
         self,
         callback: Callable[[Any], None],
@@ -445,6 +475,9 @@ def clear_model(self):
         self._scalar_arrays.clear()
         self._scalar_accumulators.clear()
         self._scalar_smoothing.clear()
+        self._array_buffers.clear()
+        self._array_dirty.clear()
+        self._clear_array_textures()
 
         super().clear_model()
 
@@ -1220,15 +1253,33 @@ def log_gaussian(
             cache["colors_uploaded"] = True
 
     @override
-    def log_array(self, name: str, array: wp.array[Any] | np.ndarray):
+    def log_array(self, name: str, array: wp.array[Any] | np.ndarray | None):
         """
-        Log a generic array for visualization (not implemented).
+        Log a numeric array for visualization.
 
         Args:
             name: Unique path/name for the array signal.
-            array: Array data to visualize.
+            array: Array data to visualize, or ``None`` to remove a previously
+                logged array.
         """
-        pass
+        if array is None:
+            self._array_buffers.pop(name, None)
+            self._array_dirty.discard(name)
+            self._delete_array_texture(name)
+            return
+
+        array_np = array.numpy() if isinstance(array, wp.array) else np.asarray(array)
+        array_np = np.asarray(array_np, dtype=np.float32)
+
+        if array_np.ndim == 0:
+            array_np = array_np.reshape(1, 1)
+        elif array_np.ndim == 1:
+            array_np = array_np.reshape(1, -1)
+        elif array_np.ndim != 2:
+            raise ValueError("ViewerGL.log_array only supports scalar, 1-D, or 2-D arrays.")
+
+        self._array_buffers[name] = np.ascontiguousarray(array_np)
+        self._array_dirty.add(name)
 
     @override
     def log_scalar(
@@ -1590,6 +1641,8 @@ def close(self):
         """
         Close the viewer and clean up resources.
         """
+        self._clear_array_textures()
+        self._invalidate_pbo()
         self.renderer.close()
 
     @property
@@ -2325,18 +2378,192 @@ def _edit_color3(
 
         imgui.end()
 
+    @staticmethod
+    def _build_heatmap_color_lut() -> np.ndarray:
+        inferno_stops = (
+            (0.0, (0.001, 0.000, 0.014)),
+            (0.2, (0.169, 0.042, 0.341)),
+            (0.4, (0.416, 0.090, 0.433)),
+            (0.6, (0.698, 0.165, 0.388)),
+            (0.8, (0.944, 0.403, 0.121)),
+            (1.0, (0.988, 0.998, 0.645)),
+        )
+        lut = np.empty((256, 4), dtype=np.uint8)
+        for index, value in enumerate(np.linspace(0.0, 1.0, 256, dtype=np.float32)):
+            for stop_index in range(len(inferno_stops) - 1):
+                t0, c0 = inferno_stops[stop_index]
+                t1, c1 = inferno_stops[stop_index + 1]
+                if value <= t1:
+                    alpha = 0.0 if t1 <= t0 else (float(value) - t0) / (t1 - t0)
+                    rgb = [round(255.0 * ((1.0 - alpha) * c0[channel] + alpha * c1[channel])) for channel in range(3)]
+                    lut[index, :3] = rgb
+                    lut[index, 3] = 255
+                    break
+            else:
+                lut[index, :3] = [round(255.0 * channel) for channel in inferno_stops[-1][1]]
+                lut[index, 3] = 255
+        return lut
+
+    @staticmethod
+    def _downsample_heatmap(array: np.ndarray, target_rows: int, target_cols: int) -> np.ndarray:
+        rows, cols = array.shape
+        if rows <= target_rows and cols <= target_cols:
+            return array
+
+        row_factor = max(1, (rows + target_rows - 1) // target_rows)
+        col_factor = max(1, (cols + target_cols - 1) // target_cols)
+        new_rows = max(1, rows // row_factor)
+        new_cols = max(1, cols // col_factor)
+        if new_rows == rows and new_cols == cols:
+            return array
+
+        trimmed = array[: new_rows * row_factor, : new_cols * col_factor]
+        finite_mask = np.isfinite(trimmed)
+        safe_values = np.where(finite_mask, trimmed, 0.0)
+        reshaped_shape = (new_rows, row_factor, new_cols, col_factor)
+        value_sum = safe_values.reshape(reshaped_shape).sum(axis=(1, 3), dtype=np.float64)
+        value_count = finite_mask.reshape(reshaped_shape).sum(axis=(1, 3))
+        downsampled = np.full((new_rows, new_cols), np.nan, dtype=np.float32)
+        np.divide(value_sum, value_count, out=downsampled, where=value_count > 0)
+        return downsampled
+
+    def _colorize_heatmap(self, array: np.ndarray) -> tuple[np.ndarray, float, float]:
+        finite_mask = np.isfinite(array)
+        if not np.any(finite_mask):
+            rgba = np.empty((*array.shape, 4), dtype=np.uint8)
+            rgba[...] = self._heatmap_nan_rgba
+            return np.ascontiguousarray(rgba[::-1]), float("nan"), float("nan")
+
+        finite_values = array[finite_mask]
+        value_min = float(np.min(finite_values))
+        value_max = float(np.max(finite_values))
+        denom = max(value_max - value_min, 1.0e-8)
+
+        normalized = np.zeros(array.shape, dtype=np.float32)
+        np.subtract(array, value_min, out=normalized, where=finite_mask)
+        np.divide(normalized, denom, out=normalized, where=finite_mask)
+        np.clip(normalized, 0.0, 1.0, out=normalized)
+
+        lut_indices = np.rint(normalized * 255.0).astype(np.uint8)
+        rgba = self._heatmap_color_lut[lut_indices].copy()
+        rgba[~finite_mask] = self._heatmap_nan_rgba
+        return np.ascontiguousarray(rgba[::-1]), value_min, value_max
+
+    def _ensure_array_texture(self, name: str, width: int, height: int) -> dict[str, Any]:
+        texture_state = self._array_textures.get(name)
+        if texture_state is not None and texture_state["size"] == (width, height):
+            return texture_state
+
+        if texture_state is not None:
+            self._delete_array_texture(name)
+
+        gl = RendererGL.gl
+        texture_id = (gl.GLuint * 1)()
+        gl.glGenTextures(1, texture_id)
+        gl.glBindTexture(gl.GL_TEXTURE_2D, texture_id[0])
+        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_NEAREST)
+        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_NEAREST)
+        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_CLAMP_TO_EDGE)
+        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_CLAMP_TO_EDGE)
+        gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
+        gl.glTexImage2D(
+            gl.GL_TEXTURE_2D,
+            0,
+            gl.GL_RGBA8,
+            width,
+            height,
+            0,
+            gl.GL_RGBA,
+            gl.GL_UNSIGNED_BYTE,
+            None,
+        )
+        gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
+
+        texture_state = {
+            "texture_id": texture_id[0],
+            "size": (width, height),
+            "source_shape": None,
+            "display_shape": None,
+            "value_min": 0.0,
+            "value_max": 0.0,
+        }
+        self._array_textures[name] = texture_state
+        return texture_state
+
+    def _update_array_texture(self, texture_id: int, rgba: np.ndarray):
+        gl = RendererGL.gl
+        gl.glBindTexture(gl.GL_TEXTURE_2D, texture_id)
+        gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
+        gl.glTexSubImage2D(
+            gl.GL_TEXTURE_2D,
+            0,
+            0,
+            0,
+            rgba.shape[1],
+            rgba.shape[0],
+            gl.GL_RGBA,
+            gl.GL_UNSIGNED_BYTE,
+            rgba.ctypes.data_as(ctypes.POINTER(ctypes.c_ubyte)),
+        )
+        gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
+
+    def _render_array_heatmap(self, name: str, array: np.ndarray, width: float):
+        imgui = self.ui.imgui
+
+        rows, cols = array.shape
+        heatmap_width = max(120.0, width)
+        heatmap_height = np.clip(heatmap_width * rows / max(cols, 1), 80.0, 220.0)
+        target_cols = max(1, min(cols, int(heatmap_width / self._heatmap_min_cell_pixels)))
+        target_rows = max(1, min(rows, int(heatmap_height / self._heatmap_min_cell_pixels)))
+        display_array = self._downsample_heatmap(array, target_rows, target_cols)
+        display_rows, display_cols = display_array.shape
+        texture_state = self._ensure_array_texture(name, display_cols, display_rows)
+
+        if (
+            name in self._array_dirty
+            or texture_state["source_shape"] != array.shape
+            or texture_state["display_shape"] != display_array.shape
+        ):
+            rgba, value_min, value_max = self._colorize_heatmap(display_array)
+            self._update_array_texture(texture_state["texture_id"], rgba)
+            texture_state["source_shape"] = array.shape
+            texture_state["display_shape"] = display_array.shape
+            texture_state["value_min"] = value_min
+            texture_state["value_max"] = value_max
+            self._array_dirty.discard(name)
+
+        draw_list = imgui.get_window_draw_list()
+        origin = imgui.get_cursor_screen_pos()
+        imgui.image(imgui.ImTextureRef(texture_state["texture_id"]), imgui.ImVec2(heatmap_width, heatmap_height))
+
+        border_color = imgui.color_convert_float4_to_u32(imgui.ImVec4(1.0, 1.0, 1.0, 0.25))
+        draw_list.add_rect(
+            imgui.ImVec2(origin.x, origin.y),
+            imgui.ImVec2(origin.x + heatmap_width, origin.y + heatmap_height),
+            border_color,
+        )
+        shape_text = f"shape {rows}x{cols}"
+        if (display_rows, display_cols) != (rows, cols):
+            shape_text += f"  shown {display_rows}x{display_cols}"
+        if np.isfinite(texture_state["value_min"]) and np.isfinite(texture_state["value_max"]):
+            range_text = f"min {texture_state['value_min']:.4g}  max {texture_state['value_max']:.4g}"
+        else:
+            range_text = "min --  max --"
+        imgui.text(f"{shape_text}  {range_text}")
+
     def _render_scalar_plots(self):
-        """Render an ImGui window with live line plots for all logged scalars."""
-        if not self._scalar_buffers:
+        """Render an ImGui window with live line plots and array heatmaps."""
+        if not self._scalar_buffers and not self._array_buffers:
             return
 
         imgui = self.ui.imgui
         io = self.ui.io
 
         window_width = 400
+        item_height = len(self._scalar_buffers) * 140 + len(self._array_buffers) * 260
         window_height = min(
             io.display_size[1] - 20,
-            len(self._scalar_buffers) * 140 + 60,
+            item_height + 60,
         )
         imgui.set_next_window_pos(
             imgui.ImVec2(io.display_size[0] - window_width - 10, 10),
@@ -2365,6 +2592,13 @@ def _render_scalar_plots(self):
                     imgui.TreeNodeFlags_.default_open.value,
                 ):
                     imgui.plot_lines(f"##{name}", arr, graph_size=graph_size, overlay_text=overlay)
+
+            for name, array in self._array_buffers.items():
+                if imgui.collapsing_header(
+                    name,
+                    imgui.TreeNodeFlags_.default_open.value,
+                ):
+                    self._render_array_heatmap(name, array, window_width - 40.0)
         imgui.end()
 
     def _render_stats_overlay(self):