feat: Add facet_row support to px.imshow
Summary
This PR adds facet_row support to plotly.express.imshow, enabling visualization of 4D and 5D data by faceting along both row and column dimensions.
Closes #4108
Motivation
Currently, px.imshow only supports facet_col for creating faceted plots, limiting users to visualizing 3D data (with one facet dimension). Many scientific and
engineering use cases require visualizing 4D tensor data (e.g., full factorial experimental designs in ML and biological sciences). This PR enables users to use both
facet_row and facet_col simultaneously, along with animation_frame, to visualize up to 5D datasets.
Changes
plotly/express/_imshow.py
- Added
facet_rowparameter to the function signature - Updated docstring to document the new parameter
- Updated slice dimension calculations to handle both
facet_rowandfacet_col - Updated xarray handling to extract
facet_rowdimension labels - Updated array reshaping to handle the additional dimension
- Updated figure building to generate row labels and place traces correctly
- Note:
facet_col_wrapis ignored whenfacet_rowis set (consistent with other px functions)
tests/test_optional/test_px/test_imshow.py
- Added
test_facet_row: Tests facet_row with numpy arrays - Added
test_facet_row_and_col: Tests using both facet_row and facet_col together - Added
test_animation_facet_row_and_col: Tests animation + facet_row + facet_col - Added
test_imshow_xarray_facet_row: Tests facet_row with xarray - Added
test_imshow_xarray_facet_row_and_col: Tests both facets with xarray
Example Usage
import plotly.express as px
# Prepare data for multi dimensional usage
df = px.data.tips()
cols = ['smoker', 'sex', 'time', 'day', 'size']
pivot = df.groupby(cols)['total_bill'].mean().reset_index()
da = pivot.set_index(cols).to_xarray()['total_bill']
# 4D data: visualize with facet_row and facet_col
fig1 = px.imshow(da.sum('day'), facet_row="smoker", facet_col="sex")
fig1.show()
# Also works with numpy
fig2 = px.imshow(da.sum('day').values, facet_row=0, facet_col=1)
fig2.show()
# 5D data: animation + facet_row + facet_col
fig3 = px.imshow(da, facet_row="smoker", facet_col="sex", animation_frame='day')
fig3.show()
Build fails are unrelated to changes afaik. #5442
Build fails are unrelated to changes afaik. #5442
@FBumann You're right, #5442 has been merged so the docs build should be fixed on main; you can update this branch and run CI again.
@emilykl Ready for review