load_tile_map: Add the new parameter 'crs' to set the CRS of the returned dataarray (#3554)

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

Author: seisman

pygmt/datasets/tile_map.py

@@ -3,25 +3,29 @@
 :class:`xarray.DataArray`.
 """
 
-import contextlib
 from collections.abc import Sequence
 from typing import Literal
 
 from packaging.version import Version
 
 try:
     import contextily
+    from rasterio.crs import CRS
     from xyzservices import TileProvider
 
     _HAS_CONTEXTILY = True
 except ImportError:
+    CRS = None
     TileProvider = None
     _HAS_CONTEXTILY = False
 
-with contextlib.suppress(ImportError):
-    # rioxarray is needed to register the rio accessor
+try:
     import rioxarray  # noqa: F401
 
+    _HAS_RIOXARRAY = True
+except ImportError:
+    _HAS_RIOXARRAY = False
+
 import numpy as np
 import xarray as xr
 
@@ -33,6 +37,7 @@ def load_tile_map(
     zoom: int | Literal["auto"] = "auto",
     source: TileProvider | str | None = None,
     lonlat: bool = True,
+    crs: str | CRS = "EPSG:3857",
     wait: int = 0,
     max_retries: int = 2,
     zoom_adjust: int | None = None,
@@ -42,7 +47,8 @@ def load_tile_map(
 
     The tiles that compose the map are merged and georeferenced into an
     :class:`xarray.DataArray` image with 3 bands (RGB). Note that the returned image is
-    in a Spherical Mercator (EPSG:3857) coordinate reference system.
+    in a Spherical Mercator (EPSG:3857) coordinate reference system (CRS) by default,
+    but can be customized using the ``crs`` parameter.
 
     Parameters
     ----------
@@ -80,6 +86,10 @@ def load_tile_map(
     lonlat
         If ``False``, coordinates in ``region`` are assumed to be Spherical Mercator as
         opposed to longitude/latitude.
+    crs
+        Coordinate reference system (CRS) of the returned :class:`xarray.DataArray`
+        image. Default is ``"EPSG:3857"`` (i.e., Spherical Mercator). The CRS can be in
+        either string or :class:`rasterio.crs.CRS` format.
     wait
         If the tile API is rate-limited, the number of seconds to wait between a failed
         request and the next try.
@@ -128,6 +138,10 @@ def load_tile_map(
     ...     raster.rio.crs.to_string()
     'EPSG:3857'
     """
+    # The CRS of the source tile provider. If the source is a TileProvider object, use
+    # its crs attribute if available. Otherwise, default to EPSG:3857.
+    _source_crs = getattr(source, "crs", "EPSG:3857")
+
     if not _HAS_CONTEXTILY:
         msg = (
             "Package `contextily` is required to be installed to use this function. "
@@ -136,28 +150,34 @@ def load_tile_map(
         )
         raise ImportError(msg)
 
-    contextily_kwargs = {}
+    if crs != _source_crs and not _HAS_RIOXARRAY:
+        msg = (
+            f"Package `rioxarray` is required if CRS is not '{_source_crs}'. "
+            "Please use `python -m pip install rioxarray` or "
+            "`mamba install -c conda-forge rioxarray` to install the package."
+        )
+        raise ImportError(msg)
+
+    # Keyword arguments for contextily.bounds2img
+    contextily_kwargs = {
+        "zoom": zoom,
+        "source": source,
+        "ll": lonlat,
+        "wait": wait,
+        "max_retries": max_retries,
+    }
     if zoom_adjust is not None:
-        contextily_kwargs["zoom_adjust"] = zoom_adjust
         if Version(contextily.__version__) < Version("1.5.0"):
             msg = (
                 "The `zoom_adjust` parameter requires `contextily>=1.5.0` to work. "
                 "Please upgrade contextily, or manually set the `zoom` level instead."
             )
-            raise TypeError(msg)
+            raise ValueError(msg)
+        contextily_kwargs["zoom_adjust"] = zoom_adjust
 
     west, east, south, north = region
     image, extent = contextily.bounds2img(
-        w=west,
-        s=south,
-        e=east,
-        n=north,
-        zoom=zoom,
-        source=source,
-        ll=lonlat,
-        wait=wait,
-        max_retries=max_retries,
-        **contextily_kwargs,
+        w=west, s=south, e=east, n=north, **contextily_kwargs
     )
 
     # Turn RGBA img from channel-last to channel-first and get 3-band RGB only
@@ -176,8 +196,12 @@ def load_tile_map(
         dims=("band", "y", "x"),
     )
 
-    # If rioxarray is installed, set the coordinate reference system
+    # If rioxarray is installed, set the coordinate reference system.
     if hasattr(dataarray, "rio"):
-        dataarray = dataarray.rio.write_crs(input_crs="EPSG:3857")
+        dataarray = dataarray.rio.write_crs(input_crs=_source_crs)
+
+        # Reproject raster image from the source CRS to the specified CRS.
+        if crs != _source_crs:
+            dataarray = dataarray.rio.reproject(dst_crs=crs)
 
     return dataarray

pygmt/src/tilemap.py

@@ -9,13 +9,9 @@
 from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
 
 try:
-    import rioxarray  # noqa: F401
     from xyzservices import TileProvider
-
-    _HAS_RIOXARRAY = True
 except ImportError:
     TileProvider = None
-    _HAS_RIOXARRAY = False
 
 
 @fmt_docstring
@@ -111,39 +107,21 @@ def tilemap(
 
     kwargs : dict
         Extra keyword arguments to pass to :meth:`pygmt.Figure.grdimage`.
-
-    Raises
-    ------
-    ImportError
-        If ``rioxarray`` is not installed. Follow
-        :doc:`install instructions for rioxarray <rioxarray:installation>`, (e.g. via
-        ``python -m pip install rioxarray``) before using this function.
     """
     kwargs = self._preprocess(**kwargs)
 
-    if not _HAS_RIOXARRAY:
-        msg = (
-            "Package `rioxarray` is required to be installed to use this function. "
-            "Please use `python -m pip install rioxarray` or "
-            "`mamba install -c conda-forge rioxarray` to install the package."
-        )
-        raise ImportError(msg)
-
     raster = load_tile_map(
         region=region,
         zoom=zoom,
         source=source,
         lonlat=lonlat,
+        crs="OGC:CRS84" if lonlat is True else "EPSG:3857",
         wait=wait,
         max_retries=max_retries,
         zoom_adjust=zoom_adjust,
     )
-
-    # Reproject raster from Spherical Mercator (EPSG:3857) to lonlat (OGC:CRS84) if
-    # bounding box region was provided in lonlat
-    if lonlat and raster.rio.crs == "EPSG:3857":
-        raster = raster.rio.reproject(dst_crs="OGC:CRS84")
-        raster.gmt.gtype = 1  # set to geographic type
+    if lonlat:
+        raster.gmt.gtype = 1  # Set to geographic type
 
     # Only set region if no_clip is None or False, so that plot is clipped to exact
     # bounding box region

pygmt/tests/test_tilemap.py

@@ -2,14 +2,24 @@
 Test Figure.tilemap.
 """
 
+import importlib
+from unittest.mock import patch
+
 import pytest
 from pygmt import Figure
 
-contextily = pytest.importorskip("contextily")
-rioxarray = pytest.importorskip("rioxarray")
+try:
+    import contextily
+
+    _HAS_CONTEXTILY = True
+except ImportError:
+    _HAS_CONTEXTILY = False
+
+_HAS_RIOXARRAY = bool(importlib.util.find_spec("rioxarray"))
 
 
 @pytest.mark.mpl_image_compare
+@pytest.mark.skipif(not _HAS_CONTEXTILY, reason="contextily is not installed")
 def test_tilemap_web_mercator():
     """
     Create a tilemap plot in Spherical Mercator projection (EPSG:3857).
@@ -27,6 +37,10 @@ def test_tilemap_web_mercator():
 
 @pytest.mark.benchmark
 @pytest.mark.mpl_image_compare
+@pytest.mark.skipif(
+    not (_HAS_CONTEXTILY and _HAS_RIOXARRAY),
+    reason="contextily and rioxarray are not installed",
+)
 def test_tilemap_ogc_wgs84():
     """
     Create a tilemap plot using longitude/latitude coordinates (OGC:WGS84), centred on
@@ -45,6 +59,10 @@ def test_tilemap_ogc_wgs84():
 
 @pytest.mark.mpl_image_compare
 @pytest.mark.parametrize("no_clip", [False, True])
+@pytest.mark.skipif(
+    not (_HAS_CONTEXTILY and _HAS_RIOXARRAY),
+    reason="contextily and rioxarray are not installed",
+)
 def test_tilemap_no_clip(no_clip):
     """
     Create a tilemap plot clipped to the Southern Hemisphere when no_clip is False, but
@@ -60,3 +78,34 @@ def test_tilemap_no_clip(no_clip):
         no_clip=no_clip,
     )
     return fig
+
+
+@pytest.mark.skipif(_HAS_CONTEXTILY, reason="contextily is installed.")
+def test_tilemap_no_contextily():
+    """
+    Raise an ImportError when contextily is not installed.
+    """
+    fig = Figure()
+    with pytest.raises(ImportError, match="Package `contextily` is required"):
+        fig.tilemap(
+            region=[-20000000.0, 20000000.0, -20000000.0, 20000000.0],
+            zoom=0,
+            lonlat=False,
+            frame="afg",
+        )
+
+
+@pytest.mark.skipif(_HAS_RIOXARRAY, reason="rioxarray is installed.")
+def test_tilemap_no_rioxarray():
+    """
+    Raise an ImportError when rioxarray is not installed and contextily is installed.
+    """
+    fig = Figure()
+    # In our CI, contextily and rioxarray are installed together, so we will see the
+    # error about contextily, not rioxarray. Here we mock contextily as installed, to
+    # make sure that we see the rioxarray error message when rioxarray is not installed.
+    with patch("pygmt.datasets.tile_map._HAS_CONTEXTILY", True):
+        with pytest.raises(ImportError, match="Package `rioxarray` is required"):
+            fig.tilemap(
+                region=[-180.0, 180.0, -90, 90], zoom=0, lonlat=True, frame="afg"
+            )