1graphcast/               # 核心程式碼
 2  graphcast.py           # GraphCast 模型架構
 3  gencast.py             # GenCast 模型架構
 4  autoregressive.py      # 自回歸推論包裝器
 5  icosahedral_mesh.py    # 正二十面體網格定義
 6  deep_typed_graph_net.py # 深層圖神經網路
 7  ...
 8graphcast_demo.ipynb     # GraphCast 展示 Notebook
 9gencast_mini_demo.ipynb  # GenCast Mini 展示 Notebook
10setup.py                 # 套件安裝

 1# 簡化的 Message Passing 虛擬碼
 2def message_passing(graph):
 3    # 1. 邊更新：根據發送節點、接收節點、邊特徵計算新邊特徵
 4    for edge in graph.edges:
 5        edge.features = MLP([sender.features, receiver.features, edge.features])
 6
 7    # 2. 節點更新：聚合所有入邊的訊息
 8    for node in graph.nodes:
 9        incoming = aggregate([e.features for e in node.incoming_edges])
10        node.features = MLP([node.features, incoming])
11
12    return graph

 1# 1. Clone Repository
 2git clone https://github.com/google-deepmind/graphcast.git
 3cd graphcast
 4
 5# 2. 建立虛擬環境 (推薦使用 uv)
 6uv venv --python 3.11 .venv
 7source .venv/bin/activate
 8
 9# 3. 安裝套件
10pip install -e .
11
12# 4. 驗證安裝
13python -c "import graphcast; print('GraphCast installed successfully')"

 1# 核心框架
 2JAX                    # Autograd + XLA 編譯器 (Google 的深度學習框架)
 3dm-haiku               # JAX 上的 Neural Network Library (神經網路函式庫)
 4jraph                  # JAX 上的 Graph Neural Network Library
 5
 6# 資料處理
 7xarray                 # 多維標籤化陣列 (適合氣象資料)
 8xarray-tensorstore     # 高效能 Zarr/TensorStore 讀取
 9numpy / scipy          # 數值計算
10pandas                 # 時間序列索引
11dask                   # 惰性 / 平行計算
12
13# 輔助工具
14chex                   # JAX 測試與除錯工具
15trimesh                # 三角網格運算 (建立正二十面體)
16dinosaur               # 球諧函數運算

1# 安裝 JAX GPU 版本 (以 CUDA 12 為例)
2pip install --upgrade "jax[cuda12]"
3
4# 驗證 GPU 可見
5python -c "import jax; print(jax.devices())"
6# 預期輸出: [GpuDevice(id=0, process_index=0)]

 1# 1. 建立 TPU VM
 2gcloud compute tpus tpu-vm create graphcast-vm \
 3    --zone=us-central2-b \
 4    --accelerator-type=v4-8 \
 5    --version=tpu-ubuntu2204-base
 6
 7# 2. SSH 連線
 8gcloud compute tpus tpu-vm ssh graphcast-vm --zone=us-central2-b
 9
10# 3. 在 VM 上安裝
11pip install graphcast
12pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

 1# 列出可用檔案
 2gsutil ls gs://dm_graphcast/graphcast/
 3gsutil ls gs://dm_graphcast/gencast/
 4
 5# 下載 GraphCast small 模型 (最小，適合測試)
 6gsutil cp gs://dm_graphcast/graphcast/GraphCast_small_params.npz ./data/
 7gsutil cp gs://dm_graphcast/graphcast/stats/ ./data/stats/ -r
 8
 9# 下載範例輸入資料
10gsutil cp gs://dm_graphcast/graphcast/dataset/ ./data/dataset/ -r

 1import dataclasses
 2import functools
 3import math
 4import xarray
 5from graphcast import (
 6    autoregressive,
 7    casting,
 8    checkpoint,
 9    data_utils,
10    graphcast,
11    normalization,
12    rollout,
13)
14
15# 1. 載入模型參數
16with open("GraphCast_small_params.npz", "rb") as f:
17    ckpt = checkpoint.load(f, graphcast.CheckPoint)
18    params = ckpt.params
19    state = {}
20    model_config = ckpt.model_config
21    task_config = ckpt.task_config
22
23print(f"Model config: {model_config}")
24print(f"Task config: {task_config}")

 1# 2. 載入範例 ERA5 資料
 2# 需要兩個連續時步 (t-6h, t) 的大氣狀態
 3example_batch = xarray.open_dataset("era5_example_data.nc")
 4
 5# 資料維度說明：
 6# - batch: 批次維度
 7# - time: 時間步 (至少 2 步)
 8# - lat: 緯度 (721 for 0.25deg)
 9# - lon: 經度 (1440 for 0.25deg)
10# - level: 氣壓層 (37 or 13)
11
12# 3. 分離 inputs / targets / forcings
13inputs, targets, forcings = data_utils.extract_inputs_targets_forcings(
14    example_batch,
15    target_lead_times=slice("6h", f"{6 * task_config.eval_steps}h"),
16    **dataclasses.asdict(task_config),
17)

 1# 4. 載入正規化統計量
 2diffs_stddev_by_level = xarray.open_dataset("stats/diffs_stddev_by_level.nc")
 3mean_by_level = xarray.open_dataset("stats/mean_by_level.nc")
 4stddev_by_level = xarray.open_dataset("stats/stddev_by_level.nc")
 5
 6# 5. 建構 Predictor 鏈
 7#    GraphCast -> Normalization -> Autoregressive -> Casting
 8def construct_wrapped_graphcast(model_config, task_config):
 9    predictor = graphcast.GraphCast(model_config, task_config)
10
11    predictor = normalization.InputsAndResiduals(
12        predictor,
13        diffs_stddev_by_level=diffs_stddev_by_level,
14        mean_by_level=mean_by_level,
15        stddev_by_level=stddev_by_level,
16    )
17
18    predictor = autoregressive.Predictor(
19        predictor, gradient_checkpointing=True
20    )
21
22    return predictor
23
24# 6. 使用 Haiku 封裝 (JAX 風格)
25import haiku as hk
26import jax
27
28@hk.transform_with_state
29def run_forward(inputs, targets_template, forcings):
30    predictor = construct_wrapped_graphcast(model_config, task_config)
31    return predictor(inputs, targets_template=targets_template, forcings=forcings)
32
33# 7. 執行推論
34predictions, _ = run_forward.apply(
35    params, state, jax.random.PRNGKey(0),
36    inputs, targets, forcings
37)
38
39print(f"Predictions shape: {predictions.dims}")
40# 預期: {'batch': 1, 'time': 40, 'lat': 721, 'lon': 1440, 'level': 37}

 1from graphcast import rollout
 2
 3# rollout 模式：不保留梯度，可跑更長時步
 4predictions = rollout.chunked_prediction(
 5    run_forward.apply,
 6    rng=jax.random.PRNGKey(0),
 7    inputs=inputs,
 8    targets_template=targets * 0,  # 只需要形狀資訊
 9    forcings=forcings,
10    num_steps_per_chunk=1,          # 每次推論 1 步
11    params=params,
12    state=state,
13)

 1# 計算訓練損失
 2@hk.transform_with_state
 3def loss_fn(inputs, targets, forcings):
 4    predictor = construct_wrapped_graphcast(model_config, task_config)
 5    loss, diagnostics = predictor.loss(
 6        inputs, targets, forcings
 7    )
 8    return loss, diagnostics
 9
10# 計算梯度
11(loss, diagnostics), grads = jax.value_and_grad(
12    loss_fn.apply, has_aux=True
13)(params, state, jax.random.PRNGKey(0), inputs, targets, forcings)
14
15print(f"Training loss: {loss:.4f}")

 1from graphcast import graphcast
 2
 3# 自訂模型配置
 4custom_config = graphcast.ModelConfig(
 5    resolution=1,           # 1 degree 解析度
 6    mesh_size=5,            # 正二十面體細分次數 (5 → 10242 mesh nodes)
 7    latent_size=512,        # 隱藏層維度
 8    gnn_msg_steps=16,       # GNN 訊息傳遞步數
 9    hidden_layers=1,        # MLP 隱藏層數
10    radius_query_fraction_edge_length=0.6,  # Grid-Mesh 連接半徑
11)
12
13# 任務配置
14custom_task = graphcast.TaskConfig(
15    input_variables=[...],   # 輸入氣象變數列表
16    target_variables=[...],  # 預測目標變數列表
17    forcing_variables=[...], # 強迫變數 (太陽輻射等)
18    pressure_levels=[...],   # 氣壓層列表
19    input_duration="12h",    # 輸入時間跨度
20)

 1# 策略 1: Gradient Checkpointing (梯度檢查點)
 2predictor = autoregressive.Predictor(
 3    predictor,
 4    gradient_checkpointing=True  # 以計算換記憶體
 5)
 6
 7# 策略 2: BFloat16 精度
 8predictor = casting.Bfloat16Cast(predictor)
 9
10# 策略 3: Chunked Rollout (分段推論)
11predictions = rollout.chunked_prediction(
12    ...,
13    num_steps_per_chunk=1,  # 每次只推論 1 步，釋放中間記憶體
14)

 1from graphcast import solar_radiation
 2
 3# 計算 Top-of-Atmosphere 入射太陽輻射
 4# 與 ERA5 的 tisr 變數相容
 5toa_radiation = solar_radiation.get_toa_incident_solar_radiation_for_xarray(
 6    single_level_dataset,
 7    datetime_dim="time",
 8    latitude_dim="lat",
 9    longitude_dim="lon",
10)

 1import jax.numpy as jnp
 2
 3# GenCast 生成多個 Ensemble Members
 4n_ensemble = 50
 5
 6predictions_ensemble = []
 7for i in range(n_ensemble):
 8    rng = jax.random.PRNGKey(i)  # 不同隨機種子產生不同成員
 9    pred = gencast_rollout(rng, inputs, forcings)
10    predictions_ensemble.append(pred)
11
12# 計算集合統計量
13ensemble_mean = jnp.mean(jnp.stack(predictions_ensemble), axis=0)
14ensemble_std = jnp.std(jnp.stack(predictions_ensemble), axis=0)
15# ensemble_std 代表預報不確定性

 1# 範例：將 GraphCast 預報整合進既有資料管線
 2import xarray as xr
 3
 4# 1. 取得 GraphCast 預報結果
 5predictions = run_graphcast_inference(current_era5_data)
 6
 7# 2. 轉換為標準 CF-conventions NetCDF
 8predictions.attrs["Conventions"] = "CF-1.8"
 9predictions.to_netcdf("graphcast_forecast.nc")
10
11# 3. 與既有 NWP 產品比較
12ecmwf_forecast = xr.open_dataset("ecmwf_hres_forecast.nc")
13bias = predictions - ecmwf_forecast
14rmse = ((bias ** 2).mean(dim=["lat", "lon"])) ** 0.5
15
16# 4. 產生決策支援資料
17high_wind_warning = predictions["10m_wind_speed"] > 17.2  # Beaufort Scale 8
18heavy_rain_alert = predictions["total_precipitation"] > 50  # mm/6h

1# WeatherBench2 提供標準化的驗證工具
2# https://weatherbench2.readthedocs.io
3from weatherbench2 import evaluation
4
5rmse = evaluation.compute_rmse(predictions, era5_truth)
6acc = evaluation.compute_acc(predictions, era5_truth, climatology)

 1@article{lam2023learning,
 2  title={Learning skillful medium-range global weather forecasting},
 3  author={Lam, Remi and Sanchez-Gonzalez, Alvaro and Willson, Matthew
 4          and Wirnsberger, Peter and Fortunato, Meire and Alet, Ferran
 5          and Ravuri, Suman and Ewalds, Timo and Eaton-Rosen, Zach
 6          and Hu, Weihua and others},
 7  journal={Science},
 8  volume={382},
 9  number={6677},
10  pages={1416--1421},
11  year={2023},
12  publisher={American Association for the Advancement of Science}
13}
14
15@article{price2023gencast,
16  title={GenCast: Diffusion-based ensemble forecasting
17         for medium-range weather},
18  author={Price, Ilan and Sanchez-Gonzalez, Alvaro and Alet, Ferran
19          and Andersson, Tom R and El-Kadi, Andrew and Masters, Dominic
20          and Ewalds, Timo and Stott, Jacklynn and Mohamed, Shakir
21          and Battaglia, Peter and Lam, Remi and Willson, Matthew},
22  journal={arXiv preprint arXiv:2312.15796},
23  year={2023}
24}