Add array heatmap logging to ViewerGL#2433
Add array heatmap logging to ViewerGL#2433han-xudong wants to merge 4 commits intonewton-physics:mainfrom
Conversation
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdd heatmap rendering for 2D numeric arrays to ViewerGL: implement Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant ViewerGL
participant CPU as CPU (downsample/LUT)
participant GPU as OpenGL
participant ImGui
Client->>ViewerGL: log_array(name, array)
ViewerGL->>ViewerGL: validate/normalize array, store in _array_buffers[name], set _array_dirty[name]=true
loop per frame
ViewerGL->>ViewerGL: check for scalar/array data
ViewerGL->>ViewerGL: if _array_dirty or shape/display changed
ViewerGL->>CPU: compute downsampled grid, normalize values, apply LUT (colorize), preserve NaNs
CPU->>GPU: allocate/update texture (glTexImage2D/glTexSubImage2D)
GPU->>ViewerGL: texture ready
ViewerGL->>ImGui: render Plots window with texture and overlays (border, shape/min/max)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
newton/_src/viewer/viewer_gl.py (1)
2386-2400: Consider downsampling large arrays for performance.The nested loop issues one
add_rect_filledcall per cell. For large arrays (e.g., 256×256 = 65,536 cells), this could cause noticeable frame rate drops in the diagnostics overlay.Consider adding downsampling for arrays exceeding a threshold (e.g., 64×64) to maintain interactive frame rates:
Sketch of downsampling approach
# At the start of _render_array_heatmap, before the loop: MAX_CELLS = 64 if rows > MAX_CELLS or cols > MAX_CELLS: # Downsample using block averaging row_factor = max(1, rows // MAX_CELLS) col_factor = max(1, cols // MAX_CELLS) new_rows = rows // row_factor new_cols = cols // col_factor array = array[:new_rows * row_factor, :new_cols * col_factor] array = array.reshape(new_rows, row_factor, new_cols, col_factor).mean(axis=(1, 3)) rows, cols = new_rows, new_cols🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@newton/_src/viewer/viewer_gl.py` around lines 2386 - 2400, The nested per-cell draw in _render_array_heatmap currently calls draw_list.add_rect_filled for every element (using variables rows, cols and cell drawing logic), which can be very slow for large arrays; implement downsampling before the nested loop when rows or cols exceed a threshold (e.g., MAX_CELLS = 64): compute integer row_factor and col_factor, trim the array to a multiple of those factors, perform block averaging to produce a smaller array and update rows and cols to the downsampled dimensions, then continue using the existing loop and add_rect_filled calls on the downsampled array (keep the original origin/cell sizing logic but base cell counts on the new rows/cols so the overlay scales correctly).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@newton/_src/viewer/viewer_gl.py`:
- Around line 1225-1248: The log_array method's signature and docstring don't
reflect that None can be passed to remove a previously logged array and the
Warp-array detection is brittle; update the type annotation of
ViewerGL.log_array to accept Optional (None) for the array parameter, update the
docstring to document that passing None removes the named array from
_array_buffers, and change the Warp-array detection from isinstance(array,
wp.array) to a hasattr(array, "numpy") check (consistent with
ViewerRerun._to_numpy()) before converting to numpy; ensure you import
typing.Optional if needed and keep the rest of the conversion/contiguity logic
the same.
---
Nitpick comments:
In `@newton/_src/viewer/viewer_gl.py`:
- Around line 2386-2400: The nested per-cell draw in _render_array_heatmap
currently calls draw_list.add_rect_filled for every element (using variables
rows, cols and cell drawing logic), which can be very slow for large arrays;
implement downsampling before the nested loop when rows or cols exceed a
threshold (e.g., MAX_CELLS = 64): compute integer row_factor and col_factor,
trim the array to a multiple of those factors, perform block averaging to
produce a smaller array and update rows and cols to the downsampled dimensions,
then continue using the existing loop and add_rect_filled calls on the
downsampled array (keep the original origin/cell sizing logic but base cell
counts on the new rows/cols so the overlay scales correctly).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: 7eab504e-c449-4ed7-a131-ab06ca1cd602
📒 Files selected for processing (2)
CHANGELOG.mdnewton/_src/viewer/viewer_gl.py
| def log_array(self, name: str, array: wp.array[Any] | np.ndarray): | ||
| """ | ||
| 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. | ||
| """ | ||
| pass | ||
| if array is None: | ||
| self._array_buffers.pop(name, None) | ||
| return | ||
|
|
||
| array_np = array.numpy() if isinstance(array, wp.array) else np.asarray(array) | ||
| array_np = np.asarray(array_np, dtype=np.float32) | ||
| array_np = np.squeeze(array_np) | ||
|
|
||
| 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) |
There was a problem hiding this comment.
Type annotation and docstring should reflect None behavior.
The method accepts None to remove a previously logged array (per PR description), but the type annotation doesn't include None and the docstring doesn't document this behavior.
Additionally, consider using the hasattr(x, "numpy") pattern for Warp array detection, which is more flexible and consistent with the pattern in ViewerRerun._to_numpy().
Proposed fix
`@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 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.
"""
if array is None:
self._array_buffers.pop(name, None)
return
- array_np = array.numpy() if isinstance(array, wp.array) else np.asarray(array)
+ array_np = array.numpy() if hasattr(array, "numpy") else np.asarray(array)
array_np = np.asarray(array_np, dtype=np.float32)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@newton/_src/viewer/viewer_gl.py` around lines 1225 - 1248, The log_array
method's signature and docstring don't reflect that None can be passed to remove
a previously logged array and the Warp-array detection is brittle; update the
type annotation of ViewerGL.log_array to accept Optional (None) for the array
parameter, update the docstring to document that passing None removes the named
array from _array_buffers, and change the Warp-array detection from
isinstance(array, wp.array) to a hasattr(array, "numpy") check (consistent with
ViewerRerun._to_numpy()) before converting to numpy; ensure you import
typing.Optional if needed and keep the rest of the conversion/contiguity logic
the same.
|
Nice addition! Having inline heatmap visualization for logged arrays is a great debugging tool for ViewerGL. The main concern is that the per-cell rendering loop runs entirely in Python every frame, which will bottleneck the viewer at moderate array sizes. |
| for row in range(rows): | ||
| for col in range(cols): | ||
| value = float(array[row, col]) | ||
| if np.isfinite(value): |
There was a problem hiding this comment.
🟡 This nested loop runs rows × cols Python-level draw calls every frame. For a 64×64 array that is 4,096 iterations of float(), np.isfinite(), _heatmap_color(), two ImVec2 constructions, and an add_rect_filled() call -- enough to drop below interactive frame rates.
_heatmap_color makes it worse by calling np.clip on a single scalar each iteration instead of plain min/max.
Suggested fix (two levels):
| Level | Change | Effort |
|---|---|---|
| Quick win | Replace np.clip(value, ...) in _heatmap_color with max(0.0, min(1.0, value)), and add a size cap that down-samples arrays larger than 128×128 before rendering |
Small |
| Proper fix | Vectorize normalization and color mapping across the whole array in NumPy, rasterize into an RGBA texture, and draw it with a single imgui.image() call |
Medium |
The proper fix moves all O(rows×cols) work out of Python and into one GPU texture upload per frame, which keeps the viewer responsive regardless of array size.
Quick-win example for _heatmap_color
# before
normalized = np.clip(value, 0.0, 1.0)
# after -- avoids numpy overhead on a single float
normalized = max(0.0, min(1.0, value))
han-xudong
force-pushed
the
feat/viewergl-log-array
branch
from
cc137e8 to
e33d05e
Compare
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@newton/_src/viewer/viewer_gl.py`:
- Line 2398: The expression int(round(...)) is redundant in the rgb list
comprehension in viewer_gl.py; replace int(round(255.0 * ((1.0 - alpha) *
c0[channel] + alpha * c1[channel]))) with round(255.0 * ((1.0 - alpha) *
c0[channel] + alpha * c1[channel])) (and the same change for the other
occurrence around line 2403) so the list comprehension producing rgb uses
round(...) directly (round returns an int when no ndigits is given).
- Around line 1270-1279: The code in ViewerGL.log_array currently calls
np.squeeze which collapses singleton dimensions and can silently transpose
column vectors and accept higher-rank arrays with singleton dims; change the
logic to inspect the original array.ndim (or isinstance(array, wp.array)
converted shape) before calling np.squeeze: if original ndim == 0 reshape to
(1,1), if ndim == 1 reshape to (1, -1), if ndim == 2 leave as-is (do not squeeze
away singleton dims that would change orientation), otherwise raise the
ValueError; update the branch around the existing array_np handling (referencing
ViewerGL.log_array and the variables array and array_np) so that squeezing does
not alter the authoritative ndim-based behavior described by the error message.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: 3ccb3bdf-c68b-4225-9380-af382c49bd33
📒 Files selected for processing (2)
CHANGELOG.mdnewton/_src/viewer/viewer_gl.py
✅ Files skipped from review due to trivial changes (1)
- CHANGELOG.md
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@newton/_src/viewer/viewer_gl.py`:
- Around line 2432-2452: In _colorize_heatmap, when the input has no finite
values (finite_mask all False) currently you return (rgba, 0.0, 0.0); change it
to return NaN for the min/max so callers can detect "no finite values" (e.g.,
return np.nan or float('nan') for value_min and value_max) while still returning
the same rgba (filled with self._heatmap_nan_rgba and
ascontiguousarray(rgba[::-1])); update the early-return branch in the
_colorize_heatmap function to set value_min/value_max to NaN instead of 0.0 so
the overlay can render "—" or similar.
- Around line 2462-2465: The glGenTextures call in viewer_gl.py uses a scalar
gl.GLuint() which isn't a ctypes-compatible pointer; change to allocate an array
like the PBO code: create texture_id = (gl.GLuint * 1)() and call
gl.glGenTextures(1, texture_id), then use texture_id[0] (not texture_id.value)
for subsequent calls such as gl.glBindTexture and anywhere the texture ID is
read; update any references to texture_id.value to texture_id[0].
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: bd9661b5-6d1e-4486-b596-a28fdfe1441c
📒 Files selected for processing (1)
newton/_src/viewer/viewer_gl.py
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
Implement texture-backed array heatmap rendering in ViewerGL.\n\nSupport scalar, 1-D, and 2-D array logging, including removal via\nNone, array downsampling, NaN handling, and texture lifecycle\ncleanup. Fold the follow-up fixes for column-vector handling, all-NaN\nmetadata display, and correct OpenGL texture allocation into the same\nchange so the PR history stays compact.
han-xudong
force-pushed
the
feat/viewergl-log-array
branch
from
49fae6d to
e126a2d
Compare
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@newton/_src/viewer/viewer_gl.py`:
- Around line 1256-1282: The ViewerGL.log_array change accepts None to remove a
logged array but the abstract and peer backends (ViewerBase.log_array,
ViewerFile.log_array, ViewerNull.log_array, ViewerRerun.log_array) still use the
old non-optional signature and lack removal semantics; update the abstract
signature in ViewerBase.log_array to accept Optional[wp.array[Any] | np.ndarray]
(or None) and update the docstring to document removal behavior, then modify
each peer override (ViewerFile.log_array, ViewerNull.log_array,
ViewerRerun.log_array) to accept None and implement the same removal/no-op
semantics (e.g., treat None by removing any stored buffer/metadata or returning
early) so the contract is consistent across all viewers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: 8827e4b6-d866-47b1-b1ac-65909a32dfe1
📒 Files selected for processing (2)
CHANGELOG.mdnewton/_src/viewer/viewer_gl.py
✅ Files skipped from review due to trivial changes (1)
- CHANGELOG.md
| 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) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Abstract base and peer backends still don't accept None — contract divergence.
The override now accepts None to remove a logged array, but ViewerBase.log_array (newton/_src/viewer/viewer.py:1205), ViewerFile.log_array, ViewerNull.log_array, and ViewerRerun.log_array all keep the old array: wp.array[Any] | np.ndarray signature and have no removal semantics. Code written against the abstract contract either can't type-safely pass None, or will silently no-op / raise on the other backends. Please align the abstract signature/docstring (and at minimum make the peer overrides treat None as a no-op removal) so removal behavior is portable across viewers.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@newton/_src/viewer/viewer_gl.py` around lines 1256 - 1282, The
ViewerGL.log_array change accepts None to remove a logged array but the abstract
and peer backends (ViewerBase.log_array, ViewerFile.log_array,
ViewerNull.log_array, ViewerRerun.log_array) still use the old non-optional
signature and lack removal semantics; update the abstract signature in
ViewerBase.log_array to accept Optional[wp.array[Any] | np.ndarray] (or None)
and update the docstring to document removal behavior, then modify each peer
override (ViewerFile.log_array, ViewerNull.log_array, ViewerRerun.log_array) to
accept None and implement the same removal/no-op semantics (e.g., treat None by
removing any stored buffer/metadata or returning early) so the contract is
consistent across all viewers.
|
@eric-heiden Thanks for the thoughtful review and for calling out the Python-side rendering cost here. You were absolutely right that the per-cell draw loop would become a bottleneck at moderate array sizes. I reworked the heatmap path toward the “proper fix” you suggested: the data is now normalized and color-mapped in NumPy, downsampled to the displayed resolution, uploaded as an RGBA texture, and rendered with a single I also folded in a few follow-up fixes that came up during review, including some CodeRabbit suggestions:
Thanks again for the push here. The performance concern was the right one to fix properly. |
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
| 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") |
There was a problem hiding this comment.
🟡 The rgba[::-1] here flips the heatmap vertically before upload, so the top row of the source array is rendered at the bottom of the image. imgui.image() with the default UVs already maps texture row 0 to the top of the rect, so this inverts the array relative to the usual NumPy/Matplotlib orientation. The same reversal is present in the non-NaN return below. Unless origin='lower' is intentional, return np.ascontiguousarray(rgba) in both places, or do the flip through uv0/uv1 and document why.
There was a problem hiding this comment.
@eric-heiden Thanks, good catch. You were right that the extra rgba[::-1] was flipping the heatmap vertically. I removed the reversal in both return paths so the texture upload now preserves the usual array orientation.
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
han-xudong
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
han-xudong
requested a deployment
to
external-pr-approval
GitHub Actions
Waiting
han-xudong
requested a deployment
to
external-pr-approval
GitHub Actions
Waiting
2 participants
Description
Add support for visualizing logged numeric arrays in
ViewerGL.This implements
ViewerGL.log_array()for scalar, 1-D, and 2-D numeric arrays.2-D arrays are rendered as heatmaps in the existing ImGui diagnostics overlay,
alongside the current scalar plots.
This makes it easier to inspect grid-based or field-like debug data directly in
the GL viewer.
Checklist
CHANGELOG.mdhas been updated (if user-facing change)Test plan
Manual verification in a local
ViewerGLsession.Verified that:
ViewerGL.log_array()accepts scalar inputs and displays them as a 1x1 arrayViewerGL.log_array()accepts 1-D arrays and displays them as a single-row heatmapViewerGL.log_array()accepts 2-D arrays and displays them as a heatmapNoneremoves a previously logged arrayNew feature / API change
Summary by CodeRabbit
New Features
Bug Fixes