Add attributes to generated variables. (#45)

Author: dcherian

docs/raster_index/design_choices.md

@@ -59,3 +59,9 @@ Thus, {py:func}`assign_index` will delete the `"GeoTransform"` attribute on the
 to construct the affine matrix.
 
 If you wish to extract the GeoTransform attribute to write it to a location of your choosing use {py:meth}`RasterIndex.as_geotransform`.
+
+## Attributes on coordinate variables
+
+Currently `assign_index` will first drop any existing coordinate variables and then recreate them as lazy coordinate variables.
+This loses any existing attributes, and new attributes appropriate to the coordinate system (generated by {py:meth}`pyproj.CRS.cs_to_cf`)
+are added.

src/rasterix/raster_index.py

@@ -494,6 +494,8 @@ def from_transform(
         height: int,
         x_dim: str = "x",
         y_dim: str = "y",
+        x_coord_name: str = "xc",
+        y_coord_name: str = "yc",
         crs: CRS | Any | None = None,
     ) -> RasterIndex:
         """Create a RasterIndex from an affine transform and raster dimensions.
@@ -507,10 +509,14 @@ def from_transform(
             Number of pixels in the x direction.
         height : int
             Number of pixels in the y direction.
-        x_dim : str, default "x"
+        x_dim : str, optional
             Name for the x dimension.
-        y_dim : str, default "y"
+        y_dim : str, optional
             Name for the y dimension.
+        x_coord_name : str, optional
+            Name for the x dimension. For non-rectilinear transforms only.
+        y_coord_name : str, optional
+            Name for the y dimension. For non-rectilinear transforms only.
         crs : :class:`pyproj.crs.CRS` or any, optional
             The coordinate reference system. Any value accepted by
             :meth:`pyproj.crs.CRS.from_user_input`.
@@ -533,21 +539,35 @@ def from_transform(
         >>> from affine import Affine
         >>> transform = Affine(1.0, 0.0, 0.0, 0.0, -1.0, 100.0)
         >>> index = RasterIndex.from_transform(transform, width=100, height=100)
+
+        Create a non-rectilinear index:
+
+        >>> index = RasterIndex.from_transform(
+        ...     Affine.rotation(45), width=100, height=100, x_coord_name="x1", y_coord_name="x2"
+        ... )
         """
         index: WrappedIndex
 
         # pixel centered coordinates
         affine = affine * Affine.translation(0.5, 0.5)
 
         if affine.is_rectilinear and affine.b == affine.d == 0:
-            x_transform = AxisAffineTransform(affine, width, "x", x_dim, is_xaxis=True)
-            y_transform = AxisAffineTransform(affine, height, "y", y_dim, is_xaxis=False)
+            x_transform = AxisAffineTransform(affine, width, x_dim, x_dim, is_xaxis=True)
+            y_transform = AxisAffineTransform(affine, height, y_dim, y_dim, is_xaxis=False)
             index = (
                 AxisAffineTransformIndex(x_transform),
                 AxisAffineTransformIndex(y_transform),
             )
         else:
-            xy_transform = AffineTransform(affine, width, height, x_dim=x_dim, y_dim=y_dim)
+            xy_transform = AffineTransform(
+                affine,
+                width,
+                height,
+                x_dim=x_dim,
+                y_dim=y_dim,
+                x_coord_name=x_coord_name,
+                y_coord_name=y_coord_name,
+            )
             index = CoordinateTransformIndex(xy_transform)
 
         return cls(index, crs=crs)
@@ -568,6 +588,18 @@ def create_variables(self, variables: Mapping[Any, Variable] | None = None) -> d
         for index in self._wrapped_indexes:
             new_variables.update(index.create_variables())
 
+        if self.crs is not None:
+            xname, yname = self.xy_coord_names
+            xattrs, yattrs = self.crs.cs_to_cf()
+            if "axis" in xattrs and "axis" in yattrs:
+                # The axis order is defined by the projection
+                # So we have to figure which is which.
+                # This is an ugly hack that works for common cases.
+                if xattrs["axis"] == "Y":
+                    xattrs, yattrs = yattrs, xattrs
+                new_variables[xname].attrs = xattrs
+                new_variables[yname].attrs = yattrs
+
         return new_variables
 
     @property

tests/test_raster_index.py

@@ -75,6 +75,11 @@ def test_raster_index_properties():
     )
     assert index3.as_geotransform(decimals=6) == "0.000000 0.707107 -0.707107 0.000000 0.707107 0.707107"
 
+    index4 = RasterIndex.from_transform(
+        Affine.rotation(45.0), width=12, height=10, x_coord_name="x1", y_coord_name="x2"
+    )
+    assert index4.xy_coord_names == ("x1", "x2")
+
 
 # TODO: parameterize over
 # 1. y points up;
@@ -107,12 +112,51 @@ def test_sel_slice():
     assert_identical(reverse.y, ds.y[:5:-1])
 
 
-def test_crs() -> None:
+def test_crs_generated_attributes() -> None:
     index = RasterIndex.from_transform(Affine.identity(), width=12, height=10)
     assert index.crs is None
+    variables = index.create_variables()
+    assert variables["x"].attrs == {}
+    assert variables["y"].attrs == {}
 
     index = RasterIndex.from_transform(Affine.identity(), width=12, height=10, crs="epsg:31370")
     assert index.crs == pyproj.CRS.from_user_input("epsg:31370")
+    variables = index.create_variables()
+    assert variables["x"].attrs == {
+        "axis": "X",
+        "long_name": "Easting",
+        "standard_name": "projection_x_coordinate",
+        "units": "metre",
+    }
+    assert variables["y"].attrs == {
+        "axis": "Y",
+        "long_name": "Northing",
+        "standard_name": "projection_y_coordinate",
+        "units": "metre",
+    }
+
+    index = RasterIndex.from_transform(
+        Affine.identity(),
+        width=12,
+        height=10,
+        x_dim="lon",
+        y_dim="lat",
+        crs="epsg:4326",
+    )
+    assert index.crs == pyproj.CRS.from_user_input("epsg:4326")
+    variables = index.create_variables()
+    assert variables["lon"].attrs == {
+        "axis": "X",
+        "long_name": "longitude coordinate",
+        "standard_name": "longitude",
+        "units": "degrees_east",
+    }
+    assert variables["lat"].attrs == {
+        "axis": "Y",
+        "long_name": "latitude coordinate",
+        "standard_name": "latitude",
+        "units": "degrees_north",
+    }
 
 
 # asserting (in)equality for both "x" and "y" is redundant but not harmful