conf/config.yaml

project:
  name: neuralgcm_develop
  task: earth_system_forecasting
  seed: 20260904

paths:
  project_root: .
  # Upstream source is supplied by the external neuralgcm package.
  official_source_dir: null
  virtual_era5_dir: data
  checkpoint_dir: data/checkpoint
  result_dir: results
  metadata_dir: metadata

model:
  # Native NeuralGCM pressure-level input contract.
  variant: weather_forecast
  grid_degrees: 0.7
  # Gaussian grid shape in [longitude, latitude] order (the model API reports
  # the same grid as (latitude, longitude) when printing sizes).
  grid_shape: [512, 256]
  profiles:
    weather_forecast:
      description: "未来2至15天天气预报"
      official_reference: models_v1_deterministic_0_7_deg.pkl
      grid_degrees: 0.7
      grid_shape: [512, 256]
    climate_scale:
      description: "气候尺度模拟"
      official_reference: models_v1_deterministic_1_4_deg.pkl
      grid_degrees: 1.4
      grid_shape: [256, 128]
    forecast_2_8_deg:
      description: "2.8度天气预报"
      official_reference: models_v1_deterministic_2_8_deg.pkl
      grid_degrees: 2.8
      grid_shape: [128, 64]
    stochastic_1_4_deg:
      description: "1.4度随机预报"
      official_reference: models_v1_stochastic_1_4_deg.pkl
      grid_degrees: 1.4
      grid_shape: [256, 128]
  pressure_levels_hpa: [1, 2, 3, 5, 7, 10, 20, 30, 50, 70, 100, 125, 150, 175, 200, 225, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700, 750, 775, 800, 825, 850, 875, 900, 925, 950, 975, 1000]
  input_variables: [geopotential, specific_humidity, temperature, u_component_of_wind, v_component_of_wind]
  optional_input_variables: [specific_cloud_ice_water_content, specific_cloud_liquid_water_content]
  forcing_variables: [sea_ice_cover, sea_surface_temperature]
  official_checkpoint: null
  load_pretrained: false

data:
  dataset_class: onescience.datapipes.climate.ERA5Dataset
  data_dir: data
  # Optional auxiliary static fields generated by fake_data.py or supplied by
  # a real ERA5 preprocessing job.  Dynamic channels remain in data/*.h5.
  static_file: data/static.nc
  # Exact Gaussian-grid static fields extracted from the four official
  # checkpoints by scripts/prepare_static_data.py. These take precedence over
  # the source-grid synthetic fallback above.
  static_files:
    weather_forecast: data/static/weather_forecast.nc
    climate_scale: data/static/climate_scale.nc
    forecast_2_8_deg: data/static/forecast_2_8_deg.nc
    stochastic_1_4_deg: data/static/stochastic_1_4_deg.nc
  field_key: fields
  time_step_hours: 6
  input_steps: 1
  # Official training consumes a time trajectory. Increase for production
  # rollouts; 1 is retained for the minimal data validation command.
  output_steps: 1
  normalize: false
  batch_size: 1
  num_workers: 0
  train_years: [1999]
  val_years: [2000]
  test_years: [2001]
  virtual:
    # Memory-conscious default: one initial frame + eight future 6-hour
    # frames. Use --forecast-steps 60 for the full official 15-day horizon.
    timesteps_per_year: 9
    forecast_steps: 8
    forecast_horizon_days: 2
    height: 721
    width: 1440
    seed: 20260904
  # Exact flattened fields order used by fake_data.py and ERA5Dataset.
  channel_order: &channel_order
    - geopotential_1
    - geopotential_2
    - geopotential_3
    - geopotential_5
    - geopotential_7
    - geopotential_10
    - geopotential_20
    - geopotential_30
    - geopotential_50
    - geopotential_70
    - geopotential_100
    - geopotential_125
    - geopotential_150
    - geopotential_175
    - geopotential_200
    - geopotential_225
    - geopotential_250
    - geopotential_300
    - geopotential_350
    - geopotential_400
    - geopotential_450
    - geopotential_500
    - geopotential_550
    - geopotential_600
    - geopotential_650
    - geopotential_700
    - geopotential_750
    - geopotential_775
    - geopotential_800
    - geopotential_825
    - geopotential_850
    - geopotential_875
    - geopotential_900
    - geopotential_925
    - geopotential_950
    - geopotential_975
    - geopotential_1000
    - specific_humidity_1
    - specific_humidity_2
    - specific_humidity_3
    - specific_humidity_5
    - specific_humidity_7
    - specific_humidity_10
    - specific_humidity_20
    - specific_humidity_30
    - specific_humidity_50
    - specific_humidity_70
    - specific_humidity_100
    - specific_humidity_125
    - specific_humidity_150
    - specific_humidity_175
    - specific_humidity_200
    - specific_humidity_225
    - specific_humidity_250
    - specific_humidity_300
    - specific_humidity_350
    - specific_humidity_400
    - specific_humidity_450
    - specific_humidity_500
    - specific_humidity_550
    - specific_humidity_600
    - specific_humidity_650
    - specific_humidity_700
    - specific_humidity_750
    - specific_humidity_775
    - specific_humidity_800
    - specific_humidity_825
    - specific_humidity_850
    - specific_humidity_875
    - specific_humidity_900
    - specific_humidity_925
    - specific_humidity_950
    - specific_humidity_975
    - specific_humidity_1000
    - temperature_1
    - temperature_2
    - temperature_3
    - temperature_5
    - temperature_7
    - temperature_10
    - temperature_20
    - temperature_30
    - temperature_50
    - temperature_70
    - temperature_100
    - temperature_125
    - temperature_150
    - temperature_175
    - temperature_200
    - temperature_225
    - temperature_250
    - temperature_300
    - temperature_350
    - temperature_400
    - temperature_450
    - temperature_500
    - temperature_550
    - temperature_600
    - temperature_650
    - temperature_700
    - temperature_750
    - temperature_775
    - temperature_800
    - temperature_825
    - temperature_850
    - temperature_875
    - temperature_900
    - temperature_925
    - temperature_950
    - temperature_975
    - temperature_1000
    - u_component_of_wind_1
    - u_component_of_wind_2
    - u_component_of_wind_3
    - u_component_of_wind_5
    - u_component_of_wind_7
    - u_component_of_wind_10
    - u_component_of_wind_20
    - u_component_of_wind_30
    - u_component_of_wind_50
    - u_component_of_wind_70
    - u_component_of_wind_100
    - u_component_of_wind_125
    - u_component_of_wind_150
    - u_component_of_wind_175
    - u_component_of_wind_200
    - u_component_of_wind_225
    - u_component_of_wind_250
    - u_component_of_wind_300
    - u_component_of_wind_350
    - u_component_of_wind_400
    - u_component_of_wind_450
    - u_component_of_wind_500
    - u_component_of_wind_550
    - u_component_of_wind_600
    - u_component_of_wind_650
    - u_component_of_wind_700
    - u_component_of_wind_750
    - u_component_of_wind_775
    - u_component_of_wind_800
    - u_component_of_wind_825
    - u_component_of_wind_850
    - u_component_of_wind_875
    - u_component_of_wind_900
    - u_component_of_wind_925
    - u_component_of_wind_950
    - u_component_of_wind_975
    - u_component_of_wind_1000
    - v_component_of_wind_1
    - v_component_of_wind_2
    - v_component_of_wind_3
    - v_component_of_wind_5
    - v_component_of_wind_7
    - v_component_of_wind_10
    - v_component_of_wind_20
    - v_component_of_wind_30
    - v_component_of_wind_50
    - v_component_of_wind_70
    - v_component_of_wind_100
    - v_component_of_wind_125
    - v_component_of_wind_150
    - v_component_of_wind_175
    - v_component_of_wind_200
    - v_component_of_wind_225
    - v_component_of_wind_250
    - v_component_of_wind_300
    - v_component_of_wind_350
    - v_component_of_wind_400
    - v_component_of_wind_450
    - v_component_of_wind_500
    - v_component_of_wind_550
    - v_component_of_wind_600
    - v_component_of_wind_650
    - v_component_of_wind_700
    - v_component_of_wind_750
    - v_component_of_wind_775
    - v_component_of_wind_800
    - v_component_of_wind_825
    - v_component_of_wind_850
    - v_component_of_wind_875
    - v_component_of_wind_900
    - v_component_of_wind_925
    - v_component_of_wind_950
    - v_component_of_wind_975
    - v_component_of_wind_1000
    - specific_cloud_ice_water_content_1
    - specific_cloud_ice_water_content_2
    - specific_cloud_ice_water_content_3
    - specific_cloud_ice_water_content_5
    - specific_cloud_ice_water_content_7
    - specific_cloud_ice_water_content_10
    - specific_cloud_ice_water_content_20
    - specific_cloud_ice_water_content_30
    - specific_cloud_ice_water_content_50
    - specific_cloud_ice_water_content_70
    - specific_cloud_ice_water_content_100
    - specific_cloud_ice_water_content_125
    - specific_cloud_ice_water_content_150
    - specific_cloud_ice_water_content_175
    - specific_cloud_ice_water_content_200
    - specific_cloud_ice_water_content_225
    - specific_cloud_ice_water_content_250
    - specific_cloud_ice_water_content_300
    - specific_cloud_ice_water_content_350
    - specific_cloud_ice_water_content_400
    - specific_cloud_ice_water_content_450
    - specific_cloud_ice_water_content_500
    - specific_cloud_ice_water_content_550
    - specific_cloud_ice_water_content_600
    - specific_cloud_ice_water_content_650
    - specific_cloud_ice_water_content_700
    - specific_cloud_ice_water_content_750
    - specific_cloud_ice_water_content_775
    - specific_cloud_ice_water_content_800
    - specific_cloud_ice_water_content_825
    - specific_cloud_ice_water_content_850
    - specific_cloud_ice_water_content_875
    - specific_cloud_ice_water_content_900
    - specific_cloud_ice_water_content_925
    - specific_cloud_ice_water_content_950
    - specific_cloud_ice_water_content_975
    - specific_cloud_ice_water_content_1000
    - specific_cloud_liquid_water_content_1
    - specific_cloud_liquid_water_content_2
    - specific_cloud_liquid_water_content_3
    - specific_cloud_liquid_water_content_5
    - specific_cloud_liquid_water_content_7
    - specific_cloud_liquid_water_content_10
    - specific_cloud_liquid_water_content_20
    - specific_cloud_liquid_water_content_30
    - specific_cloud_liquid_water_content_50
    - specific_cloud_liquid_water_content_70
    - specific_cloud_liquid_water_content_100
    - specific_cloud_liquid_water_content_125
    - specific_cloud_liquid_water_content_150
    - specific_cloud_liquid_water_content_175
    - specific_cloud_liquid_water_content_200
    - specific_cloud_liquid_water_content_225
    - specific_cloud_liquid_water_content_250
    - specific_cloud_liquid_water_content_300
    - specific_cloud_liquid_water_content_350
    - specific_cloud_liquid_water_content_400
    - specific_cloud_liquid_water_content_450
    - specific_cloud_liquid_water_content_500
    - specific_cloud_liquid_water_content_550
    - specific_cloud_liquid_water_content_600
    - specific_cloud_liquid_water_content_650
    - specific_cloud_liquid_water_content_700
    - specific_cloud_liquid_water_content_750
    - specific_cloud_liquid_water_content_775
    - specific_cloud_liquid_water_content_800
    - specific_cloud_liquid_water_content_825
    - specific_cloud_liquid_water_content_850
    - specific_cloud_liquid_water_content_875
    - specific_cloud_liquid_water_content_900
    - specific_cloud_liquid_water_content_925
    - specific_cloud_liquid_water_content_950
    - specific_cloud_liquid_water_content_975
    - specific_cloud_liquid_water_content_1000
    - sea_ice_cover
    - sea_surface_temperature

training:
  mode: weather_forecast
  max_steps: 3
  trajectory_length: 2
  # Global batch size. For --devices N it is rounded up to a multiple of N;
  # each replica then receives distinct trajectories.
  samples_per_step: 1
  # Number of local JAX devices for optional synchronous data parallelism.
  devices: 1
  shuffle: true
  drop_last: true
  # OneScience ERA5Dataset samples are prefetched on host threads while the
  # current DCU step runs. Keep the queue shallow for full 721x1440 fields.
  data_num_workers: 2
  prefetch_batches: 1
  # Full params/EMA/optimizer/reader state is always saved on clean exit. Set a
  # positive interval for periodic resumable checkpoints during long runs.
  checkpoint_interval: 0
  learning_rate: 0.0001
  optimizer:
    name: adam
    schedule: constant
    b1: 0.9
    b2: 0.95
    eps: 1.0e-6
    # Optional piecewise constant schedule. Empty boundaries use base LR.
    rates: []
    boundaries: []
  # Public Experiment tracks an EMA for evaluation/checkpointing. Set to 0 to
  # disable; otherwise this is the effective average window in optimizer steps.
  ema_num_steps: 1000
  rollout_schedule: []
  # Public NeuralGCM uses transformed trajectory losses. The private job loss
  # bindings and complete normalization tables are unavailable, so every
  # published coefficient and every auditable fallback remain explicit here.
  gradient_clip_norm: 1.0
  loss:
    backend: official
    # Supplementary G.4 deterministic objective coefficients:
    # 20*data MSE + 0.1*data spectrum MSE + 1*model MSE
    # + 0.1*model spectrum MSE + 2*batch spectral bias MSE.
    data_weight: 20.0
    data_spectrum_weight: 0.1
    model_weight: 1.0
    model_spectrum_weight: 0.1
    bias_weight: 2.0
    accuracy_time_scale_hours: 24.0
    spectral_time_scale_hours: 40.0
    spectral_cutoff_by_mode:
      weather_forecast: 120
      climate_scale: 80
      forecast_2_8_deg: 42
    # Optional exact PerVariableRescaling weights. Each value multiplies the
    # squared error. When null, factor/scale below multiplies the error.
    variable_weights: null
    # The paper uses ERA5 24-hour difference standard deviations, but does not
    # publish the complete numerical tables. These auditable fallbacks keep
    # physical variables balanced; replace them with statistics calculated
    # from the exact ERA5 training vintage for a precision reproduction.
    time_rescaling: legacy
    spectral_weight: 0.0
    variable_scales:
      z: 10000.0
      t: 30.0
      u: 30.0
      v: 30.0
      specific_humidity: 0.01
      specific_cloud_ice_water_content: 1.0e-5
      specific_cloud_liquid_water_content: 2.0e-5
      divergence: 0.1
      vorticity: 0.1
      log_surface_pressure: 0.1
      default: 1.0
    # Additional balancing factors stated explicitly in Supplementary G.3.
    variable_factors:
      z: 2.0
      specific_humidity: 0.66
      log_surface_pressure: 5.0
      specific_cloud_ice_water_content: 0.05
      specific_cloud_liquid_water_content: 0.05
      default: 1.0
    # Order 12 is exact. Absolute half-power cutoffs below are digitized from
    # Supplementary Fig. 8 because the underlying numeric table was not
    # released. Interpolation is performed at the configured output times.
    predictability_filter:
      enabled: true
      order: 12
      lead_hours: [0, 6, 12, 24, 36, 48, 60, 72]
      cutoffs:
        temperature: [80, 120, 120, 95, 45, 35, 30, 25]
        wind: [80, 120, 115, 82, 48, 36, 29, 24]
        moisture: [80, 120, 110, 52, 34, 28, 24, 21]
        divergence: [80, 120, 105, 43, 24, 19, 16, 14]
        default: [80, 120, 115, 82, 48, 36, 29, 24]
    # Optional multiplicative weights for pressure levels, ordered as the
    # configured ERA5 pressure-level list. Empty means uniform weighting.
    level_weights: []
  # Long-run reproduction settings inferred from the public paper description
  # and released training pseudocode. The paper's private job bindings are not
  # available, so these are explicit project settings rather than exact claims.
  # They are enabled only by --paper-defaults; CLI values remain highest priority.
  profiles:
    weather_forecast:
      max_steps: 25000
      learning_rate: 0.001
      optimizer: &paper_optimizer
        schedule: neuralgcm
        warmup_steps: 2000
        decay_start: 15000
        decay_steps: 10000
        decay_rate: 0.5
      rollout_schedule:
        - {trajectory_length: 2, until_step: 0}      # 6 h
        - {trajectory_length: 3, until_step: 500}    # 12 h
        - {trajectory_length: 4, until_step: 2000}   # 18 h
        - {trajectory_length: 5, until_step: 4500}   # 24 h
        - {trajectory_length: 7, until_step: 8000}   # 36 h
        - {trajectory_length: 9, until_step: 12500}  # 48 h
        - {trajectory_length: 11, until_step: 18000} # 60 h
    climate_scale:
      max_steps: 26000
      learning_rate: 0.002
      optimizer: *paper_optimizer
      rollout_schedule: &coarse_rollout_schedule
        - {trajectory_length: 3, until_step: 0}       # 12 h
        - {trajectory_length: 5, until_step: 2000}    # 24 h
        - {trajectory_length: 7, until_step: 5656}    # 36 h
        - {trajectory_length: 9, until_step: 10392}   # 48 h
        - {trajectory_length: 11, until_step: 16000}  # 60 h
        - {trajectory_length: 13, until_step: 22360}  # 72 h
    forecast_2_8_deg:
      max_steps: 38000
      learning_rate: 0.002
      optimizer: *paper_optimizer
      rollout_schedule: *coarse_rollout_schedule
    stochastic_1_4_deg:
      max_steps: 43000
      learning_rate: 0.001
      ensemble_size: 2
      optimizer: *paper_optimizer
      rollout_schedule:
        - {trajectory_length: 2, until_step: 0}       # 6 h
        - {trajectory_length: 3, until_step: 500}     # 12 h
        - {trajectory_length: 4, until_step: 2000}    # 18 h
        - {trajectory_length: 5, until_step: 4500}    # 24 h
        - {trajectory_length: 7, until_step: 8000}    # 36 h
        - {trajectory_length: 9, until_step: 12500}   # 48 h
        - {trajectory_length: 11, until_step: 18000}  # 60 h
        - {trajectory_length: 13, until_step: 24500}  # 72 h
        - {trajectory_length: 17, until_step: 32000}  # 96 h
        - {trajectory_length: 21, until_step: 40500}  # 120 h
      loss:
        backend: crps
        variable_weights: null
        variable_scale: 1.0
        nodal_time_scale_hours: 24.0
        spectral_time_scale_hours: 40.0
        spectral_max_wavenumber: 80
  checkpoint: null
  gin_config: null
  train_dataset: null
  eval_dataset: null

inference:
  mode: weather_forecast
  # Used by stochastic profiles; deterministic checkpoints ignore the key.
  seed: 20260904
  # Memory-conscious default: eight 6-hour outputs reach forecast day 2.
  # Set --steps 60 to exercise the model's full 15-day capability.
  prediction_steps: 8
  output_interval_hours: 6
  # An official-format local model_bak.pkl takes precedence when present;
  # otherwise inference selects the profile's bundled checkpoint.
  checkpoint: data/checkpoint/model_bak.pkl
  output: results/predictions.nc

runtime:
  platform: auto
  dcu_device: 0