From c86e8d54fa6cd2bdf0ffae9cbdb433eb8d79392b Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Sat, 26 Jul 2025 15:22:43 +0800 Subject: [PATCH] Add Figure.scalebar to plot a scale bar on maps --- doc/api/index.rst | 1 + pygmt/figure.py | 1 + pygmt/src/__init__.py | 1 + pygmt/src/_common.py | 7 +- pygmt/src/scalebar.py | 162 ++++++++++++++++++ pygmt/tests/baseline/test_scalebar.png.dvc | 5 + .../baseline/test_scalebar_cartesian.png.dvc | 5 + .../baseline/test_scalebar_complete.png.dvc | 5 + pygmt/tests/test_scalebar.py | 62 +++++++ 9 files changed, 246 insertions(+), 3 deletions(-) create mode 100644 pygmt/src/scalebar.py create mode 100644 pygmt/tests/baseline/test_scalebar.png.dvc create mode 100644 pygmt/tests/baseline/test_scalebar_cartesian.png.dvc create mode 100644 pygmt/tests/baseline/test_scalebar_complete.png.dvc create mode 100644 pygmt/tests/test_scalebar.py diff --git a/doc/api/index.rst b/doc/api/index.rst index 264f5a9175a..9b312ca2f9f 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -31,6 +31,7 @@ Plotting map elements Figure.inset Figure.legend Figure.logo + Figure.scalebar Figure.solar Figure.text Figure.timestamp diff --git a/pygmt/figure.py b/pygmt/figure.py index 56ad2c3d5cf..3c48080c482 100644 --- a/pygmt/figure.py +++ b/pygmt/figure.py @@ -427,6 +427,7 @@ def _repr_html_(self) -> str: plot3d, psconvert, rose, + scalebar, set_panel, shift_origin, solar, diff --git a/pygmt/src/__init__.py b/pygmt/src/__init__.py index 8905124f917..2aa4e6b5587 100644 --- a/pygmt/src/__init__.py +++ b/pygmt/src/__init__.py @@ -43,6 +43,7 @@ from pygmt.src.project import project from pygmt.src.psconvert import psconvert from pygmt.src.rose import rose +from pygmt.src.scalebar import scalebar from pygmt.src.select import select from pygmt.src.shift_origin import shift_origin from pygmt.src.solar import solar diff --git a/pygmt/src/_common.py b/pygmt/src/_common.py index 92202828df2..27366c5ab2c 100644 --- a/pygmt/src/_common.py +++ b/pygmt/src/_common.py @@ -251,7 +251,7 @@ def _parse_position( position: Position | Sequence[float | str] | str | None, kwdict: dict[str, Any], default: Position | None, -) -> Position | str: +) -> Position | str | None: """ Parse the "position" parameter for embellishment-plotting functions. @@ -269,7 +269,8 @@ def _parse_position( The keyword arguments dictionary that conflicts with ``position`` if ``position`` is given as a raw GMT command string. default - The default Position object to use if ``position`` is ``None``. + The default Position object to use if ``position`` is ``None``. If ``default`` + is ``None``, the GMT default is used. Returns ------- @@ -349,7 +350,7 @@ def _parse_position( position = Position(position, cstype="plotcoords") case Position(): # Already a Position object. pass - case None if default is not None: # Set default position. + case None: # Set default position. position = default case _: msg = f"Invalid type for parameter 'position': {type(position)}." diff --git a/pygmt/src/scalebar.py b/pygmt/src/scalebar.py new file mode 100644 index 00000000000..9b0594806bf --- /dev/null +++ b/pygmt/src/scalebar.py @@ -0,0 +1,162 @@ +""" +scalebar - Add a scale bar. +""" + +from collections.abc import Sequence +from typing import Literal + +from pygmt._typing import AnchorCode +from pygmt.alias import Alias, AliasSystem +from pygmt.clib import Session +from pygmt.exceptions import GMTInvalidInput +from pygmt.helpers import build_arg_list, fmt_docstring +from pygmt.params import Box, Position +from pygmt.src._common import _parse_position + +__doctest_skip__ = ["scalebar"] + + +@fmt_docstring +def scalebar( # noqa: PLR0913 + self, + position: Position | Sequence[float | str] | AnchorCode | None = None, + length: float | str | None = None, + height: float | str | None = None, + scale_position: float | Sequence[float] | bool = False, + label: str | bool = False, + label_alignment: Literal["left", "right", "top", "bottom"] | None = None, + unit: bool = False, + fancy: bool = False, + vertical: bool = False, + box: Box | bool = False, + verbose: Literal["quiet", "error", "warning", "timing", "info", "compat", "debug"] + | bool = False, + panel: int | Sequence[int] | bool = False, + perspective: float | Sequence[float] | str | bool = False, + transparency: float | None = None, +): + """ + Add a scale bar on the plot. + + Parameters + ---------- + position + Position of the scale bar on the plot. It can be specified in multiple ways: + + - A :class:`pygmt.params.Position` object to fully control the reference point, + anchor point, and offset. + - A sequence of two values representing the x and y coordinates in plot + coordinates, e.g., ``(1, 2)`` or ``("1c", "2c")``. + - A :doc:`2-character justification code ` for a + position inside the plot, e.g., ``"TL"`` for Top Left corner inside the plot. + + If not specified, defaults to the lower-left corner of the plot. + length + Length of the scale bar in km. Append a suffix to specify different units. Valid + units are: **e**: meters; **f**: feet; **k**: kilometers; **M**: statute mile; + **n**: nautical miles; **u**: US Survey foot. + height + Height of the scale bar. Only works when ``fancy=True``. [Default is ``"5p"``]. + scale_position + Specify the location where on a geographic map the scale applies. It can be: + + - *slat*: Map scale is calculated for latitude *slat* + - (*slon*, *slat*): Map scale is calculated for latitude *slat* and longitude + *slon*, which is useful for oblique projections. + - ``True``: Map scale is calculated for the middle of the map. + - ``False``: Default to the location of the reference point. + label + Text string to use as the scale bar label. If ``False``, no label is drawn. If + ``True``, the distance unit provided in the ``length`` parameter (default is km) + is used as the label. This parameter requires ``fancy=True``. + label_alignment + Alignment of the scale bar label. Choose from ``"left"``, ``"right"``, + ``"top"``, or ``"bottom"``. [Default is ``"top"``]. + fancy + If ``True``, draw a "fancy" scale bar, which is a segmented bar with alternating + black and white rectangles. If ``False``, draw a plain scale bar. + unit + If ``True``, append the unit to all distance annotations along the scale. For a + plain scale, this will instead select the unit to be appended to the distance + length. The unit is determined from the suffix in the ``length`` or defaults to + ``"km"``. + vertical + If ``True``, plot a vertical rather than a horizontal Cartesian scale. + box + Draw a background box behind the directional rose. If set to ``True``, a simple + rectangular box is drawn using :gmt-term:`MAP_FRAME_PEN`. To customize the box + appearance, pass a :class:`pygmt.params.Box` object to control style, fill, pen, + and other box properties. + $verbose + $panel + $perspective + $transparency + + Examples + -------- + >>> import pygmt + >>> from pygmt.params import Box, Position + >>> fig = pygmt.Figure() + >>> fig.basemap(region=[0, 80, -30, 30], projection="M10c", frame=True) + >>> fig.scalebar( + ... position=Position((10, 10), cstype="mapcoords"), + ... length=1000, + ... fancy=True, + ... label="Scale", + ... unit=True, + ... ) + >>> fig.show() + """ + self._activate_figure() + + position = _parse_position( + position, + kwdict={ + "length": length, + "height": height, + "label_alignment": label_alignment, + "scale_position": scale_position, + "fancy": fancy, + "label": label, + "unit": unit, + "vertical": vertical, + }, + default=None, # Use the GMT default. + ) + + if length is None: + msg = "Parameter 'length' must be specified." + raise GMTInvalidInput(msg) + + aliasdict = AliasSystem( + F=Alias(box, name="box"), + L=[ + Alias(position, name="position"), + Alias(length, name="length", prefix="+w"), + Alias( + label_alignment, + name="label_alignment", + prefix="+a", + mapping={"left": "l", "right": "r", "top": "t", "bottom": "b"}, + ), + Alias(scale_position, name="scale_position", prefix="+c", sep="/", size=2), + Alias(fancy, name="fancy", prefix="+f"), + Alias(label, name="label", prefix="+l"), + Alias(unit, name="unit", prefix="+u"), + Alias(vertical, name="vertical", prefix="+v"), + ], + ).add_common( + V=verbose, + c=panel, + p=perspective, + t=transparency, + ) + + confdict = {} + if height is not None: + confdict["MAP_SCALE_HEIGHT"] = height + + with Session() as lib: + lib.call_module( + module="basemap", args=build_arg_list(aliasdict, confdict=confdict) + ) diff --git a/pygmt/tests/baseline/test_scalebar.png.dvc b/pygmt/tests/baseline/test_scalebar.png.dvc new file mode 100644 index 00000000000..793d5dc0f26 --- /dev/null +++ b/pygmt/tests/baseline/test_scalebar.png.dvc @@ -0,0 +1,5 @@ +outs: +- md5: 79fc9a29d68df6467b032c38e6b0b263 + size: 10288 + hash: md5 + path: test_scalebar.png diff --git a/pygmt/tests/baseline/test_scalebar_cartesian.png.dvc b/pygmt/tests/baseline/test_scalebar_cartesian.png.dvc new file mode 100644 index 00000000000..0d01ea6cca5 --- /dev/null +++ b/pygmt/tests/baseline/test_scalebar_cartesian.png.dvc @@ -0,0 +1,5 @@ +outs: +- md5: e09a7c67f6146530ea594694853b6f98 + size: 6508 + hash: md5 + path: test_scalebar_cartesian.png diff --git a/pygmt/tests/baseline/test_scalebar_complete.png.dvc b/pygmt/tests/baseline/test_scalebar_complete.png.dvc new file mode 100644 index 00000000000..4ac6b71bb99 --- /dev/null +++ b/pygmt/tests/baseline/test_scalebar_complete.png.dvc @@ -0,0 +1,5 @@ +outs: +- md5: c018b219d3ebc719fb1b1686e074dcd9 + size: 11749 + hash: md5 + path: test_scalebar_complete.png diff --git a/pygmt/tests/test_scalebar.py b/pygmt/tests/test_scalebar.py new file mode 100644 index 00000000000..eff1c44aeb7 --- /dev/null +++ b/pygmt/tests/test_scalebar.py @@ -0,0 +1,62 @@ +""" +Test Figure.scalebar. +""" + +import pytest +from pygmt import Figure +from pygmt.exceptions import GMTInvalidInput +from pygmt.params import Position + + +@pytest.mark.mpl_image_compare +def test_scalebar(): + """ + Create a map with a scale bar. + """ + fig = Figure() + fig.basemap(region=[100, 120, 20, 30], projection="M10c", frame=True) + fig.scalebar(length=500) + return fig + + +@pytest.mark.mpl_image_compare +def test_scalebar_complete(): + """ + Test all parameters of scalebar. + """ + fig = Figure() + fig.basemap(region=[100, 120, 20, 30], projection="M10c", frame=True) + fig.scalebar( + position=Position((110, 22), cstype="mapcoords"), + length=1000, + height="10p", + fancy=True, + label="Scale", + label_alignment="left", + scale_position=(110, 25), + unit=True, + box=True, + ) + return fig + + +@pytest.mark.mpl_image_compare +def test_scalebar_cartesian(): + """ + Test scale bar in Cartesian coordinates. + """ + fig = Figure() + fig.basemap(region=[0, 10, 0, 5], projection="X10c/5c", frame=True) + fig.scalebar(position=Position((2, 1), cstype="mapcoords"), length=1) + fig.scalebar(position=Position((4, 1), cstype="mapcoords"), length=1, vertical=True) + return fig + + +def test_scalebar_no_length(): + """ + Test that an error is raised when length is not provided. + """ + fig = Figure() + fig.basemap(region=[100, 120, 20, 30], projection="M10c", frame=True) + with pytest.raises(GMTInvalidInput): + fig.scalebar(position=Position((118, 22), cstype="mapcoords"))