upload sigma level models

FuXi-sigma-base/checkpoints/checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:938f77598d0d82c647a2cd75b09bb9cc488af0eefd0524011840068e1fee85b0
+size 940

FuXi-sigma-base/checkpoints/model_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:686aba049e2c5a3df80f0d1bc2522af473ccbb1787e7738762e616f4f2060cc1
+size 5970491642

FuXi-sigma-base/configs/model_multi.yml

@@ -0,0 +1,197 @@
+# --------------------------------------------------------------------------------------------------------------------- #
+# This yaml file implements 6 hourly FuXi on NSF NCAR HPCs (casper.ucar.edu and derecho.hpc.ucar.edu)
+# Technical details are available in:
+#
+#     Sha, Y., J. Schreck, W. Chapman, D. J. Gagne II, 2025: Investigating the contribution of terrain-following
+#     coordinates and conservation schemes in AI-driven precipitation forecasts. Submitted to:
+#     Geophysical Research Letters. pre-print: https://arxiv.org/abs/2503.00332
+#
+# --------------------------------------------------------------------------------------------------------------------- #
+
+save_loc: '/glade/work/ksha/CREDIT_runs/fuxi_mlevel_dry/'
+seed: 1000
+
+data:
+    # upper-air variables
+    variables: ['specific_total_water', 'temperature', 'u_component_of_wind','v_component_of_wind']
+    save_loc: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # surface variables
+    surface_variables: ['SP', 'VAR_2T', 'VAR_10U', 'VAR_10V']
+    save_loc_surface: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # dynamic forcing variables
+    dynamic_forcing_variables: ['toa_incident_solar_radiation', 'land_sea_CI_mask']
+    save_loc_dynamic_forcing: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # diagnostic variables
+    diagnostic_variables: ['evaporation', 'total_precipitation',
+                           'surface_net_solar_radiation',
+                           'surface_net_thermal_radiation',
+                           'surface_sensible_heat_flux',
+                           'surface_latent_heat_flux',
+                           'top_net_solar_radiation',
+                           'top_net_thermal_radiation']
+    save_loc_diagnostic: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # static variables
+    static_variables: ['z_norm', 'soil_type'] 
+    save_loc_static: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+
+    # physics file
+    save_loc_physics: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+    
+    # mean / std path
+    mean_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/mean_6h_1979_2019_conserve_1deg.nc'
+    std_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/std_residual_6h_1979_2019_conserve_1deg.nc'
+    
+    # train / validation split
+    train_years: [1979, 2019]
+    valid_years: [2019, 2020]
+    
+    # data workflow
+    scaler_type: 'std_new'
+     
+    history_len: 2
+    valid_history_len: 2
+    
+    forecast_len: 1 # <----------- forecast_len increases from 1 to 11 for 12 hr to 72 hr multi-steps
+    valid_forecast_len: 11
+
+    one_shot: False
+    
+    # 1 for hourly model
+    lead_time_periods: 6
+    
+    # do not use skip_period
+    skip_periods: null
+    
+    # compatible with the old 'std'
+    static_first: True
+    
+    sst_forcing:
+        activate: False
+
+trainer:
+    type: multi-step
+    
+    mode: fsdp
+    cpu_offload: False
+    activation_checkpoint: True
+    
+    load_weights: True
+    load_optimizer: False
+    load_scaler: False
+    load_sheduler: False
+    num_epoch: 1
+    
+    skip_validation: False
+    update_learning_rate: False
+    
+    save_backup_weights: True
+    save_best_weights: True
+    save_metric_vars: True
+
+    learning_rate: 3.0e-07
+    weight_decay: 3.0e-06
+    
+    train_batch_size: 1
+    valid_batch_size: 1
+    
+    batches_per_epoch: 0
+    valid_batches_per_epoch: 0
+    stopping_patience: 999
+    
+    start_epoch: 0
+    
+    reload_epoch: True
+    epochs: &epochs 999
+    
+    use_scheduler: False
+    
+    # Automatic Mixed Precision: False
+    amp: False
+    
+    # rescale loss as loss = loss / grad_accum_every
+    grad_accum_every: 1 
+    # gradient clipping
+    grad_max_norm: 999.0
+    
+    # number of workers
+    thread_workers: 4
+    valid_thread_workers: 0
+
+model:
+    type: "fuxi"
+    
+    frames: 2                # number of input states
+    image_height: 181        # number of latitude grids
+    image_width: 360        # number of longitude grids
+    levels: 18               # number of upper-air variable levels
+    channels: 4              # upper-air variable channels
+    surface_channels: 4      # surface variable channels
+    input_only_channels: 4   # dynamic forcing, forcing, static channels
+    output_only_channels: 8  # diagnostic variable channels
+
+    # patchify layer
+    patch_height: 4         # number of latitude grids in each 3D patch
+    patch_width: 4          # number of longitude grids in each 3D patch
+    frame_patch_size: 2     # number of input states in each 3D patch
+    
+    # hidden layers
+    dim: 1536               # dimension (default: 1536)
+    num_groups: 32          # number of groups (default: 32)
+    num_heads: 8            # number of heads (default: 8)
+    window_size: 7          # window size (default: 7)
+    depth: 48               # number of swin transformers (default: 48)
+    
+    # use spectral norm
+    use_spectral_norm: True
+
+    # use interpolation to match the output size
+    interp: False
+    
+    # map boundary padding
+    padding_conf:
+        activate: True
+        mode: earth
+        pad_lat: [21, 22]
+        pad_lon: [44, 44]
+    
+    post_conf: 
+        activate: False
+
+loss: 
+    # the main training loss
+    training_loss: "mse"
+    
+    # power loss (x), spectral_loss (x)
+    use_power_loss: False 
+    use_spectral_loss: False
+    
+    # use latitude weighting
+    use_latitude_weights: True
+    latitude_weights: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+
+    # turn-off variable weighting
+    use_variable_weights: True
+    variable_weights:
+        specific_total_water: [1.36577589e-06, 1.19436047e-05, 1.03276589e-04, 4.46097338e-04, 1.26807904e-03, 2.77554864e-03, 5.13295742e-03, 8.70407041e-03, 1.41016958e-02, 4.12818192e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02]
+        temperature: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        u_component_of_wind: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        v_component_of_wind: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        SP: 0.13
+        VAR_2T: 0.13
+        VAR_10U: 0.13
+        VAR_10V: 0.13
+        surface_latent_heat_flux: 0.065
+        surface_net_solar_radiation: 0.065
+        evaporation: 0.065
+        surface_net_thermal_radiation: 0.065
+        toa_incident_solar_radiation: 0.065
+        surface_sensible_heat_flux: 0.065
+        total_precipitation: 0.065
+        top_net_thermal_radiation: 0.065
+        top_net_solar_radiation: 0.065
+
+

