fix examples

clouddrift/plotting.py

@@ -61,58 +61,71 @@
     Examples
     --------
 
-    Plot the first 100 trajectories from the gdp1h dataset, assigning
-    a different color to each trajectory:
+    Load 100 trajectories from the gdp1h dataset for the examples.
 
     >>> from clouddrift import datasets
+    >>> from clouddrift.ragged import subset
+    >>> from clouddrift.plotting import plot_ragged
+    >>> import numpy as np
     >>> import matplotlib.pyplot as plt
+    >>> from mpl_toolkits.axes_grid1 import make_axes_locatable
     >>> ds = datasets.gdp1h()
-    >>> ds = subset(ds, {"ID": ds.ID[:100].values}).load()
+    >>> ds = subset(ds, {"id": ds.id[:100].values}).load()
+
+    Plot the trajectories, assigning a different color to each trajectory:
+
     >>> fig = plt.figure()
     >>> ax = fig.add_subplot(1, 1, 1)
-
-    >>> plot_ragged(
+    >>> l = plot_ragged(
     >>>     ax,
     >>>     ds.lon,
     >>>     ds.lat,
     >>>     ds.rowsize,
     >>>     colors=np.arange(len(ds.rowsize))
     >>> )
+    >>> divider = make_axes_locatable(ax)
+    >>> cax = divider.append_axes('right', size='3%', pad=0.05)
+    >>> fig.colorbar(l, cax=cax)
 
     To plot the same trajectories, but assigning a different color to each
-    observation and specifying a colormap:
+    observation based on time and specifying a colormap:
 
     >>> fig = plt.figure()
     >>> ax = fig.add_subplot(1, 1, 1)
     >>> time = [v.astype(np.int64) / 86400 / 1e9 for v in ds.time.values]
-    >>> lc = plot_ragged(
+    >>> l = plot_ragged(
     >>>     ax,
     >>>     ds.lon,
     >>>     ds.lat,
     >>>     ds.rowsize,
     >>>     colors=np.floor(time),
     >>>     cmap="inferno"
     >>> )
-    >>> fig.colorbar(lc[0])
-    >>> ax.set_xlim([-180, 180])
-    >>> ax.set_ylim([-90, 90])
+    >>> divider = make_axes_locatable(ax)
+    >>> cax = divider.append_axes('right', size="3%", pad=0.05)
+    >>> fig.colorbar(l, cax=cax)
 
     Finally, to plot the same trajectories, but using a cartopy
     projection:
 
     >>> import cartopy.crs as ccrs
     >>> fig = plt.figure()
     >>> ax = fig.add_subplot(1, 1, 1, projection=ccrs.Mollweide())
-    >>> time = [v.astype(np.int64) / 86400 / 1e9 for v in ds.time.values]
-    >>> lc = plot_ragged(
+    >>> l = plot_ragged(
     >>>     ax,
     >>>     ds.lon,
     >>>     ds.lat,
     >>>     ds.rowsize,
     >>>     colors=np.arange(len(ds.rowsize)),
     >>>     transform=ccrs.PlateCarree(),
-    >>>     cmap=cmocean.cm.ice,
+    >>>     cmap="Blues",
     >>> )
+    >>> ax.set_extent([-180, 180, -90, 90])
+    >>> ax.coastlines()
+    >>> ax.gridlines(draw_labels=True)
+    >>> divider = make_axes_locatable(ax)
+    >>> cax = divider.append_axes('right', size="3%", pad=0.25, axes_class=plt.Axes)
+    >>> fig.colorbar(l, cax=cax)
 
     Raises
     ------
@@ -161,18 +174,21 @@
         raise ValueError("shape colors must match the shape of lon/lat or rowsize.")
 
     # define a colormap
-    cmap = kwargs.pop("cmap", cm.viridis)
+    if isinstance(cmap := kwargs.pop("cmap", cm.viridis), str):
+        cmap = plt.get_cmap(cmap)
 
     # define a normalization obtain uniform colors
     # for the sequence of lines or LineCollection
     norm = kwargs.pop(
         "norm", mcolors.Normalize(vmin=np.nanmin(colors), vmax=np.nanmax(colors))
     )
 
+    # create Mappable for colorbar
+    cb = plt.cm.ScalarMappable(norm=norm, cmap=cmap)
+
     mpl_plot = True if colors is None or len(colors) == len(rowsize) else False
     traj_idx = rowsize_to_index(rowsize)
 
-    lines = []
     for i in range(len(rowsize)):
         lon_i, lat_i = (
             longitude[traj_idx[i] : traj_idx[i + 1]],
@@ -184,7 +200,7 @@
             end = start + length
 
             if mpl_plot:
-                line = ax.plot(
+                ax.plot(
                     lon_i[start:end],
                     lat_i[start:end],
                     c=cmap(norm(colors[i])) if colors is not None else None,
@@ -201,18 +217,17 @@
                         lat_i[start + 1 : end],
                     ]
                 ).reshape(-1, 2, 2)
-                line = LineCollection(segments, cmap=cmap, norm=norm, *args, **kwargs)
-                line.set_array(
+                lc = LineCollection(segments, cmap=cmap, norm=norm, *args, **kwargs)
+                lc.set_array(
                     # color of a segment is the average of its two data points
                     np.convolve(colors_i[start:end], [0.5, 0.5], mode="valid")
                 )
-                ax.add_collection(line)
+                ax.add_collection(lc)
 
             start = end
-            lines.append(line)
 
     # set axis limits
     ax.set_xlim([np.min(longitude), np.max(longitude)])
     ax.set_ylim([np.min(latitude), np.max(latitude)])
 
-    return lines
+    return cb

tests/adapters/utils_tests.py

@@ -224,7 +224,7 @@ def test_raises_on_any_exception_and_cleanup(self):
                 patch("clouddrift.adapters.utils.open", self.open_mock),
                 patch("clouddrift.adapters.utils.concurrent.futures", futures_mock),
                 patch("clouddrift.adapters.utils.requests", self.requests_mock),
-                patch("clouddrift.adapters.utils.os", os_mock)
+                patch("clouddrift.adapters.utils.os", os_mock),
             ]
         ) as _:
             self.assertRaises(

tests/plotting_tests.py

@@ -42,15 +42,15 @@ def test_plot_colored_trajectory(self):
         l = plot_ragged(
             ax, self.lon, self.lat, self.rowsize, colors=np.arange(len(self.rowsize))
         )
-        self.assertIsInstance(l, list)
+        self.assertIsInstance(l, plt.cm.ScalarMappable)
 
     def test_plot_colored_datapoints(self):
         fig = plt.figure()
         ax = fig.add_subplot(1, 1, 1)
         l = plot_ragged(
             ax, self.lon, self.lat, self.rowsize, colors=np.arange(len(self.lat))
         )
-        self.assertIsInstance(l, list)
+        self.assertIsInstance(l, plt.cm.ScalarMappable)
 
     def test_plot_color_wrong_dimension(self):
         fig = plt.figure()
@@ -82,7 +82,7 @@ def test_plot_cartopy(self):
             colors=np.arange(len(self.rowsize)),
             transform=ccrs.PlateCarree(),
         )
-        self.assertIsInstance(l, list)
+        self.assertIsInstance(l, plt.cm.ScalarMappable)
 
     def test_plot_segments(self):
         self.lon = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
@@ -99,8 +99,7 @@ def test_plot_segments(self):
             colors=np.arange(len(self.rowsize)),
             transform=ccrs.PlateCarree(),
         )
-        self.assertIsInstance(l, list)
-        self.assertEqual(len(l), 3)
+        self.assertIsInstance(l, plt.cm.ScalarMappable)
 
     def test_plot_segments_split(self):
         self.lon = [-170, -175, -180, 175, 170]
@@ -117,8 +116,7 @@ def test_plot_segments_split(self):
             colors=np.arange(len(self.rowsize)),
             transform=ccrs.PlateCarree(),
         )
-        self.assertIsInstance(l, list)
-        self.assertEqual(len(l), 2)
+        self.assertIsInstance(l, plt.cm.ScalarMappable)
 
     def test_plot_segments_split_domain(self):
         self.lon = [-1, -2, -3, 3, 2, 1]
@@ -136,8 +134,7 @@ def test_plot_segments_split_domain(self):
             transform=ccrs.PlateCarree(),
             tolerance=5,
         )
-        self.assertIsInstance(l, list)
-        self.assertEqual(len(l), 2)
+        self.assertIsInstance(l, plt.cm.ScalarMappable)
 
     def test_matplotlib_not_installed(self):
         try: