import mosaic
import matplotlib.pyplot as plt

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

    mosaic.polypcolor(ax, descriptor, ds.indexToCellID * 0.0, ec='k', cmap='binary')
    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")

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

descriptor = mosaic.Descriptor(ds)

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

plt.show()