FuXi-sigma-base/configs/model_predict.yml

@@ -0,0 +1,195 @@
+# --------------------------------------------------------------------------------------------------------------------- #
+# This yaml file implements 6 hourly FuXi on NSF NCAR HPCs (casper.ucar.edu and derecho.hpc.ucar.edu)
+# Technical details are available in:
+#
+#     Sha, Y., J. Schreck, W. Chapman, D. J. Gagne II, 2025: Investigating the contribution of terrain-following
+#     coordinates and conservation schemes in AI-driven precipitation forecasts. Submitted to:
+#     Geophysical Research Letters. pre-print: https://arxiv.org/abs/2503.00332
+#
+# --------------------------------------------------------------------------------------------------------------------- #
+
+save_loc: '/glade/work/ksha/CREDIT_runs/fuxi_mlevel_dry/'
+seed: 1000
+
+data:
+    # upper-air variables
+    variables: ['specific_total_water', 'temperature', 'u_component_of_wind','v_component_of_wind']
+    save_loc: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # surface variables
+    surface_variables: ['SP', 'VAR_2T', 'VAR_10U', 'VAR_10V']
+    save_loc_surface: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # dynamic forcing variables
+    dynamic_forcing_variables: ['toa_incident_solar_radiation', 'land_sea_CI_mask']
+    save_loc_dynamic_forcing: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # diagnostic variables
+    diagnostic_variables: ['evaporation', 'total_precipitation',
+                           'surface_net_solar_radiation',
+                           'surface_net_thermal_radiation',
+                           'surface_sensible_heat_flux',
+                           'surface_latent_heat_flux',
+                           'top_net_solar_radiation',
+                           'top_net_thermal_radiation']
+    save_loc_diagnostic: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # static variables
+    static_variables: ['z_norm', 'soil_type'] 
+    save_loc_static: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+
+    # physics file
+    save_loc_physics: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+    
+    # mean / std path
+    mean_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/mean_6h_1979_2019_conserve_1deg.nc'
+    std_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/std_residual_6h_1979_2019_conserve_1deg.nc'
+    
+    # train / validation split
+    train_years: [1979, 2019]
+    valid_years: [2019, 2020]
+    
+    # data workflow
+    scaler_type: 'std_new'
+     
+    history_len: 2
+    valid_history_len: 2
+    
+    forecast_len: 0
+    valid_forecast_len: 0
+
+    one_shot: False
+    
+    # 1 for hourly model
+    lead_time_periods: 6
+    
+    # do not use skip_period
+    skip_periods: null
+    
+    # compatible with the old 'std'
+    static_first: True
+    
+    sst_forcing:
+        activate: False
+
+trainer:
+    type: standard
+    
+    mode: fsdp
+    cpu_offload: False
+    activation_checkpoint: True
+    
+    load_weights: True
+    load_optimizer: True
+    load_scaler: True
+    load_sheduler: True
+    num_epoch: 8
+    
+    skip_validation: False
+    update_learning_rate: False
+    
+    save_backup_weights: True
+    save_best_weights: True
+    save_metric_vars: True
+
+    learning_rate: 1.0e-03
+    weight_decay: 3.0e-06
+    
+    train_batch_size: 1
+    valid_batch_size: 1
+    
+    batches_per_epoch: 0
+    valid_batches_per_epoch: 0
+    stopping_patience: 999
+    
+    start_epoch: 0
+
+    reload_epoch: True
+    epochs: &epochs 200
+    
+    use_scheduler: True
+    scheduler: {'scheduler_type': 'cosine-annealing', 'T_max': *epochs,  'last_epoch': -1}
+
+    # Automatic Mixed Precision: False
+    amp: False
+    
+    # rescale loss as loss = loss / grad_accum_every
+    grad_accum_every: 1 
+    # gradient clipping
+    grad_max_norm: 32.0
+    
+    # number of workers
+    thread_workers: 4
+    valid_thread_workers: 0
+
+model:
+    type: "fuxi"
+    
+    frames: 2                # number of input states
+    image_height: 181        # number of latitude grids
+    image_width: 360        # number of longitude grids
+    levels: 18               # number of upper-air variable levels
+    channels: 4              # upper-air variable channels
+    surface_channels: 4      # surface variable channels
+    input_only_channels: 4   # dynamic forcing, forcing, static channels
+    output_only_channels: 8  # diagnostic variable channels
+
+    # patchify layer
+    patch_height: 4         # number of latitude grids in each 3D patch
+    patch_width: 4          # number of longitude grids in each 3D patch
+    frame_patch_size: 2     # number of input states in each 3D patch
+    
+    # hidden layers
+    dim: 1536               # dimension (default: 1536)
+    num_groups: 32          # number of groups (default: 32)
+    num_heads: 8            # number of heads (default: 8)
+    window_size: 7          # window size (default: 7)
+    depth: 48               # number of swin transformers (default: 48)
+    
+    # use spectral norm
+    use_spectral_norm: True
+
+    # use interpolation to match the output size
+    interp: False
+    
+    # map boundary padding
+    padding_conf:
+        activate: True
+        mode: earth
+        pad_lat: [21, 22]
+        pad_lon: [44, 44]
+    
+    post_conf: 
+        activate: False
+
+loss: 
+    # the main training loss
+    training_loss: "mse"
+    
+    # power loss (x), spectral_loss (x)
+    use_power_loss: False 
+    use_spectral_loss: False
+    
+    # use latitude weighting
+    use_latitude_weights: True
+    latitude_weights: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+
+    # turn-off variable weighting
+
+predict:
+    forecasts:
+        type: "custom"       # keep it as "custom"
+        start_year: 2020     # year of the first initialization (where rollout will start)
+        start_month: 1       # month of the first initialization
+        start_day: 1         # day of the first initialization
+        start_hours: [0, 12] # hour-of-day for each initialization, 0 for 00Z, 12 for 12Z
+        duration: 768        # number of days to initialize, starting from the (year, mon, day) above
+                             # duration should be divisible by the number of GPUs
+                             # (e.g., duration: 384 for 365-day rollout using 32 GPUs)
+        days: 15             # forecast lead time as days (1 means 24-hour forecast)
+
+    save_forecast: '/glade/derecho/scratch/ksha/CREDIT/RAW_OUTPUT/fuxi_mlevel_dry2/'
+    use_laplace_filter: False
+
+
+

