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.