ENH: enable passing RGB(A) to polormesh

[rcomer]

Rationale

Since Matplotlib v3.7.0, pcolormesh accepts RGBA arrays
https://matplotlib.org/stable/users/prev_whats_new/whats_new_3.7.0.html#pcolormesh-accepts-rgb-a-colors

Cartopy just needs a small tweak to be able to pass these down, and we can adapt this StackOverflow question to give

import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import numpy as np

np.random.seed(100)

x = np.arange(10, 20)
y = np.arange(0, 10)
x, y = np.meshgrid(x, y)

img = np.random.randint(low=0, high=255, size=(10, 10, 4)) / 255

projection = ccrs.Mollweide()
transform = ccrs.PlateCarree()
ax = plt.axes(projection=projection)
ax.set_global()
ax.coastlines()

ax.pcolormesh(np.linspace(0, 120, 11), np.linspace(0, 80, 11), img, transform=transform)

plt.show()

Closes #2156 and I think also resolves #2045 and resolves #2171, as the request for set_array to handle None seems to be a workaround that is not needed since mpl 3.7. I also note that the None handling is not actually documented in QuadMesh.set_array.

This also adds test coverage for and fixes #2208.

I'm not sure of the best way to test this. Does it just need a smoke test of a pcolormesh plot with suitable data?

Implications

[greglucas]

Does this also need an update of the GeoQuadMesh collection to handle the set_array() case with wrapping?

One thought for testing this would be to duplicate one of the current tests but use RGBA values from cmap(norm(arr)) values, and then pass that into the pcolormesh(rgba_arr) call and verify it is the same as pcolormesh(arr, cmap=cmap, norm=norm). It would be good to do this for a case with wrapping as well to make sure we are handling that broadcasting properly too.

[rcomer]

Thanks @greglucas, I took a look at the wrapping code and this is clearly more complicated than I first thought! I think I can see how to proceed but not sure how soon I will get to it, so I’ll put this in draft for now.

[rcomer]

Ah, I had not twigged that this uses pcolor underneath.

cartopy/lib/cartopy/mpl/geoaxes.py

Lines 1954 to 1964 in c8f1b0f

	pcolor_data = np.ma.array(np.zeros(C.shape),
	mask=~mask)
	pcolor_col = self.pcolor(coords[..., 0], coords[..., 1],
	pcolor_data, zorder=zorder,
	**kwargs)
	# Now add back in the masked data if there was any
	full_mask = ~mask if C_mask is None else ~mask | C_mask
	pcolor_data = np.ma.array(C, mask=full_mask)
	# The pcolor_col is now possibly shorter than the
	# actual collection, so grab the masked cells
	pcolor_col.set_array(pcolor_data[mask].ravel())

Does this mean that this basically can't work without changes to pcolor in Matplotlib?

[greglucas]

Ahh, you might be right that we can't quite handle RGBA in the general case yet. I do have an open PR over in MPL to make pcolor more mesh-like, which may help with the pcolor RGBA support as well if you want to try that combination out... matplotlib/matplotlib#25027

[rcomer]

OK, I see that matplotlib/matplotlib#25027 stops pcolor from ravelling arrays and compressing masked arrays, which definitely helps. I think we will then need a further change in mpl to get pcolor to handle RGB(A). Part of me is tempted to branch off your branch to see if I can get it working locally, but I think it might be more sensible to wait till matplotlib/matplotlib#25027 is merged.

[greglucas]

Part of me is tempted to branch off your branch to see if I can get it working locally, but I think it might be more sensible to wait till matplotlib/matplotlib#25027 is merged.

Yeah, good call, I wouldn't put too much effort into this yet until that has some consensus on whether the PR is desired or if it goes in a different direction. Feel free to leave comments/review over there as well if you have thoughts for how to improve it!

[rcomer]

OK this now works with your branch for the cases I have tested. There are some loose ends to tidy up:

• Currently this works with RGBA but not RGB, but I don't think it will take much to address that.
• GeoQuadMesh.set_array still only works for 1D arrays.
• Need to look at the get_array method.
• Possibly adapt some more tests for better coverage.

But I'm happy I have a working POC, and it probably doesn't make sense to polish this too much until matplotlib/matplotlib#25027 makes it into a release candidate.

There were a couple of new test failures when I picked up the latest commit to your branch. I've not looked into those yet, but they were both failing on the assertion of expected output which doesn't always mean something has gone wrong!

[rcomer]

Thanks to @kylehofmann for the latest commit, which handles the RGB case that I had previously neglected.

[rcomer]

I think all the functionality is there now. A few things are not yet tested:

