DeepMind「GenCast」次世代高精度天気予測モデル2「GraphCast」モデルを利用し全変数の可視化

# **DeepMind「GenCast」次世代高精度天気予測モデル2「GraphCast」モデルを利用し全変数の可視化**

こんにちは、しゅんです！  

今回は、**Google DeepMind**の高精度天気予測モデル「**GraphCast**」で提供されるデータセットを用いて、すべての変数を可視化する方法をご紹介します。このモデルは、膨大な気象データを効率的に処理するために設計されていますが、この記事では、**モデルのトレーニングや推論を行うのではなく、提供されたデータの内容を理解するための可視化**に焦点を当てます。

僕は**Jupyter Notebookが嫌い**なので（笑）、今回も**Pythonスクリプト形式**で実行しました！この記事を参考に、皆さんもデータセットの内容をぜひ確認してみてください。

---

## **GraphCastとは？**

### **概要**
GraphCastは、DeepMindが開発した**グラフニューラルネットワーク（GNN）を基盤とする天気予測モデル**です。このモデルは、膨大な気象データを効率的に処理し、**気温、降水量、風速など多様な気象変数を高精度に予測**することを目的としています。

具体的な特徴：
- **最大15日先までの天気予測が可能**
- **台風進路や気象現象の変化をモデル化**
- **多様な気象変数に対応**
- 柔軟性が高く、異なる解像度やデータセットでの利用が可能

### **今回の目的**
この記事では、GraphCastが提供する**過去の観測データセット**（例: ERA5 データ）の内容を可視化し、各変数がどのような情報を持つかを確認します。  
**予測モデルの利用やトレーニングはこの記事では扱いません。**

---

## **データセットの概要**

以下は、データセットに含まれる変数とその次元情報の一覧です。

| **変数名**                     | **次元情報**                             | **データの種類**                     | **説明**                                                                                                         |
|---------------------------------|-----------------------------------------|-------------------------------------|-----------------------------------------------------------------------------------------------------------------|
| `geopotential_at_surface`       | `(lat, lon)`                            | 静的データ                          | 地表での地形高度                                                                                               |
| `land_sea_mask`                 | `(lat, lon)`                            | 静的データ                          | 陸地（1）と海洋（0）の識別                                                                                     |
| `2m_temperature`                | `(batch, time, lat, lon)`               | 時間依存データ                      | 地表付近（2m）の気温                                                                                           |
| `mean_sea_level_pressure`       | `(batch, time, lat, lon)`               | 時間依存データ                      | 平均海面気圧                                                                                                   |
| `10m_v_component_of_wind`       | `(batch, time, lat, lon)`               | 時間依存データ                      | 地表付近（10m）の南北方向の風速                                                                                |
| `10m_u_component_of_wind`       | `(batch, time, lat, lon)`               | 時間依存データ                      | 地表付近（10m）の東西方向の風速                                                                                |
| `total_precipitation_6hr`       | `(batch, time, lat, lon)`               | 時間依存データ                      | 6時間ごとの総降水量                                                                                           |
| `toa_incident_solar_radiation`  | `(batch, time, lat, lon)`               | 時間依存データ                      | 大気上端での入射太陽放射量                                                                                    |
| `temperature`                   | `(batch, time, level, lat, lon)`        | 時間依存データ（レベル有り）         | 特定の大気層での温度                                                                                          |
| `geopotential`                  | `(batch, time, level, lat, lon)`        | 時間依存データ（レベル有り）         | 特定の気圧面での地形高度                                                                                      |
| `u_component_of_wind`           | `(batch, time, level, lat, lon)`        | 時間依存データ（レベル有り）         | 特定の気圧面での東西方向の風速                                                                                |
| `v_component_of_wind`           | `(batch, time, level, lat, lon)`        | 時間依存データ（レベル有り）         | 特定の気圧面での南北方向の風速                                                                                |
| `vertical_velocity`             | `(batch, time, level, lat, lon)`        | 時間依存データ（レベル有り）         | 特定の気圧面での鉛直速度                                                                                      |
| `specific_humidity`             | `(batch, time, level, lat, lon)`        | 時間依存データ（レベル有り）         | 特定の気圧面での比湿                                                                                        |

---

## **可視化のコード**

以下は、データセットを可視化するコードの全体です。

```graphcast_demo_without_train.py
import dataclasses
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import xarray
from graphcast import data_utils
from graphcast import graphcast
from graphcast import normalization
from graphcast import rollout
from graphcast import xarray_jax

# データセットのパスを設定
DATASET_PATH = "./graphcast/dataset/source-era5_date-2022-01-01_res-0.25_levels-13_steps-12.nc"

# プロット関数
def select(data, variable, level=None, max_steps=None):
    data = data[variable]
    if "batch" in data.dims:
        data = data.isel(batch=0)
    if max_steps is not None and "time" in data.sizes and max_steps < data.sizes["time"]:
        data = data.isel(time=range(0, max_steps))
    if level is not None and "level" in data.coords:
        data = data.sel(level=level)
    return data

def scale(data, center=None, robust=False):
    vmin = np.nanpercentile(data, 2 if robust else 0)
    vmax = np.nanpercentile(data, 98 if robust else 100)
    if center is not None:
        diff = max(vmax - center, center - vmin)
        vmin = center - diff
        vmax = center + diff
    return (data, matplotlib.colors.Normalize(vmin, vmax), "RdBu_r" if center is not None else "viridis")

def plot_data(data, title, plot_size=10):
    cols = len(data)
    rows = 1
    fig, axs = plt.subplots(rows, cols, figsize=(plot_size * cols, plot_size))
    if cols == 1:
        axs = [axs]
    for i, (key, (values, norm, cmap)) in enumerate(data.items()):
        ax = axs[i]
        im = ax.imshow(values, norm=norm, origin="lower", cmap=cmap)
        ax.set_title(key)
        fig.colorbar(im, ax=ax)
    plt.tight_layout()
    plt.show()

# データセットをロード
example_batch = xarray.load_dataset(DATASET_PATH).compute()

print("Dataset variables:", example_batch.data_vars.keys())
print("Dataset dimensions:", example_batch.sizes)

# すべてのデータ変数をプロット
for variable in example_batch.data_vars.keys():
    try:
        print(f"Plotting variable: {variable}")
        data = example_batch[variable]

        if "time" in data.dims:
            level = data.coords["level"].values[0] if "level" in data.dims else None
            data_to_plot = select(example_batch, variable, level)
            scaled_data = {f"{variable}": scale(data_to_plot.isel(time=0))}
            plot_data(scaled_data, f"{variable} (time=0)")
        elif "lat" in data.dims and "lon" in data.dims:
            scaled_data = {f"{variable}": scale(data)}
            plot_data(scaled_data, f"{variable} (Static Data)")
        else:
            print(f"Skipping {variable}: Unsupported dimensions {data.dims}")
    except Exception as e:
        print(f"Error while plotting {variable}: {e}")

```

