Add parallel example script using ZarrMonitor (#126)

This example script serves both as usage documentation, and as a basic parallel test of the ZarrMonitor.

external/fv3gfs-util/Makefile

@@ -41,7 +41,7 @@ test: ## run tests quickly with the default Python
 
 test_mpi:
 	mpirun -n 6 --allow-run-as-root --mca btl_vader_single_copy_mechanism none --oversubscribe pytest $(PYTEST_ARGS) tests/mpi
-	$(MAKE) -C examples/runfiles
+	$(MAKE) -C examples/mpi
 
 lint:
 	black --diff --check $(PYTHON_FILES) $(PYTHON_INIT_FILES)
@@ -50,6 +50,9 @@ lint:
 	flake8 --ignore=F401 $(PYTHON_INIT_FILES)
 	@echo "LINTING SUCCESSFUL"
 
+clean:
+	$(MAKE) -c examples/mpi clean
+
 reformat:
 	black $(PYTHON_FILES) $(PYTHON_INIT_FILES)
 
@@ -58,4 +61,4 @@ docs: ## generate Sphinx HTML documentation, including API docs
 	$(MAKE) -C docs html
 	$(BROWSER) docs/_build/html/index.html
 
-.PHONY: docs test lint reformat test_mpi coverage help
+.PHONY: docs clean test lint reformat test_mpi coverage help

external/fv3gfs-util/examples/mpi/.gitignore

@@ -0,0 +1 @@
+output

external/fv3gfs-util/examples/mpi/Makefile

@@ -0,0 +1,14 @@
+
+MPI_ENV_VARS=PMIX_MCA_gds=hash
+
+all: global_timings zarr_monitor
+
+global_timings:
+	$(MPI_ENV_VARS) mpirun -n 4 python -m mpi4py global_timings.py
+
+zarr_monitor:
+	$(MPI_ENV_VARS) mpirun -n 6 python -m mpi4py zarr_monitor.py
+
+clean:
+	$(RM) -r output/*
+	touch output/.gitkeep

external/fv3gfs-util/examples/mpi/zarr_monitor.py

@@ -0,0 +1,41 @@
+import fv3gfs.util
+from mpi4py import MPI
+import numpy as np
+from datetime import timedelta
+import cftime
+import zarr
+
+OUTPUT_PATH = "output/zarr_monitor.zarr"
+
+
+def get_example_state(time):
+    sizer = fv3gfs.util.SubtileGridSizer(
+        nx=48, ny=48, nz=70, n_halo=3, extra_dim_lengths={}
+    )
+    allocator = fv3gfs.util.QuantityFactory(sizer, np)
+    air_temperature = allocator.zeros(
+        [fv3gfs.util.X_DIM, fv3gfs.util.Y_DIM, fv3gfs.util.Z_DIM], units="degK"
+    )
+    air_temperature.view[:] = np.random.randn(*air_temperature.extent)
+    return {"time": time, "air_temperature": air_temperature}
+
+
+if __name__ == "__main__":
+    size = MPI.COMM_WORLD.Get_size()
+    # assume square tile faces
+    ranks_per_edge = int((size // 6) ** 0.5)
+    layout = (ranks_per_edge, ranks_per_edge)
+
+    store = zarr.storage.DirectoryStore(OUTPUT_PATH)
+    partitioner = fv3gfs.util.CubedSpherePartitioner(
+        fv3gfs.util.TilePartitioner(layout)
+    )
+    monitor = fv3gfs.util.ZarrMonitor(store, partitioner, mpi_comm=MPI.COMM_WORLD)
+
+    time = cftime.DatetimeJulian(2020, 1, 1)
+    timestep = timedelta(hours=1)
+
+    for i in range(10):
+        state = get_example_state(time)
+        monitor.store(state)
+        time += timestep