Fix the issue - "Empty source still have embedding network" (ecmwf#1114)

* Replace cf.rank==0 with utils.distributed.is_root

* fix empty source inputs still have embedding layer

* fix lint

* fix source empty or source exclude all

* fix source empty or source exclude all

* fix forecast mode empty source

---------

Co-authored-by: wang85 <wang85@jwlogin22.juwels>
Co-authored-by: wang85 <wang85@jwlogin24.juwels>
Co-authored-by: wang85 <wang85@jwb0149.juwels>
Co-authored-by: wang85 <wang85@jwlogin21.juwels>

Author: 5 people

src/weathergen/datasets/multi_stream_data_sampler.py

@@ -248,7 +248,9 @@ def advance(self):
     ###################################################
     def get_sources_size(self):
         return [
-            ds[0].get_source_num_channels()
+            0
+            if ds[0].get_source_num_channels() == 0
+            else ds[0].get_source_num_channels()
             + ds[0].get_geoinfo_size()
             + ds[0].get_coords_size()
             + self.tokenizer.get_size_time_embedding()

src/weathergen/model/engines.py

@@ -47,7 +47,7 @@ def __init__(self, cf: Config, sources_size) -> None:
         for i, si in enumerate(self.cf.streams):
             stream_name = si.get("name", i)
 
-            if "diagnostic" in si and si["diagnostic"]:
+            if si.get("diagnostic", False) or self.sources_size[i] == 0:
                 self.embeds.append(torch.nn.Identity())
                 continue