Figure.colorbar: Add parameters label/unit/annot/tick/grid and more to set colorbar annotations/labels (#4407)

Author: seisman

pygmt/src/colorbar.py

@@ -11,12 +11,77 @@
 from pygmt.exceptions import GMTValueError
 from pygmt.helpers import build_arg_list, fmt_docstring, use_alias
 from pygmt.helpers.utils import is_nonstr_iter
-from pygmt.params import Box, Position
+from pygmt.params import Axis, Box, Frame, Position
 from pygmt.src._common import _parse_position
 
 __doctest_skip__ = ["colorbar"]
 
 
+def _build_frame(
+    annot: float | bool = False,
+    tick: float | bool = False,
+    grid: float | bool = False,
+    annot_angle: float | None = None,
+    annot_prefix: str | None = None,
+    annot_unit: str | None = None,
+    label: str | None = None,
+    unit: str | None = None,
+    frame=None,
+):
+    """
+    Create the list of Alias objects for the -B option.
+
+    >>> list(_build_frame(annot=1, tick=0.5, label="Distance", unit="km"))
+    ['xa1f0.5+lDistance', 'y+lkm']
+
+    >>> list(
+    ...     _build_frame(
+    ...         annot=1,
+    ...         tick=0.5,
+    ...         grid=0.2,
+    ...         annot_angle=30,
+    ...         label="Distance",
+    ...         unit="km",
+    ...     )
+    ... )
+    ['xa1f0.5g0.2+lDistance+a30', 'y+lkm']
+    >>> list(_build_frame(frame=["xaf0.5+lDistance", "y+lkm"]))
+    ['xaf0.5+lDistance', 'y+lkm']
+
+    >>> _build_frame(frame="none")
+    'none'
+    >>> _build_frame()  # Passing no parameters returns None
+    """
+    # Using the old 'frame' parameter.
+    if frame is not None and frame is not False:
+        return frame
+
+    _xaxis_is_set = any(
+        v is not None and v is not False
+        for v in {annot, tick, grid, annot_angle, annot_prefix, annot_unit, label}
+    )
+    _yaxis_is_set = unit is not None
+
+    # Need to return None if no parameters are give. Otherwise, it may return "".
+    if not (_xaxis_is_set or _yaxis_is_set):
+        return None
+
+    xaxis, yaxis = None, None
+    if _xaxis_is_set:
+        xaxis = Axis(
+            annot=annot,
+            tick=tick,
+            grid=grid,
+            angle=annot_angle,
+            prefix=annot_prefix,
+            unit=annot_unit,
+            label=label,
+        )
+    if _yaxis_is_set:
+        yaxis = Axis(label=unit)
+    return Frame(xaxis=xaxis, yaxis=yaxis)
+
+
 def _alias_option_D(  # noqa: N802, PLR0913
     position=None,
     length=None,
@@ -153,6 +218,14 @@ def colorbar(  # noqa: PLR0913
     length: float | str | None = None,
     width: float | str | None = None,
     orientation: Literal["horizontal", "vertical"] | None = None,
+    label: str | None = None,
+    unit: str | None = None,
+    annot: float | bool = False,
+    tick: float | bool = False,
+    grid: float | bool = False,
+    annot_angle: float | None = None,
+    annot_prefix: str | None = None,
+    annot_unit: str | None = None,
     reverse: bool = False,
     nan: bool = False,
     nan_position: Literal["start", "end"] | None = None,
@@ -201,7 +274,6 @@ def colorbar(  # noqa: PLR0913
     Full GMT docs at :gmt-docs:`colorbar.html`.
 
     $aliases
-       - B = frame
        - F = box
        - G = truncate
        - I = shading
@@ -217,6 +289,7 @@ def colorbar(  # noqa: PLR0913
     .. hlist::
        :columns: 1
 
+       - B = label, unit, annot, tick, grid, annot_angle, annot_prefix, annot_unit
        - D = position, **+w**: length/width, **+h**/**+v**: orientation,
          **+r**: reverse, **+n**: nan/nan_position,
          **+e**: fg_triangle/bg_triangle/triangle_height,
@@ -244,6 +317,24 @@ def colorbar(  # noqa: PLR0913
         given with unit ``%`` then it is in percentage of the bar length. [Length
         defaults to 80% of the corresponding plot side dimension, and width defaults
         to 4% of the bar length].
+    label
+    unit
+        Set the label and unit for the colorbar. The label is placed along the colorbar
+        and the unit is placed at the end of the colorbar.
+    annot
+    tick
+    grid
+        Intervals for annotations, ticks, and gridlines. Refer to
+        :class:`pygmt.params.Axis` for more details on how these parameters work.
+    annot_prefix
+    annot_unit
+    annot_angle
+        The prefix, unit and angle for the annotations. The prefix is placed before the
+        annotation text; the unit is placed after the annotation text; and the angle is
+        the angle of the annotation text.
+    frame
+        Set colorbar boundary frame, labels, and axes attributes. If set to ``"none"``,
+        then no frame will be drawn.
     orientation
         Set the colorbar orientation to either ``"horizontal"`` or ``"vertical"``.
         [Default is vertical, unless ``position`` is set to bottom-center or top-center
@@ -321,9 +412,6 @@ def colorbar(  # noqa: PLR0913
         requested colorbar length.
     $projection
     $region
-    frame
-        Set colorbar boundary frame, labels, and axes attributes. If set to ``"none"``,
-        then no frame will be drawn.
     $verbose
     $panel
     $perspective
@@ -337,12 +425,7 @@ def colorbar(  # noqa: PLR0913
     >>> # Create a basemap
     >>> fig.basemap(region=[0, 10, 0, 3], projection="X10c/3c", frame=True)
     >>> # Call the colorbar method for the plot
-    >>> fig.colorbar(
-    ...     # Set cmap to the "roma" CPT
-    ...     cmap="SCM/roma",
-    ...     # Label the x-axis "Velocity" and the y-axis "m/s"
-    ...     frame=["x+lVelocity", "y+lm/s"],
-    ... )
+    >>> fig.colorbar(cmap="SCM/roma", label="Velocity", unit="m/s")
     >>> # Show the plot
     >>> fig.show()
     """
@@ -387,7 +470,17 @@ def colorbar(  # noqa: PLR0913
         Q=Alias(log, name="log"),
         W=Alias(scale, name="scale"),
     ).add_common(
-        B=frame,
+        B=_build_frame(
+            annot=annot,
+            tick=tick,
+            grid=grid,
+            annot_angle=annot_angle,
+            annot_prefix=annot_prefix,
+            annot_unit=annot_unit,
+            label=label,
+            unit=unit,
+            frame=frame,
+        ),
         J=projection,
         R=region,
         V=verbose,