FuXi-sigma-base/configs/model_single.yml

@@ -0,0 +1,197 @@
+# --------------------------------------------------------------------------------------------------------------------- #
+# This yaml file implements 6 hourly FuXi on NSF NCAR HPCs (casper.ucar.edu and derecho.hpc.ucar.edu)
+# Technical details are available in:
+#
+#     Sha, Y., J. Schreck, W. Chapman, D. J. Gagne II, 2025: Investigating the contribution of terrain-following
+#     coordinates and conservation schemes in AI-driven precipitation forecasts. Submitted to:
+#     Geophysical Research Letters. pre-print: https://arxiv.org/abs/2503.00332
+#
+# --------------------------------------------------------------------------------------------------------------------- #
+
+save_loc: '/glade/work/ksha/CREDIT_runs/fuxi_mlevel_dry/'
+seed: 1000
+
+data:
+    # upper-air variables
+    variables: ['specific_total_water', 'temperature', 'u_component_of_wind','v_component_of_wind']
+    save_loc: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # surface variables
+    surface_variables: ['SP', 'VAR_2T', 'VAR_10U', 'VAR_10V']
+    save_loc_surface: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # dynamic forcing variables
+    dynamic_forcing_variables: ['toa_incident_solar_radiation', 'land_sea_CI_mask']
+    save_loc_dynamic_forcing: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # diagnostic variables
+    diagnostic_variables: ['evaporation', 'total_precipitation',
+                           'surface_net_solar_radiation',
+                           'surface_net_thermal_radiation',
+                           'surface_sensible_heat_flux',
+                           'surface_latent_heat_flux',
+                           'top_net_solar_radiation',
+                           'top_net_thermal_radiation']
+    save_loc_diagnostic: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # static variables
+    static_variables: ['z_norm', 'soil_type'] 
+    save_loc_static: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+
+    # physics file
+    save_loc_physics: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+    
+    # mean / std path
+    mean_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/mean_6h_1979_2019_conserve_1deg.nc'
+    std_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/std_residual_6h_1979_2019_conserve_1deg.nc'
+    
+    # train / validation split
+    train_years: [1979, 2019]
+    valid_years: [2019, 2020]
+    
+    # data workflow
+    scaler_type: 'std_new'
+     
+    history_len: 2
+    valid_history_len: 2
+    
+    forecast_len: 0
+    valid_forecast_len: 0
+
+    one_shot: False
+    
+    # 1 for hourly model
+    lead_time_periods: 6
+    
+    # do not use skip_period
+    skip_periods: null
+    
+    # compatible with the old 'std'
+    static_first: True
+    
+    sst_forcing:
+        activate: False
+
+trainer:
+    type: standard
+    
+    mode: fsdp
+    cpu_offload: False
+    activation_checkpoint: True
+    
+    load_weights: True
+    load_optimizer: True
+    load_scaler: True
+    load_sheduler: True
+    num_epoch: 8
+    
+    skip_validation: False
+    update_learning_rate: False
+    
+    save_backup_weights: True
+    save_best_weights: True
+    save_metric_vars: True
+
+    learning_rate: 1.0e-03
+    weight_decay: 3.0e-06
+    
+    train_batch_size: 1
+    valid_batch_size: 1
+    
+    batches_per_epoch: 0
+    valid_batches_per_epoch: 0
+    stopping_patience: 999
+    
+    start_epoch: 0
+
+    reload_epoch: True
+    epochs: &epochs 70
+    
+    use_scheduler: True
+    scheduler: {'scheduler_type': 'cosine-annealing', 'T_max': *epochs,  'last_epoch': -1}
+
+    # Automatic Mixed Precision: False
+    amp: False
+    
+    # rescale loss as loss = loss / grad_accum_every
+    grad_accum_every: 1 
+    # gradient clipping
+    grad_max_norm: 32.0
+    
+    # number of workers
+    thread_workers: 4
+    valid_thread_workers: 0
+
+model:
+    type: "fuxi"
+    
+    frames: 2                # number of input states
+    image_height: 181        # number of latitude grids
+    image_width: 360        # number of longitude grids
+    levels: 18               # number of upper-air variable levels
+    channels: 4              # upper-air variable channels
+    surface_channels: 4      # surface variable channels
+    input_only_channels: 4   # dynamic forcing, forcing, static channels
+    output_only_channels: 8  # diagnostic variable channels
+
+    # patchify layer
+    patch_height: 4         # number of latitude grids in each 3D patch
+    patch_width: 4          # number of longitude grids in each 3D patch
+    frame_patch_size: 2     # number of input states in each 3D patch
+    
+    # hidden layers
+    dim: 1536               # dimension (default: 1536)
+    num_groups: 32          # number of groups (default: 32)
+    num_heads: 8            # number of heads (default: 8)
+    window_size: 7          # window size (default: 7)
+    depth: 48               # number of swin transformers (default: 48)
+    
+    # use spectral norm
+    use_spectral_norm: True
+
+    # use interpolation to match the output size
+    interp: False
+    
+    # map boundary padding
+    padding_conf:
+        activate: True
+        mode: earth
+        pad_lat: [21, 22]
+        pad_lon: [44, 44]
+    
+    post_conf: 
+        activate: False
+
+loss: 
+    # the main training loss
+    training_loss: "mse"
+    
+    # power loss (x), spectral_loss (x)
+    use_power_loss: False 
+    use_spectral_loss: False
+    
+    # use latitude weighting
+    use_latitude_weights: True
+    latitude_weights: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+
+    # turn-off variable weighting
+    use_variable_weights: True
+    variable_weights:
+        specific_total_water: [1.36577589e-06, 1.19436047e-05, 1.03276589e-04, 4.46097338e-04, 1.26807904e-03, 2.77554864e-03, 5.13295742e-03, 8.70407041e-03, 1.41016958e-02, 4.12818192e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02]
+        temperature: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        u_component_of_wind: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        v_component_of_wind: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        SP: 0.13
+        VAR_2T: 0.13
+        VAR_10U: 0.13
+        VAR_10V: 0.13
+        surface_latent_heat_flux: 0.065
+        surface_net_solar_radiation: 0.065
+        evaporation: 0.065
+        surface_net_thermal_radiation: 0.065
+        toa_incident_solar_radiation: 0.065
+        surface_sensible_heat_flux: 0.065
+        total_precipitation: 0.065
+        top_net_thermal_radiation: 0.065
+        top_net_solar_radiation: 0.065
+

