Use enums GridReg and GridType in gmt accessors

Author: seisman

pygmt/accessors.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 
 import xarray as xr
+from pygmt.enums import GridReg, GridType
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.src.grdinfo import grdinfo
 
@@ -15,110 +16,131 @@ 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 following properties:
+
+    - ``registration``: Grid registration type, either ``GridReg.GRIDLINE`` or
+      ``GridReg.PIXEL``.
+    - ``gtype``: Grid coordinate system type, either ``GridType.CARTESIAN`` or
+      ``GridType.GEOGRAPHIC``.
+
+    and can be accessed like ``grid.gmt.registration`` and ``grid.gmt.gtype``.
 
     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.
+    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.
 
     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 is Gridline or Pixel registration
     >>> grid.gmt.registration
-    1
-    >>> # See if grid uses Cartesian (0) or Geographic (1) coordinate system
+    <GridReg.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
+    For :class:`xarray.DataArray` grids created by yourself, ``registration`` and
+    ``gtype`` default to ``GridReg.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 GridReg, 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-registrated Cartesian grid
+    >>> grid.gmt.registration
+    <GridReg.GRIDLINE: 0>
+    >>> grid.gmt.gtype
+    <GridType.CARTESIAN: 0>
+    >>> # Manually set it to a gridline-registered geographic grid
+    >>> grid.gmt.registration = GridReg.GRIDLINE
+    >>> grid.gmt.gtype = GridType.GEOGRAPHIC
+    >>> grid.gmt.registration
+    <GridReg.GRIDLINE: 0>
+    >>> grid.gmt.gtype
+    <GridType.GEOGRAPHIC: 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.
 
     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
+    <GridReg.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
+    <GridReg.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
+    <GridReg.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
+    <GridReg.GRIDLINE: 0>
+    >>> ds.zval.gmt.gtype
+    <GridType.CARTESIAN: 0>
+    >>> # Manually set these properties won't work as expected
+    >>> ds.zval.gmt.registration = GridReg.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)
+    (<GridReg.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
+    (<GridReg.GRIDLINE: 0>, <GridType.CARTESIAN: 0>)
+    >>> zval.gmt.registration = GridReg.GRIDLINE
+    >>> zval.gmt.gtype = GridType.GEOGRAPHIC
     >>> zval.gmt.registration, zval.gmt.gtype
-    (0, 1)
+    (<GridReg.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 = GridReg.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 +153,39 @@ def __init__(self, xarray_obj):
     @property
     def registration(self):
         """
-        Registration type of the grid, either 0 (Gridline) or 1 (Pixel).
+        Registration type of the grid, either ``GridReg.GRIDLINE`` or
+        ``GridReg.PIXEL``.
         """
         return self._registration
 
     @registration.setter
     def registration(self, value):
-        if value not in {0, 1}:
+        if value in {"gridline", "pixel"}:  # Support for string-type values
+            value = GridReg[value.upper()]
+        if value not in GridReg:
             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 GridReg.GRIDLINE or GridReg.PIXEL."
             )
             raise GMTInvalidInput(msg)
-        self._registration = value
+        self._registration = GridReg(value)
 
     @property
     def gtype(self):
         """
-        Coordinate system type of the grid, either 0 (Cartesian) or 1 (Geographic).
+        Coordinate system type of the grid, either ``GridType.CARTESIAN`` or
+        ``GridType.GEOGRAPHIC``.
         """
         return self._gtype
 
     @gtype.setter
     def gtype(self, value):
-        if value not in {0, 1}:
+        if value in {"cartesian", "geographic"}:  # Support for string-type values
+            value = GridType[value.upper()]
+        if value not in GridType:
             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 or GridType.GEOGRAPHIC."
             )
             raise GMTInvalidInput(msg)
-        self._gtype = value
+        self._gtype = GridType(value)