Repository: corteva/rioxarray

Issue:
Add ability to set CRS and write CRS information to dataset
See: https://github.com/geoxarray/geoxarray/issues/9

Comments:

Patch:
diff --git a/rioxarray/exceptions.py b/rioxarray/exceptions.py
index 723a391..99a58fd 100644
--- a/rioxarray/exceptions.py
+++ b/rioxarray/exceptions.py
@@ -26,3 +26,7 @@ class TooManyDimensions(RioXarrayError):
 
 class InvalidDimensionOrder(RioXarrayError):
     """This is raised when there the dimensions are not ordered correctly."""
+
+
+class MissingCRS(RioXarrayError):
+    """Missing the CRS in the dataset."""
diff --git a/rioxarray/rioxarray.py b/rioxarray/rioxarray.py
index acefb4b..1a9b0ec 100644
--- a/rioxarray/rioxarray.py
+++ b/rioxarray/rioxarray.py
@@ -25,6 +25,7 @@ from scipy.interpolate import griddata
 
 from rioxarray.exceptions import (
     InvalidDimensionOrder,
+    MissingCRS,
     NoDataInBounds,
     OneDimensionalRaster,
     TooManyDimensions,
@@ -32,7 +33,7 @@ from rioxarray.exceptions import (
 from rioxarray.crs import crs_to_wkt
 
 FILL_VALUE_NAMES = ("_FillValue", "missing_value", "fill_value", "nodata")
-UNWANTED_RIO_ATTRS = ("nodata", "nodatavals", "crs", "is_tiled", "res")
+UNWANTED_RIO_ATTRS = ("nodatavals", "crs", "is_tiled", "res")
 DEFAULT_GRID_MAP = "spatial_ref"
 
 
@@ -123,23 +124,9 @@ def add_xy_grid_meta(coords):
 
 
 def add_spatial_ref(in_ds, dst_crs, grid_map_name):
-    # remove old grid map if exists
-    try:
-        del in_ds.coords[grid_map_name]
-    except KeyError:
-        pass
-
-    # add grid mapping variable
-    in_ds.coords[grid_map_name] = xarray.Variable((), 0)
-    match_proj = crs_to_wkt(CRS.from_user_input(dst_crs))
-
-    grid_map_attrs = dict()
-    # add grid mapping variable
-    grid_map_attrs["spatial_ref"] = match_proj
-    # http://cfconventions.org/Data/cf-conventions/cf-conventions-1.7/cf-conventions.html#appendix-grid-mappings
-    # http://desktop.arcgis.com/en/arcmap/10.3/manage-data/netcdf/spatial-reference-for-netcdf-data.htm
-    grid_map_attrs["crs_wkt"] = match_proj
-    in_ds.coords[grid_map_name].attrs = grid_map_attrs
+    in_ds.rio.write_crs(
+        input_crs=dst_crs, grid_mapping_name=grid_map_name, inplace=True
+    )
     return in_ds
 
 
@@ -230,7 +217,93 @@ class XRasterBase(object):
         """:obj:`rasterio.crs.CRS`:
             The projection of the dataset.
         """
-        return self._crs
+        raise NotImplementedError
+
+    def set_crs(self, input_crs, inplace=True):
+        """
+        Set the CRS value for the Dataset/DataArray without modifying
+        the dataset/data array.
+
+        Parameters:
+        -----------
+        input_crs: object
+            Anythong accepted by `rasterio.crs.CRS.from_user_input`.
+
+        Returns:
+        --------
+        xarray.Dataset or xarray.DataArray:
+        Dataset with crs attribute.
+
+        """
+        crs = CRS.from_user_input(input_crs)
+        if inplace:
+            self._crs = crs
+            return self._obj
+        xobj = self._obj.copy(deep=True)
+        xobj.rio._crs = crs
+        return xobj
+
+    def write_crs(
+        self, input_crs=None, grid_mapping_name=DEFAULT_GRID_MAP, inplace=False
+    ):
+        """
+        Write the CRS to the dataset in a CF compliant manner.
+
+        Parameters:
+        -----------
+        input_crs: object
+            Anythong accepted by `rasterio.crs.CRS.from_user_input`.
+        grid_mapping_name: str, optional
+            Name of the coordinate to store the CRS information in.
+        inplace: bool, optional
+            If True, it will write to the existing dataset. Default is False.
+
+        Returns:
+        --------
+        xarray.Dataset or xarray.DataArray:
+        Modified dataset with CF compliant CRS information.
+
+        """
+        if input_crs is not None:
+            data_obj = self.set_crs(input_crs, inplace=inplace)
+        elif inplace:
+            data_obj = self._obj
+        else:
+            data_obj = self._obj.copy(deep=True)
+
+        # remove old grid maping coordinate if exists
+        try:
+            del data_obj.coords[grid_mapping_name]
+        except KeyError:
+            pass
+
+        if data_obj.rio.crs is None:
+            raise MissingCRS("CRS not found. Please set the CRS with 'set_crs()'.")
+        # add grid mapping coordinate
+        data_obj.coords[grid_mapping_name] = xarray.Variable((), 0)
+        crs_wkt = crs_to_wkt(data_obj.rio.crs)
+        grid_map_attrs = dict()
+        grid_map_attrs["spatial_ref"] = crs_wkt
+        # http://cfconventions.org/Data/cf-conventions/cf-conventions-1.7/cf-conventions.html#appendix-grid-mappings
+        # http://desktop.arcgis.com/en/arcmap/10.3/manage-data/netcdf/spatial-reference-for-netcdf-data.htm
+        grid_map_attrs["crs_wkt"] = crs_wkt
+        data_obj.coords[grid_mapping_name].attrs = grid_map_attrs
+
+        # add grid mapping attribute to variables
+        if hasattr(data_obj, "data_vars"):
+            for var in data_obj.data_vars:
+                if (
+                    self.x_dim in data_obj[var].dims
+                    and self.y_dim in data_obj[var].dims
+                ):
+                    var_attrs = dict(data_obj[var].attrs)
+                    var_attrs.update(grid_mapping=grid_mapping_name)
+                    data_obj[var].attrs = var_attrs
+        else:
+            var_attrs = dict(data_obj.attrs)
+            var_attrs.update(grid_mapping=grid_mapping_name)
+            data_obj.attrs = var_attrs
+        return data_obj
 
     @property
     def shape(self):
@@ -255,17 +328,23 @@ class RasterArray(XRasterBase):
             Retrieve projection from `xarray.DataArray`
         """
         if self._crs is not None:
-            return self._crs
+            return None if self._crs is False else self._crs
 
         try:
             # look in grid_mapping
             grid_mapping_coord = self._obj.attrs["grid_mapping"]
-            self._crs = CRS.from_user_input(
-                self._obj.coords[grid_mapping_coord].attrs["spatial_ref"]
-            )
+            try:
+                crs_wkt = self._obj.coords[grid_mapping_coord].attrs["spatial_ref"]
+            except KeyError:
+                crs_wkt = self._obj.coords[grid_mapping_coord].attrs["crs_wkt"]
+            self._crs = CRS.from_user_input(crs_wkt)
         except KeyError:
-            # look in attrs for 'crs' from rasterio xarray
-            self._crs = CRS.from_user_input(self._obj.attrs["crs"])
+            try:
+                # look in attrs for 'crs' from rasterio xarray
+                self._crs = CRS.from_user_input(self._obj.attrs["crs"])
+            except KeyError:
+                self._crs = False
+                return None
         return self._crs
 
     @property
@@ -326,7 +405,7 @@ class RasterArray(XRasterBase):
             left, bottom, right, top = self._int_bounds()
 
             if width == 1 or height == 1:
-                raise ValueError(
+                raise OneDimensionalRaster(
                     "Only 1 dimenional array found. Cannot calculate the resolution."
                 )
 
diff --git a/sphinx/history.rst b/sphinx/history.rst
index c2d61e7..b739537 100644
--- a/sphinx/history.rst
+++ b/sphinx/history.rst
@@ -6,6 +6,7 @@ History
 - Find nodata and nodatavals in 'nodata' property (pull #12)
 - Added 'encoded_nodata' property to DataArray (pull #12)
 - Write the raster with encoded_nodata instead of NaN for nodata (pull #12)
+- Added methods to set and write CRS (issue #5)
 
 0.0.4
 ------


Test patch:
diff --git a/test/integration/test_integration_rioxarray.py b/test/integration/test_integration_rioxarray.py
index 1708a01..8f3311c 100644
--- a/test/integration/test_integration_rioxarray.py
+++ b/test/integration/test_integration_rioxarray.py
@@ -7,7 +7,7 @@ import xarray
 from numpy.testing import assert_almost_equal, assert_array_equal
 from rasterio.crs import CRS
 
-from rioxarray.exceptions import NoDataInBounds, OneDimensionalRaster
+from rioxarray.exceptions import MissingCRS, NoDataInBounds, OneDimensionalRaster
 from rioxarray.rioxarray import UNWANTED_RIO_ATTRS, _make_coords
 from test.conftest import TEST_COMPARE_DATA_DIR, TEST_INPUT_DATA_DIR
 
@@ -508,9 +508,9 @@ def test_make_src_affine__single_point():
         calculated_transform = tuple(xdi.rio.transform(recalc=True))
         # delete the transform to ensure it is not being used
         del xdi.attrs["transform"]
-        with pytest.raises(ValueError):
+        with pytest.raises(OneDimensionalRaster):
             xdi.rio.transform(recalc=True)
-        with pytest.raises(ValueError):
+        with pytest.raises(OneDimensionalRaster):
             xdi.rio.transform()
 
         assert_array_equal(attribute_transform, attribute_transform_func)
@@ -726,3 +726,134 @@ def test_to_raster_3d(tmpdir):
         assert_array_equal(rds.transform, xds.rio.transform())
         assert_array_equal(rds.nodata, xds.rio.nodata)
         assert_array_equal(rds.read(), xds.values)
+
+
+def test_missing_dimensions():
+    test_ds = xarray.Dataset()
+    with pytest.raises(KeyError):
+        test_ds.rio
+
+
+def test_crs_setter():
+    test_da = xarray.DataArray(
+        numpy.zeros((5, 5)),
+        dims=("y", "x"),
+        coords={"y": numpy.arange(1, 6), "x": numpy.arange(2, 7)},
+    )
+    assert test_da.rio.crs is None
+    out_ds = test_da.rio.set_crs(4326)
+    assert test_da.rio.crs.to_epsg() == 4326
+    assert out_ds.rio.crs.to_epsg() == 4326
+    test_ds = test_da.to_dataset(name="test")
+    assert test_ds.rio.crs is None
+    out_ds = test_ds.rio.set_crs(4326)
+    assert test_ds.rio.crs.to_epsg() == 4326
+    assert out_ds.rio.crs.to_epsg() == 4326
+
+
+def test_crs_setter__copy():
+    test_da = xarray.DataArray(
+        numpy.zeros((5, 5)),
+        dims=("y", "x"),
+        coords={"y": numpy.arange(1, 6), "x": numpy.arange(2, 7)},
+    )
+    assert test_da.rio.crs is None
+    out_ds = test_da.rio.set_crs(4326, inplace=False)
+    assert test_da.rio.crs is None
+    assert out_ds.rio.crs.to_epsg() == 4326
+    test_ds = test_da.to_dataset(name="test")
+    assert test_ds.rio.crs is None
+    out_ds = test_ds.rio.set_crs(4326, inplace=False)
+    assert test_ds.rio.crs is None
+    assert out_ds.rio.crs.to_epsg() == 4326
+
+
+def test_crs_writer__array__copy():
+    test_da = xarray.DataArray(
+        numpy.zeros((5, 5)),
+        dims=("y", "x"),
+        coords={"y": numpy.arange(1, 6), "x": numpy.arange(2, 7)},
+    )
+    assert test_da.rio.crs is None
+    out_da = test_da.rio.write_crs(4326, grid_mapping_name="crs")
+    assert "crs_wkt" in out_da.coords["crs"].attrs
+    assert "spatial_ref" in out_da.coords["crs"].attrs
+    out_da.rio._crs = None
+    assert out_da.rio.crs.to_epsg() == 4326
+    test_da.rio._crs = None
+    assert test_da.rio.crs is None
+    assert "crs" not in test_da.coords
+    assert out_da.attrs["grid_mapping"] == "crs"
+
+
+def test_crs_writer__array__inplace():
+    test_da = xarray.DataArray(
+        numpy.zeros((5, 5)),
+        dims=("y", "x"),
+        coords={"y": numpy.arange(1, 6), "x": numpy.arange(2, 7)},
+    )
+    assert test_da.rio.crs is None
+    out_da = test_da.rio.write_crs(4326, inplace=True)
+    assert "crs_wkt" in test_da.coords["spatial_ref"].attrs
+    assert "spatial_ref" in test_da.coords["spatial_ref"].attrs
+    assert out_da.coords["spatial_ref"] == test_da.coords["spatial_ref"]
+    test_da.rio._crs = None
+    assert test_da.rio.crs.to_epsg() == 4326
+    assert test_da.attrs["grid_mapping"] == "spatial_ref"
+    assert out_da.attrs == test_da.attrs
+    out_da.rio._crs = None
+    assert out_da.rio.crs.to_epsg() == 4326
+
+
+def test_crs_writer__dataset__copy():
+    test_da = xarray.DataArray(
+        numpy.zeros((5, 5)),
+        dims=("y", "x"),
+        coords={"y": numpy.arange(1, 6), "x": numpy.arange(2, 7)},
+    )
+    test_da = test_da.to_dataset(name="test")
+    assert test_da.rio.crs is None
+    out_da = test_da.rio.write_crs(4326, grid_mapping_name="crs")
+    assert "crs_wkt" in out_da.coords["crs"].attrs
+    assert "spatial_ref" in out_da.coords["crs"].attrs
+    out_da.test.rio._crs = None
+    assert out_da.rio.crs.to_epsg() == 4326
+    assert out_da.test.attrs["grid_mapping"] == "crs"
+    # make sure input did not change the dataset
+    test_da.test.rio._crs = None
+    test_da.rio._crs = None
+    assert test_da.rio.crs is None
+    assert "crs" not in test_da.coords
+
+
+def test_crs_writer__dataset__inplace():
+    test_da = xarray.DataArray(
+        numpy.zeros((5, 5)),
+        dims=("y", "x"),
+        coords={"y": numpy.arange(1, 6), "x": numpy.arange(2, 7)},
+    )
+    test_da = test_da.to_dataset(name="test")
+    assert test_da.rio.crs is None
+    out_da = test_da.rio.write_crs(4326, inplace=True)
+    assert "crs_wkt" in test_da.coords["spatial_ref"].attrs
+    assert "spatial_ref" in test_da.coords["spatial_ref"].attrs
+    assert out_da.coords["spatial_ref"] == test_da.coords["spatial_ref"]
+    out_da.test.rio._crs = None
+    assert out_da.rio.crs.to_epsg() == 4326
+    test_da.test.rio._crs = None
+    test_da.rio._crs = None
+    assert test_da.rio.crs.to_epsg() == 4326
+    assert out_da.test.attrs["grid_mapping"] == "spatial_ref"
+    assert out_da.test.attrs == test_da.test.attrs
+
+
+def test_crs_writer__missing():
+    test_da = xarray.DataArray(
+        numpy.zeros((5, 5)),
+        dims=("y", "x"),
+        coords={"y": numpy.arange(1, 6), "x": numpy.arange(2, 7)},
+    )
+    with pytest.raises(MissingCRS):
+        test_da.rio.write_crs()
+    with pytest.raises(MissingCRS):
+        test_da.to_dataset(name="test").rio.write_crs()