From d12d59d10e72d255b3d533c1f334ad037168a056 Mon Sep 17 00:00:00 2001 From: Tennessee Leeuwenburg Date: Thu, 22 May 2025 13:07:29 +1000 Subject: [PATCH 01/10] Work in progress including: - Linting of YAML files - Source code changes needed to support low-resolution tutorial --- .../fourcastnext/Training/configs/config.yaml | 2 +- .../Training/configs/data/module/default.yaml | 4 +- .../Training/configs/data/splits/default.yaml | 2 +- .../configs/data/splits/short_training.yaml | 10 ++ .../Training/configs/model/default.yaml | 9 +- .../Training/configs/trainer/default.yaml | 5 +- .../data/lowres.yaml | 9 ++ .../data/module/default.yaml | 2 +- .../data/splits/default.yaml | 6 +- .../data/splits/short_training_splits.yaml | 11 +++ .../limited_vars_early_stop.yaml | 2 +- .../lowres.yaml | 19 ++++ .../model/default.yaml | 9 +- .../trainer/default.yaml | 2 +- .../trainer/few_epochs.yaml | 23 +++++ .../pipelines/low_res_demo_subset.pipe | 46 ++++++++++ .../fourcastnext/src/fourcastnext/__init__.py | 2 +- .../fourcastnext/configs/Data/low_res.pipe | 46 ++++++++++ packages/data/.yamllint.yml | 2 +- packages/data/src/pyearthtools/data/data.yaml | 8 +- .../Australia/Capital/download.yaml | 1 - .../Australia/States/download.yaml | 1 - .../_geographic/Australia/download.yaml | 3 +- .../static/_geographic/World/download.yaml | 1 - .../data/transforms/RegionLookup.yaml | 2 +- packages/nci_site_archive/mkdocs.yml | 2 +- .../src/pyearthtools/pipeline/pipeline.yaml | 12 +-- packages/tutorial/pyproject.toml | 3 - .../pyearthtools/tutorial/ERA5DataClass.py | 91 +++++++++++++++++++ .../src/pyearthtools/tutorial/__init__.py | 3 + .../utils/src/pyearthtools/utils/iPython.py | 1 + packages/zoo/src/pyearthtools/zoo/model.py | 27 +++--- 32 files changed, 309 insertions(+), 57 deletions(-) create mode 100644 packages/bundled_models/fourcastnext/Training/configs/data/splits/short_training.yaml create mode 100644 packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/lowres.yaml create mode 100644 packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/splits/short_training_splits.yaml create mode 100644 packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/lowres.yaml create mode 100644 packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/trainer/few_epochs.yaml create mode 100644 packages/bundled_models/fourcastnext/Training/pipelines/low_res_demo_subset.pipe create mode 100644 packages/bundled_models/fourcastnext/src/fourcastnext/configs/Data/low_res.pipe diff --git a/packages/bundled_models/fourcastnext/Training/configs/config.yaml b/packages/bundled_models/fourcastnext/Training/configs/config.yaml index dd6c508d..cdeac2ae 100644 --- a/packages/bundled_models/fourcastnext/Training/configs/config.yaml +++ b/packages/bundled_models/fourcastnext/Training/configs/config.yaml @@ -10,7 +10,7 @@ defaults: experiment_name: ${now:%Y-%m-%d}/${now:%H-%M-%S} path: /scratch/kd24/${oc.env:USER}/ML/FourCastNeXt/Training/${experiment_name} -fit: False +fit: false hydra: verbose: false diff --git a/packages/bundled_models/fourcastnext/Training/configs/data/module/default.yaml b/packages/bundled_models/fourcastnext/Training/configs/data/module/default.yaml index 585b7705..de56ca80 100644 --- a/packages/bundled_models/fourcastnext/Training/configs/data/module/default.yaml +++ b/packages/bundled_models/fourcastnext/Training/configs/data/module/default.yaml @@ -1,2 +1,2 @@ -num_workers: 12 -batch_size: 8 +num_workers: 1 +batch_size: 1 diff --git a/packages/bundled_models/fourcastnext/Training/configs/data/splits/default.yaml b/packages/bundled_models/fourcastnext/Training/configs/data/splits/default.yaml index 1167f230..87aefd69 100644 --- a/packages/bundled_models/fourcastnext/Training/configs/data/splits/default.yaml +++ b/packages/bundled_models/fourcastnext/Training/configs/data/splits/default.yaml @@ -1,4 +1,4 @@ -random: True +random: true random_seed: 42 train: - 1980 diff --git a/packages/bundled_models/fourcastnext/Training/configs/data/splits/short_training.yaml b/packages/bundled_models/fourcastnext/Training/configs/data/splits/short_training.yaml new file mode 100644 index 00000000..f286ebcf --- /dev/null +++ b/packages/bundled_models/fourcastnext/Training/configs/data/splits/short_training.yaml @@ -0,0 +1,10 @@ +random: true +random_seed: 42 +train: + - 2000 + - 2015 + - 6 hours +valid: + - 2018 + - 2020 + - 6 hours diff --git a/packages/bundled_models/fourcastnext/Training/configs/model/default.yaml b/packages/bundled_models/fourcastnext/Training/configs/model/default.yaml index 204af0e6..49312355 100644 --- a/packages/bundled_models/fourcastnext/Training/configs/model/default.yaml +++ b/packages/bundled_models/fourcastnext/Training/configs/model/default.yaml @@ -1,10 +1,9 @@ - +--- _target_: fourcastnext.lightning_model.FourCastNextLM _recursive_: true - # Model parameters to init AFNONet model_params: - img_size : ${data.img_size} - in_channels : ${data.in_channels} - out_channels : ${data.out_channels} + img_size: ${data.img_size} + in_channels: ${data.in_channels} + out_channels: ${data.out_channels} diff --git a/packages/bundled_models/fourcastnext/Training/configs/trainer/default.yaml b/packages/bundled_models/fourcastnext/Training/configs/trainer/default.yaml index 379a8c37..d2ca183d 100644 --- a/packages/bundled_models/fourcastnext/Training/configs/trainer/default.yaml +++ b/packages/bundled_models/fourcastnext/Training/configs/trainer/default.yaml @@ -1,5 +1,6 @@ +--- precision: 16-mixed -max_epochs: 200 +max_epochs: 5 checkpointing: - monitor: "train_loss" @@ -18,5 +19,5 @@ checkpointing: mode: "max" dirpath: "{path}/Checkpoints/Epoch" filename: "model-{epoch:02d}" - save_on_train_epoch_end: True + save_on_train_epoch_end: true save_top_k: 50 diff --git a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/lowres.yaml b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/lowres.yaml new file mode 100644 index 00000000..d27caba4 --- /dev/null +++ b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/lowres.yaml @@ -0,0 +1,9 @@ +# Use the ERA5 limited variable pipeline + +pipelines: pipelines/limited_variables.pipe + +img_size: + - 64 + - 28 +in_channels: 4 +out_channels: 4 diff --git a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/module/default.yaml b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/module/default.yaml index 1a197ec1..de56ca80 100644 --- a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/module/default.yaml +++ b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/module/default.yaml @@ -1,2 +1,2 @@ -num_workers: 4 +num_workers: 1 batch_size: 1 diff --git a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/splits/default.yaml b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/splits/default.yaml index 2843fa6b..87aefd69 100644 --- a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/splits/default.yaml +++ b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/splits/default.yaml @@ -1,8 +1,8 @@ -random: True +random: true random_seed: 42 train: - - 2000 - - 2015 + - 1980 + - 2018 - 6 hours valid: - 2018 diff --git a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/splits/short_training_splits.yaml b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/splits/short_training_splits.yaml new file mode 100644 index 00000000..365610d0 --- /dev/null +++ b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/data/splits/short_training_splits.yaml @@ -0,0 +1,11 @@ +--- +random: true +random_seed: 42 +train: + - 2000 + - 2015 + - 6 hours +valid: + - 2018 + - 2020 + - 6 hours diff --git a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/limited_vars_early_stop.yaml b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/limited_vars_early_stop.yaml index 1fda4d18..b8c71ad1 100644 --- a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/limited_vars_early_stop.yaml +++ b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/limited_vars_early_stop.yaml @@ -10,7 +10,7 @@ defaults: experiment_name: ${now:%Y-%m-%d}/${now:%H-%M-%S} path: /scratch/kd24/${oc.env:USER}/ML/FourCastNeXt/Training/${experiment_name} -fit: False +fit: false hydra: verbose: false diff --git a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/lowres.yaml b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/lowres.yaml new file mode 100644 index 00000000..3a16ba65 --- /dev/null +++ b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/lowres.yaml @@ -0,0 +1,19 @@ +--- +defaults: + - _self_ + - trainer: few_epochs + - model: default + - data: lowres + + - data/splits: short_training_splits + - data/module: default + +experiment_name: ${now:%Y-%m-%d}/${now:%H-%M-%S} +path: ./Experiments/FourCastNeXt/Training/${experiment_name} + +fit: false + +hydra: + verbose: false + run: + dir: ${path} diff --git a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/model/default.yaml b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/model/default.yaml index 204af0e6..49312355 100644 --- a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/model/default.yaml +++ b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/model/default.yaml @@ -1,10 +1,9 @@ - +--- _target_: fourcastnext.lightning_model.FourCastNextLM _recursive_: true - # Model parameters to init AFNONet model_params: - img_size : ${data.img_size} - in_channels : ${data.in_channels} - out_channels : ${data.out_channels} + img_size: ${data.img_size} + in_channels: ${data.in_channels} + out_channels: ${data.out_channels} diff --git a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/trainer/default.yaml b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/trainer/default.yaml index 379a8c37..ab70ccb7 100644 --- a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/trainer/default.yaml +++ b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/trainer/default.yaml @@ -18,5 +18,5 @@ checkpointing: mode: "max" dirpath: "{path}/Checkpoints/Epoch" filename: "model-{epoch:02d}" - save_on_train_epoch_end: True + save_on_train_epoch_end: true save_top_k: 50 diff --git a/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/trainer/few_epochs.yaml b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/trainer/few_epochs.yaml new file mode 100644 index 00000000..d2ca183d --- /dev/null +++ b/packages/bundled_models/fourcastnext/Training/limited_variables_early_stopping/trainer/few_epochs.yaml @@ -0,0 +1,23 @@ +--- +precision: 16-mixed +max_epochs: 5 + +checkpointing: +- monitor: "train_loss" + mode: "min" + dirpath: "{path}/Checkpoints/Train" + filename: "model-{epoch:02d}-{step:02d}" + every_n_train_steps: 1000 + save_top_k: 10 +- monitor: "valid_loss" + mode: "min" + dirpath: "{path}/Checkpoints/Valid" + filename: "model-{epoch:02d}-{step:02d}-{valid_loss}" + every_n_train_steps: 5000 + save_top_k: 10 +- monitor: "epoch" + mode: "max" + dirpath: "{path}/Checkpoints/Epoch" + filename: "model-{epoch:02d}" + save_on_train_epoch_end: true + save_top_k: 50 diff --git a/packages/bundled_models/fourcastnext/Training/pipelines/low_res_demo_subset.pipe b/packages/bundled_models/fourcastnext/Training/pipelines/low_res_demo_subset.pipe new file mode 100644 index 00000000..744134be --- /dev/null +++ b/packages/bundled_models/fourcastnext/Training/pipelines/low_res_demo_subset.pipe @@ -0,0 +1,46 @@ +!pyearthtools@pyearthtools.pipeline.controller.Pipeline +__args: !!python/tuple +- !pyearthtools@pyearthtools.data.archive.era5_demo_subset + level_value: null + product: reanalysis + transforms: null + variables: + - msl + - 10u + - 10v + - 2t +- !pyearthtools@pyearthtools.pipeline.operations.transforms.Transforms + apply: !pyearthtools@pyearthtools.data.transforms.coordinates.StandardLongitude + type: -180-180 + transforms: null + undo: null +- !pyearthtools@fourcastnext.CropToRectangle + warn: true +- !pyearthtools@pyearthtools.pipeline.modifications.idx_modification.TemporalRetrieval + concat: true + delta_unit: null + merge_function: null + merge_kwargs: + axis: 1 + samples: !!python/tuple + - !!python/tuple + - -6 + - 1 + - !!python/tuple + - 6 + - 2 + - 6 +- !pyearthtools@pyearthtools.pipeline.operations.xarray.conversion.ToNumpy + warn: true +- !pyearthtools@pyearthtools.pipeline.operations.numpy.reshape.Rearrange + rearrange: c t h w -> t c h w + rearrange_kwargs: null + reverse_rearrange: null + skip: false +exceptions_to_ignore: null +iterator: null +sampler: !pyearthtools@pyearthtools.pipeline.samplers.Default {} + +--CONFIG-- +VERSION: 1.0.1 +import: [] diff --git a/packages/bundled_models/fourcastnext/src/fourcastnext/__init__.py b/packages/bundled_models/fourcastnext/src/fourcastnext/__init__.py index a571c343..614d40b2 100644 --- a/packages/bundled_models/fourcastnext/src/fourcastnext/__init__.py +++ b/packages/bundled_models/fourcastnext/src/fourcastnext/__init__.py @@ -39,7 +39,7 @@ def __init__(self, warn=True): def apply_func(self, dataset: xr.Dataset): subset_dataset = dataset.isel( - latitude=slice(0, -1), + latitude=slice(0, -4), ) return subset_dataset diff --git a/packages/bundled_models/fourcastnext/src/fourcastnext/configs/Data/low_res.pipe b/packages/bundled_models/fourcastnext/src/fourcastnext/configs/Data/low_res.pipe new file mode 100644 index 00000000..744134be --- /dev/null +++ b/packages/bundled_models/fourcastnext/src/fourcastnext/configs/Data/low_res.pipe @@ -0,0 +1,46 @@ +!pyearthtools@pyearthtools.pipeline.controller.Pipeline +__args: !!python/tuple +- !pyearthtools@pyearthtools.data.archive.era5_demo_subset + level_value: null + product: reanalysis + transforms: null + variables: + - msl + - 10u + - 10v + - 2t +- !pyearthtools@pyearthtools.pipeline.operations.transforms.Transforms + apply: !pyearthtools@pyearthtools.data.transforms.coordinates.StandardLongitude + type: -180-180 + transforms: null + undo: null +- !pyearthtools@fourcastnext.CropToRectangle + warn: true +- !pyearthtools@pyearthtools.pipeline.modifications.idx_modification.TemporalRetrieval + concat: true + delta_unit: null + merge_function: null + merge_kwargs: + axis: 1 + samples: !!python/tuple + - !!python/tuple + - -6 + - 1 + - !!python/tuple + - 6 + - 2 + - 6 +- !pyearthtools@pyearthtools.pipeline.operations.xarray.conversion.ToNumpy + warn: true +- !pyearthtools@pyearthtools.pipeline.operations.numpy.reshape.Rearrange + rearrange: c t h w -> t c h w + rearrange_kwargs: null + reverse_rearrange: null + skip: false +exceptions_to_ignore: null +iterator: null +sampler: !pyearthtools@pyearthtools.pipeline.samplers.Default {} + +--CONFIG-- +VERSION: 1.0.1 +import: [] diff --git a/packages/data/.yamllint.yml b/packages/data/.yamllint.yml index 7e4b2ee5..3a3e85eb 100644 --- a/packages/data/.yamllint.yml +++ b/packages/data/.yamllint.yml @@ -1,4 +1,4 @@ - +--- extends: default rules: diff --git a/packages/data/src/pyearthtools/data/data.yaml b/packages/data/src/pyearthtools/data/data.yaml index e0abddfc..039673bd 100644 --- a/packages/data/src/pyearthtools/data/data.yaml +++ b/packages/data/src/pyearthtools/data/data.yaml @@ -6,24 +6,24 @@ data: # Data configuration chunks: 'auto' combine_attrs: 'drop_conflicts' xarray_mf: - parallel: False + parallel: false save: # Save files default kwargs xarray: engine: 'netcdf4' zarr: - compute: True + compute: true search_function: 'filesystem' # Search function to use to find files - experimental: False # Experimental feature toggle + experimental: false # Experimental feature toggle series: warning_threshold: 5 # .series warning threshold - future_warning: True # Show FutureWarning of pyearthtools's development + future_warning: true # Show FutureWarning of pyearthtools's development patterns: # Pattern defaults default_extension: .nc diff --git a/packages/data/src/pyearthtools/data/static/_geographic/Australia/Capital/download.yaml b/packages/data/src/pyearthtools/data/static/_geographic/Australia/Capital/download.yaml index 326f9e57..0163a1d9 100644 --- a/packages/data/src/pyearthtools/data/static/_geographic/Australia/Capital/download.yaml +++ b/packages/data/src/pyearthtools/data/static/_geographic/Australia/Capital/download.yaml @@ -1,4 +1,3 @@ - --- # Download Config For Capital Cities Data url: https://www.abs.gov.au/statistics/standards/australian-statistical-geography-standard-asgs-edition-3/jul2021-jun2026/access-and-downloads/digital-boundary-files/GCCSA_2021_AUST_SHP_GDA2020.zip diff --git a/packages/data/src/pyearthtools/data/static/_geographic/Australia/States/download.yaml b/packages/data/src/pyearthtools/data/static/_geographic/Australia/States/download.yaml index 7fc6ba04..b2471e7a 100644 --- a/packages/data/src/pyearthtools/data/static/_geographic/Australia/States/download.yaml +++ b/packages/data/src/pyearthtools/data/static/_geographic/Australia/States/download.yaml @@ -1,4 +1,3 @@ - --- # Download Config For States Data url: https://www.abs.gov.au/statistics/standards/australian-statistical-geography-standard-asgs-pyearthtoolsion-3/jul2021-jun2026/access-and-downloads/digital-boundary-files/STE_2021_AUST_SHP_GDA2020.zip diff --git a/packages/data/src/pyearthtools/data/static/_geographic/Australia/download.yaml b/packages/data/src/pyearthtools/data/static/_geographic/Australia/download.yaml index 155e7c66..2f1eced3 100644 --- a/packages/data/src/pyearthtools/data/static/_geographic/Australia/download.yaml +++ b/packages/data/src/pyearthtools/data/static/_geographic/Australia/download.yaml @@ -1,6 +1,5 @@ - --- # Download Config For Australia Data url: https://www.abs.gov.au/statistics/standards/australian-statistical-geography-standard-asgs-edition-3/jul2021-jun2026/access-and-downloads/digital-boundary-files/AUS_2021_AUST_SHP_GDA2020.zip save_name: Australia.zip -zip: true +zip: true \ No newline at end of file diff --git a/packages/data/src/pyearthtools/data/static/_geographic/World/download.yaml b/packages/data/src/pyearthtools/data/static/_geographic/World/download.yaml index bb9912f8..21baceec 100644 --- a/packages/data/src/pyearthtools/data/static/_geographic/World/download.yaml +++ b/packages/data/src/pyearthtools/data/static/_geographic/World/download.yaml @@ -1,4 +1,3 @@ - --- # Download Config For World Data url: https://public.opendatasoft.com/api/explore/v2.1/catalog/datasets/world-administrative-boundaries/exports/shp?lang=en&timezone=Australia%2FSydney diff --git a/packages/data/src/pyearthtools/data/transforms/RegionLookup.yaml b/packages/data/src/pyearthtools/data/transforms/RegionLookup.yaml index 7697dcc9..e3f0eb27 100644 --- a/packages/data/src/pyearthtools/data/transforms/RegionLookup.yaml +++ b/packages/data/src/pyearthtools/data/transforms/RegionLookup.yaml @@ -11,7 +11,7 @@ Australia: Oceania: [-50, -5, 90, 180] -Adelaide: [-39,-30, 131, 140] +Adelaide: [-39, -30, 131, 140] Brisbane: [-30, -20, 145, 155] Sydney: [-38, -28, 145, 155] Melbourne: [-42, -32, 139, 151] diff --git a/packages/nci_site_archive/mkdocs.yml b/packages/nci_site_archive/mkdocs.yml index 93088ff2..e45fd1fc 100644 --- a/packages/nci_site_archive/mkdocs.yml +++ b/packages/nci_site_archive/mkdocs.yml @@ -57,7 +57,7 @@ plugins: docstring_style: "google" paths: [src] members_order: source - show_submodules: True + show_submodules: true options: allow_inspection: false show_source: false diff --git a/packages/pipeline/src/pyearthtools/pipeline/pipeline.yaml b/packages/pipeline/src/pyearthtools/pipeline/pipeline.yaml index 61d0bdf1..de4bf830 100644 --- a/packages/pipeline/src/pyearthtools/pipeline/pipeline.yaml +++ b/packages/pipeline/src/pyearthtools/pipeline/pipeline.yaml @@ -1,15 +1,15 @@ pipeline: - run_parallel: False # Attempt to run parallel at all, if False, will use Serial Interface + run_parallel: false # Attempt to run parallel at all, if False, will use Serial Interface parallel: # Configuration for parallel running enabled: # Which interfaces are enabled - Delayed: True - Futures: False + Delayed: true + Futures: false dask: # Configuration for dask - start: True + start: true client: # Dask client config - processes: False + processes: false config: {} # Dask config to be set after spinning up distributed cluster @@ -21,4 +21,4 @@ pipeline: default_ignore: [] # Default exceptions to ignore when iterating repr: - show_graph: True # Show graph in ipython repr + show_graph: true # Show graph in ipython repr diff --git a/packages/tutorial/pyproject.toml b/packages/tutorial/pyproject.toml index b97e8ea3..a72c5bcf 100644 --- a/packages/tutorial/pyproject.toml +++ b/packages/tutorial/pyproject.toml @@ -15,9 +15,6 @@ classifiers = [ "Operating System :: OS Independent", ] dependencies = [ - "pyearthtools.data", - "pyearthtools.pipeline", - "pyearthtools.training", "rich", "ipywidgets", "scores", diff --git a/packages/tutorial/src/pyearthtools/tutorial/ERA5DataClass.py b/packages/tutorial/src/pyearthtools/tutorial/ERA5DataClass.py index 7ec32f44..508f2542 100644 --- a/packages/tutorial/src/pyearthtools/tutorial/ERA5DataClass.py +++ b/packages/tutorial/src/pyearthtools/tutorial/ERA5DataClass.py @@ -45,10 +45,13 @@ # This tells pyearthtools what the actual resolution or time-step of the data is inside the files ERA_RESOLUTION = (1, "hour") +ERADEMO_RESOLUTION = (6, "hour") # This dictionary tells pyearthtools what variable renames to apply during load ERA5_RENAME = {"t2m": "2t", "u10": "10u", "v10": "10v", "siconc": "ci"} +ERA5DEMO_RENAME = {"t2m": "2t", "10u": "10m_u_component_of_wind", "v10": "10v"} + V_TO_PATH = { "10m_u_component_of_wind": "10m_u_component_of_wind", "10m_v_component_of_wind": "10m_v_component_of_wind", @@ -192,3 +195,91 @@ def filesystem( def _import(self): """module to import for to load this step in an Pipeline""" return "pyearthtools.tutorial" + + +@register_archive("era5_demo_subset", sample_kwargs=dict(variable="2t")) +class ERA5LowResDemoIndex(ArchiveIndex): + """ECWMF ReAnalysis v5""" + + @property + def _desc_(self): + return { + "singleline": "ECWMF ReAnalysis v5", + "range": "1970-current", + "Documentation": "https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation", + } + + @decorators.alias_arguments( + level_value=["pressure"], + variables=["variable"], + product=["resolution"], + ) + @decorators.variable_modifications(variable_keyword="variables", remove_variables=False) + @decorators.deprecated_arguments( + level="`level` is deprecated in the ERA5 index. Simply provide the variables, `level` will be autofound." + ) + def __init__( + self, + variables: list[str] | str, + *, + level_value: int | float | list[int | float] | tuple[list | int, ...] | None = None, + transforms: Transform | TransformCollection | None = None, + product=None, + ): + """ + Setup ERA5 Low-Res Indexer + + Args: + variables (list[str] | str): + Data variables to retrieve + resolution (Literal[ERA_RES], optional): + Resolution of data, must be one of 'monthly-averaged','monthly-averaged-by-hour', 'reanalysis'. + Defaults to 'reanalysis'. + level_value: (int, optional): + Level value to select if data contains levels. Defaults to None. + transforms (Transform | TransformCollection, optional): + Base Transforms to apply. + Defaults to TransformCollection(). + """ + + variables = [variables] if isinstance(variables, str) else variables + + self.resolution = ERADEMO_RESOLUTION + + self.variables = variables + base_transform = TransformCollection() + + base_transform += pyearthtools.data.transforms.attributes.Rename(ERA5DEMO_RENAME) + # base_transform += pyearthtools.data.transforms.variables.variable_trim(variables) + + self.level_value = level_value + + if level_value: + base_transform += pyearthtools.data.transforms.coordinates.Select( + {coord: level_value for coord in ["level"]}, ignore_missing=True + ) + + super().__init__( + transforms=base_transform + (transforms or TransformCollection()), + data_interval=ERADEMO_RESOLUTION, + ) + self.record_initialisation() + + def filesystem( + self, + querytime: str | Petdt, + ) -> Path | dict[str, str | Path]: + ERA5_HOME = self.ROOT_DIRECTORIES["era5lowresdemo"] + + """ + This tells pyearthtools how to go from a request for a date/time to a path containing the files + which will match that request. + """ + path = Path(ERA5_HOME) / "era5_lowres.nc" # Everything fits into a single 2 GIG file + + return [path] + + @property + def _import(self): + """module to import for to load this step in an Pipeline""" + return "pyearthtools.tutorial" diff --git a/packages/tutorial/src/pyearthtools/tutorial/__init__.py b/packages/tutorial/src/pyearthtools/tutorial/__init__.py index 8ff1b4c3..284e94f5 100644 --- a/packages/tutorial/src/pyearthtools/tutorial/__init__.py +++ b/packages/tutorial/src/pyearthtools/tutorial/__init__.py @@ -30,9 +30,12 @@ default_base = "/g/data/wb00/NCI-Weatherbench/5.625deg" # taken from NCI noteboook on github lowres_base = os.environ.get("ERA5LOWRES", default_base) +USER_HOME = os.path.expanduser("~") +lowresdemo_base = os.environ.get("ERA5LOWRESDEMO", USER_HOME) ROOT_DIRECTORIES = { "era5lowres": lowres_base, # Update this to the base dir, get var from config + "era5lowresdemo": lowresdemo_base, # Update this to the base dir, get var from config } # Register archive returns a callable which can be used to register an object diff --git a/packages/utils/src/pyearthtools/utils/iPython.py b/packages/utils/src/pyearthtools/utils/iPython.py index d63a325a..33b29f6d 100644 --- a/packages/utils/src/pyearthtools/utils/iPython.py +++ b/packages/utils/src/pyearthtools/utils/iPython.py @@ -20,6 +20,7 @@ # TODO: Determine if this file should be deleted, or else give an example of its # use in a notebook + def display_np_arrays_as_images(): def np_to_png(a): if 2 <= len(a.shape) <= 3: diff --git a/packages/zoo/src/pyearthtools/zoo/model.py b/packages/zoo/src/pyearthtools/zoo/model.py index 25e14758..19f2646d 100644 --- a/packages/zoo/src/pyearthtools/zoo/model.py +++ b/packages/zoo/src/pyearthtools/zoo/model.py @@ -66,7 +66,7 @@ def __enter__(self): def __exit__(self, *args): elapsed = time.time() - self.start log = self.logger or LOG - log.debug("%s: took %.2f seconds.", self.title, elapsed) + # log.debug("%s: took %.2f seconds.", self.title, elapsed) class BaseForecastModel: @@ -290,7 +290,7 @@ def __init__( # pylint: disable=R0913 self._data_cleanup = data_cleanup self._pipeline_name = pipeline_name - self.log.debug("Using pipeline: %r", self._pipeline_name) + # self.log.debug("Using pipeline: %r", self._pipeline_name) import pyearthtools.data # pylint: disable=C0321 @@ -380,10 +380,10 @@ def download_assets(self) -> None: pass asset.parent.mkdir(exist_ok=True, parents=True) - self.log.info("Retrieving %s", link) + # self.log.info("Retrieving %s", link) if Path(link).exists(): - self.log.debug(f"Copying {link} to {asset}.") + # self.log.debug(f"Copying {link} to {asset}.") shutil.copyfile(link, str(asset) + ".download") else: download(link, str(asset) + ".download") @@ -436,7 +436,7 @@ def get_all_config_paths(cls, config_path: os.PathLike | None) -> tuple[Path, .. if len(paths) == 0: raise ValueError("No config paths could be established.") - cls.log.debug(f"Config paths: {paths}") + # cls.log.debug(f"Config paths: {paths}") return tuple(paths) @classmethod @@ -493,7 +493,7 @@ def find_paths(path: Path) -> list[str]: paths: list[str] = list(itertools.chain(*(find_paths(config_path) for config_path in config_paths))) - cls.log.debug(f"Pipeline paths for {sub_path!r}: {paths}") + # cls.log.debug(f"Pipeline paths for {sub_path!r}: {paths}") def valid(x: str) -> bool: return "-" in str(x) if ancillary else "-" not in str(x) @@ -544,7 +544,7 @@ def _get_cache(self, cache: os.PathLike): if self._delete_cache and cache is not None: from pyearthtools.data.indexes.utilities import delete_files - self.log.debug("Deleting all data in cache at %r", cache) + # self.log.debug("Deleting all data in cache at %r", cache) delete_files.delete_path(cache) return pyearthtools.pipeline.modifications.Cache( # pylint: disable=E0110,E1120 @@ -670,7 +670,7 @@ def pipeline(self) -> "pyearthtools.pipeline.Pipeline": # type: ignore # noqa: Get pipeline as configured in the init. """ pipe = self._get_pipeline(self._pipeline_name, cache=self.cache) - self.log.debug("Using Pipeline: %r", pipe) + # self.log.debug("Using Pipeline: %r", pipe) return pipe @functools.cached_property @@ -738,9 +738,10 @@ def data(self, basetime: str) -> list[Any]: try: _ = pipeline[data_time] # type: ignore except Exception as e: # pylint: disable=W0718 - self.log.debug("An error occured when getting data from ancillary %s: %s.", key, e) + raise + # self.log.debug("An error occured when getting data from ancillary %s: %s.", key, e) data.append(d) - self.log.debug("Loaded all data.") + # self.log.debug("Loaded all data.") return data @abstractmethod @@ -776,8 +777,8 @@ def _get_index(self, cache: os.PathLike | None, **kwargs) -> "pyearthtools.train model, index_kwargs = self.load(**kwargs) # Load model and trainer - self.log.debug(f"Loading returned {model.__class__ =}") - self.log.debug(f"Loading returned {index_kwargs =}") + # self.log.debug(f"Loading returned {model.__class__ =}") + # self.log.debug(f"Loading returned {index_kwargs =}") post_transforms = index_kwargs.pop("post_transforms", pyearthtools.data.TransformCollection()) @@ -798,7 +799,7 @@ def _get_index(self, cache: os.PathLike | None, **kwargs) -> "pyearthtools.train ) + post_transforms ) - self.log.debug(f"Initialising MLDataIndex with {full_index_kwargs =}") + # self.log.debug(f"Initialising MLDataIndex with {full_index_kwargs =}") return pyearthtools.training.MLDataIndex( model, From 4653f2f710572997ec9dde9da5c04338196d30fc Mon Sep 17 00:00:00 2001 From: Tennessee Leeuwenburg Date: Thu, 22 May 2025 14:02:02 +1000 Subject: [PATCH 02/10] Training demonstration of a very small FourCastNeXt model --- notebooks/tutorial/FourCastMini Demo.ipynb | 1026 ++++++++++++++++++++ 1 file changed, 1026 insertions(+) create mode 100644 notebooks/tutorial/FourCastMini Demo.ipynb diff --git a/notebooks/tutorial/FourCastMini Demo.ipynb b/notebooks/tutorial/FourCastMini Demo.ipynb new file mode 100644 index 00000000..928c6368 --- /dev/null +++ b/notebooks/tutorial/FourCastMini Demo.ipynb @@ -0,0 +1,1026 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "8f960b6e-4219-4dc0-be5b-70a97825e4e9", + "metadata": {}, + "outputs": [], + "source": [ + "# IMPORTANT! Set this to where you want to store your copy of the data!\n", + "import os\n", + "os.environ['ERA5LOWRESDEMO'] = '/scratch/kd24/tjl548/'\n", + "\n", + "import hydra\n", + "import pathlib\n", + "import xarray as xr\n", + "\n", + "from omegaconf import OmegaConf\n", + "\n", + "import pyearthtools.data.archive\n", + "import pyearthtools.tutorial\n", + "import pyearthtools.training\n", + "import pyearthtools.pipeline\n", + "\n", + "import fourcastnext\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "365f4984-8838-4fda-8964-93c288bd9c7e", + "metadata": {}, + "outputs": [], + "source": [ + "# era5_lowres = xr.open_zarr('gs://weatherbench2/datasets/era5/1959-2022-6h-64x32_equiangular_conservative.zarr')" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "bcf32a03-c908-4000-962c-fe1d789be1dc", + "metadata": {}, + "outputs": [], + "source": [ + "# subset = era5_lowres[['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature', 'mean_sea_level_pressure']]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "50cbe29d-1ef7-4df8-899f-19e27d82a632", + "metadata": {}, + "outputs": [], + "source": [ + "# IMPORTANT! Put this somewhere sensible, then set the environment variable as indicated at the start of the tutorial\n", + "# subset.to_netcdf('/scratch/kd24/tjl548/era5_lowres.nc')" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "61aec70a-ffbb-4a37-95dc-9f6ceb59a361", + "metadata": {}, + "outputs": [], + "source": [ + "era5_loaded = xr.open_dataset('/scratch/kd24/tjl548/era5_lowres.nc')\n", + "# era5_loaded # Uncomment this if you want to " + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "217eb07c-e789-4193-b78e-e9699f45da3c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.Dataset> Size: 132kB\n",
+       "Dimensions:                  (time: 4, longitude: 64, latitude: 32)\n",
+       "Coordinates:\n",
+       "  * latitude                 (latitude) float64 256B -87.19 -81.56 ... 87.19\n",
+       "  * longitude                (longitude) float64 512B 0.0 5.625 ... 348.8 354.4\n",
+       "  * time                     (time) datetime64[ns] 32B 2010-01-01 ... 2010-01...\n",
+       "Data variables:\n",
+       "    10m_u_component_of_wind  (time, longitude, latitude) float32 33kB dask.array<chunksize=(4, 36, 18), meta=np.ndarray>\n",
+       "    10m_v_component_of_wind  (time, longitude, latitude) float32 33kB dask.array<chunksize=(4, 36, 18), meta=np.ndarray>\n",
+       "    2m_temperature           (time, longitude, latitude) float32 33kB dask.array<chunksize=(4, 36, 18), meta=np.ndarray>\n",
+       "    mean_sea_level_pressure  (time, longitude, latitude) float32 33kB dask.array<chunksize=(4, 36, 18), meta=np.ndarray>
" + ], + "text/plain": [ + " Size: 132kB\n", + "Dimensions: (time: 4, longitude: 64, latitude: 32)\n", + "Coordinates:\n", + " * latitude (latitude) float64 256B -87.19 -81.56 ... 87.19\n", + " * longitude (longitude) float64 512B 0.0 5.625 ... 348.8 354.4\n", + " * time (time) datetime64[ns] 32B 2010-01-01 ... 2010-01...\n", + "Data variables:\n", + " 10m_u_component_of_wind (time, longitude, latitude) float32 33kB dask.array\n", + " 10m_v_component_of_wind (time, longitude, latitude) float32 33kB dask.array\n", + " 2m_temperature (time, longitude, latitude) float32 33kB dask.array\n", + " mean_sea_level_pressure (time, longitude, latitude) float32 33kB dask.array" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "accessor = pyearthtools.data.archive.era5_demo_subset('foo')\n", + "accessor['2010-01-01']" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "38ff3d85-8a7e-4145-9c52-6f1516d91800", + "metadata": {}, + "outputs": [], + "source": [ + "config_dir = str((pathlib.Path(fourcastnext.__file__).parent / '../../Training/limited_variables_early_stopping').resolve())\n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "bee177a3-c883-4f6c-b1d2-0f9f446ceeec", + "metadata": {}, + "outputs": [], + "source": [ + "initialised = hydra.initialize_config_dir(version_base=None, config_dir = config_dir)\n", + "cfg = hydra.compose(config_name=\"lowres.yaml\")" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "c05733a4-e8e8-4b21-9b6d-5449dfdcc9dd", + "metadata": {}, + "outputs": [], + "source": [ + "splits = {\n", + " \"train_split\": pyearthtools.pipeline.iterators.DateRange(*cfg.data.splits.train),\n", + " \"valid_split\": pyearthtools.pipeline.iterators.DateRange(*cfg.data.splits.valid),\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "3f0e9c1c-39eb-457b-9001-9f1c9ae45a28", + "metadata": {}, + "outputs": [], + "source": [ + "pipeline_path = str((pathlib.Path(fourcastnext.__file__).parent / '../../Training/pipelines/low_res_demo_subset.pipe').resolve())" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "9eae67cf-fe58-42d8-88b6-56a317ce9e5a", + "metadata": {}, + "outputs": [], + "source": [ + "datamodule = pyearthtools.training.data.lightning.PipelineLightningDataModule(\n", + " pipeline_path,\n", + " **splits,\n", + " **cfg.data.module\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "0c475fbf-f74f-42c7-a0b5-887047677f38", + "metadata": {}, + "outputs": [], + "source": [ + "model = hydra.utils.instantiate(cfg.model)" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "771fe4e1-9550-4b2c-b6cb-b540e57b952a", + "metadata": {}, + "outputs": [], + "source": [ + "trainer = pyearthtools.training.lightning.Train(\n", + " model,\n", + " datamodule,\n", + " path=cfg.path,\n", + " trainer_kwargs={'num_sanity_val_steps': 0},\n", + " **OmegaConf.to_object(cfg.trainer)\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b072c58d-6ca0-4189-9c1a-183a09843557", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "GPU available: True (cuda), used: True\n", + "TPU available: False, using: 0 TPU cores\n", + "HPU available: False, using: 0 HPUs\n", + "/opt/conda/envs/pet/lib/python3.11/site-packages/lightning/pytorch/loops/utilities.py:73: `max_epochs` was not set. Setting it to 1000 epochs. To train without an epoch limit, set `max_epochs=-1`.\n", + "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n", + "\n", + " | Name | Type | Params | Mode \n", + "---------------------------------------------\n", + "0 | model | AFNONet | 60.6 M | train\n", + "1 | loss_obj | L1Loss | 0 | train\n", + "---------------------------------------------\n", + "60.6 M Trainable params\n", + "0 Non-trainable params\n", + "60.6 M Total params\n", + "242.261 Total estimated model params size (MB)\n", + "104 Modules in train mode\n", + "0 Modules in eval mode\n", + "/opt/conda/envs/pet/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:425: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n", + "/opt/conda/envs/pet/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:425: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "d66b77870e53410189fcc87fcf72280c", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Training: | | 0/? [00:00 Date: Thu, 22 May 2025 15:19:32 +1000 Subject: [PATCH 03/10] Update notebook to be smoother in a job-submission environment Restore commented-out log diagnostics following debugging process --- notebooks/tutorial/FourCastMini Demo.ipynb | 187 +++++++++++++----- .../Training/configs/data/module/default.yaml | 4 +- .../data/module/default.yaml | 4 +- packages/zoo/src/pyearthtools/zoo/model.py | 21 +- 4 files changed, 150 insertions(+), 66 deletions(-) diff --git a/notebooks/tutorial/FourCastMini Demo.ipynb b/notebooks/tutorial/FourCastMini Demo.ipynb index 928c6368..67fe794b 100644 --- a/notebooks/tutorial/FourCastMini Demo.ipynb +++ b/notebooks/tutorial/FourCastMini Demo.ipynb @@ -1,15 +1,32 @@ { "cells": [ + { + "cell_type": "markdown", + "id": "7d499237-9410-4f01-b517-93100f47a0bf", + "metadata": {}, + "source": [ + "# Quick Training and Inference Demo\n", + "\n", + "Overview:\n", + "- Downloading training data (takes a few minutes)\n", + "- Training a neural network to predict global weather conditions (takes around 90 minutes per epoch)\n", + "- Inferencing the network on unseen data (takes only a moment)\n", + "\n", + "This demo allows the user to download a 2.8GB file containing around 60 years of global eartht system analysis data. Analysis means the science comnunity's best estimate of historical weather conditions based on available observations. The data set used here was produced by the European Centre for Medium Range Weather Forecasting (ECMWF). This is a standard data set used in the field, however most research is done on a higher-resolution version of the data. That said, there is valuable research done at the lower resolution also. The spatial (latitude and longitude) resolution of this data is 64 pixels by 32 pixels, but the time series is very long. Only a few of the most interesting variables are downloaded in this notebook, for reasons of space.\n", + "\n", + "The data is made available by the ECMWF under license, and the conditions are described here: https://www.ecmwf.int/en/forecasts/accessing-forecasts/licenses-available . Please review this before making use of the data for anything. In this tutorial, the data is downloaded using instructions from the weatherbench2 data guide. Please see https://weatherbench2.readthedocs.io/en/latest/data-guide.html for more information on the data, open access, and accessing other resolutions of the data.\n" + ] + }, { "cell_type": "code", - "execution_count": null, - "id": "8f960b6e-4219-4dc0-be5b-70a97825e4e9", + "execution_count": 1, + "id": "c7e44d63-956f-41fd-95a6-6441ded61a3e", "metadata": {}, "outputs": [], "source": [ "# IMPORTANT! Set this to where you want to store your copy of the data!\n", "import os\n", - "os.environ['ERA5LOWRESDEMO'] = '/scratch/kd24/tjl548/'\n", + "os.environ['ERA5LOWRESDEMO'] = os.environ['PBS_JOBFS']\n", "\n", "import hydra\n", "import pathlib\n", @@ -22,7 +39,9 @@ "import pyearthtools.training\n", "import pyearthtools.pipeline\n", "\n", - "import fourcastnext\n" + "import fourcastnext\n", + "\n", + "# import torch; torch.set_default_device() # Uncomment and set this if you need to configure a non-default device." ] }, { @@ -32,44 +51,29 @@ "metadata": {}, "outputs": [], "source": [ - "# era5_lowres = xr.open_zarr('gs://weatherbench2/datasets/era5/1959-2022-6h-64x32_equiangular_conservative.zarr')" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "bcf32a03-c908-4000-962c-fe1d789be1dc", - "metadata": {}, - "outputs": [], - "source": [ - "# subset = era5_lowres[['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature', 'mean_sea_level_pressure']]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "50cbe29d-1ef7-4df8-899f-19e27d82a632", - "metadata": {}, - "outputs": [], - "source": [ + "# Uncomment the code in this cell in order to download data the first time running through.\n", + "\n", + "# era5_lowres = xr.open_zarr('gs://weatherbench2/datasets/era5/1959-2022-6h-64x32_equiangular_conservative.zarr')\n", + "# subset = era5_lowres[['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature', 'mean_sea_level_pressure']]\n", + "\n", "# IMPORTANT! Put this somewhere sensible, then set the environment variable as indicated at the start of the tutorial\n", - "# subset.to_netcdf('/scratch/kd24/tjl548/era5_lowres.nc')" + "# subset.to_netcdf(os.environ['ERA5LOWRESDEMO'] + '/era5_lowres.nc')" ] }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 3, "id": "61aec70a-ffbb-4a37-95dc-9f6ceb59a361", "metadata": {}, "outputs": [], "source": [ - "era5_loaded = xr.open_dataset('/scratch/kd24/tjl548/era5_lowres.nc')\n", + "era5_loaded = xr.open_dataset(os.environ['ERA5LOWRESDEMO'] + '/era5_lowres.nc')\n", "# era5_loaded # Uncomment this if you want to " ] }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 4, "id": "217eb07c-e789-4193-b78e-e9699f45da3c", "metadata": {}, "outputs": [ @@ -456,20 +460,20 @@ " 10m_u_component_of_wind (time, longitude, latitude) float32 33kB dask.array<chunksize=(4, 36, 18), meta=np.ndarray>\n", " 10m_v_component_of_wind (time, longitude, latitude) float32 33kB dask.array<chunksize=(4, 36, 18), meta=np.ndarray>\n", " 2m_temperature (time, longitude, latitude) float32 33kB dask.array<chunksize=(4, 36, 18), meta=np.ndarray>\n", - " mean_sea_level_pressure (time, longitude, latitude) float32 33kB dask.array<chunksize=(4, 36, 18), meta=np.ndarray>