🐛[BUG]: e2mip://graphcast potentially broken
nbren12 opened this issue · comments
Version
main
On which installation method(s) does this occur?
No response
Describe the issue
There is an error when trying to load graphcast.
python -c 'import earth2mip.networks; earth2mip.networks.get_model("e2mip://graphcast")'
/usr/local/lib/python3.10/dist-packages/gribapi/__init__.py:23: UserWarning: ecCodes 2.31.0 or higher is recommended. You are running version 2.30.0
warnings.warn(
Model description:
GraphCast model at 0.25deg resolution, with 37 pressure levels. This model is
trained on ERA5 data from 1979 to 2017, and can be causally evaluated on 2018
and later years. This model takes as inputs `total_precipitation_6hr`. This was
described in the paper
`GraphCast: Learning skillful medium-range global weather forecasting`
(https://arxiv.org/abs/2212.12794).
Model license:
The model weights are licensed under the Creative Commons
Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0). You
may obtain a copy of the License at:
https://creativecommons.org/licenses/by-nc-sa/4.0/.
The weights were trained on ERA5 data, see README for attribution statement.
Traceback (most recent call last):
File "<string>", line 1, in <module>
File "/usr/local/lib/python3.10/dist-packages/earth2mip/networks/__init__.py", line 345, in get_model
return _load_package_builtin(package, device, name=url.netloc)
File "/usr/local/lib/python3.10/dist-packages/earth2mip/networks/__init__.py", line 291, in _load_package_builtin
return inference_loader(package, device=device)
File "/usr/local/lib/python3.10/dist-packages/earth2mip/networks/graphcast.py", line 642, in load_time_loop
return _load_time_loop_from_description(
File "/usr/local/lib/python3.10/dist-packages/earth2mip/networks/graphcast.py", line 632, in _load_time_loop_from_description
stepper = load_stepper(checkpoint, dataset_path, stats_dir, device=device)
File "/usr/local/lib/python3.10/dist-packages/earth2mip/networks/graphcast.py", line 530, in load_stepper
) = data_utils.extract_inputs_targets_forcings(
File "/usr/local/lib/python3.10/dist-packages/graphcast/data_utils.py", line 288, in extract_inputs_targets_forcings
dataset = dataset.sel(level=list(pressure_levels))
File "/usr/local/lib/python3.10/dist-packages/xarray/core/dataset.py", line 3047, in sel
query_results = map_index_queries(
File "/usr/local/lib/python3.10/dist-packages/xarray/core/indexing.py", line 193, in map_index_queries
results.append(index.sel(labels, **options))
File "/usr/local/lib/python3.10/dist-packages/xarray/core/indexes.py", line 784, in sel
raise KeyError(f"not all values found in index {coord_name!r}")
KeyError: "not all values found in index 'level'"
The bug is here:
earth2mip/earth2mip/networks/graphcast.py
Line 627 in 2b6d3f8
Environment details
No response