Add Figure.vlines for plotting vertical lines

Author: seisman

doc/api/index.rst

@@ -36,6 +36,7 @@ Plotting map elements
     Figure.solar
     Figure.text
     Figure.timestamp
+    Figure.vlines
 
 Plotting tabular data
 ~~~~~~~~~~~~~~~~~~~~~

pygmt/figure.py

@@ -436,6 +436,7 @@ def _repr_html_(self) -> str:
         tilemap,
         timestamp,
         velo,
+        vlines,
         wiggle,
     )
 

pygmt/src/__init__.py

@@ -57,6 +57,7 @@
 from pygmt.src.timestamp import timestamp
 from pygmt.src.triangulate import triangulate
 from pygmt.src.velo import velo
+from pygmt.src.vlines import vlines
 from pygmt.src.which import which
 from pygmt.src.wiggle import wiggle
 from pygmt.src.x2sys_cross import x2sys_cross

pygmt/src/vlines.py

@@ -0,0 +1,132 @@
+"""
+vlines - Plot vertical lines.
+"""
+
+from collections.abc import Sequence
+
+import numpy as np
+from pygmt.exceptions import GMTInvalidInput
+
+__doctest_skip__ = ["vlines"]
+
+
+def vlines(
+    self,
+    x: float | Sequence[float],
+    ymin: float | Sequence[float] | None = None,
+    ymax: float | Sequence[float] | None = None,
+    pen: str | None = None,
+    label: str | None = None,
+    no_clip: bool = False,
+    perspective: str | bool | None = None,
+):
+    """
+    Plot one or multiple vertical line(s).
+
+    This method is a high-level wrapper around :meth:`pygmt.Figure.plot` that focuses on
+    plotting vertical lines at X-coordinates specified by the ``x`` parameter. The ``x``
+    parameter can be a single value (for a single vertical line) or a sequence of values
+    (for multiple vertical lines).
+
+    By default, the Y-coordinates of the start and end points of the lines are set to be
+    the Y-limits of the current plot, but this can be overridden by specifying the
+    ``ymin`` and ``ymax`` parameters. ``ymin`` and ``ymax`` can be either a single value
+    or a sequence of values. If a single value is provided, it is applied to all lines.
+    If a sequence is provided, the length of ``ymin`` and ``ymax`` must match the length
+    of ``x``.
+
+    The term "vertical" lines can be interpreted differently in different coordinate
+    systems:
+
+    - **Cartesian** coordinate system: lines are plotted as straight lines.
+    - **Polar** projection: lines are plotted as straight lines along radius.
+    - **Geographic** projection: lines are plotted as meridians along constant
+      longitude.
+
+    Parameters
+    ----------
+    x
+        X-coordinates to plot the lines. It can be a single value (for a single line)
+        or a sequence of values (for multiple lines).
+    ymin/ymax
+        Y-coordinates of the start/end point of the line(s). If ``None``, defaults to
+        the Y-limits of the current plot. ``ymin`` and ``ymax`` can either be a single
+        value or a sequence of values. If a single value is provided, it is applied to
+        all lines. If a sequence is provided, the length of ``ymin`` and ``ymax`` must
+        match the length of ``x``.
+    pen
+        Pen attributes for the line(s), in the format of *width,color,style*.
+    label
+        Label for the line(s), to be displayed in the legend.
+    no_clip
+        If ``True``, do not clip lines outside the plot region. Only makes sense in the
+        Cartesian coordinate system.
+    perspective
+        Select perspective view and set the azimuth and elevation angle of the
+        viewpoint. Refer to :meth:`pygmt.Figure.plot` for details.
+
+    Examples
+    --------
+    >>> import pygmt
+    >>> fig = pygmt.Figure()
+    >>> fig.basemap(region=[0, 10, 0, 10], projection="X10c/10c", frame=True)
+    >>> fig.vlines(x=1, pen="1p,black", label="Line 1")
+    >>> fig.vlines(x=2, ymin=2, ymax=8, pen="1p,red,-", label="Line at x=2")
+    >>> fig.vlines(x=[3, 4], ymin=3, ymax=7, pen="1p,black,.", label="Lines at x=3,4")
+    >>> fig.vlines(x=[5, 6], ymin=4, ymax=9, pen="1p,red", label="Lines at x=5,6")
+    >>> fig.vlines(
+    ...     x=[7, 8], ymin=[0, 1], ymax=[7, 8], pen="1p,blue", label="Lines at x=7,8"
+    ... )
+    >>> fig.legend()
+    >>> fig.show()
+    """
+    self._preprocess()
+
+    # Determine the x limits from the current plot region if not specified.
+    if ymin is None or ymax is None:
+        ylimits = self.region[2:]
+        if ymin is None:
+            ymin = ylimits[0]
+        if ymax is None:
+            ymax = ylimits[1]
+
+    # Ensure y/xmin/xmax are 1-D arrays.
+    _x = np.atleast_1d(x)
+    _ymin = np.atleast_1d(ymin)
+    _ymax = np.atleast_1d(ymax)
+
+    nlines = len(_x)  # Number of lines to plot.
+
+    # Check if ymin/ymax are scalars or have the expected length.
+    if _ymin.size not in {1, nlines} or _ymax.size not in {1, nlines}:
+        msg = (
+            f"'ymin' and 'ymax' are expected to be scalars or have lengths '{nlines}', "
+            f"but lengths '{_ymin.size}' and '{_ymax.size}' are given."
+        )
+        raise GMTInvalidInput(msg)
+
+    # Repeat ymin/ymax to match the length of x if they are scalars.
+    if nlines != 1:
+        if _ymin.size == 1:
+            _ymin = np.repeat(_ymin, nlines)
+        if _ymax.size == 1:
+            _ymax = np.repeat(_ymax, nlines)
+
+    # Call the Figure.plot method to plot the lines.
+    for i in range(nlines):
+        # Special handling for label.
+        # 1. Only specify a label when plotting the first line.
+        # 2. The -l option can accept comma-separated labels for labeling multiple lines
+        #    with auto-coloring enabled. We don't need this feature here, so we need to
+        #    replace comma with \054 if the label contains commas.
+        _label = label.replace(",", "\\054") if label and i == 0 else None
+
+        self.plot(
+            x=[_x[i], _x[i]],
+            y=[_ymin[i], _ymax[i]],
+            pen=pen,
+            label=_label,
+            no_clip=no_clip,
+            perspective=perspective,
+            straight_line="y",
+        )

