File size: 3,403 Bytes
0ba2894
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
{
  "model_name": "Prithvi WxC",
  "model_type": "prithvi_wxc",
  "architectures": [
    "PrithviWxC"
  ],
  "framework": "PyTorch",
  "domain": "climate-and-atmosphere",
  "task": "global-weather-forecasting",
  "implementation": {
    "entry_point": "model/prithvi_wxc.py",
    "scope": "YAML-driven wrapper around the vendored official Hiera-MaxViT encoder-decoder, with identity model scalers for the repository's small connectivity configuration"
  },
  "architecture": {
    "family": "Hiera-MaxViT encoder-decoder Vision Transformer",
    "attention_mechanism": "alternating local block attention and global grid attention",
    "input_format": "BTCHW atmospheric history plus BCHW static fields",
    "output_format": "BCHW next atmospheric state",
    "tokenization": "Conv2d patch embedding followed by mask-unit grouping and pixel-shuffle reconstruction",
    "activation": "GELU",
    "normalization": "LayerNorm",
    "repository_default_config": {
      "purpose": "small connectivity validation with synthetic or reduced ERA5 data",
      "in_channels": 6,
      "input_size_time": 2,
      "in_channels_static": 4,
      "grid_size": [
        32,
        64
      ],
      "patch_size_px": [
        2,
        2
      ],
      "mask_unit_size_px": [
        8,
        8
      ],
      "mask_ratio_inputs": 0.0,
      "embed_dim": 32,
      "n_blocks_encoder": 1,
      "n_blocks_decoder": 1,
      "n_heads": 4,
      "mlp_multiplier": 4.0,
      "dropout": 0.0,
      "drop_path": 0.0,
      "parameter_dropout": 0.0,
      "residual": "none",
      "masking_mode": "global",
      "positional_encoding": "absolute",
      "encoder_shifting": false,
      "decoder_shifting": false,
      "normalization_scalers": "identity in the YAML wrapper"
    },
    "paper_reference_config": {
      "grid_resolution_degrees": [
        0.5,
        0.625
      ],
      "grid_size": "approximately 360/361 x 576",
      "dynamic_channels": 160,
      "embed_dim": 2560,
      "encoder_blocks": "13 local plus 12 global blocks",
      "decoder_blocks": "3 local plus 2 global blocks",
      "n_heads": 16,
      "pretraining_mask_ratio_inputs": 0.5,
      "forecast_finetuning_mask_ratio_inputs": 0.0,
      "parameter_count": "approximately 2.3 billion",
      "compatibility_note": "official 2.3B checkpoints do not match the repository default small configuration without aligning channels, grid size, and architecture"
    }
  },
  "data": {
    "dataset": "ERA5",
    "storage": "HDF5 fields with shape TCHW",
    "temporal_interval_hours": 6,
    "input_steps": 2,
    "output_steps": 1,
    "forecast_lead_time_hours": 6,
    "dynamic_channels": [
      "10m_u_component_of_wind",
      "10m_v_component_of_wind",
      "2m_temperature",
      "mean_sea_level_pressure",
      "geopotential_500",
      "temperature_850"
    ],
    "static_channels": 4,
    "default_spatial_size": [
      32,
      64
    ],
    "default_year_splits": {
      "train": [
        1951,
        1952
      ],
      "validation": [
        1953
      ],
      "test": [
        1954
      ]
    },
    "synthetic_data_generator": "scripts/fake_data.py"
  },
  "configuration_sources": [
    "conf/config.yaml",
    "model/prithvi_wxc.py",
    "model/prithvi_wxc_official.py",
    "scripts/train.py",
    "scripts/inference.py",
    "scripts/fake_data.py",
    "README.md",
    "configuration.json"
  ]
}