Merge pull request #191 from csiro-coasts/transect-data-loading

Load transect data frame by frame

Author: mx-moth

docs/releases/development.rst

@@ -15,3 +15,8 @@ Next release (in development)
 * Fixed an issue with UGrid datasets when some of the mesh topology variables
   are not present in Dataset.data_vars as they are detected as coordinate variables
   (:issue:`159`, :pr:`188`).
+* Load data frame by frame when plotting a transect animation.
+  This has a backwards incompatible effect of not examining the whole data array to generate a reasonable *clim*.
+  Only the first frame is used to generate the clim to avoid loading more data than required.
+  The new clim parameter allows users to specify the data limits explicitly if this is insufficient
+  (:pr:`191`).

src/emsarray/transect.py

@@ -562,6 +562,7 @@ def plot_on_figure(
         clamp_to_surface: bool = True,
         bathymetry: xarray.DataArray | None = None,
         cmap: str | Colormap = 'jet',
+        clim: tuple[float, float] | None = None,
         ocean_floor_colour: str = 'black',
         landmarks: list[Landmark] | None = None,
     ) -> None:
@@ -607,6 +608,7 @@ def plot_on_figure(
             clamp_to_surface=clamp_to_surface,
             bathymetry=bathymetry,
             cmap=cmap,
+            clim=clim,
             ocean_floor_colour=ocean_floor_colour,
             landmarks=landmarks,
         )
@@ -622,6 +624,7 @@ def animate_on_figure(
         clamp_to_surface: bool = True,
         bathymetry: xarray.DataArray | None = None,
         cmap: str | Colormap = 'jet',
+        clim: tuple[float, float] | None = None,
         ocean_floor_colour: str = 'black',
         landmarks: list[Landmark] | None = None,
         coordinate: xarray.DataArray | None = None,
@@ -652,6 +655,7 @@ def animate_on_figure(
         if coordinate is None:
             coordinate = self.convention.time_coordinate
         coordinate_indexes = numpy.arange(coordinate.size)
+        animation_dimension = coordinate.dims[0]
 
         coordinate_callable: Callable[[Any], str]
         if title is None:
@@ -667,14 +671,17 @@ def animate_on_figure(
         else:
             coordinate_callable = title
 
-        axes, collection, data_array = self._plot_on_figure(
+        first_frame = data_array.isel({animation_dimension: 0})
+        first_frame.load()
+        axes, collection, _prepared_frame = self._plot_on_figure(
             figure=figure,
-            data_array=data_array,
+            data_array=first_frame,
             title=None,
             trim_nans=trim_nans,
             clamp_to_surface=clamp_to_surface,
             bathymetry=bathymetry,
             cmap=cmap,
+            clim=clim,
             ocean_floor_colour=ocean_floor_colour,
             landmarks=landmarks,
         )
@@ -686,7 +693,10 @@ def animate(index: int) -> Iterable[Artist]:
             axes.set_title(coordinate_callable(coordinate_value))
             changes.append(axes)
 
-            collection.set_array(data_array[index].values.flatten())
+            frame_data = data_array.isel({animation_dimension: index})
+            frame_data.load()
+            prepared_data = self.prepare_data_array_for_transect(frame_data)
+            collection.set_array(prepared_data.values.flatten())
             changes.append(collection)
             return changes
 
@@ -711,6 +721,7 @@ def _plot_on_figure(
         clamp_to_surface: bool = True,
         bathymetry: xarray.DataArray | None = None,
         cmap: str | Colormap = 'jet',
+        clim: tuple[float, float] | None = None,
         ocean_floor_colour: str = 'black',
         landmarks: list[Landmark] | None = None,
     ) -> tuple[Axes, PolyCollection, xarray.DataArray]:
@@ -724,6 +735,7 @@ def _plot_on_figure(
         depth = transect_dataset.coords['depth']
         distance_bounds = transect_dataset.data_vars['distance_bounds']
 
+        data_array = data_array.load()
         data_array = self.prepare_data_array_for_transect(data_array)
 
         positive_down = depth.attrs['positive'] == 'down'
@@ -769,12 +781,11 @@ def _plot_on_figure(
             cmap = colormaps[cmap].copy()
         cmap.set_bad(ocean_floor_colour)
 
-        if data_array.size != 0:
+        # Find a min/max from the data if clim isn't provided and the data array is not empty.
+        # An empty data array happens when the transect line does not intersect
+        # the dataset geometry.
+        if clim is None and data_array.size != 0:
             clim = (numpy.nanmin(data_array), numpy.nanmax(data_array))
-        else:
-            # An empty data array happens when the transect line does not
-            # intersect the dataset geometry.
-            clim = None
 
         collection = self.make_poly_collection(cmap=cmap, clim=clim, edgecolor='face')
         axes.add_collection(collection)