Merge pull request #399 from gwmod/plot_improvements

Add cross-section plot utils

Author: dbrakenhoff

docs/examples/13_plot_methods.ipynb

@@ -171,7 +171,9 @@
    "id": "fba12db0",
    "metadata": {},
    "source": [
-    "With the DatasetCrossSection in `nlmod` it is also possible to plot the layers according to the official colors of REGIS, to plot the layer names on the plot, or to plot the model grid in the cross-section. An example is shown in the plot below."
+    "With the DatasetCrossSection in `nlmod` it is also possible to plot the layers according to the official colors of REGIS, to plot the layer names on the plot, or to plot the model grid in the cross-section. An example is shown in the plot below.\n",
+    "\n",
+    "The location of the cross-section and the cross-section labels can be added using `nlmod.plot.inset_map()` and `nlmod.plot.add_xsec_line_and_labels()`."
    ]
   },
   {
@@ -185,7 +187,9 @@
     "dcs = DatasetCrossSection(ds, line=line, zmin=-200, zmax=10, ax=ax)\n",
     "colors = nlmod.read.regis.get_legend()\n",
     "dcs.plot_layers(colors=colors, min_label_area=1000)\n",
-    "dcs.plot_grid(vertical=False, linewidth=0.5);"
+    "dcs.plot_grid(vertical=False, linewidth=0.5)\n",
+    "mapax = nlmod.plot.inset_map(ax, ds.extent)\n",
+    "nlmod.plot.add_xsec_line_and_labels(line, ax, mapax)"
    ]
   },
   {

nlmod/plot/__init__.py

@@ -14,9 +14,11 @@
 )
 from .plotutil import (
     add_background_map,
+    add_xsec_line_and_labels,
     colorbar_inside,
     get_figsize,
     get_map,
+    inset_map,
     rd_ticks,
     rotate_yticklabels,
     title_inside,

nlmod/plot/dcs.py

@@ -566,9 +566,9 @@ def animate(
         elif "units" in da.attrs:
             cbar.set_label(da.units)
 
-        if da.time.dtype.kind == 'M':
+        if da.time.dtype.kind == "M":
             t = pd.Timestamp(da.time.values[iper]).strftime(date_fmt)
-        elif da.time.dtype.kind == 'O':
+        elif da.time.dtype.kind == "O":
             t = da.time.values[iper].strftime(date_fmt)
         else:
             t = f"{da.time.values[iper]} {da.time.time_units}"
@@ -584,9 +584,9 @@ def update(iper, pc, title):
             pc.set_array(array)
 
             # update title
-            if da.time.dtype.kind == 'M':
+            if da.time.dtype.kind == "M":
                 t = pd.Timestamp(da.time.values[iper]).strftime(date_fmt)
-            elif da.time.dtype.kind == 'O':
+            elif da.time.dtype.kind == "O":
                 t = da.time.values[iper].strftime(date_fmt)
             else:
                 t = f"{da.time.values[iper]} {da.time.time_units}"

nlmod/plot/plotutil.py

@@ -1,7 +1,11 @@
+from typing import Callable, Optional, Union
+
 import matplotlib.pyplot as plt
 import numpy as np
+from matplotlib import patheffects
 from matplotlib.patches import Polygon
 from matplotlib.ticker import FuncFormatter, MultipleLocator
+from shapely import LineString
 
 from ..dims.grid import get_affine_mod_to_world
 from ..epsg28992 import EPSG_28992
@@ -269,3 +273,135 @@ def title_inside(
         bbox=bbox,
         **kwargs,
     )
+
+
+def inset_map(
+    ax: plt.Axes,
+    extent: Union[tuple[float], list[float]],
+    axes_bounds: Union[tuple[float], list[float]] = (0.63, 0.025, 0.35, 0.35),
+    anchor: str = "SE",
+    provider: Optional[str] = "nlmaps.water",
+    add_to_plot: Optional[list[Callable]] = None,
+):
+    """Add an inset map to an axes.
+
+    Parameters
+    ----------
+    ax : matplotlib.Axes
+        The axes to add the inset map to.
+    extent : list of 4 floats
+        The extent of the inset map.
+    axes_bounds : list of 4 floats, optional
+        The bounds (left, right, width height) of the inset axes, default
+        is [0.63, 0.025, 0.35, 0.35]. This is rescaled according to the extent of
+        the inset map.
+    anchor : str, optional
+        The anchor point of the inset map, default is 'SE'.
+    provider : str, optional
+        Add a backgroundmap if map provider is passed, default is 'nlmaps.water'. To
+        turn off the backgroundmap set provider to None.
+    add_to_plot : list of functions, optional
+        List of functions to plot on the inset map, default is None. The functions
+        must accept an ax argument. Hint: use `functools.partial` to set plot style,
+        and pass the partial function to add_to_plot.
+
+    Returns
+    -------
+    mapax : matplotlib.Axes
+        The inset map axes.
+    """
+    mapax = ax.inset_axes(axes_bounds)
+    mapax.axis(extent)
+    mapax.set_aspect("equal", adjustable="box", anchor=anchor)
+    mapax.set_xticks([])
+    mapax.set_yticks([])
+    mapax.set_xlabel("")
+    mapax.set_ylabel("")
+
+    if provider:
+        add_background_map(mapax, map_provider=provider, attribution=False)
+
+    if add_to_plot:
+        for fplot in add_to_plot:
+            fplot(ax=mapax)
+
+    return mapax
+
+
+def add_xsec_line_and_labels(
+    line: Union[list, LineString],
+    ax: plt.Axes,
+    mapax: plt.Axes,
+    x_offset: float = 0.0,
+    y_offset: float = 0.0,
+    label: str = "A",
+    **kwargs,
+):
+    """Add a cross-section line to an overview map and label the start and end points.
+
+    Parameters
+    ----------
+    line : list or shapely LineString
+        The line to plot.
+    ax : matplotlib.Axes
+        The axes to plot the labels on.
+    mapax : matplotlib.Axes
+        The axes of the overview map to plot the line on.
+    x_offset : float, optional
+        The x offset of the labels, default is 0.0.
+    y_offset : float, optional
+        The y offset of the labels, default is 0.0.
+    kwargs : dict
+        Keyword arguments to pass to the line plot function.
+
+    Raises
+    ------
+    ValueError
+        If the line is not a list or a shapely LineString.
+    """
+    if isinstance(line, list):
+        x, y = np.array(line).T
+    elif isinstance(line, LineString):
+        x, y = line.xy
+    else:
+        raise ValueError("line should be a list or a shapely LineString")
+    mapax.plot(x, y, **kwargs)
+    stroke = [patheffects.withStroke(linewidth=2, foreground="w")]
+    mapax.text(
+        x[0] + x_offset,
+        y[0] + y_offset,
+        f"{label}",
+        fontweight="bold",
+        path_effects=stroke,
+        fontsize=7,
+    )
+    mapax.text(
+        x[-1] + x_offset,
+        y[-1] + y_offset,
+        f"{label}'",
+        fontweight="bold",
+        path_effects=stroke,
+        fontsize=7,
+    )
+    ax.text(
+        0.01,
+        0.99,
+        f"{label}",
+        transform=ax.transAxes,
+        path_effects=stroke,
+        fontsize=14,
+        ha="left",
+        va="top",
+        fontweight="bold",
+    )
+    ax.text(
+        0.99,
+        0.99,
+        f"{label}'",
+        transform=ax.transAxes,
+        path_effects=stroke,
+        fontsize=14,
+        ha="right",
+        va="top",
+        fontweight="bold",
+    )

tests/test_011_dcs.py

@@ -1,6 +1,7 @@
 import util
 
 import nlmod
+import matplotlib.pyplot as plt
 
 
 def test_dcs_structured():
@@ -21,3 +22,16 @@ def test_dcs_vertex():
     dcs.label_layers()
     dcs.plot_array(ds["kh"], alpha=0.5)
     dcs.plot_grid(vertical=False)
+
+
+def test_cross_section_utils():
+    ds = util.get_ds_vertex()
+    line = [(0, 0), (1000, 1000)]
+    _, ax = plt.subplots(1, 1, figsize=(10, 4))
+    dcs = nlmod.plot.DatasetCrossSection(ds, line, ax=ax)
+    dcs.plot_layers()
+    dcs.label_layers()
+    dcs.plot_array(ds["kh"], alpha=0.5)
+    dcs.plot_grid(vertical=False)
+    mapax = nlmod.plot.inset_map(ax, ds.extent, provider=None)
+    nlmod.plot.add_xsec_line_and_labels(line, ax, mapax)