Vertical localization parameters in ETKF (L2)

src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py

@@ -17,7 +17,6 @@ def background_ensemble(template_dict: Mapping) -> Mapping:
     horizontal_resolution = template_dict['horizontal_resolution']
 
     background_ensemble = {
-        'date': template_dict['window_begin_iso'],
         'members from template': {
             'template': {
                 'datetime': template_dict['local_background_time_iso'],

src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_solver.py

@@ -18,14 +18,17 @@ def ensemble_solver(template_dict: Mapping) -> Mapping:
     local_ensemble_inflation_rtps = template_dict['local_ensemble_inflation_rtps']
     local_ensemble_inflation_rtpp = template_dict['local_ensemble_inflation_rtpp']
     local_ensemble_inflation_mult = template_dict['local_ensemble_inflation_mult']
+    frac_retained_variance = template_dict['vertical_localization_frac_retained_variance']
+    vertical_localization_lengthscale = template_dict['vertical_localization_lengthscale']
+    vert_loc_unit = template_dict['vertical_localization_unit']
 
     ensemble_solver = {
         'solver': local_ensemble_solver,
         'use linear observer': local_ensemble_use_linear_observer,
         'vertical localization': {
-            'fraction of retained variance': 0.5,
-            'lengthscale': 1.5,
-            'lengthscale units': 'logp'
+            'fraction of retained variance': frac_retained_variance,
+            'lengthscale': vertical_localization_lengthscale,
+            'lengthscale units': vert_loc_unit
         },
         'inflation': {
             'rtps': local_ensemble_inflation_rtps,

src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml

@@ -351,8 +351,17 @@ vertical_localization_ioda_vertical_coord_group:
   options:
   - MetaData
 
+vertical_localization_unit:
+  default_value: 'logp'
+  options:
+  - 'levels'
+  - 'logp'
+
 vertical_localization_lengthscale:
-  default_value: 4
+  default_value: 2
+
+vertical_localization_frac_retained_variance:
+  default_value: 0.95
 
 vertical_localization_method:
   default_value: Vertical localization

src/swell/suites/localensembleda/suite_config.py

@@ -47,6 +47,9 @@ class SuiteConfig(QuestionContainer, Enum):
             qd.skip_ensemble_hofx(True),
             qd.local_ensemble_solver("Deterministic GETKF"),
             qd.local_ensemble_use_linear_observer(True),
+            qd.vertical_localization_unit('logp'),
+            qd.vertical_localization_lengthscale(2.0),
+            qd.vertical_localization_frac_retained_variance(0.95),
             qd.ensmean_only(False),
             qd.local_ensemble_save_posterior_mean(True),
             qd.local_ensemble_save_posterior_mean_increment(True),
@@ -91,6 +94,9 @@ class SuiteConfig(QuestionContainer, Enum):
             qd.skip_ensemble_hofx(True),
             qd.local_ensemble_solver("Deterministic GETKF"),
             qd.local_ensemble_use_linear_observer(True),
+            qd.vertical_localization_unit('logp'),
+            qd.vertical_localization_lengthscale(2.0),
+            qd.vertical_localization_frac_retained_variance(0.95),
             qd.ensmean_only(False),
             qd.local_ensemble_save_posterior_mean(True),
             qd.local_ensemble_save_posterior_mean_increment(True),

src/swell/tasks/run_jedi_etkf_observer.py

@@ -87,8 +87,12 @@ def execute(self) -> None:
                                     self.config.vertical_localization_method())
         self.jedi_rendering.add_key('vertical_localization_apply_log_transform',
                                     self.config.vertical_localization_apply_log_transform())
+        self.jedi_rendering.add_key('vertical_localization_unit',
+                                    self.config.vertical_localization_unit())
         self.jedi_rendering.add_key('vertical_localization_lengthscale',
                                     self.config.vertical_localization_lengthscale())
+        self.jedi_rendering.add_key('vertical_localization_frac_retained_variance',
+                                    self.config.vertical_localization_frac_retained_variance())
         self.jedi_rendering.add_key('vertical_localization_ioda_vertical_coord',
                                     self.config.vertical_localization_ioda_vertical_coord())
         self.jedi_rendering.add_key('vertical_localization_ioda_vertical_coord_group',

src/swell/tasks/run_jedi_etkf_solver.py

@@ -86,8 +86,12 @@ def execute(self) -> None:
                                     self.config.vertical_localization_method())
         self.jedi_rendering.add_key('vertical_localization_apply_log_transform',
                                     self.config.vertical_localization_apply_log_transform())
+        self.jedi_rendering.add_key('vertical_localization_unit',
+                                    self.config.vertical_localization_unit())
         self.jedi_rendering.add_key('vertical_localization_lengthscale',
                                     self.config.vertical_localization_lengthscale())
+        self.jedi_rendering.add_key('vertical_localization_frac_retained_variance',
+                                    self.config.vertical_localization_frac_retained_variance())
         self.jedi_rendering.add_key('vertical_localization_ioda_vertical_coord',
                                     self.config.vertical_localization_ioda_vertical_coord())
         self.jedi_rendering.add_key('vertical_localization_ioda_vertical_coord_group',

src/swell/tasks/run_jedi_local_ensemble_da_executable.py

@@ -93,8 +93,12 @@ def execute(self) -> None:
                                     self.config.vertical_localization_method())
         self.jedi_rendering.add_key('vertical_localization_apply_log_transform',
                                     self.config.vertical_localization_apply_log_transform())
+        self.jedi_rendering.add_key('vertical_localization_unit',
+                                    self.config.vertical_localization_unit())
         self.jedi_rendering.add_key('vertical_localization_lengthscale',
                                     self.config.vertical_localization_lengthscale())
+        self.jedi_rendering.add_key('vertical_localization_frac_retained_variance',
+                                    self.config.vertical_localization_frac_retained_variance())
         self.jedi_rendering.add_key('vertical_localization_ioda_vertical_coord',
                                     self.config.vertical_localization_ioda_vertical_coord())
         self.jedi_rendering.add_key('vertical_localization_ioda_vertical_coord_group',

src/swell/tasks/task_questions.py

@@ -691,6 +691,8 @@ class TaskQuestions(QuestionContainer, Enum):
             qd.vertical_localization_ioda_vertical_coord(),
             qd.vertical_localization_ioda_vertical_coord_group(),
             qd.vertical_localization_lengthscale(),
+            qd.vertical_localization_unit(),
+            qd.vertical_localization_frac_retained_variance(),
             qd.vertical_localization_method(),
             qd.perhost(),
             qd.change_vbc_to_sbc(),
@@ -746,6 +748,8 @@ class TaskQuestions(QuestionContainer, Enum):
             qd.vertical_localization_ioda_vertical_coord(),
             qd.vertical_localization_ioda_vertical_coord_group(),
             qd.vertical_localization_lengthscale(),
+            qd.vertical_localization_unit(),
+            qd.vertical_localization_frac_retained_variance(),
             qd.vertical_localization_method(),
             qd.perhost(),
             qd.change_vbc_to_sbc(),
@@ -783,6 +787,8 @@ class TaskQuestions(QuestionContainer, Enum):
             qd.vertical_localization_ioda_vertical_coord(),
             qd.vertical_localization_ioda_vertical_coord_group(),
             qd.vertical_localization_lengthscale(),
+            qd.vertical_localization_unit(),
+            qd.vertical_localization_frac_retained_variance(),
             qd.vertical_localization_method(),
             qd.perhost(),
             qd.change_vbc_to_sbc(),

src/swell/utilities/question_defaults.py

@@ -1436,7 +1436,35 @@ class vertical_localization_lengthscale(TaskQuestion):
             "geos_atmosphere"
         ])
         prompt: str = "What is the length scale for vertical covariance localization?"
-        widget_type: WType = WType.INTEGER
+        widget_type: WType = WType.FLOAT
+
+    # --------------------------------------------------------------------------------------------------
+
+    @dataclass
+    class vertical_localization_frac_retained_variance(TaskQuestion):
+        default_value: str = "defer_to_model"
+        question_name: str = "vertical_localization_frac_retained_variance"
+        ask_question: bool = True
+        options: str = "defer_to_model"
+        models: List[str] = mutable_field([
+            "geos_atmosphere"
+        ])
+        prompt: str = "What is the fraction of vertical retained variance for GETKF?"
+        widget_type: WType = WType.FLOAT
+
+    # --------------------------------------------------------------------------------------------------
+
+    @dataclass
+    class vertical_localization_unit (TaskQuestion):
+        default_value: str = "defer_to_model"
+        question_name: str = "vertical_localization_unit"
+        ask_question: bool = True
+        options: str = "defer_to_model"
+        models: List[str] = mutable_field([
+            "geos_atmosphere"
+        ])
+        prompt: str = "What is the vertical localization unit for GETKF?"
+        widget_type: WType = WType.STRING
 
     # --------------------------------------------------------------------------------------------------
 

src/swell/utilities/render_jedi_interface_files.py

@@ -125,6 +125,8 @@ def __init__(
             'vertical_localization_ioda_vertical_coord',
             'vertical_localization_ioda_vertical_coord_group',
             'vertical_localization_lengthscale',
+            'vertical_localization_unit',
+            'vertical_localization_frac_retained_variance',
             'vertical_localization_method',
             'vertical_resolution',
             'window_begin',