diff --git a/CHANGELOG.md b/CHANGELOG.md index 24800bee8..1062838b0 100644 --- a/CHANGELOG.md +++ b/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` diff --git a/newton/_src/viewer/viewer_gl.py b/newton/_src/viewer/viewer_gl.py index 4bff11831..7ae717b90 100644 --- a/newton/_src/viewer/viewer_gl.py +++ b/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), 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), 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):