Fix plt motionfields (#166)

* Fix interface to plot_geography

* Fix docstrings

* Fix plotting functions for motion fields

* Remove pixelsize as a required argument

* Add missing 'lw' argument

* Add FutureWarning and internally rename 'map' to 'plot_map'

* Rename map argument

- Since map was an argument and not a keyword, it can be renamed safely.
- The keywords arguments were explicitly declared in the plot_geography calls.

* Remove unused shape argument from plot_geography

* Speed up tests

Co-authored-by: Andres Perez Hortal <andresperezcba@gmail.com>

Author: dnerini

pysteps/tests/test_plt_motionfields.py

@@ -0,0 +1,113 @@
+# -*- coding: utf-8 -*-
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+
+from pysteps import motion
+from pysteps.visualization import plot_precip_field, quiver, streamplot
+from pysteps.tests.helpers import get_precipitation_fields
+
+
+arg_names_quiver = (
+    "source",
+    "map",
+    "drawlonlatlines",
+    "lw",
+    "axis",
+    "step",
+    "quiver_kwargs",
+    "upscale",
+)
+
+arg_values_quiver = [
+    (None, None, False, 0.5, "off", 10, None, None),
+    ("bom", None, False, 0.5, "on", 10, None, 4000),
+    ("bom", "cartopy", True, 0.5, "on", 10, None, 4000),
+    ("mch", "cartopy", False, 0.5, "on", 20, None, 2000),
+    ("bom", "basemap", False, 0.5, "off", 10, None, 4000),
+]
+
+
+@pytest.mark.parametrize(arg_names_quiver, arg_values_quiver)
+def test_visualization_motionfields_quiver(
+    source, map, drawlonlatlines, lw, axis, step, quiver_kwargs, upscale,
+):
+
+    if map == "cartopy":
+        pytest.importorskip("cartopy")
+    elif map == "basemap":
+        pytest.importorskip("basemap")
+
+    if source is not None:
+        fields, geodata = get_precipitation_fields(0, 2, False, True, upscale, source)
+        ax = plot_precip_field(fields[-1], map=map, geodata=geodata,)
+        oflow_method = motion.get_method("LK")
+        UV = oflow_method(fields)
+
+    else:
+        shape = (100, 100)
+        geodata = None
+        ax = None
+        u = np.ones(shape[1]) * shape[0]
+        v = np.arange(0, shape[0])
+        U, V = np.meshgrid(u, v)
+        UV = np.concatenate([U[None, :], V[None, :]])
+
+    __ = quiver(UV, ax, map, geodata, drawlonlatlines, lw, axis, step, quiver_kwargs)
+
+
+arg_names_streamplot = (
+    "source",
+    "map",
+    "drawlonlatlines",
+    "lw",
+    "axis",
+    "streamplot_kwargs",
+    "upscale",
+)
+
+arg_values_streamplot = [
+    (None, None, False, 0.5, "off", None, None),
+    ("bom", None, False, 0.5, "on", None, 4000),
+    ("bom", "cartopy", True, 0.5, "on", {"density": 0.1}, 4000),
+]
+
+
+@pytest.mark.parametrize(arg_names_streamplot, arg_values_streamplot)
+def test_visualization_motionfields_streamplot(
+    source, map, drawlonlatlines, lw, axis, streamplot_kwargs, upscale,
+):
+
+    if map == "cartopy":
+        pytest.importorskip("cartopy")
+    elif map == "basemap":
+        pytest.importorskip("basemap")
+
+    if source is not None:
+        fields, geodata = get_precipitation_fields(0, 2, False, True, upscale, source)
+        ax = plot_precip_field(fields[-1], map=map, geodata=geodata,)
+        oflow_method = motion.get_method("LK")
+        UV = oflow_method(fields)
+
+    else:
+        shape = (100, 100)
+        geodata = None
+        ax = None
+        u = np.ones(shape[1]) * shape[0]
+        v = np.arange(0, shape[0])
+        U, V = np.meshgrid(u, v)
+        UV = np.concatenate([U[None, :], V[None, :]])
+
+    __ = streamplot(UV, ax, map, geodata, drawlonlatlines, lw, axis, streamplot_kwargs)
+
+
+if __name__ == "__main__":
+
+    for i, args in enumerate(arg_values_quiver):
+        test_visualization_motionfields_quiver(*args)
+        plt.show()
+
+    for i, args in enumerate(arg_values_streamplot):
+        test_visualization_motionfields_streamplot(*args)
+        plt.show()

pysteps/visualization/animations.py

@@ -125,6 +125,14 @@ def animate(
     ax : fig axes
         Figure axes. Needed if one wants to add e.g. text inside the plot.
     """
+    if map is not None:
+        FutureWarning(
+            "'map' argument will be renamed to 'plot_map' in 1.4.0. Use 'plot_map' to silence this warning."
+        )
+        plot_map = map
+    else:
+        plot_map = kwargs.pop("plot_map", None)
+
     if timestamps is not None:
         startdate_str = timestamps[-1].strftime("%Y%m%d%H%M")
     else:
@@ -172,7 +180,7 @@ def animate(
                         ax = st.plt.plot_precip_field(
                             P_obs,
                             type="prob",
-                            map=map,
+                            plot_map=plot_map,
                             geodata=geodata,
                             units=units,
                             probthr=prob_thr,
@@ -183,7 +191,7 @@ def animate(
                         title += "Observed Rainfall"
                         ax = st.plt.plot_precip_field(
                             R_obs[i, :, :],
-                            map=map,
+                            plot_map=plot_map,
                             geodata=geodata,
                             units=units,
                             colorscale=colorscale,
@@ -234,7 +242,7 @@ def animate(
                         ax = st.plt.plot_precip_field(
                             P,
                             type="prob",
-                            map=map,
+                            plot_map=plot_map,
                             geodata=geodata,
                             units=units,
                             probthr=prob_thr,
@@ -248,7 +256,7 @@ def animate(
                         )
                         ax = st.plt.plot_precip_field(
                             EM,
-                            map=map,
+                            plot_map=plot_map,
                             geodata=geodata,
                             units=units,
                             title=title,
@@ -260,7 +268,7 @@ def animate(
                         title += "Forecast Rainfall"
                         ax = st.plt.plot_precip_field(
                             R_fct[n, i - n_obs, :, :],
-                            map=map,
+                            plot_map=plot_map,
                             geodata=geodata,
                             units=units,
                             title=title,

pysteps/visualization/basemaps.py

@@ -40,9 +40,7 @@
 from . import utils
 
 
-def plot_geography(
-    map, proj4str, extent, shape=None, lw=0.5, drawlonlatlines=False, **kwargs
-):
+def plot_geography(plot_map, proj4str, extent, lw=0.5, drawlonlatlines=False, **kwargs):
     """
     Plot geographical map using either cartopy_ or basemap_ in a chosen projection.
 
@@ -54,15 +52,12 @@ def plot_geography(
 
     Parameters
     ----------
-    map : {'cartopy', 'basemap'}
+    plot_map : {'cartopy', 'basemap'}
         The type of basemap.
     proj4str : str
         The PROJ.4-compatible projection string.
     extent: scalars (left, right, bottom, top)
         The bounding box in proj4str coordinates.
-    shape: scalars (n_rows, n_cols)
-        The dimensions of the image.
-        Only used if needing a fallback projection.
     lw: float, optional
         Linewidth of the map (administrative boundaries and coastlines).
     drawlonlatlines : bool, optional
@@ -94,34 +89,30 @@ def plot_geography(
         Whether the projection allows plotting a regular grid.
         Returns False in case a fall-back projection is used.
     """
-    if map is not None and map not in ["basemap", "cartopy"]:
+
+    if plot_map not in ["basemap", "cartopy"]:
         raise ValueError(
-            "unknown map method %s: must be" + " 'basemap' or 'cartopy'" % map
+            "unknown map method %s: must be" + " 'basemap' or 'cartopy'" % plot_map
         )
-    if map == "basemap" and not basemap_imported:
+    if plot_map == "basemap" and not basemap_imported:
         raise MissingOptionalDependency(
             "map='basemap' option passed to plot_geography function "
             "but the basemap package is not installed"
         )
-    if map == "cartopy" and not cartopy_imported:
+    if plot_map == "cartopy" and not cartopy_imported:
         raise MissingOptionalDependency(
             "map='cartopy' option passed to plot_geography function "
             "but the cartopy package is not installed"
         )
-    if map is not None and not pyproj_imported:
+    if plot_map is not None and not pyproj_imported:
         raise MissingOptionalDependency(
             "map!=None option passed to plot_geography function "
             "but the pyproj package is not installed"
         )
 
-    if map == "basemap":
+    if plot_map == "basemap":
         basemap_resolution = kwargs.get("resolution", "l")
         basemap_scale_args = kwargs.get("scale_args", None)
-    if map == "cartopy":
-        cartopy_scale = kwargs.get("scale", "50m")
-        cartopy_subplot = kwargs.get("subplot", (1, 1, 1))
-
-    if map == "basemap":
         pr = pyproj.Proj(proj4str)
         x1, x2, y1, y2 = extent[0], extent[1], extent[2], extent[3]
         ll_lon, ll_lat = pr(x1, y1, inverse=True)
@@ -139,6 +130,8 @@ def plot_geography(
         if basemap_scale_args is not None:
             ax.drawmapscale(*basemap_scale_args, fontsize=6, yoffset=10000)
     else:
+        cartopy_scale = kwargs.get("scale", "50m")
+        cartopy_subplot = kwargs.get("subplot", (1, 1, 1))
         crs = utils.proj4_to_cartopy(proj4str)
 
         ax = plot_map_cartopy(