DeepMindがAIとJAXを使ってサイクロンを予測する方法

1週間先の天気予報は、従来、強力な超级コンピュータに依存していました。それらは大規模なグリッドで流体力学と熱力学の方程式を解くのに何時間も費やします。Google DeepMindは異なるアプローチを取りました：エンジニアたちはニューラルネットワークを訓練し、数秒で気象プロセスを予測するようにしました。リポジトリ google-deepmind/weathernext には、WeatherNext 2、GraphCast、GenCastアーキテクチャのソースコードと重みが含まれており、熱帯低気圧を追跡するための specialized modelsも含まれています。

プロジェクトの構造、必要なハードウェアリソース、開発者や研究者がそこから何を 얻出せるかを探索しましょう。

リポジトリの中身

このプロジェクトは、DeepMindの気象ニューラルネットワークの複数の世代を組み合わせています。主な焦点はWeatherNext 2 (WN2) — 全球中期予報モデルです。空間解像度0.25度（地球表面で約30km間隔のグリッド）で動作します。

このモデルは大気の初期状態を取り、自己回帰的実行を行い、温度、様々な高度での風速、重力ポテンシャル、サイクロンの動きを段階的に予報します。

リポジトリには3つのカテゴリのモデルが含まれています：

• WeatherNext 2 (<2025): 主力モデル。運用されているECMWF HRES（高解像度予報）データでファインチューニングされているため、実際の気象解析から直接初期化できます。標準的な指標に加えて、風力発電に重要な地上100mの風速予報も可能です。
• WeatherNext Cyclones: ハリケーンと熱帯低気圧の追跡に最適化されたモデルシリーズ（2023、2024、2025年版）。2025年版のFNV3は、実際の大西洋ハリケーンシーズンで实战磨练済みです。
• WeatherNext Cyclones Mini: 解像度1度の軽量モデル。はるかに小さく、ローカルテストや单个GPUでの実行に適しています。

WN2之外、フォルダ docs/ には以前のモデルのリンクと資料が含まれています：グラフベースのGraphCastと拡散ベースのGenCast。

ハードウェアとシステム要件

0.25度のフルサイズモデルを実行するのはリソース集約的です。重たいチェックポイント推論はTPU（特にTPU v5p）に最適化されています。GPUでフルバージョンを実行しようとすると、高解像度テンソル実行に必要なメモリ量のため、NVIDIA H100レベルのカードが必要です。

探索と実験のために、作成者はMiniバージョンをリリースしました。これは開発者に優しいです：

• 無料でTPU v5e-1でGoogle Colabスムーズに実行可能。
• NVIDIA P100やT4などの古いアクセラレーターでもGPUで動作。

コードの構造と開始方法

コードはPythonで書かれており、多次元気象データ配列を操作するJAXとライブラリ xarray を使用しています。

ライブラリをインストールするには、単に実行してください：

作成者は特定のリリースに固定することを強くお勧めします。これは研究コードであり、APIの安定性は保証されていないためです。

Colabによるクイックスタート

フォルダ docs/weathernext2/wn2_demo.ipynb にはすぐに使えるノートブックが含まれています。モデルの完全なワークフローを示しています：

1. 公開Google Cloudバケットから重みをロード（dm_graphcast）。
2. 初期大気条件を準備（Zarr形式とWeatherBench2プラットフォーム経由のERA5またはHRESデータ）。
3. JAXでFGN / WN2アーキテクチャを初期化。
4. 自己回帰ロールアウトを実行して数日先の予報を生成。
5. 組み込みトラッカーを使用してサイクロン軌道を計算。
6. 損失を計算し、1つの勾配ステップを実行（はい局所的にモデルをファインチューンできます）。

概念的には、予報ループがコードでどのように見えるかです：

実践的な応用

独自の気象データでモデルをファインチューンしたり、科学的なパイプラインを構築する予定がない場合は、JAXインフラのセットアップは不要です。Googleは готовые WN2結果をオープンソースに公開しています：

• OpenMeteo API: GPUに接続せずに通常のRESTリクエストでWN2データを取得できます。
• Google Cloud / BigQuery / Vertex AI: 分析と大規模クエリ用の готовые データセット。
• WeatherLab: ハリケーン軌道を可視化するインタラクティブサービス。

再生可能エネルギー研究、気候リスク分析、独自のMLサービスの構築に取り組む人々にとって、リポジトリからのコードは готовые 基盤として機能します。セクション utils/ には、データ正規化、グラフブロックの操作、特定の気象損失の計算のための便利なユーティリティが含まれています。

まとめ

WeatherNextは、深層学習が伝統的な数値天気予報（NWP）を徐々に置き換えるか補完しているかを示す優れたデモンストレーションです。

長所：

• Colabで実行できるオープンウェイトとデモノートブック。
• 通常のハードウェアでの実験用の軽量Miniモデル。
• 実際の気象データとの готовые 統合（WeatherBench2、ERA5、HRES）。

短所：

• フル解像度推論には最上位のH100またはTPUアクセラレーターが必要です。
• APIはバージョン間の後方互換性を保証しません。

JAXで作業したり、大規模なテンソルグリッドで自己回帰パスを実行する方法を理解したい場合は、ぜひWN2デモノートブックを試してください。