pygmt/tests/baseline/test_vlines_clip.png.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: 4eb9c7fd7e3a803dcc3cde1409ad7fa7
+  size: 7361
+  hash: md5
+  path: test_vlines_clip.png

pygmt/tests/baseline/test_vlines_geographic_global.png.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: 3fb4a271c670e4cbe647838b6fee5a8c
+  size: 67128
+  hash: md5
+  path: test_vlines_geographic_global.png

pygmt/tests/baseline/test_vlines_multiple_lines.png.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: 6d189c9bbee1f294b4e43a74683a1927
+  size: 13555
+  hash: md5
+  path: test_vlines_multiple_lines.png

pygmt/tests/baseline/test_vlines_one_line.png.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: 986772f58935f81e9596736e914acb78
+  size: 13604
+  hash: md5
+  path: test_vlines_one_line.png

pygmt/tests/baseline/test_vlines_polar_projection.png.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: 1981df3bd9c57cd975b6e74946496175
+  size: 44621
+  hash: md5
+  path: test_vlines_polar_projection.png

pygmt/tests/test_vlines.py

@@ -0,0 +1,102 @@
+"""
+Tests for Figure.vlines.
+"""
+
+import pytest
+from pygmt import Figure
+from pygmt.exceptions import GMTInvalidInput
+
+
+@pytest.mark.mpl_image_compare
+def test_vlines_one_line():
+    """
+    Plot one vertical line.
+    """
+    fig = Figure()
+    fig.basemap(region=[0, 10, 0, 10], projection="X10c/10c", frame=True)
+    fig.vlines(1)
+    fig.vlines(2, ymin=1)
+    fig.vlines(3, ymax=9)
+    fig.vlines(4, ymin=3, ymax=8)
+    fig.vlines(5, ymin=4, ymax=8, pen="1p,blue", label="Line at y=5")
+    fig.vlines(6, ymin=5, ymax=7, pen="1p,red", label="Line at y=6")
+    fig.legend()
+    return fig
+
+
+@pytest.mark.mpl_image_compare
+def test_vlines_multiple_lines():
+    """
+    Plot multiple vertical lines.
+    """
+    fig = Figure()
+    fig.basemap(region=[0, 16, 0, 10], projection="X10c/10c", frame=True)
+    fig.vlines([1, 2])
+    fig.vlines([3, 4, 5], ymin=[1, 2, 3])
+    fig.vlines([6, 7, 8], ymax=[7, 8, 9])
+    fig.vlines([9, 10], ymin=[1, 2], ymax=[9, 10])
+    fig.vlines([11, 12], ymin=1, ymax=8, pen="1p,blue", label="Line at y=11,12")
+    fig.vlines(
+        [13, 14], ymin=[3, 4], ymax=[7, 8], pen="1p,red", label="Line at y=13,14"
+    )
+    fig.legend()
+    return fig
+
+
+@pytest.mark.mpl_image_compare
+def test_vlines_clip():
+    """
+    Plot vertical lines with clipping or not.
+    """
+    fig = Figure()
+    fig.basemap(region=[0, 10, 0, 4], projection="X10c/4c", frame=True)
+    fig.vlines(1, ymin=-1, ymax=5)
+    fig.vlines(2, ymin=-1, ymax=5, no_clip=True)
+    return fig
+
+
+@pytest.mark.mpl_image_compare
+def test_vlines_geographic_global():
+    """
+    Plot vertical lines in geographic coordinates.
+    """
+    fig = Figure()
+    fig.basemap(region=[-180, 180, -90, 90], projection="R15c", frame="a30g30")
+    fig.vlines(30, pen="1p")
+    fig.vlines(90, ymin=-60, pen="1p,blue")
+    fig.vlines(-90, ymax=60, pen="1p,blue")
+    fig.vlines(120, ymin=-60, ymax=60, pen="1p,blue")
+    return fig
+
+
+@pytest.mark.mpl_image_compare
+def test_vlines_polar_projection():
+    """
+    Plot vertical lines in polar projection.
+    """
+    fig = Figure()
+    fig.basemap(region=[0, 360, 0, 1], projection="P15c", frame=True)
+    fig.vlines(0, pen="1p")
+    fig.vlines(30, ymin=0, ymax=1, pen="1p")
+    fig.vlines(60, ymin=0.5, pen="1p")
+    fig.vlines(90, ymax=0.5, pen="1p")
+    fig.vlines(120, ymin=0.25, ymax=0.75, pen="1p")
+    return fig
+
+
+def test_vlines_invalid_input():
+    """
+    Test invalid input for vlines.
+    """
+    fig = Figure()
+    fig.basemap(region=[0, 10, 0, 6], projection="X10c/6c", frame=True)
+    with pytest.raises(GMTInvalidInput):
+        fig.vlines(1, ymin=2, ymax=[3, 4])
+    with pytest.raises(GMTInvalidInput):
+        fig.vlines(1, ymin=[2, 3], ymax=4)
+    with pytest.raises(GMTInvalidInput):
+        fig.vlines(1, xmin=[2, 3], xmax=[4, 5])
+    with pytest.raises(GMTInvalidInput):
+        fig.vlines([1, 2], ymin=[2, 3, 4], ymax=3)
+    with pytest.raises(GMTInvalidInput):
+        fig.vlines([1, 2], ymin=[2, 3], ymax=[4, 5, 6])