FuXi-sigma-physics/checkpoints/checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:938f77598d0d82c647a2cd75b09bb9cc488af0eefd0524011840068e1fee85b0
+size 940

FuXi-sigma-physics/checkpoints/model_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:49513fd0779aa460c5e734b0db5191e6ad40c139dd78255a5915279e9facbd06
+size 5970491642

FuXi-sigma-physics/configs/model_multi.yml

@@ -0,0 +1,248 @@
+# --------------------------------------------------------------------------------------------------------------------- #
+# This yaml file implements 6 hourly FuXi on NSF NCAR HPCs (casper.ucar.edu and derecho.hpc.ucar.edu)
+# Technical details are available in:
+#
+#     Sha, Y., J. Schreck, W. Chapman, D. J. Gagne II, 2025: Investigating the contribution of terrain-following
+#     coordinates and conservation schemes in AI-driven precipitation forecasts. Submitted to:
+#     Geophysical Research Letters. pre-print: https://arxiv.org/abs/2503.00332
+#
+# --------------------------------------------------------------------------------------------------------------------- #
+
+save_loc: '/glade/work/ksha/CREDIT_runs/fuxi_mlevel_physics/'
+seed: 1000
+
+data:
+    # upper-air variables
+    variables: ['specific_total_water', 'temperature', 'u_component_of_wind','v_component_of_wind']
+    save_loc: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # surface variables
+    surface_variables: ['SP', 'VAR_2T', 'VAR_10U', 'VAR_10V']
+    save_loc_surface: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # dynamic forcing variables
+    dynamic_forcing_variables: ['toa_incident_solar_radiation', 'land_sea_CI_mask']
+    save_loc_dynamic_forcing: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # diagnostic variables
+    diagnostic_variables: ['evaporation', 'total_precipitation',
+                           'surface_net_solar_radiation',
+                           'surface_net_thermal_radiation',
+                           'surface_sensible_heat_flux',
+                           'surface_latent_heat_flux',
+                           'top_net_solar_radiation',
+                           'top_net_thermal_radiation']
+    save_loc_diagnostic: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # static variables
+    static_variables: ['z_norm', 'soil_type'] 
+    save_loc_static: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+
+    # physics file
+    save_loc_physics: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+    
+    # mean / std path
+    mean_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/mean_6h_1979_2019_conserve_1deg.nc'
+    std_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/std_residual_6h_1979_2019_conserve_1deg.nc'
+    
+    # train / validation split
+    train_years: [1979, 2019]
+    valid_years: [2019, 2020]
+    
+    # data workflow
+    scaler_type: 'std_new'
+     
+    history_len: 2
+    valid_history_len: 2
+    
+    forecast_len: 1 # <----------- forecast_len increases from 1 to 11 for 12 hr to 72 hr multi-steps
+    valid_forecast_len: 11
+
+    one_shot: False
+    
+    # 1 for hourly model
+    lead_time_periods: 6
+    
+    # do not use skip_period
+    skip_periods: null
+    
+    # compatible with the old 'std'
+    static_first: True
+    
+    sst_forcing:
+        activate: False
+
+trainer:
+    type: multi-step
+    
+    mode: fsdp
+    cpu_offload: False
+    activation_checkpoint: True
+    
+    load_weights: True
+    load_optimizer: False
+    load_scaler: False
+    load_scheduler: False
+    num_epoch: 1
+    
+    skip_validation: False
+    update_learning_rate: False
+    
+    save_backup_weights: True
+    save_best_weights: True
+    save_metric_vars: True
+
+    learning_rate: 3.0e-07
+    weight_decay: 3.0e-06
+    
+    train_batch_size: 1
+    valid_batch_size: 1
+    
+    batches_per_epoch: 0
+    valid_batches_per_epoch: 0
+    stopping_patience: 999
+    
+    start_epoch: 0
+
+    reload_epoch: True
+    epochs: &epochs 999
+    
+    use_scheduler: False
+
+    # Automatic Mixed Precision: False
+    amp: False
+    
+    # rescale loss as loss = loss / grad_accum_every
+    grad_accum_every: 1 
+    # gradient clipping
+    grad_max_norm: 999.0
+    
+    # number of workers
+    thread_workers: 4
+    valid_thread_workers: 0
+
+model:
+    type: "fuxi"
+    
+    frames: 2                # number of input states
+    image_height: 181        # number of latitude grids
+    image_width: 360        # number of longitude grids
+    levels: 18               # number of upper-air variable levels
+    channels: 4              # upper-air variable channels
+    surface_channels: 4      # surface variable channels
+    input_only_channels: 4   # dynamic forcing, forcing, static channels
+    output_only_channels: 8  # diagnostic variable channels
+
+    # patchify layer
+    patch_height: 4         # number of latitude grids in each 3D patch
+    patch_width: 4          # number of longitude grids in each 3D patch
+    frame_patch_size: 2     # number of input states in each 3D patch
+    
+    # hidden layers
+    dim: 1536               # dimension (default: 1536)
+    num_groups: 32          # number of groups (default: 32)
+    num_heads: 8            # number of heads (default: 8)
+    window_size: 7          # window size (default: 7)
+    depth: 48               # number of swin transformers (default: 48)
+    
+    # use spectral norm
+    use_spectral_norm: True
+
+    # use interpolation to match the output size
+    interp: False
+    
+    # map boundary padding
+    padding_conf:
+        activate: True
+        mode: earth
+        pad_lat: [21, 22]
+        pad_lon: [44, 44]
+    
+    post_conf: 
+        activate: True
+        
+        skebs:
+            activate: False
+        
+        tracer_fixer:
+            activate: True
+            denorm: True
+            tracer_name: ['specific_total_water', 'total_precipitation']
+            tracer_thres: [0, 0]
+            
+        global_mass_fixer:
+            activate: True
+            activate_outside_model: True
+            simple_demo: False
+            denorm: True
+            grid_type: 'sigma'
+            midpoint: True
+            fix_level_num: 7
+            lon_lat_level_name: ['lon2d', 'lat2d', 'coef_a', 'coef_b']
+            surface_pressure_name: ['SP']
+            specific_total_water_name: ['specific_total_water']
+            
+        global_water_fixer:
+            activate: True
+            activate_outside_model: True
+            simple_demo: False
+            denorm: True
+            grid_type: 'sigma'
+            midpoint: True
+            lon_lat_level_name: ['lon2d', 'lat2d', 'coef_a', 'coef_b']
+            surface_pressure_name: ['SP']
+            specific_total_water_name: ['specific_total_water']
+            precipitation_name: ['total_precipitation']
+            evaporation_name: ['evaporation']
+            
+        global_energy_fixer:
+            activate: True
+            activate_outside_model: True
+            simple_demo: False
+            denorm: True
+            grid_type: 'sigma'
+            midpoint: True
+            lon_lat_level_name: ['lon2d', 'lat2d', 'coef_a', 'coef_b']
+            surface_pressure_name: ['SP']
+            air_temperature_name: ['temperature']
+            specific_total_water_name: ['specific_total_water']
+            u_wind_name: ['u_component_of_wind']
+            v_wind_name: ['v_component_of_wind']
+            surface_geopotential_name: ['geopotential_at_surface']
+            TOA_net_radiation_flux_name: ['top_net_solar_radiation', 'top_net_thermal_radiation']
+            surface_net_radiation_flux_name: ['surface_net_solar_radiation', 'surface_net_thermal_radiation']
+            surface_energy_flux_name: ['surface_sensible_heat_flux', 'surface_latent_heat_flux',]
+
+loss: 
+    # the main training loss
+    training_loss: "mse"
+    
+    # power loss (x), spectral_loss (x)
+    use_power_loss: False 
+    use_spectral_loss: False
+    
+    # use latitude weighting
+    use_latitude_weights: True
+    latitude_weights: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+
+    # turn-off variable weighting
+    use_variable_weights: True
+    variable_weights:
+        specific_total_water: [1.36577589e-06, 1.19436047e-05, 1.03276589e-04, 4.46097338e-04, 1.26807904e-03, 2.77554864e-03, 5.13295742e-03, 8.70407041e-03, 1.41016958e-02, 4.12818192e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02]
+        temperature: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        u_component_of_wind: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        v_component_of_wind: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        SP: 0.13
+        VAR_2T: 0.13
+        VAR_10U: 0.13
+        VAR_10V: 0.13
+        surface_latent_heat_flux: 0.065
+        surface_net_solar_radiation: 0.065
+        evaporation: 0.065
+        surface_net_thermal_radiation: 0.065
+        toa_incident_solar_radiation: 0.065
+        surface_sensible_heat_flux: 0.065
+        total_precipitation: 0.065
+        top_net_thermal_radiation: 0.065
+        top_net_solar_radiation: 0.065
+

