Add Cartopy background maps

pyproject.toml

@@ -47,6 +47,7 @@ sphinx-autodoc-typehints = { version = "*", optional = true }
 nbsphinx = { version = "*", optional = true }
 pypandoc-binary = { version = "*", optional = true }
 llvmlite = ">=0.42.0"
+cartopy = ">=0.25.0"
 
 [tool.poetry.group.dev.dependencies]
 # only required for dev work on local machine

src/poligrain/plot_map.py

@@ -2,6 +2,10 @@
 
 from __future__ import annotations
 
+import cartopy
+import cartopy.crs
+import cartopy.io.img_tiles as cimgt
+import cartopy.mpl.geoaxes
 import matplotlib.axes
 import matplotlib.patheffects as pe
 import matplotlib.pyplot as plt
@@ -27,6 +31,7 @@ def scatter_lines(
     vmax: float | None = None,
     cmap: (str | Colormap) = "viridis",
     ax: (matplotlib.axes.Axes | None) = None,
+    data_crs: (cartopy.crs.Projection | None) = None,
 ) -> LineCollection:
     """Plot lines as if you would use plt.scatter for points.
 
@@ -66,14 +71,20 @@ def scatter_lines(
     ax : matplotlib.axes.Axes  |  None, optional
         A `Axes` object on which to plot. If not supplied, a new figure with an `Axes`
         will be created. By default None.
+    data_crs : cartopy.crs.Projection | None, optional
+        The coordinate reference system of the data provided. The default is None.
+        In the default case cartopy.crs.PlateCarree will be used when plotting
+        with a `ax` that is a `cartopy.mpl.geoaxes.GeoAxes`. When plotting with `ax`
+        being a normal `matplotlib.axes.Axes` `data_crs` has to be None since the
+        coordinate transformation it implies are not supported by matplotlib.
 
     Returns
     -------
     LineCollection
         _description_
     """
     if ax is None:
-        _, ax = plt.subplots()
+        ax = plt.axes()
 
     x0 = np.atleast_1d(x0)
     y0 = np.atleast_1d(y0)
@@ -92,14 +103,30 @@ def scatter_lines(
             pe.Normal(),
         ]
 
+    style_kwargs = {
+        "linewidth": s,
+        "linestyle": line_style,
+        "capstyle": cap_style,
+        "path_effects": path_effects,
+    }
+
+    # We only allow a data_crs if we have a GeoAxes. And only in this case we
+    # add a `transform` kwarg that is passed to `LineCollection` because it
+    # seems that `transform = None` results in an empty plot when passing it
+    # to `LineCollection` with `ax` being a normal matplotilb `Axes` object.
+    if isinstance(ax, cartopy.mpl.geoaxes.GeoAxes):
+        if data_crs is None:
+            data_crs = cartopy.crs.PlateCarree()
+        style_kwargs["transform"] = data_crs
+    elif data_crs is not None:
+        msg = "data_crs has to be None if `ax` is not a cartopy.mpl.geoaxes.GeoAxes"
+        raise ValueError(msg)
+
     if data is None:
         lines = LineCollection(
             [((x0[i], y0[i]), (x1[i], y1[i])) for i in range(len(x0))],
-            linewidth=s,
-            linestyles=line_style,
-            capstyle=cap_style,
             color=c,
-            path_effects=path_effects,
+            **style_kwargs,
         )
 
     else:
@@ -112,10 +139,7 @@ def scatter_lines(
             [((x0[i], y0[i]), (x1[i], y1[i])) for i in range(len(x0))],
             norm=norm,
             cmap=cmap,
-            linewidth=s,
-            linestyles=line_style,
-            capstyle=cap_style,
-            path_effects=path_effects,
+            **style_kwargs,
         )
         lines.set_array(data)
 
@@ -139,6 +163,8 @@ def plot_lines(
     line_style: str = "-",
     cap_style: str = "round",
     ax: (matplotlib.axes.Axes | None) = None,
+    background_map: (str | None) = None,
+    projection: (cartopy.crs.Projection | None) = None,
 ) -> LineCollection:
     """Plot paths of line-based sensors like CMLs.
 
@@ -178,15 +204,16 @@ def plot_lines(
     ax : matplotlib.axes.Axes  |  None, optional
         A `Axes` object on which to plot. If not supplied, a new figure with an `Axes`
         will be created. By default None.
+    background_map : str | None, optional
+        Type of background map.
+    projection : cartopy.crs.Projection | None, optional
+        The map projection to be used.
 
     Returns
     -------
     LineCollection
 
     """
-    if ax is None:
-        _, ax = plt.subplots()
-
     try:
         color_data = cmls.data
         if len(color_data.shape) != 1:
@@ -210,6 +237,9 @@ def plot_lines(
         y0_name = "site_0_y"
         y1_name = "site_1_y"
 
+    if ax is None:
+        ax = set_up_axes(background_map=background_map, projection=projection)
+
     return scatter_lines(
         x0=cmls[x0_name].values,
         y0=cmls[y0_name].values,
@@ -228,13 +258,75 @@ def plot_lines(
     )
 
 
+def set_up_axes(background_map=None, projection=None, extent=None):
+    """Create and configure matplotlib axes based on background map type.
+
+    Set up the plot axes using Cartopy or standard matplotlib. Supports
+    different types of background maps including stock images, OpenStreetMap,
+    and Natural Earth datasets with optional geographic projections and extents.
+
+    Parameters
+    ----------
+        background_map : str or None
+            Type of background map to use. Can be 'stock' for built-in stock image,
+            'OSM' for OpenStreetMap or 'NE' for Natural Earth features.
+            If not provided creates mapltolib Axes object without any background.
+
+        projection : cartopy.crs.Projection, optional
+            Projection to use for the plot. For 'OSM' map background it is automatically
+            set. For other map backgrounds it defaults to PlateCarree() if not provided.
+
+        extent : list-like or None, optional
+            Geographic bounding box [lon_min, lat_min, lon_max, lat_max].
+            Defaults to None.
+
+    Returns
+    -------
+        matplotlib.axes.Axes or cartopy.mpl.geoaxes.GeoAxes object with appropriate
+        projection and background.
+    """
+    if background_map is None and projection is None:
+        ax = plt.axes()
+    elif background_map is None and projection is not None:
+        ax = plt.axes(projection=projection)
+    elif background_map == "stock":
+        if projection is None:
+            projection = cartopy.crs.PlateCarree()
+        ax = plt.axes(projection=projection)
+        if extent is not None:
+            ax.set_extent(extent)
+        ax.stock_img()
+    elif background_map == "OSM":
+        request = cimgt.OSM()
+        ax = plt.axes(projection=request.crs)
+        if extent is not None:
+            ax.set_extent(extent)
+        ax.add_image(request, 10)
+    elif background_map == "NE":
+        if projection is None:
+            projection = cartopy.crs.PlateCarree()
+        ax = plt.axes(projection=projection)
+        if extent is not None:
+            ax.set_extent(extent)
+        ax.add_feature(cartopy.feature.OCEAN, facecolor="lightblue")
+        ax.add_feature(cartopy.feature.LAND, facecolor="lightgrey")
+        ax.add_feature(cartopy.feature.LAKES, facecolor="lightblue", linewidth=0.00001)
+        ax.add_feature(cartopy.feature.BORDERS, linewidth=0.3)
+        ax.coastlines(resolution="10m", linewidth=0.3)
+    else:
+        msg = f"unsupported value of background_map {background_map}"
+        raise ValueError(msg)
+    return ax
+
+
 def plot_plg(
     da_grid=None,
     da_cmls=None,
     da_gauges=None,
     vmin=None,
     vmax=None,
     cmap="turbo",
+    alpha=1,
     ax=None,
     use_lon_lat=True,
     edge_color="k",
@@ -246,6 +338,9 @@ def plot_plg(
     colorbar_label="",
     kwargs_cmls_plot=None,
     kwargs_gauges_plot=None,
+    background_map: (str | None) = None,
+    projection: (cartopy.crs.Projection | None) = None,
+    extent: (list | None) = None,
 ):
     """Plot point, line and grid data.
 
@@ -277,6 +372,8 @@ def plot_plg(
         it will be derived individually for each data source when plotting.
     cmap : str, optional
         cmap for all three data sources, by default "turbo"
+    alpha : float, optional
+        Alpha values used for the gridded dataset.
     ax : _type_, optional
         Axes object from matplotlib, by default None which will create a new
         figure and return the Axes object.
@@ -311,26 +408,36 @@ def plot_plg(
         `plot_lines` for supported kwargs.
     kwargs_gauges_plot : dict or None, optional
         kwargs to be passed to plt.scatter, by default None.
-
+    data_crs........
     """
     if kwargs_cmls_plot is None:
         kwargs_cmls_plot = {}
     if kwargs_gauges_plot is None:
         kwargs_gauges_plot = {}
+    kwargs_pcolormesh = {}
 
     if ax is None:
-        _, ax = plt.subplots()
+        ax = set_up_axes(
+            background_map=background_map,
+            projection=projection,
+            extent=extent,
+        )
 
     if use_lon_lat:
         grid_x_name = "lon"
         grid_y_name = "lat"
         point_x_name = "lon"
         point_y_name = "lat"
+        data_crs = cartopy.crs.PlateCarree()
+        if isinstance(ax, cartopy.mpl.geoaxes.GeoAxes):
+            kwargs_gauges_plot["transform"] = data_crs
+            kwargs_pcolormesh["transform"] = data_crs
     else:
         grid_x_name = "x_grid"
         grid_y_name = "y_grid"
         point_x_name = "x"
         point_y_name = "y"
+        data_crs = None
 
     plotted_objects_for_cmap = []
     if da_grid is not None:
@@ -342,6 +449,8 @@ def plot_plg(
             cmap=cmap,
             ax=ax,
             add_colorbar=False,
+            alpha=alpha,
+            **kwargs_pcolormesh,
         )
         plotted_objects_for_cmap.append(pc)
     if da_cmls is not None:
@@ -356,6 +465,8 @@ def plot_plg(
             line_color=kwargs_cmls_plot.pop("line_color", line_color),
             pad_color=kwargs_cmls_plot.pop("edge_color", edge_color),
             pad_width=kwargs_cmls_plot.pop("edge_width", edge_width),
+            background_map=None,
+            projection=None,
             **kwargs_cmls_plot,
         )
         # only add line_collection in case we really want to apply a cmap