Update ERA5 download example (#845)

* Update era5 download example

* Update changelog

Author: pzharrington

CHANGELOG.md

@@ -20,6 +20,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Simplified CorrDiff config files, updated default values
 - Refactored CorrDiff losses and samplers to use the patching API
 - Support for non-square images and patches in patch-based diffusion
+- ERA5 download example updated to use current file format convention and
+  restricts global statistics computation to the training set
 - Support for training custom StormCast models and various other improvements for StormCast
 
 ### Deprecated

examples/weather/dataset_download/README.md

@@ -19,7 +19,7 @@ flexibly select different meteorological variables for their training dataset.
 ## How to Use
 
 1. Make sure you have the CDS API key setup following
-    [these instructions](https://cds.climate.copernicus.eu/api-how-to).
+    [these instructions](https://cds.climate.copernicus.eu/how-to-api).
 2. Run the main script, `python start_mirror.py`. This will perform all actions to
     generate HDF5 files needed for training. First it will download and save all
     variables as Zarr arrays. This may take a substantial amount of time. If the process

examples/weather/dataset_download/era5_mirror.py

@@ -196,6 +196,8 @@ def download_and_upload_chunk(
 
         # Download the data
         ds = self.download_chunk(variable, year, month, hours, pressure_level)
+        if "valid_time" in ds.dims:
+            ds = ds.rename({"valid_time": "time"})
 
         # Create the Zarr path
         zarr_path = self.variable_to_zarr_name(variable, pressure_level)

examples/weather/dataset_download/start_mirror.py

@@ -61,18 +61,21 @@ def main(cfg: DictConfig) -> None:
 
     # Save mean and std
     if cfg.compute_mean_std:
+        train_era5_xarray = era5_xarray.sel(
+            time=era5_xarray.time.dt.year.isin(train_years)
+        )
         stats_path = os.path.join(cfg.hdf5_store_path, "stats")
         print(f"Saving global mean and std at {stats_path}")
         if not os.path.exists(stats_path):
             os.makedirs(stats_path)
         era5_mean = np.array(
-            era5_xarray.mean(dim=("time", "latitude", "longitude")).values
+            train_era5_xarray.mean(dim=("time", "latitude", "longitude")).values
         )
         np.save(
             os.path.join(stats_path, "global_means.npy"), era5_mean.reshape(1, -1, 1, 1)
         )
         era5_std = np.array(
-            era5_xarray.std(dim=("time", "latitude", "longitude")).values
+            train_era5_xarray.std(dim=("time", "latitude", "longitude")).values
         )
         np.save(
             os.path.join(stats_path, "global_stds.npy"), era5_std.reshape(1, -1, 1, 1)