config.json

{
  "model_name": "AURORA",
  "model_type": "aurora",
  "architectures": [
    "Aurora"
  ],
  "framework": "PyTorch",
  "domain": "atmosphere",
  "task": "earth-system-forecasting",
  "implementation": {
    "entry_point": "model/aurora.py",
    "scope": "OneScience adapter over the official Microsoft Aurora implementation, with base (1.3B) and small variants and official checkpoint compatibility"
  },
  "architecture": {
    "family": "3D Perceiver encoder, 3D Swin Transformer processor, and 3D Perceiver decoder",
    "input_grid_shape": [
      721,
      1440
    ],
    "patch_size": 4,
    "window_size": [
      2,
      6,
      12
    ],
    "latent_levels": 4,
    "max_history_size": 2,
    "mlp_ratio": 4.0,
    "input_channels": 69,
    "output_channels": 69,
    "static_variables": [
      "lsm",
      "z",
      "slt"
    ],
    "variants": {
      "base": {
        "parameters_billion": 1.3,
        "embed_dim": 512,
        "encoder_depths": [
          6,
          10,
          8
        ],
        "encoder_num_heads": [
          8,
          16,
          32
        ],
        "decoder_depths": [
          8,
          10,
          6
        ],
        "decoder_num_heads": [
          32,
          16,
          8
        ],
        "aggregation_num_heads": 16
      },
      "small": {
        "purpose": "smoke test",
        "embed_dim": 256,
        "encoder_depths": [
          2,
          6,
          2
        ],
        "encoder_num_heads": [
          4,
          8,
          16
        ],
        "decoder_depths": [
          2,
          6,
          2
        ],
        "decoder_num_heads": [
          16,
          8,
          4
        ],
        "aggregation_num_heads": 8
      }
    }
  },
  "data": {
    "dataset": "ERA5",
    "spatial_resolution_degrees": 0.25,
    "time_step_hours": 6,
    "input_steps": 2,
    "output_steps": 1,
    "surface_variables": [
      "2m_temperature",
      "10m_u_component_of_wind",
      "10m_v_component_of_wind",
      "mean_sea_level_pressure"
    ],
    "atmospheric_variables": [
      "geopotential",
      "u_component_of_wind",
      "v_component_of_wind",
      "temperature",
      "specific_humidity"
    ],
    "pressure_levels_hpa": [
      50,
      100,
      150,
      200,
      250,
      300,
      400,
      500,
      600,
      700,
      850,
      925,
      1000
    ]
  },
  "configuration_sources": [
    "conf/config.yaml",
    "model/aurora.py",
    "model/aurora_network.py"
  ]
}