earth2studio-create-prognostic

クイックスタートチェックリスト

以下の手順を順番どおりに実行してください。手順を省略しないでください。

• まず、この SKILL.md を最後まで読む
• 参照スクリプトを取得する（手順 0）
• 三重継承を使って earth2studio/models/px/<name>.py を作成する
• モックテストを含む test/models/px/test_<name>.py を作成する
• 実行する：uv run pytest test/models/px/test_<name>.py -v
• モデル用エクストラ、インストール、API ドキュメント、変更履歴を追加／更新する（手順 1～2、9）
• 実行する：make format && make lint

⚠️ 重要： Python コマンドには必ず uv run を使用してください。

• ✅ uv run pytest ... ／ uv run python ...
• ❌ pytest ... ／ python ...（依存関係が不足します）

行き詰まった場合や出力が間違っている場合： 同じ修正を何度も試さないでください。続行する前に、 自己改善に従って、このスキルに修正を加えてください。

目的

サードパーティ製の ML 気象モデルと Earth2Studio を接続する予報モデル用ラッパーを実装します。予報モデルは時間方向に積分を進めます。つまり、初期状態を与えると、時間を刻みながら（例：6 時間ごとに）未来の状態を予測します。

ワークスペース

コンテキスト	場所
Harbor 評価	/workspace/output/earth2studio/models/px/... に書き込みます
Harbor ＋ --copy-repo	/workspace/repo に完全なチェックアウトがあります
ローカルクローン	pyproject.toml があるディレクトリ

evals/targets/ は絶対に読み込まないでください — 採点用の参照のみです。

参照ファイル

該当する手順で必要に応じて読み込みます。

ファイル	内容	読み込む手順
references/skeleton-template.py	FILL コメント付きの完全なモデルスケルトン	手順 3～6
references/method-templates.py	標準的なメソッド実装	手順 4～6
references/testing-guide.py	テストスケルトンとモックのパターン	手順 7
references/validation-guide.md	比較スクリプト、PR、コードレビュー	手順 10～11

---

ワークフローの手順

手順 0 — 参照スクリプトを取得する

$ARGUMENTS が指定されている場合は、それを使用します。指定されていない場合は、次のように尋ねます。

参照用の推論スクリプトの URL またはパスを指定してください。

手順 1 — 依存関係を分析して提案する

パッケージ、アーキテクチャ、入出力の形状、時間ステップ、解像度、チェックポイントを分析します。

pyproject.toml のグループを提案します（アルファベット順に並べ、all に追加）。必要なパッケージがない場合でも、すべての予報モデルにオプションの依存関係エクストラが必要です：

model-name = ["package1>=version", "package2"]
# or, when no additional packages are required:
model-name = []

[CONFIRM] 依存関係を提示し、ユーザーに承認を求めます。

手順 2 — 依存関係を追加する

pyproject.toml を編集します。モデル用エクストラを、空の場合も含めてアルファベット順に追加し、all の集約設定を更新します。

手順 3 — モデルファイルを作成する

ファイル： earth2studio/models/px/<lowercase>.py

必須の継承（3 つすべて）：

class ModelName(torch.nn.Module, AutoModelMixin, PrognosticMixin):

必須のインポート：

import numpy as np
import torch
from earth2studio.models.auto import AutoModelMixin, Package
from earth2studio.models.batch import batch_coords, batch_func
from earth2studio.models.px.base import PrognosticMixin
from earth2studio.models.utils import create_coords_from_lat_lon, handshake_dim
from earth2studio.lexicon import E2STUDIO_VOCAB
from earth2studio.utils import check_optional_dependencies
from loguru import logger

SPDX ヘッダー（すべての .py ファイルの先頭に必須）：

# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES.
# SPDX-License-Identifier: Apache-2.0

標準的なメソッドの順序：

1. __init__ 2. input_coords 3. output_coords (@batch_coords)
2. load_default_package 5. load_model 6. to（任意）
3. プライベートメソッド 8. __call__ (@batch_func) 9. _default_generator
4. create_iterator

手順 4 — 座標を実装する

input_coords のルール：

• batch：np.empty(0)
• time：np.empty(0)（動的）
• lead_time：np.timedelta64(0, "h") から開始
• lat：90 から -90（北から南）。ソースモデルが逆の順序を使用する場合でも、これは Earth2Studio で公開する際の規約です
• lon：0 から 360
• チェックポイント／モデルのコアが南から北への緯度を想定している場合は、コアモデルの前後で内部的にテンソルを反転します。input_coords や output_coords では、反転した緯度を公開しないでください
• 変数を E2STUDIO_VOCAB（earth2studio/lexicon/base.py にある 282 項目）に対応付けます

output_coords： 入力検証には handshake_dim／handshake_coords を使用し、その後 lead_time を増分します。座標チェック用ヘルパーを共通化し、モデルの実行前に output_coords、__call__、イテレーターのセットアップから呼び出すことを推奨します。

ステップ 5 — フォワードパスの実装

__call__: @batch_func でデコレートされ、形状は (batch, time, lead_time, var, lat, lon) です。 モデルの形式にリシェイプ → モデルを呼び出し → 元に戻す。

create_iterator: 必ず最初に初期条件 (ステップ 0) を yield してください。 摂動の注入には front_hook/rear_hook を使用します。

ステップ 6 — モデル読み込みの実装

