From e5f21af071116e4ff231269ca03055d282d91583 Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Sun, 5 Nov 2023 18:53:39 +0800
Subject: [PATCH] Add Figure.choropleth to plot choropleth maps

---
 examples/gallery/maps/choropleth_map.py | 37 ++++++++++---------------
 pygmt/figure.py                         |  1 +
 pygmt/src/__init__.py                   |  1 +
 pygmt/src/choropleth.py                 | 29 +++++++++++++++++++
 4 files changed, 45 insertions(+), 23 deletions(-)
 create mode 100644 pygmt/src/choropleth.py

diff --git a/examples/gallery/maps/choropleth_map.py b/examples/gallery/maps/choropleth_map.py
index 6c43d24d3dd..ebd6f3674f2 100644
--- a/examples/gallery/maps/choropleth_map.py
+++ b/examples/gallery/maps/choropleth_map.py
@@ -2,16 +2,14 @@
 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 a 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 colum name to the ``column`` parameter.
 """
 
 # %%
@@ -29,11 +27,10 @@
     frame="+tPopulation of Chicago",
 )
 
-# The dataset contains different attributes, here we select
-# the "population" column to plot.
+# The dataset contains different attributes, here we select the "population" column to
+# plot.
 
-# First, we define the colormap to fill the polygons based on
-# the "population" column.
+# First, we define the colormap to fill the polygons based on the "population" column.
 pygmt.makecpt(
     cmap="acton",
     series=[gdf["population"].min(), gdf["population"].max(), 10],
@@ -41,15 +38,9 @@
     reverse=True,
 )
 
-# Next, we plot the polygons and fill them using the defined colormap.
-# The target column is defined by the aspatial parameter.
-fig.plot(
-    data=gdf,
-    pen="0.3p,gray10",
-    fill="+z",
-    cmap=True,
-    aspatial="Z=population",
-)
+# Next, we plot the polygons and fill them using the defined colormap. The target column
+# is specified by the `column` parameter.
+fig.choropleth(data=gdf, column="population", pen="0.3p,gray10", cmap=True)
 
 # Add colorbar legend
 fig.colorbar(frame="x+lPopulation", position="jML+o-0.5c+w3.5c/0.2c")
diff --git a/pygmt/figure.py b/pygmt/figure.py
index 12aaa67f722..c7057cde436 100644
--- a/pygmt/figure.py
+++ b/pygmt/figure.py
@@ -402,6 +402,7 @@ def _repr_html_(self):
 
     from pygmt.src import (  # type: ignore [misc]
         basemap,
+        choropleth,
         coast,
         colorbar,
         contour,
diff --git a/pygmt/src/__init__.py b/pygmt/src/__init__.py
index e4db7321963..8fa13808675 100644
--- a/pygmt/src/__init__.py
+++ b/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
diff --git a/pygmt/src/choropleth.py b/pygmt/src/choropleth.py
new file mode 100644
index 00000000000..039219707a6
--- /dev/null
+++ b/pygmt/src/choropleth.py
@@ -0,0 +1,29 @@
+"""
+choropleth - Plot a choropleth map.
+"""
+
+import contextlib
+
+with contextlib.suppress(ImportError):
+    import geopandas as gpd
+
+
+def choropleth(self, data: gpd.GeoDataFrame, column: str, **kwargs):
+    """
+    Plot a choropleth map.
+
+    Parameters
+    ----------
+    data
+        A :class:`geopandas.DataFrame` object or a OGR_GMT file containing the geometry
+        and data to plot.
+    column
+        The name of the data column to use for the fill.
+    """
+    self.plot(
+        data=data,
+        close=True,
+        fill="+z",
+        aspatial=f"Z={column}",
+        **kwargs,
+    )