Skip to content

Commit 99f77a3

Browse files
Jammy2211claude
authored andcommitted
Add RGB support to plot_array and unit test
`plot_array` now detects 3- or 4-channel arrays and skips colormap, norm, and colorbar — all of which are inappropriate for RGB images. A `test__plot_array_rgb` unit test covers the new code path. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent a8e8567 commit 99f77a3

File tree

2 files changed

+42
-10
lines changed

2 files changed

+42
-10
lines changed

autoarray/plot/array.py

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,8 @@ def plot_array(
133133
if array is None or array.size == 0:
134134
return
135135

136+
is_rgb = array.ndim == 3 and array.shape[2] in (3, 4)
137+
136138
if colormap is None:
137139
from autoarray.plot.utils import _default_colormap
138140
colormap = _default_colormap()
@@ -194,21 +196,30 @@ def plot_array(
194196
h, w = array.shape[:2]
195197
_box_aspect = (w / h) if h > 0 else 1.0
196198

197-
im = ax.imshow(
198-
array,
199-
cmap=colormap,
200-
norm=norm,
201-
extent=extent,
202-
aspect="auto", # image fills the axes box; box shape set below
203-
origin=origin_imshow,
204-
)
199+
if is_rgb:
200+
im = ax.imshow(
201+
array,
202+
extent=extent,
203+
aspect="auto",
204+
origin=origin_imshow,
205+
)
206+
else:
207+
im = ax.imshow(
208+
array,
209+
cmap=colormap,
210+
norm=norm,
211+
extent=extent,
212+
aspect="auto", # image fills the axes box; box shape set below
213+
origin=origin_imshow,
214+
)
205215

206216
# Shape the axes box to match the data so there is no surrounding
207217
# whitespace when the panel is embedded in a subplot grid.
208218
ax.set_aspect(_box_aspect, adjustable="box")
209219

210-
from autoarray.plot.utils import _apply_colorbar
211-
_apply_colorbar(im, ax, cb_unit=cb_unit, is_subplot=not owns_figure)
220+
if not is_rgb:
221+
from autoarray.plot.utils import _apply_colorbar
222+
_apply_colorbar(im, ax, cb_unit=cb_unit, is_subplot=not owns_figure)
212223

213224
# --- overlays --------------------------------------------------------------
214225
if array_overlay is not None:

test_autoarray/structures/plot/test_structure_plotters.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,3 +145,24 @@ def test__array_rgb(
145145
)
146146

147147
assert path.join(plot_path, "array_rgb.png") in plot_patch.paths
148+
149+
150+
def test__plot_array_rgb(
151+
array_2d_rgb_7x7,
152+
plot_path,
153+
plot_patch,
154+
):
155+
"""
156+
`plot_array` (the high-level function) must handle `Array2DRGB` inputs without
157+
applying a colormap, norm, or colorbar — all of which would raise errors or
158+
produce nonsense for a 3-channel image.
159+
"""
160+
aplt.plot_array(
161+
array=array_2d_rgb_7x7,
162+
title="RGB Test",
163+
output_path=plot_path,
164+
output_filename="array_rgb_high_level",
165+
output_format="png",
166+
)
167+
168+
assert path.join(plot_path, "array_rgb_high_level.png") in plot_patch.paths

0 commit comments

Comments
 (0)