LETKF for single cycle (#256)

* LETKF directives from hofx - first pass; largely renaming or copying

* Top level options for LETKF

* Templates for LocalEnsembleDA.yaml

* Added task questions for LocalEnsembleDA in geos_atmosphere; rename letkf/suites_questions.yaml to letkf/suite_questions.yaml

* Close, but unfinished, letkf

* Add questions and default choices for local ensemble DA

* GetEnsemble working now.

* Corrected jinja templating

* Tier 1 tests for letkf

* Correct obs_experiment to x0048v1, not v2

* Add r2d2 path to letkf suite

* Correct jinja templating for ensemble background and LocalEnsembleDA driver

* Add window_type 3D to tier-1 overrides

* Change horizontal_resolution for letkf tier-1 tests to match resolution of ensemble members

* Matching naming convention for posterior ensemble members/mean/increments with linked background members at cycle start

* Add questions for horizontal and vertical covariance localization

* A hacky way to include localizations and distribution in ensemble-based DA schemes

* Add guard if user desires both posterior ensemble and posterior mean to be saved

* Correct jinja2 conditional templating

* pynorms

* Swell coding norms

* Try YAML lint a different way

* ignore approach to yamllint

* back to find

* dev

* test for proper failure

* undo failure

* Update src/swell/suites/letkf/flow.cylc

Co-authored-by: Dan Holdaway <[email protected]>

* Implement changes suggested in #256

* Update src/swell/tasks/task_questions.yaml

Co-authored-by: Dan Holdaway <[email protected]>

* Update src/swell/tasks/task_questions.yaml

Co-authored-by: Dan Holdaway <[email protected]>

* Update task_questions.yaml

* Update src/swell/tasks/task_questions.yaml

Co-authored-by: Dan Holdaway <[email protected]>

* Remove obs localizations from run_jedi_executables.py

* Relocate obs localizations within run_jedi_local_ensemble_da_executable.py

* Fixed py/swell norms - copying from terminal to GitHub adds whitespace.

---------

Co-authored-by: danholdaway <[email protected]>
Co-authored-by: Dan Holdaway <[email protected]>

.yamllint.yml

@@ -2,6 +2,9 @@
 
 extends: default
 
+ignore: |
+  src/swell/configuration/jedi/oops/LocalEnsembleDA.yaml
+
 rules:
   braces:
     level: warning
@@ -27,3 +30,4 @@ rules:
     level: warning
     allow-non-breakable-inline-mappings: true
   truthy: disable
+

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

@@ -6,3 +6,4 @@ executables:
   variational3D: fv3jedi_var.x
   variational4D: fv3jedi_var.x
   variational4DEnsVar: fv3jedi_var.x
+  localensembleda: fv3jedi_letkf.x

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

@@ -0,0 +1,13 @@
+date: '{{window_begin_iso}}'
+members from template:
+  template:
+    datetime: '{{local_background_time_iso}}'
+    filetype: cube sphere history
+    provider: geos
+    datapath: '.'
+    filename: 'geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4'
+    state variables: [ua,va,t,ps,delp,q,qi,ql,qr,qs,o3ppmv,phis,frocean,frlake,frseaice,
+                      sheleg,ts,soilt,soilm,u10m,v10m]
+  pattern: '%mem%'
+  nmembers: {{ensemble_num_members}}
+  zero padding: 3

src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_driver.yaml

@@ -0,0 +1,4 @@
+save posterior mean: {{local_ensemble_save_posterior_mean}}
+save posterior ensemble: {{local_ensemble_save_posterior_ensemble}}
+save posterior mean increment: {{local_ensemble_save_posterior_mean_increment}}
+save posterior ensemble increments: {{local_ensemble_save_posterior_ensemble_increments}}

src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_mean_increment_output.yaml

@@ -0,0 +1,4 @@
+filetype: auxgrid
+gridtype: latlon
+filename: '{{cycle_dir}}/geos.mean-inc.'
+

src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_mean_output.yaml

@@ -0,0 +1,3 @@
+filetype: auxgrid
+gridtype: latlon
+filename: '{{cycle_dir}}/geos.analysis.mean.'

src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_members_increment_output.yaml

@@ -0,0 +1,4 @@
+filetype: auxgrid
+gridtype: latlon
+filename: '{{cycle_dir}}/geos.mem%{member}%-inc'
+

src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_members_output.yaml

@@ -0,0 +1,3 @@
+filetype: auxgrid
+gridtype: latlon
+filename: '{{cycle_dir}}/geos.analysis.mem%{member}%.'

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

@@ -0,0 +1,5 @@
+solver: {{local_ensemble_solver}}
+inflation:
+  rtps: {{local_ensemble_inflation_rtps}}
+  rtpp: {{local_ensemble_inflation_rtpp}}
+  mult: {{local_ensemble_inflation_mult}}

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

@@ -53,6 +53,19 @@ geovals_provider:
 gradient_norm_reduction:
   default_value: 10e-5
 
+horizontal_localization_lengthscale:
+  default_value: 10000e3
+
+horizontal_localization_max_nobs:
+  default_value: 1000
+
+horizontal_localization_method:
+  default_value: Horizontal Gaspari-Cohn
+  options:
+  - Horizontal Gaspari-Cohn
+  - Horizontal SOAR
+  - Horizontal Box car
+
 horizontal_resolution:
   default_value: '361'
   options:
@@ -66,6 +79,21 @@ jedi_forecast_model:
   options:
   - pseudo-model
 
+local_ensemble_inflation_mult:
+  default_value: 1.1
+
+local_ensemble_inflation_rtpp:
+  default_value: 0.6
+
+local_ensemble_inflation_rtps:
+  default_value: 0.5
+
+local_ensemble_solver:
+  default_value: LETKF
+  options:
+  - LETKF
+  - GETKF
+
 minimizer:
   default_value: DRIPCG
   options:
@@ -157,7 +185,7 @@ observations:
   - ompsnm_npp
 
 path_to_ensemble:
-  default_value: /discover/nobackup/drholdaw/SwellTestData/letk/ensemble/Y%Y/M%m/D%d/H%H
+  default_value: /discover/nobackup/drholdaw/SwellTestData/letk/ensemble/Y%Y/M%m/D%d/H%H/geos*%Y%m%d_%H%M%Sz.nc4
 
 path_to_geos_adas_background:
   default_value: /discover/nobackup/drholdaw/SwellTestData/geosadas/bkg/*bkg_clcv_rst*
@@ -174,6 +202,29 @@ produce_geovals:
   - true
   - false
 
+vertical_localization_function:
+  default_value: Gaspari Cohn
+  options:
+  - Gaspari Cohn
+
+vertical_localization_ioda_vertical_coord:
+  default_value: pressure
+  options:
+  - pressure
+
+vertical_localization_ioda_vertical_coord_group:
+  default_value: MetaData
+  options:
+  - MetaData
+
+vertical_localization_lengthscale:
+  default_value: 4
+
+vertical_localization_method:
+  default_value: Vertical localization
+  options:
+  - Vertical localization
+
 vertical_resolution:
   default_value: '72'
   options:

src/swell/configuration/jedi/oops/LocalEnsembleDA.yaml

@@ -0,0 +1,36 @@
+geometry:
+  TASKFILLgeometry
+
+window begin: '{{window_begin_iso}}'
+window length: '{{window_length}}'
+
+background:
+  TASKFILLbackground_ensemble
+
+observations:
+  observers:
+    SPECIALobservations
+
+local ensemble DA:
+  TASKFILLensemble_solver
+
+driver:
+  TASKFILLensemble_driver
+
+{% if local_ensemble_save_posterior_mean or local_ensemble_save_posterior_ensemble %}
+output:
+  {% if local_ensemble_save_posterior_mean and not local_ensemble_save_posterior_ensemble %}
+  TASKFILLensemble_mean_output
+  {% endif %}
+  {% if not local_ensemble_save_posterior_mean and local_ensemble_save_posterior_ensemble %}
+  TASKFILLensemble_members_output
+  {% endif %}
+{% endif %}
+{% if local_ensemble_save_posterior_mean_increment %}
+output increment:
+  TASKFILLensemble_mean_increment_output
+{% endif %}
+{% if local_ensemble_save_posterior_ensemble_increments %}
+output ensemble increments:
+  TASKFILLensemble_members_increment_output
+{% endif %}

src/swell/deployment/create_experiment.py

@@ -450,8 +450,9 @@ def prepare_cylc_suite_jinja2(logger, swell_suite_path, exp_suite_path, experime
         'EvaObservations',
         'GenerateBClimatology',
         'RunJediHofxExecutable',
-        'RunJediVariationalExecutable',
+        'RunJediLocalEnsembleDaExecutable',
         'RunJediUfoTestsExecutable',
+        'RunJediVariationalExecutable',
         'RunGeosExecutable',
         ]
 

src/swell/suites/letkf/flow.cylc → src/swell/suites/localensembleda/flow.cylc

@@ -6,7 +6,7 @@
 
 # --------------------------------------------------------------------------------------------------
 
-# Cylc suite for executing JEDI-based h(x)
+# Cylc suite for executing JEDI-based LocalEnsembleDA Algorithm
 
 # --------------------------------------------------------------------------------------------------
 
@@ -57,17 +57,17 @@
             StageJediCycle-{{model_component}}
 
             # Run Jedi hofx executable
-            BuildJediByLinking[^]? | BuildJedi[^]  => RunJediLetkfExecutable-{{model_component}}
-            StageJedi-{{model_component}}[^] => RunJediLetkfExecutable-{{model_component}}
-            StageJediCycle-{{model_component}} => RunJediLetkfExecutable-{{model_component}}
-            GetEnsemble-{{model_component}} => RunJediLetkfExecutable-{{model_component}}
-            GetObservations-{{model_component}} => RunJediLetkfExecutable-{{model_component}}
+            BuildJediByLinking[^]? | BuildJedi[^]  => RunJediLocalEnsembleDaExecutable-{{model_component}}
+            StageJedi-{{model_component}}[^] => RunJediLocalEnsembleDaExecutable-{{model_component}}
+            StageJediCycle-{{model_component}} => RunJediLocalEnsembleDaExecutable-{{model_component}}
+            GetEnsemble-{{model_component}} => RunJediLocalEnsembleDaExecutable-{{model_component}}
+            GetObservations-{{model_component}} => RunJediLocalEnsembleDaExecutable-{{model_component}}
 
             # EvaObservations
-            RunJediLetkfExecutable-{{model_component}} => EvaObservations-{{model_component}}
+            RunJediLocalEnsembleDaExecutable-{{model_component}} => EvaObservations-{{model_component}}
 
             # Save observations
-            RunJediLetkfExecutable-{{model_component}} => SaveObsDiags-{{model_component}}
+            RunJediLocalEnsembleDaExecutable-{{model_component}} => SaveObsDiags-{{model_component}}
 
             # Clean up large files
             EvaObservations-{{model_component}} & SaveObsDiags-{{model_component}} =>
@@ -127,19 +127,19 @@
     [[GetObservations-{{model_component}}]]
         script = "swell task GetObservations $config -d $datetime -m {{model_component}}"
 
-    [[RunJediLetkfExecutable-{{model_component}}]]
-        script = "swell task RunJediLetkfExecutable $config -d $datetime -m {{model_component}}"
+    [[RunJediLocalEnsembleDaExecutable-{{model_component}}]]
+        script = "swell task RunJediLocalEnsembleDaExecutable $config -d $datetime -m {{model_component}}"
         platform = {{platform}}
-        execution time limit = {{scheduling["RunJediLetkfExecutable"]["execution_time_limit"]}}
+        execution time limit = {{scheduling["RunJediLocalEnsembleDaExecutable"]["execution_time_limit"]}}
         [[[directives]]]
-            --account = {{scheduling["RunJediLetkfExecutable"]["account"]}}
-            --qos = {{scheduling["RunJediLetkfExecutable"]["qos"]}}
-            --job-name = RunJediLetkfExecutable
-            --nodes={{scheduling["RunJediLetkfExecutable"]["nodes"]}}
-            --ntasks-per-node={{scheduling["RunJediLetkfExecutable"]["ntasks_per_node"]}}
-            --constraint={{scheduling["RunJediLetkfExecutable"]["constraint"]}}
-            {% if scheduling["RunJediLetkfExecutable"]["partition"] %}
-            --partition={{scheduling["RunJediLetkfExecutable"]["partition"]}}
+            --account = {{scheduling["RunJediLocalEnsembleDaExecutable"]["account"]}}
+            --qos = {{scheduling["RunJediLocalEnsembleDaExecutable"]["qos"]}}
+            --job-name = RunJediLocalEnsembleDaExecutable
+            --nodes={{scheduling["RunJediLocalEnsembleDaExecutable"]["nodes"]}}
+            --ntasks-per-node={{scheduling["RunJediLocalEnsembleDaExecutable"]["ntasks_per_node"]}}
+            --constraint={{scheduling["RunJediLocalEnsembleDaExecutable"]["constraint"]}}
+            {% if scheduling["RunJediLocalEnsembleDaExecutable"]["partition"] %}
+            --partition={{scheduling["RunJediLocalEnsembleDaExecutable"]["partition"]}}
             {% endif %}
 
     [[EvaObservations-{{model_component}}]]

src/swell/suites/letkf/suite_questions.yaml → src/swell/suites/localensembleda/suite_questions.yaml

@@ -24,3 +24,9 @@ cycle_times:
   - all
   prompt: Enter the cycle times for this model.
   type: string-check-list
+
+r2d2_local_path:
+  ask_question: False
+  default_value: defer_to_platform
+  prompt: Enter the path where R2D2 will store experiment output
+  type: string

src/swell/tasks/get_ensemble.py

@@ -8,10 +8,8 @@
 # --------------------------------------------------------------------------------------------------
 
 
-import datetime
 import glob
 import os
-import re
 
 from swell.tasks.base.task_base import taskBase
 
@@ -22,14 +20,24 @@
 class GetEnsemble(taskBase):
 
     def execute(self):
-
-        # Get the path and pattern for the background files
+        """Acquires ensemble member files for a given experiment and cycle
+
+           Parameters
+           ----------
+             All inputs are extracted from the JEDI experiment file configuration.
+             See the taskBase constructor for more information.
+        """
+        # Get the path and pattern for the ensemble members
         # -------------------------------------------------
         ensemble_path = self.config.path_to_ensemble()
 
+        # Replate ensemble_path with true path
+        # ---------------------------------------------
+        ensemble_location = self.cycle_time_dto().strftime(ensemble_path)
+
         # Fetch list of ensemble members
         # --------------------------------
-        ensemble_members = glob.glob(ensemble_path)
+        ensemble_members = glob.glob(ensemble_location)
 
         # Assert at least one ensemble member was found
         # -----------------------------------------------
@@ -44,17 +52,11 @@ def execute(self):
             # Get filename from full path
             member_file = os.path.basename(ensemble_member)
 
-            # Extract the datetime part from the string
-            datetime_part = re.search(r"\d{8}_\d{4}\w", member_file).group()
-
-            # Get datetime for the file from the filename
-            member_file_datetime = datetime.datetime.strptime(datetime_part, '%Y%m%d_%H%Mz')
-
-            # Create target filename using the datetime format
-            member_file_target = member_file_datetime.strftime('geos.mem001.%Y%m%d_%H%M%Sz.nc4')
+            # Extract the member name, excluding path
+            member_part = member_file.split('/')[-1]
 
             # Target path and filename
-            ensemble_path_file_target = os.path.join(self.cycle_dir(), member_file_target)
+            ensemble_path_file_target = os.path.join(self.cycle_dir(), member_part)
 
             # Remove target file if it exists (might be a link)
             if os.path.exists(ensemble_path_file_target):