load_default_package: HuggingFace の URL を固定します: hf://org/repo@commit

load_model: package.resolve()、map_location="cpu"、eval() モードを使用し、 @check_optional_dependencies() でデコレートします。

ステップ 7 — テストの作成

ファイル: test/models/px/test_<name>.py

必須テスト:

Function	Purpose
test_<model>_call	単一のフォワードパス (デバイス/時間をパラメータ化)
test_<model>_iter	イテレータがシーケンスを生成する
test_<model>_exceptions	無効な座標がエラーを発生させる
test_<model>_package	実重み (@pytest.mark.package)

モックテスト用に、インターフェースに一致する PhooModelName ダミーを作成します。

テストの実行:

uv run pytest test/models/px/test_<name>.py -m "not package" -v
uv run pytest test/models/px/test_<name>.py::test_<model>_package --package -v

パッケージテストを省略しないでください。任意のランダム入力が実際のチェックポイントに対して 物理的に有効でない場合は、実重みを読み込みフォワードパスを実行しつつ、 モデルに適した安定した合成入力を使用してください。

ステップ 8 — モデルの登録 (要求された場合)

• earth2studio/models/px/__init__.py に追加 (アルファベット順)
• pyproject.toml の依存関係を確認

ステップ 9 — ドキュメント

• docs/modules/models_px.rst に追加 (アルファベット順)。これは新しい予測モデルごとに 必須で、API ドキュメントに生成ページが含まれるようにします。
• モデル extra について docs/userguide/about/install.md (アルファベット順のタブ) に追加します。 extra が空の場合でも同様です。モデル固有の注意事項に加え、 pip install earth2studio[model-name] と uv add earth2studio --extra model-name の両方の手順を含めてください。
• CHANGELOG.md の ### Added の下を更新します。これは新しい予測モデルごとに 必須です。

フォーマットとリント:

make format && make lint && make license

ステップ 10 - 検証 (要求された場合)

references/validation-guide.md に従ってください。未コミットの vanilla、E2S、 比較、サニティチェック用のスクリプトを作成します。生成された出力や画像はコミットしないでください。 ユーザーが手動で画像をアップロードできるよう、プロットには PR で安全なプレースホルダーを使用します。

[CONFIRM] ユーザーは続行する前にプロットを目視で確認する必要があります。

ステップ 11 - PR (要求された場合)

references/validation-guide.md に従い、以下を使用します:

• references/pr-body-template.md
• references/pr-comment-template.md

PR を作成する前に、pyproject.toml にモデル extra があること、all extra にそれが含まれていること、 インストールドキュメントに pip と uv の両方のコマンドが含まれていること、 そして docs/modules/models_px.rst と CHANGELOG.md が更新されていることを確認してください。

PR のテキストにマシン名、絶対パス、デバイス構成、アップロードされた画像リンクを含めないでください。 代わりにプロットのプレースホルダーを使用してください。

---

例

シンプルな Identity モデル

User: Create IdentityModel - returns input unchanged, 6h step, 181x360, vars: t2m, u10m, v10m, msl

Agent: [reads SKILL.md, creates identity.py with triple inheritance,
        creates test_identity.py, runs pytest, runs make format && lint]

外部モデル (Pangu)

User: Add Pangu-Weather wrapper
      GitHub: https://github.com/198808xc/Pangu-Weather

Agent: [reads SKILL.md, fetches inference.py, creates pangu.py,
        creates test_pangu.py, runs pytest]

---

主要なパターン

座標テンプレート

@property
def input_coords(self) -> CoordSystem:
    return CoordSystem({
        "batch": np.empty(0),
        "time": np.empty(0),
        "lead_time": np.array([np.timedelta64(0, "h")]),
        "variable": np.array(["t2m", "u10m", ...]),
        # Public Earth2Studio convention is north-to-south latitude.
        "lat": np.linspace(90, -90, 181),
        "lon": np.linspace(0, 359, 360),
    })

@batch_coords()
def output_coords(self, input_coords: CoordSystem) -> CoordSystem:
    output = input_coords.copy()
    output["lead_time"] = input_coords["lead_time"] + np.timedelta64(6, "h")
    return output

イテレータテンプレート

def create_iterator(self, x, coords):
    yield x, coords  # Initial condition (step 0)
    while True:
        x, coords = self.front_hook(x, coords)
        x, coords = self(x, coords)
        x, coords = self.rear_hook(x, coords)
        yield x, coords

---

トラブルシューティング

Error	Solution
OptionalDependencyFailure	uv add --optional <group> <pkg>
座標ハンドシェイクが失敗する	handshake_dim のインデックスが次元の位置と一致しているか確認する
イテレータの形状が正しくない	ランダム入力でリシェイプロジックをデバッグする
ModuleNotFoundError: pytest	pytest ではなく uv run pytest を使用する

---

注意事項

すべきこと:

• すべての Python コマンドに uv run python を使用する
• print() ではなく loguru.logger を使用する
• torch.nn.Module + AutoModelMixin + PrognosticMixin を継承する
• create_iterator で最初に初期条件を yield する
• _default_generator で front_hook()/rear_hook() を使用する
• すべての .py ファイルに SPDX ヘッダーを含める

すべきでないこと:

• 再利用のための汎用基底クラスを作成する
• API キーや比較スクリプトをコミットする
• evals/targets/ から読み取る