earth2studio-create-diagnostic

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

以下の手順を順番どおりに実行してください。先に進んで手順を飛ばさないでください。編集を始める前に、この SKILL.md を読み、モデルの種類に応じた関連リファレンスファイルを読み込んでください。実装後、テストに合格したと報告する前に、対象を絞った pytest コマンドを実行してください。テストを実行できない場合は、成功したと主張せず、正確なコマンドと失敗内容を報告してください。

• まず、この SKILL.md を最後まで読む
• リファレンススクリプト、リポジトリ、論文、またはモデルのドキュメントを入手する（ステップ 0）
• 診断モデルを simple、AutoModel、または generative に分類する（ステップ 1）
• 依存関係ファイルを編集する前に、依存関係の追加オプションを提案する（ステップ 1）
• 診断専用 API を備えた earth2studio/models/dx/<name>.py を作成する
• モックテストを備えた test/models/dx/test_<name>.py を作成する
• 次を実行する：uv run pytest test/models/dx/test_<name>.py -m "not package" -v
• 必要に応じて、モデルの追加オプション、インストール手順のドキュメント、API ドキュメント、変更履歴を追加・更新する
• make format && make lint && make license を実行する

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

• uv run pytest ... と uv run python ... を使用してください
• リポジトリのワークフローでは、pytest や python を単独で使用しないでください

生成したモデルが正しくない場合、同じ修正を何度も試さないでください。 Self-Improvement に従って、このスキルまたはそのリファレンスを修正してから、修正後のワークフローを続けてください。

目的

サードパーティまたは派生の ML 変換を Earth2Studio に接続する診断モデルのラッパーを実装します。診断モデルは単一時点のデータを変換します。入力フィールドを受け取り、出力フィールドを返します。予報の時間積分は行いません。

前提条件

• 開発用依存関係を含め、uv で Earth2Studio をインストール済みである（uv sync --all-extras）
• Python 3.10 以降の環境
• リファレンス推論スクリプト、リポジトリ、論文、またはモデルのドキュメント
• パッケージに含めるモデルのチェックポイントの取得元とライセンス情報

制限事項

• 単一ステップの変換のみを扱います
• 時間ステップを進める予報モデルには対応していません。その場合は earth2studio-create-prognostic を使用してください
• 実際のパッケージを使ったテストでは、NGC、HuggingFace、S3、またはその他のレジストリへのネットワークアクセスが必要になる場合があります
• 生成モデルの検証では、有意義な比較のために GPU と固定シードが必要になる場合があります

診断モデルの種類

種類	継承	依存関係の追加オプション	例
単純な派生診断モデル	torch.nn.Module のみ	通常は不要	Identity、風速
パッケージ済み AutoModel 診断モデル	torch.nn.Module, AutoModelMixin	必須（空でも必要）	PrecipitationAFNO
生成診断モデル	torch.nn.Module, AutoModelMixin	必須（空でも必要）	CorrDiff

ワークスペース

コンテキスト	場所
Harbor eval	/workspace/output/earth2studio/models/dx/... に書き込む
Harbor + --copy-repo	/workspace/repo に完全なチェックアウト
ローカルクローン	pyproject.toml があるディレクトリ

evals/targets/ は絶対に読まないでください。これらのファイルは採点者の参照用です。

リファレンスファイル

該当するワークフローの途中で、必要に応じて以下のファイルを読み込んでください。

ファイル	内容	読み込む手順
references/skeleton-template.py	simple、AutoModel、generative ラッパー用の診断モデルの全体的な骨組み	ステップ 3-6
references/method-templates.py	座標、読み込み、forward、デバイスメソッドの要点をまとめたコード例	ステップ 4-6
references/testing-guide.py	モック、パッケージ、例外、サンプル、シードのテストパターン	ステップ 7
references/validation-guide.md	リファレンスとの比較、プロット、PR の衛生管理、レビュー後の対応	ステップ 10-11
references/pr-body-template.md	PR 本文のテンプレート	ステップ 11
references/pr-comment-template.md	検証コメントのテンプレート	ステップ 11

手順

ステップ 0 - 参考資料の取得

$ARGUMENTS に URL またはローカルパスが指定されている場合は、それを使用します。指定されていない場合は、次のように尋ねます。

参考となる推論スクリプト、リポジトリ、論文、またはモデルのドキュメントを提供してください。

参考モデルの入力変数、出力変数、テンソル形状、正規化、グリッド、チェックポイントの取得元、依存関係の要件、ライセンスを記録します。

ステップ 1 - 種類を分析し、依存関係を提案する

ファイルを編集する前に、要求された診断の種類を分類します。

モデルが…	使用するもの…
チェックポイントなしで派生量を計算する	シンプルな診断
Package または外部チェックポイントから重みを読み込む	AutoModel 診断
複数のサンプル、拡散出力、VAE サンプル、または確率的な超解像を生成する	生成系診断

