Sb asma/614 chemistry fine tuning

config/cams_eac4_config.yml

@@ -0,0 +1,20 @@
+streams_directory:  "./config/streams/streams_cams_eac4/"
+
+start_date: 200301010000
+# end_date: 202112310000
+# start_date_val: 202201010000
+# end_date_val: 202205300000
+
+num_epochs: 45 # 10
+
+# samples_per_epoch: 700  # HERE
+# samples_per_validation: 200  # HERE
+# shuffle: True
+
+loader_num_workers: 4
+
+masking_rate: 0.8
+forecast_offset : 1 # HERE
+
+
+with_mixed_precision: True

config/config_eval.yml

@@ -0,0 +1,161 @@
+verbose: true
+global_plotting_options:
+ image_format : "png" #options: "png", "pdf", "svg", "eps", "jpg" ..
+ dpi_val : 300
+
+summary_plots : true
+print_summary: false
+
+evaluation:
+  metrics  : ["rmse"]
+  regions: ["global", "nhem"]
+run_ids :
+  wtqfk9i5:
+    label: "CAMS EAC4 forecast epoch=25 dim_embed = 512 token_size = 16 num_blocks = 2"
+    epoch: 0
+    rank: 0
+    streams:
+      # ERA5:
+      #   channels: ["2t", "10u", "10v", "z_500", "t_850", "u_850", "v_850", "q_850", ]
+      #   evaluation:
+      #     forecast_step: "all"
+      #     sample: "all"
+      #   plotting:
+      #     sample: [0]
+      #     forecast_step: [ 1, 2, 3, 4, 5, 6, 7, 8]
+      #     plot_maps: true
+      #     plot_histograms: true
+      #     plot_animations: false
+      CAMSEAC4: 
+        channels: [
+          # Surface variables
+          'pm1', 'pm2p5', 'pm10',
+          'tc_co', 'tc_no', 'tc_no2', 'tc_o3', 'tc_so2',
+
+          # Ozone (o3)
+          'go3_1000', 'go3_925', 'go3_850', 'go3_700', 'go3_600', 'go3_500',
+          'go3_400', 'go3_300', 'go3_250', 'go3_200', 'go3_150', 'go3_100', 'go3_50',
+
+          # Sulfur dioxide (so2)
+          'so2_1000', 'so2_925', 'so2_850', 'so2_700', 'so2_600', 'so2_500',
+          'so2_400', 'so2_300', 'so2_250', 'so2_200', 'so2_150', 'so2_100', 'so2_50',
+
+          # Nitrogen monoxide (no)
+          'no_1000', 'no_925', 'no_850', 'no_700', 'no_600', 'no_500',
+          'no_400', 'no_300', 'no_250', 'no_200', 'no_150', 'no_100', 'no_50',
+
+          # Nitrogen dioxide (no2)
+          'no2_1000', 'no2_925', 'no2_850', 'no2_700', 'no2_600', 'no2_500',
+          'no2_400', 'no2_300', 'no2_250', 'no2_200', 'no2_150', 'no2_100', 'no2_50',
+
+          # Carbon monoxide (co)
+          'co_1000', 'co_925', 'co_850', 'co_700', 'co_600', 'co_500',
+          'co_400', 'co_300', 'co_250', 'co_200', 'co_150', 'co_100', 'co_50',
+        ]
+        evaluation:
+          forecast_step: "all"
+          sample: "all"
+        plotting:
+          sample: [0]
+          forecast_step: [ 1, 2, 3, 4, 5, 6, 7, 8]
+          plot_maps: true
+          plot_histograms: true
+          plot_animations: true
+  # e8fzh2t1:
+  #   label: "CAMS forecast finetune dim_embed = 256 token_size = 16 num_blocks = 2"
+  #   epoch: 0
+  #   rank: 0
+  #   streams:
+  #     # ERA5:
+  #     #   channels: ["2t", "10u", "10v", "z_500", "t_850", "u_850", "v_850", "q_850", ]
+  #     #   evaluation:
+  #     #     forecast_step: "all"
+  #     #     sample: "all"
+  #     #   plotting:
+  #     #     sample: [0]
+  #     #     forecast_step: [ 1, 2, 3, 4, 5, 6, 7, 8]
+  #     #     plot_maps: true
+  #     #     plot_histograms: true
+  #     #     plot_animations: false
+  #     CAMSEAC4: 
+  #       channels: [
+  #         # Surface variables
+  #         'pm1', 'pm2p5', 'pm10',
+  #         'tc_co', 'tc_no', 'tc_no2', 'tc_o3', 'tc_so2',
+
+  #         # Ozone (o3)
+  #         'go3_1000', 'go3_500', 'go3_250','go3_50',
+
+  #         # Sulfur dioxide (so2)
+  #         'so2_1000', 'so2_500', 'so2_250','so2_50',
+
+
+  #         # Nitrogen monoxide (no)
+  #         'no_1000', 'no_500', 'no_250','no_50',
+
+  #         # Nitrogen dioxide (no2)
+  #         'no2_1000', 'no2_500', 'no2_250','no2_50',
+
+
+  #         # Carbon monoxide (co)
+  #         'co_1000', 'co_500', 'co_250','co_50',
+
+  #       ]
+  #       evaluation:
+  #         forecast_step: "all"
+  #         sample: "all"
+  #       plotting:
+  #         sample: [0]
+  #         forecast_step: [ 1, 2, 3, 4, 5, 6, 7, 8]
+  #         plot_maps: true
+  #         plot_histograms: true
+  #         plot_animations: true
+  # z6aup4r1:
+  #   label: "CAMS forecast finetune dim_embed = 512 token_size = 32 num_blocks = 2"
+  #   epoch: 0
+  #   rank: 0
+  #   streams:
+  #     # ERA5:
+  #     #   channels: ["2t", "10u", "10v", "z_500", "t_850", "u_850", "v_850", "q_850", ]
+  #     #   evaluation:
+  #     #     forecast_step: "all"
+  #     #     sample: "all"
+  #     #   plotting:
+  #     #     sample: [0]
+  #     #     forecast_step: [ 1, 2, 3, 4, 5, 6, 7, 8]
+  #     #     plot_maps: true
+  #     #     plot_histograms: true
+  #     #     plot_animations: false
+  #     CAMSEAC4: 
+  #       channels: [
+  #         # Surface variables
+  #         'pm1', 'pm2p5', 'pm10',
+  #         'tc_co', 'tc_no', 'tc_no2', 'tc_o3', 'tc_so2',
+
+  #         # Ozone (o3)
+  #         'go3_1000', 'go3_500', 'go3_250','go3_50',
+
+  #         # Sulfur dioxide (so2)
+  #         'so2_1000', 'so2_500', 'so2_250','so2_50',
+
+
+  #         # Nitrogen monoxide (no)
+  #         'no_1000', 'no_500', 'no_250','no_50',
+
+  #         # Nitrogen dioxide (no2)
+  #         'no2_1000', 'no2_500', 'no2_250','no2_50',
+
+
+  #         # Carbon monoxide (co)
+  #         'co_1000', 'co_500', 'co_250','co_50',
+
+  #       ]
+  #       evaluation:
+  #         forecast_step: "all"
+  #         sample: "all"
+  #       plotting:
+  #         sample: [0]
+  #         forecast_step: [ 1, 2, 3, 4, 5, 6, 7, 8]
+  #         plot_maps: true
+  #         plot_histograms: true
+  #         plot_animations: true

