Source code for pygmt.src.vlines

"""
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(s) 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 at x=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 y 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 x/ymin/ymax 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",
        )