_load_remote_dataset: Refactor to write CRS information for raster images (#3678)

Author: seisman

pygmt/datasets/earth_day.py

@@ -5,17 +5,12 @@
 The images are available in various resolutions.
 """
 
-import contextlib
 from collections.abc import Sequence
 from typing import Literal
 
 import xarray as xr
 from pygmt.datasets.load_remote_dataset import _load_remote_dataset
 
-with contextlib.suppress(ImportError):
-    # rioxarray is needed to register the rio accessor
-    import rioxarray  # noqa: F401
-
 __doctest_skip__ = ["load_blue_marble"]
 
 
@@ -98,7 +93,4 @@ def load_blue_marble(
         region=region,
         registration="pixel",
     )
-    # If rioxarray is installed, set the coordinate reference system
-    if hasattr(image, "rio"):
-        image = image.rio.write_crs(input_crs="OGC:CRS84")
     return image

pygmt/datasets/earth_night.py

@@ -5,17 +5,12 @@
 The images are available in various resolutions.
 """
 
-import contextlib
 from collections.abc import Sequence
 from typing import Literal
 
 import xarray as xr
 from pygmt.datasets.load_remote_dataset import _load_remote_dataset
 
-with contextlib.suppress(ImportError):
-    # rioxarray is needed to register the rio accessor
-    import rioxarray  # noqa: F401
-
 __doctest_skip__ = ["load_black_marble"]
 
 
@@ -97,7 +92,4 @@ def load_black_marble(
         region=region,
         registration="pixel",
     )
-    # If rioxarray is installed, set the coordinate reference system
-    if hasattr(image, "rio"):
-        image = image.rio.write_crs(input_crs="OGC:CRS84")
     return image

pygmt/datasets/load_remote_dataset.py

@@ -2,6 +2,7 @@
 Internal function to load GMT remote datasets.
 """
 
+import contextlib
 from collections.abc import Sequence
 from typing import Any, Literal, NamedTuple
 
@@ -11,6 +12,10 @@
 from pygmt.helpers import build_arg_list, kwargs_to_strings
 from pygmt.src import which
 
+with contextlib.suppress(ImportError):
+    # rioxarray is needed to register the rio accessor
+    import rioxarray  # noqa: F401
+
 
 class Resolution(NamedTuple):
     """
@@ -48,13 +53,17 @@ class GMTRemoteDataset(NamedTuple):
         Dictionary of available resolution as keys and Resolution objects as values.
     extra_attributes
         A dictionary of extra or unique attributes of the dataset.
+    crs
+        The coordinate reference system of the raster image. Need to be set for images,
+        and should be ``None`` for grids.
     """
 
     description: str
     kind: Literal["grid", "image"]
     units: str | None
     resolutions: dict[str, Resolution]
     extra_attributes: dict[str, Any]
+    crs: str | None = None
 
 
 datasets = {
@@ -81,6 +90,7 @@ class GMTRemoteDataset(NamedTuple):
         description="NASA Day Images",
         kind="image",
         units=None,
+        crs="OGC:CRS84",
         extra_attributes={"long_name": "blue_marble", "horizontal_datum": "WGS84"},
         resolutions={
             "01d": Resolution("01d", registrations=["pixel"]),
@@ -300,6 +310,7 @@ class GMTRemoteDataset(NamedTuple):
         description="NASA Night Images",
         kind="image",
         units=None,
+        crs="OGC:CRS84",
         extra_attributes={"long_name": "black_marble", "horizontal_datum": "WGS84"},
         resolutions={
             "01d": Resolution("01d", registrations=["pixel"]),
@@ -598,4 +609,9 @@ def _load_remote_dataset(
     grid.attrs.pop("actual_range", None)
     for coord in grid.coords:
         grid[coord].attrs.pop("actual_range", None)
+
+    # For images, if rioxarray is installed, set the coordinate reference system.
+    if dataset.crs is not None and hasattr(grid, "rio"):
+        grid = grid.rio.write_crs(input_crs=dataset.crs)
+
     return grid