diff --git a/pyproject.toml b/pyproject.toml index 7ad5b26abc..c3647291a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ extras = [ "cartopy>=0.21.0", "dask>=2020.12.0", "netCDF4>=1.7.1.post1", + "rioxarray>=0.13.0", "shapely>=1.6.4", "boto3>=1.26.45" ] diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index 5d5ae524d3..974a2152c7 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -746,6 +746,63 @@ def assign_y_x(self, force=False, tolerance=None): new_dataarray = self._data_array.assign_coords(**{y.name: y, x.name: x}) return new_dataarray.metpy.assign_coordinates(None) + def to_geotiff(self, path, **kwargs): + """Write this DataArray to a GeoTIFF file, preserving the MetPy CRS. + + This is a convenience bridge to ``rioxarray``: it attaches the DataArray's MetPy + coordinate reference system (as assigned by ``.parse_cf`` or ``.assign_crs``) and its + x/y spatial dimension coordinates so the field is written as a properly georeferenced + raster that GIS software (e.g. QGIS, ArcGIS, GDAL-based tooling) can read directly. + + Parameters + ---------- + path : str or path-like + Destination path for the output GeoTIFF. + kwargs + Additional keyword arguments passed through to ``rioxarray``'s ``to_raster`` + (e.g. ``driver``, ``compress``, ``dtype``, ``tiled``). + + Notes + ----- + Requires the optional ``rioxarray`` dependency. A valid CRS coordinate must be present + (as assigned by ``.parse_cf`` or ``.assign_crs``), and the DataArray must have + identifiable x and y dimension coordinates. PyProj is used to communicate the CRS to + ``rioxarray``. + + """ + try: + import rioxarray # noqa: F401 + except ImportError: + raise ImportError( + 'rioxarray is required to write GeoTIFF files. Install it with ' + '`pip install rioxarray` or `conda install -c conda-forge rioxarray`.' + ) from None + + # Identify the spatial dimension coordinates that rioxarray needs + try: + x, y = self.coordinates('x', 'y') + except AttributeError: + raise ValueError( + 'Both x and y dimension coordinates are required to write a GeoTIFF. Verify ' + 'that the data has been parsed by MetPy with proper x and y dimension ' + 'coordinates.' + ) from None + + # Grab the CRS before we drop the coordinate holding it + crs = self.pyproj_crs + + # rioxarray/rasterio operate on plain arrays, so drop pint units first. Also drop the + # metpy_crs coordinate, which holds a CFProjection object that cannot be serialized as + # a raster tag--the CRS is instead communicated to rioxarray explicitly below. + data = self.dequantify() + if 'metpy_crs' in data.coords: + data = data.drop_vars('metpy_crs') + + # Tell rioxarray which dimensions are spatial and what the CRS is, then write + data = data.rio.set_spatial_dims(x_dim=x.name, y_dim=y.name, inplace=False) + data = data.rio.write_crs(crs, inplace=False) + data.rio.to_raster(path, **kwargs) + @xr.register_dataset_accessor('metpy') class MetPyDatasetAccessor: diff --git a/tests/test_xarray.py b/tests/test_xarray.py index e9fbb3cfdf..973ae789d4 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -74,6 +74,49 @@ def test_pyproj_projection(test_var): assert proj.coordinate_operation.method_name == 'Lambert Conic Conformal (1SP)' +def test_to_geotiff(test_var, tmp_path): + """Test writing a parsed DataArray to a georeferenced GeoTIFF.""" + rioxarray = pytest.importorskip('rioxarray') + + # Reduce to a single 2D spatial field so we have a plain (y, x) raster + field = test_var + for dim in field.dims: + if dim not in ('y', 'x'): + field = field.isel({dim: 0}) + + path = tmp_path / 'test.tif' + field.metpy.to_geotiff(path) + assert path.exists() + + # Read back and verify the georeferencing round-trips + result = rioxarray.open_rasterio(path) + assert result.rio.crs == pyproj.CRS(test_var.metpy.pyproj_crs) + assert result.rio.width == test_var.sizes['x'] + assert result.rio.height == test_var.sizes['y'] + + +def test_to_geotiff_no_rioxarray(test_var, monkeypatch): + """Test that a helpful error is raised when rioxarray is not installed.""" + import builtins + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name == 'rioxarray': + raise ImportError('No module named rioxarray') + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, '__import__', fake_import) + with pytest.raises(ImportError, match='rioxarray is required'): + test_var.metpy.to_geotiff('unused.tif') + + +def test_to_geotiff_no_spatial_coords(test_ds_generic, tmp_path): + """Test that writing data without x/y dimension coordinates raises a helpful error.""" + pytest.importorskip('rioxarray') + with pytest.raises(ValueError, match='x and y dimension coordinates are required'): + test_ds_generic['test'].metpy.to_geotiff(tmp_path / 'unused.tif') + + def test_no_projection(test_ds): """Test getting the crs attribute when not available produces a sensible error.""" var = test_ds.lat