GMTDataArrayAccessor: Support passing values using enums GridRegistration and GridType for grid registration and type (#3696)

Co-authored-by: Michael Grund <[email protected]>
Co-authored-by: Wei Ji <[email protected]>

Author: 3 people

pygmt/accessors.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 
 import xarray as xr
+from pygmt.enums import GridRegistration, GridType
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.src.grdinfo import grdinfo
 
@@ -15,110 +16,122 @@ class GMTDataArrayAccessor:
     """
     GMT accessor for :class:`xarray.DataArray`.
 
-    The accessor extends :class:`xarray.DataArray` to store GMT-specific
-    properties about grids, which are important for PyGMT to correctly process
-    and plot the grids.
+    The *gmt* accessor extends :class:`xarray.DataArray` to store GMT-specific
+    properties for grids, which are important for PyGMT to correctly process and plot
+    the grids. The *gmt* accessor contains the following properties:
 
-    Notes
-    -----
-
-    Due to the limitations of xarray accessors, the GMT accessors are created
-    once per :class:`xarray.DataArray` instance. You may lose these
-    GMT-specific properties when manipulating grids (e.g., arithmetic and slice
-    operations) or when accessing a :class:`xarray.DataArray` from a
-    :class:`xarray.Dataset`. In these cases, you need to manually set these
-    properties before passing the grid to PyGMT.
+    - ``registration``: Grid registration type :class:`pygmt.enums.GridRegistration`.
+    - ``gtype``: Grid coordinate system type :class:`pygmt.enums.GridType`.
 
     Examples
     --------
-
-    For GMT's built-in remote datasets, these GMT-specific properties are
-    automatically determined and you can access them as follows:
+    For GMT's built-in remote datasets, these GMT-specific properties are automatically
+    determined and you can access them as follows:
 
     >>> from pygmt.datasets import load_earth_relief
     >>> # Use the global Earth relief grid with 1 degree spacing
     >>> grid = load_earth_relief(resolution="01d", registration="pixel")
-    >>> # See if grid uses Gridline (0) or Pixel (1) registration
+    >>> # See if grid uses Gridline or Pixel registration
     >>> grid.gmt.registration
-    1
-    >>> # See if grid uses Cartesian (0) or Geographic (1) coordinate system
+    <GridRegistration.PIXEL: 1>
+    >>> # See if grid is in Cartesian or Geographic coordinate system
     >>> grid.gmt.gtype
-    1
+    <GridType.GEOGRAPHIC: 1>
 
-    For :class:`xarray.DataArray` grids created by yourself, grid properties
-    ``registration`` and ``gtype`` default to 0 (i.e., a gridline-registered,
-    Cartesian grid). You need to set the correct properties before
-    passing it to PyGMT functions:
+    For :class:`xarray.DataArray` grids created by yourself, ``registration`` and
+    ``gtype`` default to ``GridRegistration.GRIDLINE`` and ``GridType.CARTESIAN`` (i.e.,
+    a gridline-registered, Cartesian grid). You need to set the correct properties
+    before passing it to PyGMT functions:
 
     >>> import numpy as np
-    >>> import pygmt
     >>> import xarray as xr
-    >>> # create a DataArray in gridline coordinates of sin(lon) * cos(lat)
+    >>> import pygmt
+    >>> from pygmt.enums import GridRegistration, GridType
+    >>> # Create a DataArray in gridline coordinates of sin(lon) * cos(lat)
     >>> interval = 2.5
     >>> lat = np.arange(90, -90 - interval, -interval)
     >>> lon = np.arange(0, 360 + interval, interval)
     >>> longrid, latgrid = np.meshgrid(lon, lat)
     >>> data = np.sin(np.deg2rad(longrid)) * np.cos(np.deg2rad(latgrid))
     >>> grid = xr.DataArray(data, coords=[("latitude", lat), ("longitude", lon)])
-    >>> # default to a gridline-registered Cartesian grid
-    >>> grid.gmt.registration, grid.gmt.gtype
-    (0, 0)
-    >>> # set it to a gridline-registered geographic grid
-    >>> grid.gmt.registration = 0
-    >>> grid.gmt.gtype = 1
-    >>> grid.gmt.registration, grid.gmt.gtype
-    (0, 1)
-
-    Note that the accessors are created once per :class:`xarray.DataArray`
-    instance, so you may lose these GMT-specific properties after manipulating
-    your grid.
+    >>> # Default to a gridline-registered Cartesian grid
+    >>> grid.gmt.registration
+    <GridRegistration.GRIDLINE: 0>
+    >>> grid.gmt.gtype
+    <GridType.CARTESIAN: 0>
+    >>> # Manually set it to a gridline-registered geographic grid
+    >>> grid.gmt.registration = GridRegistration.GRIDLINE
+    >>> grid.gmt.gtype = GridType.GEOGRAPHIC
+    >>> grid.gmt.registration
+    <GridRegistration.GRIDLINE: 0>
+    >>> grid.gmt.gtype
+    <GridType.GEOGRAPHIC: 1>
+
+    Notes
+    -----
+    Due to the limitations of xarray accessors, the GMT accessors are created once per
+    :class:`xarray.DataArray` instance. You may lose these GMT-specific properties when
+    manipulating grids (e.g., arithmetic and slice operations) or when accessing a
+    :class:`xarray.DataArray` from a :class:`xarray.Dataset`. In these cases, you need
+    to manually set these properties before passing the grid to PyGMT.
 
     Inplace assignment operators like ``*=`` don't create new instances, so the
     properties are still kept:
 
     >>> grid *= 2.0
-    >>> grid.gmt.registration, grid.gmt.gtype
-    (0, 1)
+    >>> grid.gmt.registration
+    <GridRegistration.GRIDLINE: 0>
+    >>> grid.gmt.gtype
+    <GridType.GEOGRAPHIC: 1>
 
-    Other grid operations (e.g., arithmetic or slice operations) create new
-    instances, so the properties will be lost:
+    Other grid operations (e.g., arithmetic or slice operations) create new instances,
+    so the properties will be lost:
 
     >>> # grid2 is a slice of the original grid
     >>> grid2 = grid[0:30, 50:80]
-    >>> # properties are reset to the default values for new instance
-    >>> grid2.gmt.registration, grid2.gmt.gtype
-    (0, 0)
-    >>> # need to set these properties before passing the grid to PyGMT
+    >>> # Properties are reset to the default values for new instance
+    >>> grid2.gmt.registration
+    <GridRegistration.GRIDLINE: 0>
+    >>> grid2.gmt.gtype
+    <GridType.CARTESIAN: 0>
+    >>> # Need to set these properties before passing the grid to PyGMT
     >>> grid2.gmt.registration = grid.gmt.registration
     >>> grid2.gmt.gtype = grid.gmt.gtype
-    >>> grid2.gmt.registration, grid2.gmt.gtype
-    (0, 1)
+    >>> grid2.gmt.registration
+    <GridRegistration.GRIDLINE: 0>
+    >>> grid2.gmt.gtype
+    <GridType.GEOGRAPHIC: 1>
 
-    Accessing a :class:`xarray.DataArray` from a :class:`xarray.Dataset` always
-    creates new instances, so these properties are always lost. The workaround
-    is to assign the :class:`xarray.DataArray` into a variable:
+    Accessing a :class:`xarray.DataArray` from a :class:`xarray.Dataset` always creates
+    new instances, so these properties are always lost. The workaround is to assign the
+    :class:`xarray.DataArray` into a variable:
 
     >>> ds = xr.Dataset({"zval": grid})
+    >>> ds.zval.gmt.registration
+    <GridRegistration.GRIDLINE: 0>
+    >>> ds.zval.gmt.gtype
+    <GridType.CARTESIAN: 0>
+    >>> # Manually set these properties won't work as expected
+    >>> ds.zval.gmt.registration = GridRegistration.GRIDLINE
+    >>> ds.zval.gmt.gtype = GridType.GEOGRAPHIC
     >>> ds.zval.gmt.registration, ds.zval.gmt.gtype
-    (0, 0)
-    >>> # manually set these properties won't work as expected
-    >>> ds.zval.gmt.registration, ds.zval.gmt.gtype = 0, 1
-    >>> ds.zval.gmt.registration, ds.zval.gmt.gtype
-    (0, 0)
+    (<GridRegistration.GRIDLINE: 0>, <GridType.CARTESIAN: 0>)
     >>> # workaround: assign the DataArray into a variable
     >>> zval = ds.zval
     >>> zval.gmt.registration, zval.gmt.gtype
-    (0, 0)
-    >>> zval.gmt.registration, zval.gmt.gtype = 0, 1
+    (<GridRegistration.GRIDLINE: 0>, <GridType.CARTESIAN: 0>)
+    >>> zval.gmt.registration = GridRegistration.GRIDLINE
+    >>> zval.gmt.gtype = GridType.GEOGRAPHIC
     >>> zval.gmt.registration, zval.gmt.gtype
-    (0, 1)
+    (<GridRegistration.GRIDLINE: 0>, <GridType.GEOGRAPHIC: 1>)
     """
 
     def __init__(self, xarray_obj):
         self._obj = xarray_obj
+
         # Default to Gridline registration and Cartesian grid type
-        self._registration = 0
-        self._gtype = 0
+        self._registration = GridRegistration.GRIDLINE
+        self._gtype = GridType.CARTESIAN
 
         # If the source file exists, get grid registration and grid type from the last
         # two columns of the shortened summary information of grdinfo.
@@ -131,33 +144,35 @@ def __init__(self, xarray_obj):
     @property
     def registration(self):
         """
-        Registration type of the grid, either 0 (Gridline) or 1 (Pixel).
+        Grid registration type :class:`pygmt.enums.GridRegistration`.
         """
         return self._registration
 
     @registration.setter
     def registration(self, value):
-        if value not in {0, 1}:
+        # TODO(Python>=3.12): Simplify to `if value not in GridRegistration`.
+        if value not in GridRegistration.__members__.values():
             msg = (
-                f"Invalid grid registration value: {value}, should be either "
-                "0 for Gridline registration or 1 for Pixel registration."
+                f"Invalid grid registration: '{value}'. Should be either "
+                "GridRegistration.GRIDLINE (0) or GridRegistration.PIXEL (1)."
             )
             raise GMTInvalidInput(msg)
-        self._registration = value
+        self._registration = GridRegistration(value)
 
     @property
     def gtype(self):
         """
-        Coordinate system type of the grid, either 0 (Cartesian) or 1 (Geographic).
+        Grid coordinate system type :class:`pygmt.enums.GridType`.
         """
         return self._gtype
 
     @gtype.setter
     def gtype(self, value):
-        if value not in {0, 1}:
+        # TODO(Python>=3.12): Simplify to `if value not in GridType`.
+        if value not in GridType.__members__.values():
             msg = (
-                f"Invalid coordinate system type: {value}, should be "
-                "either 0 for Cartesian or 1 for Geographic."
+                f"Invalid grid coordinate system type: '{value}'. "
+                "Should be either GridType.CARTESIAN (0) or GridType.GEOGRAPHIC (1)."
             )
             raise GMTInvalidInput(msg)
-        self._gtype = value
+        self._gtype = GridType(value)

pygmt/datatypes/grid.py

@@ -165,7 +165,7 @@ def to_dataarray(self) -> xr.DataArray:
             axis:           Y
             actual_range:   [-24. -10.]
         >>> da.gmt.registration, da.gmt.gtype
-        (1, 1)
+        (<GridRegistration.PIXEL: 1>, <GridType.GEOGRAPHIC: 1>)
         """
         # The grid header
         header = self.header.contents

pygmt/datatypes/image.py

@@ -166,7 +166,7 @@ def to_dataarray(self) -> xr.DataArray:
             axis:          Y
             actual_range:  [-90.  90.]
         >>> da.gmt.registration, da.gmt.gtype
-        (1, 1)
+        (<GridRegistration.PIXEL: 1>, <GridType.GEOGRAPHIC: 1>)
         """
         # The image header
         header = self.header.contents

pygmt/src/tilemap.py

@@ -6,6 +6,7 @@
 
 from pygmt.clib import Session
 from pygmt.datasets.tile_map import load_tile_map
+from pygmt.enums import GridType
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
 
 try:
@@ -121,7 +122,7 @@ def tilemap(
         zoom_adjust=zoom_adjust,
     )
     if lonlat:
-        raster.gmt.gtype = 1  # Set to geographic type
+        raster.gmt.gtype = GridType.GEOGRAPHIC
 
     # Only set region if no_clip is None or False, so that plot is clipped to exact
     # bounding box region