FuXi-sigma-physics/configs/model_predict.yml

@@ -0,0 +1,243 @@
+# --------------------------------------------------------------------------------------------------------------------- #
+# This yaml file implements 6 hourly FuXi on NSF NCAR HPCs (casper.ucar.edu and derecho.hpc.ucar.edu)
+# Technical details are available in:
+#
+#     Sha, Y., J. Schreck, W. Chapman, D. J. Gagne II, 2025: Investigating the contribution of terrain-following
+#     coordinates and conservation schemes in AI-driven precipitation forecasts. Submitted to:
+#     Geophysical Research Letters. pre-print: https://arxiv.org/abs/2503.00332
+#
+# --------------------------------------------------------------------------------------------------------------------- #
+
+save_loc: '/glade/work/ksha/CREDIT_runs/fuxi_mlevel_physics/'
+seed: 1000
+
+data:
+    # upper-air variables
+    variables: ['specific_total_water', 'temperature', 'u_component_of_wind','v_component_of_wind']
+    save_loc: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # surface variables
+    surface_variables: ['SP', 'VAR_2T', 'VAR_10U', 'VAR_10V']
+    save_loc_surface: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # dynamic forcing variables
+    dynamic_forcing_variables: ['toa_incident_solar_radiation', 'land_sea_CI_mask']
+    save_loc_dynamic_forcing: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # diagnostic variables
+    diagnostic_variables: ['evaporation', 'total_precipitation',
+                           'surface_net_solar_radiation',
+                           'surface_net_thermal_radiation',
+                           'surface_sensible_heat_flux',
+                           'surface_latent_heat_flux',
+                           'top_net_solar_radiation',
+                           'top_net_thermal_radiation']
+    save_loc_diagnostic: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # static variables
+    static_variables: ['z_norm', 'soil_type'] 
+    save_loc_static: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+
+    # physics file
+    save_loc_physics: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+    
+    # mean / std path
+    mean_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/mean_6h_1979_2019_conserve_1deg.nc'
+    std_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/std_residual_6h_1979_2019_conserve_1deg.nc'
+    
+    # train / validation split
+    train_years: [1979, 2019]
+    valid_years: [2019, 2020]
+    
+    # data workflow
+    scaler_type: 'std_new'
+     
+    history_len: 2
+    valid_history_len: 2
+    
+    forecast_len: 0
+    valid_forecast_len: 0
+
+    one_shot: False
+    
+    # 1 for hourly model
+    lead_time_periods: 6
+    
+    # do not use skip_period
+    skip_periods: null
+    
+    # compatible with the old 'std'
+    static_first: True
+    
+    sst_forcing:
+        activate: False
+
+trainer:
+    type: standard
+    
+    mode: fsdp
+    cpu_offload: False
+    activation_checkpoint: True
+    
+    load_weights: True
+    load_optimizer: True
+    load_scaler: True
+    load_scheduler: True
+    num_epoch: 8
+    
+    skip_validation: False
+    update_learning_rate: False
+    
+    save_backup_weights: True
+    save_best_weights: True
+    save_metric_vars: True
+
+    learning_rate: 1.0e-03
+    weight_decay: 3.0e-06
+    
+    train_batch_size: 1
+    valid_batch_size: 1
+    
+    batches_per_epoch: 0
+    valid_batches_per_epoch: 0
+    stopping_patience: 999
+    
+    start_epoch: 0
+
+    reload_epoch: True
+    epochs: &epochs 200
+    
+    use_scheduler: True
+    scheduler: {'scheduler_type': 'cosine-annealing', 'T_max': *epochs,  'last_epoch': -1}
+
+    # Automatic Mixed Precision: False
+    amp: False
+    
+    # rescale loss as loss = loss / grad_accum_every
+    grad_accum_every: 1 
+    # gradient clipping
+    grad_max_norm: 32.0
+    
+    # number of workers
+    thread_workers: 4
+    valid_thread_workers: 0
+
+model:
+    type: "fuxi"
+    
+    frames: 2                # number of input states
+    image_height: 181        # number of latitude grids
+    image_width: 360        # number of longitude grids
+    levels: 18               # number of upper-air variable levels
+    channels: 4              # upper-air variable channels
+    surface_channels: 4      # surface variable channels
+    input_only_channels: 4   # dynamic forcing, forcing, static channels
+    output_only_channels: 8  # diagnostic variable channels
+
+    # patchify layer
+    patch_height: 4         # number of latitude grids in each 3D patch
+    patch_width: 4          # number of longitude grids in each 3D patch
+    frame_patch_size: 2     # number of input states in each 3D patch
+    
+    # hidden layers
+    dim: 1536               # dimension (default: 1536)
+    num_groups: 32          # number of groups (default: 32)
+    num_heads: 8            # number of heads (default: 8)
+    window_size: 7          # window size (default: 7)
+    depth: 48               # number of swin transformers (default: 48)
+    
+    # use spectral norm
+    use_spectral_norm: True
+
+    # use interpolation to match the output size
+    interp: False
+    
+    # map boundary padding
+    padding_conf:
+        activate: True
+        mode: earth
+        pad_lat: [21, 22]
+        pad_lon: [44, 44]
+    
+    post_conf: 
+        activate: True
+        
+        skebs:
+            activate: False
+        
+        tracer_fixer:
+            activate: True
+            denorm: True
+            tracer_name: ['specific_total_water', 'total_precipitation']
+            tracer_thres: [0, 0]
+            
+        global_mass_fixer:
+            activate: True
+            activate_outside_model: True
+            simple_demo: False
+            denorm: True
+            grid_type: 'sigma'
+            midpoint: True
+            fix_level_num: 7
+            lon_lat_level_name: ['lon2d', 'lat2d', 'coef_a', 'coef_b']
+            surface_pressure_name: ['SP']
+            specific_total_water_name: ['specific_total_water']
+            
+        global_water_fixer:
+            activate: True
+            activate_outside_model: True
+            simple_demo: False
+            denorm: True
+            grid_type: 'sigma'
+            midpoint: True
+            lon_lat_level_name: ['lon2d', 'lat2d', 'coef_a', 'coef_b']
+            surface_pressure_name: ['SP']
+            specific_total_water_name: ['specific_total_water']
+            precipitation_name: ['total_precipitation']
+            evaporation_name: ['evaporation']
+            
+        global_energy_fixer:
+            activate: True
+            activate_outside_model: True
+            simple_demo: False
+            denorm: True
+            grid_type: 'sigma'
+            midpoint: True
+            lon_lat_level_name: ['lon2d', 'lat2d', 'coef_a', 'coef_b']
+            surface_pressure_name: ['SP']
+            air_temperature_name: ['temperature']
+            specific_total_water_name: ['specific_total_water']
+            u_wind_name: ['u_component_of_wind']
+            v_wind_name: ['v_component_of_wind']
+            surface_geopotential_name: ['geopotential_at_surface']
+            TOA_net_radiation_flux_name: ['top_net_solar_radiation', 'top_net_thermal_radiation']
+            surface_net_radiation_flux_name: ['surface_net_solar_radiation', 'surface_net_thermal_radiation']
+            surface_energy_flux_name: ['surface_sensible_heat_flux', 'surface_latent_heat_flux',]
+
+loss: 
+    # the main training loss
+    training_loss: "mse"
+    
+    # power loss (x), spectral_loss (x)
+    use_power_loss: False 
+    use_spectral_loss: False
+    
+    # use latitude weighting
+    use_latitude_weights: True
+    latitude_weights: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+   
+predict:
+    forecasts:
+        type: "custom"       # keep it as "custom"
+        start_year: 2020     # year of the first initialization (where rollout will start)
+        start_month: 1       # month of the first initialization
+        start_day: 1         # day of the first initialization
+        start_hours: [0, 12] # hour-of-day for each initialization, 0 for 00Z, 12 for 12Z
+        duration: 768        # number of days to initialize, starting from the (year, mon, day) above
+                             # duration should be divisible by the number of GPUs
+                             # (e.g., duration: 384 for 365-day rollout using 32 GPUs)
+        days: 15             # forecast lead time as days (1 means 24-hour forecast)
+
+    save_forecast: '/glade/derecho/scratch/ksha/CREDIT/RAW_OUTPUT/fuxi_mlevel_physics2/'
+    use_laplace_filter: False
+

