Update the vector plotting example to use newer functionality

It was based on some examples written before the geometry update and had
not been fully adapted to the new API.

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()
+
+# %%
+#
+# Plotting every cell is straight forward for small areas:
 
-# Plot the points
 figure = plt.figure(figsize=(8, 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()