Mirror patches across periodic boundaries

[andrewdnolan]

This PR add the ability to mirror patch across periodic boundaries for planar periodic meshes. With this feature added we have the "full" periodic effect when plotting.

Much of the implementation of this PR follows the cartopy.mpl.geo_collection module pretty closely. The cartopy approach split the collections between pcolor and pcolormesh call for patches that cross projection boundaries, because pcolormesh does not support nan's in the coordinate array. Our need to keeping two collections is a little different than cartopy.

We instead keep track of a second collection, which is the mirrored patches, and the associated indices to enable plotting. We have created our own MPASCollection object, following catropy's approach, which has take special care to define a MPASCollection.set_array method so that the collection array can be updated (i.e. for animating) while doing the bookkeeping need for periodic mirroring under the hood.

[andrewdnolan]

Testing:

import mosaic
import matplotlib.pyplot as plt

def test_periodic(descriptor, ds, loc):
    fig, ax = plt.subplots()

    lc = mosaic.polypcolor(ax, descriptor, ds.indexToCellID * 0.0, ec='k', cmap='binary', zorder=0)
    pc = mosaic.polypcolor(ax, descriptor, ds[f"indexTo{loc}ID"], alpha=0.5, ec='tab:orange')

    ax.scatter(ds[f"x{loc}"], ds[f"y{loc}"], color="tab:blue")

    ax.set_title(f"{loc} Patches")
    ax.set_aspect("equal")
    fig.savefig(f”mirror/{loc}.png", dpi=300, bbox_inches='tight’)
    return ax, pc, lc

ds = mosaic.datasets.open_dataset("doubly_periodic_4x4")

descriptor = mosaic.Descriptor(ds)

for loc in ["Cell", "Edge", "Vertex"]:
    ax, pc, lc = test_periodic(descriptor, ds, loc)

plt.show()

Cell:

Edge:

Vertex:

[andrewdnolan]

Testing (cont).

Here's an example of using the MPASCollection.set_array method, which updates the data array and the mirror patch data values. Animations will be a common use case for this functionality so it's exciting to have this working!

import mosaic
import matplotlib.animation as animation
import matplotlib.pyplot as plt

ds = mosaic.datasets.open_dataset("doubly_periodic_4x4")
descriptor = mosaic.Descriptor(ds)

fig, ax = plt.subplots(constrained_layout=True)

pc = mosaic.polypcolor(ax, descriptor, ds.indexToCellID, alpha=0.5, ec='k')

def update(frame):
    pc.set_array(ds.indexToCellID - frame)
    return(pc, )

ani = animation.FuncAnimation(fig=fig, func=update, frames=10, interval=500)

[andrewdnolan]

Once perlmutter is back online I'll test with the baroclinic channel mesh from polaris. Some unit tests seem to be intermittently failing due to network issues, hopefully that will resolve itself when I test again next week.

[andrewdnolan]

I've come back to this PR, in hopes to get it merged in time to be included in the 1.2.0 release of mosaic for the upcoming e3sm-unified release.

I had continued the testing this PR by running the "baroclinic channel" test case from polaris. I had been struggling with a persistent issues where the clim of the original array and mirrored array were different. Producing a figure like:

I was able to confirm the clims were different for the original vs. mirrored collection by adding a some print statements:

--------------------------------------------------                                                       
Original Collection                                 
--------------------------------------------------                                                       
clim = (8.974999999999998, 13.025000000000002)                                                           
--------------------------------------------------                                                       
Mirrored Collection                                                                                      
--------------------------------------------------                                                       
clim = (11.824999996395533, 13.025000000000002)

It turns out that the vmin/vmax were not being passed through the _mirror_polycollection call. So the collection was scaled to the array min/max values, where the data min was not as small as the requested vmin.

[andrewdnolan]

By using the pop method, we were removing the entry from the kwargs dictionary. Which means we this was forwarded along to the _mirror_polycollection call below the entries here were all none.

[andrewdnolan]

But, PolyCollection does not accept these arguments, which is why there were originally poped.

So we still need to pop but then repack before the _mirror_polycollection so that all the kwargs are consistent.

[andrewdnolan]

This now works successfully for the "baroclinic channel" testcase. As for the examples shown above.

I've added unit test which should catch the inconsistent color limits I was dealing with earlier.

@xylar If you have other ideas about how to further test this PR, they be appreciated. Otherwise a look over would be appreciated!

[xylar]

@andrewdnolan, this looks great to me. I looked through the code and I think your testing is sufficient. Please merge when ready.