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
11 changes: 7 additions & 4 deletions tests/test_20_open_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,6 @@ def test_open_dataset(tmp_path: Path, index_node: str, download: bool) -> None:
assert set(ds.coords) == {
"areacella",
"experiment_id",
"height",
"lat",
"lat_bnds",
"lon",
Expand All @@ -77,15 +76,19 @@ def test_open_dataset(tmp_path: Path, index_node: str, download: bool) -> None:
if name != "experiment_id"
)

# Dimensionless coords
assert "height" not in ds["pr"].attrs
assert ds["tas"].attrs["height"] == 2.0

# Data vars
assert set(ds.data_vars) == {"pr", "tas"}
assert ds["pr"].coordinates == "areacella experiment_id height lat lon time"
assert ds["tas"].coordinates == "areacella experiment_id height lat lon time"
assert ds["pr"].coordinates == "areacella experiment_id lat lon time"
assert ds["tas"].coordinates == "areacella experiment_id lat lon time"

# Attributes
assert (
ds.coordinates
== "areacella experiment_id height lat lat_bnds lon lon_bnds time time_bnds"
== "areacella experiment_id lat lat_bnds lon lon_bnds time time_bnds"
)
assert ds.dataset_ids == [
"CMIP6.ScenarioMIP.EC-Earth-Consortium.EC-Earth3-CC.ssp245.r1i1p1f1.Amon.pr.gr.v20210113",
Expand Down
15 changes: 15 additions & 0 deletions xarray_esgf/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,17 @@ def add_coordinates_attr_inplace(obj: Dataset | DataArray) -> None:
add_coordinates_attr_inplace(da)


def move_dimensionless_coords_to_attrs(ds: Dataset) -> Dataset:
attrs = {}
for var, da in ds.coords.items():
if not da.dims:
attrs[var] = da.item()
ds = ds.drop_vars(list(attrs))
for da in ds.data_vars.values():
da.attrs.update(attrs)
return ds


@dataclasses.dataclass
class Client:
selection: dict[str, str | list[str]]
Expand Down Expand Up @@ -183,11 +194,15 @@ def _open_datasets(
(ds,) = datasets
else:
ds = combine_datasets(datasets)

ds = ds.set_coords([
name
for name, da in ds.variables.items()
if BOUNDS_DIMS.intersection(da.dims) or "time" not in da.dims
])

ds = move_dimensionless_coords_to_attrs(ds)

ds = ds.expand_dims({dim: [dataset_id_dict[dim]] for dim in concat_dims})
combined_datasets[dataset_id] = ds
LOGGER.debug(f"{dataset_id}: {dict(ds.sizes)}")
Expand Down