依存関係のポリシー：

• シンプルな派生診断では通常、pyproject.toml の extra は不要です。
• AutoModel 診断と生成系診断では、リストが空でも、名前付きのオプション依存関係 extra が必要です。
• [project.optional-dependencies] 内で extra をアルファベット順に追加し、all の集約リストにも含めます。
• OptionalDependencyFailure("model-extra") と @check_optional_dependencies() では、モデル extra の名前を使用します。

pyproject.toml を編集する前に、依存関係 extra の提案を提示してユーザーの承認を求めます。

model-name = ["package1>=version", "package2"]
# or, when the packaged diagnostic needs no extra runtime packages:
model-name = []

ステップ 2 - 依存関係を追加する

承認を得た後、pyproject.toml を編集します。

• extra をアルファベット順に追加します。
• all の集約リストを更新します。
• 参考パッケージのドキュメントにある最小サポートバージョンを優先します。
• 参考モデルで必要とされ、ユーザーが承認した場合を除き、広範な未固定の Git 依存関係を追加しないでください。

ステップ 3 - モデルファイルを作成する

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

既存のモデルファイルにある、リポジトリ標準の SPDX／ライセンスヘッダーを使用します。

シンプルな診断では、一般に次のものをインポートします。

from collections import OrderedDict
import numpy as np
import torch
from earth2studio.models.batch import batch_coords, batch_func
from earth2studio.utils import handshake_coords, handshake_dim
from earth2studio.utils.type import CoordSystem

パッケージ化された診断と生成系診断では、一般に次のものもインポートします。

from earth2studio.models.auto import AutoModelMixin, Package
from earth2studio.models.dx.base import DiagnosticModel
from earth2studio.utils.imports import OptionalDependencyFailure, check_optional_dependencies
from loguru import logger

標準のメソッド順：

1. __init__
2. input_coords
3. @batch_coords() でデコレートした output_coords
4. 必要に応じて __str__
5. AutoModel／生成系診断用の load_default_package
6. AutoModel／生成系診断用の load_model
7. PyTorch 以外の状態をデバイス間で移動する場合に限り to
8. プライベートメソッド／サポートメソッド
9. @torch.inference_mode() と @batch_func() でデコレートした __call__

ラッパーに複数の密接に関連するバリエーションが自然に存在し、小さな基底クラスで重複を減らせる場合を除き、共有基底クラスや広範な抽象化は避けてください。

ステップ 4 - 座標を実装する

診断の入力座標では通常、次の公開 Earth2Studio 順序を使用します。

1. batch：np.empty(0) とし、OrderedDict の先頭に置きます。
2. variable：Earth2Studio の用語に従った入力変数名
3. lat：公開緯度規約の北から南。通常は 90 から -90
4. lon：公開経度規約の 0 から 360。通常、終端値は含めません。

診断ラッパーで lead_time を公開してはいけません。診断に有効時刻のメタデータが必要な場合は、coords["time"] のサンプルごとのメタデータとして説明し、既存の dx パターンで必要な場合を除いて、テンソル次元にはしないでください。

output_coords では、handshake_dim と handshake_coords を使用して入力を検証する必要があります。その後、出力変数と、必要に応じて出力の緯度／経度の解像度を更新します。生成系診断では、batch の後に sample 次元を追加する必要があります。

ステップ 5 - フォワードパスを実装する

単一ステップの __call__ を使用し、イテレーターは作成しないでください。モデルを実行する前に座標を検証し、その後 (output_tensor, output_coords) を返します。

@torch.inference_mode()
@batch_func()
def __call__(self, x: torch.Tensor, coords: CoordSystem) -> tuple[torch.Tensor, CoordSystem]:
    output_coords = self.output_coords(coords)
    x = (x - self.center) / self.scale
    out = self.core_model(x)
    return out, output_coords

生成系診断では、バッチ次元をループして、各入力項目につき number_of_samples 個のサンプルを生成します。参考実装がシード付きサンプリングに対応している場合は、再現性のために明示的なシードを使用してください。

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

パッケージ化された診断の場合:

