xarray-esgf/tests/test_20_open_dataset.py at 03c54a767b4d262ae914caa76c3121d5ce511fdf · bopen/xarray-esgf · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
from pathlib import Path

import pytest
import xarray as xr


@pytest.mark.parametrize("download", [True, False])
def test_open_dataset(tmp_path: Path, index_node: str, download: bool) -> None:
    esgpull_path = tmp_path / "esgpull"
    selection = {
        "query": [
            '"tas_Amon_EC-Earth3-CC_ssp245_r1i1p1f1_gr_201901-201912.nc"',
            '"tas_Amon_EC-Earth3-CC_ssp245_r1i1p1f1_gr_202001-202012.nc"',
            '"tas_Amon_EC-Earth3-CC_ssp585_r1i1p1f1_gr_201901-201912.nc"',
            '"tas_Amon_EC-Earth3-CC_ssp585_r1i1p1f1_gr_202001-202012.nc"',
            '"pr_Amon_EC-Earth3-CC_ssp245_r1i1p1f1_gr_201901-201912.nc"',
            '"pr_Amon_EC-Earth3-CC_ssp245_r1i1p1f1_gr_202001-202012.nc"',
            '"pr_Amon_EC-Earth3-CC_ssp585_r1i1p1f1_gr_201901-201912.nc"',
            '"pr_Amon_EC-Earth3-CC_ssp585_r1i1p1f1_gr_202001-202012.nc"',
            '"CMIP6.ScenarioMIP.EC-Earth-Consortium.EC-Earth3-CC.ssp245.r1i1p1f1.fx.areacella.gr.v20210113.areacella_fx_EC-Earth3-CC_ssp245_r1i1p1f1_gr.nc"',
            '"CMIP6.ScenarioMIP.EC-Earth-Consortium.EC-Earth3-CC.ssp585.r1i1p1f1.fx.areacella.gr.v20210113.areacella_fx_EC-Earth3-CC_ssp585_r1i1p1f1_gr.nc"',
        ]
    }
    ds = xr.open_dataset(
        selection,  # type: ignore[arg-type]
        esgpull_path=esgpull_path,
        concat_dims="experiment_id",
        engine="esgf",
        index_node=index_node,
        download=download,
        chunks={},
    )

    assert (esgpull_path / "data" / "CMIP6").exists() is download

    # Chunks
    for dim in ds.dims:
        assert not ds[dim].chunks
    assert ds.chunksizes == {
        "experiment_id": (1, 1),
        "time": (12, 12),
        "lat": (256,),
        "lon": (512,),
        "bnds": (2,),
    }

    # Dims
    assert ds.sizes == {
        "experiment_id": 2,
        "time": 24,
        "bnds": 2,
        "lat": 256,
        "lon": 512,
    }

    # Coords
    assert set(ds.coords) == {
        "areacella",
        "experiment_id",
        "height",
        "lat",
        "lat_bnds",
        "lon",
        "lon_bnds",
        "time",
        "time_bnds",
    }
    assert all(
        "experiment_id" not in coord.dims
        for name, coord in ds.coords.items()
        if name != "experiment_id"
    )

    # Data vars
    assert set(ds.data_vars) == {"tas", "pr"}

    # Attributes
    assert ds.dataset_ids == [
        "CMIP6.ScenarioMIP.EC-Earth-Consortium.EC-Earth3-CC.ssp245.r1i1p1f1.Amon.pr.gr.v20210113",
        "CMIP6.ScenarioMIP.EC-Earth-Consortium.EC-Earth3-CC.ssp245.r1i1p1f1.Amon.tas.gr.v20210113",
        "CMIP6.ScenarioMIP.EC-Earth-Consortium.EC-Earth3-CC.ssp245.r1i1p1f1.fx.areacella.gr.v20210113",
        "CMIP6.ScenarioMIP.EC-Earth-Consortium.EC-Earth3-CC.ssp585.r1i1p1f1.Amon.pr.gr.v20210113",
        "CMIP6.ScenarioMIP.EC-Earth-Consortium.EC-Earth3-CC.ssp585.r1i1p1f1.Amon.tas.gr.v20210113",
        "CMIP6.ScenarioMIP.EC-Earth-Consortium.EC-Earth3-CC.ssp585.r1i1p1f1.fx.areacella.gr.v20210113",
    ]


def test_open_dataset_identical(tmp_path: Path, index_node: str) -> None:
    esgpull_path = tmp_path / "esgpull"
    selection = {
        "query": [
            '"tas_Amon_EC-Earth3-CC_ssp245_r1i1p1f1_gr_201901-201912.nc"',
        ]
    }
    ds_pydap = xr.open_dataset(
        selection,  # type: ignore[arg-type]
        esgpull_path=esgpull_path,
        engine="esgf",
        index_node=index_node,
        download=False,
        chunks={},
        use_pydap=True,
    )
    ds_h5netcdf = xr.open_dataset(
        selection,  # type: ignore[arg-type]
        esgpull_path=esgpull_path,
        engine="esgf",
        index_node=index_node,
        download=False,
        chunks={},
        use_pydap=False,
    )
    xr.testing.assert_identical(ds_pydap, ds_h5netcdf)