From 79375de8227ad48f4959830832effde0a5fc2e72 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 27 May 2026 13:43:10 -0600 Subject: [PATCH 01/68] step 1. add A --- .../geos_atmosphere/model/background_eda.py | 83 ++ .../model/background_error_hybrid.py | 221 +++++ .../observations/sfcship.yaml.simple | 39 + .../observations/sondes.yaml.simple | 38 + src/swell/configuration/jedi/oops/eda3D.py | 66 ++ .../eda/eva/increment-geos_atmosphere.yaml | 760 ++++++++++++++++++ .../eda/eva/jedi_log-geos_atmosphere.yaml | 89 ++ .../eda/eva/observations-geos_atmosphere.yaml | 564 +++++++++++++ src/swell/suites/eda/flow.cylc | 228 ++++++ src/swell/suites/eda/suite_config.py | 119 +++ src/swell/tasks/run_jedi_eda_executable.py | 195 +++++ src/swell/utilities/yaml_utils.py | 19 + 12 files changed, 2421 insertions(+) create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_hybrid.py create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sfcship.yaml.simple create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sondes.yaml.simple create mode 100644 src/swell/configuration/jedi/oops/eda3D.py create mode 100644 src/swell/suites/eda/eva/increment-geos_atmosphere.yaml create mode 100644 src/swell/suites/eda/eva/jedi_log-geos_atmosphere.yaml create mode 100644 src/swell/suites/eda/eva/observations-geos_atmosphere.yaml create mode 100644 src/swell/suites/eda/flow.cylc create mode 100644 src/swell/suites/eda/suite_config.py create mode 100644 src/swell/tasks/run_jedi_eda_executable.py create mode 100644 src/swell/utilities/yaml_utils.py diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py new file mode 100644 index 000000000..1087f4766 --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py @@ -0,0 +1,83 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ + field_io_names + +# -------------------------------------------------------------------------------------------------- + +state_variables = [ + 'eastward_wind', + 'northward_wind', + 'air_temperature', + 'air_pressure_at_surface', + 'air_pressure_levels', + 'water_vapor_mixing_ratio_wrt_moist_air', + 'cloud_liquid_ice', + 'cloud_liquid_water', + 'rain_water', + 'snow_water', + 'mole_fraction_of_ozone_in_air', + 'geopotential_height_times_gravity_at_surface', + 'initial_mass_fraction_of_large_scale_cloud_condensate', + 'initial_mass_fraction_of_convective_cloud_condensate', + 'convective_cloud_area_fraction', + 'fraction_of_ocean', + 'fraction_of_land', + 'isotropic_variance_of_filtered_topography', + 'surface_velocity_scale', + 'surface_buoyancy_scale', + 'planetary_boundary_layer_height', + 'surface_exchange_coefficient_for_momentum', + 'surface_exchange_coefficient_for_heat', + 'surface_exchange_coefficient_for_moisture', + 'KCBL_before_moist', + 'surface_temp_before_moist', + 'lower_index_where_Kh_greater_than_2', + 'upper_index_where_Kh_greater_than_2', + 'fraction_of_lake', + 'fraction_of_ice', + 'vtype', + 'stype', + 'vfrac', + 'sheleg', + 'skin_temperature_at_surface', + 'soilt', + 'soilm', + 'eastward_wind_at_surface', + 'northward_wind_at_surface', + # 'sea_surface_temperature', + # 'mole_fraction_of_carbon_dioxide_in_air', +] + +# -------------------------------------------------------------------------------------------------- + + +def background_eda(template_dict: Mapping) -> Mapping: + horizontal_resolution = template_dict['horizontal_resolution'] + imem=template_dict['ensemble_imember'] + + background = { + 'datetime': template_dict['local_background_time_iso'], + 'filetype': 'cube sphere history', + 'provider': 'geos', + 'compute edge pressure from surface pressure': True, + 'max allowable geometry difference': 1e-3, + 'datapath': template_dict['cycle_dir'], + 'filenames': [ + f'ebkg/mem{imem:03d}/geos.mem{imem:03d}.%yyyy%mm%dd_%hh%MM%ssz.nc4', + f'fv3-jedi/bkg/geos.crtmsrf.{horizontal_resolution}.nc4' + ], + 'state variables': state_variables, + 'field io names': field_io_names, + } + + return background + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_hybrid.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_hybrid.py new file mode 100644 index 000000000..321470de1 --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_hybrid.py @@ -0,0 +1,221 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- +from ruamel.yaml import YAML +from jinja2 import Environment +##from typing import Mapping +from collections.abc import Mapping +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ + field_io_names + +# -------------------------------------------------------------------------------------------------- + +state_variables_to_inverse = [ + 'eastward_wind', + 'northward_wind', + 'air_temperature', + 'air_pressure_at_surface', +# 'air_pressure_levels', + 'water_vapor_mixing_ratio_wrt_moist_air', + 'cloud_liquid_ice', + 'cloud_liquid_water', + 'rain_water', + 'snow_water', + 'mole_fraction_of_ozone_in_air', + 'fraction_of_ocean', + 'fraction_of_lake', + 'fraction_of_ice', + 'geopotential_height_times_gravity_at_surface', + 'skin_temperature_at_surface', +] + +# -------------------------------------------------------------------------------------------------- + +background_error_template = """ +covariance model: hybrid +components: +- covariance: + covariance model: SABER + covariance type: gsi hybrid covariance + saber central block: + saber block name: gsi hybrid covariance + read: + gsi akbk: ./fv3-jedi/fv3files/akbk{{vertical_resolution}}.nc4 + gsi error covariance file: ./fv3-jedi/gsibec/gsi-coeffs-gmao-global-l{{vertical_resolution}}x{{gsibec_nlons}}y{{gsibec_nlats}}.nc4 + gsi berror namelist file: ./fv3-jedi/gsibec/{{gsibec_configuration}}_l{{vertical_resolution}}x{{gsibec_nlons}}y{{gsibec_nlats}}.nml + processor layout x direction: {{gsibec_npx_proc}} + processor layout y direction: {{gsibec_npy_proc}} + debugging mode: false + saber outer blocks: + - saber block name: interpolation + inner geometry: + function space: StructuredColumns + custom grid matching gsi: + type: latlon + lats: {{gsibec_nlats}} + lons: {{gsibec_nlons}} + custom partitioner matching gsi: + bands: {{gsibec_npy_proc}} + halo: 1 + forward interpolator: + local interpolator type: oops unstructured grid interpolator + inverse interpolator: + local interpolator type: oops unstructured grid interpolator + state variables to inverse: &bvar + {%- for var in state_variables_to_inverse %} + - {{ var }} + {%- endfor %} + linear variable change: + linear variable change name: Control2Analysis + input variables: *bvar + weight: + value: 0.1 + +- covariance: + covariance model: ensemble + members from template: + template: + datetime: {{ local_background_time_iso }} + filetype: cube sphere history + provider: geos + state variables: *bvar + datapath: ./ebkg/mem%mem% + filename: geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4 + max allowable geometry difference: 1.e-4 + pattern: '%mem%' + nmembers: {{ ensemble_num_members }} + zero padding: 3 + localization: + localization method: SABER + saber central block: + saber block name: BUMP_NICAS + active variables: *bvar + read: + general: + universe length-scale: 2500.0e3 + drivers: + multivariate strategy: duplicated + compute nicas: true + model: + level for 2d variables: last + nicas: + resolution: 6 + explicit length-scales: true + horizontal length-scale: + - groups: + - common + profile: + - 3350515.0 + - 3350515.0 + - 3350515.0 + - 3350515.0 + - 3350515.0 + - 3350515.0 + - 3350515.0 + - 3092783.0 + - 3092783.0 + - 3092783.0 + - 3092783.0 + - 2835052.0 + - 2835052.0 + - 2835052.0 + - 2577320.0 + - 2577320.0 + - 2577320.0 + - 2577320.0 + - 2577320.0 + - 2577320.0 + - 2577320.0 + - 2448454.0 + - 2319588.0 + - 2190722.0 + - 2190722.0 + - 2061856.0 + - 1804124.0 + - 1804124.0 + - 1804124.0 + - 1804124.0 + - 1804124.0 + - 1804124.0 + - 1804124.0 + - 1804124.0 + - 1804124.0 + - 1804124.0 + - 1804124.0 + - 1804124.0 + - 1675258.0 + - 1417526.0 + - 1288660.0 + - 1237113.0 + - 1185567.0 + - 1082474.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 1030928.0 + - 2030928.0 + - 2030928.0 + - 2030928.0 + vertical length-scale: + - groups: + - common + value: 1.23 + linear variable change: + linear variable change name: Control2Analysis + input variables: *bvar + weight: + value: 0.9 +""" + + +def background_error_hybrid(template_dict: Mapping) -> Mapping: + render_context = { + **template_dict, + 'state_variables_to_inverse': state_variables_to_inverse, + 'field_io_names': field_io_names + } + + # initialize Jinja and render the template + env = Environment() + template = env.from_string(background_error_template) + + # pass the combined dictionary here + rendered_yaml_string = template.render(**render_context) + # Use ruamel.yaml to load the string instead of standard pyyaml + ruamel_yaml = YAML() + background_error = ruamel_yaml.load(rendered_yaml_string) + cov_template = background_error['components'][1]['covariance']['members from template']['template'] + cov_template['field io names'] = field_io_names + a = background_error['components'][0]['covariance']['linear variable change'] + a['output variables'] = template_dict['analysis_variables'] + a = background_error['components'][1]['covariance']['linear variable change'] + a['output variables'] = template_dict['analysis_variables'] + + return background_error + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sfcship.yaml.simple b/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sfcship.yaml.simple new file mode 100644 index 000000000..3b2854599 --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sfcship.yaml.simple @@ -0,0 +1,39 @@ +obs space: + name: Surface Marine Stations + obsdatain: + engine: + type: H5File + obsfile: '{{cycle_dir}}/sfcship.{{window_begin}}.nc4' + missing file action: warn + obsdataout: + engine: + type: H5File + obsfile: '{{cycle_dir}}/{{experiment_id}}.sfcship.{{window_begin}}.nc4' + io pool: + max pool size: 6 + simulated variables: [stationPressure] + +obs operator: + name: Composite + components: + + # Surface pressure is handled using the correction scheme + - name: SfcCorrected + variables: + - name: stationPressure + correction scheme to use: GSL + geovar_geomz: geopotential_height + geovar_sfc_geomz: height_above_mean_sea_level_at_surface + station_altitude: height + +linear obs operator: + name: Composite + components: + + # Specific humidity and station pressure is using the lowest model level (identity) + - name: Identity + variables: + - name: stationPressure + +obs error: + covariance model: diagonal diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sondes.yaml.simple b/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sondes.yaml.simple new file mode 100644 index 000000000..49d95e91c --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sondes.yaml.simple @@ -0,0 +1,38 @@ +obs space: + name: Radiosondes + obsdatain: + engine: + type: H5File + obsfile: '{{cycle_dir}}/sondes.{{window_begin}}.nc4' + missing file action: error + obsgrouping: + group variables: ["stationIdentification", "releaseTime"] + sort variable: "pressure" + sort order: "descending" + obsdataout: + engine: + type: H5File + obsfile: '{{cycle_dir}}/{{experiment_id}}.sondes.{{window_begin}}.nc4' + io pool: + max pool size: 6 + simulated variables: [airTemperature] + + +obs operator: + name: Composite + components: + + - name: VertInterp + variables: + - name: airTemperature + +linear obs operator: + name: Composite + components: + + - name: VertInterp + variables: + - name: airTemperature + +obs error: + covariance model: diagonal diff --git a/src/swell/configuration/jedi/oops/eda3D.py b/src/swell/configuration/jedi/oops/eda3D.py new file mode 100644 index 000000000..4f53b7b56 --- /dev/null +++ b/src/swell/configuration/jedi/oops/eda3D.py @@ -0,0 +1,66 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from swell.utilities.oops_config import OopsConfig + +# -------------------------------------------------------------------------------------------------- + + +class eda3D(OopsConfig): + + def render_oops(self): + oops = { + 'cost function': { + 'cost type': '3D-Var', + 'jb evaluation': False, + 'time window': { + 'begin': self.template_dict['window_begin_iso'], + 'end': self.template_dict['window_end_iso'], + 'bound to include': 'begin' + }, + 'geometry': self.interface_model('geometry'), + 'analysis variables': self.template_dict['analysis_variables'], + 'background': self.interface_model('background_eda'), + 'background error': self.interface_model('background_error_hybrid'), + 'observations': { + 'obs perturbations': False, + 'get values': self.interface_model('getvalues'), + 'observers': self.special_observations(), + } + }, + 'variational': { + 'minimizer': { + 'algorithm': self.template_dict['minimizer'] + }, + 'iterations': [{ + 'geometry': self.interface_model('geometry_inner'), + 'gradient norm reduction': float(self.template_dict['gradient_norm_reduction']), + 'ninner': self.template_dict['number_of_iterations'], + 'diagnostics': { + 'departures': 'ombg' + }, + 'online diagnostics': self.interface_model('varincrement1') + }], + }, + 'final': { + 'diagnostics': { + 'departures': 'oman' + }, + 'prints': { + 'frequency': 'PT3H' + } + }, + 'output': self.interface_model('analysis') + } + + # TODO: Implement this more cleanly in the OOPS schema + if self.jedi_interface == 'geos_cf': + oops['final']['increment'] = {'geometry': self.interface_model('geometry'), + 'output': self.interface_model('increment_cs')} + + return oops diff --git a/src/swell/suites/eda/eva/increment-geos_atmosphere.yaml b/src/swell/suites/eda/eva/increment-geos_atmosphere.yaml new file mode 100644 index 000000000..126435876 --- /dev/null +++ b/src/swell/suites/eda/eva/increment-geos_atmosphere.yaml @@ -0,0 +1,760 @@ +datasets: + +- group: increment + type: LatLon + filename: {{increment_file_path}} + name: experiment_increment + variables: [ps, ts, ua, va, t, q, lat, lon] + +graphics: + + plotting_backend: Emcpy + figure_list: + + #map plot for surface pressure increment + - batch figure: + variables: [ps] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Surface Pressure Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ps + label: PS increment + colorbar: true + cmap: 'bwr' + vmin: -100 + vmax: 100 + - batch figure: + variables: [ts] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Tskin Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ts + label: TS increment + colorbar: true + cmap: 'bwr' + vmin: -2 + vmax: 2 + #map plot for temperature increment (lowest level) + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1000.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[71,...]' + label: T increment (1000 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for temperature increment + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_850.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[62,...]' + label: T increment (850 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for temperature increment + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_500.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[49,...]' + label: T increment (500 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for temperature increment + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_200.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[42,...]' + label: T increment (200 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for temperature increment + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_10.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[24,...]' + label: T increment (10 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for temperature increment + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[14,...]' + label: T increment (1 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment (lowest level) + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1000.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[71,...]' + label: U increment (1000 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_850.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[62,...]' + label: U increment (850 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_500.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[49,...]' + label: U increment (500 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_200.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[42,...]' + label: U increment (200 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_10.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[24,...]' + label: U increment (10 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[14,...]' + label: U increment (1 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment (lowest level) + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1000.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[71,...]' + label: V increment (1000 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_850.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[62,...]' + label: V increment (850 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_500.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[49,...]' + label: V increment (500 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_200.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[42,...]' + label: V increment (200 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_10.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[24,...]' + label: V increment (10 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[14,...]' + label: V increment (1 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Specific Humidity increment (lowest level) + - batch figure: + variables: [q] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1000.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Specific Humidity Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::q + slices: '[71,...]' + label: Q increment (1000 hPa) + colorbar: true + cmap: 'bwr' + vmin: -0.001 + vmax: 0.001 + #map plot for Specific Humidity increment + - batch figure: + variables: [q] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_850.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Specific Humidity Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::q + slices: '[62,...]' + label: Q increment (850 hPa) + colorbar: true + cmap: 'bwr' + vmin: -0.001 + vmax: 0.001 + #map plot for Specific Humidity increment + - batch figure: + variables: [q] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_500.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Specific Humidity Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::q + slices: '[49,...]' + label: Q increment (500 hPa) + colorbar: true + cmap: 'bwr' + vmin: -0.001 + vmax: 0.001 + #map plot for Specific Humidity increment + - batch figure: + variables: [q] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_200.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Specific Humidity Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::q + slices: '[42,...]' + label: Q increment (200 hPa) + colorbar: true + cmap: 'bwr' + vmin: -0.001 + vmax: 0.001 + #map plot for Specific Humidity increment + - batch figure: + variables: [q] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_10.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Specific Humidity Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::q + slices: '[24,...]' + label: Q increment (10 hPa) + colorbar: true + cmap: 'bwr' + vmin: -0.001 + vmax: 0.001 diff --git a/src/swell/suites/eda/eva/jedi_log-geos_atmosphere.yaml b/src/swell/suites/eda/eva/jedi_log-geos_atmosphere.yaml new file mode 100644 index 000000000..2d848d549 --- /dev/null +++ b/src/swell/suites/eda/eva/jedi_log-geos_atmosphere.yaml @@ -0,0 +1,89 @@ +datasets: + +- type: JediLog + collection_name: JediLogTest + jedi_log_to_parse: '{{cycle_dir}}/jedi_variational_log.log' + data_to_parse: + convergence: true + +transforms: +- transform: arithmetic + new name: JediLogTest::convergence::${variable}_log + equals: log(JediLogTest::convergence::${variable}) + for: + variable: [residual_norm, norm_reduction] + +graphics: + + plotting_backend: Emcpy + figure_list: + + - figure: + layout: [3,1] + figure size: [12,10] + title: 'Residual Norm and Norm Reduction Plots' + output name: '{{cycle_dir}}/eva/jedi_log/convergence/residual_norm_reduction.png' + plots: + - add_xlabel: 'Total inner iteration number' + add_ylabel: 'Residual norm' + layers: + - type: LinePlot + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::residual_norm + color: 'black' + + - add_xlabel: 'Total inner iteration number' + add_ylabel: 'Log(norm reduction)' + layers: + - type: LinePlot + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::norm_reduction + color: 'black' + + - add_xlabel: 'Total inner iteration number' + add_ylabel: 'Log(reduction)' + add_legend: + layers: + - type: LinePlot + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::residual_norm_log + color: 'red' + label: 'Log(residual norm)' + - type: LinePlot + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::norm_reduction_log + color: 'blue' + label: 'Log norm reduction' + + - figure: + title: 'Cost Function Plot' + output name: '{{cycle_dir}}/eva/jedi_log/cost_function/cost_function.png' + plots: + - add_xlabel: 'Total inner iteration number' + add_ylabel: 'Quadratic Cost Function' + add_legend: + layers: + - type: LinePlot + label: 'jojc' + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::jojc + color: 'blue' + markersize: 2 + - type: LinePlot + label: 'jb' + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::jb + color: 'red' + markersize: 2 diff --git a/src/swell/suites/eda/eva/observations-geos_atmosphere.yaml b/src/swell/suites/eda/eva/observations-geos_atmosphere.yaml new file mode 100644 index 000000000..edee73160 --- /dev/null +++ b/src/swell/suites/eda/eva/observations-geos_atmosphere.yaml @@ -0,0 +1,564 @@ +datasets: + +- name: experiment + type: IodaObsSpace + filenames: + - {{obs_path_file}} + channels: &channels {{channels}} + groups: + - name: ObsValue + variables: &variables {{simulated_variables}} + - name: GsiHofXBc + #- name: GsiEffectiveQC + - name: hofx0 + - name: hofx1 + - name: ombg + - name: oman + - name: EffectiveQC0 + - name: EffectiveQC1 + - name: MetaData + +transforms: + +# Generate hofx0 for GSI +- transform: arithmetic + new name: experiment::ObsValueMinusGsiHofXBc::${variable} + equals: experiment::ObsValue::${variable}-experiment::GsiHofXBc::${variable} + for: + variable: *variables + +# Generate hofx0 for JEDI +- transform: arithmetic + new name: experiment::ObsValueMinusHofx0::${variable} + equals: experiment::ObsValue::${variable}-experiment::hofx0::${variable} + for: + variable: *variables + +# Generate hofx difference +- transform: arithmetic + new name: experiment::Hofx0MinusGsiHofXBc::${variable} + equals: experiment::hofx0::${variable}-experiment::GsiHofXBc::${variable} + for: + variable: *variables + +# Generate hofx that passed QC for JEDI +- transform: accept where + new name: experiment::hofx0PassedQc::${variable} + starting field: experiment::hofx0::${variable} + where: + - experiment::EffectiveQC0::${variable} == 0 + for: + variable: *variables + +# Generate GSI hofx that passed JEDI QC +- transform: accept where + new name: experiment::GsiHofXBcPassedQc::${variable} + starting field: experiment::GsiHofXBc::${variable} + where: + - experiment::EffectiveQC0::${variable} == 0 + for: + variable: *variables + +# Generate hofx0 that passed QC for JEDI +- transform: accept where + new name: experiment::ObsValueMinushofx0PassedQc::${variable} + starting field: experiment::ObsValueMinusHofx0::${variable} + where: + - experiment::EffectiveQC0::${variable} == 0 + for: + variable: *variables + +# Generate hofx0 that passed QC for GSI +- transform: accept where + new name: experiment::ObsValueMinusGsiHofXBcPassedQc::${variable} + starting field: experiment::ObsValueMinusGsiHofXBc::${variable} + where: + - experiment::EffectiveQC0::${variable} == 0 + for: + variable: *variables + +# Generate ombg that passed QC for JEDI +- transform: accept where + new name: experiment::ombgPassedQc::${variable} + starting field: experiment::ombg::${variable} + where: + - experiment::EffectiveQC0::${variable} == 0 + for: + variable: *variables + +# Generate oman that passed QC for JEDI +- transform: accept where + new name: experiment::omanPassedQc::${variable} + starting field: experiment::oman::${variable} + where: + - experiment::EffectiveQC1::${variable} == 0 + for: + variable: *variables + +# Generate obs contribution to analysis (OmA*OmA)-(OmB*OmB) +- transform: arithmetic + new name: experiment::ResidualRMSdiff::${variable} + equals: (experiment::omanPassedQc::${variable})*(experiment::omanPassedQc::${variable})-(experiment::ombgPassedQc::${variable})*(experiment::ombgPassedQc::${variable}) + for: + variable: *variables + +graphics: + + plotting_backend: Emcpy + figure_list: + + # Correlation scatter plots + # ------------------------- + + # JEDI h(x) vs Observations + - batch figure: + variables: *variables + channels: *channels + figure: + layout: [1,1] + title: 'Observations vs. JEDI h(x) | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/correlation_scatter/${variable}${channel}/jedi_hofx0_vs_obs_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'Observation Value' + add_ylabel: 'JEDI h(x)' + add_grid: + add_legend: + loc: 'upper left' + layers: + - type: Scatter + x: + variable: experiment::ObsValue::${variable} + y: + variable: experiment::hofx0::${variable} + channel: ${channel} + markersize: 5 + color: 'black' + label: 'JEDI h(x) versus obs (all obs)' + - type: Scatter + x: + variable: experiment::ObsValue::${variable} + y: + variable: experiment::hofx0PassedQc::${variable} + channel: ${channel} + markersize: 5 + color: 'red' + label: 'JEDI h(x) versus obs (passed QC in JEDI)' + + # GSI h(x) vs Observations + - batch figure: + variables: *variables + channels: *channels + figure: + layout: [1,1] + title: 'Observations vs. GSI h(x) | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/correlation_scatter/${variable}${channel}/gsi_hofx0_vs_obs_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'Observation Value' + add_ylabel: 'GSI h(x)' + add_grid: + add_legend: + loc: 'upper left' + layers: + - type: Scatter + x: + variable: experiment::ObsValue::${variable} + y: + variable: experiment::GsiHofXBc::${variable} + channel: ${channel} + markersize: 5 + color: 'black' + label: 'GSI h(x) versus obs (all obs)' + - type: Scatter + x: + variable: experiment::ObsValue::${variable} + y: + variable: experiment::GsiHofXBcPassedQc::${variable} + channel: ${channel} + markersize: 5 + color: 'red' + label: 'GSI h(x) versus obs (passed QC in JEDI)' + + # JEDI h(x) vs GSI h(x) + - batch figure: + variables: *variables + channels: *channels + figure: + layout: [1,1] + title: 'JEDI h(x) vs. GSI h(x) | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/correlation_scatter/${variable}${channel}/gsi_hofx_vs_jedi0_hofx_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'GSI h(x)' + add_ylabel: 'JEDI h(x)' + add_grid: + add_legend: + loc: 'upper left' + layers: + - type: Scatter + x: + variable: experiment::GsiHofXBc::${variable} + y: + variable: experiment::hofx0::${variable} + channel: ${channel} + markersize: 5 + color: 'black' + label: 'JEDI h(x) versus GSI h(x)' + - type: Scatter + x: + variable: experiment::GsiHofXBcPassedQc::${variable} + y: + variable: experiment::hofx0PassedQc::${variable} + channel: ${channel} + markersize: 5 + color: 'red' + label: 'JEDI h(x) versus GSI h(x) (passed QC in JEDI)' + + # JEDI hofx0 vs GSI hofx0 + - batch figure: + variables: *variables + channels: *channels + figure: + layout: [1,1] + title: 'JEDI hofx0 vs. GSI hofx0 | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/correlation_scatter/${variable}${channel}/gsi_hofx0_vs_jedi_hofx0_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'GSI observation minus h(x)' + add_ylabel: 'JEDI observation minus h(x)' + add_grid: + add_legend: + loc: 'upper left' + layers: + - type: Scatter + x: + variable: experiment::ObsValueMinusGsiHofXBc::${variable} + y: + variable: experiment::ObsValueMinusHofx0::${variable} + channel: ${channel} + markersize: 5 + color: 'black' + label: 'GSI hofx0 vs JEDI hofx0 (all obs)' + - type: Scatter + x: + variable: experiment::ObsValueMinusGsiHofXBcPassedQc::${variable} + y: + variable: experiment::ObsValueMinushofx0PassedQc::${variable} + channel: ${channel} + markersize: 5 + color: 'red' + label: 'GSI hofx0 vs JEDI hofx0 (passed QC in JEDI)' + + # JEDI oma vs omb + - batch figure: + variables: *variables + channels: *channels + figure: + layout: [1,1] + title: 'OmA vs. OmB | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/correlation_scatter/${variable}${channel}/jedi_oma_vs_omb_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'OmA' + add_ylabel: 'OmB' + add_grid: + add_legend: + loc: 'upper left' + layers: + - type: Scatter + x: + variable: experiment::oman::${variable} + y: + variable: experiment::ombg::${variable} + channel: ${channel} + markersize: 5 + color: 'black' + label: 'OmA versus OmB (all residuals)' + - type: Scatter + x: + variable: experiment::omanPassedQc::${variable} + y: + variable: experiment::ombgPassedQc::${variable} + channel: ${channel} + markersize: 5 + color: 'red' + label: 'OmA versus OmB (passed QC)' + +# Map plots# --------- + + # Observations + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: vminvmaxcmap + channel: ${channel} + data variable: experiment::ObsValue::${variable} + figure: + figure size: [20,10] + layout: [1,1] + title: 'Observations | {{instrument_title}} | Obs Value' + output name: '{{cycle_dir}}/eva/{{instrument}}/map_plots/${variable}${channel}/observations_{{instrument}}_${variable}${channel}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: ObsValue + add_grid: + layers: + - type: MapScatter + longitude: + variable: experiment::MetaData::longitude + latitude: + variable: experiment::MetaData::latitude + data: + variable: experiment::ObsValue::${variable} + channel: ${channel} + markersize: 2 + label: ObsValue + colorbar: true + cmap: ${dynamic_cmap} + vmin: ${dynamic_vmin} + vmax: ${dynamic_vmax} + + # hofx0 jedi + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: vminvmaxcmap + channel: ${channel} + data variable: experiment::ObsValueMinusHofx0::${variable} + figure: + figure size: [20,10] + layout: [1,1] + title: 'JEDI hofx0 | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/map_plots/${variable}${channel}/hofx0_jedi_{{instrument}}_${variable}${channel}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: '${variable}' + add_grid: + layers: + - type: MapScatter + longitude: + variable: experiment::MetaData::longitude + latitude: + variable: experiment::MetaData::latitude + data: + variable: experiment::ObsValueMinusHofx0::${variable} + channel: ${channel} + markersize: 2 + label: '${variable}' + colorbar: true + cmap: ${dynamic_cmap} + vmin: ${dynamic_vmin} + vmax: ${dynamic_vmax} + + # hofx0 gsi + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: vminvmaxcmap + channel: ${channel} + data variable: experiment::ObsValueMinusGsiHofXBc::${variable} + figure: + figure size: [20,10] + layout: [1,1] + title: 'GSI hofx0 | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/map_plots/${variable}${channel}/hofx0_gsi_{{instrument}}_${variable}${channel}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: '${variable}' + add_grid: + layers: + - type: MapScatter + longitude: + variable: experiment::MetaData::longitude + latitude: + variable: experiment::MetaData::latitude + data: + variable: experiment::ObsValueMinusGsiHofXBc::${variable} + channel: ${channel} + markersize: 2 + label: '${variable}' + colorbar: true + cmap: ${dynamic_cmap} + vmin: ${dynamic_vmin} + vmax: ${dynamic_vmax} + + # hofx difference + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: vminvmaxcmap + channel: ${channel} + data variable: experiment::Hofx0MinusGsiHofXBc::${variable} + figure: + figure size: [20,10] + layout: [1,1] + title: 'Hofx0 Difference | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/map_plots/${variable}${channel}/hofx0_difference_{{instrument}}_${variable}${channel}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: '${variable}' + add_grid: + layers: + - type: MapScatter + longitude: + variable: experiment::MetaData::longitude + latitude: + variable: experiment::MetaData::latitude + data: + variable: experiment::Hofx0MinusGsiHofXBc::${variable} + channel: ${channel} + markersize: 2 + label: '${variable}' + colorbar: true + cmap: ${dynamic_cmap} + vmin: ${dynamic_vmin} + vmax: ${dynamic_vmax} + + + # RMS(oma)-RMS(omb) difference + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: vminvmaxcmap + channel: ${channel} + data variable: experiment::ResidualRMSdiff::${variable} + figure: + figure size: [20,10] + layout: [1,1] + title: 'RMS Residual Difference | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/map_plots/${variable}${channel}/rmsres_difference_{{instrument}}_${variable}${channel}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: '${variable}' + add_grid: + layers: + - type: MapScatter + longitude: + variable: experiment::MetaData::longitude + latitude: + variable: experiment::MetaData::latitude + data: + variable: experiment::ResidualRMSdiff::${variable} + channel: ${channel} + markersize: 2 + label: '${variable}' + colorbar: true + cmap: ${dynamic_cmap} + vmin: ${dynamic_vmin} + vmax: ${dynamic_vmax} +# Histogram plots# --------------- + + # hofx0 vs hofx0 + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: histogram_bins + channel: ${channel} + number of bins rule: sturges + data variable: experiment::ObsValueMinusHofx0::${variable} + figure: + layout: [1,1] + title: 'JEDI hofx0 vs. GSI hofx0 | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/histograms/${variable}${channel}/gsi_hofx0_vs_jedi_hofx0_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'Observation minus h(x)' + add_ylabel: 'Count' + add_legend: + loc: 'upper left' + layers: + - type: Histogram + data: + variable: experiment::ObsValueMinusGsiHofXBc::${variable} + channel: ${channel} + color: 'blue' + label: 'GSI hofx0 (all obs)' + bins: ${dynamic_bins} + alpha: 0.5 + - type: Histogram + data: + variable: experiment::ObsValueMinusHofx0::${variable} + channel: ${channel} + color: 'red' + label: 'JEDI hofx0 (all obs)' + bins: ${dynamic_bins} + alpha: 0.5 + + # JEDI omb vs oma + - batch figure: + variables: *variables + dynamic options: + - type: histogram_bins + data variable: experiment::omanPassedQc::${variable} + number of bins rule: 'rice' + figure: + layout: [1,1] + title: 'OmB vs. OmA | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/histograms/${variable}/ombg_oman_{{instrument}}_${variable}.png' + plots: + - add_xlabel: 'Difference' + add_ylabel: 'Count' + set_xlim: [-3, 3] + add_legend: + loc: 'upper left' + statistics: + fields: + - field_name: experiment::ombgPassedQc::${variable} + xloc: 0.5 + yloc: -0.10 + kwargs: + color: 'black' + fontsize: 8 + fontfamily: monospace + - field_name: experiment::omanPassedQc::${variable} + xloc: 0.5 + yloc: -0.13 + kwargs: + color: 'red' + fontsize: 8 + fontfamily: monospace + statistics_variables: + - n + - min + - mean + - max + - std + layers: + - type: Histogram + data: + variable: experiment::ombgPassedQc::${variable} + color: 'red' + label: 'observations minus background ' + bins: ${dynamic_bins} + alpha: 0.5 + density: true + - type: Histogram + data: + variable: experiment::omanPassedQc::${variable} + color: 'blue' + label: 'observations minus analysis' + bins: ${dynamic_bins} + alpha: 0.5 + density: true + diff --git a/src/swell/suites/eda/flow.cylc b/src/swell/suites/eda/flow.cylc new file mode 100644 index 000000000..9f1e31357 --- /dev/null +++ b/src/swell/suites/eda/flow.cylc @@ -0,0 +1,228 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +# Cylc suite for executing JEDI-based non-cycling variational data assimilation + +# -------------------------------------------------------------------------------------------------- + +[scheduler] + UTC mode = True + allow implicit tasks = False + +# -------------------------------------------------------------------------------------------------- + +[scheduling] + + initial cycle point = {{start_cycle_point}} + final cycle point = {{final_cycle_point}} + runahead limit = {{runahead_limit}} + + [[graph]] + R1 = """ + # Triggers for non cycle time dependent tasks + # ------------------------------------------- + # Clone JEDI source code + CloneJedi + + # Build JEDI source code by linking + CloneJedi => BuildJediByLinking? + + # If not able to link to build create the build + BuildJediByLinking:fail? => BuildJedi + + {% for model_component in model_components %} + # Clone geos ana for generating observing system records + CloneGeosMksi-{{model_component}} + {% endfor %} + """ + + {% for cycle_time in cycle_times %} + {{cycle_time.cycle_time}} = """ + {% for model_component in model_components %} + {% if cycle_time[model_component] %} + + # Task triggers for: {{model_component}} + # ------------------ + # Generate satellite channel records + CloneGeosMksi-{{model_component}}[^] => GenerateObservingSystemRecords-{{model_component}} + + # Get background, provide a way to get background directly from GEOS experiment + GetBackgroundGeosExperiment-{{model_component}} :fail? => GetBackground-{{model_component}} + + # Get observations + {% if cycling_varbc %} + # Cycling VarBC is active, biases from the previous cycle will be used + + RunJediEdaExecutable-{{model_component}}[-PT6H] => GetObservations-{{model_component}} + {% else %} + + # Cycling VarBC is inactive, static bias files will be used + GetObsNotInR2d2-{{model_component}}: fail? => GetObservations-{{model_component}} + {% endif %} + + # Perform staging that is cycle dependent + StageJediCycle-{{model_component}} + + # Run Jedi variational executable + BuildJediByLinking[^]? | BuildJedi[^] => SP1-{{model_component}} + CloneJedi[^] => StageJediCycle-{{model_component}} + StageJediCycle-{{model_component}} => SP1-{{model_component}} + GetBackgroundGeosExperiment-{{model_component}}? | GetBackground-{{model_component}} => SP1-{{model_component}} + GetObsNotInR2d2-{{model_component}}? | GetObservations-{{model_component}} => SP1-{{model_component}} + SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} + GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} + RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} + {% for i in range(1, 3) %} + SP2-{{model_component}} => RunJediEdaExecutable_mem{{i}}-{{model_component}} => SP3-{{model_component}} + {% endfor %} + + # EvaObservations + SP3-{{model_component}} => EvaObservations-{{model_component}} + + # EvaJediLog + SP3-{{model_component}} => EvaJediLog-{{model_component}} + + # EvaIncrement + SP3-{{model_component}} => EvaIncrement-{{model_component}} + + # Clean up large files + EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} & + EvaObservations-{{model_component}} => + CleanCycle-{{model_component}} + + {% endif %} + {% endfor %} + """ + {% endfor %} + +# -------------------------------------------------------------------------------------------------- + +[runtime] + + # Task defaults + # ------------- + [[root]] + pre-script = "source $CYLC_SUITE_DEF_PATH/modules" + + [[[environment]]] + datetime = $CYLC_TASK_CYCLE_POINT + config = $CYLC_SUITE_DEF_PATH/experiment.yaml + + # Tasks + # ----- + [[CloneJedi]] + script = "swell task CloneJedi $config" + + [[BuildJediByLinking]] + script = "swell task BuildJediByLinking $config" + + [[BuildJedi]] + script = "swell task BuildJedi $config" + platform = {{platform}} + execution time limit = {{scheduling["BuildJedi"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["BuildJedi"]["directives"]["all"].items() %} + --{{key}} = {{value}} + {%- endfor %} + + {% for model_component in model_components %} + + [[CloneGeosMksi-{{model_component}}]] + script = "swell task CloneGeosMksi $config -m {{model_component}}" + + [[GenerateObservingSystemRecords-{{model_component}}]] + script = "swell task GenerateObservingSystemRecords $config -d $datetime -m {{model_component}}" + + [[GetObsNotInR2d2-{{model_component}}]] + script = "swell task GetObsNotInR2d2 $config -d $datetime -m {{model_component}}" + + [[StageJediCycle-{{model_component}}]] + script = "swell task StageJedi $config -d $datetime -m {{model_component}}" + + [[GetBackground-{{model_component}} ]] + script = "swell task GetBackground $config -d $datetime -m {{model_component}}" + + [[GetBackgroundGeosExperiment-{{model_component}} ]] + script = "swell task GetBackgroundGeosExperiment $config -d $datetime -m {{model_component}}" + + [[GetObservations-{{model_component}}]] + script = "swell task GetObservations $config -d $datetime -m {{model_component}}" + + [[RenderJediObservations-{{model_component}}]] + script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" + + {% for i in range(1, 3) %} + [[RunJediEdaExecutable_mem{{i}}-{{model_component}}]] + script = "swell task RunJediEdaExecutable $config -d $datetime -m {{model_component}} -imem {{i}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEdaExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEdaExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + {% endfor %} + + [[RunJediEdaExecutable-{{model_component}}]] + script = "swell task RunJediEdaExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEdaExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEdaExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[RunJediObsfiltersExecutable-{{model_component}}]] + script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediObsfiltersExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediObsfiltersExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[RunJediEnsembleMeanVariance-{{model_component}}]] + script = "swell task RunJediEnsembleMeanVariance $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEnsembleMeanVariance"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEnsembleMeanVariance"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[EvaJediLog-{{model_component}}]] + script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}" + + [[EvaIncrement-{{model_component}}]] + script = "swell task EvaIncrement $config -d $datetime -m {{model_component}}" + + [[EvaObservations-{{model_component}}]] + script = "swell task EvaObservations $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["EvaObservations"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["EvaObservations"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[SaveObsDiags-{{model_component}}]] + script = "swell task SaveObsDiags $config -d $datetime -m {{model_component}}" + + [[CleanCycle-{{model_component}}]] + script = "swell task CleanCycle $config -d $datetime -m {{model_component}}" + + [[SP1-{{model_component}}]] + script = true + + [[SP2-{{model_component}}]] + script = true + + [[SP3-{{model_component}}]] + script = true + {% endfor %} + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda/suite_config.py new file mode 100644 index 000000000..a85dbd183 --- /dev/null +++ b/src/swell/suites/eda/suite_config.py @@ -0,0 +1,119 @@ +# -------------------------------------------------------------------------------------------------- +# @package configuration +# +# Class containing the configuration. This is a dictionary that is converted from +# an input yaml configuration file. Various function are included for interacting with the +# dictionary. +# +# -------------------------------------------------------------------------------------------------- + + +from swell.utilities.swell_questions import QuestionContainer, QuestionList +from swell.utilities.question_defaults import QuestionDefaults as qd +from swell.suites.suite_questions import SuiteQuestions as sq + +from enum import Enum + + +# -------------------------------------------------------------------------------------------------- + +class SuiteConfig(QuestionContainer, Enum): + + # -------------------------------------------------------------------------------------------------- + + eda_atmos_tier1 = QuestionList( + list_name="eda_atmos_tier1", + questions=[ + sq.common, + qd.start_cycle_point("2023-10-10T00:00:00Z"), + qd.final_cycle_point("2023-10-10T06:00:00Z"), + qd.runahead_limit("P2"), + qd.jedi_build_method("use_existing"), + qd.model_components(['geos_atmosphere']), + qd.ensemble_num_members(3), + ], + geos_atmosphere=[ + qd.cycle_times([ + "T00", + "T06", + "T12", + "T18" + ]), + qd.geos_x_background_directory("/discover/nobackup/projects/gmao/" + "dadev/rtodling/archive/Restarts/JEDI/541x"), + qd.window_length("PT6H"), + qd.window_type("3D"), + qd.horizontal_resolution("91"), + qd.gsibec_nlats("91"), + qd.gsibec_nlons("144"), + qd.vertical_resolution("72"), +# qd.ensemble_num_members(3), + qd.obs_pert_amplitude(0.5), + qd.analysis_variables([ + "eastward_wind", + "northward_wind", + "air_temperature", + "water_vapor_mixing_ratio_wrt_moist_air", + "air_pressure_at_surface", +# "air_pressure_levels", + "cloud_liquid_ice", + "cloud_liquid_water", + "rain_water", + "snow_water", + "mole_fraction_of_ozone_in_air", + "geopotential_height_times_gravity_at_surface", + "fraction_of_ocean", + "fraction_of_lake", + "fraction_of_ice", + "skin_temperature_at_surface" + ]), + qd.observations([ + "aircraft_temperature", + "aircraft_wind", + "airs_aqua", + "amsr2_gcom-w1", + "amsua_aqua", + "amsua_metop-b", + "amsua_metop-c", + "amsua_n15", + "amsua_n18", + "amsua_n19", + "atms_n20", + "atms_npp", + "avhrr3_metop-b", + "avhrr3_n18", + "avhrr3_n19", + "cris-fsr_n20", + "cris-fsr_npp", + "gmi_gpm", + "gps", + "iasi_metop-b", + "iasi_metop-c", + "mhs_metop-b", + "mhs_metop-c", + "mhs_n19", + "mls55_aura", + "omi_aura", + "ompsnm_npp", + "pibal", + "satwind", + "scatwind", + "sfcship", + "sfc", + "sondes", + "ssmis_f17" + ]), + qd.clean_patterns(['*.txt', '*.csv']), + ] + ) + + # -------------------------------------------------------------------------------------------------- + + eda = QuestionList( + list_name="eda", + questions=[ + eda_atmos_tier1 + ] + ) + + # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py new file mode 100644 index 000000000..27dd6a68b --- /dev/null +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -0,0 +1,195 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + + +# -------------------------------------------------------------------------------------------------- + +import os +import copy +from ruamel.yaml import YAML + +from swell.tasks.base.task_base import taskBase +from swell.utilities.run_jedi_executables import run_executable + + +# -------------------------------------------------------------------------------------------------- + + +class RunJediEdaExecutable(taskBase): + + # ---------------------------------------------------------------------------------------------- + + def execute(self) -> None: + + # Jedi application name + # --------------------- + jedi_application = 'eda' + + # Parse configuration + # ------------------- + window_type = self.config.window_type() + window_length = self.config.window_length() + forecast_length = self.config.forecast_length(window_length) + background_time_offset = self.config.background_time_offset() + number_of_iterations = self.config.number_of_iterations() + jedi_forecast_model = self.config.jedi_forecast_model(None) + generate_yaml_and_exit = self.config.generate_yaml_and_exit(False) + perhost = self.config.perhost(None) + + # Set the observing system records path + self.jedi_rendering.set_obs_records_path(self.config.observing_system_records_path(None)) + + gsibec_nlats = self.config.gsibec_nlats(None) + gsibec_nlons = self.config.gsibec_nlons(None) + gsibec_configuration = self.config.gsibec_configuration(None) + npx_proc = self.config.npx_proc(None) + npy_proc = self.config.npy_proc(None) + npx = self.config.npx(None) + npy = self.config.npy(None) + + # Compute data assimilation window parameters + # -------------------------------------------- + background_time = self.da_window_params.background_time(background_time_offset) + local_background_time = self.da_window_params.local_background_time(window_length, + window_type) + local_background_time_iso = self.da_window_params.local_background_time_iso(window_length, + window_type) + window_begin = self.da_window_params.window_begin(window_length) + window_begin_iso = self.da_window_params.window_begin_iso(window_length) + window_end_iso = self.da_window_params.window_end_iso(window_length) + nmember = self.config.ensemble_num_members() + imember = self.get_ensemble_imember() +### print( f' imember = {imember}' ) + + # Populate jedi interface templates dictionary + # -------------------------------------------- + self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) + self.jedi_rendering.add_key('window_end_iso', window_end_iso) + self.jedi_rendering.add_key('window_length', window_length) + self.jedi_rendering.add_key('forecast_length', forecast_length) + self.jedi_rendering.add_key('minimizer', self.config.minimizer()) + self.jedi_rendering.add_key('number_of_iterations', number_of_iterations[0]) + self.jedi_rendering.add_key('analysis_variables', self.config.analysis_variables()) + self.jedi_rendering.add_key('saber_central_block', self.config.saber_central_block()) + self.jedi_rendering.add_key('saber_outer_block', self.config.saber_outer_block()) + self.jedi_rendering.add_key('gradient_norm_reduction', + self.config.gradient_norm_reduction()) + self.jedi_rendering.add_key('marine_models', self.config.marine_models(None)) + + # Background + # ---------- + self.jedi_rendering.add_key('horizontal_resolution', self.config.horizontal_resolution()) + self.jedi_rendering.add_key('local_background_time', local_background_time) + self.jedi_rendering.add_key('local_background_time_iso', local_background_time_iso) + self.jedi_rendering.add_key('ensemble_num_members', self.config.ensemble_num_members()) + self.jedi_rendering.add_key('ensemble_imember', imember) + + # Geometry + # -------- + self.jedi_rendering.add_key('vertical_resolution', self.config.vertical_resolution()) + self.jedi_rendering.add_key('gsibec_nlats', gsibec_nlats) + self.jedi_rendering.add_key('gsibec_nlons', gsibec_nlons) + self.jedi_rendering.add_key('npx_proc', npx_proc) + self.jedi_rendering.add_key('npy_proc', npy_proc) + self.jedi_rendering.add_key('npx', npx) + self.jedi_rendering.add_key('npy', npy) + self.jedi_rendering.add_key('total_processors', self.config.total_processors(None)) + + # Observations + # ------------ + self.jedi_rendering.add_key('background_time', background_time) + self.jedi_rendering.add_key('crtm_coeff_dir', self.config.crtm_coeff_dir(None)) + self.jedi_rendering.add_key('window_begin', window_begin) + + # Atmosphere background error model + # --------------------------------- + if gsibec_configuration is not None: + self.jedi_rendering.add_key('gsibec_configuration', gsibec_configuration) + self.jedi_rendering.add_key('gsibec_nlats', gsibec_nlats) + self.jedi_rendering.add_key('gsibec_nlons', gsibec_nlons) + self.jedi_rendering.add_key('gsibec_npx_proc', npx_proc) + self.jedi_rendering.add_key('gsibec_npy_proc', 6*npy_proc) + + # Model + # ----- + if window_type == '4D': + self.jedi_rendering.add_key('background_frequency', self.config.background_frequency()) + + + # Jedi configuration file + # ----------------------- + jedi_config_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}{window_type}_config_mem{imember:03d}.yaml') + + # Output log file + # --------------- + output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}{window_type}_log.log') + + # Open the JEDI config file and fill initial templates + # ---------------------------------------------------- + jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}{window_type}', + window_type, + jedi_forecast_model) + + # imember: specify yaml + # ---------------------------------------------------- + + if imember == 1: + jedi_config_dict['cost function']['observations'].pop('obs perturbations', None) + else: + jedi_config_dict['cost function']['observations'].update({'obs perturbations': True}) + + for observer in jedi_config_dict['cost function']['observations']['observers']: + # Get observation name + observation = observer['observation_name'] + if imember > 1: + observer['obs space'].update({'obs perturbations seed': imember}) + observer['obs error'].update({'obs perturbations amplitude': 0.5}) + hxout = observer['obs space']['obsdataout']['engine']['obsfile'] + dir1, fname = os.path.split(hxout) + hxout = os.path.join(dir1, f'diag/mem{imember:003d}', fname) + observer['obs space']['obsdataout']['engine']['obsfile'] = hxout + print (f'hxout = {hxout}') + + for it in jedi_config_dict['variational']['iterations']: + it['online diagnostics']['increment']['state component']['filename'] = f'eda.mem{imember:03d}.increment-iter1.' + jedi_config_dict['output']['filename'] = f'eda.mem{imember:03d}.analysis.%yyyy%mm%dd_%hh%MM%ssz.nc4' + + ruamel_yaml = YAML() + ruamel_yaml.default_flow_style = False + + # Write the ordered dictionary to YAML file + with open(jedi_config_file, 'w') as jedi_config_file_open: + ruamel_yaml.dump(jedi_config_dict, jedi_config_file_open) + + # Get the JEDI interface metadata + # ------------------------------- + model_component_meta = self.jedi_rendering.render_interface_meta() + + # Compute number of processors + # ---------------------------- + np = eval(str(model_component_meta['total_processors'])) + + # Jedi executable name + # -------------------- + jedi_executable = model_component_meta['executables'][f'variational{window_type}'] + jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', 'build', 'bin', + jedi_executable) + + # Run the JEDI executable + # ----------------------- + if not generate_yaml_and_exit: + self.logger.info('Running '+jedi_executable_path+' with '+str(np)+' processors.') + run_executable(self.logger, self.cycle_dir(), np, jedi_executable_path, + jedi_config_file, output_log_file, perhost) + else: + mpi_command = "mpirun" + if not (perhost is None or perhost == "None"): + mpi_command += f" -perhost {perhost}" + mpi_command += f" -np {np} {jedi_executable_path} {jedi_config_file} {output_log_file}" + print(f'intended mpi_command = {mpi_command}') + self.logger.info('YAML generated, now exiting.') + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/utilities/yaml_utils.py b/src/swell/utilities/yaml_utils.py new file mode 100644 index 000000000..fc73d28e8 --- /dev/null +++ b/src/swell/utilities/yaml_utils.py @@ -0,0 +1,19 @@ + +# -------------------------------------------------------------------------------------------------- + +def replace_key(obj, old_key, new_key): + """ + Recursively replace dictionary keys in nested dictionaries/lists. + """ + if isinstance(obj, dict): + new_dict = {} + for k, v in obj.items(): + new_k = new_key if k == old_key else k + new_dict[new_k] = replace_key(v, old_key, new_key) + return new_dict + elif isinstance(obj, list): + return [replace_key(item, old_key, new_key) for item in obj] + else: + return obj + +# -------------------------------------------------------------------------------------------------- From 278409bb9dcd6738a0ea638e401af1b80dabb187 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 27 May 2026 15:23:44 -0600 Subject: [PATCH 02/68] merged M files --- .../geos_atmosphere/geos_atmosphere.yaml | 2 ++ .../model/background_ensemble.py | 1 - src/swell/swell.py | 6 ++++-- src/swell/tasks/base/task_base.py | 17 +++++++++++++---- src/swell/tasks/task_questions.py | 12 ++++++++++++ src/swell/utilities/question_defaults.py | 13 +++++++++++++ .../utilities/render_jedi_interface_files.py | 1 + src/swell/utilities/run_jedi_executables.py | 2 +- src/swell/utilities/slurm.py | 1 + 9 files changed, 47 insertions(+), 8 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml b/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml index 2fc64df8e..ac17a05fc 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml @@ -11,3 +11,5 @@ executables: localensembleda: fv3jedi_letkf.x ensmeanvariance: fv3jedi_ensmeanvariance.x obsfilters: test_ObsFilters.x + eda3D: fv3jedi_var.x + eda4D: fv3jedi_var.x diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py index 9b904fb41..351ec1644 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py +++ b/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'], diff --git a/src/swell/swell.py b/src/swell/swell.py index 025327a8c..8e1d8accc 100644 --- a/src/swell/swell.py +++ b/src/swell/swell.py @@ -194,12 +194,14 @@ def launch( @click.option('-d', '--datetime', 'datetime', default=None, help=datetime_help) @click.option('-m', '--model', 'model', default=None, help=model_help) @click.option('-p', '--ensemblePacket', 'ensemblePacket', default=None, help=ensemble_help) +@click.option('-imem', '--ensemble_imember', 'imember', type=int, default=None, help=ensemble_help) def task( task: str, config: str, datetime: Optional[str], model: Optional[str], - ensemblePacket: Optional[str] + ensemblePacket: Optional[str], + imember: Optional[int] ) -> None: """ Run a workflow task @@ -211,7 +213,7 @@ def task( config (str): Path to the configuration file for the task.\n """ - task_wrapper(task, config, datetime, model, ensemblePacket) + task_wrapper(task, config, datetime, model, ensemblePacket, imember) # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/base/task_base.py b/src/swell/tasks/base/task_base.py index 109b3af74..2134cda8f 100644 --- a/src/swell/tasks/base/task_base.py +++ b/src/swell/tasks/base/task_base.py @@ -41,6 +41,7 @@ def __init__( datetime_input: Optional[str], model: str, ensemblePacket: Optional[str], + imember: Optional[int], task_name: str ) -> None: @@ -65,6 +66,7 @@ def __init__( # Keep copy of ensemblePacket # --------------------------- self.__ensemble_packet__ = ensemblePacket + self.__ensemble_imember__ = imember # Keep copy of model directive # ---------------------------- @@ -168,6 +170,11 @@ def get_ensemble_packet(self) -> Optional[str]: # ---------------------------------------------------------------------------------------------- + def get_ensemble_imember(self) -> int | None: + return self.__ensemble_imember__ + + # ---------------------------------------------------------------------------------------------- + def get_model(self) -> str: return self.__model__ @@ -274,7 +281,8 @@ def create_task( config: str, datetime: Union[str, dt, None], model: str, - ensemblePacket: Optional[str] + ensemblePacket: Optional[str], + imember: Optional[int] ) -> taskBase: # Convert camel case string to snake case @@ -303,7 +311,7 @@ def create_task( factory_logger.info(f'Using module swell.tasks.{task_lower}') # Return task object - return task_class(config, datetime, model, ensemblePacket, task) + return task_class(config, datetime, model, ensemblePacket, imember, task) # -------------------------------------------------------------------------------------------------- @@ -334,13 +342,14 @@ def task_wrapper( config: str, datetime: Union[str, dt, None], model: Optional[str], - ensemblePacket: Optional[str] + ensemblePacket: Optional[str], + imember: Optional[int] ) -> None: # Create the object constrc_start = time.perf_counter() creator = taskFactory() - task_object = creator.create_task(task, config, datetime, model, ensemblePacket) + task_object = creator.create_task(task, config, datetime, model, ensemblePacket, imember) constrc_final = time.perf_counter() constrc_time = f'Constructed in {constrc_final - constrc_start:0.4f} seconds' diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index b8cf4f755..75494dce6 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -785,6 +785,18 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + RunJediEdaExecutable = QuestionList( + list_name="RunJediEdaExecutable", + questions=[ + run_jedi_executable, + qd.ensemble_num_members(), + qd.perhost(), + qd.comparison_log_type('variational'), + ] + ) + + # -------------------------------------------------------------------------------------------------- + SaveObsDiags = QuestionList( list_name="SaveObsDiags", questions=[ diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index 741918fab..d36c9d108 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -462,6 +462,19 @@ class ensemble_num_members(TaskQuestion): # -------------------------------------------------------------------------------------------------- + @dataclass + class obs_pert_amplitude(TaskQuestion): + default_value: str = "defer_to_model" + question_name: str = "obs_pert_amplitude" + options: str = "defer_to_model" + models: List[str] = mutable_field([ + "geos_atmosphere" + ]) + prompt: str = "Enter obs perturbation amplitude for EDA:" + widget_type: WType = WType.FLOAT + + # -------------------------------------------------------------------------------------------------- + @dataclass class ensmean_only(TaskQuestion): default_value: bool = False diff --git a/src/swell/utilities/render_jedi_interface_files.py b/src/swell/utilities/render_jedi_interface_files.py index e303d754a..4db12b2c0 100644 --- a/src/swell/utilities/render_jedi_interface_files.py +++ b/src/swell/utilities/render_jedi_interface_files.py @@ -81,6 +81,7 @@ def __init__( 'ensemble_hofx_packets', 'ensemble_hofx_strategy', 'ensemble_num_members', + 'ensemble_imember', 'ensmean_only', 'ensmeanvariance_only', 'experiment_id', diff --git a/src/swell/utilities/run_jedi_executables.py b/src/swell/utilities/run_jedi_executables.py index c4457a9c7..46cce2797 100644 --- a/src/swell/utilities/run_jedi_executables.py +++ b/src/swell/utilities/run_jedi_executables.py @@ -74,7 +74,7 @@ def run_executable( # Run the JEDI executable # ----------------------- - if perhost is None: + if (perhost is None or perhost == "None"): logger.info(f"Running {jedi_executable_path} with {str(np)} processors.") command = [ 'mpirun', diff --git a/src/swell/utilities/slurm.py b/src/swell/utilities/slurm.py index 787650d81..a3d6be04d 100644 --- a/src/swell/utilities/slurm.py +++ b/src/swell/utilities/slurm.py @@ -79,6 +79,7 @@ def prepare_scheduling_dict( 'EvaTimeseries', 'GenerateBClimatology', 'RunGeos', + 'RunJediEdaExecutable', 'RunJediEnsembleMeanVariance', 'RunJediConvertStateSoca2ciceExecutable', 'RunJediFgatExecutable', From e87abbbe62a2be49182865ab430d62aa8bc54138 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 28 May 2026 00:19:42 -0600 Subject: [PATCH 03/68] + self.jedi_rendering.add_key('saber_central_block', self.config.saber_central_block(None)) + self.jedi_rendering.add_key('saber_outer_block', self.config.saber_outer_block(None)) --- src/swell/tasks/run_jedi_eda_executable.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 27dd6a68b..22604b4c1 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -73,8 +73,8 @@ def execute(self) -> None: self.jedi_rendering.add_key('minimizer', self.config.minimizer()) self.jedi_rendering.add_key('number_of_iterations', number_of_iterations[0]) self.jedi_rendering.add_key('analysis_variables', self.config.analysis_variables()) - self.jedi_rendering.add_key('saber_central_block', self.config.saber_central_block()) - self.jedi_rendering.add_key('saber_outer_block', self.config.saber_outer_block()) + self.jedi_rendering.add_key('saber_central_block', self.config.saber_central_block(None)) + self.jedi_rendering.add_key('saber_outer_block', self.config.saber_outer_block(None)) self.jedi_rendering.add_key('gradient_norm_reduction', self.config.gradient_norm_reduction()) self.jedi_rendering.add_key('marine_models', self.config.marine_models(None)) From 3918af38a7655b1590051a540c4788e2f9933fc6 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 28 May 2026 23:34:01 -0600 Subject: [PATCH 04/68] update new --- src/swell/suites/eda/suite_config.py | 8 +++++--- src/swell/tasks/run_jedi_eda_executable.py | 2 +- src/swell/tasks/task_questions.py | 1 + 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda/suite_config.py index a85dbd183..24c019999 100644 --- a/src/swell/suites/eda/suite_config.py +++ b/src/swell/suites/eda/suite_config.py @@ -27,10 +27,11 @@ class SuiteConfig(QuestionContainer, Enum): sq.common, qd.start_cycle_point("2023-10-10T00:00:00Z"), qd.final_cycle_point("2023-10-10T06:00:00Z"), - qd.runahead_limit("P2"), + qd.runahead_limit("P1"), qd.jedi_build_method("use_existing"), qd.model_components(['geos_atmosphere']), - qd.ensemble_num_members(3), + ## + ## qd.ensemble_num_members(2), ], geos_atmosphere=[ qd.cycle_times([ @@ -47,7 +48,7 @@ class SuiteConfig(QuestionContainer, Enum): qd.gsibec_nlats("91"), qd.gsibec_nlons("144"), qd.vertical_resolution("72"), -# qd.ensemble_num_members(3), + qd.ensemble_num_members(2), qd.obs_pert_amplitude(0.5), qd.analysis_variables([ "eastward_wind", @@ -103,6 +104,7 @@ class SuiteConfig(QuestionContainer, Enum): "sondes", "ssmis_f17" ]), + qd.obs_thinning_rej_fraction(0.75), qd.clean_patterns(['*.txt', '*.csv']), ] ) diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 22604b4c1..380287a6b 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -125,7 +125,7 @@ def execute(self) -> None: # Output log file # --------------- - output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}{window_type}_log.log') + output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}{window_type}_log_mem{imember:03d}.log') # Open the JEDI config file and fill initial templates # ---------------------------------------------------- diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index 75494dce6..e80bd5a7f 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -790,6 +790,7 @@ class TaskQuestions(QuestionContainer, Enum): questions=[ run_jedi_executable, qd.ensemble_num_members(), + qd.obs_thinning_rej_fraction(), qd.perhost(), qd.comparison_log_type('variational'), ] From d83347c726273f44ac9d49ed32204bbe865951ba Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Fri, 29 May 2026 23:49:00 -0600 Subject: [PATCH 05/68] update eda fn from variational ymal input ens_mean_var read output from obs perturbed 3d-Var --- .../jedi/interfaces/geos_atmosphere/model/analysis.py | 5 +++++ .../jedi/interfaces/geos_atmosphere/model/ensemble_block.py | 6 ++++++ src/swell/tasks/run_jedi_eda_executable.py | 5 ++--- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/analysis.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/analysis.py index ca7a53983..dfa4e8ae5 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/analysis.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/analysis.py @@ -23,6 +23,11 @@ def analysis(template_dict: Mapping) -> Mapping: 'field io names': field_io_names, } + # EDA case: override filename + imem = template_dict.get('ensemble_imember', None) + if imem is not None: + analysis['filename'] = f'eda.mem{imem:03d}.analysis.%yyyy%mm%dd_%hh%MM%ssz.nc4' + return analysis # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_block.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_block.py index bc6f365d1..5bf6d6b33 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_block.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_block.py @@ -77,6 +77,12 @@ def ensemble_block(template_dict: Mapping) -> Mapping: } } + imem = template_dict.get('ensemble_imem', None) + if imem is not None: + # EDA diagnostic case + fn = ensemble_block['members from template']['template']['filename'] + fn = f'eda.mem{imem:03d}.analysis.%yyyy%mm%dd_%hh%MM%ssz.nc4' + return ensemble_block # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 380287a6b..916a4b195 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -153,9 +153,8 @@ def execute(self) -> None: observer['obs space']['obsdataout']['engine']['obsfile'] = hxout print (f'hxout = {hxout}') - for it in jedi_config_dict['variational']['iterations']: - it['online diagnostics']['increment']['state component']['filename'] = f'eda.mem{imember:03d}.increment-iter1.' - jedi_config_dict['output']['filename'] = f'eda.mem{imember:03d}.analysis.%yyyy%mm%dd_%hh%MM%ssz.nc4' + for i, iter in enumerate(jedi_config_dict['variational']['iterations']): + iter['online diagnostics']['increment']['state component']['filename'] = f'eda.mem{imember:03d}.increment-iter{i+1}.' ruamel_yaml = YAML() ruamel_yaml.default_flow_style = False From 83725568aedbeaf044f927ed18143af13347600c Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Mon, 1 Jun 2026 14:15:03 -0600 Subject: [PATCH 06/68] add suite_name to utilities/render_jedi_interface_files.py --- .../model/ensemble_latlon_prior_mean_output.py | 8 +++++++- .../model/ensemble_latlon_prior_variance_output.py | 8 +++++++- src/swell/utilities/render_jedi_interface_files.py | 1 + 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_mean_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_mean_output.py index 38c1005be..ac3bd27f2 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_mean_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_mean_output.py @@ -14,11 +14,17 @@ def ensemble_latlon_prior_mean_output(template_dict: Mapping) -> Mapping: + suite_name = template_dict.get('suite_name') + if suite_name == 'eda': + fn = 'eda.mean.' + else: + fn = 'geos.prior.mean.' + ensemble_latlon_prior_mean_output = { 'filetype': 'auxgrid', 'gridtype': 'latlon', 'datapath': template_dict['cycle_dir'], - 'filename': 'geos.prior.mean.', + 'filename': fn, 'field io names': field_io_names_ensemble } diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_variance_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_variance_output.py index 75f1f77d8..087c03065 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_variance_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_variance_output.py @@ -14,11 +14,17 @@ def ensemble_latlon_prior_variance_output(template_dict: Mapping) -> Mapping: + suite_name = template_dict.get('suite_name') + if suite_name == 'eda': + fn = 'eda.variance.' + else: + fn = 'geos.prior.variance.' + ensemble_latlon_prior_variance_output = { 'filetype': 'auxgrid', 'gridtype': 'latlon', 'datapath': template_dict['cycle_dir'], - 'filename': 'geos.prior.variance.', + 'filename': fn, 'field io names': field_io_names_ensemble } diff --git a/src/swell/utilities/render_jedi_interface_files.py b/src/swell/utilities/render_jedi_interface_files.py index 4db12b2c0..58ca94906 100644 --- a/src/swell/utilities/render_jedi_interface_files.py +++ b/src/swell/utilities/render_jedi_interface_files.py @@ -137,6 +137,7 @@ def __init__( 'window_end_iso', 'window_length', 'forecast_length', + 'suite_name', ] # List of all potential valid dynamic keys that can be used in templates From 46846b86440a482a2fc8059a8196831a611cb9fb Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 2 Jun 2026 11:11:19 -0600 Subject: [PATCH 07/68] swell/suites/eda/flow.cylc {% for i in range( 1, models[model_component]['ensemble_num_members'] + 1 ) %} {% for i in range(1, 5) %} it is only 50% working --- src/swell/suites/eda/flow.cylc | 22 ++++++++++++---------- src/swell/suites/eda/suite_config.py | 2 +- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/src/swell/suites/eda/flow.cylc b/src/swell/suites/eda/flow.cylc index 9f1e31357..0e22098ec 100644 --- a/src/swell/suites/eda/flow.cylc +++ b/src/swell/suites/eda/flow.cylc @@ -77,7 +77,8 @@ SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} - {% for i in range(1, 3) %} + + {% for i in range( 1, models[model_component]['ensemble_num_members'] + 1 ) %} SP2-{{model_component}} => RunJediEdaExecutable_mem{{i}}-{{model_component}} => SP3-{{model_component}} {% endfor %} @@ -156,7 +157,7 @@ [[RenderJediObservations-{{model_component}}]] script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" - {% for i in range(1, 3) %} + {% for i in range(1, 5) %} [[RunJediEdaExecutable_mem{{i}}-{{model_component}}]] script = "swell task RunJediEdaExecutable $config -d $datetime -m {{model_component}} -imem {{i}}" platform = {{platform}} @@ -167,14 +168,14 @@ {%- endfor %} {% endfor %} - [[RunJediEdaExecutable-{{model_component}}]] - script = "swell task RunJediEdaExecutable $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediEdaExecutable"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediEdaExecutable"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} +# [[RunJediEdaExecutable-{{model_component}}]] +# script = "swell task RunJediEdaExecutable $config -d $datetime -m {{model_component}}" +# platform = {{platform}} +# execution time limit = {{scheduling["RunJediEdaExecutable"]["execution_time_limit"]}} +# [[[directives]]] +# {%- for key, value in scheduling["RunJediEdaExecutable"]["directives"][model_component].items() %} +# --{{key}} = {{value}} +# {%- endfor %} [[RunJediObsfiltersExecutable-{{model_component}}]] script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" @@ -223,6 +224,7 @@ [[SP3-{{model_component}}]] script = true + {% endfor %} # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda/suite_config.py index 24c019999..4b1626409 100644 --- a/src/swell/suites/eda/suite_config.py +++ b/src/swell/suites/eda/suite_config.py @@ -48,7 +48,7 @@ class SuiteConfig(QuestionContainer, Enum): qd.gsibec_nlats("91"), qd.gsibec_nlons("144"), qd.vertical_resolution("72"), - qd.ensemble_num_members(2), + qd.ensemble_num_members(4), qd.obs_pert_amplitude(0.5), qd.analysis_variables([ "eastward_wind", From f3b6fb68eb5f6090b430b197e278566a59bfefe4 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 2 Jun 2026 23:54:29 -0400 Subject: [PATCH 08/68] Add: GetEnsembleGeosExperiment qd.geos_x_ensemble_directory qd.background_time_offset("PT3H") --- .../observations/sfcship.yaml.simple | 39 ------------------- .../observations/sondes.yaml.simple | 38 ------------------ src/swell/suites/eda/flow.cylc | 8 ++-- src/swell/suites/eda/suite_config.py | 8 ++-- 4 files changed, 7 insertions(+), 86 deletions(-) delete mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sfcship.yaml.simple delete mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sondes.yaml.simple diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sfcship.yaml.simple b/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sfcship.yaml.simple deleted file mode 100644 index 3b2854599..000000000 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sfcship.yaml.simple +++ /dev/null @@ -1,39 +0,0 @@ -obs space: - name: Surface Marine Stations - obsdatain: - engine: - type: H5File - obsfile: '{{cycle_dir}}/sfcship.{{window_begin}}.nc4' - missing file action: warn - obsdataout: - engine: - type: H5File - obsfile: '{{cycle_dir}}/{{experiment_id}}.sfcship.{{window_begin}}.nc4' - io pool: - max pool size: 6 - simulated variables: [stationPressure] - -obs operator: - name: Composite - components: - - # Surface pressure is handled using the correction scheme - - name: SfcCorrected - variables: - - name: stationPressure - correction scheme to use: GSL - geovar_geomz: geopotential_height - geovar_sfc_geomz: height_above_mean_sea_level_at_surface - station_altitude: height - -linear obs operator: - name: Composite - components: - - # Specific humidity and station pressure is using the lowest model level (identity) - - name: Identity - variables: - - name: stationPressure - -obs error: - covariance model: diagonal diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sondes.yaml.simple b/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sondes.yaml.simple deleted file mode 100644 index 49d95e91c..000000000 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/observations/sondes.yaml.simple +++ /dev/null @@ -1,38 +0,0 @@ -obs space: - name: Radiosondes - obsdatain: - engine: - type: H5File - obsfile: '{{cycle_dir}}/sondes.{{window_begin}}.nc4' - missing file action: error - obsgrouping: - group variables: ["stationIdentification", "releaseTime"] - sort variable: "pressure" - sort order: "descending" - obsdataout: - engine: - type: H5File - obsfile: '{{cycle_dir}}/{{experiment_id}}.sondes.{{window_begin}}.nc4' - io pool: - max pool size: 6 - simulated variables: [airTemperature] - - -obs operator: - name: Composite - components: - - - name: VertInterp - variables: - - name: airTemperature - -linear obs operator: - name: Composite - components: - - - name: VertInterp - variables: - - name: airTemperature - -obs error: - covariance model: diagonal diff --git a/src/swell/suites/eda/flow.cylc b/src/swell/suites/eda/flow.cylc index 0e22098ec..e1f107eda 100644 --- a/src/swell/suites/eda/flow.cylc +++ b/src/swell/suites/eda/flow.cylc @@ -51,9 +51,6 @@ # Generate satellite channel records CloneGeosMksi-{{model_component}}[^] => GenerateObservingSystemRecords-{{model_component}} - # Get background, provide a way to get background directly from GEOS experiment - GetBackgroundGeosExperiment-{{model_component}} :fail? => GetBackground-{{model_component}} - # Get observations {% if cycling_varbc %} # Cycling VarBC is active, biases from the previous cycle will be used @@ -72,7 +69,7 @@ BuildJediByLinking[^]? | BuildJedi[^] => SP1-{{model_component}} CloneJedi[^] => StageJediCycle-{{model_component}} StageJediCycle-{{model_component}} => SP1-{{model_component}} - GetBackgroundGeosExperiment-{{model_component}}? | GetBackground-{{model_component}} => SP1-{{model_component}} + GetEnsembleGeosExperiment-{{model_component}} => SP1-{{model_component}} GetObsNotInR2d2-{{model_component}}? | GetObservations-{{model_component}} => SP1-{{model_component}} SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} @@ -151,6 +148,9 @@ [[GetBackgroundGeosExperiment-{{model_component}} ]] script = "swell task GetBackgroundGeosExperiment $config -d $datetime -m {{model_component}}" + [[GetEnsembleGeosExperiment-{{model_component}}]] + script = "swell task GetEnsembleGeosExperiment $config -d $datetime -m {{model_component}}" + [[GetObservations-{{model_component}}]] script = "swell task GetObservations $config -d $datetime -m {{model_component}}" diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda/suite_config.py index 4b1626409..3fbf679be 100644 --- a/src/swell/suites/eda/suite_config.py +++ b/src/swell/suites/eda/suite_config.py @@ -30,19 +30,17 @@ class SuiteConfig(QuestionContainer, Enum): qd.runahead_limit("P1"), qd.jedi_build_method("use_existing"), qd.model_components(['geos_atmosphere']), - ## - ## qd.ensemble_num_members(2), ], geos_atmosphere=[ qd.cycle_times([ "T00", - "T06", - "T12", - "T18" ]), qd.geos_x_background_directory("/discover/nobackup/projects/gmao/" "dadev/rtodling/archive/Restarts/JEDI/541x"), + qd.geos_x_ensemble_directory('/discover/nobackup/projects/gmao/dadev/' + 'rtodling/archive/541/Milan'), qd.window_length("PT6H"), + qd.background_time_offset("PT3H"), qd.window_type("3D"), qd.horizontal_resolution("91"), qd.gsibec_nlats("91"), From 097276228fa06edb444c217dc6898e9a9652cfc8 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 4 Jun 2026 13:33:50 -0600 Subject: [PATCH 09/68] add `ensmeanvariance_spec` to question defaults --- .../geos_atmosphere/model/analysis.py | 5 - ...ior_mean_output.py => comp_mean_output.py} | 33 +++--- ...ance_output.py => comp_variance_output.py} | 33 +++--- .../geos_atmosphere/model/eda_analysis.py | 30 +++++ .../model/eda_varincrement1.py | 47 ++++++++ .../geos_atmosphere/model/ensemble_block.py | 10 +- src/swell/configuration/jedi/oops/eda3D.py | 4 +- .../jedi/oops/ensmeanvariance.py | 4 +- src/swell/suites/eda/suite_config.py | 3 +- src/swell/tasks/run_jedi_eda_executable.py | 27 +++-- .../tasks/run_jedi_ensemble_mean_variance.py | 110 +++++++++++------- src/swell/tasks/task_questions.py | 1 + src/swell/utilities/question_defaults.py | 26 +++++ .../utilities/render_jedi_interface_files.py | 1 + 14 files changed, 233 insertions(+), 101 deletions(-) rename src/swell/configuration/jedi/interfaces/geos_atmosphere/model/{ensemble_latlon_prior_mean_output.py => comp_mean_output.py} (58%) rename src/swell/configuration/jedi/interfaces/geos_atmosphere/model/{ensemble_latlon_prior_variance_output.py => comp_variance_output.py} (57%) create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_varincrement1.py diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/analysis.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/analysis.py index dfa4e8ae5..ca7a53983 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/analysis.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/analysis.py @@ -23,11 +23,6 @@ def analysis(template_dict: Mapping) -> Mapping: 'field io names': field_io_names, } - # EDA case: override filename - imem = template_dict.get('ensemble_imember', None) - if imem is not None: - analysis['filename'] = f'eda.mem{imem:03d}.analysis.%yyyy%mm%dd_%hh%MM%ssz.nc4' - return analysis # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_mean_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py similarity index 58% rename from src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_mean_output.py rename to src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py index ac3bd27f2..8c408fd67 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_mean_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py @@ -12,23 +12,22 @@ # -------------------------------------------------------------------------------------------------- -def ensemble_latlon_prior_mean_output(template_dict: Mapping) -> Mapping: - - suite_name = template_dict.get('suite_name') - if suite_name == 'eda': - fn = 'eda.mean.' - else: - fn = 'geos.prior.mean.' - - ensemble_latlon_prior_mean_output = { - 'filetype': 'auxgrid', - 'gridtype': 'latlon', - 'datapath': template_dict['cycle_dir'], - 'filename': fn, - 'field io names': field_io_names_ensemble - } - - return ensemble_latlon_prior_mean_output +def comp_mean_output(template_dict: Mapping) -> Mapping: + + fn_output_mean = template_dict['ensmeanvariance_spec_item'].get('fn_output_mean') + geom = template_dict['ensmeanvariance_spec_item'].get('geom') + + mean_output = {} + if geom == 'latlon': + mean_output = { + 'filetype': 'auxgrid', + 'gridtype': 'latlon', + 'datapath': template_dict['cycle_dir'], + 'filename': fn_output_mean, + 'field io names': field_io_names_ensemble + } + + return ensemble_mean_output # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_variance_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py similarity index 57% rename from src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_variance_output.py rename to src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py index 087c03065..d9fc6a924 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_latlon_prior_variance_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py @@ -12,23 +12,22 @@ # -------------------------------------------------------------------------------------------------- -def ensemble_latlon_prior_variance_output(template_dict: Mapping) -> Mapping: - - suite_name = template_dict.get('suite_name') - if suite_name == 'eda': - fn = 'eda.variance.' - else: - fn = 'geos.prior.variance.' - - ensemble_latlon_prior_variance_output = { - 'filetype': 'auxgrid', - 'gridtype': 'latlon', - 'datapath': template_dict['cycle_dir'], - 'filename': fn, - 'field io names': field_io_names_ensemble - } - - return ensemble_latlon_prior_variance_output +def ensemble_variance_output(template_dict: Mapping) -> Mapping: + + fn_output_variance = template_dict['ensmeanvariance_spec_item'].get('fn_output_variance') + geom = template_dict['ensmeanvariance_spec_item'].get('geom') + + variance_output = {} + if geom == 'latlon': + variance_output = { + 'filetype': 'auxgrid', + 'gridtype': 'latlon', + 'datapath': template_dict['cycle_dir'], + 'filename': fn_output_variance, + 'field io names': field_io_names_ensemble + } + + return variance_output # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py new file mode 100644 index 000000000..f40195457 --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py @@ -0,0 +1,30 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import field_io_names +import os +# -------------------------------------------------------------------------------------------------- + + +def eda_analysis(template_dict: Mapping) -> Mapping: + + imem = template_dict.get('ensemble_imember', None) + analysis = { + 'filetype': 'cube sphere history', + 'provider': 'geos', + 'datapath': f'./analysis/mem{imem:03d}', + 'filename': f'{template_dict["experiment_id"]}.ana.mem{imem:03d}.%yyyy%mm%dd_%hh%MM%ssz.nc4', + 'first': 'PT0H', + 'frequency': 'PT1H', + 'field io names': field_io_names, + } + + return analysis + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_varincrement1.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_varincrement1.py new file mode 100644 index 000000000..df885008d --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_varincrement1.py @@ -0,0 +1,47 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping +import os +# -------------------------------------------------------------------------------------------------- + + +def eda_varincrement1(template_dict: Mapping) -> Mapping: + + experiment_id = template_dict['experiment_id'] + imem = template_dict.get('ensemble_imember', None) + varincrement1 = { + 'write increment': True, + 'increment': { + 'state component': { + 'filetype': 'auxgrid', + 'gridtype': 'latlon', + 'datapath': '.', + 'filename': f'analysis/mem{imem:03d}/{experiment_id}.inc_iter1.mem{imem:03d}.', + 'field io names': { + 'eastward_wind': 'ua', + 'northward_wind': 'va', + 'air_temperature': 't', + 'air_pressure_levels': 'pe', + 'water_vapor_mixing_ratio_wrt_moist_air': 'q', + 'cloud_liquid_ice': 'qi', + 'cloud_liquid_water': 'ql', + 'rain_water': 'qr', + 'snow_water': 'qs', + 'mole_fraction_of_ozone_in_air': 'o3ppmv', + 'geopotential_height_times_gravity_at_surface': 'phis', + 'skin_temperature_at_surface': 'ts', + 'air_pressure_at_surface': 'ps', + } + } + } + } + + return varincrement1 + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_block.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_block.py index 5bf6d6b33..96980149b 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_block.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_block.py @@ -13,6 +13,8 @@ def ensemble_block(template_dict: Mapping) -> Mapping: + fn_input = template_dict['ensmeanvariance_spec_item'].get('fn_input') + ensemble_block = { 'members from template': { 'template': { @@ -22,7 +24,7 @@ def ensemble_block(template_dict: Mapping) -> Mapping: 'compute edge pressure from surface pressure': True, 'max allowable geometry difference': 1e-3, 'datapath': template_dict['cycle_dir'], - 'filename': 'ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4', + 'filename': fn_input, 'state variables': [ 'eastward_wind', 'northward_wind', @@ -77,12 +79,6 @@ def ensemble_block(template_dict: Mapping) -> Mapping: } } - imem = template_dict.get('ensemble_imem', None) - if imem is not None: - # EDA diagnostic case - fn = ensemble_block['members from template']['template']['filename'] - fn = f'eda.mem{imem:03d}.analysis.%yyyy%mm%dd_%hh%MM%ssz.nc4' - return ensemble_block # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/oops/eda3D.py b/src/swell/configuration/jedi/oops/eda3D.py index 4f53b7b56..54da555c3 100644 --- a/src/swell/configuration/jedi/oops/eda3D.py +++ b/src/swell/configuration/jedi/oops/eda3D.py @@ -44,7 +44,7 @@ def render_oops(self): 'diagnostics': { 'departures': 'ombg' }, - 'online diagnostics': self.interface_model('varincrement1') + 'online diagnostics': self.interface_model('eda_varincrement1') }], }, 'final': { @@ -55,7 +55,7 @@ def render_oops(self): 'frequency': 'PT3H' } }, - 'output': self.interface_model('analysis') + 'output': self.interface_model('eda_analysis') } # TODO: Implement this more cleanly in the OOPS schema diff --git a/src/swell/configuration/jedi/oops/ensmeanvariance.py b/src/swell/configuration/jedi/oops/ensmeanvariance.py index 5e12051a4..ca87acc91 100644 --- a/src/swell/configuration/jedi/oops/ensmeanvariance.py +++ b/src/swell/configuration/jedi/oops/ensmeanvariance.py @@ -18,8 +18,8 @@ def render_oops(self): oops = { 'geometry': self.interface_model('geometry'), 'ensemble': self.interface_model('ensemble_block'), - 'variance output': self.interface_model('ensemble_latlon_prior_variance_output'), - 'mean output': self.interface_model('ensemble_latlon_prior_mean_output') + 'variance output': self.interface_model('comp_variance_output'), + 'mean output': self.interface_model('comp_mean_output') } return oops diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda/suite_config.py index 3fbf679be..4404fe805 100644 --- a/src/swell/suites/eda/suite_config.py +++ b/src/swell/suites/eda/suite_config.py @@ -27,7 +27,7 @@ class SuiteConfig(QuestionContainer, Enum): sq.common, qd.start_cycle_point("2023-10-10T00:00:00Z"), qd.final_cycle_point("2023-10-10T06:00:00Z"), - qd.runahead_limit("P1"), + qd.runahead_limit("P2"), qd.jedi_build_method("use_existing"), qd.model_components(['geos_atmosphere']), ], @@ -40,7 +40,6 @@ class SuiteConfig(QuestionContainer, Enum): qd.geos_x_ensemble_directory('/discover/nobackup/projects/gmao/dadev/' 'rtodling/archive/541/Milan'), qd.window_length("PT6H"), - qd.background_time_offset("PT3H"), qd.window_type("3D"), qd.horizontal_resolution("91"), qd.gsibec_nlats("91"), diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 916a4b195..6c4931d07 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -135,7 +135,6 @@ def execute(self) -> None: # imember: specify yaml # ---------------------------------------------------- - if imember == 1: jedi_config_dict['cost function']['observations'].pop('obs perturbations', None) else: @@ -145,16 +144,30 @@ def execute(self) -> None: # Get observation name observation = observer['observation_name'] if imember > 1: - observer['obs space'].update({'obs perturbations seed': imember}) - observer['obs error'].update({'obs perturbations amplitude': 0.5}) + obs_cov_model = observer.get('obs error', {}).get('covariance model') + print( f'{observation}: obs_cov_model = {obs_cov_model}') + if obs_cov_model and 'cross variable covariances' in obs_cov_model: + print(f"Found cross covariance obs: {obs_cov_model}, skip perturbation") + else: + print(f"No cross varaible covariance found for {observation}, Obs Error Diagonal") + obs_error_dict = { + 'covariance model': 'diagonal', + 'zero-mean perturbations': True, + 'member': imember, + 'number of members': nmember + } + observer.update({'obs error': obs_error_dict}) + observer['obs space'].update({'obs perturbations seed': imember}) + hxout = observer['obs space']['obsdataout']['engine']['obsfile'] dir1, fname = os.path.split(hxout) - hxout = os.path.join(dir1, f'diag/mem{imember:003d}', fname) + hxout = os.path.join(dir1, f'analysis/mem{imember:003d}', fname) observer['obs space']['obsdataout']['engine']['obsfile'] = hxout - print (f'hxout = {hxout}') + # print (f'hxout = {hxout}') - for i, iter in enumerate(jedi_config_dict['variational']['iterations']): - iter['online diagnostics']['increment']['state component']['filename'] = f'eda.mem{imember:03d}.increment-iter{i+1}.' + # create subdir + xdir = os.path.join(self.cycle_dir(), f'analysis/mem{imember:003d}') + os.makedirs(xdir, exist_ok=True) ruamel_yaml = YAML() ruamel_yaml.default_flow_style = False diff --git a/src/swell/tasks/run_jedi_ensemble_mean_variance.py b/src/swell/tasks/run_jedi_ensemble_mean_variance.py index 7d98c327c..a514998ab 100644 --- a/src/swell/tasks/run_jedi_ensemble_mean_variance.py +++ b/src/swell/tasks/run_jedi_ensemble_mean_variance.py @@ -51,6 +51,7 @@ def execute(self) -> None: # Ensemble # ------------------------ + # Populate jedi interface templates dictionary # -------------------------------------------- self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) @@ -71,6 +72,24 @@ def execute(self) -> None: # Ensemble self.jedi_rendering.add_key('ensemble_num_members', self.config.ensemble_num_members(None)) + + print( f'self.config.ensmeanvariance_spec = {self.config.ensmeanvariance_spec()}') + meanvar_spec_dict = self.config.ensmeanvariance_spec() +# # Loop directly over the items in the list +# for spec in config_dict: +# # Unpack variables directly from the current dictionary item +# fn_input = spec["fn_input"] +# fn_output_mean = spec["fn_output_mean"] +# fn_output_variance = spec["fn_output_variance"] +# +# # Use the variables directly +# print(f"Input Path: {fn_input}") +# print(f"Output Mean: {fn_output_mean}") +# print(f"Output Var: {fn_output_variance}") +# print("-" * 40) +# + + # Add placeholder names if mock experiment # ---------------------------------------- if self.config.mock_experiment(False): @@ -78,49 +97,56 @@ def execute(self) -> None: self.jedi_rendering.add_key('experiment_id', 'experiment_id') self.jedi_rendering.add_key('cycle_dir', 'cycle_dir') - # Jedi configuration file - # ----------------------- - jedi_config_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_config.yaml') - - # Output log file - # --------------- - output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}.log') - - # Open the JEDI config file and fill templates - # -------------------------------------------- - jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}', - window_type, - jedi_forecast_model) - - yaml = YAML() - yaml.default_flow_style = False - - # Write the expanded dictionary to YAML file - # ------------------------------------------ - with open(jedi_config_file, 'w') as jedi_config_file_open: - yaml.dump(jedi_config_dict, jedi_config_file_open) - - # Get the JEDI interface metadata - # ------------------------------- - model_component_meta = self.jedi_rendering.render_interface_meta() - - # Compute number of processors - # ---------------------------- - np = eval(str(model_component_meta['total_processors'])) - - # Jedi executable name - # -------------------- - jedi_executable = model_component_meta['executables'][f'{jedi_application}'] - jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', 'build', 'bin', - jedi_executable) - # Run the JEDI executable + # loop item in ensmeanvariance_spec # ----------------------- - if not generate_yaml_and_exit: - self.logger.info('Running '+jedi_executable_path+' with '+str(np)+' processors.') - run_executable(self.logger, self.cycle_dir(), np, jedi_executable_path, - jedi_config_file, output_log_file) - else: - self.logger.info('YAML generated, now exiting.') + for idx, spec in enumerate(meanvar_spec_dict, start=1): + + self.jedi_rendering.add_key('ensmeanvariance_spec_item', spec) + + # Jedi configuration file + # ----------------------- + jedi_config_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_config_{idx}.yaml') + + # Output log file + # --------------- + output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_{idx}.log') + + # Open the JEDI config file and fill templates + # -------------------------------------------- + jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}', + window_type, + jedi_forecast_model) + + yaml = YAML() + yaml.default_flow_style = False + + # Write the expanded dictionary to YAML file + # ------------------------------------------ + with open(jedi_config_file, 'w') as jedi_config_file_open: + yaml.dump(jedi_config_dict, jedi_config_file_open) + + # Get the JEDI interface metadata + # ------------------------------- + model_component_meta = self.jedi_rendering.render_interface_meta() + + # Compute number of processors + # ---------------------------- + np = eval(str(model_component_meta['total_processors'])) + + # Jedi executable name + # -------------------- + jedi_executable = model_component_meta['executables'][f'{jedi_application}'] + jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', 'build', 'bin', + jedi_executable) + + # Run the JEDI executable + # ----------------------- + if not generate_yaml_and_exit: + self.logger.info('Running '+jedi_executable_path+' with '+str(np)+' processors.') + run_executable(self.logger, self.cycle_dir(), np, jedi_executable_path, + jedi_config_file, output_log_file) + else: + self.logger.info('YAML generated, now exiting.') # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index e80bd5a7f..de923990d 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -638,6 +638,7 @@ class TaskQuestions(QuestionContainer, Enum): window_questions, qd.analysis_variables(), qd.ensemble_num_members(), + qd.ensmeanvariance_spec(), qd.generate_yaml_and_exit(), qd.jedi_forecast_model(), qd.observations(), diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index 2baa5e1bb..73e9ab1f6 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -536,6 +536,32 @@ class ensmeanvariance_only(TaskQuestion): # -------------------------------------------------------------------------------------------------- + @dataclass + class ensmeanvariance_spec(TaskQuestion): + default_value: List[Dict[str, str]] = mutable_field([ + { + "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "geos.prior.mean", + "fn_output_variance": "geos.prior.variance", + "geom": "latlon" + }, + { + "fn_input": "analysis/mem%mem%/eda.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "eda.ana.mean", + "fn_output_variance": "eda.ana.variance", + "geom": "latlon" + }, + ]) + question_name: str = "ensmeanvariance_spec" + models: List[str] = mutable_field([ + "all_models" + ]) + ask_question: bool = True + prompt: str = "Configure the ensemble mean and variance specifications: [e.g. ebkg, eda.analysis, eda.inc, etc.]" + widget_type: WType = WType.STRING + + # -------------------------------------------------------------------------------------------------- + @dataclass class existing_geos_gcm_build_path(TaskQuestion): default_value: str = "defer_to_platform" diff --git a/src/swell/utilities/render_jedi_interface_files.py b/src/swell/utilities/render_jedi_interface_files.py index 58ca94906..6532913cb 100644 --- a/src/swell/utilities/render_jedi_interface_files.py +++ b/src/swell/utilities/render_jedi_interface_files.py @@ -84,6 +84,7 @@ def __init__( 'ensemble_imember', 'ensmean_only', 'ensmeanvariance_only', + 'ensmeanvariance_spec_item', 'experiment_id', 'experiment_root', 'final_cycle_point', From 02eb8bd53e32cd605e7618e01b930f82045c1647 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Mon, 8 Jun 2026 11:32:45 -0600 Subject: [PATCH 10/68] I find there are problems with using ensemble_block.py two_states.py state_ensemble.py IP. --- .../geos_atmosphere/geos_atmosphere.yaml | 1 + .../geos_atmosphere/model/comp_mean_output.py | 34 +++-- .../model/comp_variance_output.py | 29 +++- .../model/diffstates_output.py | 42 ++++++ .../geos_atmosphere/model/geometry_wofms.py | 32 +++++ .../geos_atmosphere/model/state_ensemble.py | 42 ++++++ .../geos_atmosphere/model/two_states.py | 45 ++++++ .../configuration/jedi/oops/diffstates.py | 29 ++++ .../jedi/oops/ensmeanvariance.py | 2 +- src/swell/suites/eda/flow.cylc | 23 ++- src/swell/suites/eda/suite_config.py | 20 +++ src/swell/tasks/run_jedi_diffstates.py | 136 ++++++++++++++++++ .../tasks/run_jedi_ensemble_mean_variance.py | 97 +++++++------ src/swell/tasks/task_questions.py | 16 +++ src/swell/utilities/question_defaults.py | 34 ++--- .../utilities/render_jedi_interface_files.py | 3 + src/swell/utilities/slurm.py | 1 + 17 files changed, 504 insertions(+), 82 deletions(-) create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/geometry_wofms.py create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/state_ensemble.py create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py create mode 100644 src/swell/configuration/jedi/oops/diffstates.py create mode 100644 src/swell/tasks/run_jedi_diffstates.py diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml b/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml index ac17a05fc..763690911 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml @@ -13,3 +13,4 @@ executables: obsfilters: test_ObsFilters.x eda3D: fv3jedi_var.x eda4D: fv3jedi_var.x + diffstates: fv3jedi_diffstates.x diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py index 8c408fd67..5ecc0d812 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py @@ -7,27 +7,45 @@ # -------------------------------------------------------------------------------------------------- from collections.abc import Mapping -from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import field_io_names_ensemble +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ + field_io_names, field_io_names_ensemble # -------------------------------------------------------------------------------------------------- def comp_mean_output(template_dict: Mapping) -> Mapping: - fn_output_mean = template_dict['ensmeanvariance_spec_item'].get('fn_output_mean') - geom = template_dict['ensmeanvariance_spec_item'].get('geom') + prefix_output_mean = template_dict['ensmeanvariance_spec_item'].get('fn_output_mean') + state = template_dict['ensmeanvariance_spec_item'].get('state') + grid_type= template_dict['ensmeanvariance_spec_gridtype'] mean_output = {} - if geom == 'latlon': + if grid_type == 'cs': + mean_output = { + 'filetype': 'cube sphere history', + 'provider': 'geos', + 'datapath': template_dict['cycle_dir'], + 'filename': f'{prefix_output_mean}.%yyyy%mm%dd_%hh%MM%ssz.nc4', + 'first': 'PT0H', + 'frequency': 'PT1H', + 'field io names': [] + } + elif grid_type == 'latlon': mean_output = { 'filetype': 'auxgrid', 'gridtype': 'latlon', 'datapath': template_dict['cycle_dir'], - 'filename': fn_output_mean, - 'field io names': field_io_names_ensemble - } + 'filename': f'{prefix_output_mean}.ll.', + 'field io names': [] + } + + if state in ['bkg', 'analysis']: + mean_output['field io names'] = field_io_names + else: + mean_output['field io names'] = field_io_names_ensemble - return ensemble_mean_output + + return mean_output # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py index d9fc6a924..af0d471ad 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py @@ -7,26 +7,41 @@ # -------------------------------------------------------------------------------------------------- from collections.abc import Mapping -from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import field_io_names_ensemble +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ + field_io_names, field_io_names_ensemble # -------------------------------------------------------------------------------------------------- -def ensemble_variance_output(template_dict: Mapping) -> Mapping: +def comp_variance_output(template_dict: Mapping) -> Mapping: - fn_output_variance = template_dict['ensmeanvariance_spec_item'].get('fn_output_variance') - geom = template_dict['ensmeanvariance_spec_item'].get('geom') + prefix_output_variance = template_dict['ensmeanvariance_spec_item'].get('fn_output_variance') + state = template_dict['ensmeanvariance_spec_item'].get('state') + grid_type= template_dict['ensmeanvariance_spec_gridtype'] variance_output = {} - if geom == 'latlon': + if grid_type == 'cs': + variance_output = { + 'filetype': 'cube sphere history', + 'provider': 'geos', + 'datapath': template_dict['cycle_dir'], + 'filename': f'{prefix_output_variance}.%yyyy%mm%dd_%hh%MM%ssz.nc4', + 'field io names': [] + } + elif grid_type == 'latlon': variance_output = { 'filetype': 'auxgrid', 'gridtype': 'latlon', 'datapath': template_dict['cycle_dir'], - 'filename': fn_output_variance, - 'field io names': field_io_names_ensemble + 'filename': f'{prefix_output_variance}.ll.', + 'field io names': [] } + if state in ['bkg', 'analysis']: + variance_output['field io names'] = field_io_names + else: + variance_output['field io names'] = field_io_names_ensemble + return variance_output diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py new file mode 100644 index 000000000..db3c2e789 --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py @@ -0,0 +1,42 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ + field_io_names, field_io_names_ensemble + +# -------------------------------------------------------------------------------------------------- + + +def diffstates_output(template_dict: Mapping) -> Mapping: + + # diff states (inc) + grid_type= template_dict['diffstates_spec_gridtype'] + prefix_output = template_dict['diffstates_spec'].get('state_diff', {}).get('fn_output') + + output = {} + if grid_type == 'cs': + output = { + 'filetype': 'cube sphere history', + 'provider': 'geos', + 'datapath': template_dict['cycle_dir'], + 'filename': f'{prefix_output}.%yyyy%mm%dd_%hh%MM%ssz.nc4', + 'field io names': field_io_names_ensemble + } + elif grid_type == 'latlon': + output = { + 'filetype': 'auxgrid', + 'gridtype': 'latlon', + 'datapath': template_dict['cycle_dir'], + 'filename': f'{prefix_output}.ll.', + 'field io names': field_io_names_ensemble + } + + return output + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/geometry_wofms.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/geometry_wofms.py new file mode 100644 index 000000000..8c2c8bd2b --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/geometry_wofms.py @@ -0,0 +1,32 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping + +# -------------------------------------------------------------------------------------------------- + + +def geometry_wofms(template_dict: Mapping) -> Mapping: + + npx_proc = template_dict['npx_proc'] + npy_proc = template_dict['npy_proc'] + + horizontal_resolution = template_dict['horizontal_resolution'] + vertical_resolution = template_dict['vertical_resolution'] + + geometry = { + 'akbk': f'./fv3-jedi/fv3files/akbk{vertical_resolution}.nc4', + 'layout': [npx_proc, npy_proc], + 'npx': horizontal_resolution, + 'npy': horizontal_resolution, + 'npz': vertical_resolution + } + + return geometry + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/state_ensemble.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/state_ensemble.py new file mode 100644 index 000000000..608dd51f2 --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/state_ensemble.py @@ -0,0 +1,42 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ + state_variables, field_io_names + +# -------------------------------------------------------------------------------------------------- + + +def state_ensemble(template_dict: Mapping) -> Mapping: + + # full fname including subdir path + fn_input = template_dict['ensmeanvariance_spec_item'].get('fn_input') + + state_ensemble_dict = { + 'members from template': { + 'template': { + 'datetime': template_dict['local_background_time_iso'], + 'filetype': 'cube sphere history', + 'provider': 'geos', + 'compute edge pressure from surface pressure': True, + 'max allowable geometry difference': 1e-3, + 'datapath': template_dict['cycle_dir'], + 'filename': fn_input, + 'state variables': state_variables, + 'field io names': field_io_names + }, + 'pattern': '%mem%', + 'nmembers': template_dict['ensemble_num_members'], + 'zero padding': 3 + } + } + + return state_ensemble_dict + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py new file mode 100644 index 000000000..b116e40e6 --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py @@ -0,0 +1,45 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ + state_variables, field_io_names + +# -------------------------------------------------------------------------------------------------- + +def two_states(template_dict: Mapping) -> Mapping: + horizontal_resolution = template_dict['horizontal_resolution'] + f1 = template_dict['diffstates_spec']['state1']['fn_input'] + f2 = template_dict['diffstates_spec']['state2']['fn_input'] + + # state1: bkg + state1 = { + 'datetime': template_dict['local_background_time_iso'], + 'filetype': 'cube sphere history', + 'provider': 'geos', + 'datapath': template_dict['cycle_dir'], + 'filename': f1, + 'state variables': state_variables, + 'field io names': field_io_names, + } + + # state2: analysis + state2 = { + 'datetime': template_dict['local_background_time_iso'], + 'filetype': 'cube sphere history', + 'provider': 'geos', + 'datapath': template_dict['cycle_dir'], + 'filenames': f2, + 'state variables': state_variables, + 'field io names': field_io_names, + } + + combo = {'state1': state1, 'state2': state2} + return combo + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/oops/diffstates.py b/src/swell/configuration/jedi/oops/diffstates.py new file mode 100644 index 000000000..22efc092b --- /dev/null +++ b/src/swell/configuration/jedi/oops/diffstates.py @@ -0,0 +1,29 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from swell.utilities.oops_config import OopsConfig + +# -------------------------------------------------------------------------------------------------- + + +class diffstates(OopsConfig): + + def render_oops(self): + + two_states = self.interface_model('two_states') + oops = { + 'state geometry': self.interface_model('geometry'), + 'increment geometry': self.interface_model('geometry_wofms'), + **two_states, + 'output': self.interface_model('diffstates_output') + } + + return oops + + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/oops/ensmeanvariance.py b/src/swell/configuration/jedi/oops/ensmeanvariance.py index ca87acc91..1c51acd12 100644 --- a/src/swell/configuration/jedi/oops/ensmeanvariance.py +++ b/src/swell/configuration/jedi/oops/ensmeanvariance.py @@ -17,7 +17,7 @@ def render_oops(self): oops = { 'geometry': self.interface_model('geometry'), - 'ensemble': self.interface_model('ensemble_block'), + 'ensemble': self.interface_model('state_ensemble'), 'variance output': self.interface_model('comp_variance_output'), 'mean output': self.interface_model('comp_mean_output') } diff --git a/src/swell/suites/eda/flow.cylc b/src/swell/suites/eda/flow.cylc index e1f107eda..88052268a 100644 --- a/src/swell/suites/eda/flow.cylc +++ b/src/swell/suites/eda/flow.cylc @@ -79,14 +79,18 @@ SP2-{{model_component}} => RunJediEdaExecutable_mem{{i}}-{{model_component}} => SP3-{{model_component}} {% endfor %} + # bkg/ana : mean/variance + SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} + RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} + # EvaObservations - SP3-{{model_component}} => EvaObservations-{{model_component}} + SP4-{{model_component}} => EvaObservations-{{model_component}} # EvaJediLog - SP3-{{model_component}} => EvaJediLog-{{model_component}} + SP4-{{model_component}} => EvaJediLog-{{model_component}} # EvaIncrement - SP3-{{model_component}} => EvaIncrement-{{model_component}} + SP4-{{model_component}} => EvaIncrement-{{model_component}} # Clean up large files EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} & @@ -195,6 +199,15 @@ --{{key}} = {{value}} {%- endfor %} + [[RunJediDiffstates-{{model_component}}]] + script = "swell task RunJediDiffstates $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediDiffstates"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediDiffstates"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + [[EvaJediLog-{{model_component}}]] script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}" @@ -216,6 +229,7 @@ [[CleanCycle-{{model_component}}]] script = "swell task CleanCycle $config -d $datetime -m {{model_component}}" + # SP: Sync Point [[SP1-{{model_component}}]] script = true @@ -225,6 +239,9 @@ [[SP3-{{model_component}}]] script = true + [[SP4-{{model_component}}]] + script = true + {% endfor %} # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda/suite_config.py index 4404fe805..b6a060918 100644 --- a/src/swell/suites/eda/suite_config.py +++ b/src/swell/suites/eda/suite_config.py @@ -102,6 +102,26 @@ class SuiteConfig(QuestionContainer, Enum): "ssmis_f17" ]), qd.obs_thinning_rej_fraction(0.75), + qd.ensmeanvariance_spec([ + {"state": "bkg", + "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "geos.prior.mean", + "fn_output_variance": "geos.prior.variance", + "grid_type": ['cs', 'latlon']}, + {"state": "analysis", + "fn_input": "analysis/mem%mem%/eda.ana.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "eda.ana.mean", + "fn_output_variance": "eda.ana.variance", + "grid_type": ['cs', 'latlon']}, + ]), + qd.diffstates_spec({ + "state1": + {"fn_input": "geos.prior.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + "state2": + {"fn_input": "eda.ana.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + "state_diff": + {"fn_output": "eda.inc.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + }), qd.clean_patterns(['*.txt', '*.csv']), ] ) diff --git a/src/swell/tasks/run_jedi_diffstates.py b/src/swell/tasks/run_jedi_diffstates.py new file mode 100644 index 000000000..2bb115d80 --- /dev/null +++ b/src/swell/tasks/run_jedi_diffstates.py @@ -0,0 +1,136 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + + +# -------------------------------------------------------------------------------------------------- + + +import os +from ruamel.yaml import YAML + +from swell.tasks.base.task_base import taskBase +from swell.utilities.run_jedi_executables import run_executable + + +# -------------------------------------------------------------------------------------------------- + + +class RunJediDiffstates(taskBase): + + # ---------------------------------------------------------------------------------------------- + + def execute(self) -> None: + + # Jedi application name + # --------------------- + jedi_application = 'diffstates' + + # Parse configuration + # ------------------- + jedi_forecast_model = self.config.jedi_forecast_model(None) + generate_yaml_and_exit = self.config.generate_yaml_and_exit(False) + npx_proc = self.config.npx_proc(None) + npy_proc = self.config.npy_proc(None) + + # Compute data assimilation window parameters + window_type = self.config.window_type() + window_length = self.config.window_length() + local_background_time = self.da_window_params.local_background_time(window_length, + window_type) + local_background_time_iso = self.da_window_params.local_background_time_iso(window_length, + window_type) + window_begin_iso = self.da_window_params.window_begin_iso(window_length) + window_end_iso = self.da_window_params.window_end_iso(window_length) + + # Set the observing system records path + self.jedi_rendering.set_obs_records_path(self.config.observing_system_records_path(None)) + + + # Populate jedi interface templates dictionary + # -------------------------------------------- + self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) + self.jedi_rendering.add_key('window_end_iso', window_end_iso) + self.jedi_rendering.add_key('window_length', window_length) + self.jedi_rendering.add_key('analysis_variables', self.config.analysis_variables()) + + # Background + self.jedi_rendering.add_key('horizontal_resolution', self.config.horizontal_resolution()) + self.jedi_rendering.add_key('local_background_time', local_background_time) + self.jedi_rendering.add_key('local_background_time_iso', local_background_time_iso) + + # Geometry + self.jedi_rendering.add_key('vertical_resolution', self.config.vertical_resolution()) + self.jedi_rendering.add_key('npx_proc', npx_proc) + self.jedi_rendering.add_key('npy_proc', npy_proc) + + # Diffstates + spec_dict = self.config.diffstates_spec() + self.jedi_rendering.add_key('diffstates_spec', spec_dict) + print( f'diffstates = {spec_dict}') + + # Add placeholder names if mock experiment + # ---------------------------------------- + if self.config.mock_experiment(False): + self.jedi_rendering.add_key('experiment_root', 'experiment_root') + self.jedi_rendering.add_key('experiment_id', 'experiment_id') + self.jedi_rendering.add_key('cycle_dir', 'cycle_dir') + + # loop output grid_type: + grid_type = spec_dict['state_diff'].get('grid_type') + print (f'grid_type = {grid_type}') + + if grid_type is None: + grid_type = ['latlon'] + + for output_grid in grid_type: + + self.jedi_rendering.add_key('diffstates_spec_gridtype', output_grid) + + # Jedi configuration file + # ----------------------- + jedi_config_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_output_{output_grid}_config.yaml') + + # Output log file + # --------------- + output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_output_{output_grid}.log') + + # Open the JEDI config file and fill templates + # -------------------------------------------- + jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}', + window_type, + jedi_forecast_model) + yaml = YAML() + yaml.default_flow_style = False + + # Write the expanded dictionary to YAML file + # ------------------------------------------ + with open(jedi_config_file, 'w') as jedi_config_file_open: + yaml.dump(jedi_config_dict, jedi_config_file_open) + + # Get the JEDI interface metadata + # ------------------------------- + model_component_meta = self.jedi_rendering.render_interface_meta() + + # Compute number of processors + # ---------------------------- + np = eval(str(model_component_meta['total_processors'])) + + # Jedi executable name + # -------------------- + jedi_executable = model_component_meta['executables'][f'{jedi_application}'] + jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', 'build', 'bin', + jedi_executable) + + # Run the JEDI executable + # ----------------------- + if not generate_yaml_and_exit: + self.logger.info('Running '+jedi_executable_path+' with '+str(np)+' processors.') + run_executable(self.logger, self.cycle_dir(), np, jedi_executable_path, + jedi_config_file, output_log_file) + else: + self.logger.info('YAML generated, now exiting.') + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/run_jedi_ensemble_mean_variance.py b/src/swell/tasks/run_jedi_ensemble_mean_variance.py index a514998ab..1e32aebb4 100644 --- a/src/swell/tasks/run_jedi_ensemble_mean_variance.py +++ b/src/swell/tasks/run_jedi_ensemble_mean_variance.py @@ -98,55 +98,60 @@ def execute(self) -> None: self.jedi_rendering.add_key('cycle_dir', 'cycle_dir') - # loop item in ensmeanvariance_spec + # loop item in ensmeanvariance_spec: state + grid_type # ----------------------- for idx, spec in enumerate(meanvar_spec_dict, start=1): self.jedi_rendering.add_key('ensmeanvariance_spec_item', spec) - - # Jedi configuration file - # ----------------------- - jedi_config_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_config_{idx}.yaml') - - # Output log file - # --------------- - output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_{idx}.log') - - # Open the JEDI config file and fill templates - # -------------------------------------------- - jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}', - window_type, - jedi_forecast_model) - - yaml = YAML() - yaml.default_flow_style = False - - # Write the expanded dictionary to YAML file - # ------------------------------------------ - with open(jedi_config_file, 'w') as jedi_config_file_open: - yaml.dump(jedi_config_dict, jedi_config_file_open) - - # Get the JEDI interface metadata - # ------------------------------- - model_component_meta = self.jedi_rendering.render_interface_meta() - - # Compute number of processors - # ---------------------------- - np = eval(str(model_component_meta['total_processors'])) - - # Jedi executable name - # -------------------- - jedi_executable = model_component_meta['executables'][f'{jedi_application}'] - jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', 'build', 'bin', - jedi_executable) - - # Run the JEDI executable - # ----------------------- - if not generate_yaml_and_exit: - self.logger.info('Running '+jedi_executable_path+' with '+str(np)+' processors.') - run_executable(self.logger, self.cycle_dir(), np, jedi_executable_path, - jedi_config_file, output_log_file) - else: - self.logger.info('YAML generated, now exiting.') + state_name = spec.get('state') + + for output_grid in spec.get('grid_type'): + + self.jedi_rendering.add_key('ensmeanvariance_spec_gridtype', output_grid) + + # Jedi configuration file + # ----------------------- + jedi_config_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_config_{state_name}_{output_grid}.yaml') + + # Output log file + # --------------- + output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_{state_name}_{output_grid}.log') + + # Open the JEDI config file and fill templates + # -------------------------------------------- + jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}', + window_type, + jedi_forecast_model) + + yaml = YAML() + yaml.default_flow_style = False + + # Write the expanded dictionary to YAML file + # ------------------------------------------ + with open(jedi_config_file, 'w') as jedi_config_file_open: + yaml.dump(jedi_config_dict, jedi_config_file_open) + + # Get the JEDI interface metadata + # ------------------------------- + model_component_meta = self.jedi_rendering.render_interface_meta() + + # Compute number of processors + # ---------------------------- + np = eval(str(model_component_meta['total_processors'])) + + # Jedi executable name + # -------------------- + jedi_executable = model_component_meta['executables'][f'{jedi_application}'] + jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', 'build', 'bin', + jedi_executable) + + # Run the JEDI executable + # ----------------------- + if not generate_yaml_and_exit: + self.logger.info('Running '+jedi_executable_path+' with '+str(np)+' processors.') + run_executable(self.logger, self.cycle_dir(), np, jedi_executable_path, + jedi_config_file, output_log_file) + else: + self.logger.info('YAML generated, now exiting.') # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index de923990d..db44736fc 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -650,6 +650,22 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + RunJediDiffstates = QuestionList( + list_name="RunJediDiffstates", + questions=[ + np_proc_resolution, + window_questions, + qd.analysis_variables(), + qd.diffstates_spec(), + qd.generate_yaml_and_exit(), + qd.jedi_forecast_model(), + qd.comparison_log_type('diffstates'), + qd.mock_experiment() + ] + ) + + # -------------------------------------------------------------------------------------------------- + RunJediFgatExecutable = QuestionList( list_name="RunJediFgatExecutable", questions=[ diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index 73e9ab1f6..2856861b2 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -8,8 +8,8 @@ # -------------------------------------------------------------------------------------------------- -from dataclasses import dataclass -from typing import List, Dict +from dataclasses import dataclass, field +from typing import List, Dict, Any from swell.utilities.swell_questions import SuiteQuestion, TaskQuestion from swell.utilities.swell_questions import WidgetType as WType @@ -538,26 +538,26 @@ class ensmeanvariance_only(TaskQuestion): @dataclass class ensmeanvariance_spec(TaskQuestion): - default_value: List[Dict[str, str]] = mutable_field([ - { - "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", - "fn_output_mean": "geos.prior.mean", - "fn_output_variance": "geos.prior.variance", - "geom": "latlon" - }, - { - "fn_input": "analysis/mem%mem%/eda.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", - "fn_output_mean": "eda.ana.mean", - "fn_output_variance": "eda.ana.variance", - "geom": "latlon" - }, - ]) + default_value: List[Dict[str, str]] = field(default_factory=lambda: [{}]) question_name: str = "ensmeanvariance_spec" models: List[str] = mutable_field([ "all_models" ]) ask_question: bool = True - prompt: str = "Configure the ensemble mean and variance specifications: [e.g. ebkg, eda.analysis, eda.inc, etc.]" + prompt: str = "Configure the ensemble mean and variance specifications: [ebkg, eda.analysis]" + widget_type: WType = WType.STRING + + # -------------------------------------------------------------------------------------------------- + + @dataclass + class diffstates_spec(TaskQuestion): + default_value: Dict[str, Any] = field(default_factory=dict) + question_name: str = "diffstates_spec" + models: List[str] = mutable_field([ + "all_models" + ]) + ask_question: bool = True + prompt: str = "Configure the diffstates specifications: [state1, state2]" widget_type: WType = WType.STRING # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/utilities/render_jedi_interface_files.py b/src/swell/utilities/render_jedi_interface_files.py index 6532913cb..8a887e367 100644 --- a/src/swell/utilities/render_jedi_interface_files.py +++ b/src/swell/utilities/render_jedi_interface_files.py @@ -85,6 +85,9 @@ def __init__( 'ensmean_only', 'ensmeanvariance_only', 'ensmeanvariance_spec_item', + 'ensmeanvariance_spec_gridtype', + 'diffstates_spec', + 'diffstates_spec_gridtype', 'experiment_id', 'experiment_root', 'final_cycle_point', diff --git a/src/swell/utilities/slurm.py b/src/swell/utilities/slurm.py index a3d6be04d..6e9d1576e 100644 --- a/src/swell/utilities/slurm.py +++ b/src/swell/utilities/slurm.py @@ -81,6 +81,7 @@ def prepare_scheduling_dict( 'RunGeos', 'RunJediEdaExecutable', 'RunJediEnsembleMeanVariance', + 'RunJediDiffstates', 'RunJediConvertStateSoca2ciceExecutable', 'RunJediFgatExecutable', 'RunJediHofxEnsembleExecutable', From 359b48b9d5aa4e692edc3982b13dd1c2a8513a5c Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Mon, 8 Jun 2026 17:38:44 -0600 Subject: [PATCH 11/68] wip: eorr in diffstates --- .../geos_atmosphere/model/shared.py | 39 +++++++++++++++++++ .../geos_atmosphere/model/two_states.py | 15 ++++--- .../jedi/oops/ensmeanvariance.py | 6 ++- 3 files changed, 53 insertions(+), 7 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py index 4c32cf09c..dd9bcd9c7 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py @@ -101,7 +101,46 @@ 'water_vapor_mixing_ratio_wrt_moist_air': 'q', 'cloud_liquid_ice': 'qi', 'cloud_liquid_water': 'ql', + 'rain_water': 'qr', + 'snow_water': 'qs', 'mole_fraction_of_ozone_in_air': 'o3ppmv', + 'geopotential_height_times_gravity_at_surface': 'phis', + 'fraction_of_ocean': 'frocean', + 'fraction_of_lake': 'frlake', + 'fraction_of_ice': 'frseaice', + 'skin_temperature_at_surface': 'ts', + 'eastward_wind_at_surface': 'u10m', + 'northward_wind_at_surface': 'v10m', + # 'sea_surface_temperature': 'ts_found', + # 'mole_fraction_of_carbon_dioxide_in_air': 'co2', } # -------------------------------------------------------------------------------------------------- + +state_variables_ensemble = [ + 'eastward_wind', + 'northward_wind', + 'air_temperature', + 'air_pressure_at_surface', + 'air_pressure_levels', + 'water_vapor_mixing_ratio_wrt_moist_air', + 'cloud_liquid_ice', + 'cloud_liquid_water', + 'rain_water', + 'snow_water', + 'mole_fraction_of_ozone_in_air', + 'geopotential_height_times_gravity_at_surface', + 'fraction_of_ocean', + 'fraction_of_lake', + 'fraction_of_ice', + 'sheleg', + 'skin_temperature_at_surface', + 'soilt', + 'soilm', + 'eastward_wind_at_surface', + 'northward_wind_at_surface', + # 'sea_surface_temperature', + # 'mole_fraction_of_carbon_dioxide_in_air', + ] + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py index b116e40e6..b614cfcb5 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py @@ -8,14 +8,17 @@ from collections.abc import Mapping from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ - state_variables, field_io_names + state_variables, field_io_names, state_variables_ensemble, field_io_names_ensemble # -------------------------------------------------------------------------------------------------- def two_states(template_dict: Mapping) -> Mapping: horizontal_resolution = template_dict['horizontal_resolution'] f1 = template_dict['diffstates_spec']['state1']['fn_input'] - f2 = template_dict['diffstates_spec']['state2']['fn_input'] + f2 = template_dict['diffstates_spec']['state2']['fn_input'] + + state_variables_loc = state_variables_ensemble + field_io_names_loc = field_io_names_ensemble # state1: bkg state1 = { @@ -24,8 +27,8 @@ def two_states(template_dict: Mapping) -> Mapping: 'provider': 'geos', 'datapath': template_dict['cycle_dir'], 'filename': f1, - 'state variables': state_variables, - 'field io names': field_io_names, + 'state variables': state_variables_loc, + 'field io names': field_io_names_loc, } # state2: analysis @@ -35,8 +38,8 @@ def two_states(template_dict: Mapping) -> Mapping: 'provider': 'geos', 'datapath': template_dict['cycle_dir'], 'filenames': f2, - 'state variables': state_variables, - 'field io names': field_io_names, + 'state variables': state_variables_loc, + 'field io names': field_io_names_loc, } combo = {'state1': state1, 'state2': state2} diff --git a/src/swell/configuration/jedi/oops/ensmeanvariance.py b/src/swell/configuration/jedi/oops/ensmeanvariance.py index 1c51acd12..b2272ad2f 100644 --- a/src/swell/configuration/jedi/oops/ensmeanvariance.py +++ b/src/swell/configuration/jedi/oops/ensmeanvariance.py @@ -17,7 +17,11 @@ def render_oops(self): oops = { 'geometry': self.interface_model('geometry'), - 'ensemble': self.interface_model('state_ensemble'), + 'ensemble': self.interface_model('ensemble_block'), + # + # full state variables do not work in ensemble + # 'ensemble': self.interface_model('state_ensemble'), + # 'variance output': self.interface_model('comp_variance_output'), 'mean output': self.interface_model('comp_mean_output') } From 98f6ee69c252d593a7327e8788c3e9729ff01829 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Mon, 8 Jun 2026 20:19:59 -0600 Subject: [PATCH 12/68] For configuration/jedi/interfaces/geos_atmosphere/model/two_states.py - select minimal state variables / field_io_names - add `if state == ensemble` selection - correct typo filenames -> filename --- .../model/diffstates_output.py | 16 ++++++---- .../geos_atmosphere/model/shared.py | 29 +++++++++++++++++++ .../geos_atmosphere/model/two_states.py | 15 ++++++---- src/swell/tasks/run_jedi_diffstates.py | 11 ++++--- .../utilities/render_jedi_interface_files.py | 3 +- 5 files changed, 59 insertions(+), 15 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py index db3c2e789..d78a0322a 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py @@ -8,7 +8,7 @@ from collections.abc import Mapping from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ - field_io_names, field_io_names_ensemble + field_io_names, field_io_names_sa1 # -------------------------------------------------------------------------------------------------- @@ -19,6 +19,12 @@ def diffstates_output(template_dict: Mapping) -> Mapping: grid_type= template_dict['diffstates_spec_gridtype'] prefix_output = template_dict['diffstates_spec'].get('state_diff', {}).get('fn_output') + stateType = template_dict.get('diffstates_spec_statetype') + if [ stateType == 'ensemble' ]: + field_io_names_loc = field_io_names_sa1 + else: + field_io_names_loc = field_io_names + output = {} if grid_type == 'cs': output = { @@ -26,7 +32,7 @@ def diffstates_output(template_dict: Mapping) -> Mapping: 'provider': 'geos', 'datapath': template_dict['cycle_dir'], 'filename': f'{prefix_output}.%yyyy%mm%dd_%hh%MM%ssz.nc4', - 'field io names': field_io_names_ensemble + 'field io names': field_io_names_loc } elif grid_type == 'latlon': output = { @@ -34,9 +40,9 @@ def diffstates_output(template_dict: Mapping) -> Mapping: 'gridtype': 'latlon', 'datapath': template_dict['cycle_dir'], 'filename': f'{prefix_output}.ll.', - 'field io names': field_io_names_ensemble - } - + 'field io names': field_io_names_loc + } + return output # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py index dd9bcd9c7..939563f8b 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py @@ -144,3 +144,32 @@ ] # -------------------------------------------------------------------------------------------------- + +# stand alone 1 +state_variables_sa1 = [ + 'eastward_wind', + 'northward_wind', + 'air_temperature', + 'air_pressure_at_surface', + 'air_pressure_levels', + 'water_vapor_mixing_ratio_wrt_moist_air', + 'cloud_liquid_ice', + 'cloud_liquid_water', + 'mole_fraction_of_ozone_in_air' + ] + +# -------------------------------------------------------------------------------------------------- + +field_io_names_sa1 = { + 'eastward_wind': 'ua', + 'northward_wind': 'va', + 'air_temperature': 't', + 'air_pressure_at_surface': 'ps', + 'air_pressure_levels': 'pe', + 'water_vapor_mixing_ratio_wrt_moist_air': 'q', + 'cloud_liquid_ice': 'qi', + 'cloud_liquid_water': 'ql', + 'mole_fraction_of_ozone_in_air': 'o3ppmv', +} + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py index b614cfcb5..035244347 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py @@ -8,7 +8,7 @@ from collections.abc import Mapping from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ - state_variables, field_io_names, state_variables_ensemble, field_io_names_ensemble + state_variables, field_io_names, state_variables_sa1, field_io_names_sa1 # -------------------------------------------------------------------------------------------------- @@ -17,8 +17,13 @@ def two_states(template_dict: Mapping) -> Mapping: f1 = template_dict['diffstates_spec']['state1']['fn_input'] f2 = template_dict['diffstates_spec']['state2']['fn_input'] - state_variables_loc = state_variables_ensemble - field_io_names_loc = field_io_names_ensemble + stateType = template_dict.get('diffstates_spec_statetype') + if [ stateType == 'ensemble' ]: + state_variables_loc = state_variables_sa1 + field_io_names_loc = field_io_names_sa1 + else: + state_variables_loc = state_variables + field_io_names_loc = field_io_names # state1: bkg state1 = { @@ -37,11 +42,11 @@ def two_states(template_dict: Mapping) -> Mapping: 'filetype': 'cube sphere history', 'provider': 'geos', 'datapath': template_dict['cycle_dir'], - 'filenames': f2, + 'filename': f2, 'state variables': state_variables_loc, 'field io names': field_io_names_loc, } - + combo = {'state1': state1, 'state2': state2} return combo diff --git a/src/swell/tasks/run_jedi_diffstates.py b/src/swell/tasks/run_jedi_diffstates.py index 2bb115d80..25eb046fd 100644 --- a/src/swell/tasks/run_jedi_diffstates.py +++ b/src/swell/tasks/run_jedi_diffstates.py @@ -68,7 +68,7 @@ def execute(self) -> None: # Diffstates spec_dict = self.config.diffstates_spec() - self.jedi_rendering.add_key('diffstates_spec', spec_dict) + self.jedi_rendering.add_key('diffstates_spec', spec_dict) print( f'diffstates = {spec_dict}') # Add placeholder names if mock experiment @@ -78,13 +78,16 @@ def execute(self) -> None: self.jedi_rendering.add_key('experiment_id', 'experiment_id') self.jedi_rendering.add_key('cycle_dir', 'cycle_dir') + state_type = spec_dict.get('state_type') + self.jedi_rendering.add_key('diffstates_spec_statetype', state_type) + # loop output grid_type: grid_type = spec_dict['state_diff'].get('grid_type') print (f'grid_type = {grid_type}') if grid_type is None: grid_type = ['latlon'] - + for output_grid in grid_type: self.jedi_rendering.add_key('diffstates_spec_gridtype', output_grid) @@ -96,7 +99,7 @@ def execute(self) -> None: # Output log file # --------------- output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_output_{output_grid}.log') - + # Open the JEDI config file and fill templates # -------------------------------------------- jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}', @@ -109,7 +112,7 @@ def execute(self) -> None: # ------------------------------------------ with open(jedi_config_file, 'w') as jedi_config_file_open: yaml.dump(jedi_config_dict, jedi_config_file_open) - + # Get the JEDI interface metadata # ------------------------------- model_component_meta = self.jedi_rendering.render_interface_meta() diff --git a/src/swell/utilities/render_jedi_interface_files.py b/src/swell/utilities/render_jedi_interface_files.py index 8a887e367..7e0c6029d 100644 --- a/src/swell/utilities/render_jedi_interface_files.py +++ b/src/swell/utilities/render_jedi_interface_files.py @@ -87,7 +87,8 @@ def __init__( 'ensmeanvariance_spec_item', 'ensmeanvariance_spec_gridtype', 'diffstates_spec', - 'diffstates_spec_gridtype', + 'diffstates_spec_gridtype', + 'diffstates_spec_statetype', 'experiment_id', 'experiment_root', 'final_cycle_point', From 87a3ca3c3c6c2bc6c4e038307034a5901d782a62 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Mon, 8 Jun 2026 22:52:44 -0600 Subject: [PATCH 13/68] checked swell create, filter, eda imem, mean_variance, diffstate --- .../interfaces/geos_atmosphere/model/diffstates_output.py | 4 ++-- .../jedi/interfaces/geos_atmosphere/model/shared.py | 4 +++- src/swell/suites/eda/flow.cylc | 6 +++++- src/swell/suites/eda/suite_config.py | 5 +++-- src/swell/tasks/eva_increment.py | 2 ++ 5 files changed, 15 insertions(+), 6 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py index d78a0322a..578f75f06 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py @@ -31,7 +31,7 @@ def diffstates_output(template_dict: Mapping) -> Mapping: 'filetype': 'cube sphere history', 'provider': 'geos', 'datapath': template_dict['cycle_dir'], - 'filename': f'{prefix_output}.%yyyy%mm%dd_%hh%MM%ssz.nc4', + 'filename': f'{prefix_output}.cs.%yyyy%mm%dd_%hh%MM%ssz.nc4', 'field io names': field_io_names_loc } elif grid_type == 'latlon': @@ -39,7 +39,7 @@ def diffstates_output(template_dict: Mapping) -> Mapping: 'filetype': 'auxgrid', 'gridtype': 'latlon', 'datapath': template_dict['cycle_dir'], - 'filename': f'{prefix_output}.ll.', + 'filename': f'{prefix_output}.', 'field io names': field_io_names_loc } diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py index 939563f8b..c5a250d7a 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py @@ -155,7 +155,8 @@ 'water_vapor_mixing_ratio_wrt_moist_air', 'cloud_liquid_ice', 'cloud_liquid_water', - 'mole_fraction_of_ozone_in_air' + 'mole_fraction_of_ozone_in_air', + 'skin_temperature_at_surface', ] # -------------------------------------------------------------------------------------------------- @@ -170,6 +171,7 @@ 'cloud_liquid_ice': 'qi', 'cloud_liquid_water': 'ql', 'mole_fraction_of_ozone_in_air': 'o3ppmv', + 'skin_temperature_at_surface': 'ts', } # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/eda/flow.cylc b/src/swell/suites/eda/flow.cylc index 88052268a..d62bac10a 100644 --- a/src/swell/suites/eda/flow.cylc +++ b/src/swell/suites/eda/flow.cylc @@ -75,9 +75,11 @@ GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} + {% if models[model_component]['ensemble_num_members'] is defined %} {% for i in range( 1, models[model_component]['ensemble_num_members'] + 1 ) %} SP2-{{model_component}} => RunJediEdaExecutable_mem{{i}}-{{model_component}} => SP3-{{model_component}} {% endfor %} + {% endif %} # bkg/ana : mean/variance SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} @@ -161,7 +163,8 @@ [[RenderJediObservations-{{model_component}}]] script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" - {% for i in range(1, 5) %} + {% if models[model_component]['ensemble_num_members'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_num_members'] + 1 ) %} [[RunJediEdaExecutable_mem{{i}}-{{model_component}}]] script = "swell task RunJediEdaExecutable $config -d $datetime -m {{model_component}} -imem {{i}}" platform = {{platform}} @@ -171,6 +174,7 @@ --{{key}} = {{value}} {%- endfor %} {% endfor %} + {% endif %} # [[RunJediEdaExecutable-{{model_component}}]] # script = "swell task RunJediEdaExecutable $config -d $datetime -m {{model_component}}" diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda/suite_config.py index b6a060918..3ff33b71f 100644 --- a/src/swell/suites/eda/suite_config.py +++ b/src/swell/suites/eda/suite_config.py @@ -101,7 +101,7 @@ class SuiteConfig(QuestionContainer, Enum): "sondes", "ssmis_f17" ]), - qd.obs_thinning_rej_fraction(0.75), + qd.obs_thinning_rej_fraction(0.98), qd.ensmeanvariance_spec([ {"state": "bkg", "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", @@ -120,7 +120,8 @@ class SuiteConfig(QuestionContainer, Enum): "state2": {"fn_input": "eda.ana.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, "state_diff": - {"fn_output": "eda.inc.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + {"fn_output": "eda.mean-inc", "grid_type": ['cs', 'latlon']}, + "state_type": "ensemble" }), qd.clean_patterns(['*.txt', '*.csv']), ] diff --git a/src/swell/tasks/eva_increment.py b/src/swell/tasks/eva_increment.py index ae798fd3d..10a53764f 100644 --- a/src/swell/tasks/eva_increment.py +++ b/src/swell/tasks/eva_increment.py @@ -68,6 +68,8 @@ def execute(self) -> None: incr_file = f'{self.experiment_id()}.increment-iter{iter_no}.{cycle_time_reformat}.nc4' if self.suite_name() == 'localensembleda': incr_file = f'geos.mean-inc.{local_bkg_time}.nc4' + elif self.suite_name() == 'eda': + incr_file = f'eda.mean-inc.{local_bkg_time}.nc4' if window_type == '4D' and 'atmos' in self.suite_name(): incr_file = f'{self.experiment_id()}.increment-iter{iter_no}.{window_begin}.nc4' From d92c5e55b4fbc0192f07df44d1f22027c43c647d Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 9 Jun 2026 13:04:28 -0600 Subject: [PATCH 14/68] Add ebkg_time_offset to get ensemble geos experiment --- .../interfaces/geos_atmosphere/task_questions.yaml | 3 +++ src/swell/tasks/get_ensemble_geos_experiment.py | 2 +- src/swell/tasks/task_questions.py | 1 + src/swell/utilities/question_defaults.py | 13 +++++++++++++ 4 files changed, 18 insertions(+), 1 deletion(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml b/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml index 636e49e1c..fb83711a5 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml @@ -52,6 +52,9 @@ background_frequency: background_time_offset: default_value: PT9H +ebkg_time_offset: + default_value: PT3H + bufr_obs_classes: default_value: - disc_airs_bufr diff --git a/src/swell/tasks/get_ensemble_geos_experiment.py b/src/swell/tasks/get_ensemble_geos_experiment.py index 11d69feab..178967987 100644 --- a/src/swell/tasks/get_ensemble_geos_experiment.py +++ b/src/swell/tasks/get_ensemble_geos_experiment.py @@ -45,7 +45,7 @@ def execute(self): # ------------ background_experiment = self.config.background_experiment() geos_x_ensemble_directory = self.config.geos_x_ensemble_directory() - background_time_offset = self.config.background_time_offset() + background_time_offset = self.config.ebkg_time_offset() # Since this is an optional task, check if the geos_x_ensemble_directory is # set to /dev/null, if so fail the task diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index 121174439..570227112 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -392,6 +392,7 @@ class TaskQuestions(QuestionContainer, Enum): questions=[ qd.background_experiment(), qd.background_time_offset(), + qd.ebkg_time_offset(), qd.geos_x_ensemble_directory() ] ) diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index 3a0ea5b2c..e02f7900f 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -356,6 +356,19 @@ class background_time_offset(TaskQuestion): # -------------------------------------------------------------------------------------------------- + @dataclass + class ebkg_time_offset(TaskQuestion): + default_value: str = "defer_to_model" + question_name: str = "ebkg_time_offset" + models: List[str] = mutable_field([ + "all_models" + ]) + prompt: str = ("How long before the middle of the analysis window did" + " the ensemble background providing forecast begin?") + widget_type: WType = WType.ISO_DURATION + + # -------------------------------------------------------------------------------------------------- + @dataclass class rst_experiment(TaskQuestion): default_value: str = "defer_to_model" From aeeb8dd3b9d60cc4da5b13cba954c106501a7f3d Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 9 Jun 2026 18:07:32 -0400 Subject: [PATCH 15/68] adjust 3dvar config, e.g., npx, number_of_iteration, etc --- src/swell/suites/eda/suite_config.py | 12 ++++++++---- src/swell/utilities/slurm.py | 2 +- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda/suite_config.py index 3ff33b71f..899e7f3de 100644 --- a/src/swell/suites/eda/suite_config.py +++ b/src/swell/suites/eda/suite_config.py @@ -39,14 +39,18 @@ class SuiteConfig(QuestionContainer, Enum): "dadev/rtodling/archive/Restarts/JEDI/541x"), qd.geos_x_ensemble_directory('/discover/nobackup/projects/gmao/dadev/' 'rtodling/archive/541/Milan'), + qd.npx_proc(4), + qd.npy_proc(5), qd.window_length("PT6H"), qd.window_type("3D"), qd.horizontal_resolution("91"), qd.gsibec_nlats("91"), qd.gsibec_nlons("144"), qd.vertical_resolution("72"), - qd.ensemble_num_members(4), + qd.ensemble_num_members(16), qd.obs_pert_amplitude(0.5), + qd.number_of_iterations([25]), + qd.gradient_norm_reduction(1.e-8), qd.analysis_variables([ "eastward_wind", "northward_wind", @@ -101,7 +105,7 @@ class SuiteConfig(QuestionContainer, Enum): "sondes", "ssmis_f17" ]), - qd.obs_thinning_rej_fraction(0.98), + qd.obs_thinning_rej_fraction(0.8), qd.ensmeanvariance_spec([ {"state": "bkg", "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", @@ -117,9 +121,9 @@ class SuiteConfig(QuestionContainer, Enum): qd.diffstates_spec({ "state1": {"fn_input": "geos.prior.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, - "state2": + "state2": {"fn_input": "eda.ana.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, - "state_diff": + "state_diff": {"fn_output": "eda.mean-inc", "grid_type": ['cs', 'latlon']}, "state_type": "ensemble" }), diff --git a/src/swell/utilities/slurm.py b/src/swell/utilities/slurm.py index 6e9d1576e..f412bd2b1 100644 --- a/src/swell/utilities/slurm.py +++ b/src/swell/utilities/slurm.py @@ -81,7 +81,7 @@ def prepare_scheduling_dict( 'RunGeos', 'RunJediEdaExecutable', 'RunJediEnsembleMeanVariance', - 'RunJediDiffstates', + 'RunJediDiffstates', 'RunJediConvertStateSoca2ciceExecutable', 'RunJediFgatExecutable', 'RunJediHofxEnsembleExecutable', From 90726ee7c3cabeffd3b6af0e6ed01da30a340e14 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 9 Jun 2026 19:22:31 -0600 Subject: [PATCH 16/68] move obs and fv3-jedi files into analysis/mem00x and change corresponding input/output dirnames in yaml --- .../geos_atmosphere/model/background_eda.py | 4 +- src/swell/tasks/run_jedi_eda_executable.py | 76 +++++++++++++++++-- 2 files changed, 73 insertions(+), 7 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py index 1087f4766..61de175f0 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py @@ -71,8 +71,8 @@ def background_eda(template_dict: Mapping) -> Mapping: 'max allowable geometry difference': 1e-3, 'datapath': template_dict['cycle_dir'], 'filenames': [ - f'ebkg/mem{imem:03d}/geos.mem{imem:03d}.%yyyy%mm%dd_%hh%MM%ssz.nc4', - f'fv3-jedi/bkg/geos.crtmsrf.{horizontal_resolution}.nc4' + f'./ebkg/mem{imem:03d}/geos.mem{imem:03d}.%yyyy%mm%dd_%hh%MM%ssz.nc4', + f'./fv3-jedi/bkg/geos.crtmsrf.{horizontal_resolution}.nc4' ], 'state variables': state_variables, 'field io names': field_io_names, diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 6c4931d07..e5aeb174c 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -9,6 +9,8 @@ import os import copy +import glob +import shutil from ruamel.yaml import YAML from swell.tasks.base.task_base import taskBase @@ -140,9 +142,23 @@ def execute(self) -> None: else: jedi_config_dict['cost function']['observations'].update({'obs perturbations': True}) + # create subdir + mem_dir = f'analysis/mem{imember:003d}/' + xdir = os.path.join(self.cycle_dir(), mem_dir) + os.makedirs(xdir, exist_ok=True) + for observer in jedi_config_dict['cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] + + print( f'ob= {observation}' ) + # copy obs input file to avoid multi MPI reading the same file + files = glob.glob(os.path.join(self.cycle_dir(), f'{observation}*.nc4')) + \ + glob.glob(os.path.join(self.cycle_dir(), f'{observation}*.txt')) + for src_file in files: + print( f'f= {src_file}' ) + shutil.copy(src_file, xdir) + if imember > 1: obs_cov_model = observer.get('obs error', {}).get('covariance model') print( f'{observation}: obs_cov_model = {obs_cov_model}') @@ -161,13 +177,37 @@ def execute(self) -> None: hxout = observer['obs space']['obsdataout']['engine']['obsfile'] dir1, fname = os.path.split(hxout) - hxout = os.path.join(dir1, f'analysis/mem{imember:003d}', fname) + hxout = os.path.join(dir1, mem_dir, fname) observer['obs space']['obsdataout']['engine']['obsfile'] = hxout - # print (f'hxout = {hxout}') - # create subdir - xdir = os.path.join(self.cycle_dir(), f'analysis/mem{imember:003d}') - os.makedirs(xdir, exist_ok=True) + obsFileIn = observer['obs space']['obsdatain']['engine']['obsfile'] + dir1, fname = os.path.split(obsFileIn) + obsFileIn = os.path.join(dir1, mem_dir, fname) + observer['obs space']['obsdatain']['engine']['obsfile'] = obsFileIn + + obs_bias = observer.get('obs bias') + if obs_bias is not None: + File = obs_bias['input file'] + dir1, fname = os.path.split(File) + File = os.path.join(dir1, mem_dir, fname) + obs_bias['input file'] = File + # + File = obs_bias['output file'] + dir1, fname = os.path.split(File) + File = os.path.join(dir1, mem_dir, fname) + obs_bias['output file'] = File + # + File = obs_bias.get('covariance', {}).get('output file') + if File is not None: + dir1, fname = os.path.split(File) + File = os.path.join(dir1, mem_dir, fname) + obs_bias['covariance']['output file'] = File + # + File = obs_bias.get('covariance', {}).get('prior',{}).get('input file') + if File is not None: + dir1, fname = os.path.split(File) + File = os.path.join(dir1, mem_dir, fname) + obs_bias['covariance']['prior']['input file'] = File ruamel_yaml = YAML() ruamel_yaml.default_flow_style = False @@ -176,6 +216,22 @@ def execute(self) -> None: with open(jedi_config_file, 'w') as jedi_config_file_open: ruamel_yaml.dump(jedi_config_dict, jedi_config_file_open) + # copy fv3-jedi dir, update dir names + d1 = os.path.join(self.cycle_dir(), 'fv3-jedi') + d2 = os.path.join(self.cycle_dir(), mem_dir, 'fv3-jedi') + shutil.copytree(d1, d2, dirs_exist_ok=True) + + with open(jedi_config_file, 'r') as f: + yaml_content = f.read() + + dir_list = ["bkg", "fv3files", "gsibec", "rcov"] + for i in dir_list: + j = f"fv3-jedi/{i}" + k = f"{mem_dir}{j}" # Result: analysis/mem002/fv3-jedi/rcov + yaml_content = yaml_content.replace(j, k) + with open(jedi_config_file, 'w') as f: + f.write(yaml_content) + # Get the JEDI interface metadata # ------------------------------- model_component_meta = self.jedi_rendering.render_interface_meta() @@ -204,4 +260,14 @@ def execute(self) -> None: print(f'intended mpi_command = {mpi_command}') self.logger.info('YAML generated, now exiting.') + +# clean up dir +# for observer in jedi_config_dict['cost function']['observations']['observers']: +# # Get observation name +# observation = observer['observation_name'] +# # Delete input obsfile .nc4 files +# for file_path in glob.glob(os.path.join(mem_dir, f'{observation}*.nc4')): +# os.remove(file_path) +# print(f"Deleted: {file_path}") + # -------------------------------------------------------------------------------------------------- From 1fa54c8d20b2ccdad8ae4a060ea1f52f04d6e4eb Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 10 Jun 2026 12:52:50 -0400 Subject: [PATCH 17/68] eda: iter=100; norm_reduction=1.e-8 --- src/swell/suites/eda/suite_config.py | 4 ++-- src/swell/tasks/run_jedi_eda_executable.py | 3 +-- src/swell/utilities/slurm.py | 3 ++- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda/suite_config.py index 899e7f3de..f5276285c 100644 --- a/src/swell/suites/eda/suite_config.py +++ b/src/swell/suites/eda/suite_config.py @@ -47,9 +47,9 @@ class SuiteConfig(QuestionContainer, Enum): qd.gsibec_nlats("91"), qd.gsibec_nlons("144"), qd.vertical_resolution("72"), - qd.ensemble_num_members(16), + qd.ensemble_num_members(32), qd.obs_pert_amplitude(0.5), - qd.number_of_iterations([25]), + qd.number_of_iterations([100]), qd.gradient_norm_reduction(1.e-8), qd.analysis_variables([ "eastward_wind", diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index e5aeb174c..0d6178104 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -153,8 +153,7 @@ def execute(self) -> None: print( f'ob= {observation}' ) # copy obs input file to avoid multi MPI reading the same file - files = glob.glob(os.path.join(self.cycle_dir(), f'{observation}*.nc4')) + \ - glob.glob(os.path.join(self.cycle_dir(), f'{observation}*.txt')) + files = glob.glob(os.path.join(self.cycle_dir(), f'{observation}.*')) for src_file in files: print( f'f= {src_file}' ) shutil.copy(src_file, xdir) diff --git a/src/swell/utilities/slurm.py b/src/swell/utilities/slurm.py index f412bd2b1..6e7a7d6b8 100644 --- a/src/swell/utilities/slurm.py +++ b/src/swell/utilities/slurm.py @@ -46,7 +46,8 @@ def prepare_scheduling_dict( task_defaults = { "RunJediVariationalExecutable": {"all": {"nodes": 3}}, "RunJediUfoTestsExecutable": {"all": {"ntasks-per-node": 1}}, - "RunJediConvertStateSoca2ciceExecutable": {"all": {"nodes": 1}} + "RunJediConvertStateSoca2ciceExecutable": {"all": {"nodes": 1}}, + "RunJediEdaExecutable": {"all": {"ntasks-per-node": 126}} } # Global SLURM settings stored in $HOME/.swell/swell-slurm.yaml From 465190f1c9d9f1070c9c8f349bc4c93e41ed274b Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 10 Jun 2026 12:59:25 -0600 Subject: [PATCH 18/68] Add background_error_eda_gsiB.py --- .../model/background_error_eda_gsiB.py | 90 +++++++++++++++++++ ...rid.py => background_error_eda_hybridB.py} | 3 +- src/swell/configuration/jedi/oops/eda3D.py | 6 +- 3 files changed, 97 insertions(+), 2 deletions(-) create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py rename src/swell/configuration/jedi/interfaces/geos_atmosphere/model/{background_error_hybrid.py => background_error_eda_hybridB.py} (98%) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py new file mode 100644 index 000000000..586aa28f0 --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py @@ -0,0 +1,90 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping + +# -------------------------------------------------------------------------------------------------- + +state_variables_to_inverse = [ + 'eastward_wind', + 'northward_wind', + 'air_temperature', + 'air_pressure_at_surface', +# 'air_pressure_levels', + 'water_vapor_mixing_ratio_wrt_moist_air', + 'cloud_liquid_ice', + 'cloud_liquid_water', + 'rain_water', + 'snow_water', + 'mole_fraction_of_ozone_in_air', + 'fraction_of_ocean', + 'fraction_of_lake', + 'fraction_of_ice', + 'geopotential_height_times_gravity_at_surface', + 'skin_temperature_at_surface', +] + +# -------------------------------------------------------------------------------------------------- + + +def background_error_eda_gsiB(template_dict: Mapping) -> Mapping: + vertical_resolution = template_dict['vertical_resolution'] + gsibec_nlons = template_dict['gsibec_nlons'] + gsibec_nlats = template_dict['gsibec_nlats'] + gsibec_configuration = template_dict['gsibec_configuration'] + gsibec_npx_proc = template_dict['gsibec_npx_proc'] + gsibec_npy_proc = template_dict['gsibec_npy_proc'] + + background_error = { + 'covariance model': 'SABER', + 'covariance type': 'gsi hybrid covariance', + 'saber central block': { + 'saber block name': 'gsi hybrid covariance', + 'read': { + 'gsi akbk': f'./fv3-jedi/fv3files/akbk{vertical_resolution}.nc4', + 'gsi error covariance file': + f'./fv3-jedi/gsibec/gsi-coeffs-gmao-global-l{vertical_resolution}x{gsibec_nlons}y{gsibec_nlats}.nc4', # noqa + 'gsi berror namelist file': + f'./fv3-jedi/gsibec/{gsibec_configuration}_l{vertical_resolution}x{gsibec_nlons}y{gsibec_nlats}.nml', # noqa + 'processor layout x direction': gsibec_npx_proc, + 'processor layout y direction': gsibec_npy_proc, + 'debugging mode': False + } + }, + 'saber outer blocks': [ + {'saber block name': 'interpolation', + 'inner geometry': { + 'function space': 'StructuredColumns', + 'custom grid matching gsi': { + 'type': 'latlon', + 'lats': int(gsibec_nlats), + 'lons': int(gsibec_nlons), + }, + 'custom partitioner matching gsi': { + 'bands': gsibec_npy_proc + }, + 'halo': 1, + }, + 'forward interpolator': { + 'local interpolator type': 'oops unstructured grid interpolator', + }, + 'inverse interpolator': { + 'local interpolator type': 'oops unstructured grid interpolator', + }, + 'state variables to inverse': state_variables_to_inverse} + ], + 'linear variable change': { + 'linear variable change name': 'Control2Analysis', + 'input variables': state_variables_to_inverse, + 'output variables': template_dict['analysis_variables'] + } + } + + return background_error + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_hybrid.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py similarity index 98% rename from src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_hybrid.py rename to src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py index 321470de1..8cc8d1a8c 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_hybrid.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py @@ -13,6 +13,7 @@ field_io_names # -------------------------------------------------------------------------------------------------- +# Hybrid-B: not tested state_variables_to_inverse = [ 'eastward_wind', @@ -193,7 +194,7 @@ """ -def background_error_hybrid(template_dict: Mapping) -> Mapping: +def background_error_eda_hybridB(template_dict: Mapping) -> Mapping: render_context = { **template_dict, 'state_variables_to_inverse': state_variables_to_inverse, diff --git a/src/swell/configuration/jedi/oops/eda3D.py b/src/swell/configuration/jedi/oops/eda3D.py index 54da555c3..bb77e3fff 100644 --- a/src/swell/configuration/jedi/oops/eda3D.py +++ b/src/swell/configuration/jedi/oops/eda3D.py @@ -26,7 +26,11 @@ def render_oops(self): 'geometry': self.interface_model('geometry'), 'analysis variables': self.template_dict['analysis_variables'], 'background': self.interface_model('background_eda'), - 'background error': self.interface_model('background_error_hybrid'), + #GSI:B + 'background error': self.interface_model('background_error_eda_gsiB'), + # + #Hybrid-B: not tested + #'background error': self.interface_model('background_error_eda_hybridB'), 'observations': { 'obs perturbations': False, 'get values': self.interface_model('getvalues'), From abe060a152a64f1b53e381439c50e476bd16618f Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 10 Jun 2026 17:28:21 -0600 Subject: [PATCH 19/68] add tasks/clean_eda.py to eda/flow.cylc --- .../geos_atmosphere/model/eda_analysis.py | 2 +- src/swell/suites/eda/flow.cylc | 6 + src/swell/tasks/clean_eda.py | 156 ++++++++++++++++++ src/swell/tasks/task_questions.py | 10 ++ 4 files changed, 173 insertions(+), 1 deletion(-) create mode 100644 src/swell/tasks/clean_eda.py diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py index f40195457..6205e0818 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py @@ -19,7 +19,7 @@ def eda_analysis(template_dict: Mapping) -> Mapping: 'filetype': 'cube sphere history', 'provider': 'geos', 'datapath': f'./analysis/mem{imem:03d}', - 'filename': f'{template_dict["experiment_id"]}.ana.mem{imem:03d}.%yyyy%mm%dd_%hh%MM%ssz.nc4', + 'filename': f'eda.ana.mem{imem:03d}.%yyyy%mm%dd_%hh%MM%ssz.nc4', 'first': 'PT0H', 'frequency': 'PT1H', 'field io names': field_io_names, diff --git a/src/swell/suites/eda/flow.cylc b/src/swell/suites/eda/flow.cylc index d62bac10a..6726c49d3 100644 --- a/src/swell/suites/eda/flow.cylc +++ b/src/swell/suites/eda/flow.cylc @@ -85,6 +85,9 @@ SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} + # CleanupEda directory + SP4-{{model_component}} => CleanEdaFiles-{{model_component}} + # EvaObservations SP4-{{model_component}} => EvaObservations-{{model_component}} @@ -212,6 +215,9 @@ --{{key}} = {{value}} {%- endfor %} + [[CleanEdaFiles-{{model_component}}]] + script = "swell task CleanEda $config -d $datetime -m {{model_component}}" + [[EvaJediLog-{{model_component}}]] script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}" diff --git a/src/swell/tasks/clean_eda.py b/src/swell/tasks/clean_eda.py new file mode 100644 index 000000000..ad230dddb --- /dev/null +++ b/src/swell/tasks/clean_eda.py @@ -0,0 +1,156 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + + +# -------------------------------------------------------------------------------------------------- + +import os +import copy +import glob +import shutil +from ruamel.yaml import YAML + +from swell.tasks.base.task_base import taskBase +from swell.utilities.run_jedi_executables import run_executable + + +# -------------------------------------------------------------------------------------------------- + + +class CleanEda(taskBase): + + # ---------------------------------------------------------------------------------------------- + + def execute(self) -> None: + + # Jedi application name + # --------------------- + jedi_application = 'eda' + + # Parse configuration + # ------------------- + window_type = self.config.window_type() + window_length = self.config.window_length() + forecast_length = self.config.forecast_length(window_length) + background_time_offset = self.config.background_time_offset() + number_of_iterations = self.config.number_of_iterations() + jedi_forecast_model = self.config.jedi_forecast_model(None) + generate_yaml_and_exit = self.config.generate_yaml_and_exit(False) + perhost = self.config.perhost(None) + + # Set the observing system records path + self.jedi_rendering.set_obs_records_path(self.config.observing_system_records_path(None)) + + gsibec_nlats = self.config.gsibec_nlats(None) + gsibec_nlons = self.config.gsibec_nlons(None) + gsibec_configuration = self.config.gsibec_configuration(None) + npx_proc = self.config.npx_proc(None) + npy_proc = self.config.npy_proc(None) + npx = self.config.npx(None) + npy = self.config.npy(None) + + # Compute data assimilation window parameters + # -------------------------------------------- + background_time = self.da_window_params.background_time(background_time_offset) + local_background_time = self.da_window_params.local_background_time(window_length, + window_type) + local_background_time_iso = self.da_window_params.local_background_time_iso(window_length, + window_type) + window_begin = self.da_window_params.window_begin(window_length) + window_begin_iso = self.da_window_params.window_begin_iso(window_length) + window_end_iso = self.da_window_params.window_end_iso(window_length) + nmember = self.config.ensemble_num_members() + imember = self.get_ensemble_imember() + + # Populate jedi interface templates dictionary + # -------------------------------------------- + self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) + self.jedi_rendering.add_key('window_end_iso', window_end_iso) + self.jedi_rendering.add_key('window_length', window_length) + self.jedi_rendering.add_key('forecast_length', forecast_length) + self.jedi_rendering.add_key('minimizer', self.config.minimizer()) + self.jedi_rendering.add_key('number_of_iterations', number_of_iterations[0]) + self.jedi_rendering.add_key('analysis_variables', self.config.analysis_variables()) + self.jedi_rendering.add_key('saber_central_block', self.config.saber_central_block(None)) + self.jedi_rendering.add_key('saber_outer_block', self.config.saber_outer_block(None)) + self.jedi_rendering.add_key('gradient_norm_reduction', + self.config.gradient_norm_reduction()) + self.jedi_rendering.add_key('marine_models', self.config.marine_models(None)) + + # Background + # ---------- + self.jedi_rendering.add_key('horizontal_resolution', self.config.horizontal_resolution()) + self.jedi_rendering.add_key('local_background_time', local_background_time) + self.jedi_rendering.add_key('local_background_time_iso', local_background_time_iso) + self.jedi_rendering.add_key('ensemble_num_members', self.config.ensemble_num_members()) + self.jedi_rendering.add_key('ensemble_imember', 1) # pass an integer to jinja2 + + # Geometry + # -------- + self.jedi_rendering.add_key('vertical_resolution', self.config.vertical_resolution()) + self.jedi_rendering.add_key('gsibec_nlats', gsibec_nlats) + self.jedi_rendering.add_key('gsibec_nlons', gsibec_nlons) + self.jedi_rendering.add_key('npx_proc', npx_proc) + self.jedi_rendering.add_key('npy_proc', npy_proc) + self.jedi_rendering.add_key('npx', npx) + self.jedi_rendering.add_key('npy', npy) + self.jedi_rendering.add_key('total_processors', self.config.total_processors(None)) + + # Observations + # ------------ + self.jedi_rendering.add_key('background_time', background_time) + self.jedi_rendering.add_key('crtm_coeff_dir', self.config.crtm_coeff_dir(None)) + self.jedi_rendering.add_key('window_begin', window_begin) + + # Atmosphere background error model + # --------------------------------- + if gsibec_configuration is not None: + self.jedi_rendering.add_key('gsibec_configuration', gsibec_configuration) + self.jedi_rendering.add_key('gsibec_nlats', gsibec_nlats) + self.jedi_rendering.add_key('gsibec_nlons', gsibec_nlons) + self.jedi_rendering.add_key('gsibec_npx_proc', npx_proc) + self.jedi_rendering.add_key('gsibec_npy_proc', 6*npy_proc) + + # Model + # ----- + if window_type == '4D': + self.jedi_rendering.add_key('background_frequency', self.config.background_frequency()) + + # Open the JEDI config file and fill initial templates + # ---------------------------------------------------- + jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}{window_type}', + window_type, + jedi_forecast_model) + + for imem in range(1, nmember+1): + mem_dir = f'analysis/mem{imem:003d}/' + d1 = os.path.join(self.cycle_dir(), mem_dir) + d2 = os.path.join(d1, 'fv3-jedi') + if os.path.exists(d2): + if os.path.islink(d2): + # Only remove the symlink itself, never follow it + os.unlink(d2) + print(f"Deleted symlink dir: {d2}") + else: + shutil.rmtree(d2, ignore_errors=True) + print(f"Deleted directory: {d2}") + for observer in jedi_config_dict['cost function']['observations']['observers']: + # Get observation name + observation = observer['observation_name'] + # Delete input obsfile .nc4 files + for file_path in glob.glob(os.path.join(d1, f'{observation}*.nc4')): + if os.path.islink(file_path): + os.unlink(file_path) # safe: removes only the link + print(f"Deleted symlink file: {file_path}") + else: + try: + os.remove(file_path) + print(f"Deleted file: {file_path}") + except FileNotFoundError: + pass # file disappeared between glob and remove + + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index 570227112..b5c4e5b77 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -849,6 +849,16 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + CleanEda = QuestionList( + list_name="CleanEda", + questions=[ + run_jedi_executable, + qd.ensemble_num_members(), + ] + ) + + # -------------------------------------------------------------------------------------------------- + SaveObsDiags = QuestionList( list_name="SaveObsDiags", questions=[ From b1a607533e18c36c1a66a27340af739ff7657e5a Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 10 Jun 2026 20:52:36 -0400 Subject: [PATCH 20/68] comment out EvaJediLog and EvaObservations from eda/flow.cylc --- src/swell/suites/eda/flow.cylc | 20 ++++++++++---------- src/swell/tasks/clean_eda.py | 4 ++-- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/swell/suites/eda/flow.cylc b/src/swell/suites/eda/flow.cylc index 6726c49d3..6fe6b3b81 100644 --- a/src/swell/suites/eda/flow.cylc +++ b/src/swell/suites/eda/flow.cylc @@ -88,19 +88,19 @@ # CleanupEda directory SP4-{{model_component}} => CleanEdaFiles-{{model_component}} - # EvaObservations - SP4-{{model_component}} => EvaObservations-{{model_component}} - - # EvaJediLog - SP4-{{model_component}} => EvaJediLog-{{model_component}} - # EvaIncrement SP4-{{model_component}} => EvaIncrement-{{model_component}} - # Clean up large files - EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} & - EvaObservations-{{model_component}} => - CleanCycle-{{model_component}} +# # EvaObservations +# SP4-{{model_component}} => EvaObservations-{{model_component}} +# +# # EvaJediLog +# SP4-{{model_component}} => EvaJediLog-{{model_component}} +# +# # Clean up large files +# EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} & +# EvaObservations-{{model_component}} => +# CleanCycle-{{model_component}} {% endif %} {% endfor %} diff --git a/src/swell/tasks/clean_eda.py b/src/swell/tasks/clean_eda.py index ad230dddb..079337c77 100644 --- a/src/swell/tasks/clean_eda.py +++ b/src/swell/tasks/clean_eda.py @@ -140,8 +140,8 @@ def execute(self) -> None: for observer in jedi_config_dict['cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] - # Delete input obsfile .nc4 files - for file_path in glob.glob(os.path.join(d1, f'{observation}*.nc4')): + # Delete input obsfile in analysis/mem00x (.nc4 .tlapse.txt acftbias acftbias_cov) + for file_path in glob.glob(os.path.join(d1, f'{observation}.*')): if os.path.islink(file_path): os.unlink(file_path) # safe: removes only the link print(f"Deleted symlink file: {file_path}") From 1c01bc793d419baebab89265de339d1cb3b5c656 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 11 Jun 2026 12:29:28 -0600 Subject: [PATCH 21/68] fix issues from pycodestyle_run --- .../geos_atmosphere/model/background_eda.py | 2 +- .../model/background_error_eda_gsiB.py | 1 - .../model/background_error_eda_hybridB.py | 17 ++++++----- .../geos_atmosphere/model/comp_mean_output.py | 8 ++--- .../model/comp_variance_output.py | 3 +- .../model/diffstates_output.py | 4 +-- .../geos_atmosphere/model/shared.py | 20 ++++++------- .../geos_atmosphere/model/two_states.py | 3 +- src/swell/configuration/jedi/oops/eda3D.py | 6 ++-- src/swell/suites/eda/suite_config.py | 1 - src/swell/tasks/run_jedi_diffstates.py | 15 +++++----- src/swell/tasks/run_jedi_eda_executable.py | 22 +++++++------- .../tasks/run_jedi_ensemble_mean_variance.py | 29 +++++++++---------- src/swell/utilities/question_defaults.py | 2 +- 14 files changed, 65 insertions(+), 68 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py index 61de175f0..71dd9d893 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda.py @@ -61,7 +61,7 @@ def background_eda(template_dict: Mapping) -> Mapping: horizontal_resolution = template_dict['horizontal_resolution'] - imem=template_dict['ensemble_imember'] + imem = template_dict['ensemble_imember'] background = { 'datetime': template_dict['local_background_time_iso'], diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py index 586aa28f0..7ff47d587 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py @@ -15,7 +15,6 @@ 'northward_wind', 'air_temperature', 'air_pressure_at_surface', -# 'air_pressure_levels', 'water_vapor_mixing_ratio_wrt_moist_air', 'cloud_liquid_ice', 'cloud_liquid_water', diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py index 8cc8d1a8c..5c5cfa8ee 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py @@ -7,20 +7,18 @@ # -------------------------------------------------------------------------------------------------- from ruamel.yaml import YAML from jinja2 import Environment -##from typing import Mapping from collections.abc import Mapping from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ field_io_names # -------------------------------------------------------------------------------------------------- -# Hybrid-B: not tested +# Hybrid-B from rtodling: not tested state_variables_to_inverse = [ 'eastward_wind', 'northward_wind', 'air_temperature', 'air_pressure_at_surface', -# 'air_pressure_levels', 'water_vapor_mixing_ratio_wrt_moist_air', 'cloud_liquid_ice', 'cloud_liquid_water', @@ -31,7 +29,7 @@ 'fraction_of_lake', 'fraction_of_ice', 'geopotential_height_times_gravity_at_surface', - 'skin_temperature_at_surface', + 'skin_temperature_at_surface' ] # -------------------------------------------------------------------------------------------------- @@ -46,8 +44,12 @@ saber block name: gsi hybrid covariance read: gsi akbk: ./fv3-jedi/fv3files/akbk{{vertical_resolution}}.nc4 - gsi error covariance file: ./fv3-jedi/gsibec/gsi-coeffs-gmao-global-l{{vertical_resolution}}x{{gsibec_nlons}}y{{gsibec_nlats}}.nc4 - gsi berror namelist file: ./fv3-jedi/gsibec/{{gsibec_configuration}}_l{{vertical_resolution}}x{{gsibec_nlons}}y{{gsibec_nlats}}.nml + gsi error covariance file: >- + ./fv3-jedi/gsibec/gsi-coeffs-gmao-global- + l{{vertical_resolution}}x{{gsibec_nlons}}y{{gsibec_nlats}}.nc4 + gsi berror namelist file: >- + ./fv3-jedi/gsibec/{{gsibec_configuration}}_ + l{{vertical_resolution}}x{{gsibec_nlons}}y{{gsibec_nlats}}.nml processor layout x direction: {{gsibec_npx_proc}} processor layout y direction: {{gsibec_npy_proc}} debugging mode: false @@ -210,7 +212,8 @@ def background_error_eda_hybridB(template_dict: Mapping) -> Mapping: # Use ruamel.yaml to load the string instead of standard pyyaml ruamel_yaml = YAML() background_error = ruamel_yaml.load(rendered_yaml_string) - cov_template = background_error['components'][1]['covariance']['members from template']['template'] + cov_template = background_error['components'][1]['covariance'] + ['members from template']['template'] cov_template['field io names'] = field_io_names a = background_error['components'][0]['covariance']['linear variable change'] a['output variables'] = template_dict['analysis_variables'] diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py index 5ecc0d812..9ce9625ef 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_mean_output.py @@ -14,10 +14,10 @@ def comp_mean_output(template_dict: Mapping) -> Mapping: - + prefix_output_mean = template_dict['ensmeanvariance_spec_item'].get('fn_output_mean') state = template_dict['ensmeanvariance_spec_item'].get('state') - grid_type= template_dict['ensmeanvariance_spec_gridtype'] + grid_type = template_dict['ensmeanvariance_spec_gridtype'] mean_output = {} if grid_type == 'cs': @@ -37,15 +37,13 @@ def comp_mean_output(template_dict: Mapping) -> Mapping: 'datapath': template_dict['cycle_dir'], 'filename': f'{prefix_output_mean}.ll.', 'field io names': [] - } + } if state in ['bkg', 'analysis']: mean_output['field io names'] = field_io_names else: mean_output['field io names'] = field_io_names_ensemble - return mean_output - # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py index af0d471ad..4c359d11a 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/comp_variance_output.py @@ -17,7 +17,7 @@ def comp_variance_output(template_dict: Mapping) -> Mapping: prefix_output_variance = template_dict['ensmeanvariance_spec_item'].get('fn_output_variance') state = template_dict['ensmeanvariance_spec_item'].get('state') - grid_type= template_dict['ensmeanvariance_spec_gridtype'] + grid_type = template_dict['ensmeanvariance_spec_gridtype'] variance_output = {} if grid_type == 'cs': @@ -44,5 +44,4 @@ def comp_variance_output(template_dict: Mapping) -> Mapping: return variance_output - # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py index 578f75f06..36119c395 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py @@ -16,11 +16,11 @@ def diffstates_output(template_dict: Mapping) -> Mapping: # diff states (inc) - grid_type= template_dict['diffstates_spec_gridtype'] + grid_type = template_dict['diffstates_spec_gridtype'] prefix_output = template_dict['diffstates_spec'].get('state_diff', {}).get('fn_output') stateType = template_dict.get('diffstates_spec_statetype') - if [ stateType == 'ensemble' ]: + if [stateType == 'ensemble']: field_io_names_loc = field_io_names_sa1 else: field_io_names_loc = field_io_names diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py index c5a250d7a..e2e0335a5 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/shared.py @@ -147,16 +147,16 @@ # stand alone 1 state_variables_sa1 = [ - 'eastward_wind', - 'northward_wind', - 'air_temperature', - 'air_pressure_at_surface', - 'air_pressure_levels', - 'water_vapor_mixing_ratio_wrt_moist_air', - 'cloud_liquid_ice', - 'cloud_liquid_water', - 'mole_fraction_of_ozone_in_air', - 'skin_temperature_at_surface', + 'eastward_wind', + 'northward_wind', + 'air_temperature', + 'air_pressure_at_surface', + 'air_pressure_levels', + 'water_vapor_mixing_ratio_wrt_moist_air', + 'cloud_liquid_ice', + 'cloud_liquid_water', + 'mole_fraction_of_ozone_in_air', + 'skin_temperature_at_surface', ] # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py index 035244347..92d936658 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py @@ -12,13 +12,14 @@ # -------------------------------------------------------------------------------------------------- + def two_states(template_dict: Mapping) -> Mapping: horizontal_resolution = template_dict['horizontal_resolution'] f1 = template_dict['diffstates_spec']['state1']['fn_input'] f2 = template_dict['diffstates_spec']['state2']['fn_input'] stateType = template_dict.get('diffstates_spec_statetype') - if [ stateType == 'ensemble' ]: + if [stateType == 'ensemble']: state_variables_loc = state_variables_sa1 field_io_names_loc = field_io_names_sa1 else: diff --git a/src/swell/configuration/jedi/oops/eda3D.py b/src/swell/configuration/jedi/oops/eda3D.py index bb77e3fff..e811ec356 100644 --- a/src/swell/configuration/jedi/oops/eda3D.py +++ b/src/swell/configuration/jedi/oops/eda3D.py @@ -26,11 +26,11 @@ def render_oops(self): 'geometry': self.interface_model('geometry'), 'analysis variables': self.template_dict['analysis_variables'], 'background': self.interface_model('background_eda'), - #GSI:B + # GSI:B 'background error': self.interface_model('background_error_eda_gsiB'), # - #Hybrid-B: not tested - #'background error': self.interface_model('background_error_eda_hybridB'), + # Hybrid-B: not tested + # 'background error': self.interface_model('background_error_eda_hybridB'), 'observations': { 'obs perturbations': False, 'get values': self.interface_model('getvalues'), diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda/suite_config.py index f5276285c..29723c323 100644 --- a/src/swell/suites/eda/suite_config.py +++ b/src/swell/suites/eda/suite_config.py @@ -57,7 +57,6 @@ class SuiteConfig(QuestionContainer, Enum): "air_temperature", "water_vapor_mixing_ratio_wrt_moist_air", "air_pressure_at_surface", -# "air_pressure_levels", "cloud_liquid_ice", "cloud_liquid_water", "rain_water", diff --git a/src/swell/tasks/run_jedi_diffstates.py b/src/swell/tasks/run_jedi_diffstates.py index 25eb046fd..40fc2747b 100644 --- a/src/swell/tasks/run_jedi_diffstates.py +++ b/src/swell/tasks/run_jedi_diffstates.py @@ -48,7 +48,6 @@ def execute(self) -> None: # Set the observing system records path self.jedi_rendering.set_obs_records_path(self.config.observing_system_records_path(None)) - # Populate jedi interface templates dictionary # -------------------------------------------- self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) @@ -69,7 +68,7 @@ def execute(self) -> None: # Diffstates spec_dict = self.config.diffstates_spec() self.jedi_rendering.add_key('diffstates_spec', spec_dict) - print( f'diffstates = {spec_dict}') + print(f'diffstates = {spec_dict}') # Add placeholder names if mock experiment # ---------------------------------------- @@ -83,7 +82,7 @@ def execute(self) -> None: # loop output grid_type: grid_type = spec_dict['state_diff'].get('grid_type') - print (f'grid_type = {grid_type}') + print(f'grid_type = {grid_type}') if grid_type is None: grid_type = ['latlon'] @@ -94,11 +93,13 @@ def execute(self) -> None: # Jedi configuration file # ----------------------- - jedi_config_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_output_{output_grid}_config.yaml') + jedi_config_file = os.path.join( + self.cycle_dir(), f'jedi_{jedi_application}_output_{output_grid}_config.yaml') # Output log file # --------------- - output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_output_{output_grid}.log') + output_log_file = os.path.join(self.cycle_dir(), + f'jedi_{jedi_application}_output_{output_grid}.log') # Open the JEDI config file and fill templates # -------------------------------------------- @@ -124,8 +125,8 @@ def execute(self) -> None: # Jedi executable name # -------------------- jedi_executable = model_component_meta['executables'][f'{jedi_application}'] - jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', 'build', 'bin', - jedi_executable) + jedi_executable_path = os.path.join(self.experiment_path(), + 'jedi_bundle', 'build', 'bin', jedi_executable) # Run the JEDI executable # ----------------------- diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 0d6178104..4afa9a883 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -64,7 +64,6 @@ def execute(self) -> None: window_end_iso = self.da_window_params.window_end_iso(window_length) nmember = self.config.ensemble_num_members() imember = self.get_ensemble_imember() -### print( f' imember = {imember}' ) # Populate jedi interface templates dictionary # -------------------------------------------- @@ -120,14 +119,15 @@ def execute(self) -> None: if window_type == '4D': self.jedi_rendering.add_key('background_frequency', self.config.background_frequency()) - # Jedi configuration file # ----------------------- - jedi_config_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}{window_type}_config_mem{imember:03d}.yaml') + jedi_config_file = os.path.join( + self.cycle_dir(), f'jedi_{jedi_application}{window_type}_config_mem{imember:03d}.yaml') # Output log file # --------------- - output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}{window_type}_log_mem{imember:03d}.log') + output_log_file = os.path.join( + self.cycle_dir(), f'jedi_{jedi_application}{window_type}_log_mem{imember:03d}.log') # Open the JEDI config file and fill initial templates # ---------------------------------------------------- @@ -150,21 +150,20 @@ def execute(self) -> None: for observer in jedi_config_dict['cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] - - print( f'ob= {observation}' ) + print(f'ob= {observation}') # copy obs input file to avoid multi MPI reading the same file files = glob.glob(os.path.join(self.cycle_dir(), f'{observation}.*')) for src_file in files: - print( f'f= {src_file}' ) - shutil.copy(src_file, xdir) + print(f'f= {src_file}') + shutil.copy(src_file, xdir) if imember > 1: obs_cov_model = observer.get('obs error', {}).get('covariance model') - print( f'{observation}: obs_cov_model = {obs_cov_model}') + print(f'{observation}: obs_cov_model = {obs_cov_model}') if obs_cov_model and 'cross variable covariances' in obs_cov_model: print(f"Found cross covariance obs: {obs_cov_model}, skip perturbation") else: - print(f"No cross varaible covariance found for {observation}, Obs Error Diagonal") + print(f"No cross varaible covariance found: {observation}, Obs Error Diagonal") obs_error_dict = { 'covariance model': 'diagonal', 'zero-mean perturbations': True, @@ -201,8 +200,7 @@ def execute(self) -> None: dir1, fname = os.path.split(File) File = os.path.join(dir1, mem_dir, fname) obs_bias['covariance']['output file'] = File - # - File = obs_bias.get('covariance', {}).get('prior',{}).get('input file') + File = obs_bias.get('covariance', {}).get('prior', {}).get('input file') if File is not None: dir1, fname = os.path.split(File) File = os.path.join(dir1, mem_dir, fname) diff --git a/src/swell/tasks/run_jedi_ensemble_mean_variance.py b/src/swell/tasks/run_jedi_ensemble_mean_variance.py index 1e32aebb4..0e5a4f553 100644 --- a/src/swell/tasks/run_jedi_ensemble_mean_variance.py +++ b/src/swell/tasks/run_jedi_ensemble_mean_variance.py @@ -51,7 +51,6 @@ def execute(self) -> None: # Ensemble # ------------------------ - # Populate jedi interface templates dictionary # -------------------------------------------- self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) @@ -72,8 +71,7 @@ def execute(self) -> None: # Ensemble self.jedi_rendering.add_key('ensemble_num_members', self.config.ensemble_num_members(None)) - - print( f'self.config.ensmeanvariance_spec = {self.config.ensmeanvariance_spec()}') + print(f'self.config.ensmeanvariance_spec = {self.config.ensmeanvariance_spec()}') meanvar_spec_dict = self.config.ensmeanvariance_spec() # # Loop directly over the items in the list # for spec in config_dict: @@ -87,8 +85,6 @@ def execute(self) -> None: # print(f"Output Mean: {fn_output_mean}") # print(f"Output Var: {fn_output_variance}") # print("-" * 40) -# - # Add placeholder names if mock experiment # ---------------------------------------- @@ -97,25 +93,27 @@ def execute(self) -> None: self.jedi_rendering.add_key('experiment_id', 'experiment_id') self.jedi_rendering.add_key('cycle_dir', 'cycle_dir') - - # loop item in ensmeanvariance_spec: state + grid_type + # loop item in ensmeanvariance_spec: state + grid_type # ----------------------- for idx, spec in enumerate(meanvar_spec_dict, start=1): self.jedi_rendering.add_key('ensmeanvariance_spec_item', spec) state_name = spec.get('state') - for output_grid in spec.get('grid_type'): + for output_grid in spec.get('grid_type'): + + self.jedi_rendering.add_key('ensmeanvariance_spec_gridtype', output_grid) - self.jedi_rendering.add_key('ensmeanvariance_spec_gridtype', output_grid) - # Jedi configuration file # ----------------------- - jedi_config_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_config_{state_name}_{output_grid}.yaml') + jedi_config_file = os.path.join( + self.cycle_dir(), + f'jedi_{jedi_application}_config_{state_name}_{output_grid}.yaml') # Output log file # --------------- - output_log_file = os.path.join(self.cycle_dir(), f'jedi_{jedi_application}_{state_name}_{output_grid}.log') + output_log_file = os.path.join( + self.cycle_dir(), f'jedi_{jedi_application}_{state_name}_{output_grid}.log') # Open the JEDI config file and fill templates # -------------------------------------------- @@ -142,13 +140,14 @@ def execute(self) -> None: # Jedi executable name # -------------------- jedi_executable = model_component_meta['executables'][f'{jedi_application}'] - jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', 'build', 'bin', - jedi_executable) + jedi_executable_path = os.path.join(self.experiment_path(), + 'jedi_bundle', 'build', 'bin', jedi_executable) # Run the JEDI executable # ----------------------- if not generate_yaml_and_exit: - self.logger.info('Running '+jedi_executable_path+' with '+str(np)+' processors.') + self.logger.info( + 'Running '+jedi_executable_path+' with '+str(np)+' processors.') run_executable(self.logger, self.cycle_dir(), np, jedi_executable_path, jedi_config_file, output_log_file) else: diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index e02f7900f..9628427ee 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -599,7 +599,7 @@ class ensmeanvariance_spec(TaskQuestion): "all_models" ]) ask_question: bool = True - prompt: str = "Configure the ensemble mean and variance specifications: [ebkg, eda.analysis]" + prompt: str = "Configure the ensemble mean and variance specifications:" widget_type: WType = WType.STRING # -------------------------------------------------------------------------------------------------- From 26337c696f97204de2f8a0b613131d9ce36818bc Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 11 Jun 2026 15:50:38 -0600 Subject: [PATCH 22/68] fix some ctest errors --- src/swell/tasks/clean_eda.py | 8 +------- src/swell/tasks/run_jedi_eda_executable.py | 1 - src/swell/utilities/mock_jedi_config.py | 4 ++-- 3 files changed, 3 insertions(+), 10 deletions(-) diff --git a/src/swell/tasks/clean_eda.py b/src/swell/tasks/clean_eda.py index 079337c77..85aab398f 100644 --- a/src/swell/tasks/clean_eda.py +++ b/src/swell/tasks/clean_eda.py @@ -8,13 +8,9 @@ # -------------------------------------------------------------------------------------------------- import os -import copy import glob import shutil -from ruamel.yaml import YAML - from swell.tasks.base.task_base import taskBase -from swell.utilities.run_jedi_executables import run_executable # -------------------------------------------------------------------------------------------------- @@ -38,8 +34,6 @@ def execute(self) -> None: background_time_offset = self.config.background_time_offset() number_of_iterations = self.config.number_of_iterations() jedi_forecast_model = self.config.jedi_forecast_model(None) - generate_yaml_and_exit = self.config.generate_yaml_and_exit(False) - perhost = self.config.perhost(None) # Set the observing system records path self.jedi_rendering.set_obs_records_path(self.config.observing_system_records_path(None)) @@ -63,7 +57,7 @@ def execute(self) -> None: window_begin_iso = self.da_window_params.window_begin_iso(window_length) window_end_iso = self.da_window_params.window_end_iso(window_length) nmember = self.config.ensemble_num_members() - imember = self.get_ensemble_imember() + ## imember = self.get_ensemble_imember() # Populate jedi interface templates dictionary # -------------------------------------------- diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 4afa9a883..1a9620608 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -8,7 +8,6 @@ # -------------------------------------------------------------------------------------------------- import os -import copy import glob import shutil from ruamel.yaml import YAML diff --git a/src/swell/utilities/mock_jedi_config.py b/src/swell/utilities/mock_jedi_config.py index c851c4513..30133a169 100644 --- a/src/swell/utilities/mock_jedi_config.py +++ b/src/swell/utilities/mock_jedi_config.py @@ -57,7 +57,7 @@ def mock_jedi_config(suite: str, f'{experiment_id}-suite', 'experiment.yaml') task_wrapper('RenderJediObservations', experiment_yaml, datetime, - model, ensemblePacket=None) + model, ensemblePacket=None, imember=None) if executable_type == 'localensembleda': executable_name = 'LocalEnsembleDa' @@ -65,7 +65,7 @@ def mock_jedi_config(suite: str, executable_name = executable_type.capitalize() task_wrapper(f'RunJedi{executable_name}Executable', experiment_yaml, datetime, - model, ensemblePacket=None) + model, ensemblePacket=None, imember=None) cycle_dir = os.path.join(tempdir, experiment_id, 'run', datetime, model) From 9f23c0805b0c1db730a938c8de35215285ef5d85 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 11 Jun 2026 22:45:23 -0600 Subject: [PATCH 23/68] 1. 'local_background_time_iso' 2. qd.window_length("PT6H"), 3. geos_atmosphere background_time_offset: default_value: PT9H --- .../interfaces/geos_atmosphere/model/background_ensemble.py | 4 +++- src/swell/suites/localensembleda/suite_config.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py index 9b904fb41..f6c279aa6 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py @@ -17,7 +17,9 @@ def background_ensemble(template_dict: Mapping) -> Mapping: horizontal_resolution = template_dict['horizontal_resolution'] background_ensemble = { - 'date': template_dict['window_begin_iso'], +# 4D-window +# 'date': template_dict['window_begin_iso'], + 'date': template_dict['local_background_time_iso'], 'members from template': { 'template': { 'datetime': template_dict['local_background_time_iso'], diff --git a/src/swell/suites/localensembleda/suite_config.py b/src/swell/suites/localensembleda/suite_config.py index e4c670272..e80f67cb0 100644 --- a/src/swell/suites/localensembleda/suite_config.py +++ b/src/swell/suites/localensembleda/suite_config.py @@ -42,7 +42,6 @@ class SuiteConfig(QuestionContainer, Enum): qd.npx_proc(3), qd.npy_proc(3), qd.cycle_times(['T00']), - qd.background_time_offset("PT3H"), qd.ensemble_num_members(3), qd.skip_ensemble_hofx(True), qd.local_ensemble_solver("Deterministic GETKF"), @@ -56,6 +55,7 @@ class SuiteConfig(QuestionContainer, Enum): qd.observations([ "atms_n20", ]), + qd.window_length("PT6H"), qd.window_type("3D"), qd.clean_patterns(['*.txt']) ] @@ -83,7 +83,6 @@ class SuiteConfig(QuestionContainer, Enum): qd.npy_proc(4), # qd.perhost(32), qd.cycle_times(['T00']), - qd.background_time_offset("PT3H"), qd.ensemble_num_members(16), qd.skip_ensemble_hofx(True), qd.local_ensemble_solver("Deterministic GETKF"), @@ -123,6 +122,7 @@ class SuiteConfig(QuestionContainer, Enum): "amsua_metop-b", "amsua_metop-c" ]), + qd.window_length("PT6H"), qd.window_type("3D"), qd.clean_patterns(['*.txt']) ] From 04e5f710f5648a24cb53f62d35fe3b07b0beb6ab Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 11 Jun 2026 22:55:28 -0600 Subject: [PATCH 24/68] . --- .../interfaces/geos_atmosphere/model/background_ensemble.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py index f6c279aa6..b4027016c 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_ensemble.py @@ -17,8 +17,6 @@ def background_ensemble(template_dict: Mapping) -> Mapping: horizontal_resolution = template_dict['horizontal_resolution'] background_ensemble = { -# 4D-window -# 'date': template_dict['window_begin_iso'], 'date': template_dict['local_background_time_iso'], 'members from template': { 'template': { @@ -43,4 +41,7 @@ def background_ensemble(template_dict: Mapping) -> Mapping: return background_ensemble +# Note: 4D-window +# 'date': template_dict['window_begin_iso'], + # -------------------------------------------------------------------------------------------------- From ee7f59847efb312b337e0a5bff172c9ba6583b82 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 11 Jun 2026 23:34:17 -0600 Subject: [PATCH 25/68] . --- .../geos_atmosphere/model/ensemble_mean_increment_output.py | 4 ++-- .../jedi/interfaces/geos_atmosphere/model/two_states.py | 1 - src/swell/tasks/clean_eda.py | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_mean_increment_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_mean_increment_output.py index 0c69152ce..c48884257 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_mean_increment_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/ensemble_mean_increment_output.py @@ -7,7 +7,7 @@ # -------------------------------------------------------------------------------------------------- from collections.abc import Mapping -from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import field_io_names_ensemble +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import field_io_names_sa1 # -------------------------------------------------------------------------------------------------- @@ -20,7 +20,7 @@ def ensemble_mean_increment_output(template_dict: Mapping) -> Mapping: 'filetype': 'auxgrid', 'gridtype': 'latlon', 'filename': f'{cycle_dir}/geos.mean-inc.', - 'field io names': field_io_names_ensemble + 'field io names': field_io_names_sa1 } return ensemble_mean_increment_output diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py index 92d936658..e81146ba5 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py @@ -14,7 +14,6 @@ def two_states(template_dict: Mapping) -> Mapping: - horizontal_resolution = template_dict['horizontal_resolution'] f1 = template_dict['diffstates_spec']['state1']['fn_input'] f2 = template_dict['diffstates_spec']['state2']['fn_input'] diff --git a/src/swell/tasks/clean_eda.py b/src/swell/tasks/clean_eda.py index 85aab398f..be5ba32ad 100644 --- a/src/swell/tasks/clean_eda.py +++ b/src/swell/tasks/clean_eda.py @@ -57,7 +57,7 @@ def execute(self) -> None: window_begin_iso = self.da_window_params.window_begin_iso(window_length) window_end_iso = self.da_window_params.window_end_iso(window_length) nmember = self.config.ensemble_num_members() - ## imember = self.get_ensemble_imember() + # imember = self.get_ensemble_imember() # Populate jedi interface templates dictionary # -------------------------------------------- From 35e846ce2371c9ac360ed13cf852082928bf922e Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Fri, 12 Jun 2026 00:24:48 -0600 Subject: [PATCH 26/68] Due to forthcoming implementation changes in swell localensembleda, 1. remove localensembleda from tier-1 test 2. disable localensembleda from code test --- .../jedi/interfaces/geos_atmosphere/model/eda_analysis.py | 2 +- .../interfaces/geos_atmosphere/model/eda_varincrement1.py | 2 +- src/swell/swell.py | 4 ++-- src/swell/utilities/scripts/create_mock_configs.py | 7 ++++--- 4 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py index 6205e0818..924957ca9 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis.py @@ -8,7 +8,7 @@ from collections.abc import Mapping from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import field_io_names -import os + # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_varincrement1.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_varincrement1.py index df885008d..d6f02a837 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_varincrement1.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_varincrement1.py @@ -7,7 +7,7 @@ # -------------------------------------------------------------------------------------------------- from collections.abc import Mapping -import os + # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/swell.py b/src/swell/swell.py index e7c669171..2515d05a7 100644 --- a/src/swell/swell.py +++ b/src/swell/swell.py @@ -270,9 +270,9 @@ def test(test: str) -> None: @click.option('-p', '--platform', 'platform', type=click.Choice(get_platforms()), default="nccs_discover_sles15", help=platform_help) @click.argument('suite', type=click.Choice(("hofx", "3dvar_marine", "3dvar_atmos", - "localensembleda", "3dvar_cycle"))) + "3dvar_cycle"))) def t1test( - suite: Literal["hofx", "3dvar_marine", "3dvar_atmos", "localensembleda", "3dvar_cycle"], + suite: Literal["hofx", "3dvar_marine", "3dvar_atmos", "3dvar_cycle"], platform: Optional[str] = "nccs_discover_sles15" ) -> None: """ diff --git a/src/swell/utilities/scripts/create_mock_configs.py b/src/swell/utilities/scripts/create_mock_configs.py index 7faad4d3b..3d180d300 100644 --- a/src/swell/utilities/scripts/create_mock_configs.py +++ b/src/swell/utilities/scripts/create_mock_configs.py @@ -51,9 +51,10 @@ 'model': 'geos_cf', 'executable_type': 'variational'} -defaults_dict['localensembleda'] = {'datetime': atmosphere_default_datetime, - 'model': 'geos_atmosphere', - 'executable_type': 'localensembleda'} +# WIP: not ready for ctest +# defaults_dict['localensembleda'] = {'datetime': atmosphere_default_datetime, +# 'model': 'geos_atmosphere', +# 'executable_type': 'localensembleda'} # -------------------------------------------------------------------------------------------------- From 7cf82724ff17dd5b4f4192e78066f65cea40d821 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 17 Jun 2026 19:51:39 -0400 Subject: [PATCH 27/68] deleted replica codes and commented lines; fix typos add notes to logic tree in flow.cylc --- src/swell/suites/eda/flow.cylc | 26 +++++-------------- src/swell/tasks/run_jedi_eda_executable.py | 12 +-------- .../tasks/run_jedi_ensemble_mean_variance.py | 12 --------- .../run_jedi_local_ensemble_da_executable.py | 17 +----------- .../utilities/scripts/create_mock_configs.py | 5 ---- 5 files changed, 8 insertions(+), 64 deletions(-) diff --git a/src/swell/suites/eda/flow.cylc b/src/swell/suites/eda/flow.cylc index 6fe6b3b81..da3221f0c 100644 --- a/src/swell/suites/eda/flow.cylc +++ b/src/swell/suites/eda/flow.cylc @@ -46,6 +46,12 @@ {% for model_component in model_components %} {% if cycle_time[model_component] %} + # logic tree: + # prep (clone, stage, build, get bkg, get obs) -> SP1 + # -> mksi/observation + filter -> SP2 + # -> [ submit/wait for N member runs ]-> SP3 + # -> mean/variance + diffstate -> SP4 -> eva + # Task triggers for: {{model_component}} # ------------------ # Generate satellite channel records @@ -91,17 +97,6 @@ # EvaIncrement SP4-{{model_component}} => EvaIncrement-{{model_component}} -# # EvaObservations -# SP4-{{model_component}} => EvaObservations-{{model_component}} -# -# # EvaJediLog -# SP4-{{model_component}} => EvaJediLog-{{model_component}} -# -# # Clean up large files -# EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} & -# EvaObservations-{{model_component}} => -# CleanCycle-{{model_component}} - {% endif %} {% endfor %} """ @@ -179,15 +174,6 @@ {% endfor %} {% endif %} -# [[RunJediEdaExecutable-{{model_component}}]] -# script = "swell task RunJediEdaExecutable $config -d $datetime -m {{model_component}}" -# platform = {{platform}} -# execution time limit = {{scheduling["RunJediEdaExecutable"]["execution_time_limit"]}} -# [[[directives]]] -# {%- for key, value in scheduling["RunJediEdaExecutable"]["directives"][model_component].items() %} -# --{{key}} = {{value}} -# {%- endfor %} - [[RunJediObsfiltersExecutable-{{model_component}}]] script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" platform = {{platform}} diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 1a9620608..2486972d6 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -162,7 +162,7 @@ def execute(self) -> None: if obs_cov_model and 'cross variable covariances' in obs_cov_model: print(f"Found cross covariance obs: {obs_cov_model}, skip perturbation") else: - print(f"No cross varaible covariance found: {observation}, Obs Error Diagonal") + print(f"No cross variable covariance found: {observation}, Obs Error Diagonal") obs_error_dict = { 'covariance model': 'diagonal', 'zero-mean perturbations': True, @@ -256,14 +256,4 @@ def execute(self) -> None: print(f'intended mpi_command = {mpi_command}') self.logger.info('YAML generated, now exiting.') - -# clean up dir -# for observer in jedi_config_dict['cost function']['observations']['observers']: -# # Get observation name -# observation = observer['observation_name'] -# # Delete input obsfile .nc4 files -# for file_path in glob.glob(os.path.join(mem_dir, f'{observation}*.nc4')): -# os.remove(file_path) -# print(f"Deleted: {file_path}") - # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/run_jedi_ensemble_mean_variance.py b/src/swell/tasks/run_jedi_ensemble_mean_variance.py index 0e5a4f553..24f571e82 100644 --- a/src/swell/tasks/run_jedi_ensemble_mean_variance.py +++ b/src/swell/tasks/run_jedi_ensemble_mean_variance.py @@ -73,18 +73,6 @@ def execute(self) -> None: print(f'self.config.ensmeanvariance_spec = {self.config.ensmeanvariance_spec()}') meanvar_spec_dict = self.config.ensmeanvariance_spec() -# # Loop directly over the items in the list -# for spec in config_dict: -# # Unpack variables directly from the current dictionary item -# fn_input = spec["fn_input"] -# fn_output_mean = spec["fn_output_mean"] -# fn_output_variance = spec["fn_output_variance"] -# -# # Use the variables directly -# print(f"Input Path: {fn_input}") -# print(f"Output Mean: {fn_output_mean}") -# print(f"Output Var: {fn_output_variance}") -# print("-" * 40) # Add placeholder names if mock experiment # ---------------------------------------- diff --git a/src/swell/tasks/run_jedi_local_ensemble_da_executable.py b/src/swell/tasks/run_jedi_local_ensemble_da_executable.py index 30bc0db45..754df2319 100644 --- a/src/swell/tasks/run_jedi_local_ensemble_da_executable.py +++ b/src/swell/tasks/run_jedi_local_ensemble_da_executable.py @@ -14,26 +14,11 @@ from swell.swell_path import get_swell_path from swell.tasks.base.task_base import taskBase from swell.utilities.run_jedi_executables import run_executable +from swell.utilities.yaml_utils import replace_key # -------------------------------------------------------------------------------------------------- -def replace_key(obj, old_key, new_key): - """ - Recursively replace dictionary keys in nested dictionaries/lists. - """ - if isinstance(obj, dict): - new_dict = {} - for k, v in obj.items(): - new_k = new_key if k == old_key else k - new_dict[new_k] = replace_key(v, old_key, new_key) - return new_dict - elif isinstance(obj, list): - return [replace_key(item, old_key, new_key) for item in obj] - else: - return obj - - class RunJediLocalEnsembleDaExecutable(taskBase): # ---------------------------------------------------------------------------------------------- diff --git a/src/swell/utilities/scripts/create_mock_configs.py b/src/swell/utilities/scripts/create_mock_configs.py index 3d180d300..079b80b6e 100644 --- a/src/swell/utilities/scripts/create_mock_configs.py +++ b/src/swell/utilities/scripts/create_mock_configs.py @@ -51,11 +51,6 @@ 'model': 'geos_cf', 'executable_type': 'variational'} -# WIP: not ready for ctest -# defaults_dict['localensembleda'] = {'datetime': atmosphere_default_datetime, -# 'model': 'geos_atmosphere', -# 'executable_type': 'localensembleda'} - # -------------------------------------------------------------------------------------------------- From e885d367cdf4861f852d8f1925a0db9fe72bfbf7 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 25 Jun 2026 08:57:08 -0600 Subject: [PATCH 28/68] update with additional_param --- .../interfaces/geos_atmosphere/model/geometry_wofms.py | 1 + src/swell/swell.py | 2 +- src/swell/tasks/base/task_base.py | 9 +++++---- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/geometry_wofms.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/geometry_wofms.py index 3a14a568c..3ca37c3f6 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/geometry_wofms.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/geometry_wofms.py @@ -10,6 +10,7 @@ # -------------------------------------------------------------------------------------------------- + # geometry wo 'fms initialization' segment def geometry_wofms(template_dict: Mapping) -> Mapping: diff --git a/src/swell/swell.py b/src/swell/swell.py index 1657e20bd..02501f4a6 100644 --- a/src/swell/swell.py +++ b/src/swell/swell.py @@ -216,7 +216,7 @@ def task( datetime: Optional[str], model: Optional[str], additional_parameter: Optional[str], - ensemblePacket: Optional[str] + ensemblePacket: Optional[str], imember: Optional[int] ) -> None: """ diff --git a/src/swell/tasks/base/task_base.py b/src/swell/tasks/base/task_base.py index cff3dc562..06e365eb0 100644 --- a/src/swell/tasks/base/task_base.py +++ b/src/swell/tasks/base/task_base.py @@ -42,8 +42,8 @@ def __init__( model: str, ensemblePacket: Optional[str], additional_parameter: Optional[str], + imember: int | None, task_name: str - imember: int | None = None, ) -> None: # Create message logger @@ -293,7 +293,7 @@ def create_task( model: str, additional_parameter: str | None, ensemblePacket: Optional[str], - imember: int | None = None, + imember: int | None = None, ) -> taskBase: # Convert camel case string to snake case @@ -322,7 +322,8 @@ def create_task( factory_logger.info(f'Using module swell.tasks.{task_lower}') # Return task object - return task_class(config, datetime, model, ensemblePacket, additional_parameter, task, imember=imember) + return task_class(config, datetime, model, ensemblePacket, + additional_parameter, imember, task) # -------------------------------------------------------------------------------------------------- @@ -355,7 +356,7 @@ def task_wrapper( model: Optional[str], additional_parameter: str | None, ensemblePacket: Optional[str], - imember: int | None = None, + imember: int | None = None, ) -> None: # Create the object From 542bb0a2ce21910389125717c97a34711557e56a Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 7 Jul 2026 15:22:40 -0600 Subject: [PATCH 29/68] step 1 --- src/swell/suites/eda_controlpert/flow.cylc | 243 +++++++++++++++++ .../suites/eda_controlpert/suite_config.py | 143 ++++++++++ .../run_jedi_eda_control_pert_executable.py | 249 ++++++++++++++++++ src/swell/tasks/task_questions.py | 15 ++ 4 files changed, 650 insertions(+) create mode 100644 src/swell/suites/eda_controlpert/flow.cylc create mode 100644 src/swell/suites/eda_controlpert/suite_config.py create mode 100644 src/swell/tasks/run_jedi_eda_control_pert_executable.py diff --git a/src/swell/suites/eda_controlpert/flow.cylc b/src/swell/suites/eda_controlpert/flow.cylc new file mode 100644 index 000000000..a6c59507f --- /dev/null +++ b/src/swell/suites/eda_controlpert/flow.cylc @@ -0,0 +1,243 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +# Cylc suite for executing JEDI-based non-cycling variational data assimilation + +# -------------------------------------------------------------------------------------------------- + +[scheduler] + UTC mode = True + allow implicit tasks = False + +# -------------------------------------------------------------------------------------------------- + +[scheduling] + + initial cycle point = {{start_cycle_point}} + final cycle point = {{final_cycle_point}} + runahead limit = {{runahead_limit}} + + [[graph]] + R1 = """ + # Triggers for non cycle time dependent tasks + # ------------------------------------------- + # Clone JEDI source code + CloneJedi + + # Build JEDI source code by linking + CloneJedi => BuildJediByLinking? + + # If not able to link to build create the build + BuildJediByLinking:fail? => BuildJedi + + {% for model_component in model_components %} + # Clone geos ana for generating observing system records + CloneGeosMksi-{{model_component}} + {% endfor %} + """ + + {% for cycle_time in cycle_times %} + {{cycle_time.cycle_time}} = """ + {% for model_component in model_components %} + {% if cycle_time[model_component] %} + + # logic tree: + # prep (clone, stage, build, get bkg, get obs) -> SP1 + # -> mksi/observation + filter -> SP2 + # -> [ submit/wait for N chunk runs ]-> SP3 + # -> mean/variance + diffstate -> SP4 -> eva + + # Task triggers for: {{model_component}} + # ------------------ + # Generate satellite channel records + CloneGeosMksi-{{model_component}}[^] => GenerateObservingSystemRecords-{{model_component}} + + # Get observations + {% if cycling_varbc %} + # Cycling VarBC is active, biases from the previous cycle will be used +## ygyu error +## RunJediEdaControlPertExecutable-{{model_component}}[-PT6H] => GetObservations-{{model_component}} + {% else %} + + # Cycling VarBC is inactive, static bias files will be used + GetObsNotInR2d2-{{model_component}}: fail? => GetObservations-{{model_component}} + {% endif %} + + # Perform staging that is cycle dependent + StageJediCycle-{{model_component}} + + # Run Jedi variational executable + BuildJediByLinking[^]? | BuildJedi[^] => SP1-{{model_component}} + CloneJedi[^] => StageJediCycle-{{model_component}} + StageJediCycle-{{model_component}} => SP1-{{model_component}} + GetEnsembleGeosExperiment-{{model_component}} => SP1-{{model_component}} + GetObsNotInR2d2-{{model_component}}? | GetObservations-{{model_component}} => SP1-{{model_component}} + SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} + GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} + RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} + + {% if models[model_component]['ensemble_eda_chunk'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_eda_chunk'] + 1 ) %} + SP2-{{model_component}} => RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}} => SP3-{{model_component}} + {% endfor %} + {% endif %} + + # bkg/ana : mean/variance + SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} + RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} + + # CleanupEda directory + SP4-{{model_component}} => CleanEdaFiles-{{model_component}} + + # EvaIncrement + SP4-{{model_component}} => EvaIncrement-{{model_component}} + + {% endif %} + {% endfor %} + """ + {% endfor %} + +# -------------------------------------------------------------------------------------------------- + +[runtime] + + # Task defaults + # ------------- + [[root]] + pre-script = "source $CYLC_SUITE_DEF_PATH/modules" + + [[[environment]]] + datetime = $CYLC_TASK_CYCLE_POINT + config = $CYLC_SUITE_DEF_PATH/experiment.yaml + + # Tasks + # ----- + [[CloneJedi]] + script = "swell task CloneJedi $config" + + [[BuildJediByLinking]] + script = "swell task BuildJediByLinking $config" + + [[BuildJedi]] + script = "swell task BuildJedi $config" + platform = {{platform}} + execution time limit = {{scheduling["BuildJedi"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["BuildJedi"]["directives"]["all"].items() %} + --{{key}} = {{value}} + {%- endfor %} + + {% for model_component in model_components %} + + [[CloneGeosMksi-{{model_component}}]] + script = "swell task CloneGeosMksi $config -m {{model_component}}" + + [[GenerateObservingSystemRecords-{{model_component}}]] + script = "swell task GenerateObservingSystemRecords $config -d $datetime -m {{model_component}}" + + [[GetObsNotInR2d2-{{model_component}}]] + script = "swell task GetObsNotInR2d2 $config -d $datetime -m {{model_component}}" + + [[StageJediCycle-{{model_component}}]] + script = "swell task StageJedi $config -d $datetime -m {{model_component}}" + + [[GetBackground-{{model_component}} ]] + script = "swell task GetBackground $config -d $datetime -m {{model_component}}" + + [[GetBackgroundGeosExperiment-{{model_component}} ]] + script = "swell task GetBackgroundGeosExperiment $config -d $datetime -m {{model_component}}" + + [[GetEnsembleGeosExperiment-{{model_component}}]] + script = "swell task GetEnsembleGeosExperiment $config -d $datetime -m {{model_component}}" + + [[GetObservations-{{model_component}}]] + script = "swell task GetObservations $config -d $datetime -m {{model_component}}" + + [[RenderJediObservations-{{model_component}}]] + script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" + + {% if models[model_component]['ensemble_eda_chunk'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_eda_chunk'] + 1 ) %} + [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] + script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}} -ichunk {{i}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEdaControlPertExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEdaControlPertExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + {% endfor %} + {% endif %} + + [[RunJediObsfiltersExecutable-{{model_component}}]] + script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediObsfiltersExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediObsfiltersExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[RunJediEnsembleMeanVariance-{{model_component}}]] + script = "swell task RunJediEnsembleMeanVariance $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEnsembleMeanVariance"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEnsembleMeanVariance"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[RunJediDiffstates-{{model_component}}]] + script = "swell task RunJediDiffstates $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediDiffstates"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediDiffstates"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[CleanEdaFiles-{{model_component}}]] + script = "swell task CleanEda $config -d $datetime -m {{model_component}}" + + [[EvaJediLog-{{model_component}}]] + script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}" + + [[EvaIncrement-{{model_component}}]] + script = "swell task EvaIncrement $config -d $datetime -m {{model_component}}" + + [[EvaObservations-{{model_component}}]] + script = "swell task EvaObservations $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["EvaObservations"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["EvaObservations"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[SaveObsDiags-{{model_component}}]] + script = "swell task SaveObsDiags $config -d $datetime -m {{model_component}}" + + [[CleanCycle-{{model_component}}]] + script = "swell task CleanCycle $config -d $datetime -m {{model_component}}" + + # SP: Sync Point + [[SP1-{{model_component}}]] + script = true + + [[SP2-{{model_component}}]] + script = true + + [[SP3-{{model_component}}]] + script = true + + [[SP4-{{model_component}}]] + script = true + + {% endfor %} + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/eda_controlpert/suite_config.py b/src/swell/suites/eda_controlpert/suite_config.py new file mode 100644 index 000000000..b165fefc0 --- /dev/null +++ b/src/swell/suites/eda_controlpert/suite_config.py @@ -0,0 +1,143 @@ +# -------------------------------------------------------------------------------------------------- +# @package configuration +# +# Class containing the configuration. This is a dictionary that is converted from +# an input yaml configuration file. Various function are included for interacting with the +# dictionary. +# +# -------------------------------------------------------------------------------------------------- + + +from swell.utilities.swell_questions import QuestionContainer, QuestionList +from swell.utilities.question_defaults import QuestionDefaults as qd +from swell.suites.suite_questions import SuiteQuestions as sq + +from enum import Enum + + +# -------------------------------------------------------------------------------------------------- + +class SuiteConfig(QuestionContainer, Enum): + + # -------------------------------------------------------------------------------------------------- + + eda_controlpert_atmos_tier1 = QuestionList( + list_name="eda_controlpert_atmos_tier1", + questions=[ + sq.common, + qd.start_cycle_point("2023-10-10T00:00:00Z"), + qd.final_cycle_point("2023-10-10T06:00:00Z"), + qd.runahead_limit("P2"), + qd.jedi_build_method("use_existing"), + qd.model_components(['geos_atmosphere']), + ], + geos_atmosphere=[ + qd.cycle_times([ + "T00", + ]), + qd.geos_x_background_directory("/discover/nobackup/projects/gmao/" + "dadev/rtodling/archive/Restarts/JEDI/541x"), + qd.geos_x_ensemble_directory('/discover/nobackup/projects/gmao/dadev/' + 'rtodling/archive/541/Milan'), + qd.npx_proc(4), + qd.npy_proc(5), + qd.window_length("PT6H"), + qd.window_type("3D"), + qd.horizontal_resolution("91"), + qd.gsibec_nlats("91"), + qd.gsibec_nlons("144"), + qd.vertical_resolution("72"), + qd.ensemble_num_members(32), + qd.ensemble_eda_chunk(8), + qd.obs_pert_amplitude(0.5), + qd.number_of_iterations([100]), + qd.gradient_norm_reduction(1.e-8), + qd.analysis_variables([ + "eastward_wind", + "northward_wind", + "air_temperature", + "water_vapor_mixing_ratio_wrt_moist_air", + "air_pressure_at_surface", + "cloud_liquid_ice", + "cloud_liquid_water", + "rain_water", + "snow_water", + "mole_fraction_of_ozone_in_air", + "geopotential_height_times_gravity_at_surface", + "fraction_of_ocean", + "fraction_of_lake", + "fraction_of_ice", + "skin_temperature_at_surface" + ]), + qd.observations([ + "aircraft_temperature", + "aircraft_wind", + "airs_aqua", + "amsr2_gcom-w1", + "amsua_aqua", + "amsua_metop-b", + "amsua_metop-c", + "amsua_n15", + "amsua_n18", + "amsua_n19", + "atms_n20", + "atms_npp", + "avhrr3_metop-b", + "avhrr3_n18", + "avhrr3_n19", + "cris-fsr_n20", + "cris-fsr_npp", + "gmi_gpm", + "gps", + "iasi_metop-b", + "iasi_metop-c", + "mhs_metop-b", + "mhs_metop-c", + "mhs_n19", + "mls55_aura", + "omi_aura", + "ompsnm_npp", + "pibal", + "satwind", + "scatwind", + "sfcship", + "sfc", + "sondes", + "ssmis_f17" + ]), + qd.obs_thinning_rej_fraction(0.8), + qd.ensmeanvariance_spec([ + {"state": "bkg", + "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "geos.prior.mean", + "fn_output_variance": "geos.prior.variance", + "grid_type": ['cs', 'latlon']}, + {"state": "analysis", + "fn_input": "analysis/mem%mem%/eda.ana.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "eda.ana.mean", + "fn_output_variance": "eda.ana.variance", + "grid_type": ['cs', 'latlon']}, + ]), + qd.diffstates_spec({ + "state1": + {"fn_input": "geos.prior.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + "state2": + {"fn_input": "eda.ana.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + "state_diff": + {"fn_output": "eda.mean-inc", "grid_type": ['cs', 'latlon']}, + "state_type": "ensemble" + }), + qd.clean_patterns(['*.txt', '*.csv']), + ] + ) + + # -------------------------------------------------------------------------------------------------- + + eda_controlpert = QuestionList( + list_name="eda_controlpert", + questions=[ + eda_controlpert_atmos_tier1 + ] + ) + + # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/run_jedi_eda_control_pert_executable.py b/src/swell/tasks/run_jedi_eda_control_pert_executable.py new file mode 100644 index 000000000..34b0df835 --- /dev/null +++ b/src/swell/tasks/run_jedi_eda_control_pert_executable.py @@ -0,0 +1,249 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + + +# -------------------------------------------------------------------------------------------------- + +import os +import glob +import shutil +from ruamel.yaml import YAML + +from swell.tasks.base.task_base import taskBase +from swell.utilities.run_jedi_executables import run_executable + + +# -------------------------------------------------------------------------------------------------- + + +class RunJediEdaControlPertExecutable(taskBase): + + # ---------------------------------------------------------------------------------------------- + + def execute(self) -> None: + + # Jedi application name + # --------------------- + jedi_application = 'eda_control_pert' + + # Parse configuration + # ------------------- + window_type = self.config.window_type() + window_length = self.config.window_length() + forecast_length = self.config.forecast_length(window_length) + background_time_offset = self.config.background_time_offset() + number_of_iterations = self.config.number_of_iterations() + jedi_forecast_model = self.config.jedi_forecast_model(None) + generate_yaml_and_exit = self.config.generate_yaml_and_exit(False) + perhost = self.config.perhost(None) + + # Set the observing system records path + self.jedi_rendering.set_obs_records_path(self.config.observing_system_records_path(None)) + + gsibec_nlats = self.config.gsibec_nlats(None) + gsibec_nlons = self.config.gsibec_nlons(None) + gsibec_configuration = self.config.gsibec_configuration(None) + npx_proc = self.config.npx_proc(None) + npy_proc = self.config.npy_proc(None) + npx = self.config.npx(None) + npy = self.config.npy(None) + + # Compute data assimilation window parameters + # -------------------------------------------- + background_time = self.da_window_params.background_time(background_time_offset) + local_background_time = self.da_window_params.local_background_time(window_length, + window_type) + local_background_time_iso = self.da_window_params.local_background_time_iso(window_length, + window_type) + window_begin = self.da_window_params.window_begin(window_length) + window_begin_iso = self.da_window_params.window_begin_iso(window_length) + window_end_iso = self.da_window_params.window_end_iso(window_length) + nmember = self.config.ensemble_num_members() + nchunk = self.config.ensemble_eda_chunk() + ichunk = self.get_ensemble_ichunk() + + + # Populate jedi interface templates dictionary + # -------------------------------------------- + self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) + self.jedi_rendering.add_key('window_end_iso', window_end_iso) + self.jedi_rendering.add_key('window_length', window_length) + self.jedi_rendering.add_key('forecast_length', forecast_length) + self.jedi_rendering.add_key('minimizer', self.config.minimizer()) + self.jedi_rendering.add_key('number_of_iterations', number_of_iterations[0]) + self.jedi_rendering.add_key('analysis_variables', self.config.analysis_variables()) + self.jedi_rendering.add_key('saber_central_block', self.config.saber_central_block(None)) + self.jedi_rendering.add_key('saber_outer_block', self.config.saber_outer_block(None)) + self.jedi_rendering.add_key('gradient_norm_reduction', + self.config.gradient_norm_reduction()) + self.jedi_rendering.add_key('marine_models', self.config.marine_models(None)) + + # Background + # ---------- + self.jedi_rendering.add_key('horizontal_resolution', self.config.horizontal_resolution()) + self.jedi_rendering.add_key('local_background_time', local_background_time) + self.jedi_rendering.add_key('local_background_time_iso', local_background_time_iso) + self.jedi_rendering.add_key('ensemble_num_members', self.config.ensemble_num_members()) + self.jedi_rendering.add_key('ensemble_ichunk', ichunk) + + # Geometry + # -------- + self.jedi_rendering.add_key('vertical_resolution', self.config.vertical_resolution()) + self.jedi_rendering.add_key('gsibec_nlats', gsibec_nlats) + self.jedi_rendering.add_key('gsibec_nlons', gsibec_nlons) + self.jedi_rendering.add_key('npx_proc', npx_proc) + self.jedi_rendering.add_key('npy_proc', npy_proc) + self.jedi_rendering.add_key('npx', npx) + self.jedi_rendering.add_key('npy', npy) + self.jedi_rendering.add_key('total_processors', self.config.total_processors(None)) + + # Observations + # ------------ + self.jedi_rendering.add_key('background_time', background_time) + self.jedi_rendering.add_key('crtm_coeff_dir', self.config.crtm_coeff_dir(None)) + self.jedi_rendering.add_key('window_begin', window_begin) + + # Atmosphere background error model + # --------------------------------- + if gsibec_configuration is not None: + self.jedi_rendering.add_key('gsibec_configuration', gsibec_configuration) + self.jedi_rendering.add_key('gsibec_nlats', gsibec_nlats) + self.jedi_rendering.add_key('gsibec_nlons', gsibec_nlons) + self.jedi_rendering.add_key('gsibec_npx_proc', npx_proc) + self.jedi_rendering.add_key('gsibec_npy_proc', 6*npy_proc) + + # Model + # ----- + if window_type == '4D': + self.jedi_rendering.add_key('background_frequency', self.config.background_frequency()) + # Jedi configuration file + # ----------------------- + str=f'jedi_{jedi_application}{window_type}_config_chunk{{ichunk:03d}}.yaml' + jedi_config_file = os.path.join(self.cycle_dir(), str) + + # Output log file + # --------------- + output_log_file = os.path.join( + self.cycle_dir(), f"jedi_{jedi_application}{window_type}_log_chunk{{ichunk:03d}}.log") + + # Open the JEDI config file and fill initial templates + # ---------------------------------------------------- + jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}', + window_type, + jedi_forecast_model) + + # ichunk: specify yaml + # ---------------------------------------------------- + + # link ebkg to ebkg_chunk + ebkg_chunk_dir = f'ebkg_chunk/chunk{ichunk:003d}/' + xdir = os.path.join(self.cycle_dir(), ebkg_chunk_dir) + os.makedirs(xdir, exist_ok=True) + + npert = nmember / nchunk + mem_s = npert * (ichunk -1) + 1 + + for imem in range(1, npert+1): + id = mem_s + imem - 1 + f1 = f'ebkg/mem{id:003d}/' + + + # create subdir for analysis + chunk_dir = f'analysis/chunk{ichunk:003d}/' + xdir = os.path.join(self.cycle_dir(), chunk_dir) + os.makedirs(xdir, exist_ok=True) + + for observer in jedi_config_dict['cost function']['observations']['observers']: + # Get observation name + observation = observer['observation_name'] + print(f'ob= {observation}') + # copy obs input file to avoid multi MPI reading the same file + files = glob.glob(os.path.join(self.cycle_dir(), f'{observation}.*')) + for src_file in files: + print(f'f= {src_file}') + shutil.copy(src_file, xdir) + + obs_bias = observer.get('obs bias') + if obs_bias is not None: + File = obs_bias['input file'] + dir1, fname = os.path.split(File) + File = os.path.join(dir1, chunk_dir, fname) + obs_bias['input file'] = File + # + File = obs_bias['output file'] + dir1, fname = os.path.split(File) + File = os.path.join(dir1, chunk_dir, fname) + obs_bias['output file'] = File + # + File = obs_bias.get('covariance', {}).get('output file') + if File is not None: + dir1, fname = os.path.split(File) + File = os.path.join(dir1, chunk_dir, fname) + obs_bias['covariance']['output file'] = File + File = obs_bias.get('covariance', {}).get('prior', {}).get('input file') + if File is not None: + dir1, fname = os.path.split(File) + File = os.path.join(dir1, chunk_dir, fname) + obs_bias['covariance']['prior']['input file'] = File + + ruamel_yaml = YAML() + ruamel_yaml.default_flow_style = False + + # Write the ordered dictionary to YAML file + with open(jedi_config_file, 'w') as jedi_config_file_open: + ruamel_yaml.dump(jedi_config_dict, jedi_config_file_open) + + # copy fv3-jedi dir, update dir names + d1 = os.path.join(self.cycle_dir(), 'fv3-jedi') + d2 = os.path.join(self.cycle_dir(), chunk_dir, 'fv3-jedi') + shutil.copytree(d1, d2, dirs_exist_ok=True) + + with open(jedi_config_file, 'r') as f: + yaml_content = f.read() + + dir_list = ["bkg", "fv3files", "gsibec", "rcov"] + for i in dir_list: + j = f"fv3-jedi/{i}" + k = f"{chunk_dir}{j}" # Result: analysis/mem002/fv3-jedi/rcov ??? analysis/chunk/mem002/fv3-jedi/rcov + yaml_content = yaml_content.replace(j, k) + with open(jedi_config_file, 'w') as f: + f.write(yaml_content) + + # Get the JEDI interface metadata + # ------------------------------- + model_component_meta = self.jedi_rendering.render_interface_meta() + + # Compute number of processors + # ---------------------------- + np = eval(str(model_component_meta['total_processors'])) + # modify np by (nmember / nchunk) + 1 + if ichunk == 1: + np = np * (nmember / nchunk) + else: + np = np * (nmember / nchunk + 1) + + + # Jedi executable name + # -------------------- + jedi_executable = model_component_meta['executables']['edaControlPert'] + jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', 'build', 'bin', + jedi_executable) + + # Run the JEDI executable + # ----------------------- + if not generate_yaml_and_exit: + self.logger.info('Running '+jedi_executable_path+' with '+str(np)+' processors.') + run_executable(self.logger, self.cycle_dir(), np, jedi_executable_path, + jedi_config_file, output_log_file, perhost) + else: + mpi_command = "mpirun" + if not (perhost is None or perhost == "None"): + mpi_command += f" -perhost {perhost}" + mpi_command += f" -np {np} {jedi_executable_path} {jedi_config_file} {output_log_file}" + print(f'intended mpi_command = {mpi_command}') + self.logger.info('YAML generated, now exiting.') + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index cf27cf10c..6798e883e 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -819,6 +819,21 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + RunJediEdaControlPertExecutable = QuestionList( + list_name="RunJediEdaControlPertExecutable", + questions=[ + run_jedi_executable, + qd.ensemble_num_members(), + qd.ensemble_num_chunks(), + qd.obs_thinning_rej_fraction(), + qd.perhost(), + qd.comparison_log_type('variational'), + qd.mock_experiment() + ] + ) + + # -------------------------------------------------------------------------------------------------- + SaveObsDiags = QuestionList( list_name="SaveObsDiags", questions=[ From 79ad11daefc27d32f2a0b0d7b97eb30494f23431 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 7 Jul 2026 15:38:51 -0600 Subject: [PATCH 30/68] s1 --- src/swell/suites/eda_controlpert/flow.cylc | 8 ++-- .../suites/eda_controlpert/suite_config.py | 45 +++++++++---------- src/swell/utilities/question_defaults.py | 13 ++++++ 3 files changed, 39 insertions(+), 27 deletions(-) diff --git a/src/swell/suites/eda_controlpert/flow.cylc b/src/swell/suites/eda_controlpert/flow.cylc index a6c59507f..b0dbe1b5b 100644 --- a/src/swell/suites/eda_controlpert/flow.cylc +++ b/src/swell/suites/eda_controlpert/flow.cylc @@ -81,8 +81,8 @@ GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} - {% if models[model_component]['ensemble_eda_chunk'] is defined %} - {% for i in range( 1, models[model_component]['ensemble_eda_chunk'] + 1 ) %} + {% if models[model_component]['ensemble_num_chunks'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} SP2-{{model_component}} => RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}} => SP3-{{model_component}} {% endfor %} {% endif %} @@ -161,8 +161,8 @@ [[RenderJediObservations-{{model_component}}]] script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" - {% if models[model_component]['ensemble_eda_chunk'] is defined %} - {% for i in range( 1, models[model_component]['ensemble_eda_chunk'] + 1 ) %} + {% if models[model_component]['ensemble_num_chunks'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}} -ichunk {{i}}" platform = {{platform}} diff --git a/src/swell/suites/eda_controlpert/suite_config.py b/src/swell/suites/eda_controlpert/suite_config.py index b165fefc0..c8c4889d2 100644 --- a/src/swell/suites/eda_controlpert/suite_config.py +++ b/src/swell/suites/eda_controlpert/suite_config.py @@ -48,8 +48,7 @@ class SuiteConfig(QuestionContainer, Enum): qd.gsibec_nlons("144"), qd.vertical_resolution("72"), qd.ensemble_num_members(32), - qd.ensemble_eda_chunk(8), - qd.obs_pert_amplitude(0.5), + qd.ensemble_num_chunks(8), qd.number_of_iterations([100]), qd.gradient_norm_reduction(1.e-8), qd.analysis_variables([ @@ -106,27 +105,27 @@ class SuiteConfig(QuestionContainer, Enum): "ssmis_f17" ]), qd.obs_thinning_rej_fraction(0.8), - qd.ensmeanvariance_spec([ - {"state": "bkg", - "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", - "fn_output_mean": "geos.prior.mean", - "fn_output_variance": "geos.prior.variance", - "grid_type": ['cs', 'latlon']}, - {"state": "analysis", - "fn_input": "analysis/mem%mem%/eda.ana.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", - "fn_output_mean": "eda.ana.mean", - "fn_output_variance": "eda.ana.variance", - "grid_type": ['cs', 'latlon']}, - ]), - qd.diffstates_spec({ - "state1": - {"fn_input": "geos.prior.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, - "state2": - {"fn_input": "eda.ana.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, - "state_diff": - {"fn_output": "eda.mean-inc", "grid_type": ['cs', 'latlon']}, - "state_type": "ensemble" - }), +# qd.ensmeanvariance_spec([ +# {"state": "bkg", +# "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", +# "fn_output_mean": "geos.prior.mean", +# "fn_output_variance": "geos.prior.variance", +# "grid_type": ['cs', 'latlon']}, +# {"state": "analysis", +# "fn_input": "analysis/mem%mem%/eda.ana.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", +# "fn_output_mean": "eda.ana.mean", +# "fn_output_variance": "eda.ana.variance", +# "grid_type": ['cs', 'latlon']}, +# ]), +# qd.diffstates_spec({ +# "state1": +# {"fn_input": "geos.prior.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, +# "state2": +# {"fn_input": "eda.ana.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, +# "state_diff": +# {"fn_output": "eda.mean-inc", "grid_type": ['cs', 'latlon']}, +# "state_type": "ensemble" +# }), qd.clean_patterns(['*.txt', '*.csv']), ] ) diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index 742262d16..44145917b 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -520,6 +520,19 @@ class ensemble_hofx_strategy(TaskQuestion): # -------------------------------------------------------------------------------------------------- + @dataclass + class ensemble_num_chunks(TaskQuestion): + default_value: str = "defer_to_model" + question_name: str = "ensemble_num_chunks" + options: str = "defer_to_model" + models: List[str] = mutable_field([ + "geos_atmosphere" + ]) + prompt: str = "How many chunks shall be used in EDA control pert calculations?" + widget_type: WType = WType.INTEGER + + # -------------------------------------------------------------------------------------------------- + @dataclass class ensemble_num_members(TaskQuestion): default_value: str = "defer_to_model" From 6664a9284cc050cbe93b8c8782036493dbb91308 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 7 Jul 2026 20:23:05 -0600 Subject: [PATCH 31/68] - where did model dependent task disappear in for loop - experiment.yaml + flow.cylc --- .../geos_atmosphere/task_questions.yaml | 4 + src/swell/deployment/create_experiment.py | 7 + .../prepare_config_and_suite.py | 27 +- src/swell/suites/eda_controlpert/flow.cylc | 192 +------------- .../suites/eda_controlpert/flow.cylc.complete | 237 ++++++++++++++++++ .../suites/eda_controlpert/flow.cylc.simple | 48 ++++ src/swell/suites/eda_controlpert/flow.cylc.v2 | 210 ++++++++++++++++ 7 files changed, 539 insertions(+), 186 deletions(-) create mode 100644 src/swell/suites/eda_controlpert/flow.cylc.complete create mode 100644 src/swell/suites/eda_controlpert/flow.cylc.simple create mode 100644 src/swell/suites/eda_controlpert/flow.cylc.v2 diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml b/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml index 3c5261f11..88700667f 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml @@ -125,6 +125,10 @@ ensemble_num_members: default_value: 10 options: None +ensemble_num_chunks: + default_value: 8 + options: None + geovals_experiment: default_value: x0050-geovals options: diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 46596ce78..3758d4277 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -256,9 +256,13 @@ def create_experiment_directory( # Experiment ID and root from the user input # ------------------------------------------ + + #ygyu test experiment_id = dict_get(logger, experiment_dict, 'experiment_id') experiment_root = dict_get(logger, experiment_dict, 'experiment_root') + + # Write out some info # ------------------- logger.info(f'Creating experiment: \'{experiment_id}\' in \'{experiment_root}\'') @@ -483,6 +487,9 @@ def prepare_cylc_suite_jinja2( # -------------------------- with open(os.path.join(swell_suite_path, 'flow.cylc'), 'r') as file: suite_file = file.read() + print(f'nail6. suite_file in prepare_cylc_suite_jinja2 = {suite_file}') + + # Copy the experiment dictionary to the rendering dictionary # ---------------------------------------------------------- diff --git a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py index 69987bf5f..ae198147c 100644 --- a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py +++ b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py @@ -92,6 +92,11 @@ def __init__( self.model_ind_tasks = self.get_suite_task_list_model_ind(self.suite_str) self.all_model_dep_tasks = self.get_all_model_dep_tasks(self.suite_str) + print(f'nail1 model_ind_tasks = {self.model_ind_tasks}') + print(f'nail2 all_model_dep_tasks = {self.all_model_dep_tasks}') + print(f'nail3 .suite_str = {self.suite_str}') + + # Perform the assembly of the dictionaries that contain all the questions that can possibly # be asked. This @@ -121,6 +126,8 @@ def prepare_question_dictionaries(self) -> None: suite_config_obj = AllSuites.get_config(self.suite_config) suite_question_list = suite_config_obj.expand_question_list() + print(f'suite q list : {suite_question_list}') + # Allow for adding extra tasks manually from configuration # For dynamic suite creation (e.g. comparison tests) @@ -267,6 +274,11 @@ def prepare_question_dictionaries(self) -> None: for key in keys_to_remove: del self.question_dictionary_model_dep[model][key] +# print(f'nail1.5 question_dictionary_model_dep = {self.question_dictionary_model_dep}') +# for model in self.possible_model_components: +# x = self.question_dictionary_model_dep[model]['ensemble_num_chunks'] +# print(f'x = {x}') + # ---------------------------------------------------------------------------------------------- def override_with_defaults(self) -> None: @@ -426,6 +438,13 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # 2. Perform a non-exhaustive resolving of suite file templates # ------------------------------------------------------------- suite_str = template_string_jinja2(self.logger, self.suite_str, self.experiment_dict, True) + + print(f'nail4.0 self.exp dict = {self.experiment_dict}') + + print(f'nail4.1 self.suite str = {self.suite_str}') + + + print(f'nail4.2 jinja2 suite str = {suite_str}') # 3. Get a list of tasks that do not depend on the model component # ---------------------------------------------------------------- @@ -499,7 +518,9 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # 9.1 Build a list of tasks for each model component # ------------------------------------------------- + print(f'ck suite_str = {suite_str}') model_dep_tasks = self.get_suite_task_list_model_dep(suite_str) + # 9.2 Iterate over the model_dep dictionary and ask task questions # ---------------------------------------------------------------- @@ -627,11 +648,13 @@ def get_all_model_dep_tasks(self, suite_str: str) -> list: swell_task_lines = [line.replace('"', '') for line in swell_task_lines] swell_task_lines = [line.strip() for line in swell_task_lines] + print(f'def all model dep tasks: swell_task_lines = {swell_task_lines}') + # All tasks all_tasks = [] for line in swell_task_lines: - all_tasks.append(line.split('swell task ')[1].split(' ')[0]) + all_tasks.append(line.split('swell task')[1].strip().split(' ')[0]) # Ensure all_tasks are unique all_tasks = list(set(all_tasks)) @@ -650,6 +673,8 @@ def get_suite_task_list_model_dep(self, suite_str: str) -> dict: swell_task_lines = [line.replace('"', '') for line in swell_task_lines] swell_task_lines = [line.strip() for line in swell_task_lines] + print(f'def get_suite_task_list_model_dep swell_task_lines = {swell_task_lines}') + # Now get the model part models = [] for line in swell_task_lines: diff --git a/src/swell/suites/eda_controlpert/flow.cylc b/src/swell/suites/eda_controlpert/flow.cylc index b0dbe1b5b..6f657c020 100644 --- a/src/swell/suites/eda_controlpert/flow.cylc +++ b/src/swell/suites/eda_controlpert/flow.cylc @@ -1,15 +1,3 @@ -# (C) Copyright 2021- United States Government as represented by the Administrator of the -# National Aeronautics and Space Administration. All Rights Reserved. -# -# This software is licensed under the terms of the Apache Licence Version 2.0 -# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. - -# -------------------------------------------------------------------------------------------------- - -# Cylc suite for executing JEDI-based non-cycling variational data assimilation - -# -------------------------------------------------------------------------------------------------- - [scheduler] UTC mode = True allow implicit tasks = False @@ -27,80 +15,14 @@ # Triggers for non cycle time dependent tasks # ------------------------------------------- # Clone JEDI source code - CloneJedi - - # Build JEDI source code by linking - CloneJedi => BuildJediByLinking? - - # If not able to link to build create the build - BuildJediByLinking:fail? => BuildJedi - {% for model_component in model_components %} - # Clone geos ana for generating observing system records - CloneGeosMksi-{{model_component}} - {% endfor %} - """ - - {% for cycle_time in cycle_times %} - {{cycle_time.cycle_time}} = """ - {% for model_component in model_components %} - {% if cycle_time[model_component] %} - - # logic tree: - # prep (clone, stage, build, get bkg, get obs) -> SP1 - # -> mksi/observation + filter -> SP2 - # -> [ submit/wait for N chunk runs ]-> SP3 - # -> mean/variance + diffstate -> SP4 -> eva - - # Task triggers for: {{model_component}} - # ------------------ - # Generate satellite channel records - CloneGeosMksi-{{model_component}}[^] => GenerateObservingSystemRecords-{{model_component}} - - # Get observations - {% if cycling_varbc %} - # Cycling VarBC is active, biases from the previous cycle will be used -## ygyu error -## RunJediEdaControlPertExecutable-{{model_component}}[-PT6H] => GetObservations-{{model_component}} - {% else %} - - # Cycling VarBC is inactive, static bias files will be used - GetObsNotInR2d2-{{model_component}}: fail? => GetObservations-{{model_component}} - {% endif %} - - # Perform staging that is cycle dependent - StageJediCycle-{{model_component}} - - # Run Jedi variational executable - BuildJediByLinking[^]? | BuildJedi[^] => SP1-{{model_component}} - CloneJedi[^] => StageJediCycle-{{model_component}} - StageJediCycle-{{model_component}} => SP1-{{model_component}} - GetEnsembleGeosExperiment-{{model_component}} => SP1-{{model_component}} - GetObsNotInR2d2-{{model_component}}? | GetObservations-{{model_component}} => SP1-{{model_component}} - SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} - GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} - RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} {% if models[model_component]['ensemble_num_chunks'] is defined %} - {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} - SP2-{{model_component}} => RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}} => SP3-{{model_component}} - {% endfor %} + RunJediEdaControlPertExecutable-{{model_component}} {% endif %} - # bkg/ana : mean/variance - SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} - RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} - - # CleanupEda directory - SP4-{{model_component}} => CleanEdaFiles-{{model_component}} - - # EvaIncrement - SP4-{{model_component}} => EvaIncrement-{{model_component}} - - {% endif %} - {% endfor %} + {% endfor %} """ - {% endfor %} # -------------------------------------------------------------------------------------------------- @@ -117,126 +39,26 @@ # Tasks # ----- + [[CloneJedi]] script = "swell task CloneJedi $config" - [[BuildJediByLinking]] - script = "swell task BuildJediByLinking $config" - - [[BuildJedi]] - script = "swell task BuildJedi $config" - platform = {{platform}} - execution time limit = {{scheduling["BuildJedi"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["BuildJedi"]["directives"]["all"].items() %} - --{{key}} = {{value}} - {%- endfor %} - {% for model_component in model_components %} - [[CloneGeosMksi-{{model_component}}]] - script = "swell task CloneGeosMksi $config -m {{model_component}}" - - [[GenerateObservingSystemRecords-{{model_component}}]] - script = "swell task GenerateObservingSystemRecords $config -d $datetime -m {{model_component}}" - - [[GetObsNotInR2d2-{{model_component}}]] - script = "swell task GetObsNotInR2d2 $config -d $datetime -m {{model_component}}" - - [[StageJediCycle-{{model_component}}]] - script = "swell task StageJedi $config -d $datetime -m {{model_component}}" + {% if models[model_component]['ensemble_num_chunks'] is defined %} + [[RunJediEdaControlPertExecutable-{{model_component}}]] + script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} [[GetBackground-{{model_component}} ]] script = "swell task GetBackground $config -d $datetime -m {{model_component}}" - [[GetBackgroundGeosExperiment-{{model_component}} ]] - script = "swell task GetBackgroundGeosExperiment $config -d $datetime -m {{model_component}}" - - [[GetEnsembleGeosExperiment-{{model_component}}]] - script = "swell task GetEnsembleGeosExperiment $config -d $datetime -m {{model_component}}" - - [[GetObservations-{{model_component}}]] - script = "swell task GetObservations $config -d $datetime -m {{model_component}}" - - [[RenderJediObservations-{{model_component}}]] - script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" - - {% if models[model_component]['ensemble_num_chunks'] is defined %} - {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} - [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] - script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}} -ichunk {{i}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediEdaControlPertExecutable"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediEdaControlPertExecutable"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - {% endfor %} {% endif %} - [[RunJediObsfiltersExecutable-{{model_component}}]] - script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediObsfiltersExecutable"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediObsfiltersExecutable"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - - [[RunJediEnsembleMeanVariance-{{model_component}}]] - script = "swell task RunJediEnsembleMeanVariance $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediEnsembleMeanVariance"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediEnsembleMeanVariance"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - - [[RunJediDiffstates-{{model_component}}]] - script = "swell task RunJediDiffstates $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediDiffstates"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediDiffstates"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - [[CleanEdaFiles-{{model_component}}]] - script = "swell task CleanEda $config -d $datetime -m {{model_component}}" - - [[EvaJediLog-{{model_component}}]] - script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}" - - [[EvaIncrement-{{model_component}}]] - script = "swell task EvaIncrement $config -d $datetime -m {{model_component}}" - - [[EvaObservations-{{model_component}}]] - script = "swell task EvaObservations $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["EvaObservations"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["EvaObservations"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - - [[SaveObsDiags-{{model_component}}]] - script = "swell task SaveObsDiags $config -d $datetime -m {{model_component}}" - - [[CleanCycle-{{model_component}}]] - script = "swell task CleanCycle $config -d $datetime -m {{model_component}}" - - # SP: Sync Point [[SP1-{{model_component}}]] script = true - [[SP2-{{model_component}}]] - script = true - - [[SP3-{{model_component}}]] - script = true - - [[SP4-{{model_component}}]] - script = true {% endfor %} diff --git a/src/swell/suites/eda_controlpert/flow.cylc.complete b/src/swell/suites/eda_controlpert/flow.cylc.complete new file mode 100644 index 000000000..70cd0433c --- /dev/null +++ b/src/swell/suites/eda_controlpert/flow.cylc.complete @@ -0,0 +1,237 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +# Cylc suite for executing JEDI-based non-cycling variational data assimilation + +# -------------------------------------------------------------------------------------------------- + +[scheduler] + UTC mode = True + allow implicit tasks = False + +# -------------------------------------------------------------------------------------------------- + +[scheduling] + + initial cycle point = {{start_cycle_point}} + final cycle point = {{final_cycle_point}} + runahead limit = {{runahead_limit}} + + [[graph]] + R1 = """ + # Triggers for non cycle time dependent tasks + # ------------------------------------------- + # Clone JEDI source code + CloneJedi + + # Build JEDI source code by linking + CloneJedi => BuildJediByLinking? + + # If not able to link to build create the build + BuildJediByLinking:fail? => BuildJedi + + {% for model_component in model_components %} + # Clone geos ana for generating observing system records + CloneGeosMksi-{{model_component}} + {% endfor %} + """ + + {% for cycle_time in cycle_times %} + {{cycle_time.cycle_time}} = """ + {% for model_component in model_components %} + {% if cycle_time[model_component] %} + + # logic tree: + # prep (clone, stage, build, get bkg, get obs) -> SP1 + # -> mksi/observation + filter -> SP2 + # -> [ submit/wait for N chunk runs ]-> SP3 + # -> mean/variance + diffstate -> SP4 -> eva + + # Task triggers for: {{model_component}} + # ------------------ + # Generate satellite channel records + CloneGeosMksi-{{model_component}}[^] => GenerateObservingSystemRecords-{{model_component}} + + # ygyu fix at the end + # Get observations + # Cycling VarBC is inactive, static bias files will be used + GetObsNotInR2d2-{{model_component}}: fail? => GetObservations-{{model_component}} + + # Perform staging that is cycle dependent + StageJediCycle-{{model_component}} + + # Run Jedi variational executable + BuildJediByLinking[^]? | BuildJedi[^] => SP1-{{model_component}} + CloneJedi[^] => StageJediCycle-{{model_component}} + StageJediCycle-{{model_component}} => SP1-{{model_component}} + GetEnsembleGeosExperiment-{{model_component}} => SP1-{{model_component}} + GetObsNotInR2d2-{{model_component}}? | GetObservations-{{model_component}} => SP1-{{model_component}} + SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} + GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} + RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} + + {% if models[model_component]['ensemble_num_chunks'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} + SP2-{{model_component}} => RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}} => SP3-{{model_component}} + {% endfor %} + {% endif %} + + # bkg/ana : mean/variance + SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} + RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} + + # CleanupEda directory + SP4-{{model_component}} => CleanEdaFiles-{{model_component}} + + # EvaIncrement + SP4-{{model_component}} => EvaIncrement-{{model_component}} + + {% endif %} + {% endfor %} + """ + {% endfor %} + +# -------------------------------------------------------------------------------------------------- + +[runtime] + + # Task defaults + # ------------- + [[root]] + pre-script = "source $CYLC_SUITE_DEF_PATH/modules" + + [[[environment]]] + datetime = $CYLC_TASK_CYCLE_POINT + config = $CYLC_SUITE_DEF_PATH/experiment.yaml + + # Tasks + # ----- + [[CloneJedi]] + script = "swell task CloneJedi $config" + + [[BuildJediByLinking]] + script = "swell task BuildJediByLinking $config" + + [[BuildJedi]] + script = "swell task BuildJedi $config" + platform = {{platform}} + execution time limit = {{scheduling["BuildJedi"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["BuildJedi"]["directives"]["all"].items() %} + --{{key}} = {{value}} + {%- endfor %} + + {% for model_component in model_components %} + + [[CloneGeosMksi-{{model_component}}]] + script = "swell task CloneGeosMksi $config -m {{model_component}}" + + [[GenerateObservingSystemRecords-{{model_component}}]] + script = "swell task GenerateObservingSystemRecords $config -d $datetime -m {{model_component}}" + + [[GetObsNotInR2d2-{{model_component}}]] + script = "swell task GetObsNotInR2d2 $config -d $datetime -m {{model_component}}" + + [[StageJediCycle-{{model_component}}]] + script = "swell task StageJedi $config -d $datetime -m {{model_component}}" + + [[GetBackground-{{model_component}} ]] + script = "swell task GetBackground $config -d $datetime -m {{model_component}}" + + [[GetBackgroundGeosExperiment-{{model_component}} ]] + script = "swell task GetBackgroundGeosExperiment $config -d $datetime -m {{model_component}}" + + [[GetEnsembleGeosExperiment-{{model_component}}]] + script = "swell task GetEnsembleGeosExperiment $config -d $datetime -m {{model_component}}" + + [[GetObservations-{{model_component}}]] + script = "swell task GetObservations $config -d $datetime -m {{model_component}}" + + [[RenderJediObservations-{{model_component}}]] + script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" + + {% if models[model_component]['ensemble_num_chunks'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} + [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] + script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}} -ichunk {{i}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEdaControlPertExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEdaControlPertExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + {% endfor %} + {% endif %} + + [[RunJediObsfiltersExecutable-{{model_component}}]] + script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediObsfiltersExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediObsfiltersExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[RunJediEnsembleMeanVariance-{{model_component}}]] + script = "swell task RunJediEnsembleMeanVariance $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEnsembleMeanVariance"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEnsembleMeanVariance"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[RunJediDiffstates-{{model_component}}]] + script = "swell task RunJediDiffstates $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediDiffstates"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediDiffstates"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[CleanEdaFiles-{{model_component}}]] + script = "swell task CleanEda $config -d $datetime -m {{model_component}}" + + [[EvaJediLog-{{model_component}}]] + script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}" + + [[EvaIncrement-{{model_component}}]] + script = "swell task EvaIncrement $config -d $datetime -m {{model_component}}" + + [[EvaObservations-{{model_component}}]] + script = "swell task EvaObservations $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["EvaObservations"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["EvaObservations"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[SaveObsDiags-{{model_component}}]] + script = "swell task SaveObsDiags $config -d $datetime -m {{model_component}}" + + [[CleanCycle-{{model_component}}]] + script = "swell task CleanCycle $config -d $datetime -m {{model_component}}" + + # SP: Sync Point + [[SP1-{{model_component}}]] + script = true + + [[SP2-{{model_component}}]] + script = true + + [[SP3-{{model_component}}]] + script = true + + [[SP4-{{model_component}}]] + script = true + + {% endfor %} + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/eda_controlpert/flow.cylc.simple b/src/swell/suites/eda_controlpert/flow.cylc.simple new file mode 100644 index 000000000..133dba8f6 --- /dev/null +++ b/src/swell/suites/eda_controlpert/flow.cylc.simple @@ -0,0 +1,48 @@ +[scheduler] + UTC mode = True + allow implicit tasks = False + +# -------------------------------------------------------------------------------------------------- + +[scheduling] + + initial cycle point = {{start_cycle_point}} + final cycle point = {{final_cycle_point}} + runahead limit = {{runahead_limit}} + + [[graph]] + R1 = """ + # Triggers for non cycle time dependent tasks + # ------------------------------------------- + # Clone JEDI source code + {% for model_component in model_components %} + + RunJediEdaControlPertExecutable-{{model_component}} + + {% endfor %} + """ + +# -------------------------------------------------------------------------------------------------- + +[runtime] + + # Task defaults + # ------------- + [[root]] + pre-script = "source $CYLC_SUITE_DEF_PATH/modules" + + [[[environment]]] + datetime = $CYLC_TASK_CYCLE_POINT + config = $CYLC_SUITE_DEF_PATH/experiment.yaml + + # Tasks + # ----- + + {% for model_component in model_components %} + [[RunJediEdaControlPertExecutable-{{model_component}}]] + script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} + + {% endfor %} + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/eda_controlpert/flow.cylc.v2 b/src/swell/suites/eda_controlpert/flow.cylc.v2 new file mode 100644 index 000000000..b7e1b0fb0 --- /dev/null +++ b/src/swell/suites/eda_controlpert/flow.cylc.v2 @@ -0,0 +1,210 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +# Cylc suite for executing JEDI-based non-cycling variational data assimilation + +# -------------------------------------------------------------------------------------------------- + +[scheduler] + UTC mode = True + allow implicit tasks = False + +# -------------------------------------------------------------------------------------------------- + +[scheduling] + + initial cycle point = {{start_cycle_point}} + final cycle point = {{final_cycle_point}} + runahead limit = {{runahead_limit}} + + [[graph]] + R1 = """ + # Triggers for non cycle time dependent tasks + # ------------------------------------------- + # Clone JEDI source code + CloneJedi + + # Build JEDI source code by linking + CloneJedi => BuildJediByLinking? + + # If not able to link to build create the build + BuildJediByLinking:fail? => BuildJedi + + {% for model_component in model_components %} + # Clone geos ana for generating observing system records + CloneGeosMksi-{{model_component}} + {% endfor %} + """ + + {% for cycle_time in cycle_times %} + {{cycle_time.cycle_time}} = """ + {% for model_component in model_components %} + {% if cycle_time[model_component] %} + + # logic tree: + # prep (clone, stage, build, get bkg, get obs) -> SP1 + # -> mksi/observation + filter -> SP2 + # -> [ submit/wait for N chunk runs ]-> SP3 + # -> mean/variance + diffstate -> SP4 -> eva + + # Task triggers for: {{model_component}} + # ------------------ + # Generate satellite channel records + CloneGeosMksi-{{model_component}}[^] => GenerateObservingSystemRecords-{{model_component}} + + # ygyu fix at the end + # Get observations + # Cycling VarBC is inactive, static bias files will be used + GetObsNotInR2d2-{{model_component}}: fail? => GetObservations-{{model_component}} + + # Perform staging that is cycle dependent + StageJediCycle-{{model_component}} + + # Run Jedi variational executable + BuildJediByLinking[^]? | BuildJedi[^] => SP1-{{model_component}} + CloneJedi[^] => StageJediCycle-{{model_component}} + StageJediCycle-{{model_component}} => SP1-{{model_component}} + GetEnsembleGeosExperiment-{{model_component}} => SP1-{{model_component}} + GetObsNotInR2d2-{{model_component}}? | GetObservations-{{model_component}} => SP1-{{model_component}} + SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} + GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} + RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} + + {% if models[model_component]['ensemble_num_chunks'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} + SP2-{{model_component}} => RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}} => SP3-{{model_component}} + {% endfor %} + {% endif %} + + # EvaIncrement + SP3-{{model_component}} => EvaIncrement-{{model_component}} + + {% endif %} + {% endfor %} + """ + {% endfor %} + +# -------------------------------------------------------------------------------------------------- + +[runtime] + + # Task defaults + # ------------- + [[root]] + pre-script = "source $CYLC_SUITE_DEF_PATH/modules" + + [[[environment]]] + datetime = $CYLC_TASK_CYCLE_POINT + config = $CYLC_SUITE_DEF_PATH/experiment.yaml + + # Tasks + # ----- + [[CloneJedi]] + script = "swell task CloneJedi $config" + + [[BuildJediByLinking]] + script = "swell task BuildJediByLinking $config" + + [[BuildJedi]] + script = "swell task BuildJedi $config" + platform = {{platform}} + execution time limit = {{scheduling["BuildJedi"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["BuildJedi"]["directives"]["all"].items() %} + --{{key}} = {{value}} + {%- endfor %} + + {% for model_component in model_components %} + + [[CloneGeosMksi-{{model_component}}]] + script = "swell task CloneGeosMksi $config -m {{model_component}}" + + [[GenerateObservingSystemRecords-{{model_component}}]] + script = "swell task GenerateObservingSystemRecords $config -d $datetime -m {{model_component}}" + + [[GetObsNotInR2d2-{{model_component}}]] + script = "swell task GetObsNotInR2d2 $config -d $datetime -m {{model_component}}" + + [[StageJediCycle-{{model_component}}]] + script = "swell task StageJedi $config -d $datetime -m {{model_component}}" + + [[GetBackground-{{model_component}} ]] + script = "swell task GetBackground $config -d $datetime -m {{model_component}}" + + [[GetBackgroundGeosExperiment-{{model_component}} ]] + script = "swell task GetBackgroundGeosExperiment $config -d $datetime -m {{model_component}}" + + [[GetEnsembleGeosExperiment-{{model_component}}]] + script = "swell task GetEnsembleGeosExperiment $config -d $datetime -m {{model_component}}" + + [[GetObservations-{{model_component}}]] + script = "swell task GetObservations $config -d $datetime -m {{model_component}}" + + [[RenderJediObservations-{{model_component}}]] + script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" + + {% if models[model_component]['ensemble_num_chunks'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} + [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] + script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEdaControlPertExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEdaControlPertExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + {% endfor %} + {% endif %} + + [[RunJediObsfiltersExecutable-{{model_component}}]] + script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediObsfiltersExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediObsfiltersExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + + [[EvaJediLog-{{model_component}}]] + script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}" + + [[EvaIncrement-{{model_component}}]] + script = "swell task EvaIncrement $config -d $datetime -m {{model_component}}" + + [[EvaObservations-{{model_component}}]] + script = "swell task EvaObservations $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["EvaObservations"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["EvaObservations"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[SaveObsDiags-{{model_component}}]] + script = "swell task SaveObsDiags $config -d $datetime -m {{model_component}}" + + [[CleanCycle-{{model_component}}]] + script = "swell task CleanCycle $config -d $datetime -m {{model_component}}" + + # SP: Sync Point + [[SP1-{{model_component}}]] + script = true + + [[SP2-{{model_component}}]] + script = true + + [[SP3-{{model_component}}]] + script = true + + [[SP4-{{model_component}}]] + script = true + + {% endfor %} + +# -------------------------------------------------------------------------------------------------- From c1846eed6dd4df8e849de324373ae2b46db9c48b Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 8 Jul 2026 08:49:13 -0600 Subject: [PATCH 32/68] fixed the problem for parsing in flow.cylc {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} --- .../prepare_config_and_suite.py | 23 ++- src/swell/suites/eda_controlpert/flow.cylc | 186 +++++++++++++++++- .../suites/eda_controlpert/flow.cylc.simple | 17 ++ 3 files changed, 215 insertions(+), 11 deletions(-) diff --git a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py index ae198147c..9f58dba70 100644 --- a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py +++ b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py @@ -443,7 +443,6 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: print(f'nail4.1 self.suite str = {self.suite_str}') - print(f'nail4.2 jinja2 suite str = {suite_str}') # 3. Get a list of tasks that do not depend on the model component @@ -476,10 +475,12 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # At this point the user should have provided the model components answer. Check that it is # in the experiment dictionary and retrieve the response + render_dict = copy.deepcopy(self.experiment_dict) if 'model_components' not in self.experiment_dict: self.logger.abort('The model components question has not been answered.') for model in self.experiment_dict['model_components']: + print(f'nail 5.1 model is {model}') model_dict = self.question_dictionary_model_dep[model] @@ -492,6 +493,15 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # Ask the question self.ask_a_question(model_dict, question_key, model) + # use special task questions for jinja2 + if model_dict[question_key]['question_type'] == 'task': + + ensemble_list = ['ensemble_'+s for s in ['num_members', 'num_chunks']] + if question_key in ensemble_list: + render_dict.update({'models':{model: {question_key: model_dict[question_key]['default_value']}}}) + print(f'render_dict = {render_dict}') + + # 7. Perform a more exhaustive resolving of suite file templates # -------------------------------------------------------------- # Note that we reset the suite file to avoid templates having been left unresolved @@ -499,9 +509,14 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # of templates because there are things related to scheduling that are not yet able to be # resolved. In the future it might be good to bring some of that information into the # sphere of suite questions but that requires some careful thought so as not to overload - # the user with questions. - suite_str = template_string_jinja2(self.logger, self.suite_str, self.experiment_dict, - True) + # the user with questions. A few model dep variables are needed to corectly parse tasks, + # e.g., in Eda ControlPert. Hence we use a render_dict for jinja2 here, only to parse + # the for loop around model task in .cycl file, e.g., + # {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} + # [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] + # {% endfor %} + + suite_str = template_string_jinja2(self.logger, self.suite_str, render_dict, True) # 8. Ask the new task questions that do not actually depend on the model # ----------------------------------------------------------------------- diff --git a/src/swell/suites/eda_controlpert/flow.cylc b/src/swell/suites/eda_controlpert/flow.cylc index 6f657c020..70cd0433c 100644 --- a/src/swell/suites/eda_controlpert/flow.cylc +++ b/src/swell/suites/eda_controlpert/flow.cylc @@ -1,3 +1,15 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +# Cylc suite for executing JEDI-based non-cycling variational data assimilation + +# -------------------------------------------------------------------------------------------------- + [scheduler] UTC mode = True allow implicit tasks = False @@ -15,14 +27,74 @@ # Triggers for non cycle time dependent tasks # ------------------------------------------- # Clone JEDI source code + CloneJedi + + # Build JEDI source code by linking + CloneJedi => BuildJediByLinking? + + # If not able to link to build create the build + BuildJediByLinking:fail? => BuildJedi + {% for model_component in model_components %} + # Clone geos ana for generating observing system records + CloneGeosMksi-{{model_component}} + {% endfor %} + """ + + {% for cycle_time in cycle_times %} + {{cycle_time.cycle_time}} = """ + {% for model_component in model_components %} + {% if cycle_time[model_component] %} + + # logic tree: + # prep (clone, stage, build, get bkg, get obs) -> SP1 + # -> mksi/observation + filter -> SP2 + # -> [ submit/wait for N chunk runs ]-> SP3 + # -> mean/variance + diffstate -> SP4 -> eva + + # Task triggers for: {{model_component}} + # ------------------ + # Generate satellite channel records + CloneGeosMksi-{{model_component}}[^] => GenerateObservingSystemRecords-{{model_component}} + + # ygyu fix at the end + # Get observations + # Cycling VarBC is inactive, static bias files will be used + GetObsNotInR2d2-{{model_component}}: fail? => GetObservations-{{model_component}} + + # Perform staging that is cycle dependent + StageJediCycle-{{model_component}} + + # Run Jedi variational executable + BuildJediByLinking[^]? | BuildJedi[^] => SP1-{{model_component}} + CloneJedi[^] => StageJediCycle-{{model_component}} + StageJediCycle-{{model_component}} => SP1-{{model_component}} + GetEnsembleGeosExperiment-{{model_component}} => SP1-{{model_component}} + GetObsNotInR2d2-{{model_component}}? | GetObservations-{{model_component}} => SP1-{{model_component}} + SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} + GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} + RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} {% if models[model_component]['ensemble_num_chunks'] is defined %} - RunJediEdaControlPertExecutable-{{model_component}} + {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} + SP2-{{model_component}} => RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}} => SP3-{{model_component}} + {% endfor %} {% endif %} - {% endfor %} + # bkg/ana : mean/variance + SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} + RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} + + # CleanupEda directory + SP4-{{model_component}} => CleanEdaFiles-{{model_component}} + + # EvaIncrement + SP4-{{model_component}} => EvaIncrement-{{model_component}} + + {% endif %} + {% endfor %} """ + {% endfor %} # -------------------------------------------------------------------------------------------------- @@ -39,26 +111,126 @@ # Tasks # ----- - [[CloneJedi]] script = "swell task CloneJedi $config" - {% for model_component in model_components %} + [[BuildJediByLinking]] + script = "swell task BuildJediByLinking $config" - {% if models[model_component]['ensemble_num_chunks'] is defined %} - [[RunJediEdaControlPertExecutable-{{model_component}}]] - script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}}" + [[BuildJedi]] + script = "swell task BuildJedi $config" platform = {{platform}} + execution time limit = {{scheduling["BuildJedi"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["BuildJedi"]["directives"]["all"].items() %} + --{{key}} = {{value}} + {%- endfor %} + + {% for model_component in model_components %} + + [[CloneGeosMksi-{{model_component}}]] + script = "swell task CloneGeosMksi $config -m {{model_component}}" + + [[GenerateObservingSystemRecords-{{model_component}}]] + script = "swell task GenerateObservingSystemRecords $config -d $datetime -m {{model_component}}" + + [[GetObsNotInR2d2-{{model_component}}]] + script = "swell task GetObsNotInR2d2 $config -d $datetime -m {{model_component}}" + + [[StageJediCycle-{{model_component}}]] + script = "swell task StageJedi $config -d $datetime -m {{model_component}}" [[GetBackground-{{model_component}} ]] script = "swell task GetBackground $config -d $datetime -m {{model_component}}" + [[GetBackgroundGeosExperiment-{{model_component}} ]] + script = "swell task GetBackgroundGeosExperiment $config -d $datetime -m {{model_component}}" + + [[GetEnsembleGeosExperiment-{{model_component}}]] + script = "swell task GetEnsembleGeosExperiment $config -d $datetime -m {{model_component}}" + + [[GetObservations-{{model_component}}]] + script = "swell task GetObservations $config -d $datetime -m {{model_component}}" + + [[RenderJediObservations-{{model_component}}]] + script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" + + {% if models[model_component]['ensemble_num_chunks'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} + [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] + script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}} -ichunk {{i}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEdaControlPertExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEdaControlPertExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + {% endfor %} {% endif %} + [[RunJediObsfiltersExecutable-{{model_component}}]] + script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediObsfiltersExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediObsfiltersExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[RunJediEnsembleMeanVariance-{{model_component}}]] + script = "swell task RunJediEnsembleMeanVariance $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEnsembleMeanVariance"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEnsembleMeanVariance"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[RunJediDiffstates-{{model_component}}]] + script = "swell task RunJediDiffstates $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediDiffstates"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediDiffstates"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[CleanEdaFiles-{{model_component}}]] + script = "swell task CleanEda $config -d $datetime -m {{model_component}}" + + [[EvaJediLog-{{model_component}}]] + script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}" + + [[EvaIncrement-{{model_component}}]] + script = "swell task EvaIncrement $config -d $datetime -m {{model_component}}" + + [[EvaObservations-{{model_component}}]] + script = "swell task EvaObservations $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["EvaObservations"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["EvaObservations"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[SaveObsDiags-{{model_component}}]] + script = "swell task SaveObsDiags $config -d $datetime -m {{model_component}}" + [[CleanCycle-{{model_component}}]] + script = "swell task CleanCycle $config -d $datetime -m {{model_component}}" + + # SP: Sync Point [[SP1-{{model_component}}]] script = true + [[SP2-{{model_component}}]] + script = true + + [[SP3-{{model_component}}]] + script = true + + [[SP4-{{model_component}}]] + script = true {% endfor %} diff --git a/src/swell/suites/eda_controlpert/flow.cylc.simple b/src/swell/suites/eda_controlpert/flow.cylc.simple index 133dba8f6..6f657c020 100644 --- a/src/swell/suites/eda_controlpert/flow.cylc.simple +++ b/src/swell/suites/eda_controlpert/flow.cylc.simple @@ -17,7 +17,9 @@ # Clone JEDI source code {% for model_component in model_components %} + {% if models[model_component]['ensemble_num_chunks'] is defined %} RunJediEdaControlPertExecutable-{{model_component}} + {% endif %} {% endfor %} """ @@ -38,11 +40,26 @@ # Tasks # ----- + [[CloneJedi]] + script = "swell task CloneJedi $config" + {% for model_component in model_components %} + + {% if models[model_component]['ensemble_num_chunks'] is defined %} [[RunJediEdaControlPertExecutable-{{model_component}}]] script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}}" platform = {{platform}} + [[GetBackground-{{model_component}} ]] + script = "swell task GetBackground $config -d $datetime -m {{model_component}}" + + {% endif %} + + + [[SP1-{{model_component}}]] + script = true + + {% endfor %} # -------------------------------------------------------------------------------------------------- From 625232dbf4f6b5059069f22ffa33020c7f22871f Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 8 Jul 2026 11:09:52 -0600 Subject: [PATCH 33/68] Add models[model_component]['ensemble_num_chunks'] to dict before templating suite_str for flow.cylc, in workflow case: {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] {% endfor %} --- .../prepare_config_and_suite.py | 45 +++++++++++++++++-- 1 file changed, 42 insertions(+), 3 deletions(-) diff --git a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py index 69987bf5f..1eef13165 100644 --- a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py +++ b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py @@ -457,6 +457,7 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # At this point the user should have provided the model components answer. Check that it is # in the experiment dictionary and retrieve the response + render_dict = copy.deepcopy(self.experiment_dict) if 'model_components' not in self.experiment_dict: self.logger.abort('The model components question has not been answered.') @@ -473,6 +474,39 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # Ask the question self.ask_a_question(model_dict, question_key, model) + # pass special model-dep task questions/variables to render_dict + # for suite_str (jinja2) in step 7. + if model_dict[question_key]['question_type'] == 'task': + + ensemble_list = ['ensemble_'+s for s in ['num_members', 'num_chunks']] + if question_key in ensemble_list: + + # a. Ensure the 'models' list exists + if 'models' not in render_dict: + render_dict['models'] = [] + + # b. Check if this model is already in our list + model_found = False + for m in render_dict['models']: + if model in m: + # c. If found, incrementally add the new key! + default_val = model_dict[question_key]['default_value'] + m[model][question_key] = default_val + model_found = True + break + + # d. If not found, add the model to the list for the first time + if not model_found: + default_val = model_dict[question_key]['default_value'] + render_dict['models'].append({ + model: { + question_key: default_val + } + }) + + # debug + # print(f'render_dict = {render_dict}') + # 7. Perform a more exhaustive resolving of suite file templates # -------------------------------------------------------------- # Note that we reset the suite file to avoid templates having been left unresolved @@ -480,9 +514,14 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # of templates because there are things related to scheduling that are not yet able to be # resolved. In the future it might be good to bring some of that information into the # sphere of suite questions but that requires some careful thought so as not to overload - # the user with questions. - suite_str = template_string_jinja2(self.logger, self.suite_str, self.experiment_dict, - True) + # the user with questions. A few model dep variables are needed to corectly parse tasks, + # e.g., in Eda ControlPert. Hence we use a render_dict for jinja2 here, only to parse + # the for loop around model task in .cycl file, e.g., + # {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} + # [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] + # {% endfor %} + + suite_str = template_string_jinja2(self.logger, self.suite_str, render_dict, True) # 8. Ask the new task questions that do not actually depend on the model # ----------------------------------------------------------------------- From 48de41f7fe36ae6078c362461c3ff118e7670c82 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 9 Jul 2026 09:33:03 -0600 Subject: [PATCH 34/68] add ensemble_num_chunks to render_dictionary --- src/swell/deployment/create_experiment.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 46596ce78..2cd3a2924 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -537,7 +537,8 @@ def prepare_cylc_suite_jinja2( 'dictionary.') # Check if 'ensemble_hofx_strategy' appears anywhere in suite_file - ensemble_list = ['ensemble_'+s for s in ['num_members', 'hofx_strategy', 'hofx_packets']] + ensemble_list = ['ensemble_'+s for s in + ['num_members', 'num_chunks', 'hofx_strategy', 'hofx_packets']] ensemble_list = ensemble_list + ['skip_ensemble_hofx'] for ensemble_aspect in ensemble_list: if ensemble_aspect in suite_file: From 84f9cbdb672a0978c270eb2418ea628342f6e169 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 9 Jul 2026 11:40:57 -0600 Subject: [PATCH 35/68] add interfaces/geos_atmosphere/model/eda_analysis_control_pert.py --- .../geos_atmosphere/geos_atmosphere.yaml | 1 + .../model/background_eda_control_pert.py | 83 +++++++++++++++++ .../model/background_error_eda_gsiB.py | 89 +++++++++++++++++++ .../model/eda_analysis_control_pert.py | 31 +++++++ .../jedi/oops/eda_control_pert.py | 81 +++++++++++++++++ src/swell/swell.py | 6 +- src/swell/tasks/base/task_base.py | 18 +++- .../run_jedi_eda_control_pert_executable.py | 27 +++--- .../utilities/render_jedi_interface_files.py | 3 + 9 files changed, 321 insertions(+), 18 deletions(-) create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda_control_pert.py create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py create mode 100644 src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis_control_pert.py create mode 100644 src/swell/configuration/jedi/oops/eda_control_pert.py diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml b/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml index 2fc64df8e..edd40b620 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/geos_atmosphere.yaml @@ -11,3 +11,4 @@ executables: localensembleda: fv3jedi_letkf.x ensmeanvariance: fv3jedi_ensmeanvariance.x obsfilters: test_ObsFilters.x + edaControlPert: fv3jedi_controlpert.x diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda_control_pert.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda_control_pert.py new file mode 100644 index 000000000..8fe781366 --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_eda_control_pert.py @@ -0,0 +1,83 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import \ + field_io_names + +# -------------------------------------------------------------------------------------------------- + +state_variables = [ + 'eastward_wind', + 'northward_wind', + 'air_temperature', + 'air_pressure_at_surface', + 'air_pressure_levels', + 'water_vapor_mixing_ratio_wrt_moist_air', + 'cloud_liquid_ice', + 'cloud_liquid_water', + 'rain_water', + 'snow_water', + 'mole_fraction_of_ozone_in_air', + 'geopotential_height_times_gravity_at_surface', + 'initial_mass_fraction_of_large_scale_cloud_condensate', + 'initial_mass_fraction_of_convective_cloud_condensate', + 'convective_cloud_area_fraction', + 'fraction_of_ocean', + 'fraction_of_land', + 'isotropic_variance_of_filtered_topography', + 'surface_velocity_scale', + 'surface_buoyancy_scale', + 'planetary_boundary_layer_height', + 'surface_exchange_coefficient_for_momentum', + 'surface_exchange_coefficient_for_heat', + 'surface_exchange_coefficient_for_moisture', + 'KCBL_before_moist', + 'surface_temp_before_moist', + 'lower_index_where_Kh_greater_than_2', + 'upper_index_where_Kh_greater_than_2', + 'fraction_of_lake', + 'fraction_of_ice', + 'vtype', + 'stype', + 'vfrac', + 'sheleg', + 'skin_temperature_at_surface', + 'soilt', + 'soilm', + 'eastward_wind_at_surface', + 'northward_wind_at_surface', + # 'sea_surface_temperature', + # 'mole_fraction_of_carbon_dioxide_in_air', +] + +# -------------------------------------------------------------------------------------------------- + + +def background_eda_control_pert(template_dict: Mapping) -> Mapping: + horizontal_resolution = template_dict['horizontal_resolution'] + ichunk = template_dict['ensemble_ichunk'] + cycle_dir = template_dict['cycle_dir'] + background = { + 'datetime': template_dict['local_background_time_iso'], + 'filetype': 'cube sphere history', + 'provider': 'geos', + 'compute edge pressure from surface pressure': True, + 'max allowable geometry difference': 1e-3, + 'datapath': f'{cycle_dir}/ebkg_chunk/chunk{ichunk:03d}/geos.mem%mem_pad%/', + 'filenames': [ + f'%yyyy%mm%dd_%hh%MM%ssz.nc4', + f'../../../fv3-jedi/bkg/geos.crtmsrf.{horizontal_resolution}.nc4' + ], + 'state variables': state_variables, + 'field io names': field_io_names, + } + + return background + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py new file mode 100644 index 000000000..7ff47d587 --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py @@ -0,0 +1,89 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping + +# -------------------------------------------------------------------------------------------------- + +state_variables_to_inverse = [ + 'eastward_wind', + 'northward_wind', + 'air_temperature', + 'air_pressure_at_surface', + 'water_vapor_mixing_ratio_wrt_moist_air', + 'cloud_liquid_ice', + 'cloud_liquid_water', + 'rain_water', + 'snow_water', + 'mole_fraction_of_ozone_in_air', + 'fraction_of_ocean', + 'fraction_of_lake', + 'fraction_of_ice', + 'geopotential_height_times_gravity_at_surface', + 'skin_temperature_at_surface', +] + +# -------------------------------------------------------------------------------------------------- + + +def background_error_eda_gsiB(template_dict: Mapping) -> Mapping: + vertical_resolution = template_dict['vertical_resolution'] + gsibec_nlons = template_dict['gsibec_nlons'] + gsibec_nlats = template_dict['gsibec_nlats'] + gsibec_configuration = template_dict['gsibec_configuration'] + gsibec_npx_proc = template_dict['gsibec_npx_proc'] + gsibec_npy_proc = template_dict['gsibec_npy_proc'] + + background_error = { + 'covariance model': 'SABER', + 'covariance type': 'gsi hybrid covariance', + 'saber central block': { + 'saber block name': 'gsi hybrid covariance', + 'read': { + 'gsi akbk': f'./fv3-jedi/fv3files/akbk{vertical_resolution}.nc4', + 'gsi error covariance file': + f'./fv3-jedi/gsibec/gsi-coeffs-gmao-global-l{vertical_resolution}x{gsibec_nlons}y{gsibec_nlats}.nc4', # noqa + 'gsi berror namelist file': + f'./fv3-jedi/gsibec/{gsibec_configuration}_l{vertical_resolution}x{gsibec_nlons}y{gsibec_nlats}.nml', # noqa + 'processor layout x direction': gsibec_npx_proc, + 'processor layout y direction': gsibec_npy_proc, + 'debugging mode': False + } + }, + 'saber outer blocks': [ + {'saber block name': 'interpolation', + 'inner geometry': { + 'function space': 'StructuredColumns', + 'custom grid matching gsi': { + 'type': 'latlon', + 'lats': int(gsibec_nlats), + 'lons': int(gsibec_nlons), + }, + 'custom partitioner matching gsi': { + 'bands': gsibec_npy_proc + }, + 'halo': 1, + }, + 'forward interpolator': { + 'local interpolator type': 'oops unstructured grid interpolator', + }, + 'inverse interpolator': { + 'local interpolator type': 'oops unstructured grid interpolator', + }, + 'state variables to inverse': state_variables_to_inverse} + ], + 'linear variable change': { + 'linear variable change name': 'Control2Analysis', + 'input variables': state_variables_to_inverse, + 'output variables': template_dict['analysis_variables'] + } + } + + return background_error + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis_control_pert.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis_control_pert.py new file mode 100644 index 000000000..002f2ecee --- /dev/null +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis_control_pert.py @@ -0,0 +1,31 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from collections.abc import Mapping +from swell.configuration.jedi.interfaces.geos_atmosphere.model.shared import field_io_names + +# -------------------------------------------------------------------------------------------------- + + +def eda_analysis_control_pert(template_dict: Mapping) -> Mapping: + + ichunk = template_dict.get('ensemble_ichunk', None) + analysis = { + 'filetype': 'cube sphere history', + 'provider': 'geos', + 'datapath': f'./analysis/chunk{ichunk:03d}/mem%mem_pad%', +## 'filename': f'eda.ana.mem{imem:03d}.%yyyy%mm%dd_%hh%MM%ssz.nc4', +## 'filename': f'eda.ana.%yyyy%mm%dd_%hh%MM%ssz.nc4', + 'first': 'PT0H', + 'frequency': 'PT1H', + 'field io names': field_io_names, + } + + return analysis + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/configuration/jedi/oops/eda_control_pert.py b/src/swell/configuration/jedi/oops/eda_control_pert.py new file mode 100644 index 000000000..68aa72980 --- /dev/null +++ b/src/swell/configuration/jedi/oops/eda_control_pert.py @@ -0,0 +1,81 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + +# -------------------------------------------------------------------------------------------------- + +from swell.utilities.oops_config import OopsConfig + +# -------------------------------------------------------------------------------------------------- + + +class eda_control_pert(OopsConfig): + + def render_oops(self): + nmember = self.template_dict['ensemble_num_members'] + nchunk = self.template_dict['ensemble_num_chunks'] + ichunk = self.template_dict['ensemble_ichunk'] + nstate = int ( nmember / nchunk ) + if ichunk == 1: + num_pert_mem = nstate - 1 + pert_start_index = 2 + else: + num_pert_mem = nstate + pert_start_index = 1 + + oops = { + 'assimilation': { + 'cost function': { + 'cost type': '3D-Var', + 'jb evaluation': False, + 'time window': { + 'begin': self.template_dict['window_begin_iso'], + 'end': self.template_dict['window_end_iso'], + 'bound to include': 'begin' + }, + 'geometry': self.interface_model('geometry'), + 'analysis variables': self.template_dict['analysis_variables'], + 'background': self.interface_model('background_eda_control_pert'), + 'background error': self.interface_model('background_error_eda_gsiB'), + 'observations': { + 'get values': self.interface_model('getvalues'), + 'observers': self.special_observations(), + } + }, + 'variational': { + 'minimizer': { + 'algorithm': self.template_dict['minimizer'] + }, + 'iterations': [{ + 'geometry': self.interface_model('geometry_inner'), + 'gradient norm reduction': float(self.template_dict['gradient_norm_reduction']), + 'ninner': self.template_dict['number_of_iterations'], + }], + }, + 'final': { + 'diagnostics': { + 'departures': 'oman' + }, + 'prints': { + 'frequency': 'PT3H' + } + }, + 'output': self.interface_model('eda_analysis_control_pert') + }, + 'template': { + 'pattern with zero padding': "%mem_pad%", + 'pattern without zero padding': "%mem_wo_pad%", + 'number of pert members': num_pert_mem, + 'first pert member index': pert_start_index, + 'run pert members only': False + } + } + + # TODO: Implement this more cleanly in the OOPS schema + if self.jedi_interface == 'geos_cf': + oops['final']['increment'] = {'geometry': self.interface_model('geometry'), + 'output': self.interface_model('increment_cs')} + + return oops diff --git a/src/swell/swell.py b/src/swell/swell.py index 9beb550f9..e6546a287 100644 --- a/src/swell/swell.py +++ b/src/swell/swell.py @@ -209,13 +209,15 @@ def launch( @click.option('-a', '--additional-parameter', 'additional_parameter', default=None, help=additional_parameter_help) @click.option('-p', '--ensemblePacket', 'ensemblePacket', default=None, help=ensemble_help) +@click.option('-ichunk', '--ensemble_ichunk', 'ichunk', type=int, default=None, help=ensemble_help) def task( task: str, config: str, datetime: Optional[str], model: Optional[str], additional_parameter: Optional[str], - ensemblePacket: Optional[str] + ensemblePacket: Optional[str], + ichunk: Optional[int] ) -> None: """ Run a workflow task @@ -228,7 +230,7 @@ def task( """ task_wrapper(task, config, datetime, model, additional_parameter, - ensemblePacket) + ensemblePacket, ichunk) # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/base/task_base.py b/src/swell/tasks/base/task_base.py index c590262cb..30c5aa553 100644 --- a/src/swell/tasks/base/task_base.py +++ b/src/swell/tasks/base/task_base.py @@ -42,6 +42,7 @@ def __init__( model: str, ensemblePacket: Optional[str], additional_parameter: Optional[str], + ichunk: int | None, task_name: str ) -> None: @@ -70,6 +71,7 @@ def __init__( # Keep copy of ensemblePacket # --------------------------- self.__ensemble_packet__ = ensemblePacket + self.__ensemble_ichunk__ = ichunk # Keep copy of model directive # ---------------------------- @@ -173,6 +175,11 @@ def get_ensemble_packet(self) -> Optional[str]: # ---------------------------------------------------------------------------------------------- + def get_ensemble_ichunk(self) -> int | None: + return self.__ensemble_ichunk__ + + # ---------------------------------------------------------------------------------------------- + def get_model(self) -> str: return self.__model__ @@ -285,7 +292,8 @@ def create_task( datetime: Union[str, dt, None], model: str, additional_parameter: str | None, - ensemblePacket: Optional[str] + ensemblePacket: Optional[str], + ichunk: int | None = None, ) -> taskBase: # Convert camel case string to snake case @@ -314,7 +322,8 @@ def create_task( factory_logger.info(f'Using module swell.tasks.{task_lower}') # Return task object - return task_class(config, datetime, model, ensemblePacket, additional_parameter, task) + return task_class(config, datetime, model, ensemblePacket, + additional_parameter, ichunk, task) # -------------------------------------------------------------------------------------------------- @@ -346,14 +355,15 @@ def task_wrapper( datetime: Union[str, dt, None], model: Optional[str], additional_parameter: str | None, - ensemblePacket: Optional[str] + ensemblePacket: Optional[str], + ichunk: int | None = None, ) -> None: # Create the object constrc_start = time.perf_counter() creator = taskFactory() task_object = creator.create_task(task, config, datetime, model, additional_parameter, - ensemblePacket) + ensemblePacket, ichunk=ichunk) constrc_final = time.perf_counter() constrc_time = f'Constructed in {constrc_final - constrc_start:0.4f} seconds' diff --git a/src/swell/tasks/run_jedi_eda_control_pert_executable.py b/src/swell/tasks/run_jedi_eda_control_pert_executable.py index 34b0df835..dc6a9fe80 100644 --- a/src/swell/tasks/run_jedi_eda_control_pert_executable.py +++ b/src/swell/tasks/run_jedi_eda_control_pert_executable.py @@ -62,10 +62,9 @@ def execute(self) -> None: window_begin_iso = self.da_window_params.window_begin_iso(window_length) window_end_iso = self.da_window_params.window_end_iso(window_length) nmember = self.config.ensemble_num_members() - nchunk = self.config.ensemble_eda_chunk() + nchunk = self.config.ensemble_num_chunks() ichunk = self.get_ensemble_ichunk() - # Populate jedi interface templates dictionary # -------------------------------------------- self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) @@ -86,8 +85,9 @@ def execute(self) -> None: self.jedi_rendering.add_key('horizontal_resolution', self.config.horizontal_resolution()) self.jedi_rendering.add_key('local_background_time', local_background_time) self.jedi_rendering.add_key('local_background_time_iso', local_background_time_iso) - self.jedi_rendering.add_key('ensemble_num_members', self.config.ensemble_num_members()) - self.jedi_rendering.add_key('ensemble_ichunk', ichunk) + self.jedi_rendering.add_key('ensemble_num_members', nmember) + self.jedi_rendering.add_key('ensemble_num_chunks', nchunk) + self.jedi_rendering.add_key('ensemble_ichunk', ichunk) # Geometry # -------- @@ -121,13 +121,13 @@ def execute(self) -> None: self.jedi_rendering.add_key('background_frequency', self.config.background_frequency()) # Jedi configuration file # ----------------------- - str=f'jedi_{jedi_application}{window_type}_config_chunk{{ichunk:03d}}.yaml' - jedi_config_file = os.path.join(self.cycle_dir(), str) + fname=f'jedi_{jedi_application}{window_type}_config_chunk{ichunk:03d}.yaml' + jedi_config_file = os.path.join(self.cycle_dir(), fname) # Output log file # --------------- output_log_file = os.path.join( - self.cycle_dir(), f"jedi_{jedi_application}{window_type}_log_chunk{{ichunk:03d}}.log") + self.cycle_dir(), f"jedi_{jedi_application}{window_type}_log_chunk{ichunk:03d}.log") # Open the JEDI config file and fill initial templates # ---------------------------------------------------- @@ -143,9 +143,10 @@ def execute(self) -> None: xdir = os.path.join(self.cycle_dir(), ebkg_chunk_dir) os.makedirs(xdir, exist_ok=True) - npert = nmember / nchunk + npert = int( nmember / nchunk ) mem_s = npert * (ichunk -1) + 1 - + + print(f'npert, mem_s = {npert} {mem_s}') for imem in range(1, npert+1): id = mem_s + imem - 1 f1 = f'ebkg/mem{id:003d}/' @@ -156,7 +157,7 @@ def execute(self) -> None: xdir = os.path.join(self.cycle_dir(), chunk_dir) os.makedirs(xdir, exist_ok=True) - for observer in jedi_config_dict['cost function']['observations']['observers']: + for observer in jedi_config_dict['assimilation']['cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] print(f'ob= {observation}') @@ -201,6 +202,7 @@ def execute(self) -> None: d2 = os.path.join(self.cycle_dir(), chunk_dir, 'fv3-jedi') shutil.copytree(d1, d2, dirs_exist_ok=True) + with open(jedi_config_file, 'r') as f: yaml_content = f.read() @@ -216,14 +218,15 @@ def execute(self) -> None: # ------------------------------- model_component_meta = self.jedi_rendering.render_interface_meta() + print(f"proc = {model_component_meta['total_processors']}") # Compute number of processors # ---------------------------- np = eval(str(model_component_meta['total_processors'])) # modify np by (nmember / nchunk) + 1 if ichunk == 1: - np = np * (nmember / nchunk) + np = np * npert else: - np = np * (nmember / nchunk + 1) + np = np * (npert + 1) # Jedi executable name diff --git a/src/swell/utilities/render_jedi_interface_files.py b/src/swell/utilities/render_jedi_interface_files.py index e303d754a..30cf40f46 100644 --- a/src/swell/utilities/render_jedi_interface_files.py +++ b/src/swell/utilities/render_jedi_interface_files.py @@ -81,6 +81,8 @@ def __init__( 'ensemble_hofx_packets', 'ensemble_hofx_strategy', 'ensemble_num_members', + 'ensemble_num_chunks', + 'ensemble_ichunk', 'ensmean_only', 'ensmeanvariance_only', 'experiment_id', @@ -305,6 +307,7 @@ def render_interface_meta(self, model_component_in: Union[str, dict, None] = Non # Path to configuration file config_file = os.path.join(self.jedi_config_path, 'interfaces', model_component, f'{model_component}.yaml') + print( f'config_file = {config_file}') # Render templates in file and return dictionary return self.__open_file_render_to_dict__(config_file) From c5b02773197409670546471b17d4d1644eeefee0 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Fri, 10 Jul 2026 12:53:58 -0600 Subject: [PATCH 36/68] link or copy ebkg_chunk [analysis_chunk] / chunkx /memy / xb, yo --- .../run_jedi_eda_control_pert_executable.py | 155 ++++++++++++------ 1 file changed, 101 insertions(+), 54 deletions(-) diff --git a/src/swell/tasks/run_jedi_eda_control_pert_executable.py b/src/swell/tasks/run_jedi_eda_control_pert_executable.py index dc6a9fe80..3167b66af 100644 --- a/src/swell/tasks/run_jedi_eda_control_pert_executable.py +++ b/src/swell/tasks/run_jedi_eda_control_pert_executable.py @@ -8,7 +8,7 @@ # -------------------------------------------------------------------------------------------------- import os -import glob +from glob import glob import shutil from ruamel.yaml import YAML @@ -123,6 +123,7 @@ def execute(self) -> None: # ----------------------- fname=f'jedi_{jedi_application}{window_type}_config_chunk{ichunk:03d}.yaml' jedi_config_file = os.path.join(self.cycle_dir(), fname) + print(f'file = {jedi_config_file}') # Output log file # --------------- @@ -134,61 +135,105 @@ def execute(self) -> None: jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}', window_type, jedi_forecast_model) + print(f'jedi file') +# print(jedi_config_dict) + #exit() - # ichunk: specify yaml - # ---------------------------------------------------- - - # link ebkg to ebkg_chunk - ebkg_chunk_dir = f'ebkg_chunk/chunk{ichunk:003d}/' - xdir = os.path.join(self.cycle_dir(), ebkg_chunk_dir) - os.makedirs(xdir, exist_ok=True) npert = int( nmember / nchunk ) - mem_s = npert * (ichunk -1) + 1 - - print(f'npert, mem_s = {npert} {mem_s}') - for imem in range(1, npert+1): - id = mem_s + imem - 1 - f1 = f'ebkg/mem{id:003d}/' - - - # create subdir for analysis - chunk_dir = f'analysis/chunk{ichunk:003d}/' - xdir = os.path.join(self.cycle_dir(), chunk_dir) - os.makedirs(xdir, exist_ok=True) - - for observer in jedi_config_dict['assimilation']['cost function']['observations']['observers']: - # Get observation name - observation = observer['observation_name'] - print(f'ob= {observation}') - # copy obs input file to avoid multi MPI reading the same file - files = glob.glob(os.path.join(self.cycle_dir(), f'{observation}.*')) - for src_file in files: - print(f'f= {src_file}') - shutil.copy(src_file, xdir) - - obs_bias = observer.get('obs bias') - if obs_bias is not None: - File = obs_bias['input file'] - dir1, fname = os.path.split(File) - File = os.path.join(dir1, chunk_dir, fname) - obs_bias['input file'] = File - # - File = obs_bias['output file'] - dir1, fname = os.path.split(File) - File = os.path.join(dir1, chunk_dir, fname) - obs_bias['output file'] = File - # - File = obs_bias.get('covariance', {}).get('output file') - if File is not None: + if ichunk == 1: + is = 1 # mem: [1, npert] : [ctrl, pert] + else: + is = 0 # mem: [0, 1 ... npert]: [ctrl, all pert] + ie = npert + + print(f'npert = {npert}') + for imem in range(is, ie+1): + if ichunk == 1: + id = imem + else: + if imem == 0: + id = 1 + else: + id = npert * (ichunk -1) + imem + + + # link ebkg to ebkg_chunk + # ---------------------------------------------------- + ebkg_chunk_dir = f'ebkg_chunk/chunk{ichunk:003d}/' + xdir = os.path.join(self.cycle_dir(), ebkg_chunk_dir, f'mem{id:03d}') + os.makedirs(xdir, exist_ok=True) + + # copy bkg files to imem dir + f1_list = glob(os.path.join(self.cycle_dir(), f'ebkg/mem{id:03d}/geos.mem*.nc4')) + f1 = f1_list[0] + print(f'f1 = {f1}') + if f1: + f2 = os.path.basename(f1) + f2 = f2.split('.')[2:] + f2 = '.'.join(f2) + f2 = os.path.join(xdir, f2) + print(f'f2 = {f2}') + else: + f2 = None # Handle the case where no file is found + self.logger.error(f"bkg dir is empty for member id: {id}") + if os.path.lexists(f2): + os.remove(f2) + os.symlink(f1, f2) + + + # create subdir for analysis so that mem00x has its own obs files + # ---------------------------------------------------- + chunk_dir = f'analysis_chunk/chunk{ichunk:003d}' + mem_dir = f'analysis_chunk/chunk{ichunk:003d}/mem{id:03d}' + xdir = os.path.join(self.cycle_dir(), mem_dir) + os.makedirs(xdir, exist_ok=True) + + for observer in jedi_config_dict['assimilation']['cost function']['observations']['observers']: + # Get observation name + observation = observer['observation_name'] + print(f'ob= {observation}') + # copy obs input file to avoid multi MPI reading the same file + files = glob(os.path.join(self.cycle_dir(), f'{observation}.*')) + for src_file in files: + print(f'f= {src_file}') + shutil.copy(src_file, xdir) + + + hxout = observer['obs space']['obsdataout']['engine']['obsfile'] + dir1, fname = os.path.split(hxout) + hxout = os.path.join(dir1, mem_dir, fname) + observer['obs space']['obsdataout']['engine']['obsfile'] = hxout + + obsFileIn = observer['obs space']['obsdatain']['engine']['obsfile'] + dir1, fname = os.path.split(obsFileIn) + obsFileIn = os.path.join(dir1, mem_dir, fname) + observer['obs space']['obsdatain']['engine']['obsfile'] = obsFileIn + + + obs_bias = observer.get('obs bias') + if obs_bias is not None: + File = obs_bias['input file'] dir1, fname = os.path.split(File) - File = os.path.join(dir1, chunk_dir, fname) - obs_bias['covariance']['output file'] = File - File = obs_bias.get('covariance', {}).get('prior', {}).get('input file') - if File is not None: + File = os.path.join(dir1, mem_dir, fname) + obs_bias['input file'] = File + # + File = obs_bias['output file'] dir1, fname = os.path.split(File) - File = os.path.join(dir1, chunk_dir, fname) - obs_bias['covariance']['prior']['input file'] = File + File = os.path.join(dir1, mem_dir, fname) + obs_bias['output file'] = File + # + File = obs_bias.get('covariance', {}).get('output file') + if File is not None: + dir1, fname = os.path.split(File) + File = os.path.join(dir1, mem_dir, fname) + obs_bias['covariance']['output file'] = File + # + File = obs_bias.get('covariance', {}).get('prior', {}).get('input file') + if File is not None: + dir1, fname = os.path.split(File) + File = os.path.join(dir1, mem_dir, fname) + obs_bias['covariance']['prior']['input file'] = File ruamel_yaml = YAML() ruamel_yaml.default_flow_style = False @@ -199,17 +244,19 @@ def execute(self) -> None: # copy fv3-jedi dir, update dir names d1 = os.path.join(self.cycle_dir(), 'fv3-jedi') - d2 = os.path.join(self.cycle_dir(), chunk_dir, 'fv3-jedi') + d2 = os.path.join(self.cycle_dir(), mem_dir, 'fv3-jedi') shutil.copytree(d1, d2, dirs_exist_ok=True) + # point dir to newly created chunks dir + # ---------------------------------------------------- + with open(jedi_config_file, 'r') as f: yaml_content = f.read() - dir_list = ["bkg", "fv3files", "gsibec", "rcov"] for i in dir_list: j = f"fv3-jedi/{i}" - k = f"{chunk_dir}{j}" # Result: analysis/mem002/fv3-jedi/rcov ??? analysis/chunk/mem002/fv3-jedi/rcov + k = f"{mem_dir}{j}" # Result: analysis/chunk00x/mem00y/fv3-jedi/rcov yaml_content = yaml_content.replace(j, k) with open(jedi_config_file, 'w') as f: f.write(yaml_content) From a84926daf5f8b81abc708897bdb5dbcaad5a7e80 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Mon, 13 Jul 2026 14:19:46 -0600 Subject: [PATCH 37/68] . --- .../model/eda_analysis_control_pert.py | 5 +- .../jedi/oops/eda_control_pert.py | 7 +- .../suites/eda_controlpert/suite_config.py | 42 ++--- .../run_jedi_eda_control_pert_executable.py | 171 ++++++++++++------ src/swell/tasks/task_questions.py | 1 + src/swell/utilities/question_defaults.py | 26 +++ 6 files changed, 168 insertions(+), 84 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis_control_pert.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis_control_pert.py index 002f2ecee..c37718f20 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis_control_pert.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/eda_analysis_control_pert.py @@ -18,9 +18,8 @@ def eda_analysis_control_pert(template_dict: Mapping) -> Mapping: analysis = { 'filetype': 'cube sphere history', 'provider': 'geos', - 'datapath': f'./analysis/chunk{ichunk:03d}/mem%mem_pad%', -## 'filename': f'eda.ana.mem{imem:03d}.%yyyy%mm%dd_%hh%MM%ssz.nc4', -## 'filename': f'eda.ana.%yyyy%mm%dd_%hh%MM%ssz.nc4', + 'datapath': f'./analysis_chunk/chunk{ichunk:03d}/mem%mem_pad%', + 'filename': 'eda.ana.mem%mem_pad%.%yyyy%mm%dd_%hh%MM%ssz.nc4', 'first': 'PT0H', 'frequency': 'PT1H', 'field io names': field_io_names, diff --git a/src/swell/configuration/jedi/oops/eda_control_pert.py b/src/swell/configuration/jedi/oops/eda_control_pert.py index 68aa72980..d8426d7db 100644 --- a/src/swell/configuration/jedi/oops/eda_control_pert.py +++ b/src/swell/configuration/jedi/oops/eda_control_pert.py @@ -16,14 +16,13 @@ class eda_control_pert(OopsConfig): def render_oops(self): nmember = self.template_dict['ensemble_num_members'] nchunk = self.template_dict['ensemble_num_chunks'] - ichunk = self.template_dict['ensemble_ichunk'] + ichunk = self.template_dict['ensemble_ichunk'] nstate = int ( nmember / nchunk ) if ichunk == 1: num_pert_mem = nstate - 1 - pert_start_index = 2 else: num_pert_mem = nstate - pert_start_index = 1 + pert_start_index = 1 oops = { 'assimilation': { @@ -32,7 +31,7 @@ def render_oops(self): 'jb evaluation': False, 'time window': { 'begin': self.template_dict['window_begin_iso'], - 'end': self.template_dict['window_end_iso'], + 'length': self.template_dict['window_length'], 'bound to include': 'begin' }, 'geometry': self.interface_model('geometry'), diff --git a/src/swell/suites/eda_controlpert/suite_config.py b/src/swell/suites/eda_controlpert/suite_config.py index c8c4889d2..96b0322b6 100644 --- a/src/swell/suites/eda_controlpert/suite_config.py +++ b/src/swell/suites/eda_controlpert/suite_config.py @@ -105,27 +105,27 @@ class SuiteConfig(QuestionContainer, Enum): "ssmis_f17" ]), qd.obs_thinning_rej_fraction(0.8), -# qd.ensmeanvariance_spec([ -# {"state": "bkg", -# "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", -# "fn_output_mean": "geos.prior.mean", -# "fn_output_variance": "geos.prior.variance", -# "grid_type": ['cs', 'latlon']}, -# {"state": "analysis", -# "fn_input": "analysis/mem%mem%/eda.ana.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", -# "fn_output_mean": "eda.ana.mean", -# "fn_output_variance": "eda.ana.variance", -# "grid_type": ['cs', 'latlon']}, -# ]), -# qd.diffstates_spec({ -# "state1": -# {"fn_input": "geos.prior.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, -# "state2": -# {"fn_input": "eda.ana.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, -# "state_diff": -# {"fn_output": "eda.mean-inc", "grid_type": ['cs', 'latlon']}, -# "state_type": "ensemble" -# }), + qd.ensmeanvariance_spec([ + {"state": "bkg", + "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "geos.prior.mean", + "fn_output_variance": "geos.prior.variance", + "grid_type": ['cs', 'latlon']}, + {"state": "analysis", + "fn_input": "analysis/mem%mem%/eda.ana.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "eda.ana.mean", + "fn_output_variance": "eda.ana.variance", + "grid_type": ['cs', 'latlon']}, + ]), + qd.diffstates_spec({ + "state1": + {"fn_input": "geos.prior.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + "state2": + {"fn_input": "eda.ana.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + "state_diff": + {"fn_output": "eda.mean-inc", "grid_type": ['cs', 'latlon']}, + "state_type": "ensemble" + }), qd.clean_patterns(['*.txt', '*.csv']), ] ) diff --git a/src/swell/tasks/run_jedi_eda_control_pert_executable.py b/src/swell/tasks/run_jedi_eda_control_pert_executable.py index 3167b66af..7d5422616 100644 --- a/src/swell/tasks/run_jedi_eda_control_pert_executable.py +++ b/src/swell/tasks/run_jedi_eda_control_pert_executable.py @@ -8,8 +8,9 @@ # -------------------------------------------------------------------------------------------------- import os -from glob import glob import shutil +from glob import glob +from pathlib import Path from ruamel.yaml import YAML from swell.tasks.base.task_base import taskBase @@ -22,7 +23,11 @@ class RunJediEdaControlPertExecutable(taskBase): # ---------------------------------------------------------------------------------------------- - + # + # ichunk: + # 1:nchunk: normal chunking + # 0 : reorder chunk ana file to analysis/mem00x/ana.nc4 by symlink + # def execute(self) -> None: # Jedi application name @@ -65,6 +70,55 @@ def execute(self) -> None: nchunk = self.config.ensemble_num_chunks() ichunk = self.get_ensemble_ichunk() + # exit execute if ichunk=-1: meaning reoder analysis files + # ------------------------------------------------------- + if ichunk == 0: + npert = int( nmember / nchunk ) + istart = 0 + imem = 0 + for xchunk in range (1, nchunk+1): + if xchunk == 1: + iend = npert -1 # mem: [0, 1 ... npert-1] : [ctrl, all pert] + else: + iend = npert # mem: [0, 1 ... npert] : [ctrl, all pert] + for imem_in_chunk in range(istart, iend+1): + # skip extra ctrl + if xchunk > 1 and imem_in_chunk == 0: + continue + else: + imem += 1 + dir_a = f'analysis_chunk/chunk{xchunk:03d}/mem{imem_in_chunk:03d}' + dir_b = f'analysis/mem{imem:03d}' + dir_a_full = os.path.join(self.cycle_dir(), dir_a, f'eda.ana.mem{imem_in_chunk:03d}.*.nc4') + # print(f'dir_a_full = {dir_a_full}') + fa_list = glob(dir_a_full) + # print(f'fa_list = {fa_list}') + if fa_list: + fa = fa_list[0] + else: + fa = '' + self.logger.error( + f"analysis files not found ichunk={ichunk}, imem_in_chunk={imem_in_chunk}") + print(f'fa = {fa}') + tail = '.'.join(os.path.basename(fa).split('.')[-2:]) + fb = os.path.join(self.cycle_dir(), dir_b, f'eda.ana.mem{imem:03d}.{tail}') + print(f'fb = {fb}') + # link fa to fb + # Convert strings to Path objects + fa = Path(fa) + fb = Path(fb) + # a1. Make the directory if it does not exist + fb.parent.mkdir(parents=True, exist_ok=True) + # a2. If fb exists (or is a broken symlink), safely remove it + if fb.is_symlink() or fb.exists(): + fb.unlink() + # a3. Create the symbolic link + # .resolve() gets the absolute path (like realpath in bash) + fb.symlink_to(fa.resolve()) + print(f"Successfully linked {fa} to {fb}") + return + + # Populate jedi interface templates dictionary # -------------------------------------------- self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) @@ -136,32 +190,34 @@ def execute(self) -> None: window_type, jedi_forecast_model) print(f'jedi file') -# print(jedi_config_dict) - #exit() npert = int( nmember / nchunk ) + istart = 0 if ichunk == 1: - is = 1 # mem: [1, npert] : [ctrl, pert] + iend = npert -1 # mem: [0, 1 ... npert-1] : [ctrl, all pert] else: - is = 0 # mem: [0, 1 ... npert]: [ctrl, all pert] - ie = npert + iend = npert # mem: [0, 1 ... npert] : [ctrl, all pert] print(f'npert = {npert}') - for imem in range(is, ie+1): + + # round-1: link bkg, copy obs, B and R + # ---------------------------------------------------- + chunk_dir = f'analysis_chunk/chunk{ichunk:003d}' + mem_temp_dir = f'analysis_chunk/chunk{ichunk:003d}/mem%mem_pad%' + for imem in range(istart, iend+1): if ichunk == 1: - id = imem + id = imem + 1 else: if imem == 0: id = 1 else: id = npert * (ichunk -1) + imem - # link ebkg to ebkg_chunk - # ---------------------------------------------------- + # ----------------------- ebkg_chunk_dir = f'ebkg_chunk/chunk{ichunk:003d}/' - xdir = os.path.join(self.cycle_dir(), ebkg_chunk_dir, f'mem{id:03d}') + xdir = os.path.join(self.cycle_dir(), ebkg_chunk_dir, f'geos.mem{imem:03d}') os.makedirs(xdir, exist_ok=True) # copy bkg files to imem dir @@ -181,11 +237,9 @@ def execute(self) -> None: os.remove(f2) os.symlink(f1, f2) - - # create subdir for analysis so that mem00x has its own obs files + # analysis_chunk / chunk00x / mem00y will have its own obs, B and R # ---------------------------------------------------- - chunk_dir = f'analysis_chunk/chunk{ichunk:003d}' - mem_dir = f'analysis_chunk/chunk{ichunk:003d}/mem{id:03d}' + mem_dir = f'analysis_chunk/chunk{ichunk:003d}/mem{imem:03d}' xdir = os.path.join(self.cycle_dir(), mem_dir) os.makedirs(xdir, exist_ok=True) @@ -199,41 +253,51 @@ def execute(self) -> None: print(f'f= {src_file}') shutil.copy(src_file, xdir) + # copy fv3-jedi dir, update dir names + d1 = os.path.join(self.cycle_dir(), 'fv3-jedi') + d2 = os.path.join(self.cycle_dir(), mem_dir, 'fv3-jedi') + shutil.copytree(d1, d2, dirs_exist_ok=True) - hxout = observer['obs space']['obsdataout']['engine']['obsfile'] - dir1, fname = os.path.split(hxout) - hxout = os.path.join(dir1, mem_dir, fname) - observer['obs space']['obsdataout']['engine']['obsfile'] = hxout - - obsFileIn = observer['obs space']['obsdatain']['engine']['obsfile'] - dir1, fname = os.path.split(obsFileIn) - obsFileIn = os.path.join(dir1, mem_dir, fname) - observer['obs space']['obsdatain']['engine']['obsfile'] = obsFileIn - - - obs_bias = observer.get('obs bias') - if obs_bias is not None: - File = obs_bias['input file'] + # round-2: modify dir keys in yaml to chunk00x / mem%mem_pad% + # ---------------------------------------------------- + for observer in jedi_config_dict['assimilation']['cost function']['observations']['observers']: + # Get observation name + observation = observer['observation_name'] + print(f'ob= {observation}') + hxout = observer['obs space']['obsdataout']['engine']['obsfile'] + dir1, fname = os.path.split(hxout) + hxout = os.path.join(dir1, mem_temp_dir, fname) + observer['obs space']['obsdataout']['engine']['obsfile'] = hxout + + obsFileIn = observer['obs space']['obsdatain']['engine']['obsfile'] + dir1, fname = os.path.split(obsFileIn) + obsFileIn = os.path.join(dir1, mem_temp_dir, fname) + observer['obs space']['obsdatain']['engine']['obsfile'] = obsFileIn + + obs_bias = observer.get('obs bias') + if obs_bias is not None: + File = obs_bias['input file'] + dir1, fname = os.path.split(File) + File = os.path.join(dir1, mem_temp_dir, fname) + obs_bias['input file'] = File + # + File = obs_bias['output file'] + dir1, fname = os.path.split(File) + File = os.path.join(dir1, mem_temp_dir, fname) + obs_bias['output file'] = File + # + File = obs_bias.get('covariance', {}).get('output file') + if File is not None: dir1, fname = os.path.split(File) - File = os.path.join(dir1, mem_dir, fname) - obs_bias['input file'] = File - # - File = obs_bias['output file'] + File = os.path.join(dir1, mem_temp_dir, fname) + obs_bias['covariance']['output file'] = File + # + File = obs_bias.get('covariance', {}).get('prior', {}).get('input file') + if File is not None: dir1, fname = os.path.split(File) - File = os.path.join(dir1, mem_dir, fname) - obs_bias['output file'] = File - # - File = obs_bias.get('covariance', {}).get('output file') - if File is not None: - dir1, fname = os.path.split(File) - File = os.path.join(dir1, mem_dir, fname) - obs_bias['covariance']['output file'] = File - # - File = obs_bias.get('covariance', {}).get('prior', {}).get('input file') - if File is not None: - dir1, fname = os.path.split(File) - File = os.path.join(dir1, mem_dir, fname) - obs_bias['covariance']['prior']['input file'] = File + File = os.path.join(dir1, mem_temp_dir, fname) + obs_bias['covariance']['prior']['input file'] = File + ruamel_yaml = YAML() ruamel_yaml.default_flow_style = False @@ -242,25 +306,20 @@ def execute(self) -> None: with open(jedi_config_file, 'w') as jedi_config_file_open: ruamel_yaml.dump(jedi_config_dict, jedi_config_file_open) - # copy fv3-jedi dir, update dir names - d1 = os.path.join(self.cycle_dir(), 'fv3-jedi') - d2 = os.path.join(self.cycle_dir(), mem_dir, 'fv3-jedi') - shutil.copytree(d1, d2, dirs_exist_ok=True) - - # point dir to newly created chunks dir + # round-3: point dir to newly created chunks dir # ---------------------------------------------------- - with open(jedi_config_file, 'r') as f: yaml_content = f.read() dir_list = ["bkg", "fv3files", "gsibec", "rcov"] for i in dir_list: j = f"fv3-jedi/{i}" - k = f"{mem_dir}{j}" # Result: analysis/chunk00x/mem00y/fv3-jedi/rcov + k = f"{mem_temp_dir}/{j}" # Result: analysis/chunk00x/mem_pad/fv3-jedi/rcov yaml_content = yaml_content.replace(j, k) with open(jedi_config_file, 'w') as f: f.write(yaml_content) + # Get the JEDI interface metadata # ------------------------------- model_component_meta = self.jedi_rendering.render_interface_meta() diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index 6798e883e..144ba792f 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -672,6 +672,7 @@ class TaskQuestions(QuestionContainer, Enum): window_questions, qd.analysis_variables(), qd.ensemble_num_members(), + qd.ensmeanvariance_spec(), qd.generate_yaml_and_exit(), qd.jedi_forecast_model(), qd.observations(), diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index 44145917b..0ca9fc44f 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -578,6 +578,32 @@ class ensmeanvariance_only(TaskQuestion): # -------------------------------------------------------------------------------------------------- + @dataclass + class ensmeanvariance_spec(TaskQuestion): + default_value: List[Dict[str, str]] = field(default_factory=lambda: [{}]) + question_name: str = "ensmeanvariance_spec" + models: List[str] = mutable_field([ + "all_models" + ]) + ask_question: bool = True + prompt: str = "Configure the ensemble mean and variance specifications:" + widget_type: WType = WType.STRING + + # -------------------------------------------------------------------------------------------------- + + @dataclass + class diffstates_spec(TaskQuestion): + default_value: Dict[str, Any] = field(default_factory=dict) + question_name: str = "diffstates_spec" + models: List[str] = mutable_field([ + "all_models" + ]) + ask_question: bool = True + prompt: str = "Configure the diffstates specifications: [state1, state2]" + widget_type: WType = WType.STRING + + # -------------------------------------------------------------------------------------------------- + @dataclass class existing_geos_gcm_build_path(TaskQuestion): default_value: str = "defer_to_platform" From 067dba5e4a691b3284a2ce9413d9660320dc2c46 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 14 Jul 2026 09:46:39 -0600 Subject: [PATCH 38/68] Only related to ensmeanvariance_spec / diffstates_spec diffstates_spec: is not added here. waiting for EDA PR to be merged first, then use the infrastructure from there. --- src/swell/tasks/run_jedi_diffstates.py | 140 ++++++++++++++++++ src/swell/tasks/task_questions.py | 16 ++ src/swell/utilities/question_defaults.py | 4 +- .../utilities/render_jedi_interface_files.py | 5 + 4 files changed, 163 insertions(+), 2 deletions(-) create mode 100644 src/swell/tasks/run_jedi_diffstates.py diff --git a/src/swell/tasks/run_jedi_diffstates.py b/src/swell/tasks/run_jedi_diffstates.py new file mode 100644 index 000000000..40fc2747b --- /dev/null +++ b/src/swell/tasks/run_jedi_diffstates.py @@ -0,0 +1,140 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + + +# -------------------------------------------------------------------------------------------------- + + +import os +from ruamel.yaml import YAML + +from swell.tasks.base.task_base import taskBase +from swell.utilities.run_jedi_executables import run_executable + + +# -------------------------------------------------------------------------------------------------- + + +class RunJediDiffstates(taskBase): + + # ---------------------------------------------------------------------------------------------- + + def execute(self) -> None: + + # Jedi application name + # --------------------- + jedi_application = 'diffstates' + + # Parse configuration + # ------------------- + jedi_forecast_model = self.config.jedi_forecast_model(None) + generate_yaml_and_exit = self.config.generate_yaml_and_exit(False) + npx_proc = self.config.npx_proc(None) + npy_proc = self.config.npy_proc(None) + + # Compute data assimilation window parameters + window_type = self.config.window_type() + window_length = self.config.window_length() + local_background_time = self.da_window_params.local_background_time(window_length, + window_type) + local_background_time_iso = self.da_window_params.local_background_time_iso(window_length, + window_type) + window_begin_iso = self.da_window_params.window_begin_iso(window_length) + window_end_iso = self.da_window_params.window_end_iso(window_length) + + # Set the observing system records path + self.jedi_rendering.set_obs_records_path(self.config.observing_system_records_path(None)) + + # Populate jedi interface templates dictionary + # -------------------------------------------- + self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) + self.jedi_rendering.add_key('window_end_iso', window_end_iso) + self.jedi_rendering.add_key('window_length', window_length) + self.jedi_rendering.add_key('analysis_variables', self.config.analysis_variables()) + + # Background + self.jedi_rendering.add_key('horizontal_resolution', self.config.horizontal_resolution()) + self.jedi_rendering.add_key('local_background_time', local_background_time) + self.jedi_rendering.add_key('local_background_time_iso', local_background_time_iso) + + # Geometry + self.jedi_rendering.add_key('vertical_resolution', self.config.vertical_resolution()) + self.jedi_rendering.add_key('npx_proc', npx_proc) + self.jedi_rendering.add_key('npy_proc', npy_proc) + + # Diffstates + spec_dict = self.config.diffstates_spec() + self.jedi_rendering.add_key('diffstates_spec', spec_dict) + print(f'diffstates = {spec_dict}') + + # Add placeholder names if mock experiment + # ---------------------------------------- + if self.config.mock_experiment(False): + self.jedi_rendering.add_key('experiment_root', 'experiment_root') + self.jedi_rendering.add_key('experiment_id', 'experiment_id') + self.jedi_rendering.add_key('cycle_dir', 'cycle_dir') + + state_type = spec_dict.get('state_type') + self.jedi_rendering.add_key('diffstates_spec_statetype', state_type) + + # loop output grid_type: + grid_type = spec_dict['state_diff'].get('grid_type') + print(f'grid_type = {grid_type}') + + if grid_type is None: + grid_type = ['latlon'] + + for output_grid in grid_type: + + self.jedi_rendering.add_key('diffstates_spec_gridtype', output_grid) + + # Jedi configuration file + # ----------------------- + jedi_config_file = os.path.join( + self.cycle_dir(), f'jedi_{jedi_application}_output_{output_grid}_config.yaml') + + # Output log file + # --------------- + output_log_file = os.path.join(self.cycle_dir(), + f'jedi_{jedi_application}_output_{output_grid}.log') + + # Open the JEDI config file and fill templates + # -------------------------------------------- + jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}', + window_type, + jedi_forecast_model) + yaml = YAML() + yaml.default_flow_style = False + + # Write the expanded dictionary to YAML file + # ------------------------------------------ + with open(jedi_config_file, 'w') as jedi_config_file_open: + yaml.dump(jedi_config_dict, jedi_config_file_open) + + # Get the JEDI interface metadata + # ------------------------------- + model_component_meta = self.jedi_rendering.render_interface_meta() + + # Compute number of processors + # ---------------------------- + np = eval(str(model_component_meta['total_processors'])) + + # Jedi executable name + # -------------------- + jedi_executable = model_component_meta['executables'][f'{jedi_application}'] + jedi_executable_path = os.path.join(self.experiment_path(), + 'jedi_bundle', 'build', 'bin', jedi_executable) + + # Run the JEDI executable + # ----------------------- + if not generate_yaml_and_exit: + self.logger.info('Running '+jedi_executable_path+' with '+str(np)+' processors.') + run_executable(self.logger, self.cycle_dir(), np, jedi_executable_path, + jedi_config_file, output_log_file) + else: + self.logger.info('YAML generated, now exiting.') + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index 144ba792f..0c2aece30 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -684,6 +684,22 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + RunJediDiffstates = QuestionList( + list_name="RunJediDiffstates", + questions=[ + np_proc_resolution, + window_questions, + qd.analysis_variables(), + qd.diffstates_spec(), + qd.generate_yaml_and_exit(), + qd.jedi_forecast_model(), + qd.comparison_log_type('diffstates'), + qd.mock_experiment() + ] + ) + + # -------------------------------------------------------------------------------------------------- + RunJediFgatExecutable = QuestionList( list_name="RunJediFgatExecutable", questions=[ diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index 0ca9fc44f..bb4132e5e 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -8,8 +8,8 @@ # -------------------------------------------------------------------------------------------------- -from dataclasses import dataclass -from typing import List, Dict +from dataclasses import dataclass, field +from typing import List, Dict, Any from swell.utilities.swell_questions import SuiteQuestion, TaskQuestion from swell.utilities.swell_questions import WidgetType as WType diff --git a/src/swell/utilities/render_jedi_interface_files.py b/src/swell/utilities/render_jedi_interface_files.py index 30cf40f46..1f2a32485 100644 --- a/src/swell/utilities/render_jedi_interface_files.py +++ b/src/swell/utilities/render_jedi_interface_files.py @@ -85,6 +85,11 @@ def __init__( 'ensemble_ichunk', 'ensmean_only', 'ensmeanvariance_only', + 'ensmeanvariance_spec_item', + 'ensmeanvariance_spec_gridtype', + 'diffstates_spec', + 'diffstates_spec_gridtype', + 'diffstates_spec_statetype', 'experiment_id', 'experiment_root', 'final_cycle_point', From 39580fb1bd6135bc2093070882b904b74e400fb8 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 15 Jul 2026 09:50:00 -0600 Subject: [PATCH 39/68] - rename suite/eda to suite/eda_atmos - change gradient_norm_reduction to from 1.e-7 to 1.e-3 --- .../{eda => eda_atmos}/eva/increment-geos_atmosphere.yaml | 0 .../suites/{eda => eda_atmos}/eva/jedi_log-geos_atmosphere.yaml | 0 .../{eda => eda_atmos}/eva/observations-geos_atmosphere.yaml | 0 src/swell/suites/{eda => eda_atmos}/flow.cylc | 0 src/swell/suites/{eda => eda_atmos}/suite_config.py | 2 +- 5 files changed, 1 insertion(+), 1 deletion(-) rename src/swell/suites/{eda => eda_atmos}/eva/increment-geos_atmosphere.yaml (100%) rename src/swell/suites/{eda => eda_atmos}/eva/jedi_log-geos_atmosphere.yaml (100%) rename src/swell/suites/{eda => eda_atmos}/eva/observations-geos_atmosphere.yaml (100%) rename src/swell/suites/{eda => eda_atmos}/flow.cylc (100%) rename src/swell/suites/{eda => eda_atmos}/suite_config.py (99%) diff --git a/src/swell/suites/eda/eva/increment-geos_atmosphere.yaml b/src/swell/suites/eda_atmos/eva/increment-geos_atmosphere.yaml similarity index 100% rename from src/swell/suites/eda/eva/increment-geos_atmosphere.yaml rename to src/swell/suites/eda_atmos/eva/increment-geos_atmosphere.yaml diff --git a/src/swell/suites/eda/eva/jedi_log-geos_atmosphere.yaml b/src/swell/suites/eda_atmos/eva/jedi_log-geos_atmosphere.yaml similarity index 100% rename from src/swell/suites/eda/eva/jedi_log-geos_atmosphere.yaml rename to src/swell/suites/eda_atmos/eva/jedi_log-geos_atmosphere.yaml diff --git a/src/swell/suites/eda/eva/observations-geos_atmosphere.yaml b/src/swell/suites/eda_atmos/eva/observations-geos_atmosphere.yaml similarity index 100% rename from src/swell/suites/eda/eva/observations-geos_atmosphere.yaml rename to src/swell/suites/eda_atmos/eva/observations-geos_atmosphere.yaml diff --git a/src/swell/suites/eda/flow.cylc b/src/swell/suites/eda_atmos/flow.cylc similarity index 100% rename from src/swell/suites/eda/flow.cylc rename to src/swell/suites/eda_atmos/flow.cylc diff --git a/src/swell/suites/eda/suite_config.py b/src/swell/suites/eda_atmos/suite_config.py similarity index 99% rename from src/swell/suites/eda/suite_config.py rename to src/swell/suites/eda_atmos/suite_config.py index 29723c323..7a9c1c6f1 100644 --- a/src/swell/suites/eda/suite_config.py +++ b/src/swell/suites/eda_atmos/suite_config.py @@ -50,7 +50,7 @@ class SuiteConfig(QuestionContainer, Enum): qd.ensemble_num_members(32), qd.obs_pert_amplitude(0.5), qd.number_of_iterations([100]), - qd.gradient_norm_reduction(1.e-8), + qd.gradient_norm_reduction(1.e-3), qd.analysis_variables([ "eastward_wind", "northward_wind", From 7d36f4b1850eedfc9465e1329c8b5dec5c6ff710 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 15 Jul 2026 10:29:39 -0600 Subject: [PATCH 40/68] - add to minizer DRPLanczos --- .../jedi/interfaces/geos_atmosphere/task_questions.yaml | 1 + src/swell/suites/eda_controlpert/suite_config.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml b/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml index 88700667f..71d4d575f 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml @@ -212,6 +212,7 @@ minimizer: default_value: DRPCG options: - DRPCG + - DRPLanczos ncdiag_experiments: default_value: x0050_fgat diff --git a/src/swell/suites/eda_controlpert/suite_config.py b/src/swell/suites/eda_controlpert/suite_config.py index 96b0322b6..5b672ef92 100644 --- a/src/swell/suites/eda_controlpert/suite_config.py +++ b/src/swell/suites/eda_controlpert/suite_config.py @@ -50,7 +50,8 @@ class SuiteConfig(QuestionContainer, Enum): qd.ensemble_num_members(32), qd.ensemble_num_chunks(8), qd.number_of_iterations([100]), - qd.gradient_norm_reduction(1.e-8), + qd.gradient_norm_reduction(1.e-4), + qd.minimizer("DRPLanczos"), qd.analysis_variables([ "eastward_wind", "northward_wind", From 310c355b04f5debf9412d8529532fab9488b9006 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 15 Jul 2026 12:35:58 -0600 Subject: [PATCH 41/68] - add eda_atmos = QuestionList to suite_config.py - change print to logger in clean_eda.py --- src/swell/suites/eda_atmos/suite_config.py | 4 ++-- src/swell/tasks/clean_eda.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/swell/suites/eda_atmos/suite_config.py b/src/swell/suites/eda_atmos/suite_config.py index 7a9c1c6f1..1fbf238fc 100644 --- a/src/swell/suites/eda_atmos/suite_config.py +++ b/src/swell/suites/eda_atmos/suite_config.py @@ -132,8 +132,8 @@ class SuiteConfig(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- - eda = QuestionList( - list_name="eda", + eda_atmos = QuestionList( + list_name="eda_atmos", questions=[ eda_atmos_tier1 ] diff --git a/src/swell/tasks/clean_eda.py b/src/swell/tasks/clean_eda.py index be5ba32ad..7b745ffda 100644 --- a/src/swell/tasks/clean_eda.py +++ b/src/swell/tasks/clean_eda.py @@ -127,10 +127,10 @@ def execute(self) -> None: if os.path.islink(d2): # Only remove the symlink itself, never follow it os.unlink(d2) - print(f"Deleted symlink dir: {d2}") + self.logger.info(f"Deleted symlink dir: {d2}") else: shutil.rmtree(d2, ignore_errors=True) - print(f"Deleted directory: {d2}") + self.logger.info(f"Deleted directory: {d2}") for observer in jedi_config_dict['cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] @@ -138,11 +138,11 @@ def execute(self) -> None: for file_path in glob.glob(os.path.join(d1, f'{observation}.*')): if os.path.islink(file_path): os.unlink(file_path) # safe: removes only the link - print(f"Deleted symlink file: {file_path}") + self.logger.info(f"Deleted symlink file: {file_path}") else: try: os.remove(file_path) - print(f"Deleted file: {file_path}") + self.logger.info(f"Deleted file: {file_path}") except FileNotFoundError: pass # file disappeared between glob and remove From 8a71c327420d15e6496aebe8ca297c1a1685ceba Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 15 Jul 2026 22:53:45 -0600 Subject: [PATCH 42/68] add back `air_pressure_levels` to analysis variables in eda_atmos --- .../geos_atmosphere/model/background_error_eda_gsiB.py | 1 + src/swell/suites/eda_atmos/suite_config.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py index 7ff47d587..6350dbedd 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_gsiB.py @@ -15,6 +15,7 @@ 'northward_wind', 'air_temperature', 'air_pressure_at_surface', + 'air_pressure_levels', 'water_vapor_mixing_ratio_wrt_moist_air', 'cloud_liquid_ice', 'cloud_liquid_water', diff --git a/src/swell/suites/eda_atmos/suite_config.py b/src/swell/suites/eda_atmos/suite_config.py index 1fbf238fc..b5156e20f 100644 --- a/src/swell/suites/eda_atmos/suite_config.py +++ b/src/swell/suites/eda_atmos/suite_config.py @@ -57,6 +57,7 @@ class SuiteConfig(QuestionContainer, Enum): "air_temperature", "water_vapor_mixing_ratio_wrt_moist_air", "air_pressure_at_surface", + "air_pressure_levels", "cloud_liquid_ice", "cloud_liquid_water", "rain_water", From ded0a296f3e36a54d1aad316b70e19abcf0078fd Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 16 Jul 2026 11:45:23 -0600 Subject: [PATCH 43/68] many modifications/corrections from reviewers' suggestions --- .../geos_atmosphere/model/background_error_eda_hybridB.py | 4 ++-- .../interfaces/geos_atmosphere/model/diffstates_output.py | 2 +- .../jedi/interfaces/geos_atmosphere/model/two_states.py | 2 +- src/swell/configuration/jedi/oops/eda3D.py | 2 ++ src/swell/suites/eda_atmos/flow.cylc | 2 +- src/swell/tasks/run_jedi_diffstates.py | 4 ++-- src/swell/tasks/run_jedi_eda_executable.py | 6 ++++++ src/swell/utilities/render_jedi_interface_files.py | 4 ++++ 8 files changed, 19 insertions(+), 7 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py index 5c5cfa8ee..e1c2acc68 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/background_error_eda_hybridB.py @@ -212,8 +212,8 @@ def background_error_eda_hybridB(template_dict: Mapping) -> Mapping: # Use ruamel.yaml to load the string instead of standard pyyaml ruamel_yaml = YAML() background_error = ruamel_yaml.load(rendered_yaml_string) - cov_template = background_error['components'][1]['covariance'] - ['members from template']['template'] + cov_template = (background_error['components'][1]['covariance'] + ['members from template']['template']) cov_template['field io names'] = field_io_names a = background_error['components'][0]['covariance']['linear variable change'] a['output variables'] = template_dict['analysis_variables'] diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py index 36119c395..403d3319b 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/diffstates_output.py @@ -20,7 +20,7 @@ def diffstates_output(template_dict: Mapping) -> Mapping: prefix_output = template_dict['diffstates_spec'].get('state_diff', {}).get('fn_output') stateType = template_dict.get('diffstates_spec_statetype') - if [stateType == 'ensemble']: + if stateType == 'ensemble': field_io_names_loc = field_io_names_sa1 else: field_io_names_loc = field_io_names diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py index e81146ba5..95586185a 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/two_states.py @@ -18,7 +18,7 @@ def two_states(template_dict: Mapping) -> Mapping: f2 = template_dict['diffstates_spec']['state2']['fn_input'] stateType = template_dict.get('diffstates_spec_statetype') - if [stateType == 'ensemble']: + if stateType == 'ensemble': state_variables_loc = state_variables_sa1 field_io_names_loc = field_io_names_sa1 else: diff --git a/src/swell/configuration/jedi/oops/eda3D.py b/src/swell/configuration/jedi/oops/eda3D.py index e811ec356..82df403de 100644 --- a/src/swell/configuration/jedi/oops/eda3D.py +++ b/src/swell/configuration/jedi/oops/eda3D.py @@ -14,6 +14,8 @@ class eda3D(OopsConfig): def render_oops(self): + assert self.template_dict['ensemble_imember'] != None, \ + f'ensemble_imember is None and eda3D.py fails' oops = { 'cost function': { 'cost type': '3D-Var', diff --git a/src/swell/suites/eda_atmos/flow.cylc b/src/swell/suites/eda_atmos/flow.cylc index da3221f0c..b52ac4ff0 100644 --- a/src/swell/suites/eda_atmos/flow.cylc +++ b/src/swell/suites/eda_atmos/flow.cylc @@ -89,7 +89,7 @@ # bkg/ana : mean/variance SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} - RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} + RunJediEnsembleMeanVariance-{{model_component}} => RunJediDiffstates-{{model_component}} => SP4-{{model_component}} # CleanupEda directory SP4-{{model_component}} => CleanEdaFiles-{{model_component}} diff --git a/src/swell/tasks/run_jedi_diffstates.py b/src/swell/tasks/run_jedi_diffstates.py index 40fc2747b..e86fc7319 100644 --- a/src/swell/tasks/run_jedi_diffstates.py +++ b/src/swell/tasks/run_jedi_diffstates.py @@ -68,7 +68,7 @@ def execute(self) -> None: # Diffstates spec_dict = self.config.diffstates_spec() self.jedi_rendering.add_key('diffstates_spec', spec_dict) - print(f'diffstates = {spec_dict}') + self.logger.info(f'diffstates = {spec_dict}') # Add placeholder names if mock experiment # ---------------------------------------- @@ -82,7 +82,7 @@ def execute(self) -> None: # loop output grid_type: grid_type = spec_dict['state_diff'].get('grid_type') - print(f'grid_type = {grid_type}') + self.logger.info(f'grid_type = {grid_type}') if grid_type is None: grid_type = ['latlon'] diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 2486972d6..e1c6616eb 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -205,6 +205,12 @@ def execute(self) -> None: File = os.path.join(dir1, mem_dir, fname) obs_bias['covariance']['prior']['input file'] = File + + print('jedi_config_dict') + print(jedi_config_dict) + print('end jedi_config_dict') + exit() + ruamel_yaml = YAML() ruamel_yaml.default_flow_style = False diff --git a/src/swell/utilities/render_jedi_interface_files.py b/src/swell/utilities/render_jedi_interface_files.py index 7e0c6029d..068534cb0 100644 --- a/src/swell/utilities/render_jedi_interface_files.py +++ b/src/swell/utilities/render_jedi_interface_files.py @@ -65,6 +65,10 @@ def __init__( self.__template_dict__['experiment_id'] = experiment_id self.__template_dict__['experiment_root'] = experiment_root + # Predefine eda variables to dictionary + self.__template_dict__['ensemble_imember'] = None + self.__template_dict__['ensemble_ichunk'] = None + # List of all potential valid keys that can be used in templates self.valid_template_keys = [ 'analysis_variables', From 06b7a0469d453c208ddd583af021c46e615c62b4 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 16 Jul 2026 16:22:53 -0600 Subject: [PATCH 44/68] small corrections to codes --- .../geos_atmosphere/task_questions.yaml | 6 ++- src/swell/tasks/run_jedi_eda_executable.py | 41 +++++++------------ .../tasks/run_jedi_ensemble_mean_variance.py | 2 +- src/swell/tasks/task_questions.py | 1 + src/swell/utilities/question_defaults.py | 2 +- src/swell/utilities/yaml_utils.py | 24 +++++++++++ 6 files changed, 47 insertions(+), 29 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml b/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml index 7a413d5de..0d698410c 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml @@ -125,7 +125,11 @@ ensemble_hofx_strategy: - parallel ensemble_num_members: - default_value: 10 + default_value: 16 + options: None + +obs_pert_amplitude: + default_value: 1.0 options: None geovals_experiment: diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index e1c6616eb..9a93b183b 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -14,7 +14,7 @@ from swell.tasks.base.task_base import taskBase from swell.utilities.run_jedi_executables import run_executable - +from swell.utilities.yaml_utils import replace_string_value # -------------------------------------------------------------------------------------------------- @@ -149,20 +149,19 @@ def execute(self) -> None: for observer in jedi_config_dict['cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] - print(f'ob= {observation}') # copy obs input file to avoid multi MPI reading the same file files = glob.glob(os.path.join(self.cycle_dir(), f'{observation}.*')) for src_file in files: - print(f'f= {src_file}') + self.logger.info(f'f= {src_file}') shutil.copy(src_file, xdir) if imember > 1: obs_cov_model = observer.get('obs error', {}).get('covariance model') - print(f'{observation}: obs_cov_model = {obs_cov_model}') + self.logger.info(f'{observation}: obs_cov_model = {obs_cov_model}') if obs_cov_model and 'cross variable covariances' in obs_cov_model: - print(f"Found cross covariance obs: {obs_cov_model}, skip perturbation") + self.logger.info(f"Found cross covariance obs: {obs_cov_model}, skip perturbation") else: - print(f"No cross variable covariance found: {observation}, Obs Error Diagonal") + self.logger.info(f"No cross variable covariance found: {observation}, Obs Error Diagonal") obs_error_dict = { 'covariance model': 'diagonal', 'zero-mean perturbations': True, @@ -205,34 +204,24 @@ def execute(self) -> None: File = os.path.join(dir1, mem_dir, fname) obs_bias['covariance']['prior']['input file'] = File - - print('jedi_config_dict') - print(jedi_config_dict) - print('end jedi_config_dict') - exit() - - ruamel_yaml = YAML() - ruamel_yaml.default_flow_style = False - - # Write the ordered dictionary to YAML file - with open(jedi_config_file, 'w') as jedi_config_file_open: - ruamel_yaml.dump(jedi_config_dict, jedi_config_file_open) - # copy fv3-jedi dir, update dir names d1 = os.path.join(self.cycle_dir(), 'fv3-jedi') d2 = os.path.join(self.cycle_dir(), mem_dir, 'fv3-jedi') shutil.copytree(d1, d2, dirs_exist_ok=True) - with open(jedi_config_file, 'r') as f: - yaml_content = f.read() - dir_list = ["bkg", "fv3files", "gsibec", "rcov"] for i in dir_list: j = f"fv3-jedi/{i}" - k = f"{mem_dir}{j}" # Result: analysis/mem002/fv3-jedi/rcov - yaml_content = yaml_content.replace(j, k) - with open(jedi_config_file, 'w') as f: - f.write(yaml_content) + k = f"{mem_dir}{j}" + jedi_config_dict = replace_string_value(jedi_config_dict, j, k) + # Result: e.g., analysis/mem002/fv3-jedi/rcov + + ruamel_yaml = YAML() + ruamel_yaml.default_flow_style = False + + # Write the ordered dictionary to YAML file + with open(jedi_config_file, 'w') as jedi_config_file_open: + ruamel_yaml.dump(jedi_config_dict, jedi_config_file_open) # Get the JEDI interface metadata # ------------------------------- diff --git a/src/swell/tasks/run_jedi_ensemble_mean_variance.py b/src/swell/tasks/run_jedi_ensemble_mean_variance.py index 24f571e82..d57444650 100644 --- a/src/swell/tasks/run_jedi_ensemble_mean_variance.py +++ b/src/swell/tasks/run_jedi_ensemble_mean_variance.py @@ -71,7 +71,7 @@ def execute(self) -> None: # Ensemble self.jedi_rendering.add_key('ensemble_num_members', self.config.ensemble_num_members(None)) - print(f'self.config.ensmeanvariance_spec = {self.config.ensmeanvariance_spec()}') + self.logger.info(f'self.config.ensmeanvariance_spec = {self.config.ensmeanvariance_spec()}') meanvar_spec_dict = self.config.ensmeanvariance_spec() # Add placeholder names if mock experiment diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index 4542dd4e7..7c8449ef7 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -842,6 +842,7 @@ class TaskQuestions(QuestionContainer, Enum): questions=[ run_jedi_executable, qd.ensemble_num_members(), + qd.obs_pert_amplitude(), qd.obs_thinning_rej_fraction(), qd.perhost(), qd.comparison_log_type('variational'), diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index 75190d0a8..1804effbb 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -552,7 +552,7 @@ class obs_pert_amplitude(TaskQuestion): question_name: str = "obs_pert_amplitude" options: str = "defer_to_model" models: List[str] = mutable_field([ - "geos_atmosphere" + "all_models" ]) prompt: str = "Enter obs perturbation amplitude for EDA:" widget_type: WType = WType.FLOAT diff --git a/src/swell/utilities/yaml_utils.py b/src/swell/utilities/yaml_utils.py index fc73d28e8..c15810552 100644 --- a/src/swell/utilities/yaml_utils.py +++ b/src/swell/utilities/yaml_utils.py @@ -16,4 +16,28 @@ def replace_key(obj, old_key, new_key): else: return obj + +def replace_string_value(data, sa, sb): + """ + Recursively search through a dictionary or list and replaces specific + sa by sb in the value content. sa/sb = string_a/b + """ + if isinstance(data, dict): + # If it's a dictionary, apply recursively to all values + return {k: replace_string_value(v, sa, sb) for k, v in data.items()} + + elif isinstance(data, list): + # If it's a list, apply recursively to all items + return [replace_string_value(item, sa, sb) for item in data] + + elif isinstance(data, str): + # If it's a string, check for our target keys and replace + if sa in data: + # Replaces sa with sb + data = data.replace(sa, sb) + return data + else: + # Return integers, booleans, etc., as-is + return data + # -------------------------------------------------------------------------------------------------- From 20e0c79231868ad2bad8ada1864be2e3805637f9 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Fri, 17 Jul 2026 11:28:12 -0600 Subject: [PATCH 45/68] Add obs_pert_amplitude to config --- src/swell/tasks/run_jedi_eda_executable.py | 2 ++ src/swell/tasks/task_questions.py | 17 ++++++++++++++++- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 9a93b183b..fa1db9adf 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -63,6 +63,7 @@ def execute(self) -> None: window_end_iso = self.da_window_params.window_end_iso(window_length) nmember = self.config.ensemble_num_members() imember = self.get_ensemble_imember() + obs_pert_amplitude = self.obs_pert_amplitude() # Populate jedi interface templates dictionary # -------------------------------------------- @@ -164,6 +165,7 @@ def execute(self) -> None: self.logger.info(f"No cross variable covariance found: {observation}, Obs Error Diagonal") obs_error_dict = { 'covariance model': 'diagonal', + 'obs perturbations amplitude': obs_pert_amplitude, 'zero-mean perturbations': True, 'member': imember, 'number of members': nmember diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index 7c8449ef7..e8dcda2c9 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -840,7 +840,22 @@ class TaskQuestions(QuestionContainer, Enum): RunJediEdaExecutable = QuestionList( list_name="RunJediEdaExecutable", questions=[ - run_jedi_executable, +# runJediExecutable + background_crtm_obs, + np_proc_resolution, + window_questions, + qd.analysis_variables(), + qd.background_frequency(), + qd.generate_yaml_and_exit(), + qd.gradient_norm_reduction(), + qd.gsibec_configuration(), + qd.jedi_forecast_model(), + qd.minimizer(), + qd.gsibec_nlats(), + qd.gsibec_nlons(), + qd.number_of_iterations(), + qd.total_processors(), +# add qd.ensemble_num_members(), qd.obs_pert_amplitude(), qd.obs_thinning_rej_fraction(), From 7e1b083a0862162f43b697e438f096cac7cab479 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Fri, 17 Jul 2026 12:31:55 -0600 Subject: [PATCH 46/68] bypass the problem in flow.cylc which reads {% if models[model_component]['ensemble_num_members'] is defined %} {% for i in range( 1, models[model_component]['ensemble_num_members'] + 1 ) %} --- .../geos_atmosphere/task_questions.yaml | 4 ---- src/swell/suites/eda_atmos/flow.cylc | 11 +++++++++++ src/swell/tasks/run_jedi_eda_executable.py | 2 +- src/swell/tasks/task_questions.py | 17 +---------------- 4 files changed, 13 insertions(+), 21 deletions(-) diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml b/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml index 0d698410c..0ed059e2e 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/task_questions.yaml @@ -128,10 +128,6 @@ ensemble_num_members: default_value: 16 options: None -obs_pert_amplitude: - default_value: 1.0 - options: None - geovals_experiment: default_value: x0050-geovals options: diff --git a/src/swell/suites/eda_atmos/flow.cylc b/src/swell/suites/eda_atmos/flow.cylc index b52ac4ff0..4f866af21 100644 --- a/src/swell/suites/eda_atmos/flow.cylc +++ b/src/swell/suites/eda_atmos/flow.cylc @@ -174,6 +174,17 @@ {% endfor %} {% endif %} + + [[RunJediEdaExecutable-{{model_component}}]] + script = "swell task RunJediEdaExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEdaExecutable"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunJediEdaExecutable"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[RunJediObsfiltersExecutable-{{model_component}}]] script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" platform = {{platform}} diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index fa1db9adf..3a8fe95c0 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -61,9 +61,9 @@ def execute(self) -> None: window_begin = self.da_window_params.window_begin(window_length) window_begin_iso = self.da_window_params.window_begin_iso(window_length) window_end_iso = self.da_window_params.window_end_iso(window_length) + obs_pert_amplitude = self.config.obs_pert_amplitude() nmember = self.config.ensemble_num_members() imember = self.get_ensemble_imember() - obs_pert_amplitude = self.obs_pert_amplitude() # Populate jedi interface templates dictionary # -------------------------------------------- diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index e8dcda2c9..7c8449ef7 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -840,22 +840,7 @@ class TaskQuestions(QuestionContainer, Enum): RunJediEdaExecutable = QuestionList( list_name="RunJediEdaExecutable", questions=[ -# runJediExecutable - background_crtm_obs, - np_proc_resolution, - window_questions, - qd.analysis_variables(), - qd.background_frequency(), - qd.generate_yaml_and_exit(), - qd.gradient_norm_reduction(), - qd.gsibec_configuration(), - qd.jedi_forecast_model(), - qd.minimizer(), - qd.gsibec_nlats(), - qd.gsibec_nlons(), - qd.number_of_iterations(), - qd.total_processors(), -# add + run_jedi_executable, qd.ensemble_num_members(), qd.obs_pert_amplitude(), qd.obs_thinning_rej_fraction(), From 166aabcaed88f4f09049842cbb72af5caa462b8c Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Fri, 17 Jul 2026 16:32:27 -0600 Subject: [PATCH 47/68] fix pycodestyle issues --- src/swell/configuration/jedi/oops/eda3D.py | 3 ++- src/swell/tasks/run_jedi_eda_executable.py | 8 ++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/swell/configuration/jedi/oops/eda3D.py b/src/swell/configuration/jedi/oops/eda3D.py index 82df403de..e87d0a57f 100644 --- a/src/swell/configuration/jedi/oops/eda3D.py +++ b/src/swell/configuration/jedi/oops/eda3D.py @@ -14,8 +14,9 @@ class eda3D(OopsConfig): def render_oops(self): - assert self.template_dict['ensemble_imember'] != None, \ + assert self.template_dict['ensemble_imember'] is not None, ( f'ensemble_imember is None and eda3D.py fails' + ) oops = { 'cost function': { 'cost type': '3D-Var', diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 3a8fe95c0..7a0936b63 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -160,9 +160,13 @@ def execute(self) -> None: obs_cov_model = observer.get('obs error', {}).get('covariance model') self.logger.info(f'{observation}: obs_cov_model = {obs_cov_model}') if obs_cov_model and 'cross variable covariances' in obs_cov_model: - self.logger.info(f"Found cross covariance obs: {obs_cov_model}, skip perturbation") + self.logger.info( + f"Found cross covariance obs: {obs_cov_model}, skip perturbation" + ) else: - self.logger.info(f"No cross variable covariance found: {observation}, Obs Error Diagonal") + self.logger.info( + f"No cross variable covariance found: {observation}, Obs Error Diagonal" + ) obs_error_dict = { 'covariance model': 'diagonal', 'obs perturbations amplitude': obs_pert_amplitude, From 5ee5427b39405cec8b42522290640a6c3be3e3e3 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Mon, 20 Jul 2026 15:07:03 -0600 Subject: [PATCH 48/68] many debug info for prepare_config_and_suite --- src/swell/deployment/create_experiment.py | 10 +++++---- .../prepare_config_and_suite.py | 21 +++++++++++++++---- src/swell/suites/eda_controlpert/flow.cylc | 8 +++++++ src/swell/swell.py | 2 +- src/swell/utilities/yaml_utils.py | 15 +++++++++++++ 5 files changed, 47 insertions(+), 9 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 1a1fc2637..aa4cf8fb6 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -26,7 +26,7 @@ from swell.utilities.logger import Logger, get_logger from swell.utilities.slurm import prepare_scheduling_dict from swell.utilities.check_da_params import check_da_params - +from swell.utilities.yaml_utils import print_dict # -------------------------------------------------------------------------------------------------- @@ -261,8 +261,6 @@ def create_experiment_directory( experiment_id = dict_get(logger, experiment_dict, 'experiment_id') experiment_root = dict_get(logger, experiment_dict, 'experiment_root') - - # Write out some info # ------------------- logger.info(f'Creating experiment: \'{experiment_id}\' in \'{experiment_root}\'') @@ -488,7 +486,9 @@ def prepare_cylc_suite_jinja2( with open(os.path.join(swell_suite_path, 'flow.cylc'), 'r') as file: suite_file = file.read() print(f'nail6. suite_file in prepare_cylc_suite_jinja2 = {suite_file}') - + print(f'experiment_dict =') + print_dict(experiment_dict) + exit() # Copy the experiment dictionary to the rendering dictionary @@ -547,6 +547,8 @@ def prepare_cylc_suite_jinja2( ensemble_list = ['ensemble_'+s for s in ['num_members', 'num_chunks', 'hofx_strategy', 'hofx_packets']] ensemble_list = ensemble_list + ['skip_ensemble_hofx'] + print (f'ensemble_list = {ensemble_list}') + for ensemble_aspect in ensemble_list: if ensemble_aspect in suite_file: if len(model_components) == 0: diff --git a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py index 275ebc32b..5859c59c7 100644 --- a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py +++ b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py @@ -23,6 +23,7 @@ from swell.utilities.dictionary import update_dict from swell.tasks.task_questions import TaskQuestions as task_questions from swell.suites.all_suites import AllSuites +from swell.utilities.yaml_utils import print_dict # -------------------------------------------------------------------------------------------------- @@ -94,7 +95,7 @@ def __init__( print(f'nail1 model_ind_tasks = {self.model_ind_tasks}') print(f'nail2 all_model_dep_tasks = {self.all_model_dep_tasks}') - print(f'nail3 .suite_str = {self.suite_str}') + print(f'nail3 self.suite_str = {self.suite_str}') # Perform the assembly of the dictionaries that contain all the questions that can possibly @@ -127,6 +128,7 @@ def prepare_question_dictionaries(self) -> None: suite_config_obj = AllSuites.get_config(self.suite_config) suite_question_list = suite_config_obj.expand_question_list() print(f'suite q list : {suite_question_list}') + print_dict(suite_question_list) # Allow for adding extra tasks manually from configuration @@ -209,6 +211,8 @@ def prepare_question_dictionaries(self) -> None: if 'model_components' not in question_dictionary.keys(): self.suite_needs_model_components = False + print(f'self.suite_needs_model_components = {self.suite_needs_model_components}') + # Create copy of the question_dictionary for model independent questions question_dictionary_model_ind = copy.deepcopy(question_dictionary) @@ -440,7 +444,8 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: suite_str = template_string_jinja2(self.logger, self.suite_str, self.experiment_dict, True) print(f'nail4.0 self.exp dict = {self.experiment_dict}') - + print_dict(self.experiment_dict) + print(f'nail4.1 self.suite str = {self.suite_str}') print(f'nail4.2 jinja2 suite str = {suite_str}') @@ -476,6 +481,10 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # in the experiment dictionary and retrieve the response render_dict = copy.deepcopy(self.experiment_dict) + print('r1: render_dict') + print_dict(render_dict) + + if 'model_components' not in self.experiment_dict: self.logger.abort('The model components question has not been answered.') @@ -524,8 +533,12 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: }) # debug - # print(f'render_dict = {render_dict}') +# print(f'render_dict = {render_dict}') + print('r2: render_dict') + print_dict(render_dict) + + # 7. Perform a more exhaustive resolving of suite file templates # -------------------------------------------------------------- # Note that we reset the suite file to avoid templates having been left unresolved @@ -559,8 +572,8 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # ------------------------------------------------- print(f'ck suite_str = {suite_str}') model_dep_tasks = self.get_suite_task_list_model_dep(suite_str) + print(f'ck show: model_dep_tasks = {model_dep_tasks}') - # 9.2 Iterate over the model_dep dictionary and ask task questions # ---------------------------------------------------------------- for model in self.experiment_dict['model_components']: diff --git a/src/swell/suites/eda_controlpert/flow.cylc b/src/swell/suites/eda_controlpert/flow.cylc index 70cd0433c..ecf9a0d56 100644 --- a/src/swell/suites/eda_controlpert/flow.cylc +++ b/src/swell/suites/eda_controlpert/flow.cylc @@ -168,6 +168,14 @@ {% endfor %} {% endif %} + + [[RunJediEdaControlPertExecutable_]] + script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunJediEdaControlPertExecutable"]["execution_time_limit"]}} + [[[directives]]] + + [[RunJediObsfiltersExecutable-{{model_component}}]] script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" platform = {{platform}} diff --git a/src/swell/swell.py b/src/swell/swell.py index 5a621fb5f..496dc8d3b 100644 --- a/src/swell/swell.py +++ b/src/swell/swell.py @@ -218,7 +218,7 @@ def task( model: Optional[str], additional_parameter: Optional[str], ensemblePacket: Optional[str], - ichunk: Optional[int] + ichunk: Optional[int], imember: Optional[int] ) -> None: """ diff --git a/src/swell/utilities/yaml_utils.py b/src/swell/utilities/yaml_utils.py index c15810552..a37c5982f 100644 --- a/src/swell/utilities/yaml_utils.py +++ b/src/swell/utilities/yaml_utils.py @@ -1,3 +1,4 @@ +import yaml # -------------------------------------------------------------------------------------------------- @@ -41,3 +42,17 @@ def replace_string_value(data, sa, sb): return data # -------------------------------------------------------------------------------------------------- + +def print_dict(*args, **kwargs): + return print_dict_as_yaml(*args, **kwargs) + +def print_dict_as_yaml(data_dict): + # Convert the dictionary to a YAML string + # default_flow_style=False ensures block formatting instead of inline JSON-like formatting + # sort_keys=False preserves your dictionary's original key order (Python 3.7+) + yaml_string = yaml.dump(data_dict, default_flow_style=False, sort_keys=False) + + # Output to the terminal + print(yaml_string) + +# -------------------------------------------------------------------------------------------------- From f0c78ee832f7f03fc7370589e6b2ac0abef6fa3c Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Mon, 20 Jul 2026 17:03:11 -0600 Subject: [PATCH 49/68] . --- .../{eda_controlpert => eda_controlpert_atmos}/flow.cylc | 0 .../flow.cylc.complete | 0 .../flow.cylc.simple | 0 .../{eda_controlpert => eda_controlpert_atmos}/flow.cylc.v2 | 0 .../suite_config.py | 6 +++--- 5 files changed, 3 insertions(+), 3 deletions(-) rename src/swell/suites/{eda_controlpert => eda_controlpert_atmos}/flow.cylc (100%) rename src/swell/suites/{eda_controlpert => eda_controlpert_atmos}/flow.cylc.complete (100%) rename src/swell/suites/{eda_controlpert => eda_controlpert_atmos}/flow.cylc.simple (100%) rename src/swell/suites/{eda_controlpert => eda_controlpert_atmos}/flow.cylc.v2 (100%) rename src/swell/suites/{eda_controlpert => eda_controlpert_atmos}/suite_config.py (97%) diff --git a/src/swell/suites/eda_controlpert/flow.cylc b/src/swell/suites/eda_controlpert_atmos/flow.cylc similarity index 100% rename from src/swell/suites/eda_controlpert/flow.cylc rename to src/swell/suites/eda_controlpert_atmos/flow.cylc diff --git a/src/swell/suites/eda_controlpert/flow.cylc.complete b/src/swell/suites/eda_controlpert_atmos/flow.cylc.complete similarity index 100% rename from src/swell/suites/eda_controlpert/flow.cylc.complete rename to src/swell/suites/eda_controlpert_atmos/flow.cylc.complete diff --git a/src/swell/suites/eda_controlpert/flow.cylc.simple b/src/swell/suites/eda_controlpert_atmos/flow.cylc.simple similarity index 100% rename from src/swell/suites/eda_controlpert/flow.cylc.simple rename to src/swell/suites/eda_controlpert_atmos/flow.cylc.simple diff --git a/src/swell/suites/eda_controlpert/flow.cylc.v2 b/src/swell/suites/eda_controlpert_atmos/flow.cylc.v2 similarity index 100% rename from src/swell/suites/eda_controlpert/flow.cylc.v2 rename to src/swell/suites/eda_controlpert_atmos/flow.cylc.v2 diff --git a/src/swell/suites/eda_controlpert/suite_config.py b/src/swell/suites/eda_controlpert_atmos/suite_config.py similarity index 97% rename from src/swell/suites/eda_controlpert/suite_config.py rename to src/swell/suites/eda_controlpert_atmos/suite_config.py index 5b672ef92..e3389bb12 100644 --- a/src/swell/suites/eda_controlpert/suite_config.py +++ b/src/swell/suites/eda_controlpert_atmos/suite_config.py @@ -50,7 +50,7 @@ class SuiteConfig(QuestionContainer, Enum): qd.ensemble_num_members(32), qd.ensemble_num_chunks(8), qd.number_of_iterations([100]), - qd.gradient_norm_reduction(1.e-4), + qd.gradient_norm_reduction(1.e-3), qd.minimizer("DRPLanczos"), qd.analysis_variables([ "eastward_wind", @@ -133,8 +133,8 @@ class SuiteConfig(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- - eda_controlpert = QuestionList( - list_name="eda_controlpert", + eda_controlpert_atmos = QuestionList( + list_name="eda_controlpert_atmos", questions=[ eda_controlpert_atmos_tier1 ] From 527b3fdd343fd03dc48187a4fc714f44665b0136 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 21 Jul 2026 09:51:27 -0600 Subject: [PATCH 50/68] - ichunk=1 for ShuffleAnaMembers - use dict to replace string i by j for jedi config dict --- src/swell/deployment/create_experiment.py | 1 - .../suites/eda_controlpert_atmos/flow.cylc | 25 +++++++++------ .../eda_controlpert_atmos/suite_config.py | 1 + .../run_jedi_eda_control_pert_executable.py | 31 +++++++++---------- 4 files changed, 31 insertions(+), 27 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index aa4cf8fb6..30ed44337 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -488,7 +488,6 @@ def prepare_cylc_suite_jinja2( print(f'nail6. suite_file in prepare_cylc_suite_jinja2 = {suite_file}') print(f'experiment_dict =') print_dict(experiment_dict) - exit() # Copy the experiment dictionary to the rendering dictionary diff --git a/src/swell/suites/eda_controlpert_atmos/flow.cylc b/src/swell/suites/eda_controlpert_atmos/flow.cylc index ecf9a0d56..da246ddf1 100644 --- a/src/swell/suites/eda_controlpert_atmos/flow.cylc +++ b/src/swell/suites/eda_controlpert_atmos/flow.cylc @@ -57,10 +57,21 @@ # Generate satellite channel records CloneGeosMksi-{{model_component}}[^] => GenerateObservingSystemRecords-{{model_component}} - # ygyu fix at the end # Get observations + {% if cycling_varbc %} + # Cycling VarBC is active, biases from the previous cycle will be used + + {% if models[model_component]['ensemble_num_chunks'] is defined %} + {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} + RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}[-PT6H] => GetObservations-{{model_component}} + {% endfor %} + {% endif %} + + {% else %} + # Cycling VarBC is inactive, static bias files will be used GetObsNotInR2d2-{{model_component}}: fail? => GetObservations-{{model_component}} + {% endif %} # Perform staging that is cycle dependent StageJediCycle-{{model_component}} @@ -82,7 +93,8 @@ {% endif %} # bkg/ana : mean/variance - SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} + SP3-{{model_component}} => ShuffleAnaMembers-{{model_component}} + ShuffleAnaMembers-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} # CleanupEda directory @@ -168,13 +180,8 @@ {% endfor %} {% endif %} - - [[RunJediEdaControlPertExecutable_]] - script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediEdaControlPertExecutable"]["execution_time_limit"]}} - [[[directives]]] - + [[ShuffleAnaMembers-{{model_component}}]] + script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}} -ichunk -1" [[RunJediObsfiltersExecutable-{{model_component}}]] script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" diff --git a/src/swell/suites/eda_controlpert_atmos/suite_config.py b/src/swell/suites/eda_controlpert_atmos/suite_config.py index e3389bb12..fbbda16dd 100644 --- a/src/swell/suites/eda_controlpert_atmos/suite_config.py +++ b/src/swell/suites/eda_controlpert_atmos/suite_config.py @@ -58,6 +58,7 @@ class SuiteConfig(QuestionContainer, Enum): "air_temperature", "water_vapor_mixing_ratio_wrt_moist_air", "air_pressure_at_surface", + "air_pressure_levels", "cloud_liquid_ice", "cloud_liquid_water", "rain_water", diff --git a/src/swell/tasks/run_jedi_eda_control_pert_executable.py b/src/swell/tasks/run_jedi_eda_control_pert_executable.py index 7d5422616..4885e098b 100644 --- a/src/swell/tasks/run_jedi_eda_control_pert_executable.py +++ b/src/swell/tasks/run_jedi_eda_control_pert_executable.py @@ -15,7 +15,7 @@ from swell.tasks.base.task_base import taskBase from swell.utilities.run_jedi_executables import run_executable - +from swell.utilities.yaml_utils import replace_string_value # -------------------------------------------------------------------------------------------------- @@ -72,7 +72,7 @@ def execute(self) -> None: # exit execute if ichunk=-1: meaning reoder analysis files # ------------------------------------------------------- - if ichunk == 0: + if ichunk == -1: npert = int( nmember / nchunk ) istart = 0 imem = 0 @@ -238,7 +238,7 @@ def execute(self) -> None: os.symlink(f1, f2) # analysis_chunk / chunk00x / mem00y will have its own obs, B and R - # ---------------------------------------------------- + # ----------------------------------------------------------------- mem_dir = f'analysis_chunk/chunk{ichunk:003d}/mem{imem:03d}' xdir = os.path.join(self.cycle_dir(), mem_dir) os.makedirs(xdir, exist_ok=True) @@ -259,7 +259,7 @@ def execute(self) -> None: shutil.copytree(d1, d2, dirs_exist_ok=True) # round-2: modify dir keys in yaml to chunk00x / mem%mem_pad% - # ---------------------------------------------------- + # ------------------------------------------------------------ for observer in jedi_config_dict['assimilation']['cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] @@ -298,26 +298,23 @@ def execute(self) -> None: File = os.path.join(dir1, mem_temp_dir, fname) obs_bias['covariance']['prior']['input file'] = File - - ruamel_yaml = YAML() - ruamel_yaml.default_flow_style = False - - # Write the ordered dictionary to YAML file - with open(jedi_config_file, 'w') as jedi_config_file_open: - ruamel_yaml.dump(jedi_config_dict, jedi_config_file_open) - + observer['obs space']['obs perturbations seed shift'] = (ichunk-1)*npert # round-3: point dir to newly created chunks dir # ---------------------------------------------------- - with open(jedi_config_file, 'r') as f: - yaml_content = f.read() dir_list = ["bkg", "fv3files", "gsibec", "rcov"] for i in dir_list: j = f"fv3-jedi/{i}" k = f"{mem_temp_dir}/{j}" # Result: analysis/chunk00x/mem_pad/fv3-jedi/rcov - yaml_content = yaml_content.replace(j, k) - with open(jedi_config_file, 'w') as f: - f.write(yaml_content) + jedi_config_dict = replace_string_value(jedi_config_dict, j, k) + # Result: e.g., analysis/mem002/fv3-jedi/rcov + + ruamel_yaml = YAML() + ruamel_yaml.default_flow_style = False + + # Write the ordered dictionary to YAML file + with open(jedi_config_file, 'w') as jedi_config_file_open: + ruamel_yaml.dump(jedi_config_dict, jedi_config_file_open) # Get the JEDI interface metadata From 5485008b9a2532de529a1ab5f4e8f8fff1567828 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 21 Jul 2026 12:26:21 -0600 Subject: [PATCH 51/68] add eva dir to eda_controlpert_atmos --- .../eva/increment-geos_atmosphere.yaml | 760 ++++++++++++++++++ .../eva/jedi_log-geos_atmosphere.yaml | 89 ++ .../eva/observations-geos_atmosphere.yaml | 564 +++++++++++++ src/swell/tasks/eva_increment.py | 2 +- 4 files changed, 1414 insertions(+), 1 deletion(-) create mode 100644 src/swell/suites/eda_controlpert_atmos/eva/increment-geos_atmosphere.yaml create mode 100644 src/swell/suites/eda_controlpert_atmos/eva/jedi_log-geos_atmosphere.yaml create mode 100644 src/swell/suites/eda_controlpert_atmos/eva/observations-geos_atmosphere.yaml diff --git a/src/swell/suites/eda_controlpert_atmos/eva/increment-geos_atmosphere.yaml b/src/swell/suites/eda_controlpert_atmos/eva/increment-geos_atmosphere.yaml new file mode 100644 index 000000000..126435876 --- /dev/null +++ b/src/swell/suites/eda_controlpert_atmos/eva/increment-geos_atmosphere.yaml @@ -0,0 +1,760 @@ +datasets: + +- group: increment + type: LatLon + filename: {{increment_file_path}} + name: experiment_increment + variables: [ps, ts, ua, va, t, q, lat, lon] + +graphics: + + plotting_backend: Emcpy + figure_list: + + #map plot for surface pressure increment + - batch figure: + variables: [ps] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Surface Pressure Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ps + label: PS increment + colorbar: true + cmap: 'bwr' + vmin: -100 + vmax: 100 + - batch figure: + variables: [ts] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Tskin Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ts + label: TS increment + colorbar: true + cmap: 'bwr' + vmin: -2 + vmax: 2 + #map plot for temperature increment (lowest level) + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1000.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[71,...]' + label: T increment (1000 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for temperature increment + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_850.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[62,...]' + label: T increment (850 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for temperature increment + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_500.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[49,...]' + label: T increment (500 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for temperature increment + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_200.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[42,...]' + label: T increment (200 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for temperature increment + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_10.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[24,...]' + label: T increment (10 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for temperature increment + - batch figure: + variables: [t] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Temperature Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::t + slices: '[14,...]' + label: T increment (1 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment (lowest level) + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1000.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[71,...]' + label: U increment (1000 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_850.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[62,...]' + label: U increment (850 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_500.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[49,...]' + label: U increment (500 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_200.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[42,...]' + label: U increment (200 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_10.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[24,...]' + label: U increment (10 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Zonal Wind increment + - batch figure: + variables: [ua] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Zonal Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::ua + slices: '[14,...]' + label: U increment (1 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment (lowest level) + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1000.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[71,...]' + label: V increment (1000 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_850.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[62,...]' + label: V increment (850 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_500.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[49,...]' + label: V increment (500 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_200.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[42,...]' + label: V increment (200 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_10.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[24,...]' + label: V increment (10 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Meridional Wind increment + - batch figure: + variables: [va] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Meridional Wind Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::va + slices: '[14,...]' + label: V increment (1 hPa) + colorbar: true + cmap: 'bwr' + vmin: -1 + vmax: 1 + #map plot for Specific Humidity increment (lowest level) + - batch figure: + variables: [q] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_1000.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Specific Humidity Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::q + slices: '[71,...]' + label: Q increment (1000 hPa) + colorbar: true + cmap: 'bwr' + vmin: -0.001 + vmax: 0.001 + #map plot for Specific Humidity increment + - batch figure: + variables: [q] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_850.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Specific Humidity Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::q + slices: '[62,...]' + label: Q increment (850 hPa) + colorbar: true + cmap: 'bwr' + vmin: -0.001 + vmax: 0.001 + #map plot for Specific Humidity increment + - batch figure: + variables: [q] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_500.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Specific Humidity Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::q + slices: '[49,...]' + label: Q increment (500 hPa) + colorbar: true + cmap: 'bwr' + vmin: -0.001 + vmax: 0.001 + #map plot for Specific Humidity increment + - batch figure: + variables: [q] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_200.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Specific Humidity Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::q + slices: '[42,...]' + label: Q increment (200 hPa) + colorbar: true + cmap: 'bwr' + vmin: -0.001 + vmax: 0.001 + #map plot for Specific Humidity increment + - batch figure: + variables: [q] + figure: + figure size: [20,10] + layout: [1,1] + title: 'Increment from JEDI' + output name: '{{cycle_dir}}/eva/increment/map_plots/${variable}/inc_${variable}_10.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: Specific Humidity Increment + add_grid: + layers: + - type: MapGridded + longitude: + variable: experiment_increment::increment::lon + latitude: + variable: experiment_increment::increment::lat + data: + variable: experiment_increment::increment::q + slices: '[24,...]' + label: Q increment (10 hPa) + colorbar: true + cmap: 'bwr' + vmin: -0.001 + vmax: 0.001 diff --git a/src/swell/suites/eda_controlpert_atmos/eva/jedi_log-geos_atmosphere.yaml b/src/swell/suites/eda_controlpert_atmos/eva/jedi_log-geos_atmosphere.yaml new file mode 100644 index 000000000..2d848d549 --- /dev/null +++ b/src/swell/suites/eda_controlpert_atmos/eva/jedi_log-geos_atmosphere.yaml @@ -0,0 +1,89 @@ +datasets: + +- type: JediLog + collection_name: JediLogTest + jedi_log_to_parse: '{{cycle_dir}}/jedi_variational_log.log' + data_to_parse: + convergence: true + +transforms: +- transform: arithmetic + new name: JediLogTest::convergence::${variable}_log + equals: log(JediLogTest::convergence::${variable}) + for: + variable: [residual_norm, norm_reduction] + +graphics: + + plotting_backend: Emcpy + figure_list: + + - figure: + layout: [3,1] + figure size: [12,10] + title: 'Residual Norm and Norm Reduction Plots' + output name: '{{cycle_dir}}/eva/jedi_log/convergence/residual_norm_reduction.png' + plots: + - add_xlabel: 'Total inner iteration number' + add_ylabel: 'Residual norm' + layers: + - type: LinePlot + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::residual_norm + color: 'black' + + - add_xlabel: 'Total inner iteration number' + add_ylabel: 'Log(norm reduction)' + layers: + - type: LinePlot + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::norm_reduction + color: 'black' + + - add_xlabel: 'Total inner iteration number' + add_ylabel: 'Log(reduction)' + add_legend: + layers: + - type: LinePlot + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::residual_norm_log + color: 'red' + label: 'Log(residual norm)' + - type: LinePlot + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::norm_reduction_log + color: 'blue' + label: 'Log norm reduction' + + - figure: + title: 'Cost Function Plot' + output name: '{{cycle_dir}}/eva/jedi_log/cost_function/cost_function.png' + plots: + - add_xlabel: 'Total inner iteration number' + add_ylabel: 'Quadratic Cost Function' + add_legend: + layers: + - type: LinePlot + label: 'jojc' + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::jojc + color: 'blue' + markersize: 2 + - type: LinePlot + label: 'jb' + x: + variable: JediLogTest::convergence::total_iteration + y: + variable: JediLogTest::convergence::jb + color: 'red' + markersize: 2 diff --git a/src/swell/suites/eda_controlpert_atmos/eva/observations-geos_atmosphere.yaml b/src/swell/suites/eda_controlpert_atmos/eva/observations-geos_atmosphere.yaml new file mode 100644 index 000000000..edee73160 --- /dev/null +++ b/src/swell/suites/eda_controlpert_atmos/eva/observations-geos_atmosphere.yaml @@ -0,0 +1,564 @@ +datasets: + +- name: experiment + type: IodaObsSpace + filenames: + - {{obs_path_file}} + channels: &channels {{channels}} + groups: + - name: ObsValue + variables: &variables {{simulated_variables}} + - name: GsiHofXBc + #- name: GsiEffectiveQC + - name: hofx0 + - name: hofx1 + - name: ombg + - name: oman + - name: EffectiveQC0 + - name: EffectiveQC1 + - name: MetaData + +transforms: + +# Generate hofx0 for GSI +- transform: arithmetic + new name: experiment::ObsValueMinusGsiHofXBc::${variable} + equals: experiment::ObsValue::${variable}-experiment::GsiHofXBc::${variable} + for: + variable: *variables + +# Generate hofx0 for JEDI +- transform: arithmetic + new name: experiment::ObsValueMinusHofx0::${variable} + equals: experiment::ObsValue::${variable}-experiment::hofx0::${variable} + for: + variable: *variables + +# Generate hofx difference +- transform: arithmetic + new name: experiment::Hofx0MinusGsiHofXBc::${variable} + equals: experiment::hofx0::${variable}-experiment::GsiHofXBc::${variable} + for: + variable: *variables + +# Generate hofx that passed QC for JEDI +- transform: accept where + new name: experiment::hofx0PassedQc::${variable} + starting field: experiment::hofx0::${variable} + where: + - experiment::EffectiveQC0::${variable} == 0 + for: + variable: *variables + +# Generate GSI hofx that passed JEDI QC +- transform: accept where + new name: experiment::GsiHofXBcPassedQc::${variable} + starting field: experiment::GsiHofXBc::${variable} + where: + - experiment::EffectiveQC0::${variable} == 0 + for: + variable: *variables + +# Generate hofx0 that passed QC for JEDI +- transform: accept where + new name: experiment::ObsValueMinushofx0PassedQc::${variable} + starting field: experiment::ObsValueMinusHofx0::${variable} + where: + - experiment::EffectiveQC0::${variable} == 0 + for: + variable: *variables + +# Generate hofx0 that passed QC for GSI +- transform: accept where + new name: experiment::ObsValueMinusGsiHofXBcPassedQc::${variable} + starting field: experiment::ObsValueMinusGsiHofXBc::${variable} + where: + - experiment::EffectiveQC0::${variable} == 0 + for: + variable: *variables + +# Generate ombg that passed QC for JEDI +- transform: accept where + new name: experiment::ombgPassedQc::${variable} + starting field: experiment::ombg::${variable} + where: + - experiment::EffectiveQC0::${variable} == 0 + for: + variable: *variables + +# Generate oman that passed QC for JEDI +- transform: accept where + new name: experiment::omanPassedQc::${variable} + starting field: experiment::oman::${variable} + where: + - experiment::EffectiveQC1::${variable} == 0 + for: + variable: *variables + +# Generate obs contribution to analysis (OmA*OmA)-(OmB*OmB) +- transform: arithmetic + new name: experiment::ResidualRMSdiff::${variable} + equals: (experiment::omanPassedQc::${variable})*(experiment::omanPassedQc::${variable})-(experiment::ombgPassedQc::${variable})*(experiment::ombgPassedQc::${variable}) + for: + variable: *variables + +graphics: + + plotting_backend: Emcpy + figure_list: + + # Correlation scatter plots + # ------------------------- + + # JEDI h(x) vs Observations + - batch figure: + variables: *variables + channels: *channels + figure: + layout: [1,1] + title: 'Observations vs. JEDI h(x) | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/correlation_scatter/${variable}${channel}/jedi_hofx0_vs_obs_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'Observation Value' + add_ylabel: 'JEDI h(x)' + add_grid: + add_legend: + loc: 'upper left' + layers: + - type: Scatter + x: + variable: experiment::ObsValue::${variable} + y: + variable: experiment::hofx0::${variable} + channel: ${channel} + markersize: 5 + color: 'black' + label: 'JEDI h(x) versus obs (all obs)' + - type: Scatter + x: + variable: experiment::ObsValue::${variable} + y: + variable: experiment::hofx0PassedQc::${variable} + channel: ${channel} + markersize: 5 + color: 'red' + label: 'JEDI h(x) versus obs (passed QC in JEDI)' + + # GSI h(x) vs Observations + - batch figure: + variables: *variables + channels: *channels + figure: + layout: [1,1] + title: 'Observations vs. GSI h(x) | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/correlation_scatter/${variable}${channel}/gsi_hofx0_vs_obs_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'Observation Value' + add_ylabel: 'GSI h(x)' + add_grid: + add_legend: + loc: 'upper left' + layers: + - type: Scatter + x: + variable: experiment::ObsValue::${variable} + y: + variable: experiment::GsiHofXBc::${variable} + channel: ${channel} + markersize: 5 + color: 'black' + label: 'GSI h(x) versus obs (all obs)' + - type: Scatter + x: + variable: experiment::ObsValue::${variable} + y: + variable: experiment::GsiHofXBcPassedQc::${variable} + channel: ${channel} + markersize: 5 + color: 'red' + label: 'GSI h(x) versus obs (passed QC in JEDI)' + + # JEDI h(x) vs GSI h(x) + - batch figure: + variables: *variables + channels: *channels + figure: + layout: [1,1] + title: 'JEDI h(x) vs. GSI h(x) | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/correlation_scatter/${variable}${channel}/gsi_hofx_vs_jedi0_hofx_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'GSI h(x)' + add_ylabel: 'JEDI h(x)' + add_grid: + add_legend: + loc: 'upper left' + layers: + - type: Scatter + x: + variable: experiment::GsiHofXBc::${variable} + y: + variable: experiment::hofx0::${variable} + channel: ${channel} + markersize: 5 + color: 'black' + label: 'JEDI h(x) versus GSI h(x)' + - type: Scatter + x: + variable: experiment::GsiHofXBcPassedQc::${variable} + y: + variable: experiment::hofx0PassedQc::${variable} + channel: ${channel} + markersize: 5 + color: 'red' + label: 'JEDI h(x) versus GSI h(x) (passed QC in JEDI)' + + # JEDI hofx0 vs GSI hofx0 + - batch figure: + variables: *variables + channels: *channels + figure: + layout: [1,1] + title: 'JEDI hofx0 vs. GSI hofx0 | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/correlation_scatter/${variable}${channel}/gsi_hofx0_vs_jedi_hofx0_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'GSI observation minus h(x)' + add_ylabel: 'JEDI observation minus h(x)' + add_grid: + add_legend: + loc: 'upper left' + layers: + - type: Scatter + x: + variable: experiment::ObsValueMinusGsiHofXBc::${variable} + y: + variable: experiment::ObsValueMinusHofx0::${variable} + channel: ${channel} + markersize: 5 + color: 'black' + label: 'GSI hofx0 vs JEDI hofx0 (all obs)' + - type: Scatter + x: + variable: experiment::ObsValueMinusGsiHofXBcPassedQc::${variable} + y: + variable: experiment::ObsValueMinushofx0PassedQc::${variable} + channel: ${channel} + markersize: 5 + color: 'red' + label: 'GSI hofx0 vs JEDI hofx0 (passed QC in JEDI)' + + # JEDI oma vs omb + - batch figure: + variables: *variables + channels: *channels + figure: + layout: [1,1] + title: 'OmA vs. OmB | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/correlation_scatter/${variable}${channel}/jedi_oma_vs_omb_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'OmA' + add_ylabel: 'OmB' + add_grid: + add_legend: + loc: 'upper left' + layers: + - type: Scatter + x: + variable: experiment::oman::${variable} + y: + variable: experiment::ombg::${variable} + channel: ${channel} + markersize: 5 + color: 'black' + label: 'OmA versus OmB (all residuals)' + - type: Scatter + x: + variable: experiment::omanPassedQc::${variable} + y: + variable: experiment::ombgPassedQc::${variable} + channel: ${channel} + markersize: 5 + color: 'red' + label: 'OmA versus OmB (passed QC)' + +# Map plots# --------- + + # Observations + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: vminvmaxcmap + channel: ${channel} + data variable: experiment::ObsValue::${variable} + figure: + figure size: [20,10] + layout: [1,1] + title: 'Observations | {{instrument_title}} | Obs Value' + output name: '{{cycle_dir}}/eva/{{instrument}}/map_plots/${variable}${channel}/observations_{{instrument}}_${variable}${channel}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: ObsValue + add_grid: + layers: + - type: MapScatter + longitude: + variable: experiment::MetaData::longitude + latitude: + variable: experiment::MetaData::latitude + data: + variable: experiment::ObsValue::${variable} + channel: ${channel} + markersize: 2 + label: ObsValue + colorbar: true + cmap: ${dynamic_cmap} + vmin: ${dynamic_vmin} + vmax: ${dynamic_vmax} + + # hofx0 jedi + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: vminvmaxcmap + channel: ${channel} + data variable: experiment::ObsValueMinusHofx0::${variable} + figure: + figure size: [20,10] + layout: [1,1] + title: 'JEDI hofx0 | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/map_plots/${variable}${channel}/hofx0_jedi_{{instrument}}_${variable}${channel}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: '${variable}' + add_grid: + layers: + - type: MapScatter + longitude: + variable: experiment::MetaData::longitude + latitude: + variable: experiment::MetaData::latitude + data: + variable: experiment::ObsValueMinusHofx0::${variable} + channel: ${channel} + markersize: 2 + label: '${variable}' + colorbar: true + cmap: ${dynamic_cmap} + vmin: ${dynamic_vmin} + vmax: ${dynamic_vmax} + + # hofx0 gsi + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: vminvmaxcmap + channel: ${channel} + data variable: experiment::ObsValueMinusGsiHofXBc::${variable} + figure: + figure size: [20,10] + layout: [1,1] + title: 'GSI hofx0 | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/map_plots/${variable}${channel}/hofx0_gsi_{{instrument}}_${variable}${channel}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: '${variable}' + add_grid: + layers: + - type: MapScatter + longitude: + variable: experiment::MetaData::longitude + latitude: + variable: experiment::MetaData::latitude + data: + variable: experiment::ObsValueMinusGsiHofXBc::${variable} + channel: ${channel} + markersize: 2 + label: '${variable}' + colorbar: true + cmap: ${dynamic_cmap} + vmin: ${dynamic_vmin} + vmax: ${dynamic_vmax} + + # hofx difference + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: vminvmaxcmap + channel: ${channel} + data variable: experiment::Hofx0MinusGsiHofXBc::${variable} + figure: + figure size: [20,10] + layout: [1,1] + title: 'Hofx0 Difference | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/map_plots/${variable}${channel}/hofx0_difference_{{instrument}}_${variable}${channel}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: '${variable}' + add_grid: + layers: + - type: MapScatter + longitude: + variable: experiment::MetaData::longitude + latitude: + variable: experiment::MetaData::latitude + data: + variable: experiment::Hofx0MinusGsiHofXBc::${variable} + channel: ${channel} + markersize: 2 + label: '${variable}' + colorbar: true + cmap: ${dynamic_cmap} + vmin: ${dynamic_vmin} + vmax: ${dynamic_vmax} + + + # RMS(oma)-RMS(omb) difference + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: vminvmaxcmap + channel: ${channel} + data variable: experiment::ResidualRMSdiff::${variable} + figure: + figure size: [20,10] + layout: [1,1] + title: 'RMS Residual Difference | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/map_plots/${variable}${channel}/rmsres_difference_{{instrument}}_${variable}${channel}.png' + plots: + - mapping: + projection: plcarr + domain: global + add_map_features: ['coastline'] + add_colorbar: + label: '${variable}' + add_grid: + layers: + - type: MapScatter + longitude: + variable: experiment::MetaData::longitude + latitude: + variable: experiment::MetaData::latitude + data: + variable: experiment::ResidualRMSdiff::${variable} + channel: ${channel} + markersize: 2 + label: '${variable}' + colorbar: true + cmap: ${dynamic_cmap} + vmin: ${dynamic_vmin} + vmax: ${dynamic_vmax} +# Histogram plots# --------------- + + # hofx0 vs hofx0 + - batch figure: + variables: *variables + channels: *channels + dynamic options: + - type: histogram_bins + channel: ${channel} + number of bins rule: sturges + data variable: experiment::ObsValueMinusHofx0::${variable} + figure: + layout: [1,1] + title: 'JEDI hofx0 vs. GSI hofx0 | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/histograms/${variable}${channel}/gsi_hofx0_vs_jedi_hofx0_{{instrument}}_${variable}${channel}.png' + plots: + - add_xlabel: 'Observation minus h(x)' + add_ylabel: 'Count' + add_legend: + loc: 'upper left' + layers: + - type: Histogram + data: + variable: experiment::ObsValueMinusGsiHofXBc::${variable} + channel: ${channel} + color: 'blue' + label: 'GSI hofx0 (all obs)' + bins: ${dynamic_bins} + alpha: 0.5 + - type: Histogram + data: + variable: experiment::ObsValueMinusHofx0::${variable} + channel: ${channel} + color: 'red' + label: 'JEDI hofx0 (all obs)' + bins: ${dynamic_bins} + alpha: 0.5 + + # JEDI omb vs oma + - batch figure: + variables: *variables + dynamic options: + - type: histogram_bins + data variable: experiment::omanPassedQc::${variable} + number of bins rule: 'rice' + figure: + layout: [1,1] + title: 'OmB vs. OmA | {{instrument_title}} | ${variable_title}' + output name: '{{cycle_dir}}/eva/{{instrument}}/histograms/${variable}/ombg_oman_{{instrument}}_${variable}.png' + plots: + - add_xlabel: 'Difference' + add_ylabel: 'Count' + set_xlim: [-3, 3] + add_legend: + loc: 'upper left' + statistics: + fields: + - field_name: experiment::ombgPassedQc::${variable} + xloc: 0.5 + yloc: -0.10 + kwargs: + color: 'black' + fontsize: 8 + fontfamily: monospace + - field_name: experiment::omanPassedQc::${variable} + xloc: 0.5 + yloc: -0.13 + kwargs: + color: 'red' + fontsize: 8 + fontfamily: monospace + statistics_variables: + - n + - min + - mean + - max + - std + layers: + - type: Histogram + data: + variable: experiment::ombgPassedQc::${variable} + color: 'red' + label: 'observations minus background ' + bins: ${dynamic_bins} + alpha: 0.5 + density: true + - type: Histogram + data: + variable: experiment::omanPassedQc::${variable} + color: 'blue' + label: 'observations minus analysis' + bins: ${dynamic_bins} + alpha: 0.5 + density: true + diff --git a/src/swell/tasks/eva_increment.py b/src/swell/tasks/eva_increment.py index 10a53764f..711e42473 100644 --- a/src/swell/tasks/eva_increment.py +++ b/src/swell/tasks/eva_increment.py @@ -68,7 +68,7 @@ def execute(self) -> None: incr_file = f'{self.experiment_id()}.increment-iter{iter_no}.{cycle_time_reformat}.nc4' if self.suite_name() == 'localensembleda': incr_file = f'geos.mean-inc.{local_bkg_time}.nc4' - elif self.suite_name() == 'eda': + elif self.suite_name().startswith('eda_') and 'atmos' in self.suite_name(): incr_file = f'eda.mean-inc.{local_bkg_time}.nc4' if window_type == '4D' and 'atmos' in self.suite_name(): incr_file = f'{self.experiment_id()}.increment-iter{iter_no}.{window_begin}.nc4' From abae041c76f169b7b60118b4267b99c410f8cb42 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 22 Jul 2026 09:28:38 -0600 Subject: [PATCH 52/68] fix eda error on suite_name in eva_increment --- src/swell/tasks/eva_increment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/swell/tasks/eva_increment.py b/src/swell/tasks/eva_increment.py index 10a53764f..711e42473 100644 --- a/src/swell/tasks/eva_increment.py +++ b/src/swell/tasks/eva_increment.py @@ -68,7 +68,7 @@ def execute(self) -> None: incr_file = f'{self.experiment_id()}.increment-iter{iter_no}.{cycle_time_reformat}.nc4' if self.suite_name() == 'localensembleda': incr_file = f'geos.mean-inc.{local_bkg_time}.nc4' - elif self.suite_name() == 'eda': + elif self.suite_name().startswith('eda_') and 'atmos' in self.suite_name(): incr_file = f'eda.mean-inc.{local_bkg_time}.nc4' if window_type == '4D' and 'atmos' in self.suite_name(): incr_file = f'{self.experiment_id()}.increment-iter{iter_no}.{window_begin}.nc4' From 8ea73bc75dd9611a6d148eb816570bf76ed0ebec Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 22 Jul 2026 12:45:49 -0400 Subject: [PATCH 53/68] . --- src/swell/suites/eda_controlpert_atmos/flow.cylc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/swell/suites/eda_controlpert_atmos/flow.cylc b/src/swell/suites/eda_controlpert_atmos/flow.cylc index da246ddf1..5ff90d761 100644 --- a/src/swell/suites/eda_controlpert_atmos/flow.cylc +++ b/src/swell/suites/eda_controlpert_atmos/flow.cylc @@ -98,7 +98,7 @@ RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} # CleanupEda directory - SP4-{{model_component}} => CleanEdaFiles-{{model_component}} + #SP4-{{model_component}} => CleanEdaFiles-{{model_component}} # EvaIncrement SP4-{{model_component}} => EvaIncrement-{{model_component}} From 93d9f4a53f83bed7d5488bdb545a202ddb14d953 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 22 Jul 2026 14:24:13 -0600 Subject: [PATCH 54/68] set slurm directives --- src/swell/deployment/create_experiment.py | 4 ++++ src/swell/utilities/slurm.py | 4 +++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 30ed44337..52ac2f24e 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -589,6 +589,10 @@ def prepare_cylc_suite_jinja2( render_dictionary['scheduling']['RunGeos']['execution_time_limit'] = 'PT30M' render_dictionary['scheduling']['RunJediLocalEnsembleDaExecutable'][ 'execution_time_limit'] = 'PT1H' + render_dictionary['scheduling']['RunJediEdaExecutable'][ + 'execution_time_limit'] = 'PT30M' + render_dictionary['scheduling']['RunJediEdaControlPertExecutable'][ + 'execution_time_limit'] = 'PT30M' # Set jinja templated string to use upon runtime # ---------------------------------------------- diff --git a/src/swell/utilities/slurm.py b/src/swell/utilities/slurm.py index 26622af72..eb2f697ae 100644 --- a/src/swell/utilities/slurm.py +++ b/src/swell/utilities/slurm.py @@ -47,7 +47,8 @@ def prepare_scheduling_dict( "RunJediVariationalExecutable": {"all": {"nodes": 3}}, "RunJediUfoTestsExecutable": {"all": {"ntasks-per-node": 1}}, "RunJediConvertStateSoca2ciceExecutable": {"all": {"nodes": 1}}, - "RunJediEdaExecutable": {"all": {"ntasks-per-node": 126}} + "RunJediEdaExecutable": {"all": {"ntasks-per-node": 126}}, + "RunJediEdaControlPertExecutable": {"all": {"ntasks-per-node": 126}} } # Global SLURM settings stored in $HOME/.swell/swell-slurm.yaml @@ -81,6 +82,7 @@ def prepare_scheduling_dict( 'GenerateBClimatology', 'RunGeos', 'RunJediEdaExecutable', + 'RunJediEdaControlPertExecutable', 'RunJediEnsembleMeanVariance', 'RunJediDiffstates', 'RunJediConvertStateSoca2ciceExecutable', From d82a798ef140ac6e6bcf9c6193e871196cbd9203 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 22 Jul 2026 21:59:53 -0600 Subject: [PATCH 55/68] link multiple ebkg files --- .../tasks/run_jedi_eda_control_pert_executable.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/src/swell/tasks/run_jedi_eda_control_pert_executable.py b/src/swell/tasks/run_jedi_eda_control_pert_executable.py index 4885e098b..dda902dc1 100644 --- a/src/swell/tasks/run_jedi_eda_control_pert_executable.py +++ b/src/swell/tasks/run_jedi_eda_control_pert_executable.py @@ -222,20 +222,17 @@ def execute(self) -> None: # copy bkg files to imem dir f1_list = glob(os.path.join(self.cycle_dir(), f'ebkg/mem{id:03d}/geos.mem*.nc4')) - f1 = f1_list[0] - print(f'f1 = {f1}') - if f1: + if not f1_list: + self.logger.error(f"ebkg dir is empty for member id: {id}") + for f1 in f1_list: f2 = os.path.basename(f1) f2 = f2.split('.')[2:] f2 = '.'.join(f2) f2 = os.path.join(xdir, f2) print(f'f2 = {f2}') - else: - f2 = None # Handle the case where no file is found - self.logger.error(f"bkg dir is empty for member id: {id}") - if os.path.lexists(f2): - os.remove(f2) - os.symlink(f1, f2) + if os.path.lexists(f2): + os.remove(f2) + os.symlink(f1, f2) # analysis_chunk / chunk00x / mem00y will have its own obs, B and R # ----------------------------------------------------------------- From d0718598a45bbb37c5128c18db2796beae99b9cb Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 22 Jul 2026 23:08:22 -0600 Subject: [PATCH 56/68] Handle eda control pert case in clean_eda.py --- src/swell/tasks/clean_eda.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/src/swell/tasks/clean_eda.py b/src/swell/tasks/clean_eda.py index 7b745ffda..6df022e6b 100644 --- a/src/swell/tasks/clean_eda.py +++ b/src/swell/tasks/clean_eda.py @@ -10,6 +10,7 @@ import os import glob import shutil +from pathlib import Path from swell.tasks.base.task_base import taskBase @@ -19,6 +20,8 @@ class CleanEda(taskBase): # ---------------------------------------------------------------------------------------------- + # this file handles both eda and eda_controlpert cases + # def execute(self) -> None: @@ -119,10 +122,18 @@ def execute(self) -> None: window_type, jedi_forecast_model) - for imem in range(1, nmember+1): - mem_dir = f'analysis/mem{imem:003d}/' - d1 = os.path.join(self.cycle_dir(), mem_dir) - d2 = os.path.join(d1, 'fv3-jedi') + # + # handle eda case: analysis/mem00x + # eda controlpert case: analysis_chunk/chunk00x/mem00y + # + d1 = os.path.join(self.cycle_dir(), 'analysis_chunk') # control pert case + if not os.path.exists(d1): + d1 = os.path.join(self.cycle_dir(), 'analysis') # eda case + target_dirs = [str(p) for p in Path(d1).rglob('mem*') if p.is_dir()] + print (f'target_dirs = {target_dirs}') + + for mem_dir in target_dirs: + d2 = os.path.join(mem_dir, 'fv3-jedi') if os.path.exists(d2): if os.path.islink(d2): # Only remove the symlink itself, never follow it @@ -134,10 +145,10 @@ def execute(self) -> None: for observer in jedi_config_dict['cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] - # Delete input obsfile in analysis/mem00x (.nc4 .tlapse.txt acftbias acftbias_cov) - for file_path in glob.glob(os.path.join(d1, f'{observation}.*')): + # Delete input obsfile (.nc4 .tlapse.txt acftbias acftbias_cov) + for file_path in glob.glob(os.path.join(mem_dir, f'{observation}.*')): if os.path.islink(file_path): - os.unlink(file_path) # safe: removes only the link + os.unlink(file_path) # safe: removes only the link self.logger.info(f"Deleted symlink file: {file_path}") else: try: From a8b6673ffcd92f651513cc46e19bd0369e5502c9 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 22 Jul 2026 23:13:06 -0600 Subject: [PATCH 57/68] . --- .../suites/eda_controlpert_atmos/flow.cylc | 2 +- .../eda_controlpert_atmos/flow.cylc.complete | 237 ------------------ .../eda_controlpert_atmos/flow.cylc.simple | 65 ----- .../suites/eda_controlpert_atmos/flow.cylc.v2 | 210 ---------------- 4 files changed, 1 insertion(+), 513 deletions(-) delete mode 100644 src/swell/suites/eda_controlpert_atmos/flow.cylc.complete delete mode 100644 src/swell/suites/eda_controlpert_atmos/flow.cylc.simple delete mode 100644 src/swell/suites/eda_controlpert_atmos/flow.cylc.v2 diff --git a/src/swell/suites/eda_controlpert_atmos/flow.cylc b/src/swell/suites/eda_controlpert_atmos/flow.cylc index 5ff90d761..da246ddf1 100644 --- a/src/swell/suites/eda_controlpert_atmos/flow.cylc +++ b/src/swell/suites/eda_controlpert_atmos/flow.cylc @@ -98,7 +98,7 @@ RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} # CleanupEda directory - #SP4-{{model_component}} => CleanEdaFiles-{{model_component}} + SP4-{{model_component}} => CleanEdaFiles-{{model_component}} # EvaIncrement SP4-{{model_component}} => EvaIncrement-{{model_component}} diff --git a/src/swell/suites/eda_controlpert_atmos/flow.cylc.complete b/src/swell/suites/eda_controlpert_atmos/flow.cylc.complete deleted file mode 100644 index 70cd0433c..000000000 --- a/src/swell/suites/eda_controlpert_atmos/flow.cylc.complete +++ /dev/null @@ -1,237 +0,0 @@ -# (C) Copyright 2021- United States Government as represented by the Administrator of the -# National Aeronautics and Space Administration. All Rights Reserved. -# -# This software is licensed under the terms of the Apache Licence Version 2.0 -# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. - -# -------------------------------------------------------------------------------------------------- - -# Cylc suite for executing JEDI-based non-cycling variational data assimilation - -# -------------------------------------------------------------------------------------------------- - -[scheduler] - UTC mode = True - allow implicit tasks = False - -# -------------------------------------------------------------------------------------------------- - -[scheduling] - - initial cycle point = {{start_cycle_point}} - final cycle point = {{final_cycle_point}} - runahead limit = {{runahead_limit}} - - [[graph]] - R1 = """ - # Triggers for non cycle time dependent tasks - # ------------------------------------------- - # Clone JEDI source code - CloneJedi - - # Build JEDI source code by linking - CloneJedi => BuildJediByLinking? - - # If not able to link to build create the build - BuildJediByLinking:fail? => BuildJedi - - {% for model_component in model_components %} - # Clone geos ana for generating observing system records - CloneGeosMksi-{{model_component}} - {% endfor %} - """ - - {% for cycle_time in cycle_times %} - {{cycle_time.cycle_time}} = """ - {% for model_component in model_components %} - {% if cycle_time[model_component] %} - - # logic tree: - # prep (clone, stage, build, get bkg, get obs) -> SP1 - # -> mksi/observation + filter -> SP2 - # -> [ submit/wait for N chunk runs ]-> SP3 - # -> mean/variance + diffstate -> SP4 -> eva - - # Task triggers for: {{model_component}} - # ------------------ - # Generate satellite channel records - CloneGeosMksi-{{model_component}}[^] => GenerateObservingSystemRecords-{{model_component}} - - # ygyu fix at the end - # Get observations - # Cycling VarBC is inactive, static bias files will be used - GetObsNotInR2d2-{{model_component}}: fail? => GetObservations-{{model_component}} - - # Perform staging that is cycle dependent - StageJediCycle-{{model_component}} - - # Run Jedi variational executable - BuildJediByLinking[^]? | BuildJedi[^] => SP1-{{model_component}} - CloneJedi[^] => StageJediCycle-{{model_component}} - StageJediCycle-{{model_component}} => SP1-{{model_component}} - GetEnsembleGeosExperiment-{{model_component}} => SP1-{{model_component}} - GetObsNotInR2d2-{{model_component}}? | GetObservations-{{model_component}} => SP1-{{model_component}} - SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} - GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} - RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} - - {% if models[model_component]['ensemble_num_chunks'] is defined %} - {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} - SP2-{{model_component}} => RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}} => SP3-{{model_component}} - {% endfor %} - {% endif %} - - # bkg/ana : mean/variance - SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} - RunJediEnsembleMeanVariance-{{model_component}} ==> RunJediDiffstates-{{model_component}} => SP4-{{model_component}} - - # CleanupEda directory - SP4-{{model_component}} => CleanEdaFiles-{{model_component}} - - # EvaIncrement - SP4-{{model_component}} => EvaIncrement-{{model_component}} - - {% endif %} - {% endfor %} - """ - {% endfor %} - -# -------------------------------------------------------------------------------------------------- - -[runtime] - - # Task defaults - # ------------- - [[root]] - pre-script = "source $CYLC_SUITE_DEF_PATH/modules" - - [[[environment]]] - datetime = $CYLC_TASK_CYCLE_POINT - config = $CYLC_SUITE_DEF_PATH/experiment.yaml - - # Tasks - # ----- - [[CloneJedi]] - script = "swell task CloneJedi $config" - - [[BuildJediByLinking]] - script = "swell task BuildJediByLinking $config" - - [[BuildJedi]] - script = "swell task BuildJedi $config" - platform = {{platform}} - execution time limit = {{scheduling["BuildJedi"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["BuildJedi"]["directives"]["all"].items() %} - --{{key}} = {{value}} - {%- endfor %} - - {% for model_component in model_components %} - - [[CloneGeosMksi-{{model_component}}]] - script = "swell task CloneGeosMksi $config -m {{model_component}}" - - [[GenerateObservingSystemRecords-{{model_component}}]] - script = "swell task GenerateObservingSystemRecords $config -d $datetime -m {{model_component}}" - - [[GetObsNotInR2d2-{{model_component}}]] - script = "swell task GetObsNotInR2d2 $config -d $datetime -m {{model_component}}" - - [[StageJediCycle-{{model_component}}]] - script = "swell task StageJedi $config -d $datetime -m {{model_component}}" - - [[GetBackground-{{model_component}} ]] - script = "swell task GetBackground $config -d $datetime -m {{model_component}}" - - [[GetBackgroundGeosExperiment-{{model_component}} ]] - script = "swell task GetBackgroundGeosExperiment $config -d $datetime -m {{model_component}}" - - [[GetEnsembleGeosExperiment-{{model_component}}]] - script = "swell task GetEnsembleGeosExperiment $config -d $datetime -m {{model_component}}" - - [[GetObservations-{{model_component}}]] - script = "swell task GetObservations $config -d $datetime -m {{model_component}}" - - [[RenderJediObservations-{{model_component}}]] - script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" - - {% if models[model_component]['ensemble_num_chunks'] is defined %} - {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} - [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] - script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}} -ichunk {{i}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediEdaControlPertExecutable"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediEdaControlPertExecutable"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - {% endfor %} - {% endif %} - - [[RunJediObsfiltersExecutable-{{model_component}}]] - script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediObsfiltersExecutable"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediObsfiltersExecutable"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - - [[RunJediEnsembleMeanVariance-{{model_component}}]] - script = "swell task RunJediEnsembleMeanVariance $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediEnsembleMeanVariance"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediEnsembleMeanVariance"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - - [[RunJediDiffstates-{{model_component}}]] - script = "swell task RunJediDiffstates $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediDiffstates"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediDiffstates"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - - [[CleanEdaFiles-{{model_component}}]] - script = "swell task CleanEda $config -d $datetime -m {{model_component}}" - - [[EvaJediLog-{{model_component}}]] - script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}" - - [[EvaIncrement-{{model_component}}]] - script = "swell task EvaIncrement $config -d $datetime -m {{model_component}}" - - [[EvaObservations-{{model_component}}]] - script = "swell task EvaObservations $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["EvaObservations"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["EvaObservations"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - - [[SaveObsDiags-{{model_component}}]] - script = "swell task SaveObsDiags $config -d $datetime -m {{model_component}}" - - [[CleanCycle-{{model_component}}]] - script = "swell task CleanCycle $config -d $datetime -m {{model_component}}" - - # SP: Sync Point - [[SP1-{{model_component}}]] - script = true - - [[SP2-{{model_component}}]] - script = true - - [[SP3-{{model_component}}]] - script = true - - [[SP4-{{model_component}}]] - script = true - - {% endfor %} - -# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/eda_controlpert_atmos/flow.cylc.simple b/src/swell/suites/eda_controlpert_atmos/flow.cylc.simple deleted file mode 100644 index 6f657c020..000000000 --- a/src/swell/suites/eda_controlpert_atmos/flow.cylc.simple +++ /dev/null @@ -1,65 +0,0 @@ -[scheduler] - UTC mode = True - allow implicit tasks = False - -# -------------------------------------------------------------------------------------------------- - -[scheduling] - - initial cycle point = {{start_cycle_point}} - final cycle point = {{final_cycle_point}} - runahead limit = {{runahead_limit}} - - [[graph]] - R1 = """ - # Triggers for non cycle time dependent tasks - # ------------------------------------------- - # Clone JEDI source code - {% for model_component in model_components %} - - {% if models[model_component]['ensemble_num_chunks'] is defined %} - RunJediEdaControlPertExecutable-{{model_component}} - {% endif %} - - {% endfor %} - """ - -# -------------------------------------------------------------------------------------------------- - -[runtime] - - # Task defaults - # ------------- - [[root]] - pre-script = "source $CYLC_SUITE_DEF_PATH/modules" - - [[[environment]]] - datetime = $CYLC_TASK_CYCLE_POINT - config = $CYLC_SUITE_DEF_PATH/experiment.yaml - - # Tasks - # ----- - - [[CloneJedi]] - script = "swell task CloneJedi $config" - - {% for model_component in model_components %} - - {% if models[model_component]['ensemble_num_chunks'] is defined %} - [[RunJediEdaControlPertExecutable-{{model_component}}]] - script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}}" - platform = {{platform}} - - [[GetBackground-{{model_component}} ]] - script = "swell task GetBackground $config -d $datetime -m {{model_component}}" - - {% endif %} - - - [[SP1-{{model_component}}]] - script = true - - - {% endfor %} - -# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/eda_controlpert_atmos/flow.cylc.v2 b/src/swell/suites/eda_controlpert_atmos/flow.cylc.v2 deleted file mode 100644 index b7e1b0fb0..000000000 --- a/src/swell/suites/eda_controlpert_atmos/flow.cylc.v2 +++ /dev/null @@ -1,210 +0,0 @@ -# (C) Copyright 2021- United States Government as represented by the Administrator of the -# National Aeronautics and Space Administration. All Rights Reserved. -# -# This software is licensed under the terms of the Apache Licence Version 2.0 -# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. - -# -------------------------------------------------------------------------------------------------- - -# Cylc suite for executing JEDI-based non-cycling variational data assimilation - -# -------------------------------------------------------------------------------------------------- - -[scheduler] - UTC mode = True - allow implicit tasks = False - -# -------------------------------------------------------------------------------------------------- - -[scheduling] - - initial cycle point = {{start_cycle_point}} - final cycle point = {{final_cycle_point}} - runahead limit = {{runahead_limit}} - - [[graph]] - R1 = """ - # Triggers for non cycle time dependent tasks - # ------------------------------------------- - # Clone JEDI source code - CloneJedi - - # Build JEDI source code by linking - CloneJedi => BuildJediByLinking? - - # If not able to link to build create the build - BuildJediByLinking:fail? => BuildJedi - - {% for model_component in model_components %} - # Clone geos ana for generating observing system records - CloneGeosMksi-{{model_component}} - {% endfor %} - """ - - {% for cycle_time in cycle_times %} - {{cycle_time.cycle_time}} = """ - {% for model_component in model_components %} - {% if cycle_time[model_component] %} - - # logic tree: - # prep (clone, stage, build, get bkg, get obs) -> SP1 - # -> mksi/observation + filter -> SP2 - # -> [ submit/wait for N chunk runs ]-> SP3 - # -> mean/variance + diffstate -> SP4 -> eva - - # Task triggers for: {{model_component}} - # ------------------ - # Generate satellite channel records - CloneGeosMksi-{{model_component}}[^] => GenerateObservingSystemRecords-{{model_component}} - - # ygyu fix at the end - # Get observations - # Cycling VarBC is inactive, static bias files will be used - GetObsNotInR2d2-{{model_component}}: fail? => GetObservations-{{model_component}} - - # Perform staging that is cycle dependent - StageJediCycle-{{model_component}} - - # Run Jedi variational executable - BuildJediByLinking[^]? | BuildJedi[^] => SP1-{{model_component}} - CloneJedi[^] => StageJediCycle-{{model_component}} - StageJediCycle-{{model_component}} => SP1-{{model_component}} - GetEnsembleGeosExperiment-{{model_component}} => SP1-{{model_component}} - GetObsNotInR2d2-{{model_component}}? | GetObservations-{{model_component}} => SP1-{{model_component}} - SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} - GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} - RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} - - {% if models[model_component]['ensemble_num_chunks'] is defined %} - {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} - SP2-{{model_component}} => RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}} => SP3-{{model_component}} - {% endfor %} - {% endif %} - - # EvaIncrement - SP3-{{model_component}} => EvaIncrement-{{model_component}} - - {% endif %} - {% endfor %} - """ - {% endfor %} - -# -------------------------------------------------------------------------------------------------- - -[runtime] - - # Task defaults - # ------------- - [[root]] - pre-script = "source $CYLC_SUITE_DEF_PATH/modules" - - [[[environment]]] - datetime = $CYLC_TASK_CYCLE_POINT - config = $CYLC_SUITE_DEF_PATH/experiment.yaml - - # Tasks - # ----- - [[CloneJedi]] - script = "swell task CloneJedi $config" - - [[BuildJediByLinking]] - script = "swell task BuildJediByLinking $config" - - [[BuildJedi]] - script = "swell task BuildJedi $config" - platform = {{platform}} - execution time limit = {{scheduling["BuildJedi"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["BuildJedi"]["directives"]["all"].items() %} - --{{key}} = {{value}} - {%- endfor %} - - {% for model_component in model_components %} - - [[CloneGeosMksi-{{model_component}}]] - script = "swell task CloneGeosMksi $config -m {{model_component}}" - - [[GenerateObservingSystemRecords-{{model_component}}]] - script = "swell task GenerateObservingSystemRecords $config -d $datetime -m {{model_component}}" - - [[GetObsNotInR2d2-{{model_component}}]] - script = "swell task GetObsNotInR2d2 $config -d $datetime -m {{model_component}}" - - [[StageJediCycle-{{model_component}}]] - script = "swell task StageJedi $config -d $datetime -m {{model_component}}" - - [[GetBackground-{{model_component}} ]] - script = "swell task GetBackground $config -d $datetime -m {{model_component}}" - - [[GetBackgroundGeosExperiment-{{model_component}} ]] - script = "swell task GetBackgroundGeosExperiment $config -d $datetime -m {{model_component}}" - - [[GetEnsembleGeosExperiment-{{model_component}}]] - script = "swell task GetEnsembleGeosExperiment $config -d $datetime -m {{model_component}}" - - [[GetObservations-{{model_component}}]] - script = "swell task GetObservations $config -d $datetime -m {{model_component}}" - - [[RenderJediObservations-{{model_component}}]] - script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" - - {% if models[model_component]['ensemble_num_chunks'] is defined %} - {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} - [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] - script = "swell task RunJediEdaControlPertExecutable $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediEdaControlPertExecutable"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediEdaControlPertExecutable"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - {% endfor %} - {% endif %} - - [[RunJediObsfiltersExecutable-{{model_component}}]] - script = "swell task RunJediObsfiltersExecutable $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["RunJediObsfiltersExecutable"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["RunJediObsfiltersExecutable"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - - - [[EvaJediLog-{{model_component}}]] - script = "swell task EvaJediLog $config -d $datetime -m {{model_component}}" - - [[EvaIncrement-{{model_component}}]] - script = "swell task EvaIncrement $config -d $datetime -m {{model_component}}" - - [[EvaObservations-{{model_component}}]] - script = "swell task EvaObservations $config -d $datetime -m {{model_component}}" - platform = {{platform}} - execution time limit = {{scheduling["EvaObservations"]["execution_time_limit"]}} - [[[directives]]] - {%- for key, value in scheduling["EvaObservations"]["directives"][model_component].items() %} - --{{key}} = {{value}} - {%- endfor %} - - [[SaveObsDiags-{{model_component}}]] - script = "swell task SaveObsDiags $config -d $datetime -m {{model_component}}" - - [[CleanCycle-{{model_component}}]] - script = "swell task CleanCycle $config -d $datetime -m {{model_component}}" - - # SP: Sync Point - [[SP1-{{model_component}}]] - script = true - - [[SP2-{{model_component}}]] - script = true - - [[SP3-{{model_component}}]] - script = true - - [[SP4-{{model_component}}]] - script = true - - {% endfor %} - -# -------------------------------------------------------------------------------------------------- From c0c927bb9602928597f90d4c46a90a0c660eb88d Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 23 Jul 2026 11:10:15 -0600 Subject: [PATCH 58/68] add EDA_start, EDA_end as sync point to workflow --- src/swell/suites/eda_atmos/flow.cylc | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/src/swell/suites/eda_atmos/flow.cylc b/src/swell/suites/eda_atmos/flow.cylc index 4f866af21..0745ef14c 100644 --- a/src/swell/suites/eda_atmos/flow.cylc +++ b/src/swell/suites/eda_atmos/flow.cylc @@ -48,9 +48,9 @@ # logic tree: # prep (clone, stage, build, get bkg, get obs) -> SP1 - # -> mksi/observation + filter -> SP2 - # -> [ submit/wait for N member runs ]-> SP3 - # -> mean/variance + diffstate -> SP4 -> eva + # -> mksi/observation + filter -> EDA_start + # -> [ submit/wait for N member runs ]-> EDA_end + # -> mean/variance + diffstate -> SP2 -> eva # Task triggers for: {{model_component}} # ------------------ @@ -79,23 +79,23 @@ GetObsNotInR2d2-{{model_component}}? | GetObservations-{{model_component}} => SP1-{{model_component}} SP1-{{model_component}} => GenerateObservingSystemRecords-{{model_component}} GenerateObservingSystemRecords-{{model_component}} => RenderJediObservations-{{model_component}} - RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => SP2-{{model_component}} + RenderJediObservations-{{model_component}} => RunJediObsfiltersExecutable-{{model_component}} => EDA_start-{{model_component}} {% if models[model_component]['ensemble_num_members'] is defined %} {% for i in range( 1, models[model_component]['ensemble_num_members'] + 1 ) %} - SP2-{{model_component}} => RunJediEdaExecutable_mem{{i}}-{{model_component}} => SP3-{{model_component}} + EDA_start-{{model_component}} => RunJediEdaExecutable_mem{{i}}-{{model_component}} => EDA_end-{{model_component}} {% endfor %} {% endif %} # bkg/ana : mean/variance - SP3-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} - RunJediEnsembleMeanVariance-{{model_component}} => RunJediDiffstates-{{model_component}} => SP4-{{model_component}} + EDA_end-{{model_component}} => RunJediEnsembleMeanVariance-{{model_component}} + RunJediEnsembleMeanVariance-{{model_component}} => RunJediDiffstates-{{model_component}} => SP2-{{model_component}} # CleanupEda directory - SP4-{{model_component}} => CleanEdaFiles-{{model_component}} + SP2-{{model_component}} => CleanEdaFiles-{{model_component}} # EvaIncrement - SP4-{{model_component}} => EvaIncrement-{{model_component}} + SP2-{{model_component}} => EvaIncrement-{{model_component}} {% endif %} {% endfor %} @@ -240,13 +240,13 @@ [[SP1-{{model_component}}]] script = true - [[SP2-{{model_component}}]] + [[EDA_start-{{model_component}}]] script = true - [[SP3-{{model_component}}]] + [[EDA_end-{{model_component}}]] script = true - [[SP4-{{model_component}}]] + [[SP2-{{model_component}}]] script = true {% endfor %} From 2154cb97f60b9c98d719c2ab4d5b8230b4487102 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 23 Jul 2026 11:30:37 -0600 Subject: [PATCH 59/68] change back deployment to develop --- src/swell/deployment/create_experiment.py | 17 +--- .../prepare_config_and_suite.py | 87 ++----------------- 2 files changed, 7 insertions(+), 97 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 52ac2f24e..46596ce78 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -26,7 +26,7 @@ from swell.utilities.logger import Logger, get_logger from swell.utilities.slurm import prepare_scheduling_dict from swell.utilities.check_da_params import check_da_params -from swell.utilities.yaml_utils import print_dict + # -------------------------------------------------------------------------------------------------- @@ -256,8 +256,6 @@ def create_experiment_directory( # Experiment ID and root from the user input # ------------------------------------------ - - #ygyu test experiment_id = dict_get(logger, experiment_dict, 'experiment_id') experiment_root = dict_get(logger, experiment_dict, 'experiment_root') @@ -485,10 +483,6 @@ def prepare_cylc_suite_jinja2( # -------------------------- with open(os.path.join(swell_suite_path, 'flow.cylc'), 'r') as file: suite_file = file.read() - print(f'nail6. suite_file in prepare_cylc_suite_jinja2 = {suite_file}') - print(f'experiment_dict =') - print_dict(experiment_dict) - # Copy the experiment dictionary to the rendering dictionary # ---------------------------------------------------------- @@ -543,11 +537,8 @@ def prepare_cylc_suite_jinja2( 'dictionary.') # Check if 'ensemble_hofx_strategy' appears anywhere in suite_file - ensemble_list = ['ensemble_'+s for s in - ['num_members', 'num_chunks', 'hofx_strategy', 'hofx_packets']] + ensemble_list = ['ensemble_'+s for s in ['num_members', 'hofx_strategy', 'hofx_packets']] ensemble_list = ensemble_list + ['skip_ensemble_hofx'] - print (f'ensemble_list = {ensemble_list}') - for ensemble_aspect in ensemble_list: if ensemble_aspect in suite_file: if len(model_components) == 0: @@ -589,10 +580,6 @@ def prepare_cylc_suite_jinja2( render_dictionary['scheduling']['RunGeos']['execution_time_limit'] = 'PT30M' render_dictionary['scheduling']['RunJediLocalEnsembleDaExecutable'][ 'execution_time_limit'] = 'PT1H' - render_dictionary['scheduling']['RunJediEdaExecutable'][ - 'execution_time_limit'] = 'PT30M' - render_dictionary['scheduling']['RunJediEdaControlPertExecutable'][ - 'execution_time_limit'] = 'PT30M' # Set jinja templated string to use upon runtime # ---------------------------------------------- diff --git a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py index 5859c59c7..69987bf5f 100644 --- a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py +++ b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py @@ -23,7 +23,6 @@ from swell.utilities.dictionary import update_dict from swell.tasks.task_questions import TaskQuestions as task_questions from swell.suites.all_suites import AllSuites -from swell.utilities.yaml_utils import print_dict # -------------------------------------------------------------------------------------------------- @@ -93,11 +92,6 @@ def __init__( self.model_ind_tasks = self.get_suite_task_list_model_ind(self.suite_str) self.all_model_dep_tasks = self.get_all_model_dep_tasks(self.suite_str) - print(f'nail1 model_ind_tasks = {self.model_ind_tasks}') - print(f'nail2 all_model_dep_tasks = {self.all_model_dep_tasks}') - print(f'nail3 self.suite_str = {self.suite_str}') - - # Perform the assembly of the dictionaries that contain all the questions that can possibly # be asked. This @@ -127,9 +121,6 @@ def prepare_question_dictionaries(self) -> None: suite_config_obj = AllSuites.get_config(self.suite_config) suite_question_list = suite_config_obj.expand_question_list() - print(f'suite q list : {suite_question_list}') - print_dict(suite_question_list) - # Allow for adding extra tasks manually from configuration # For dynamic suite creation (e.g. comparison tests) @@ -211,8 +202,6 @@ def prepare_question_dictionaries(self) -> None: if 'model_components' not in question_dictionary.keys(): self.suite_needs_model_components = False - print(f'self.suite_needs_model_components = {self.suite_needs_model_components}') - # Create copy of the question_dictionary for model independent questions question_dictionary_model_ind = copy.deepcopy(question_dictionary) @@ -278,11 +267,6 @@ def prepare_question_dictionaries(self) -> None: for key in keys_to_remove: del self.question_dictionary_model_dep[model][key] -# print(f'nail1.5 question_dictionary_model_dep = {self.question_dictionary_model_dep}') -# for model in self.possible_model_components: -# x = self.question_dictionary_model_dep[model]['ensemble_num_chunks'] -# print(f'x = {x}') - # ---------------------------------------------------------------------------------------------- def override_with_defaults(self) -> None: @@ -442,13 +426,6 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # 2. Perform a non-exhaustive resolving of suite file templates # ------------------------------------------------------------- suite_str = template_string_jinja2(self.logger, self.suite_str, self.experiment_dict, True) - - print(f'nail4.0 self.exp dict = {self.experiment_dict}') - print_dict(self.experiment_dict) - - print(f'nail4.1 self.suite str = {self.suite_str}') - - print(f'nail4.2 jinja2 suite str = {suite_str}') # 3. Get a list of tasks that do not depend on the model component # ---------------------------------------------------------------- @@ -480,16 +457,10 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # At this point the user should have provided the model components answer. Check that it is # in the experiment dictionary and retrieve the response - render_dict = copy.deepcopy(self.experiment_dict) - print('r1: render_dict') - print_dict(render_dict) - - if 'model_components' not in self.experiment_dict: self.logger.abort('The model components question has not been answered.') for model in self.experiment_dict['model_components']: - print(f'nail 5.1 model is {model}') model_dict = self.question_dictionary_model_dep[model] @@ -502,43 +473,6 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # Ask the question self.ask_a_question(model_dict, question_key, model) - # pass special model-dep task questions/variables to render_dict - # for suite_str (jinja2) in step 7. - if model_dict[question_key]['question_type'] == 'task': - - ensemble_list = ['ensemble_'+s for s in ['num_members', 'num_chunks']] - if question_key in ensemble_list: - - # a. Ensure the 'models' list exists - if 'models' not in render_dict: - render_dict['models'] = [] - - # b. Check if this model is already in our list - model_found = False - for m in render_dict['models']: - if model in m: - # c. If found, incrementally add the new key! - default_val = model_dict[question_key]['default_value'] - m[model][question_key] = default_val - model_found = True - break - - # d. If not found, add the model to the list for the first time - if not model_found: - default_val = model_dict[question_key]['default_value'] - render_dict['models'].append({ - model: { - question_key: default_val - } - }) - - # debug -# print(f'render_dict = {render_dict}') - print('r2: render_dict') - print_dict(render_dict) - - - # 7. Perform a more exhaustive resolving of suite file templates # -------------------------------------------------------------- # Note that we reset the suite file to avoid templates having been left unresolved @@ -546,14 +480,9 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # of templates because there are things related to scheduling that are not yet able to be # resolved. In the future it might be good to bring some of that information into the # sphere of suite questions but that requires some careful thought so as not to overload - # the user with questions. A few model dep variables are needed to corectly parse tasks, - # e.g., in Eda ControlPert. Hence we use a render_dict for jinja2 here, only to parse - # the for loop around model task in .cycl file, e.g., - # {% for i in range( 1, models[model_component]['ensemble_num_chunks'] + 1 ) %} - # [[RunJediEdaControlPertExecutable_chunk{{i}}-{{model_component}}]] - # {% endfor %} - - suite_str = template_string_jinja2(self.logger, self.suite_str, render_dict, True) + # the user with questions. + suite_str = template_string_jinja2(self.logger, self.suite_str, self.experiment_dict, + True) # 8. Ask the new task questions that do not actually depend on the model # ----------------------------------------------------------------------- @@ -570,10 +499,8 @@ def ask_questions_and_configure_suite(self) -> Tuple[dict, dict]: # 9.1 Build a list of tasks for each model component # ------------------------------------------------- - print(f'ck suite_str = {suite_str}') model_dep_tasks = self.get_suite_task_list_model_dep(suite_str) - print(f'ck show: model_dep_tasks = {model_dep_tasks}') - + # 9.2 Iterate over the model_dep dictionary and ask task questions # ---------------------------------------------------------------- for model in self.experiment_dict['model_components']: @@ -700,13 +627,11 @@ def get_all_model_dep_tasks(self, suite_str: str) -> list: swell_task_lines = [line.replace('"', '') for line in swell_task_lines] swell_task_lines = [line.strip() for line in swell_task_lines] - print(f'def all model dep tasks: swell_task_lines = {swell_task_lines}') - # All tasks all_tasks = [] for line in swell_task_lines: - all_tasks.append(line.split('swell task')[1].strip().split(' ')[0]) + all_tasks.append(line.split('swell task ')[1].split(' ')[0]) # Ensure all_tasks are unique all_tasks = list(set(all_tasks)) @@ -725,8 +650,6 @@ def get_suite_task_list_model_dep(self, suite_str: str) -> dict: swell_task_lines = [line.replace('"', '') for line in swell_task_lines] swell_task_lines = [line.strip() for line in swell_task_lines] - print(f'def get_suite_task_list_model_dep swell_task_lines = {swell_task_lines}') - # Now get the model part models = [] for line in swell_task_lines: From 216b3c468f60429c382daad8045a9b2578f8915c Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 23 Jul 2026 11:48:27 -0600 Subject: [PATCH 60/68] add comment to clean_eda --- src/swell/tasks/clean_eda.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/swell/tasks/clean_eda.py b/src/swell/tasks/clean_eda.py index 6df022e6b..c630c69bd 100644 --- a/src/swell/tasks/clean_eda.py +++ b/src/swell/tasks/clean_eda.py @@ -122,16 +122,18 @@ def execute(self) -> None: window_type, jedi_forecast_model) - # + # This special design works with either eda or eda_control_pert # handle eda case: analysis/mem00x # eda controlpert case: analysis_chunk/chunk00x/mem00y - # + # This avoids blindly search for mem0xx dir, as some of them + # only contains linked analysis files and donot require clean up + #--------------------------------------------------------------- d1 = os.path.join(self.cycle_dir(), 'analysis_chunk') # control pert case if not os.path.exists(d1): d1 = os.path.join(self.cycle_dir(), 'analysis') # eda case target_dirs = [str(p) for p in Path(d1).rglob('mem*') if p.is_dir()] - print (f'target_dirs = {target_dirs}') - + self.logger.info(f'target_dirs = {target_dirs}') + for mem_dir in target_dirs: d2 = os.path.join(mem_dir, 'fv3-jedi') if os.path.exists(d2): From ed675e00c6d33a028bfcd7c36775f3144f6744f8 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Thu, 23 Jul 2026 13:26:10 -0600 Subject: [PATCH 61/68] increase node to 4 for control pert --- src/swell/utilities/slurm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/swell/utilities/slurm.py b/src/swell/utilities/slurm.py index eb2f697ae..809e8a6da 100644 --- a/src/swell/utilities/slurm.py +++ b/src/swell/utilities/slurm.py @@ -47,8 +47,8 @@ def prepare_scheduling_dict( "RunJediVariationalExecutable": {"all": {"nodes": 3}}, "RunJediUfoTestsExecutable": {"all": {"ntasks-per-node": 1}}, "RunJediConvertStateSoca2ciceExecutable": {"all": {"nodes": 1}}, - "RunJediEdaExecutable": {"all": {"ntasks-per-node": 126}}, - "RunJediEdaControlPertExecutable": {"all": {"ntasks-per-node": 126}} + "RunJediEdaExecutable": {"all": {"ntasks-per-node": 120}}, + "RunJediEdaControlPertExecutable": {"all": {"nodes": 4, "ntasks-per-node": 120}} } # Global SLURM settings stored in $HOME/.swell/swell-slurm.yaml From ba0b9eb26efab1d0c6611237c67071146ed06fcb Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 29 Jul 2026 23:30:28 -0400 Subject: [PATCH 62/68] set controlpert to nchunk=16 and with sonde only nnode=3, nx/ny=4/5; perhost=120 --- src/swell/deployment/create_experiment.py | 2 + .../eda_controlpert_atmos/suite_config.py | 71 ++++++++++--------- src/swell/utilities/slurm.py | 4 +- 3 files changed, 40 insertions(+), 37 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 46596ce78..e9b385fbf 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -580,6 +580,8 @@ def prepare_cylc_suite_jinja2( render_dictionary['scheduling']['RunGeos']['execution_time_limit'] = 'PT30M' render_dictionary['scheduling']['RunJediLocalEnsembleDaExecutable'][ 'execution_time_limit'] = 'PT1H' + render_dictionary['scheduling']['RunJediEdaControlPertExecutable'][ + 'execution_time_limit'] = 'PT10M' # Set jinja templated string to use upon runtime # ---------------------------------------------- diff --git a/src/swell/suites/eda_controlpert_atmos/suite_config.py b/src/swell/suites/eda_controlpert_atmos/suite_config.py index fbbda16dd..ea25f37d5 100644 --- a/src/swell/suites/eda_controlpert_atmos/suite_config.py +++ b/src/swell/suites/eda_controlpert_atmos/suite_config.py @@ -41,6 +41,7 @@ class SuiteConfig(QuestionContainer, Enum): 'rtodling/archive/541/Milan'), qd.npx_proc(4), qd.npy_proc(5), + qd.perhost(120), qd.window_length("PT6H"), qd.window_type("3D"), qd.horizontal_resolution("91"), @@ -48,8 +49,8 @@ class SuiteConfig(QuestionContainer, Enum): qd.gsibec_nlons("144"), qd.vertical_resolution("72"), qd.ensemble_num_members(32), - qd.ensemble_num_chunks(8), - qd.number_of_iterations([100]), + qd.ensemble_num_chunks(16), + qd.number_of_iterations([50]), qd.gradient_norm_reduction(1.e-3), qd.minimizer("DRPLanczos"), qd.analysis_variables([ @@ -71,40 +72,40 @@ class SuiteConfig(QuestionContainer, Enum): "skin_temperature_at_surface" ]), qd.observations([ - "aircraft_temperature", - "aircraft_wind", - "airs_aqua", - "amsr2_gcom-w1", - "amsua_aqua", - "amsua_metop-b", - "amsua_metop-c", - "amsua_n15", - "amsua_n18", - "amsua_n19", - "atms_n20", - "atms_npp", - "avhrr3_metop-b", - "avhrr3_n18", - "avhrr3_n19", - "cris-fsr_n20", - "cris-fsr_npp", - "gmi_gpm", - "gps", - "iasi_metop-b", - "iasi_metop-c", - "mhs_metop-b", - "mhs_metop-c", - "mhs_n19", - "mls55_aura", - "omi_aura", - "ompsnm_npp", - "pibal", - "satwind", - "scatwind", - "sfcship", - "sfc", +# "aircraft_temperature", +# "aircraft_wind", +# "airs_aqua", +# "amsr2_gcom-w1", +# "amsua_aqua", +# "amsua_metop-b", +# "amsua_metop-c", +# "amsua_n15", +# "amsua_n18", +# "amsua_n19", +# "atms_n20", +# "atms_npp", +# "avhrr3_metop-b", +# "avhrr3_n18", +# "avhrr3_n19", +# "cris-fsr_n20", +# "cris-fsr_npp", +# "gmi_gpm", +# "gps", +# "iasi_metop-b", +# "iasi_metop-c", +# "mhs_metop-b", +# "mhs_metop-c", +# "mhs_n19", +# "mls55_aura", +# "omi_aura", +# "ompsnm_npp", +# "pibal", +# "satwind", +# "scatwind", +# "sfcship", +# "sfc", "sondes", - "ssmis_f17" +# "ssmis_f17" ]), qd.obs_thinning_rej_fraction(0.8), qd.ensmeanvariance_spec([ diff --git a/src/swell/utilities/slurm.py b/src/swell/utilities/slurm.py index 809e8a6da..750fa83b2 100644 --- a/src/swell/utilities/slurm.py +++ b/src/swell/utilities/slurm.py @@ -47,8 +47,8 @@ def prepare_scheduling_dict( "RunJediVariationalExecutable": {"all": {"nodes": 3}}, "RunJediUfoTestsExecutable": {"all": {"ntasks-per-node": 1}}, "RunJediConvertStateSoca2ciceExecutable": {"all": {"nodes": 1}}, - "RunJediEdaExecutable": {"all": {"ntasks-per-node": 120}}, - "RunJediEdaControlPertExecutable": {"all": {"nodes": 4, "ntasks-per-node": 120}} + "RunJediEdaExecutable": {"all": {"ntasks-per-node": 126}}, + "RunJediEdaControlPertExecutable": {"all": {"nodes": 3, "ntasks-per-node": 126}} } # Global SLURM settings stored in $HOME/.swell/swell-slurm.yaml From 6bddb75a04d07333ead1940f20d4ca6776b85668 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 4 Aug 2026 13:24:00 -0400 Subject: [PATCH 63/68] add eda_atmos_tier1_fast to qd --- src/swell/suites/eda_atmos/suite_config.py | 83 +++++++++++++++++++++- 1 file changed, 81 insertions(+), 2 deletions(-) diff --git a/src/swell/suites/eda_atmos/suite_config.py b/src/swell/suites/eda_atmos/suite_config.py index b5156e20f..a9d1f7ea7 100644 --- a/src/swell/suites/eda_atmos/suite_config.py +++ b/src/swell/suites/eda_atmos/suite_config.py @@ -21,6 +21,85 @@ class SuiteConfig(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + eda_atmos_tier1_fast = QuestionList( + list_name="eda_atmos_tier1_fast", + questions=[ + sq.common, + qd.start_cycle_point("2023-10-10T00:00:00Z"), + qd.final_cycle_point("2023-10-10T06:00:00Z"), + qd.runahead_limit("P2"), + qd.jedi_build_method("use_existing"), + qd.model_components(['geos_atmosphere']), + ], + geos_atmosphere=[ + qd.cycle_times([ + "T00", + ]), + qd.geos_x_background_directory("/discover/nobackup/projects/gmao/" + "dadev/rtodling/archive/Restarts/JEDI/541x"), + qd.geos_x_ensemble_directory('/discover/nobackup/projects/gmao/dadev/' + 'rtodling/archive/541/Milan'), + qd.npx_proc(4), + qd.npy_proc(5), + qd.window_length("PT6H"), + qd.window_type("3D"), + qd.horizontal_resolution("91"), + qd.gsibec_nlats("91"), + qd.gsibec_nlons("144"), + qd.vertical_resolution("72"), + qd.ensemble_num_members(3), + qd.obs_pert_amplitude(0.5), + qd.number_of_iterations([5]), + qd.gradient_norm_reduction(1.e-3), + qd.analysis_variables([ + "eastward_wind", + "northward_wind", + "air_temperature", + "water_vapor_mixing_ratio_wrt_moist_air", + "air_pressure_at_surface", + "air_pressure_levels", + "cloud_liquid_ice", + "cloud_liquid_water", + "rain_water", + "snow_water", + "mole_fraction_of_ozone_in_air", + "geopotential_height_times_gravity_at_surface", + "fraction_of_ocean", + "fraction_of_lake", + "fraction_of_ice", + "skin_temperature_at_surface" + ]), + qd.observations([ + "aircraft_temperature", + "aircraft_wind", + ]), + qd.obs_thinning_rej_fraction(0.8), + qd.ensmeanvariance_spec([ + {"state": "bkg", + "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "geos.prior.mean", + "fn_output_variance": "geos.prior.variance", + "grid_type": ['cs', 'latlon']}, + {"state": "analysis", + "fn_input": "analysis/mem%mem%/eda.ana.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "eda.ana.mean", + "fn_output_variance": "eda.ana.variance", + "grid_type": ['cs', 'latlon']}, + ]), + qd.diffstates_spec({ + "state1": + {"fn_input": "geos.prior.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + "state2": + {"fn_input": "eda.ana.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + "state_diff": + {"fn_output": "eda.mean-inc", "grid_type": ['cs', 'latlon']}, + "state_type": "ensemble" + }), + qd.clean_patterns(['*.txt', '*.csv']), + ] + ) + + eda_atmos_tier1 = QuestionList( list_name="eda_atmos_tier1", questions=[ @@ -130,13 +209,13 @@ class SuiteConfig(QuestionContainer, Enum): qd.clean_patterns(['*.txt', '*.csv']), ] ) - + # -------------------------------------------------------------------------------------------------- eda_atmos = QuestionList( list_name="eda_atmos", questions=[ - eda_atmos_tier1 + eda_atmos_tier1_fast ] ) From 257a8043b928de6d1d3191acf6d2a2f7ac652809 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 4 Aug 2026 17:02:38 -0400 Subject: [PATCH 64/68] fix errors found by pycodestyle --- src/swell/suites/eda_atmos/suite_config.py | 23 +++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/swell/suites/eda_atmos/suite_config.py b/src/swell/suites/eda_atmos/suite_config.py index a9d1f7ea7..a632d8b86 100644 --- a/src/swell/suites/eda_atmos/suite_config.py +++ b/src/swell/suites/eda_atmos/suite_config.py @@ -35,10 +35,11 @@ class SuiteConfig(QuestionContainer, Enum): qd.cycle_times([ "T00", ]), - qd.geos_x_background_directory("/discover/nobackup/projects/gmao/" - "dadev/rtodling/archive/Restarts/JEDI/541x"), - qd.geos_x_ensemble_directory('/discover/nobackup/projects/gmao/dadev/' - 'rtodling/archive/541/Milan'), + qd.background_experiment('x0050'), + qd.geos_x_background_directory("/discover/nobackup/projects/gmao/dadev" + "rtodling/archive/Restarts/JEDI/541x"), + qd.geos_x_ensemble_directory("/discover/nobackup/projects/gmao/dadev/" + "rtodling/archive/541/Milan"), qd.npx_proc(4), qd.npy_proc(5), qd.window_length("PT6H"), @@ -99,7 +100,6 @@ class SuiteConfig(QuestionContainer, Enum): ] ) - eda_atmos_tier1 = QuestionList( list_name="eda_atmos_tier1", questions=[ @@ -114,10 +114,11 @@ class SuiteConfig(QuestionContainer, Enum): qd.cycle_times([ "T00", ]), - qd.geos_x_background_directory("/discover/nobackup/projects/gmao/" - "dadev/rtodling/archive/Restarts/JEDI/541x"), - qd.geos_x_ensemble_directory('/discover/nobackup/projects/gmao/dadev/' - 'rtodling/archive/541/Milan'), + qd.background_experiment('x0050'), + qd.geos_x_background_directory("/discover/nobackup/projects/gmao/dadev" + "rtodling/archive/Restarts/JEDI/541x"), + qd.geos_x_ensemble_directory("/discover/nobackup/projects/gmao/dadev/" + "rtodling/archive/541/Milan"), qd.npx_proc(4), qd.npy_proc(5), qd.window_length("PT6H"), @@ -209,7 +210,7 @@ class SuiteConfig(QuestionContainer, Enum): qd.clean_patterns(['*.txt', '*.csv']), ] ) - + # -------------------------------------------------------------------------------------------------- eda_atmos = QuestionList( @@ -218,5 +219,5 @@ class SuiteConfig(QuestionContainer, Enum): eda_atmos_tier1_fast ] ) - +# normal run: eda_atmos_tier1 # -------------------------------------------------------------------------------------------------- From 6566ccae9dd3b688168dddb6b0b2f7b7e7caacac Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 4 Aug 2026 21:58:27 -0400 Subject: [PATCH 65/68] proof concept: add eda_controlpert_atmos_tier1_fast --- .../eda_controlpert_atmos/suite_config.py | 89 ++++++++++++++++++- 1 file changed, 85 insertions(+), 4 deletions(-) diff --git a/src/swell/suites/eda_controlpert_atmos/suite_config.py b/src/swell/suites/eda_controlpert_atmos/suite_config.py index ea25f37d5..b478e8710 100644 --- a/src/swell/suites/eda_controlpert_atmos/suite_config.py +++ b/src/swell/suites/eda_controlpert_atmos/suite_config.py @@ -21,6 +21,86 @@ class SuiteConfig(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + eda_controlpert_atmos_tier1_fast = QuestionList( + list_name="eda_controlpert_atmos_tier1", + questions=[ + sq.common, + qd.start_cycle_point("2023-10-10T00:00:00Z"), + qd.final_cycle_point("2023-10-10T06:00:00Z"), + qd.runahead_limit("P2"), + qd.jedi_build_method("use_existing"), + qd.model_components(['geos_atmosphere']), + ], + geos_atmosphere=[ + qd.cycle_times([ + "T00", + ]), + qd.background_experiment('x0050'), + qd.geos_x_background_directory("/discover/nobackup/projects/gmao/dadev/" + "rtodling/archive/Restarts/JEDI/541x"), + qd.geos_x_ensemble_directory("/discover/nobackup/projects/gmao/dadev/" + 'rtodling/archive/541/Milan'), + qd.npx_proc(4), + qd.npy_proc(5), + qd.perhost(120), + qd.window_length("PT6H"), + qd.window_type("3D"), + qd.horizontal_resolution("91"), + qd.gsibec_nlats("91"), + qd.gsibec_nlons("144"), + qd.vertical_resolution("72"), + qd.ensemble_num_members(4), + qd.ensemble_num_chunks(2), + qd.number_of_iterations([10]), + qd.gradient_norm_reduction(1.e-3), + qd.minimizer("DRPLanczos"), + qd.analysis_variables([ + "eastward_wind", + "northward_wind", + "air_temperature", + "water_vapor_mixing_ratio_wrt_moist_air", + "air_pressure_at_surface", + "air_pressure_levels", + "cloud_liquid_ice", + "cloud_liquid_water", + "rain_water", + "snow_water", + "mole_fraction_of_ozone_in_air", + "geopotential_height_times_gravity_at_surface", + "fraction_of_ocean", + "fraction_of_lake", + "fraction_of_ice", + "skin_temperature_at_surface" + ]), + qd.observations([ + "sondes", + ]), + qd.obs_thinning_rej_fraction(0.8), + qd.ensmeanvariance_spec([ + {"state": "bkg", + "fn_input": "ebkg/mem%mem%/geos.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "geos.prior.mean", + "fn_output_variance": "geos.prior.variance", + "grid_type": ['cs', 'latlon']}, + {"state": "analysis", + "fn_input": "analysis/mem%mem%/eda.ana.mem%mem%.%yyyy%mm%dd_%hh%MM%ssz.nc4", + "fn_output_mean": "eda.ana.mean", + "fn_output_variance": "eda.ana.variance", + "grid_type": ['cs', 'latlon']}, + ]), + qd.diffstates_spec({ + "state1": + {"fn_input": "geos.prior.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + "state2": + {"fn_input": "eda.ana.mean.%yyyy%mm%dd_%hh%MM%ssz.nc4"}, + "state_diff": + {"fn_output": "eda.mean-inc", "grid_type": ['cs', 'latlon']}, + "state_type": "ensemble" + }), + qd.clean_patterns(['*.txt', '*.csv']), + ] + ) + eda_controlpert_atmos_tier1 = QuestionList( list_name="eda_controlpert_atmos_tier1", questions=[ @@ -35,9 +115,10 @@ class SuiteConfig(QuestionContainer, Enum): qd.cycle_times([ "T00", ]), - qd.geos_x_background_directory("/discover/nobackup/projects/gmao/" - "dadev/rtodling/archive/Restarts/JEDI/541x"), - qd.geos_x_ensemble_directory('/discover/nobackup/projects/gmao/dadev/' + qd.background_experiment('x0050'), + qd.geos_x_background_directory("/discover/nobackup/projects/gmao/dadev/" + "rtodling/archive/Restarts/JEDI/541x"), + qd.geos_x_ensemble_directory("/discover/nobackup/projects/gmao/dadev/" 'rtodling/archive/541/Milan'), qd.npx_proc(4), qd.npy_proc(5), @@ -138,7 +219,7 @@ class SuiteConfig(QuestionContainer, Enum): eda_controlpert_atmos = QuestionList( list_name="eda_controlpert_atmos", questions=[ - eda_controlpert_atmos_tier1 + eda_controlpert_atmos_tier1_fast ] ) From 1d0c259c42d97589dabeca18f92c7d181ad75954 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Tue, 4 Aug 2026 23:37:43 -0400 Subject: [PATCH 66/68] delete obsdataout --- src/swell/tasks/run_jedi_eda_executable.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/swell/tasks/run_jedi_eda_executable.py b/src/swell/tasks/run_jedi_eda_executable.py index 7a0936b63..a9011f723 100644 --- a/src/swell/tasks/run_jedi_eda_executable.py +++ b/src/swell/tasks/run_jedi_eda_executable.py @@ -177,10 +177,12 @@ def execute(self) -> None: observer.update({'obs error': obs_error_dict}) observer['obs space'].update({'obs perturbations seed': imember}) - hxout = observer['obs space']['obsdataout']['engine']['obsfile'] - dir1, fname = os.path.split(hxout) - hxout = os.path.join(dir1, mem_dir, fname) - observer['obs space']['obsdataout']['engine']['obsfile'] = hxout + del observer['obs space']['obsdataout'] + # donot output obsdataout + # hxout = observer['obs space']['obsdataout']['engine']['obsfile'] + # dir1, fname = os.path.split(hxout) + # hxout = os.path.join(dir1, mem_dir, fname) + # observer['obs space']['obsdataout']['engine']['obsfile'] = hxout obsFileIn = observer['obs space']['obsdatain']['engine']['obsfile'] dir1, fname = os.path.split(obsFileIn) From d7f2e3227227a979981ffeae0658f98185eafa13 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 5 Aug 2026 10:14:01 -0400 Subject: [PATCH 67/68] fix warnings from python code style --- .../jedi/oops/eda_control_pert.py | 5 +- .../eda_controlpert_atmos/suite_config.py | 49 ++++++----------- src/swell/tasks/clean_eda.py | 8 +-- .../run_jedi_eda_control_pert_executable.py | 52 +++++++++---------- .../utilities/render_jedi_interface_files.py | 1 - src/swell/utilities/yaml_utils.py | 5 +- 6 files changed, 52 insertions(+), 68 deletions(-) diff --git a/src/swell/configuration/jedi/oops/eda_control_pert.py b/src/swell/configuration/jedi/oops/eda_control_pert.py index d8426d7db..5834aaa0d 100644 --- a/src/swell/configuration/jedi/oops/eda_control_pert.py +++ b/src/swell/configuration/jedi/oops/eda_control_pert.py @@ -17,7 +17,7 @@ def render_oops(self): nmember = self.template_dict['ensemble_num_members'] nchunk = self.template_dict['ensemble_num_chunks'] ichunk = self.template_dict['ensemble_ichunk'] - nstate = int ( nmember / nchunk ) + nstate = int(nmember/nchunk) if ichunk == 1: num_pert_mem = nstate - 1 else: @@ -49,7 +49,8 @@ def render_oops(self): }, 'iterations': [{ 'geometry': self.interface_model('geometry_inner'), - 'gradient norm reduction': float(self.template_dict['gradient_norm_reduction']), + 'gradient norm reduction': float( + self.template_dict['gradient_norm_reduction']), 'ninner': self.template_dict['number_of_iterations'], }], }, diff --git a/src/swell/suites/eda_controlpert_atmos/suite_config.py b/src/swell/suites/eda_controlpert_atmos/suite_config.py index b478e8710..54443095d 100644 --- a/src/swell/suites/eda_controlpert_atmos/suite_config.py +++ b/src/swell/suites/eda_controlpert_atmos/suite_config.py @@ -152,41 +152,22 @@ class SuiteConfig(QuestionContainer, Enum): "fraction_of_ice", "skin_temperature_at_surface" ]), + # + # Report first crash point: Bound-1 + # qd.observations([ -# "aircraft_temperature", -# "aircraft_wind", -# "airs_aqua", -# "amsr2_gcom-w1", -# "amsua_aqua", -# "amsua_metop-b", -# "amsua_metop-c", -# "amsua_n15", -# "amsua_n18", -# "amsua_n19", -# "atms_n20", -# "atms_npp", -# "avhrr3_metop-b", -# "avhrr3_n18", -# "avhrr3_n19", -# "cris-fsr_n20", -# "cris-fsr_npp", -# "gmi_gpm", -# "gps", -# "iasi_metop-b", -# "iasi_metop-c", -# "mhs_metop-b", -# "mhs_metop-c", -# "mhs_n19", -# "mls55_aura", -# "omi_aura", -# "ompsnm_npp", -# "pibal", -# "satwind", -# "scatwind", -# "sfcship", -# "sfc", - "sondes", -# "ssmis_f17" + "aircraft_temperature", + "aircraft_wind", + "airs_aqua", + "amsr2_gcom-w1", + "amsua_aqua", + "amsua_metop-b", + "amsua_metop-c", + "amsua_n15", + "amsua_n18", + "amsua_n19", + "atms_n20", + "atms_npp", ]), qd.obs_thinning_rej_fraction(0.8), qd.ensmeanvariance_spec([ diff --git a/src/swell/tasks/clean_eda.py b/src/swell/tasks/clean_eda.py index c630c69bd..eb9e1205e 100644 --- a/src/swell/tasks/clean_eda.py +++ b/src/swell/tasks/clean_eda.py @@ -125,15 +125,15 @@ def execute(self) -> None: # This special design works with either eda or eda_control_pert # handle eda case: analysis/mem00x # eda controlpert case: analysis_chunk/chunk00x/mem00y - # This avoids blindly search for mem0xx dir, as some of them + # This avoids blindly search for mem0xx dir, as some of them # only contains linked analysis files and donot require clean up - #--------------------------------------------------------------- + # -------------------------------------------------------------- d1 = os.path.join(self.cycle_dir(), 'analysis_chunk') # control pert case if not os.path.exists(d1): d1 = os.path.join(self.cycle_dir(), 'analysis') # eda case - target_dirs = [str(p) for p in Path(d1).rglob('mem*') if p.is_dir()] + target_dirs = [str(p) for p in Path(d1).rglob('mem*') if p.is_dir()] self.logger.info(f'target_dirs = {target_dirs}') - + for mem_dir in target_dirs: d2 = os.path.join(mem_dir, 'fv3-jedi') if os.path.exists(d2): diff --git a/src/swell/tasks/run_jedi_eda_control_pert_executable.py b/src/swell/tasks/run_jedi_eda_control_pert_executable.py index dda902dc1..849e869ea 100644 --- a/src/swell/tasks/run_jedi_eda_control_pert_executable.py +++ b/src/swell/tasks/run_jedi_eda_control_pert_executable.py @@ -70,26 +70,27 @@ def execute(self) -> None: nchunk = self.config.ensemble_num_chunks() ichunk = self.get_ensemble_ichunk() - # exit execute if ichunk=-1: meaning reoder analysis files - # ------------------------------------------------------- + # exit execute if ichunk=-1: meaning reoder analysis files and return + # ------------------------------------------------------------------- if ichunk == -1: - npert = int( nmember / nchunk ) + npert = int(nmember/nchunk) istart = 0 imem = 0 - for xchunk in range (1, nchunk+1): + for xchunk in range(1, nchunk+1): if xchunk == 1: - iend = npert -1 # mem: [0, 1 ... npert-1] : [ctrl, all pert] + iend = npert - 1 # mem: [0, 1 ... npert-1] : [ctrl, all pert] else: - iend = npert # mem: [0, 1 ... npert] : [ctrl, all pert] + iend = npert # mem: [0, 1 ... npert] : [ctrl, all pert] for imem_in_chunk in range(istart, iend+1): - # skip extra ctrl + # skip extra ctrl, because of ichunk=-1 if xchunk > 1 and imem_in_chunk == 0: continue else: imem += 1 dir_a = f'analysis_chunk/chunk{xchunk:03d}/mem{imem_in_chunk:03d}' dir_b = f'analysis/mem{imem:03d}' - dir_a_full = os.path.join(self.cycle_dir(), dir_a, f'eda.ana.mem{imem_in_chunk:03d}.*.nc4') + dir_a_full = os.path.join(self.cycle_dir(), dir_a, + f'eda.ana.mem{imem_in_chunk:03d}.*.nc4') # print(f'dir_a_full = {dir_a_full}') fa_list = glob(dir_a_full) # print(f'fa_list = {fa_list}') @@ -98,10 +99,11 @@ def execute(self) -> None: else: fa = '' self.logger.error( - f"analysis files not found ichunk={ichunk}, imem_in_chunk={imem_in_chunk}") + f"analysis files not found ichunk={ichunk}, " + f"imem_in_chunk={imem_in_chunk}") print(f'fa = {fa}') - tail = '.'.join(os.path.basename(fa).split('.')[-2:]) - fb = os.path.join(self.cycle_dir(), dir_b, f'eda.ana.mem{imem:03d}.{tail}') + tail = '.'.join(os.path.basename(fa).split('.')[-2:]) + fb = os.path.join(self.cycle_dir(), dir_b, f'eda.ana.mem{imem:03d}.{tail}') print(f'fb = {fb}') # link fa to fb # Convert strings to Path objects @@ -118,7 +120,6 @@ def execute(self) -> None: print(f"Successfully linked {fa} to {fb}") return - # Populate jedi interface templates dictionary # -------------------------------------------- self.jedi_rendering.add_key('window_begin_iso', window_begin_iso) @@ -175,7 +176,7 @@ def execute(self) -> None: self.jedi_rendering.add_key('background_frequency', self.config.background_frequency()) # Jedi configuration file # ----------------------- - fname=f'jedi_{jedi_application}{window_type}_config_chunk{ichunk:03d}.yaml' + fname = f'jedi_{jedi_application}{window_type}_config_chunk{ichunk:03d}.yaml' jedi_config_file = os.path.join(self.cycle_dir(), fname) print(f'file = {jedi_config_file}') @@ -191,13 +192,12 @@ def execute(self) -> None: jedi_forecast_model) print(f'jedi file') - - npert = int( nmember / nchunk ) + npert = int(nmember/nchunk) istart = 0 if ichunk == 1: - iend = npert -1 # mem: [0, 1 ... npert-1] : [ctrl, all pert] + iend = npert - 1 # mem: [0, 1 ... npert-1] : [ctrl, all pert] else: - iend = npert # mem: [0, 1 ... npert] : [ctrl, all pert] + iend = npert # mem: [0, 1 ... npert] : [ctrl, all pert] print(f'npert = {npert}') @@ -212,7 +212,7 @@ def execute(self) -> None: if imem == 0: id = 1 else: - id = npert * (ichunk -1) + imem + id = npert*(ichunk-1) + imem # link ebkg to ebkg_chunk # ----------------------- @@ -240,7 +240,8 @@ def execute(self) -> None: xdir = os.path.join(self.cycle_dir(), mem_dir) os.makedirs(xdir, exist_ok=True) - for observer in jedi_config_dict['assimilation']['cost function']['observations']['observers']: + for observer in jedi_config_dict['assimilation'][ + 'cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] print(f'ob= {observation}') @@ -257,7 +258,8 @@ def execute(self) -> None: # round-2: modify dir keys in yaml to chunk00x / mem%mem_pad% # ------------------------------------------------------------ - for observer in jedi_config_dict['assimilation']['cost function']['observations']['observers']: + for observer in jedi_config_dict['assimilation'][ + 'cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] print(f'ob= {observation}') @@ -313,7 +315,6 @@ def execute(self) -> None: with open(jedi_config_file, 'w') as jedi_config_file_open: ruamel_yaml.dump(jedi_config_dict, jedi_config_file_open) - # Get the JEDI interface metadata # ------------------------------- model_component_meta = self.jedi_rendering.render_interface_meta() @@ -324,16 +325,15 @@ def execute(self) -> None: np = eval(str(model_component_meta['total_processors'])) # modify np by (nmember / nchunk) + 1 if ichunk == 1: - np = np * npert + np = np * npert else: - np = np * (npert + 1) - + np = np * (npert + 1) # Jedi executable name # -------------------- jedi_executable = model_component_meta['executables']['edaControlPert'] - jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', 'build', 'bin', - jedi_executable) + jedi_executable_path = os.path.join(self.experiment_path(), 'jedi_bundle', + 'build', 'bin', jedi_executable) # Run the JEDI executable # ----------------------- diff --git a/src/swell/utilities/render_jedi_interface_files.py b/src/swell/utilities/render_jedi_interface_files.py index 83c59c9f8..e94383635 100644 --- a/src/swell/utilities/render_jedi_interface_files.py +++ b/src/swell/utilities/render_jedi_interface_files.py @@ -318,7 +318,6 @@ def render_interface_meta(self, model_component_in: Union[str, dict, None] = Non # Path to configuration file config_file = os.path.join(self.jedi_config_path, 'interfaces', model_component, f'{model_component}.yaml') - print( f'config_file = {config_file}') # Render templates in file and return dictionary return self.__open_file_render_to_dict__(config_file) diff --git a/src/swell/utilities/yaml_utils.py b/src/swell/utilities/yaml_utils.py index a37c5982f..d404f33ed 100644 --- a/src/swell/utilities/yaml_utils.py +++ b/src/swell/utilities/yaml_utils.py @@ -2,6 +2,7 @@ # -------------------------------------------------------------------------------------------------- + def replace_key(obj, old_key, new_key): """ Recursively replace dictionary keys in nested dictionaries/lists. @@ -43,15 +44,17 @@ def replace_string_value(data, sa, sb): # -------------------------------------------------------------------------------------------------- + def print_dict(*args, **kwargs): return print_dict_as_yaml(*args, **kwargs) + def print_dict_as_yaml(data_dict): # Convert the dictionary to a YAML string # default_flow_style=False ensures block formatting instead of inline JSON-like formatting # sort_keys=False preserves your dictionary's original key order (Python 3.7+) yaml_string = yaml.dump(data_dict, default_flow_style=False, sort_keys=False) - + # Output to the terminal print(yaml_string) From 56a47032b0d6560ef1f87f5231398c2230225f2d Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 5 Aug 2026 11:44:25 -0400 Subject: [PATCH 68/68] fix warnings from code tests --- .../eda_controlpert_atmos/suite_config.py | 2 ++ src/swell/tasks/clean_eda.py | 2 -- .../run_jedi_eda_control_pert_executable.py | 31 +++++++------------ 3 files changed, 13 insertions(+), 22 deletions(-) diff --git a/src/swell/suites/eda_controlpert_atmos/suite_config.py b/src/swell/suites/eda_controlpert_atmos/suite_config.py index 54443095d..b6c234e9a 100644 --- a/src/swell/suites/eda_controlpert_atmos/suite_config.py +++ b/src/swell/suites/eda_controlpert_atmos/suite_config.py @@ -51,6 +51,7 @@ class SuiteConfig(QuestionContainer, Enum): qd.vertical_resolution("72"), qd.ensemble_num_members(4), qd.ensemble_num_chunks(2), + qd.obs_pert_amplitude(0.5), qd.number_of_iterations([10]), qd.gradient_norm_reduction(1.e-3), qd.minimizer("DRPLanczos"), @@ -131,6 +132,7 @@ class SuiteConfig(QuestionContainer, Enum): qd.vertical_resolution("72"), qd.ensemble_num_members(32), qd.ensemble_num_chunks(16), + qd.obs_pert_amplitude(0.5), qd.number_of_iterations([50]), qd.gradient_norm_reduction(1.e-3), qd.minimizer("DRPLanczos"), diff --git a/src/swell/tasks/clean_eda.py b/src/swell/tasks/clean_eda.py index eb9e1205e..b0558c021 100644 --- a/src/swell/tasks/clean_eda.py +++ b/src/swell/tasks/clean_eda.py @@ -59,8 +59,6 @@ def execute(self) -> None: window_begin = self.da_window_params.window_begin(window_length) window_begin_iso = self.da_window_params.window_begin_iso(window_length) window_end_iso = self.da_window_params.window_end_iso(window_length) - nmember = self.config.ensemble_num_members() - # imember = self.get_ensemble_imember() # Populate jedi interface templates dictionary # -------------------------------------------- diff --git a/src/swell/tasks/run_jedi_eda_control_pert_executable.py b/src/swell/tasks/run_jedi_eda_control_pert_executable.py index 849e869ea..ec9ca2e1d 100644 --- a/src/swell/tasks/run_jedi_eda_control_pert_executable.py +++ b/src/swell/tasks/run_jedi_eda_control_pert_executable.py @@ -91,9 +91,7 @@ def execute(self) -> None: dir_b = f'analysis/mem{imem:03d}' dir_a_full = os.path.join(self.cycle_dir(), dir_a, f'eda.ana.mem{imem_in_chunk:03d}.*.nc4') - # print(f'dir_a_full = {dir_a_full}') fa_list = glob(dir_a_full) - # print(f'fa_list = {fa_list}') if fa_list: fa = fa_list[0] else: @@ -101,10 +99,8 @@ def execute(self) -> None: self.logger.error( f"analysis files not found ichunk={ichunk}, " f"imem_in_chunk={imem_in_chunk}") - print(f'fa = {fa}') tail = '.'.join(os.path.basename(fa).split('.')[-2:]) fb = os.path.join(self.cycle_dir(), dir_b, f'eda.ana.mem{imem:03d}.{tail}') - print(f'fb = {fb}') # link fa to fb # Convert strings to Path objects fa = Path(fa) @@ -117,7 +113,7 @@ def execute(self) -> None: # a3. Create the symbolic link # .resolve() gets the absolute path (like realpath in bash) fb.symlink_to(fa.resolve()) - print(f"Successfully linked {fa} to {fb}") + self.logger.info(f"Successfully linked {fa} to {fb}") return # Populate jedi interface templates dictionary @@ -178,7 +174,6 @@ def execute(self) -> None: # ----------------------- fname = f'jedi_{jedi_application}{window_type}_config_chunk{ichunk:03d}.yaml' jedi_config_file = os.path.join(self.cycle_dir(), fname) - print(f'file = {jedi_config_file}') # Output log file # --------------- @@ -190,7 +185,6 @@ def execute(self) -> None: jedi_config_dict = self.jedi_rendering.render_oops_file(f'{jedi_application}', window_type, jedi_forecast_model) - print(f'jedi file') npert = int(nmember/nchunk) istart = 0 @@ -199,11 +193,8 @@ def execute(self) -> None: else: iend = npert # mem: [0, 1 ... npert] : [ctrl, all pert] - print(f'npert = {npert}') - - # round-1: link bkg, copy obs, B and R + # round-1: link bkg, copy obs, fv3jedi dir (B and R) # ---------------------------------------------------- - chunk_dir = f'analysis_chunk/chunk{ichunk:003d}' mem_temp_dir = f'analysis_chunk/chunk{ichunk:003d}/mem%mem_pad%' for imem in range(istart, iend+1): if ichunk == 1: @@ -229,7 +220,6 @@ def execute(self) -> None: f2 = f2.split('.')[2:] f2 = '.'.join(f2) f2 = os.path.join(xdir, f2) - print(f'f2 = {f2}') if os.path.lexists(f2): os.remove(f2) os.symlink(f1, f2) @@ -244,11 +234,10 @@ def execute(self) -> None: 'cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] - print(f'ob= {observation}') # copy obs input file to avoid multi MPI reading the same file files = glob(os.path.join(self.cycle_dir(), f'{observation}.*')) for src_file in files: - print(f'f= {src_file}') + self.loggerin.info(f'f= {src_file} is copied to {xdir}') shutil.copy(src_file, xdir) # copy fv3-jedi dir, update dir names @@ -262,11 +251,13 @@ def execute(self) -> None: 'cost function']['observations']['observers']: # Get observation name observation = observer['observation_name'] - print(f'ob= {observation}') - hxout = observer['obs space']['obsdataout']['engine']['obsfile'] - dir1, fname = os.path.split(hxout) - hxout = os.path.join(dir1, mem_temp_dir, fname) - observer['obs space']['obsdataout']['engine']['obsfile'] = hxout + + # for now delete obsdataout + del observer['obs space']['obsdataout'] + # hxout = observer['obs space']['obsdataout']['engine']['obsfile'] + # dir1, fname = os.path.split(hxout) + # hxout = os.path.join(dir1, mem_temp_dir, fname) + # observer['obs space']['obsdataout']['engine']['obsfile'] = hxout obsFileIn = observer['obs space']['obsdatain']['engine']['obsfile'] dir1, fname = os.path.split(obsFileIn) @@ -319,7 +310,7 @@ def execute(self) -> None: # ------------------------------- model_component_meta = self.jedi_rendering.render_interface_meta() - print(f"proc = {model_component_meta['total_processors']}") + self.logger.info(f"num of processors = {model_component_meta['total_processors']}") # Compute number of processors # ---------------------------- np = eval(str(model_component_meta['total_processors']))