imshow fails on xarrays when the facet_col dimension is before the animation_frame dimension

[Karl-Krauth]

Minimal example:

arr = xr.DataArray(np.random.randn(2, 2, 3, 3), dims=("A", "B", "C", "D"))
px.imshow(arr, facet_col="A", animation_frame="B")

which outputs this error message:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[124], line 2
      1 arr = xr.DataArray(np.random.randn(2, 2, 3, 3), dims=("A", "B", "C", "D"))
----> 2 px.imshow(arr, facet_col="A", animation_frame="B")

File ~/.bin/mambaforge/lib/python3.10/site-packages/plotly/express/_imshow.py:387, in imshow(img, zmin, zmax, origin, labels, x, y, animation_frame, facet_col, facet_col_wrap, facet_col_spacing, facet_row_spacing, color_continuous_scale, color_continuous_midpoint, range_color, title, template, width, height, aspect, contrast_rescaling, binary_string, binary_backend, binary_compression_level, binary_format, text_auto)
    385 y_index = slice_dimensions
    386 if y is not None and img.shape[y_index] != len(y):
--> 387     raise ValueError(
    388         "The length of the y vector must match the length of the first "
    389         + "dimension of the img matrix."
    390     )
    391 x_index = slice_dimensions + 1
    392 if x is not None and img.shape[x_index] != len(x):

ValueError: The length of the y vector must match the length of the first dimension of the img matrix.

I had a look through the code and it looks like it's because this line:
https://github.com/plotly/plotly.py/blob/master/packages/python/plotly/plotly/express/_imshow.py#L274
assumes that the animation_frame index doesn't get shifted when pop gets called for the facet_col dim.

I'd be happy to submit a pull request addressing this and #4329 if someone on the core dev team would be willing to merge it. :)

[alexcjohnson]

Nice sleuthing @Karl-Krauth - yes we would love a PR, particularly if it comes with a test (which can use exactly the code you posted above)

[gvwilson]

Hi - we are tidying up stale issues and PRs in Plotly's public repositories so that we can focus on things that are still important to our community. Since this one has been sitting for a while, I'm going to close it; if it is still a concern, please add a comment letting us know what recent version of our software you've checked it with so that I can reopen it and add it to our backlog. If you'd like to submit a PR, we'd be happy to prioritize a review, and if it's a request for tech support, please post in our community forum. Thank you - @gvwilson