Add reduced gaussian grid mapping support (#613)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: dcherian

cf_xarray/accessor.py

@@ -589,7 +589,44 @@ def _create_grid_mapping(
     cf_name = var.attrs.get("grid_mapping_name", var_name)
 
     # Create CRS from the grid mapping variable
-    if cf_name == "healpix":
+    if cf_name == "reduced_gaussian":
+        # pyproj does not recognize "reduced_gaussian" as a grid mapping name,
+        # but the grid uses geographic (lat/lon) coordinates on a sphere or
+        # spheroid. Build a geographic CRS from the earth shape parameters.
+        crs = pyproj.CRS.from_json_dict(
+            {
+                "$schema": "https://proj.org/schemas/v0.6/projjson.schema.json",
+                "type": "GeographicCRS",
+                "name": "Reduced Gaussian Grid",
+                "datum": {
+                    "type": "GeodeticReferenceFrame",
+                    "name": "Unknown",
+                    "ellipsoid": {
+                        "name": "Custom",
+                        "semi_major_axis": var.attrs.get("semi_major_axis", 6371229.0),
+                        "semi_minor_axis": var.attrs.get("semi_minor_axis", 6371229.0),
+                    },
+                },
+                "coordinate_system": {
+                    "subtype": "ellipsoidal",
+                    "axis": [
+                        {
+                            "name": "Latitude",
+                            "abbreviation": "lat",
+                            "direction": "north",
+                            "unit": "degree",
+                        },
+                        {
+                            "name": "Longitude",
+                            "abbreviation": "lon",
+                            "direction": "east",
+                            "unit": "degree",
+                        },
+                    ],
+                },
+            }
+        )
+    elif cf_name == "healpix":
         # pyproj does not recognize "healpix" as a grid mapping name,
         # but the grid uses geographic (lat/lon) coordinates on a sphere.
         # Build a geographic CRS from the earth_radius parameter.
@@ -640,19 +677,46 @@ def _create_grid_mapping(
     # """
     if not coordinates and len(grid_mapping_dict) == 1:
         if cf_name == "healpix":
-            # For HEALPix grids, coordinates are typically latitude/longitude
-            # (if present), but pixel indices are the primary indexing mechanism.
-            xname, yname = "longitude", "latitude"
-        elif crs.to_cf().get("grid_mapping_name") == "rotated_latitude_longitude":
-            xname, yname = "grid_longitude", "grid_latitude"
-        elif crs.is_geographic:
-            xname, yname = "longitude", "latitude"
-        elif crs.is_projected:
-            xname, yname = "projection_x_coordinate", "projection_y_coordinate"
-
-        x = apply_mapper(_get_with_standard_name, ds, xname, error=False, default=[[]])
-        y = apply_mapper(_get_with_standard_name, ds, yname, error=False, default=[[]])
-        coordinates = list(itertools.chain(x, y))
+            # For HEALPix grids, the primary coordinate is the pixel index.
+            coords_found = apply_mapper(
+                _get_with_standard_name, ds, "healpix_index", error=False, default=[[]]
+            )
+            coordinates = list(itertools.chain(coords_found))
+        elif cf_name == "reduced_gaussian":
+            # For reduced gaussian grids, the primary coordinate is the grid
+            # point index. For compressed subsets, also look for the gather
+            # variable (with compress attribute).
+            idx_coords = apply_mapper(
+                _get_with_standard_name,
+                ds,
+                "reduced_gaussian_index",
+                error=False,
+                default=[[]],
+            )
+            coordinates = list(itertools.chain(idx_coords))
+            # Also find any compress/gather variable that references
+            # the detected index coordinate(s)
+            for vname in ds.coords:
+                if "compress" in ds[vname].attrs:
+                    compress_target = ds[vname].attrs["compress"]
+                    if any(c in compress_target for c in coordinates):
+                        if vname not in coordinates:
+                            coordinates.append(vname)
+        else:
+            if crs.to_cf().get("grid_mapping_name") == "rotated_latitude_longitude":
+                xname, yname = "grid_longitude", "grid_latitude"
+            elif crs.is_geographic:
+                xname, yname = "longitude", "latitude"
+            elif crs.is_projected:
+                xname, yname = "projection_x_coordinate", "projection_y_coordinate"
+
+            x = apply_mapper(
+                _get_with_standard_name, ds, xname, error=False, default=[[]]
+            )
+            y = apply_mapper(
+                _get_with_standard_name, ds, yname, error=False, default=[[]]
+            )
+            coordinates = list(itertools.chain(x, y))
 
     return GridMapping(name=cf_name, crs=crs, array=da, coordinates=tuple(coordinates))
 

cf_xarray/datasets.py

@@ -861,3 +861,145 @@ def encoded_point_dataset():
         {"geometry": "geometry_container"},
     )
     return ds
+
+
+# --- Reduced Gaussian Grid test fixtures ---
+# A tiny O2 octahedral reduced gaussian grid with 4 latitudes and 40 total points.
+
+
+def _create_reduced_gaussian_global():
+    """Create a small O2 reduced gaussian grid dataset (full global)."""
+    lat = np.array([59.44, 19.47, -19.47, -59.44])
+    pl = np.array([8, 12, 12, 8], dtype=np.int32)
+    total = int(np.sum(pl))  # 40
+
+    rng = np.random.default_rng(42)
+    temp_data = rng.standard_normal((1, 1, total)).astype(np.float32)
+
+    ds = xr.Dataset(
+        {
+            "air_temperature": xr.DataArray(
+                temp_data,
+                dims=["time", "height", "reduced_gaussian_index"],
+                attrs={
+                    "grid_mapping": "reduced_gaussian",
+                    "coordinates": "reduced_gaussian_index",
+                    "standard_name": "air_temperature",
+                    "units": "K",
+                },
+            ),
+            "reduced_gaussian": xr.DataArray(
+                np.int32(0),
+                attrs={
+                    "grid_mapping_name": "reduced_gaussian",
+                    "grid_subtype": "octahedral",
+                    "points_per_latitude": "pl",
+                    "latitude_dimension": "lat",
+                    "semi_major_axis": 6371229.0,
+                    "semi_minor_axis": 6371229.0,
+                },
+            ),
+        },
+        coords={
+            "lat": xr.DataArray(
+                lat,
+                dims=["lat"],
+                attrs={"standard_name": "latitude", "units": "degrees_north"},
+            ),
+            "pl": xr.DataArray(
+                pl,
+                dims=["lat"],
+                attrs={"long_name": "number of points per latitude"},
+            ),
+            "reduced_gaussian_index": xr.DataArray(
+                np.arange(total, dtype=np.int32),
+                dims=["reduced_gaussian_index"],
+                attrs={"standard_name": "reduced_gaussian_index"},
+            ),
+            "time": xr.DataArray(
+                [0.0], dims=["time"], attrs={"units": "hours since 2024-01-01"}
+            ),
+            "height": xr.DataArray([2.0], dims=["height"], attrs={"units": "m"}),
+        },
+    )
+    return ds
+
+
+def _create_reduced_gaussian_land():
+    """Create a small O2 reduced gaussian grid with compression by gathering (land subset)."""
+    global_ds = _create_reduced_gaussian_global()
+    land_indices = np.array([2, 5, 10, 15, 20, 25, 30, 35], dtype=np.int32)
+
+    rng = np.random.default_rng(43)
+    temp_data = rng.standard_normal((1, 1, len(land_indices))).astype(np.float32)
+
+    ds = xr.Dataset(
+        {
+            "air_temperature": xr.DataArray(
+                temp_data,
+                dims=["time", "height", "grid_points"],
+                attrs={
+                    "grid_mapping": "reduced_gaussian",
+                    "coordinates": "time height reduced_gaussian_index",
+                },
+            ),
+            "reduced_gaussian": global_ds["reduced_gaussian"],
+        },
+        coords={
+            "lat": global_ds["lat"],
+            "pl": global_ds["pl"],
+            "reduced_gaussian_index": global_ds["reduced_gaussian_index"],
+            "grid_points": xr.DataArray(
+                land_indices,
+                dims=["grid_points"],
+                attrs={"compress": "reduced_gaussian_index"},
+            ),
+            "time": global_ds["time"],
+            "height": global_ds["height"],
+        },
+    )
+    return ds
+
+
+def _create_reduced_gaussian_region():
+    """Create a small O2 reduced gaussian grid with compression (regional subset)."""
+    global_ds = _create_reduced_gaussian_global()
+    region_indices = np.array([8, 9, 10, 11, 12, 13, 14], dtype=np.int32)
+
+    rng = np.random.default_rng(44)
+    temp_data = rng.standard_normal((1, 1, len(region_indices))).astype(np.float32)
+
+    ds = xr.Dataset(
+        {
+            "air_temperature": xr.DataArray(
+                temp_data,
+                dims=["time", "height", "grid_points"],
+                attrs={
+                    "grid_mapping": "reduced_gaussian",
+                    "coordinates": "time height reduced_gaussian_index",
+                },
+            ),
+            "reduced_gaussian": global_ds["reduced_gaussian"],
+        },
+        coords={
+            "lat": global_ds["lat"],
+            "pl": global_ds["pl"],
+            "reduced_gaussian_index": global_ds["reduced_gaussian_index"],
+            "grid_points": xr.DataArray(
+                region_indices,
+                dims=["grid_points"],
+                attrs={
+                    "standard_name": "reduced_gaussian_index",
+                    "compress": "reduced_gaussian_index",
+                },
+            ),
+            "time": global_ds["time"],
+            "height": global_ds["height"],
+        },
+    )
+    return ds
+
+
+reduced_gaussian_global_ds = _create_reduced_gaussian_global()
+reduced_gaussian_land_ds = _create_reduced_gaussian_land()
+reduced_gaussian_region_ds = _create_reduced_gaussian_region()

cf_xarray/tests/test_accessor.py

@@ -1264,6 +1264,132 @@ def test_grid_mappings_coordinates_attribute():
             )
 
 
+@requires_pyproj
+def test_reduced_gaussian_grid_mapping_global():
+    """Test GridMapping integration for a full reduced gaussian grid."""
+    from ..datasets import reduced_gaussian_global_ds as ds
+
+    # Grid mapping discovery
+    assert ds.cf.grid_mapping_names == {"reduced_gaussian": ["reduced_gaussian"]}
+
+    # CF indexing returns the grid mapping variable
+    gm_var = ds.cf["reduced_gaussian"]
+    assert gm_var.attrs["grid_mapping_name"] == "reduced_gaussian"
+    assert gm_var.attrs["grid_subtype"] == "octahedral"
+
+    # Grid mapping propagation to data variables
+    da = ds.cf["air_temperature"]
+    assert "reduced_gaussian" in da.coords
+
+    # grid_mapping_name property
+    assert da.cf.grid_mapping_name == "reduced_gaussian"
+
+    # .cf.grid_mappings property returns valid GridMapping
+    gms = ds.cf.grid_mappings
+    assert len(gms) == 1
+    gm = gms[0]
+    assert gm.name == "reduced_gaussian"
+    assert gm.crs is not None
+    assert gm.crs.is_geographic
+    assert gm.array.name == "reduced_gaussian"
+    assert gm.array.shape == ()  # scalar variable
+    assert isinstance(gm.coordinates, tuple)
+    # Should detect reduced_gaussian_index via standard_name
+    assert "reduced_gaussian_index" in gm.coordinates
+
+    # DataArray grid_mappings should also work
+    da_gms = da.cf.grid_mappings
+    assert len(da_gms) == 1
+    assert da_gms[0].name == "reduced_gaussian"
+    assert da_gms[0].crs.is_geographic
+
+    # Repr should include reduced_gaussian
+    assert "reduced_gaussian" in ds.cf.__repr__()
+
+
+@requires_pyproj
+def test_reduced_gaussian_grid_mapping_land():
+    """Test GridMapping for a land-only subset using CF compress/gather."""
+    from ..datasets import reduced_gaussian_land_ds as ds
+
+    # Grid mapping discovery
+    assert ds.cf.grid_mapping_names == {"reduced_gaussian": ["reduced_gaussian"]}
+
+    # Grid mapping propagation
+    da = ds.cf["air_temperature"]
+    assert "reduced_gaussian" in da.coords
+
+    # .cf.grid_mappings property
+    gms = ds.cf.grid_mappings
+    assert len(gms) == 1
+    gm = gms[0]
+    assert gm.name == "reduced_gaussian"
+    assert gm.crs.is_geographic
+    assert gm.array.attrs["grid_subtype"] == "octahedral"
+
+    # Coordinates should include the compress/gather variable
+    assert "grid_points" in gm.coordinates
+
+    # Verify the dataset structure: data on grid_points, with compress attribute
+    assert "grid_points" in ds.dims
+    assert "compress" in ds.grid_points.attrs
+    assert ds.grid_points.attrs["compress"] == "reduced_gaussian_index"
+
+
+@requires_pyproj
+def test_reduced_gaussian_grid_mapping_region():
+    """Test GridMapping for a regional subset using CF compress/gather."""
+    from ..datasets import reduced_gaussian_region_ds as ds
+
+    # Grid mapping discovery
+    assert ds.cf.grid_mapping_names == {"reduced_gaussian": ["reduced_gaussian"]}
+
+    # Grid mapping propagation
+    da = ds.cf["air_temperature"]
+    assert "reduced_gaussian" in da.coords
+
+    # .cf.grid_mappings property
+    gms = ds.cf.grid_mappings
+    assert len(gms) == 1
+    gm = gms[0]
+    assert gm.name == "reduced_gaussian"
+    assert gm.crs.is_geographic
+
+    # Coordinates should include both reduced_gaussian_index and grid_points
+    assert "reduced_gaussian_index" in gm.coordinates
+    assert "grid_points" in gm.coordinates
+
+    # Verify compress structure
+    assert "grid_points" in ds.dims
+    assert ds.grid_points.attrs["compress"] == "reduced_gaussian_index"
+
+    # CRS earth parameters should match the grid mapping variable
+    gm_attrs = ds.reduced_gaussian.attrs
+    # The CRS should be built from semi_major/minor_axis
+    assert gm.crs.ellipsoid is not None
+    assert gm.crs.ellipsoid.semi_major_metre == gm_attrs["semi_major_axis"]
+
+
+@requires_pyproj
+def test_reduced_gaussian_crs_properties():
+    """Test that the CRS built for reduced_gaussian has correct properties."""
+    from ..datasets import reduced_gaussian_global_ds as ds
+
+    gm = ds.cf.grid_mappings[0]
+    crs = gm.crs
+
+    # Must be geographic (lat/lon on a sphere)
+    assert crs.is_geographic
+    assert not crs.is_projected
+
+    # Earth shape parameters from the grid mapping variable
+    assert crs.ellipsoid.semi_major_metre == 6371229.0
+    assert crs.ellipsoid.semi_minor_metre == 6371229.0
+
+    # Should not have an EPSG code (custom sphere)
+    assert crs.to_epsg() is None
+
+
 @requires_pyproj
 def test_bad_grid_mapping_attribute():
     ds = rotds.copy(deep=False)
@@ -1310,6 +1436,8 @@ def test_healpix_grid_mapping():
     assert gm.array.name == "healpix"
     assert gm.array.shape == ()  # scalar variable
     assert isinstance(gm.coordinates, tuple)
+    # Should detect healpix_index via standard_name
+    assert "healpix_index" in gm.coordinates
 
     # DataArray grid_mappings should also work
     da_gms = da.cf.grid_mappings