config/streams/icon/icon.yml

@@ -10,8 +10,8 @@
 ICON :
   type : icon
   filenames : ['icon-art-NWP_OH_CHEMISTRY-chem_DOM01_ML_daily_repeat_reduced_levels.zarr']
-  source : ['u_00', 'v_00', 'w_80', 'temp_00']
-  target : ['u_00', 'v_00', 'w_80', 'temp_00']
+  source_channels : ['u_00', 'v_00', 'w_80', 'temp_00']
+  target_channels : ['u_00', 'v_00', 'w_80', 'temp_00']
   loss_weight : 1.
   diagnostic : False
   masking_rate : 0.6

config/streams/streams_cams_eac4/cams_eac4.yml

@@ -0,0 +1,112 @@
+# (C) Copyright 2024 WeatherGenerator contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+CAMSEAC4 :
+  type : camseac4
+  filenames : ['cams_eac4_2003_2024.zarr']
+  source : [
+      # Surface variables
+      'pm1', 'pm2p5', 'pm10',
+      'tc_co', 'tc_no', 'tc_no2', 'tc_o3', 'tc_so2',
+
+      # Ozone (o3)
+      'go3_1000', 'go3_925', 'go3_850', 'go3_700', 'go3_600', 'go3_500',
+      'go3_400', 'go3_300', 'go3_250', 'go3_200', 'go3_150', 'go3_100', 'go3_50',
+
+      # Sulfur dioxide (so2)
+      'so2_1000', 'so2_925', 'so2_850', 'so2_700', 'so2_600', 'so2_500',
+      'so2_400', 'so2_300', 'so2_250', 'so2_200', 'so2_150', 'so2_100', 'so2_50',
+
+      # Nitrogen monoxide (no)
+      'no_1000', 'no_925', 'no_850', 'no_700', 'no_600', 'no_500',
+      'no_400', 'no_300', 'no_250', 'no_200', 'no_150', 'no_100', 'no_50',
+
+      # Nitrogen dioxide (no2)
+      'no2_1000', 'no2_925', 'no2_850', 'no2_700', 'no2_600', 'no2_500',
+      'no2_400', 'no2_300', 'no2_250', 'no2_200', 'no2_150', 'no2_100', 'no2_50',
+
+      # Carbon monoxide (co)
+      'co_1000', 'co_925', 'co_850', 'co_700', 'co_600', 'co_500',
+      'co_400', 'co_300', 'co_250', 'co_200', 'co_150', 'co_100', 'co_50',
+  ]
+  target : [
+      # Surface variables
+      'pm1', 'pm2p5', 'pm10',
+      'tc_co', 'tc_no', 'tc_no2', 'tc_o3', 'tc_so2',
+
+      # Ozone (o3)
+      'go3_1000', 'go3_925', 'go3_850', 'go3_700', 'go3_600', 'go3_500',
+      'go3_400', 'go3_300', 'go3_250', 'go3_200', 'go3_150', 'go3_100', 'go3_50',
+
+      # Sulfur dioxide (so2)
+      'so2_1000', 'so2_925', 'so2_850', 'so2_700', 'so2_600', 'so2_500',
+      'so2_400', 'so2_300', 'so2_250', 'so2_200', 'so2_150', 'so2_100', 'so2_50',
+
+      # Nitrogen monoxide (no)
+      'no_1000', 'no_925', 'no_850', 'no_700', 'no_600', 'no_500',
+      'no_400', 'no_300', 'no_250', 'no_200', 'no_150', 'no_100', 'no_50',
+
+      # Nitrogen dioxide (no2)
+      'no2_1000', 'no2_925', 'no2_850', 'no2_700', 'no2_600', 'no2_500',
+      'no2_400', 'no2_300', 'no2_250', 'no2_200', 'no2_150', 'no2_100', 'no2_50',
+
+      # Carbon monoxide (co)
+      'co_1000', 'co_925', 'co_850', 'co_700', 'co_600', 'co_500',
+      'co_400', 'co_300', 'co_250', 'co_200', 'co_150', 'co_100', 'co_50',
+  ]
+  # source_exclude : []
+  # target_exclude : []
+  variables: [
+      # Surface variables
+      'pm1', 'pm2p5', 'pm10',
+      'tc_co', 'tc_no', 'tc_no2', 'tc_o3', 'tc_so2',
+
+      # Ozone (o3)
+      'go3_1000', 'go3_925', 'go3_850', 'go3_700', 'go3_600', 'go3_500',
+      'go3_400', 'go3_300', 'go3_250', 'go3_200', 'go3_150', 'go3_100', 'go3_50',
+
+      # Sulfur dioxide (so2)
+      'so2_1000', 'so2_925', 'so2_850', 'so2_700', 'so2_600', 'so2_500',
+      'so2_400', 'so2_300', 'so2_250', 'so2_200', 'so2_150', 'so2_100', 'so2_50',
+
+      # Nitrogen monoxide (no)
+      'no_1000', 'no_925', 'no_850', 'no_700', 'no_600', 'no_500',
+      'no_400', 'no_300', 'no_250', 'no_200', 'no_150', 'no_100', 'no_50',
+
+      # Nitrogen dioxide (no2)
+      'no2_1000', 'no2_925', 'no2_850', 'no2_700', 'no2_600', 'no2_500',
+      'no2_400', 'no2_300', 'no2_250', 'no2_200', 'no2_150', 'no2_100', 'no2_50',
+
+      # Carbon monoxide (co)
+      'co_1000', 'co_925', 'co_850', 'co_700', 'co_600', 'co_500',
+      'co_400', 'co_300', 'co_250', 'co_200', 'co_150', 'co_100', 'co_50',
+  ]
+  pressure_levels : ["50", "100", "150", "200", "250", "300", "400", "500", "600", "700", "850", "925", "1000"]
+  loss_weight : 1.
+  masking_rate : 0.6
+  masking_rate_none : 0.05
+  token_size : 16
+  tokenize_spacetime : True
+  embed : 
+    net : transformer
+    num_tokens : 1
+    num_heads : 8
+    dim_embed : 512
+    num_blocks : 2
+  embed_target_coords :
+    net : linear
+    dim_embed : 512
+  target_readout :
+    type : 'obs_value'  # token or obs_value
+    num_layers : 2
+    num_heads : 4
+    # sampling_rate : 0.2
+  pred_head :
+    ens_size : 1
+    num_layers : 1

