Merge pull request #215 from csiro-coasts/update-vector-examples

Update the vector plotting example to use newer functionality

Author: mx-moth

docs/releases/development.rst

@@ -10,6 +10,6 @@ Next release (in development)
   (:pr:`212`).
 * Add example showing
   :ref:`multiple ways of plotting vector data <sphx_glr_examples_plot-vector-methods.py>`
-  (:pr:`213`).
+  (:pr:`213`, :pr:`215`).
 * Add :attr:`.Grid.centroid_coordinates` attribute
   (:pr:`214`).

examples/plot-vector-methods.py

@@ -7,7 +7,6 @@
 In the `GBR4` tutorial dataset
 the u and v variables contain the x and y components of the vector
 for each cell in the dataset.
-Plotting every cell is straight forward for small areas:
 """
 
 import matplotlib.pyplot as plt
@@ -20,19 +19,23 @@
 from emsarray import plot
 from emsarray.operations import point_extraction
 
-dataset_url = "https://thredds.nci.org.au/thredds/dodsC/fx3/gbr4_bgc_GBR4_H2p0_B3p1_Cfur_Dnrt/gbr4_bgc_simple_2024-01-16.nc"
-dataset = emsarray.open_dataset(dataset_url, decode_timedelta=False)
+dataset = emsarray.tutorial.open_dataset('gbr4')
 
 surface_currents = dataset.ems.select_variables(["temp", "u", "v"]).isel(time=0, k=-1)
+surface_currents.load()
 
-# Plot the points
-figure = plt.figure(figsize=(8, 8), layout="constrained")
+# %%
+#
+# Plotting every cell is straight forward for small areas:
+
+figure = plt.figure(figsize=(7, 8), layout="constrained")
 axes = figure.add_subplot(projection=dataset.ems.data_crs)
-axes.set_title("Surface water temperature and currents near the Whitsundays")
+figure.suptitle("Surface water temperature and currents near the Whitsundays")
 
 temp = surface_currents["temp"]
-temp.attrs['units'] = '°C'
-temp_artist = dataset.ems.make_artist(axes, temp, cmap="coolwarm", clim=(27, 30), edgecolor="face")
+temp_artist = dataset.ems.make_artist(
+    axes, temp,
+    cmap="coolwarm", clim=(24, 27))
 
 current_artist = dataset.ems.make_artist(
     axes, (surface_currents["u"], surface_currents["v"]),
@@ -43,23 +46,18 @@
 
 # Just show a small area over the Whitsundays
 view_box = box(148.7, -20.4, 149.6, -19.8)
-axes.set_extent(plot.bounds_to_extent(view_box.bounds))
 axes.set_aspect("equal", adjustable="datalim")
+axes.set_extent(plot.bounds_to_extent(view_box.bounds))
 
 # %%
+#
 # Plotting the entire dataset this way leads to the current vectors becoming a confusing mess however:
 
-dataset = emsarray.tutorial.open_dataset('gbr4')
-
-surface_currents = dataset.ems.select_variables(["temp", "u", "v"]).isel(time=0, k=-1)
-
-# Plot the points
-figure = plt.figure(figsize=(8, 10), layout="constrained")
+figure = plt.figure(figsize=(7, 10), layout="constrained")
 axes = figure.add_subplot(projection=dataset.ems.data_crs)
-axes.set_title("A bad plot of surface water temperature and currents over the entire reef")
+figure.suptitle("A bad plot of surface water temperature and currents over the entire reef")
 
 temp = surface_currents["temp"]
-temp.attrs['units'] = '°C'
 temp_artist = dataset.ems.make_artist(axes, temp, cmap="coolwarm")
 
 current_artist = dataset.ems.make_artist(
@@ -70,33 +68,38 @@
 plot.add_gridlines(axes)
 
 # Show the entire model domain
-axes.autoscale()
 axes.set_aspect("equal", adjustable="datalim")
+axes.autoscale()
 
 # %%
-# For gridded datasets like this we can sample the current data at regular intervals to display only a subset of the vectors:
-
-dataset = emsarray.tutorial.open_dataset('gbr4')
-
-# Make an empty array of the same shape as the data, then select every nth cell in there
-samples = xarray.DataArray(numpy.full(dataset.ems.grid_size[dataset.ems.default_grid_kind], False))
-samples = dataset.ems.wind(samples)
+#
+# Sampling gridded datasets
+# =========================
+#
+# For datasets on a two dimensional grid like this
+# we can sample the current data at regular intervals
+# to display only a subset of the vectors.
+# The sampled points will follow the geometry of the dataset.
+# For curvilinear datasets such as `GBR4` the vectors will follow the curves of the dataset shape:
+
+# Make an empty array of the same shape as the data, then select every 10th cell in there.
+face_grid = dataset.ems.get_grid(surface_currents["u"])
+samples = xarray.DataArray(numpy.full(face_grid.size, False))
+samples = face_grid.wind(samples)
 samples[::10, ::10] = True
 samples = dataset.ems.ravel(samples)
 
 # Select the (x, y) coordinates and the (u, v) components of the sampled cells
-surface_currents = dataset.ems.select_variables(["temp", "u", "v"]).isel(time=0, k=-1)
-x, y = dataset.ems.face_centres[samples].T
+x, y = face_grid.centroid_coordinates[samples].T
 u = dataset.ems.ravel(surface_currents["u"]).values[samples]
 v = dataset.ems.ravel(surface_currents["v"]).values[samples]
 
 # Plot the points
-figure = plt.figure(figsize=(8, 10), layout="constrained")
+figure = plt.figure(figsize=(7, 10), layout="constrained")
 axes = figure.add_subplot(projection=dataset.ems.data_crs)
-axes.set_title("Surface water temperature and currents across the entire model domain")
+figure.suptitle("Surface water temperature and currents across the entire dataset domain")
 
 temp = surface_currents["temp"]
-temp.attrs['units'] = '°C'
 temp_artist = dataset.ems.make_artist(axes, temp, cmap="coolwarm")
 
 quiver = plt.quiver(
@@ -108,15 +111,21 @@
 plot.add_gridlines(axes)
 
 # Show the entire model domain
-axes.autoscale()
 axes.set_aspect("equal", adjustable="datalim")
+axes.autoscale()
 
 # %%
-# Another approach is to plot vectors at regular points across the domain. This means that the vector locations are not tied to the grid geometry.
-
-dataset = emsarray.tutorial.open_dataset('gbr4')
-
-# Generate a mesh of points across the model domain
+#
+# Sampling the dataset domain
+# ===========================
+#
+# Another approach is to plot vectors at regular points across the dataset domain
+# by sampling at regular intervals.
+# The vector locations are not tied to the grid geometry.
+# This approach will work with unstructured grids unlike the previous method.
+# :func:`.point_extraction.extract_dataframe` can be used for this:
+
+# Generate a mesh of points across the dataset domain
 domain = box(*dataset.ems.bounds)
 x = numpy.arange(domain.bounds[0], domain.bounds[2], 0.4)
 y = numpy.arange(domain.bounds[1], domain.bounds[3], 0.4)
@@ -127,18 +136,15 @@
 })
 
 # Extract the surface current components at these locations
-surface_currents = dataset.ems.select_variables(["temp", "u", "v"]).isel(time=0, k=-1)
-surface_currents.load()
 components = point_extraction.extract_dataframe(
     surface_currents, points, ('x', 'y'), missing_points='drop')
 
 # Plot the points
-figure = plt.figure(figsize=(8, 10), layout="constrained")
+figure = plt.figure(figsize=(7, 10), layout="constrained")
 axes = figure.add_subplot(projection=dataset.ems.data_crs)
-axes.set_title("Surface water temperature and currents across the entire model domain")
+figure.suptitle("Surface water temperature and currents across the entire dataset domain")
 
 temp = surface_currents["temp"]
-temp.attrs['units'] = '°C'
 temp_artist = dataset.ems.make_artist(axes, temp, cmap="coolwarm")
 
 quiver = plt.quiver(
@@ -149,6 +155,5 @@
 plot.add_coast(axes)
 plot.add_gridlines(axes)
 
-axes.autoscale()
 axes.set_aspect("equal", adjustable="datalim")
-
+axes.autoscale()