Clessig/develop/channel logging 282 (ecmwf#615)

* Fix bug with seed being divided by 0 for worker ID=0

* Fix bug causing crash when secrets aren't in private config

* Implement logging losses per channel

* Fix issue with empty targets

* Rework loss logging

* ruff

* Remove computing max_channels

* Change variables names

* ruffed

* Remove redundant enumerations

* Use stages for logging

* Add type hints

* Apply the review

* ruff

* fix

* Fix type hints

* ruff

* Implement sending tensors of different shapes

* ruff

* Fix merge

* Fix docstring

* rerun workflow

* Review

* Change default colums name

* Fix merge

* - Added ddp_average_nan that is robust to NaN/0 entries when computing mean
- Switched from all_gather to this function in trainer to robustly average
- Some code cleanup

* use all_to_all communication

* Fixing problem with single-worker (non-DDP) training

* Ruffed

* Re-enabled validation loss output in terminal

* Simplified handling of dist initalized

---------

Co-authored-by: Kacper Nowak <kacper.nowak@awi.de>
Co-authored-by: Tim Hunter <tim.hunter@ecmwf.int>
Co-authored-by: clessig <christian.lessig@ecwmf.int>

Author: 4 people

.gitignore

@@ -204,10 +204,8 @@ output/
 logs/
 models/
 results/
+plots/
 models
 results
 playground/
-plots/
 .config/
-
-

config/streams/streams_ocean/fesom.yml

@@ -9,7 +9,7 @@
 
 FESOM :
   type : fesom
-  filenames : ['fesom_ifs_awi']
+  filenames : ['test4.zarr']
   loss_weight : 1.
   source : null
   target : ['sst']

src/weathergen/datasets/data_reader_fesom.py

@@ -143,7 +143,7 @@ def __init__(
 
         self.mean = np.concatenate((np.array([0, 0]), np.array(self.ds.data.attrs["means"])))
         self.stdev = np.sqrt(
-            np.concatenate((np.array([1, 1]), np.array(self.ds.data.attrs["vars"])))
+            np.concatenate((np.array([1, 1]), np.array(self.ds.data.attrs["std"])))
         )
 
         source_channels = stream_info.get("source")

src/weathergen/train/loss.py

@@ -11,6 +11,8 @@
 import numpy as np
 import torch
 
+stat_loss_fcts = ["stats", "kernel_crps"]  # Names of loss functions that need std computed
+
 
 ####################################################################################################
 def gaussian(x, mu=0.0, std_dev=1.0):