config/streams/streams_cams_eac4/era5.yml

@@ -0,0 +1,37 @@
+# (C) Copyright 2024 WeatherGenerator contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+ERA5 :
+  type : anemoi
+  filenames : ['aifs-ea-an-oper-0001-mars-o96-1979-2022-6h-v6.zarr']
+  source_exclude : ['w_', 'skt', 'tcw', 'cp', 'tp']
+  target_exclude : ['w_', 'slor', 'sdor', 'tcw', 'cp', 'tp']
+  loss_weight : 1.
+  masking_rate : 0.6
+  masking_rate_none : 0.05
+  token_size : 8
+  tokenize_spacetime : True
+  max_num_targets: -1
+  embed : 
+    net : transformer
+    num_tokens : 1
+    num_heads : 8
+    dim_embed : 256
+    num_blocks : 2
+  embed_target_coords :
+    net : linear
+    dim_embed : 256
+  target_readout :
+    type : 'obs_value'  # token or obs_value
+    num_layers : 2
+    num_heads : 4
+    # sampling_rate : 0.2
+  pred_head :
+    ens_size : 1
+    num_layers : 1

packages/common/src/weathergen/common/config.py

@@ -64,6 +64,7 @@ def load_model_config(run_id: str, epoch: int | None, model_path: str | None) ->
     Load a configuration file from a given run_id and epoch.
     If run_id is a full path, loads it from the full path.
     """
+    print(f"model_path = {model_path}")
     if Path(run_id).exists():  # load from the full path if a full path is provided
         fname = Path(run_id)
         _logger.info(f"Loading config from provided full run_id path: {fname}")
@@ -76,6 +77,7 @@ def load_model_config(run_id: str, epoch: int | None, model_path: str | None) ->
             )
         model_path = Path(model_path)
         fname = model_path / run_id / _get_model_config_file_name(run_id, epoch)
+        print(f"fname = {fname}")
         assert fname.exists(), (
             "The fallback path to the model does not exist. Please provide a `model_path`."
         )

packages/evaluate/src/weathergen/evaluate/plotter.py

@@ -15,7 +15,7 @@
 from weathergen.common.config import _load_private_conf
 from weathergen.evaluate.plot_utils import DefaultMarkerSize
 
-work_dir = Path(_load_private_conf(None)["path_shared_working_dir"]) / "assets/cartopy"
+work_dir = Path("./assets/cartopy" ) # Path(_load_private_conf(None)["path_shared_working_dir"]) / "assets/cartopy"
 
 cartopy.config["data_dir"] = str(work_dir)
 cartopy.config["pre_existing_data_dir"] = str(work_dir)

src/weathergen/datasets/data_reader_base.py

@@ -327,7 +327,6 @@ def __len__(self) -> int:
         -------
         length of dataset
         """
-
         return self.length()
 
     def get_source(self, idx: TIndex) -> ReaderData: