Refactor the _load_remote_dataset function to load tiled and non-tiled grids in a consistent way (#3120)

Author: seisman

pygmt/datasets/load_remote_dataset.py

@@ -4,12 +4,12 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, ClassVar, NamedTuple
+from typing import TYPE_CHECKING, ClassVar, Literal, NamedTuple
 
+from pygmt.clib import Session
 from pygmt.exceptions import GMTInvalidInput
-from pygmt.helpers import kwargs_to_strings
-from pygmt.io import load_dataarray
-from pygmt.src import grdcut, which
+from pygmt.helpers import build_arg_list, kwargs_to_strings
+from pygmt.src import which
 
 if TYPE_CHECKING:
     import xarray as xr
@@ -344,7 +344,7 @@ def _load_remote_dataset(
     dataset_prefix: str,
     resolution: str,
     region: str | list,
-    registration: str,
+    registration: Literal["gridline", "pixel", None],
 ) -> xr.DataArray:
     r"""
     Load GMT remote datasets.
@@ -370,54 +370,68 @@ def _load_remote_dataset(
 
     Returns
     -------
-    grid : :class:`xarray.DataArray`
+    grid
         The GMT remote dataset grid.
 
     Note
     ----
-    The returned :class:`xarray.DataArray` doesn't support slice operation for tiled
-    grids.
+    The registration and coordinate system type of the returned
+    :class:`xarray.DataArray` grid can be accessed via the GMT accessors (i.e.,
+    ``grid.gmt.registration`` and ``grid.gmt.gtype`` respectively). However, these
+    properties may be lost after specific grid operations (such as slicing) and will
+    need to be manually set before passing the grid to any PyGMT data processing or
+    plotting functions. Refer to :class:`pygmt.GMTDataArrayAccessor` for detailed
+    explanations and workarounds.
     """
     dataset = datasets[dataset_name]
 
+    # Check resolution
     if resolution not in dataset.resolutions:
         raise GMTInvalidInput(
             f"Invalid resolution '{resolution}' for {dataset.title} dataset. "
             f"Available resolutions are: {', '.join(dataset.resolutions)}."
         )
+    resinfo = dataset.resolutions[resolution]
 
-    # check registration
-    valid_registrations = dataset.resolutions[resolution].registrations
+    # Check registration
     if registration is None:
-        # use gridline registration unless only pixel registration is available
-        registration = "gridline" if "gridline" in valid_registrations else "pixel"
+        # Use gridline registration unless only pixel registration is available
+        registration = "gridline" if "gridline" in resinfo.registrations else "pixel"
     elif registration in ("pixel", "gridline"):
-        if registration not in valid_registrations:
+        if registration not in resinfo.registrations:
             raise GMTInvalidInput(
                 f"{registration} registration is not available for the "
                 f"{resolution} {dataset.title} dataset. Only "
-                f"{valid_registrations[0]} registration is available."
+                f"{resinfo.registrations[0]} registration is available."
             )
     else:
         raise GMTInvalidInput(
             f"Invalid grid registration: '{registration}', should be either 'pixel', "
             "'gridline' or None. Default is None, where a gridline-registered grid is "
             "returned unless only the pixel-registered grid is available."
         )
-    reg = f"_{registration[0]}"
 
-    # different ways to load tiled and non-tiled grids.
-    # Known issue: tiled grids don't support slice operation
-    # See https://github.com/GenericMappingTools/pygmt/issues/524
-    if region is None:
-        if dataset.resolutions[resolution].tiled:
-            raise GMTInvalidInput(
-                f"'region' is required for {dataset.title} resolution '{resolution}'."
+    fname = f"@{dataset_prefix}{resolution}_{registration[0]}"
+    if resinfo.tiled and region is None:
+        raise GMTInvalidInput(
+            f"'region' is required for {dataset.title} resolution '{resolution}'."
+        )
+
+    # Currently, only grids are supported. Will support images in the future.
+    kwdict = {"T": "g", "R": region}  # region can be None
+    with Session() as lib:
+        with lib.virtualfile_out(kind="grid") as voutgrd:
+            lib.call_module(
+                module="read",
+                args=[fname, voutgrd, *build_arg_list(kwdict)],
             )
-        fname = which(f"@{dataset_prefix}{resolution}{reg}", download="a")
-        grid = load_dataarray(fname, engine="netcdf4")
-    else:
-        grid = grdcut(f"@{dataset_prefix}{resolution}{reg}", region=region)
+            grid = lib.virtualfile_to_raster(outgrid=None, vfname=voutgrd)
+
+    # Full path to the grid if not tiled grids.
+    source = which(fname, download="a") if not resinfo.tiled else None
+    # Manually add source to xarray.DataArray encoding to make the GMT accessors work.
+    if source:
+        grid.encoding["source"] = source
 
     # Add some metadata to the grid
     grid.name = dataset.name

pygmt/tests/test_accessor.py

@@ -115,9 +115,8 @@ def test_accessor_grid_source_file_not_exist():
     # Registration and gtype are correct
     assert grid.gmt.registration == 1
     assert grid.gmt.gtype == 1
-    # The source grid file is defined but doesn't exist
-    assert grid.encoding["source"].endswith(".nc")
-    assert not Path(grid.encoding["source"]).exists()
+    # The source grid file is undefined.
+    assert grid.encoding.get("source") is None
 
     # For a sliced grid, fallback to default registration and gtype,
     # because the source grid file doesn't exist.