Add Figure.choropleth to plot choropleth maps (#2798)

Author: seisman

doc/api/index.rst

@@ -25,6 +25,7 @@ Plotting map elements
     :toctree: generated
 
     Figure.basemap
+    Figure.choropleth
     Figure.coast
     Figure.colorbar
     Figure.directional_rose

examples/gallery/maps/choropleth_map.py

@@ -2,15 +2,15 @@
 Choropleth map
 ==============
 
-The :meth:`pygmt.Figure.plot` method allows us to plot geographical data such as
-polygons which are stored in a :class:`geopandas.GeoDataFrame` object. Use
-:func:`geopandas.read_file` to load data from any supported OGR format such as a
-shapefile (.shp), GeoJSON (.geojson), geopackage (.gpkg), etc. You can also use a full
-URL pointing to your desired data source. Then, pass the :class:`geopandas.GeoDataFrame`
-as an argument to the ``data`` parameter of :meth:`pygmt.Figure.plot`, and style the
-geometry using the ``pen`` parameter. To fill the polygons based on a corresponding
-column you need to set ``fill="+z"`` as well as select the appropriate column using the
-``aspatial`` parameter as shown in the example below.
+The :meth:`pygmt.Figure.choropleth` method allows us to plot geographical data such as
+polygons which are stored in a :class:`geopandas.GeoDataFrame` object or an OGR_GMT
+file. Use :func:`geopandas.read_file` to load data from any supported OGR formats such
+as a shapefile (.shp), GeoJSON (.geojson), geopackage (.gpkg), etc. You can also use a
+full URL pointing to your desired data source. Then, pass the
+:class:`geopandas.GeoDataFrame` as an argument to the ``data`` parameter of
+:meth:`pygmt.Figure.choropleth`, and style the geometry using the ``pen`` parameter. To
+fill the polygons based on a corresponding column you need to specify the column name to
+the ``column`` parameter.
 """
 
 # %%
@@ -34,7 +34,7 @@
 
 # Next, we plot the polygons and fill them using the defined colormap. The target column
 # is defined by the aspatial parameter.
-fig.plot(data=africa, pen="0.8p,gray50", fill="+z", cmap=True, aspatial="Z=POP_EST")
+fig.choropleth(data=africa, column="POP_EST", pen="0.8p,gray50")
 
 # Add colorbar legend.
 fig.colorbar(
@@ -45,5 +45,4 @@
     triangle_height=0.2,
     move_text="label",
 )
-
 fig.show()

pygmt/_typing.py

@@ -20,8 +20,14 @@
 with contextlib.suppress(ImportError):
     StringArrayTypes |= importlib.import_module(name="pyarrow").StringArray
 
-# PathLike and TableLike types
+# PathLike, GeoLike, and TableLike types
 PathLike = str | os.PathLike
 TableLike = dict | np.ndarray | pd.DataFrame | xr.Dataset
-with contextlib.suppress(ImportError):
-    TableLike |= importlib.import_module(name="geopandas").GeoDataFrame
+try:
+    import geopandas
+
+    GeoLike = geopandas.GeoDataFrame
+except ImportError:
+    GeoLike = None
+if GeoLike is not None:
+    TableLike |= GeoLike

pygmt/figure.py

@@ -410,6 +410,7 @@ def _repr_html_(self) -> str:
 
     from pygmt.src import (  # type: ignore[misc] # noqa: PLC0415
         basemap,
+        choropleth,
         coast,
         colorbar,
         contour,

pygmt/src/__init__.py

@@ -5,6 +5,7 @@
 from pygmt.src.basemap import basemap
 from pygmt.src.binstats import binstats
 from pygmt.src.blockm import blockmean, blockmedian, blockmode
+from pygmt.src.choropleth import choropleth
 from pygmt.src.coast import coast
 from pygmt.src.colorbar import colorbar
 from pygmt.src.config import config

pygmt/src/choropleth.py

@@ -0,0 +1,59 @@
+"""
+choropleth - Plot a choropleth map.
+"""
+
+from pygmt._typing import GeoLike, PathLike
+
+__doctest_skip__ = ["choropleth"]
+
+
+def choropleth(
+    self,
+    data: GeoLike | PathLike,
+    column: str,
+    cmap: str | bool = True,
+    **kwargs,
+):
+    """
+    Plot a choropleth map.
+
+    This method is a thin wrapper around :meth:`pygmt.Figure.plot` that sets the
+    appropriate parameters for creating a choropleth map by filling polygons based on
+    values in a specified data column. It requires the input data to be a geo-like
+    Python object that implements ``__geo_interface__`` (e.g. a
+    :class:`geopandas.GeoDataFrame`), or an OGR_GMT file containing the geometry and
+    data to plot.
+
+    Parameters
+    ----------
+    data
+        A geo-like Python object which implements ``__geo_interface__`` (e.g. a
+        :class:`geopandas.GeoDataFrame` or :class:`shapely.geometry`), or an OGR_GMT
+        file containing the geometry and data to plot.
+    column
+        The name of the data column to use for the fill.
+    cmap
+        The CPT to use for filling the polygons. If set to ``True``, the current CPT
+        will be used.
+    **kwargs
+        Additional keyword arguments passed to :meth:`pygmt.Figure.plot`.
+
+    Examples
+    --------
+    >>> import geopandas as gpd
+    >>> import pygmt
+    >>> world = gpd.read_file(
+    ...     "https://naciscdn.org/naturalearth/110m/cultural/ne_110m_admin_0_countries.zip"
+    ... )
+    >>> world["POP_EST"] *= 1e-6  # Population in millions
+
+    >>> fig = pygmt.Figure()
+    >>> fig.basemap(region=[-19.5, 53, -38, 37.5], projection="M15c", frame=True)
+    >>> pygmt.makecpt(cmap="bilbao", series=(0, 270, 10), reverse=True)
+    >>> fig.choropleth(world, column="POP_EST", pen="0.3p,gray10")
+    >>> fig.colorbar(frame=True)
+    >>> fig.show()
+    """
+    self.plot(
+        data=data, close=True, fill="+z", cmap=cmap, aspatial=f"Z={column}", **kwargs
+    )

pygmt/tests/baseline/test_choropleth.png.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: ed720d3682b43c4474cd21f49e788a98
+  size: 155020
+  hash: md5
+  path: test_choropleth.png

pygmt/tests/test_choropleth.py

@@ -0,0 +1,33 @@
+"""
+Test Figure.choropleth.
+"""
+
+import pytest
+from pygmt import Figure, makecpt
+
+gpd = pytest.importorskip("geopandas")
+
+
+@pytest.fixture(scope="module", name="world")
+def fixture_world():
+    """
+    Download and cache the Natural Earth countries dataset for testing.
+    """
+    return gpd.read_file(
+        "https://naciscdn.org/naturalearth/110m/cultural/ne_110m_admin_0_countries.zip"
+    )
+
+
+@pytest.mark.mpl_image_compare
+def test_choropleth(world):
+    """
+    Test Figure.choropleth method.
+    """
+    world["POP_EST"] *= 1e-6  # Population in millions
+
+    fig = Figure()
+    fig.basemap(region=[-19.5, 53, -38, 37.5], projection="M15c", frame=True)
+    makecpt(cmap="bilbao", series=(0, 270, 10), reverse=True)
+    fig.choropleth(world, column="POP_EST", pen="0.3p,gray10")
+    fig.colorbar(frame=True)
+    return fig