/earth2studio-create-diagnostic

name: earth2studio-create-diagnostic
version: 0.16.0
license: Apache-2.0
metadata:
  author: NVIDIA Earth-2 Team <agent-skills@nvidia.com>
  tags: [earth2studio, diagnostic-model, python]
description: >
  Create Earth2Studio diagnostic model wrappers for single-step data
  transformations, including simple derived diagnostics, packaged AutoModel
  diagnostics, and generative or diffusion diagnostics. Do NOT use for
  prognostic time-stepping models, data sources, or installation.
argument-hint: URL or local path to reference inference script (optional)

Quick Start Checklist

Do these steps in order. Do not skip ahead. Before editing, read this SKILL.md and load the relevant reference files for the model type. After implementation, run the focused pytest command before saying tests pass. If tests cannot run, report the exact command and failure instead of claiming success.

• [ ] Read this SKILL.md completely first
• [ ] Get the reference script, repo, paper, or model documentation (Step 0)
• [ ] Classify the diagnostic as simple, AutoModel, or generative (Step 1)
• [ ] Propose dependency extras before editing dependency files (Step 1)
• [ ] Create `earth2studio/models/dx/<name>.py` with diagnostic-only APIs
• [ ] Create `test/models/dx/test_<name>.py` with mock tests
• [ ] Run: `uv run pytest test/models/dx/test_<name>.py -m "not package" -v`
• [ ] Add/update model extra, install docs, API docs, and changelog when required
• [ ] Run: `make format && make lint && make license`

Critical command rule: always use `uv run` for Python commands:

• Use `uv run pytest ...` and `uv run python ...`
• Do not use bare `pytest` or `python` in repo workflows

If the generated model is wrong, do not keep retrying the same fix. Follow [Self-Improvement](#self-improvement), patch this skill or its references, then continue with the corrected workflow.

Purpose

Implement a diagnostic model wrapper connecting third-party or derived ML transforms to Earth2Studio. Diagnostic models transform data at a single time point: input fields in, output fields out, no forecast integration.

Prerequisites

• Earth2Studio installed via `uv` with dev dependencies (`uv sync --all-extras`)
• Python 3.10+ environment
• Reference inference script, repo, paper, or model documentation
• Checkpoint source and license information for packaged models

Limitations

• Handles single-step transformations only
• Does not support time-stepping forecast models; use `earth2studio-create-prognostic`
• Real package tests can require network access to NGC, HuggingFace, S3, or other registries
• Generative validation can require GPU and fixed seeds for meaningful comparison

Diagnostic Model Types

| Type | Inheritance | Dependency extra | Example | |------|-------------|------------------|---------| | Simple derived diagnostic | `torch.nn.Module` only | Usually none | `Identity`, wind speed | | Packaged AutoModel diagnostic | `torch.nn.Module, AutoModelMixin` | Required, even if empty | `PrecipitationAFNO` | | Generative diagnostic | `torch.nn.Module, AutoModelMixin` | Required, even if empty | `CorrDiff` |

Workspace

| Context | Location | |---------|----------| | Harbor eval | Write to `/workspace/output/earth2studio/models/dx/...` | | Harbor + `--copy-repo` | Full checkout at `/workspace/repo` | | Local clone | Directory with `pyproject.toml` |

Never read `evals/targets/`; those files are grader references only.

Reference Files

Load these files on demand during the matching workflow:

| File | Content | Load at | |------|---------|---------| | `references/skeleton-template.py` | Full diagnostic skeletons for simple, AutoModel, and generative wrappers | Steps 3-6 | | `references/method-templates.py` | Focused coordinate, loading, forward, and device method snippets | Steps 4-6 | | `references/testing-guide.py` | Mock, package, exception, sample, and seed test patterns | Step 7 | | `references/validation-guide.md` | Reference comparison, plots, PR hygiene, and review follow-up | Steps 10-11 | | `references/pr-body-template.md` | PR body template | Step 11 | | `references/pr-comment-template.md` | Validation comment template | Step 11 |

Instructions

Step 0 - Get Reference Material

If `$ARGUMENTS` provides a URL or local path, use it. Otherwise ask:

> Please provide a reference inference script, repository, paper, or model documentation.

Capture the reference model's input variables, output variables, tensor shapes, normalization, grid, checkpoint source, dependency requirements, and license.

Step 1 - Analyze Type and Propose Dependencies

Classify the requested diagnostic before editing files:

| If the model... | Then use... | |-----------------|-------------| | Computes a derived quantity with no checkpoint | Simple diagnostic | | Loads weights from `Package` or an external checkpoint | AutoModel diagnostic | | Produces multiple samples, diffusion outputs, VAE samples, or stochastic super-resolution | Generative diagnostic |

Dependency policy:

• Simple derived diagnostics usually do not need a `pyproject.toml` extra.
• AutoModel and generative diagnostics must have a named optional dependency extra, even if the list is empty.
• Add the extra alphabetically under `[project.optional-dependencies]` and include it in the `all` aggregate.
• Use the model-extra name in `OptionalDependencyFailure("model-extra")` and `@check_optional_dependencies()`.

Present the proposed dependency extra and ask the user to approve before editing `pyproject.toml`:

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

Step 2 - Add Dependencies

After approval, edit `pyproject.toml`:

• Add the extra alphabetically.
• Update the `all` aggregate.
• Prefer minimum supported versions from the reference package documentation.
• Do not add broad unpinned Git dependencies unless the reference model requires them and the