LGETKF: add function to tune horizontal localization for each observer (

#437)

* Add interfaces/geos_atmosphere/observations/localization

* - run_jedi_local_ensemble_da_executable.py can use yaml set in observations/localization
- add observations/localization

* removed *log  from clean_patterns: in suite_tests/localensembleda-tier1.yaml

* Separate conv obs from Sat obs to test LGETKF

* .

* Add in more observations to localensembleda-tier1.yaml

* python code test pass

* fix ntasks per node for SLES15

* fix mock test

---------

Co-authored-by: Doruk Ardağ <[email protected]>

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/aircraft.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/airs_aqua.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/amsr2_gcom-w1.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/amsua_aqua.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/amsua_metop-b.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/amsua_metop-c.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/amsua_n15.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/amsua_n18.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/amsua_n19.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/atms_n20.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/atms_npp.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/avhrr3_metop-b.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/avhrr3_n18.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/avhrr3_n19.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/cris-fsr_n20.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/cris-fsr_npp.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/gmi_gpm.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/gps.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/iasi_metop-b.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/iasi_metop-c.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/mhs_metop-b.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/mhs_metop-c.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/mhs_n19.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/mls55_aura.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/obsop_name_map.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/omi_aura.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/ompsnm_npp.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/pibal.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/satwind.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/scatwind.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/sfc.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/sfcship.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/sondes.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/localization/ssmis_f17.yaml

@@ -0,0 +1,4 @@
+obs localizations:
+  localization method: Horizontal Gaspari-Cohn
+  lengthscale: 1000.e3
+  max nobs: 100

src/swell/tasks/run_jedi_local_ensemble_da_executable.py

@@ -11,10 +11,10 @@
 import os
 import yaml
 
+from swell.swell_path import get_swell_path
 from swell.tasks.base.task_base import taskBase
 from swell.utilities.run_jedi_executables import jedi_dictionary_iterator, run_executable
 
-
 # --------------------------------------------------------------------------------------------------
 
 
@@ -153,10 +153,6 @@ def execute(self) -> None:
 
         # Assemble localizations
         # ----------------------
-        horizLoc = {'localization method': self.config.horizontal_localization_method(),
-                    'lengthscale': self.config.horizontal_localization_lengthscale(),
-                    'max nobs': self.config.horizontal_localization_max_nobs()}
-        localizations = [horizLoc]
         # # Vertical localizations have bug(s) - Commented out for now...
         # vertLoc = {'localization method': self.config.vertical_localization_method(),
         #            'apply log transformation':
@@ -171,9 +167,24 @@ def execute(self) -> None:
 
         # Include ensemble localizations and halo types with each observation
         # -------------------------------------------------------------------
-        for observer in jedi_config_dict['observations']['observers']:
-            observer.update({'obs localizations': localizations})
-            observer['obs space'].update({'distribution': {'name': 'Halo', 'halo size': 5000e3}})
+
+        swell_path = get_swell_path()
+        localization_path = os.path.join(swell_path,
+                                         f'configuration/jedi/interfaces/geos_atmosphere'
+                                         f'/observations/localization')
+        if self.config.local_ensemble_use_linear_observer():
+            for index, observation in enumerate(observations):
+                # Get pointer to observer (ref to list)
+                observer = jedi_config_dict['observations']['observers'][index]
+                print('ob=', observation)
+                config_file = os.path.join(localization_path, f'{observation}.yaml')
+                with open(config_file, 'r') as f:
+                    loc_list = yaml.safe_load(f)
+                    horizLoc = loc_list['obs localizations']
+                localization = [horizLoc]
+                observer.update({'obs localizations': localization})
+                observer['obs space'].update(
+                    {'distribution': {'name': 'Halo', 'halo size': 5000.e3}})
 
         # Write the expanded dictionary to YAML file
         # ------------------------------------------

src/swell/test/code_tests/slurm_test.py

@@ -73,8 +73,6 @@ def test_slurm_config(self, platform_mocked: Mock, mock_global_defaults: Mock) -
             for mc in ["all", "geos_atmosphere", "geos_ocean"]:
                 # Hard-coded task-specific defaults
                 self.assertEqual(sd["RunJediVariationalExecutable"]["directives"][mc]["nodes"], 3)
-                self.assertEqual(sd["RunJediVariationalExecutable"]["directives"][mc]
-                                 ["ntasks-per-node"], 36)
                 self.assertEqual(sd["RunJediUfoTestsExecutable"]["directives"][mc]
                                  ["ntasks-per-node"], 1)
                 # Global defaults from experiment dict

src/swell/test/suite_tests/localensembleda-tier1.yaml

@@ -4,28 +4,49 @@ model_components: ['geos_atmosphere']
 models:
   geos_atmosphere:
     horizontal_resolution: 91
-    npx_proc: 2
-    npy_proc: 2
+    npx_proc: 4
+    npy_proc: 4
     cycle_times:
     - T00
-    - T06
-    - T12
-    - T18
-    local_ensemble_solver: GETKF
+#    - T06
+#    - T12
+#    - T18
+    ensemble_num_members:  32
     skip_ensemble_hofx: true
+    local_ensemble_solver: GETKF
     local_ensemble_use_linear_observer: true
-    ensemble_num_members: 2
     local_ensemble_save_posterior_mean: false
     local_ensemble_save_posterior_ensemble_increments: true
     ensmean_only: false
     path_to_ensemble: /discover/nobackup/projects/gmao/advda/SwellTestData/letk/ensemble/91/Y%Y/M%m/D%d/H%H/geos*%Y%m%d_%H%M%Sz.nc4
     observations:
     - sondes
+    - amsua_aqua
     - amsua_metop-b
+#    - amsua_metop-c
+#    - amsua_n15
+#    - amsua_n18
+#    - amsua_n19
+#    - amsr2_gcom-w1
+#    - atms_n20
+#    - atms_npp
+#    - avhrr3_metop-b
+#    - avhrr3_n18
+#    - avhrr3_n19
+#    - scatwind
+#    - sfcship
+#    - sfc
+#    - mhs_metop-b
+#    - mhs_metop-c
+#    - mhs_n19
+#    - mls55_aura
+#    - omi_aura
+#    - ompsnm_npp
+#    - pibal
+#    - ssmis_f17
+
     obs_experiment: x0048v2
     background_experiment: x0048
     window_type: 3D
     clean_patterns:
-    - 'geos.mem*.nc4'
-    - '*log*'
     - '*.txt'

src/swell/utilities/slurm.py

@@ -46,7 +46,7 @@ def prepare_scheduling_dict(
     # Hard-coded SLURM defaults for certain tasks
     # -------------------------------------------
     task_defaults = {
-        "RunJediVariationalExecutable": {"all": {"nodes": 3, "ntasks-per-node": 36}},
+        "RunJediVariationalExecutable": {"all": {"nodes": 3}},
         "RunJediUfoTestsExecutable": {"all": {"ntasks-per-node": 1}},
         "RunJediConvertStateSoca2ciceExecutable": {"all": {"nodes": 1}}
     }