• load_default_package は、HuggingFace の URL をコミットに固定するか (hf://org/repo@commit)、NGC/S3 のバージョンをイミュータブルなリリースに固定する必要があります。
• load_model は package.resolve(...) を呼び出し、まず CPU 上でチェックポイントを読み込み、モジュールを eval() に設定し、適切な箇所で勾配を無効化する必要があります。
• ピクル化された完全な PyTorch オブジェクトの読み込みが必要な場合にのみ、weights_only=False を使用してください。
• オプションのモデルクラスと load_model に @check_optional_dependencies() を付けてください。
• 有用な読み込みメッセージには loguru.logger を使用してください。earth2studio/ 内で print() を使用しないでください。

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

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

必須のテスト:

関数	目的
test_<model>_call	モックまたは単純なモデルによるフォワードパス
test_<model>_exceptions	無効な座標順序、値、または変数がエラーを発生させる
test_<model>_package	AutoModel/生成診断向けの @pytest.mark.package を使用した実重み

生成診断には、サンプル数と決定論的シードのテストも必要です。 references/testing-guide.py を使用してください。実際のコアモデルのインターフェースに一致し、 決定論的な出力を生成する Phoo<ModelName> ダミーを作成してください。

フォーカスしたテストの実行:

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

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

ステップ 8 - モデルの登録

公開モデルの場合は、earth2studio/models/dx/__init__.py をアルファベット順に更新してください。 エクスポートされるべきでない内部用または実験用ファイルをユーザーが明示的に希望する場合にのみ、 登録をスキップしてください。

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

公開モデルの場合:

• API ドキュメントに生成されたページが含まれるように、docs/modules/models_dx.rst にアルファベット順で追加してください。
• モデルの extra が存在する場合は、docs/userguide/about/install.md に追加してください。モデルの注意事項に加えて、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、 Earth2Studio、比較、およびサニティチェック用のスクリプトを作成してください。生成された 出力、チェックポイント、画像、またはローカル検証スクリプトをコミットしないでください。

生成診断の場合は、シードを固定して一致するサンプルを比較するか、完全な一致が不可能な場合は 統計的または許容範囲ベースの一致を報告してください。 先に進む前に、ユーザーにプロットを目視で確認するよう依頼してください。

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

references/validation-guide.md に従い、以下を使用してください:

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

PR を作成する前に、依存関係の extra、all、インストールドキュメント、API ドキュメント、 変更履歴、テスト、および検証成果物が一貫していることを確認してください。PR のテキストに マシン名、ホスト名、絶対パス、キャッシュパス、デバイスインベントリ、またはアップロードされた 画像リンクを含めないでください。手動での画像アップロードにはプロットのプレースホルダーを 使用してください。

例

シンプルな診断

User: Create a diagnostic that computes wind speed from u10m and v10m.
Agent: Reads SKILL.md, classifies as simple, creates windspeed.py with only
       torch.nn.Module, writes call and exception tests, runs focused pytest.

AutoModel 診断

User: Add a precipitation estimator from this reference script.
Agent: Reads SKILL.md and references, proposes dependency extra, creates a
       torch.nn.Module + AutoModelMixin wrapper, writes mock/package tests,
       updates docs/changelog/dependencies, and runs validation commands.

生成診断

User: Wrap this diffusion super-resolution model.
Agent: Classifies as generative, adds sample output coordinates, supports seed
       handling, writes sample and deterministic-seed tests, and prepares seeded
       validation comparisons.

トラブルシューティング

エラー	解決策
OptionalDependencyFailure	uv sync --extra <model-extra> でインストールするか、extra 名を修正してください
座標ハンドシェイクが失敗する	OrderedDict の順序と handshake_dim のインデックスを確認してください
出力形状が正しくない	output_coords の長さが返されるテンソルの形状と一致することを確認してください
ModuleNotFoundError: pytest	素の pytest ではなく uv run pytest を使用してください
パッケージテストがランダム入力で失敗する	実際の重みを読み込みながら、安定した物理的に妥当な入力を使用してください

リマインダー

行うこと：

• すべての Python コマンドで uv run python と uv run pytest を使用してください。
• output_coords に @batch_coords() を使用してください。
• __call__ に @torch.inference_mode() と @batch_func() を使用してください。
• input_coords では、batch を最初の座標とし、np.empty(0) を指定してください。
• handshake_dim() と handshake_coords() で座標を検証してください。
• 生成系の output_coords に sample を追加してください。
• すべての Python ファイルに、リポジトリ標準の SPDX／ライセンスヘッダーを含めてください。
• earth2studio/ 内では print() を使用せず、loguru.logger を使用してください。

行わないこと：

• PrognosticMixin を継承しないでください。
• lead_time 座標を含めないでください。
• create_iterator を作成しないでください。
• 明確な複数バリアントの必要性がない限り、単一のラッパー用に汎用基底クラスを作成しないでください。
• API キー、認証情報、検証スクリプト、プロット、生成された出力をコミットしないでください。
• evals/targets/ から読み込まないでください。

自己改善

このスキルが誤った出力を生成した場合は、作業を続ける前に更新してください。

1. 生成されたコードまたはワークフローの問題を特定します。
2. ガイダンスを修正するため、SKILL.md または references/ 内の該当ファイルを編集します。
3. 変更したスキルファイルに対して、対象を絞った検証を実行します。
4. コミットが必要なブランチで作業している場合は、スキルの修正を個別にコミットします。
5. 修正したワークフローでモデルの実装を続けます。