---

## **実行結果**

コードを実行すると、データセット内の各変数がプロットされます。以下はその一部です。表を参照：

1. **`2m_temperature`**: 地表付近の気温
2. **`mean_sea_level_pressure`**: 平均海面気圧
3. **`geopotential`**: 特定の気圧面での地形高度
    ```
    Dataset variables: KeysView(Data variables:
        geopotential_at_surface       (lat, lon) float32 4MB 2.735e+04 ... -0.07617
        land_sea_mask                 (lat, lon) float32 4MB 1.0 1.0 1.0 ... 0.0 0.0
        2m_temperature                (batch, time, lat, lon) float32 58MB 250.7 ...
        mean_sea_level_pressure       (batch, time, lat, lon) float32 58MB 9.931e...
        10m_v_component_of_wind       (batch, time, lat, lon) float32 58MB -0.439...
        10m_u_component_of_wind       (batch, time, lat, lon) float32 58MB 1.309 ...
        total_precipitation_6hr       (batch, time, lat, lon) float32 58MB 0.0004...
        toa_incident_solar_radiation  (batch, time, lat, lon) float32 58MB 1.981e...
        temperature                   (batch, time, level, lat, lon) float32 756MB ...
        geopotential                  (batch, time, level, lat, lon) float32 756MB ...
        u_component_of_wind           (batch, time, level, lat, lon) float32 756MB ...
        v_component_of_wind           (batch, time, level, lat, lon) float32 756MB ...
        vertical_velocity             (batch, time, level, lat, lon) float32 756MB ...
        specific_humidity             (batch, time, level, lat, lon) float32 756MB ...)
    Dataset dimensions: Frozen({'lon': 1440, 'lat': 721, 'time': 14, 'level': 13, 'batch': 1})
    Plotting variable: geopotential_at_surface
    Plotting variable: land_sea_mask
    Plotting variable: 2m_temperature
    Plotting variable: mean_sea_level_pressure
    Plotting variable: 10m_v_component_of_wind
    Plotting variable: 10m_u_component_of_wind
    Plotting variable: total_precipitation_6hr
    Plotting variable: toa_incident_solar_radiation
    Plotting variable: temperature
    Plotting variable: geopotential
    Plotting variable: u_component_of_wind
    Plotting variable: v_component_of_wind
    Plotting variable: vertical_velocity
    Plotting variable: specific_humidity
    
    ```
---

## **注意点**

- このコードは**データセット内の観測値や初期条件を可視化**することを目的としています。
- **未来の予測やモデルの推論結果は含まれていません**。

---

# 画像

![10m_u_component_of_wind.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/ac05ce59-ab6c-4af1-5d0f-c9a1793475d8.png)

![10m_v_component_of_wind.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/4ae25f62-f03c-500b-b693-f9e92b8aa4a4.png)

![geopotential.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/87902538-1edb-aa76-a1a2-2166a970d894.png)

![geopotential_at_surface.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/94bbeda1-e055-b698-1e52-2c38b3a606b3.png)

![land_sea_mask.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/42753c9c-539f-47ff-0df4-db30fec5cef0.png)

![mean_sea_level_pressure.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/5f06a4e4-1b17-6ba3-af7e-ef21053cd6a4.png)

![specific_humidity.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/b518a972-1342-20b2-b642-9c5934751d89.png)

![temperature.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/248acf37-3353-59ed-0b46-61e6b21b95b8.png)

![toa_incident_solar_radiation.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/0bd4f1d6-1a77-5dca-8b92-402cd5ce4f36.png)

![total_precipitation_6hr.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/d6b29f9a-b117-07c8-3bb5-7258a80c623e.png)

![u_component_of_wind.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/e1fd8f66-d8ec-2188-5a1d-00d7f1624949.png)

![v_component_of_wind.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/2a744a0c-bf85-6773-616e-5e59e3b73806.png)

![vertical_velocity.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/e4c314df-32fc-2f39-f4b4-d2c992a5c0f7.png)


## **結論**

GraphCast データセットには多様な気象変数が含まれており、その可視化を通じて気象現象を直感的に理解することができます。次回は、GraphCast モデルの予測プロセスや推論結果の分析に挑戦する予定です！  

最後まで読んでいただきありがとうございました！この記事を参考に、ぜひ GraphCast のデータセットを試してみてください。