FuXi-sigma-physics/configs/model_single.yml

@@ -0,0 +1,249 @@
+# --------------------------------------------------------------------------------------------------------------------- #
+# This yaml file implements 6 hourly FuXi on NSF NCAR HPCs (casper.ucar.edu and derecho.hpc.ucar.edu) 
+# Technical details are available in:
+#
+#     Sha, Y., J. Schreck, W. Chapman, D. J. Gagne II, 2025: Investigating the contribution of terrain-following 
+#     coordinates and conservation schemes in AI-driven precipitation forecasts. Submitted to: 
+#     Geophysical Research Letters. pre-print: https://arxiv.org/abs/2503.00332
+#
+# --------------------------------------------------------------------------------------------------------------------- #
+
+save_loc: '/glade/work/ksha/CREDIT_runs/fuxi_mlevel_physics/'
+seed: 1000
+
+data:
+    # upper-air variables
+    variables: ['specific_total_water', 'temperature', 'u_component_of_wind','v_component_of_wind']
+    save_loc: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # surface variables
+    surface_variables: ['SP', 'VAR_2T', 'VAR_10U', 'VAR_10V']
+    save_loc_surface: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # dynamic forcing variables
+    dynamic_forcing_variables: ['toa_incident_solar_radiation', 'land_sea_CI_mask']
+    save_loc_dynamic_forcing: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # diagnostic variables
+    diagnostic_variables: ['evaporation', 'total_precipitation',
+                           'surface_net_solar_radiation',
+                           'surface_net_thermal_radiation',
+                           'surface_sensible_heat_flux',
+                           'surface_latent_heat_flux',
+                           'top_net_solar_radiation',
+                           'top_net_thermal_radiation']
+    save_loc_diagnostic: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/all_in_one/ERA5_mlevel_1deg_6h_subset_*_conserve.zarr'
+    
+    # static variables
+    static_variables: ['z_norm', 'soil_type'] 
+    save_loc_static: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+
+    # physics file
+    save_loc_physics: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+    
+    # mean / std path
+    mean_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/mean_6h_1979_2019_conserve_1deg.nc'
+    std_path: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/mean_std/std_residual_6h_1979_2019_conserve_1deg.nc'
+    
+    # train / validation split
+    train_years: [1979, 2019]
+    valid_years: [2019, 2020]
+    
+    # data workflow
+    scaler_type: 'std_new'
+     
+    history_len: 2
+    valid_history_len: 2
+    
+    forecast_len: 0
+    valid_forecast_len: 0
+
+    one_shot: False
+    
+    # 1 for hourly model
+    lead_time_periods: 6
+    
+    # do not use skip_period
+    skip_periods: null
+    
+    # compatible with the old 'std'
+    static_first: True
+    
+    sst_forcing:
+        activate: False
+
+trainer:
+    type: standard
+    
+    mode: fsdp
+    cpu_offload: False
+    activation_checkpoint: True
+    
+    load_weights: True
+    load_optimizer: False
+    load_scaler: False
+    load_scheduler: False
+    num_epoch: 8
+    
+    skip_validation: False
+    update_learning_rate: False
+    
+    save_backup_weights: True
+    save_best_weights: True
+    save_metric_vars: True
+
+    learning_rate: 1.0e-03
+    weight_decay: 3.0e-06
+    
+    train_batch_size: 1
+    valid_batch_size: 1
+    
+    batches_per_epoch: 0
+    valid_batches_per_epoch: 0
+    stopping_patience: 999
+    
+    start_epoch: 0
+
+    reload_epoch: True
+    epochs: &epochs 70
+    
+    use_scheduler: True
+    scheduler: {'scheduler_type': 'cosine-annealing', 'T_max': *epochs,  'last_epoch': -1}
+
+    # Automatic Mixed Precision: False
+    amp: False
+    
+    # rescale loss as loss = loss / grad_accum_every
+    grad_accum_every: 1 
+    # gradient clipping
+    grad_max_norm: 32.0
+    
+    # number of workers
+    thread_workers: 4
+    valid_thread_workers: 0
+
+model:
+    type: "fuxi"
+    
+    frames: 2                # number of input states
+    image_height: 181        # number of latitude grids
+    image_width: 360        # number of longitude grids
+    levels: 18               # number of upper-air variable levels
+    channels: 4              # upper-air variable channels
+    surface_channels: 4      # surface variable channels
+    input_only_channels: 4   # dynamic forcing, forcing, static channels
+    output_only_channels: 8  # diagnostic variable channels
+
+    # patchify layer
+    patch_height: 4         # number of latitude grids in each 3D patch
+    patch_width: 4          # number of longitude grids in each 3D patch
+    frame_patch_size: 2     # number of input states in each 3D patch
+    
+    # hidden layers
+    dim: 1536               # dimension (default: 1536)
+    num_groups: 32          # number of groups (default: 32)
+    num_heads: 8            # number of heads (default: 8)
+    window_size: 7          # window size (default: 7)
+    depth: 48               # number of swin transformers (default: 48)
+    
+    # use spectral norm
+    use_spectral_norm: True
+
+    # use interpolation to match the output size
+    interp: False
+    
+    # map boundary padding
+    padding_conf:
+        activate: True
+        mode: earth
+        pad_lat: [21, 22]
+        pad_lon: [44, 44]
+    
+    post_conf: 
+        activate: True
+        
+        skebs:
+            activate: False
+        
+        tracer_fixer:
+            activate: True
+            denorm: True
+            tracer_name: ['specific_total_water', 'total_precipitation']
+            tracer_thres: [0, 0]
+            
+        global_mass_fixer:
+            activate: True
+            activate_outside_model: False
+            simple_demo: False
+            denorm: True
+            grid_type: 'sigma'
+            midpoint: True
+            fix_level_num: 7
+            lon_lat_level_name: ['lon2d', 'lat2d', 'coef_a', 'coef_b']
+            surface_pressure_name: ['SP']
+            specific_total_water_name: ['specific_total_water']
+            
+        global_water_fixer:
+            activate: True
+            activate_outside_model: False
+            simple_demo: False
+            denorm: True
+            grid_type: 'sigma'
+            midpoint: True
+            lon_lat_level_name: ['lon2d', 'lat2d', 'coef_a', 'coef_b']
+            surface_pressure_name: ['SP']
+            specific_total_water_name: ['specific_total_water']
+            precipitation_name: ['total_precipitation']
+            evaporation_name: ['evaporation']
+            
+        global_energy_fixer:
+            activate: True
+            activate_outside_model: False
+            simple_demo: False
+            denorm: True
+            grid_type: 'sigma'
+            midpoint: True
+            lon_lat_level_name: ['lon2d', 'lat2d', 'coef_a', 'coef_b']
+            surface_pressure_name: ['SP']
+            air_temperature_name: ['temperature']
+            specific_total_water_name: ['specific_total_water']
+            u_wind_name: ['u_component_of_wind']
+            v_wind_name: ['v_component_of_wind']
+            surface_geopotential_name: ['geopotential_at_surface']
+            TOA_net_radiation_flux_name: ['top_net_solar_radiation', 'top_net_thermal_radiation']
+            surface_net_radiation_flux_name: ['surface_net_solar_radiation', 'surface_net_thermal_radiation']
+            surface_energy_flux_name: ['surface_sensible_heat_flux', 'surface_latent_heat_flux',]
+
+loss: 
+    # the main training loss
+    training_loss: "mse"
+    
+    # power loss (x), spectral_loss (x)
+    use_power_loss: False 
+    use_spectral_loss: False
+    
+    # use latitude weighting
+    use_latitude_weights: True
+    latitude_weights: '/glade/derecho/scratch/ksha/CREDIT_data/ERA5_mlevel_1deg/static/ERA5_mlevel_1deg_static_subset.zarr'
+    
+    # turn-off variable weighting
+    use_variable_weights: True
+    variable_weights:
+        specific_total_water: [1.36577589e-06, 1.19436047e-05, 1.03276589e-04, 4.46097338e-04, 1.26807904e-03, 2.77554864e-03, 5.13295742e-03, 8.70407041e-03, 1.41016958e-02, 4.12818192e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02, 6.84619425e-02]
+        temperature: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        u_component_of_wind: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        v_component_of_wind: [1.36577589e-06, 2.38872094e-05, 2.06553177e-04, 8.92194675e-04, 2.53615808e-03, 5.55109728e-03, 1.02659148e-02, 1.74081408e-02, 2.82033917e-02, 8.25636383e-02, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01, 1.36923885e-01]
+        SP: 0.13
+        VAR_2T: 0.13
+        VAR_10U: 0.13
+        VAR_10V: 0.13
+        surface_latent_heat_flux: 0.065
+        surface_net_solar_radiation: 0.065
+        evaporation: 0.065
+        surface_net_thermal_radiation: 0.065
+        toa_incident_solar_radiation: 0.065
+        surface_sensible_heat_flux: 0.065
+        total_precipitation: 0.065
+        top_net_thermal_radiation: 0.065
+        top_net_solar_radiation: 0.065
+