• GeoQuadMesh.get_array for RGB(A)
• [maybe not] GeoQuadMesh.set_array when the input is different to before, e.g. if the original pcolormesh call used a 2D array but now we're passing RGBA. This works in Matplotlib and, from a quick check, seems to work here too. Edit: actually it doesn't work in matplotlib if you initially call pcolormesh with RGBA and then try to set a 2D array - it updates but you do not get the same plot that you would have if you passed the 2D array to begin with. I don't even know what the use-case would be for this so maybe we should leave it until and unless someone asks for it.
• GeoQuadMesh.set_array when there is no wrapping. I guess this didn't need to be tested before as previously it just passed straight to QuadMesh.set_array for this case.

@kylehofmann I'm sorry but I changed my mind about supporting the masked case for RGB(A). As far as I can tell Matplotlib's pcolormesh just ignores masks on RGB(A) arrays, and I think it makes sense to be consistent with that. Dropping support for that case also made the code simpler.

[greglucas]

As far as I can tell Matplotlib's pcolormesh just ignores masks on RGB(A) arrays, and I think it makes sense to be consistent with that

I agree, the masked RGBA case should follow Matplotlib. It explicitly states the mask gets ignored in the docstring of to_rgba:
https://matplotlib.org/stable/api/cm_api.html#matplotlib.cm.ScalarMappable.to_rgba

With that being said, I think it would be fairly straight forward to implement handling of the masked arrays if we wanted to make them fully transparent? Something like the following patch.

index daefbe08e9..0c4a99496a 100644
--- a/lib/matplotlib/cm.py
+++ b/lib/matplotlib/cm.py
@@ -493,6 +493,9 @@ class ScalarMappable:
                 else:
                     raise ValueError("Image RGB array must be uint8 or "
                                      "floating point; found %s" % xx.dtype)
+                # Account for any masked entries in the original array
+                if np.ma.is_masked(x):
+                    xx[np.any(np.ma.getmaskarray(x), axis=2), 4] = 0
                 return xx
         except AttributeError:
             # e.g., x is not an ndarray; so try mapping it

[rcomer]

@greglucas thanks, I was just wondering about that. With your mpl branch, pcolor does respect the mask for the RGBA cases I tried. So if pcolormesh did as well then we could probably simplify the logic here and just use masks for all cases rather than messing with the alpha channel.

[kylehofmann]

@rcomer: No offense taken! Consistency with matplotlib is important, and as you say, the code is simpler. (A big plus!) My only suggestion is to change RBG to RGB.

[rcomer]

Have I understood this right? I think this test relates to the need to initially call pcolor with an array of zeros, which will no longer be necessary with the new PolyQuadMesh array handling.

[greglucas]

I think that the number of paths should still be short in the new version, so I think this should still be tested. i.e. we only want to draw the small number of cells, we don't want to draw all pcolor cells and make them invisible.

[rcomer]

Thanks to matplotlib/matplotlib#26096, Matplotlib's pcolormesh will respect the masks on RGB(A) data from v3.8. So I have simplified the handling to always use masks and also reinstated the masked RGB tests.

[rcomer]

I think this one is finally ready for review!

[greglucas]

This is really nice work @rcomer! Thank you for tackling this because there are a lot of paths to keep track of with the different versions hanging around everywhere. Eventually, this will slim down the code ;)

[greglucas]

I think that the number of paths should still be short in the new version, so I think this should still be tested. i.e. we only want to draw the small number of cells, we don't want to draw all pcolor cells and make them invisible.

[rcomer]

I think that the number of paths should still be short in the new version, so I think this should still be tested. i.e. we only want to draw the small number of cells, we don't want to draw all pcolor cells and make them invisible.

That test now doesn't get any paths back, so I have updated to assert that from v3.8. But I think I am still not understanding quite right.

[greglucas]

I think that the number of paths should still be short in the new version, so I think this should still be tested. i.e. we only want to draw the small number of cells, we don't want to draw all pcolor cells and make them invisible.

That test now doesn't get any paths back, so I have updated to assert that from v3.8. But I think I am still not understanding quite right.

I think you actually had a better grasp on this than I did! I think you can actually skip this after 3.8 like you had before because it doesn't make sense. Let me try and write down what I think is happening and see if you agree.

This test passes in one all-nan cell that spans the wrap-line. That actually becomes two wrapped cells in the pcolor-collection. Before MPL3.8 we wanted those cells to show up because we needed to be able to update them later. After MPL3.8 the PolyQuadMesh collection takes care of the shrink/expansion of the paths automatically for us. So, with nan's it is empty, but if we updated the mesh we'd expect two paths to be present (PolyQuadMesh.set_array(np.array([[1]])) should create those two paths I think). I suppose that case of setting a new data array could be added as an extra assertion as well if that makes sense...

[rcomer]

Yes, I think we're on the same page. It makes sense to me to always make sure the wrapped cells can be updated, so I have added that extra assertion.

[rcomer]

Thanks for all your guidance and mpl updates @greglucas!