Figure.savefig: Support generating GeoTIFF file (with extension '.tiff') (#2698)

Co-authored-by: Yvonne Fröhlich <[email protected]>
Co-authored-by: Wei Ji <[email protected]>

Author: 3 people

pygmt/figure.py

@@ -257,12 +257,18 @@ def savefig(
         """
         Save the figure to a file.
 
-        This method implements a matplotlib-like interface for
-        :meth:`pygmt.Figure.psconvert`.
+        Supported file formats and their extensions:
 
-        Supported formats: PNG (``.png``), JPEG (``.jpg`` or ``.jpeg``),
-        PDF (``.pdf``), BMP (``.bmp``), TIFF (``.tif``), EPS (``.eps``), and
-        KML (``.kml``). The KML output generates a companion PNG file.
+        - PNG (``.png``)
+        - JPEG (``.jpg`` or ``.jpeg``)
+        - PDF (``.pdf``)
+        - BMP (``.bmp``)
+        - TIFF (``.tif``)
+        - GeoTIFF (``.tiff``)
+        - EPS (``.eps``)
+        - KML (``.kml``)
+
+        For KML format, a companion PNG file is also generated.
 
         You can pass in any keyword arguments that
         :meth:`pygmt.Figure.psconvert` accepts.
@@ -279,10 +285,10 @@ def savefig(
             If ``True``, will crop the figure canvas (page) to the plot area.
         anti_alias: bool
             If ``True``, will use anti-aliasing when creating raster images
-            (PNG, JPG, TIFF). More specifically, it passes arguments ``t2``
-            and ``g2`` to the ``anti_aliasing`` parameter of
-            :meth:`pygmt.Figure.psconvert`. Ignored if creating vector
-            graphics.
+            (BMP, PNG, JPEG, TIFF, and GeoTIFF). More specifically, it passes
+            the arguments ``"t2"`` and ``"g2"`` to the ``anti_aliasing``
+            parameter of :meth:`pygmt.Figure.psconvert`. Ignored if creating
+            vector graphics.
         show: bool
             If ``True``, will open the figure in an external viewer.
         dpi : int
@@ -301,15 +307,20 @@ def savefig(
             "bmp": "b",
             "eps": "e",
             "tif": "t",
+            "tiff": None,  # GeoTIFF doesn't need the -T option
             "kml": "g",
         }
 
         fname = Path(fname)
         prefix, suffix = fname.with_suffix("").as_posix(), fname.suffix
         ext = suffix[1:].lower()  # Remove the . and normalize to lowercase
-        # alias jpeg to jpg
-        if ext == "jpeg":
+
+        if ext == "jpeg":  # Alias jpeg to jpg
             ext = "jpg"
+        elif ext == "tiff":  # GeoTIFF
+            kwargs["W"] = "+g"
+        elif ext == "kml":  # KML
+            kwargs["W"] = "+k"
 
         if ext not in fmts:
             if ext == "ps":
@@ -328,11 +339,15 @@ def savefig(
         if anti_alias:
             kwargs["Qt"] = 2
             kwargs["Qg"] = 2
-        if ext == "kml":
-            kwargs["W"] = "+k"
 
         self.psconvert(prefix=prefix, fmt=fmt, crop=crop, **kwargs)
 
+        # Remove the .pgw world file if exists
+        # Not necessary after GMT 6.5.0.
+        # See upstream fix https://github.com/GenericMappingTools/gmt/pull/7865
+        if ext == "tiff" and fname.with_suffix(".pgw").exists():
+            fname.with_suffix(".pgw").unlink()
+
         # Rename if file extension doesn't match the input file suffix
         if ext != suffix[1:]:
             fname.with_suffix("." + ext).rename(fname)

pygmt/tests/test_figure.py

@@ -92,6 +92,72 @@ def test_figure_savefig_exists():
         fname.unlink()
 
 
+def test_figure_savefig_geotiff():
+    """
+    Make sure .tif generates a normal TIFF file and .tiff generates a GeoTIFF
+    file.
+    """
+    fig = Figure()
+    fig.basemap(region=[0, 10, 0, 10], projection="M10c", frame=True)
+
+    # Save as GeoTIFF
+    geofname = Path("test_figure_savefig_geotiff.tiff")
+    fig.savefig(geofname)
+    assert geofname.exists()
+    # The .pgw should not exist
+    assert not geofname.with_suffix(".pgw").exists()
+
+    # Save as TIFF
+    fname = Path("test_figure_savefig_tiff.tif")
+    fig.savefig(fname)
+    assert fname.exists()
+
+    # Check if a TIFF is georeferenced or not
+    try:
+        # pylint: disable=import-outside-toplevel
+        import rioxarray
+        from rasterio.errors import NotGeoreferencedWarning
+        from rasterio.transform import Affine
+
+        # GeoTIFF
+        with rioxarray.open_rasterio(geofname) as xds:
+            assert xds.rio.crs is not None
+            npt.assert_allclose(
+                actual=xds.rio.bounds(),
+                desired=(
+                    -661136.0621116752,
+                    -54631.82709660966,
+                    592385.4459661598,
+                    1129371.7360144067,
+                ),
+            )
+            assert xds.rio.shape == (1257, 1331)
+            assert xds.rio.transform() == Affine(
+                a=941.789262267344,
+                b=0.0,
+                c=-661136.0621116752,
+                d=0.0,
+                e=-941.92805338983,
+                f=1129371.7360144067,
+            )
+        # TIFF
+        with pytest.warns(expected_warning=NotGeoreferencedWarning) as record:
+            with rioxarray.open_rasterio(fname) as xds:
+                assert xds.rio.crs is None
+                npt.assert_allclose(
+                    actual=xds.rio.bounds(), desired=(0.0, 0.0, 1331.0, 1257.0)
+                )
+                assert xds.rio.shape == (1257, 1331)
+                assert xds.rio.transform() == Affine(
+                    a=1.0, b=0.0, c=0.0, d=0.0, e=1.0, f=0.0
+                )
+            assert len(record) == 1
+    except ImportError:
+        pass
+    geofname.unlink()
+    fname.unlink()
+
+
 def test_figure_savefig_directory_nonexists():
     """
     Make sure that Figure.savefig() raises a FileNotFoundError when the parent