diff --git a/pygmt/datasets/tile_map.py b/pygmt/datasets/tile_map.py index 0c9f1b6271d..d18bce423f8 100644 --- a/pygmt/datasets/tile_map.py +++ b/pygmt/datasets/tile_map.py @@ -3,7 +3,6 @@ :class:`xarray.DataArray`. """ -import contextlib from collections.abc import Sequence from typing import Literal @@ -11,17 +10,22 @@ 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 diff --git a/pygmt/src/tilemap.py b/pygmt/src/tilemap.py index 82edc1e72ec..ab65c3ac3c9 100644 --- a/pygmt/src/tilemap.py +++ b/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 `, (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 diff --git a/pygmt/tests/test_tilemap.py b/pygmt/tests/test_tilemap.py index 704a96093af..6ec04decf95 100644 --- a/pygmt/tests/test_tilemap.py +++ b/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" + )