Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 78 additions & 14 deletions cf_xarray/accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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))

Expand Down
142 changes: 142 additions & 0 deletions cf_xarray/datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
128 changes: 128 additions & 0 deletions cf_xarray/tests/test_accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
Loading