credit.datasets.gen_2.era5
==========================

.. py:module:: credit.datasets.gen_2.era5

.. autoapi-nested-parse::

   era5.py
   -------
   ARCOERA5Dataset: PyTorch Dataset for streaming ERA5 data from the Google Cloud
   ARCO ERA5 public Zarr store.

   Sample structure returned by __getitem__::

       {
           "input": {
               "{source_name}/prognostic/3d/temperature":              tensor,  # (n_levels, 1, lat, lon)
               "{source_name}/prognostic/2d/surface_pressure":         tensor,  # (1,        1, lat, lon)
               "{source_name}/dynamic_forcing/2d/toa_incident_solar_radiation": tensor,
               "{source_name}/static/2d/land_sea_mask":                tensor,
               ...
           },
           "target": {                                  # only when return_target=True
               "{source_name}/prognostic/3d/temperature":      tensor,
               "{source_name}/prognostic/2d/surface_pressure": tensor,
               ...
           },
           "metadata": {
               "input_datetime":  int,                  # nanoseconds since epoch
               "target_datetime": int,                  # only when return_target=True
           },
       }

   Output key format (flat, slash-delimited)::

       "{source_name}/{field_type}/{dim}/{varname}"

       field_type: "prognostic" | "dynamic_forcing" | "static" | "diagnostic"
       dim       : "2d"  (surface / single-level)
                   "3d"  (multi-level; level_coord = "level" or "hybrid")
       varname   : variable name as in the ARCO ERA5 Zarr store



Attributes
----------

.. autoapisummary::

   credit.datasets.gen_2.era5.logger


Classes
-------

.. autoapisummary::

   credit.datasets.gen_2.era5.ARCOERA5Dataset
   credit.datasets.gen_2.era5.WeatherBench2ERA5Dataset


Module Contents
---------------

.. py:data:: logger

.. py:class:: ARCOERA5Dataset(data_config: dict[str, Any], return_target: bool = False)

   Bases: :py:obj:`credit.datasets.gen_2.base_dataset.BaseDataset`


   PyTorch Dataset for Google Cloud ARCO ERA5 data with nested input/target structure.

   See the module docstring for a full description of the output format and file naming.

   Example YAML configuration::

       data:
         source:
           Example_ARCOERA5:  # User-provided name (arbitrary key)
             dataset_type: "arco_era5"
             level_coord: "hybrid"
             levels: [10, 30, 40, 50, 60, 70, 80, 90, 95, 100, 105, 110, 120, 130, 136, 137]
             variables:
               prognostic:
                 vars_3D: ["temperature", "u_component_of_wind", "v_component_of_wind", "specific_humidity"]
                 vars_2D: ["surface_pressure"]
               dynamic_forcing:
                 vars_2D: ["toa_incident_solar_radiation"]
               static:
                 vars_2D: ["land_sea_mask"]
               diagnostic:
                 vars_2D: ["total_precipitation"]

         start_datetime: "2017-01-01"
         end_datetime: "2019-12-31"
         timestep: "6h"
         forecast_len: 1

   Assumptions:
       1. A "time" dimension / coordinate is present for non-static fields.
       2. A level coordinate (name given by ``level_coord``) represents the
          vertical axis of 3D variables.
       3. Dimension order: (time, level, latitude, longitude) for 3D;
          (time, latitude, longitude) for 2D; (latitude, longitude) for static.


   .. py:attribute:: dataset_type
      :value: 'arco_era5'



   .. py:attribute:: pressure_lev_era5_path
      :value: 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'



   .. py:attribute:: model_lev_era5_path
      :value: 'gs://gcp-public-data-arco-era5/ar/model-level-1h-0p25deg.zarr-v1'



   .. py:attribute:: model_lev_vars
      :value: ['divergence', 'fraction_of_cloud_cover', 'geopotential', 'ozone_mass_mixing_ratio',...



   .. py:attribute:: level_coord
      :type:  str


   .. py:attribute:: mod_level_store
      :value: None



   .. py:attribute:: pres_level_store
      :value: None



   .. py:attribute:: static_metadata
      :type:  dict[str, Any]


   .. py:attribute:: mode
      :value: 'remote'



.. py:class:: WeatherBench2ERA5Dataset(data_config: dict, return_target: bool = False, resolution: str = '1440x721')

   Bases: :py:obj:`credit.datasets.gen_2.base_dataset.BaseDataset`


   PyTorch Dataset for WeatherBench2 ERA5 data on Google Cloud Storage.

   Provides access to ERA5 reanalysis data prepared for the WeatherBench2
   benchmark at multiple resolutions. All data is read lazily from public
   Google Cloud Storage zarr stores (anonymous access, no credentials required).

   Available resolutions::

       +----------------+-------------------+------------+------------------+
       | ``resolution`` | Grid              | Approx deg | Timestep         |
       +================+===================+============+==================+
       | ``"1440x721"`` | 1440 × 721 global | 0.25°      | 6-hourly, 13 lev |
       +----------------+-------------------+------------+------------------+
       | ``"240x121"``  | 240 × 121 global  | 1.5°       | 6-hourly, 13 lev |
       +----------------+-------------------+------------+------------------+
       | ``"64x32"``    | 64 × 32 global    | ~5.6°      | 6-hourly, 13 lev |
       +----------------+-------------------+------------+------------------+
       | ``"full"``     | 1440 × 721 global | 0.25°      | hourly, 37 lev   |
       +----------------+-------------------+------------+------------------+

   See ``_WB2_ERA5_DEFAULT_LEVELS`` for default pressure levels per resolution.

   Example YAML configuration::

       data:
         source:
           WeatherBench2_ERA5:
             dataset_type: "weatherbench2_era5"
             resolution: "1440x721"   # optional; overridden by the resolution kwarg
             level_coord: "level"
             levels: [50, 100, 200, 500, 850, 1000]  # optional; defaults to all available
             variables:
               prognostic:
                 vars_3D: ["temperature", "u_component_of_wind", "v_component_of_wind",
                            "specific_humidity"]
                 vars_2D: ["surface_pressure", "2m_temperature"]
               dynamic_forcing:
                 vars_2D: ["total_precipitation_6hr"]
               static:
                 vars_2D: ["geopotential_at_surface"]
               diagnostic: null

         start_datetime: "2017-01-01"
         end_datetime:   "2019-12-31"
         timestep: "6h"
         forecast_len: 1

   Output key format::

       "weatherbench2_era5/{field_type}/{dim}/{varname}"

   Assumptions:
       1. Non-static variables have a "time" dimension in the zarr store.
       2. 3D pressure-level variables have a "level" coordinate (hPa).
       3. Dimension order: (time, level, latitude, longitude) for 3D;
          (time, latitude, longitude) for 2D; (latitude, longitude) for static.


   .. py:attribute:: dataset_type
      :type:  str
      :value: 'weatherbench2_era5'



   .. py:attribute:: resolution
      :type:  str


   .. py:attribute:: store_path
      :type:  str


   .. py:attribute:: level_coord
      :type:  str


   .. py:attribute:: levels
      :type:  list[int]


   .. py:attribute:: static_metadata
      :type:  dict


   .. py:attribute:: store
      :value: None



   .. py:attribute:: mode
      :value: 'remote'