earth2studio-create-diagnostic

Files (20)

• BENCHMARK.md
• SKILL.md
• config.yml · evals/
• Dockerfile · evals/environment/
• bootstrap.sh · evals/environment/setup/
• evals.json · evals/
• eval_1_target.py · evals/targets/
• eval_2_target.py · evals/targets/
• eval_3_target.py · evals/targets/
• eval_1_test_target.py · evals/targets/test/
• eval_2_test_target.py · evals/targets/test/
• eval_3_test_target.py · evals/targets/test/
• method-templates.py · references/
• pr-body-template.md · references/
• pr-comment-template.md · references/
• skeleton-template.py · references/
• testing-guide.py · references/
• validation-guide.md · references/
• skill-card.md
• skill.oms.sig

SKILL.md (first lines)

## 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`:

```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`:

From NVIDIA/skills, Apache-2.0.

More from NVIDIA/skills