nvidia/earth2studio-create-diagnostic

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

• 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 user approves.

Step 3 - Create Model File

File: earth2studio/models/dx/<lowercase>.py

Use the repo-standard SPDX/license header shown in existing model files.

Simple diagnostic imports commonly include:

python
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

Packaged and generative diagnostics commonly also include:

python
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

Canonical method order:

1. __init__
2. input_coords
3. output_coords decorated with @batch_coords()
4. __str__ if useful
5. load_default_package for AutoModel/generative diagnostics
6. load_model for AutoModel/generative diagnostics
7. to only when non-PyTorch state must move devices
8. Private/support methods
9. __call__ decorated with @torch.inference_mode() and @batch_func()

Avoid shared base classes or broad abstractions unless the wrapper naturally has multiple closely related variants where a small base class reduces duplication.

Step 4 - Implement Coordinates

Diagnostic input coordinates usually use this public Earth2Studio order:

1. batch: np.empty(0) and first in the OrderedDict
2. variable: input variable names using Earth2Studio vocabulary names
3. lat: public latitude convention north-to-south, usually 90 to -90
4. lon: public longitude convention 0 to 360, endpoint normally false

No diagnostic wrapper should expose lead_time. If a diagnostic needs validity time metadata, document it as per-sample metadata in coords["time"]; do not make it a tensor dimension unless an existing dx pattern requires it.

output_coords must validate inputs with handshake_dim and handshake_coords. Then update output variables and, when needed, output lat/lon resolution. Generative diagnostics must add a sample dimension after batch.

Step 5 - Implement Forward Pass

Use a single-step __call__; never create an iterator. Validate coordinates before model execution, then return (output_tensor, output_coords).

python
@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

For generative diagnostics, loop over the batch dimension and generate number_of_samples per input item. Use explicit seeds for reproducibility when the reference implementation supports seeded sampling.

Step 6 - Implement Model Loading

For packaged diagnostics:

• load_default_package should lock HuggingFace URLs to a commit (hf://org/repo@commit) or NGC/S3 versions to an immutable release.
• load_model should call package.resolve(...), load checkpoints on CPU first, set modules to eval(), and disable gradients where appropriate.
• Use weights_only=False only when loading a pickled full PyTorch object is required.
• Decorate optional model classes and load_model with @check_optional_dependencies().
• Use loguru.logger for useful loading messages; do not use print() inside earth2studio/.

Step 7 - Write Tests

File: test/models/dx/test_<name>.py

Required tests:

| Function | Purpose | |----------|---------| | test_<model>_call | Forward pass with mock or simple model | | test_<model>_exceptions | Invalid coordinate order, values, or variables raise errors | | test_<model>_package | Real weights with @pytest.mark.package for AutoModel/generative diagnostics |

Generative diagnostics also require sample-count and deterministic-seed tests. Use references/testing-guide.py. Create a Phoo<ModelName> dummy that matches the real core model's interface and produces deterministic output.

Run focused tests:

bash
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

Do not omit package tests for packaged models. If arbitrary random inputs are not physically valid for the real checkpoint, build a stable model-appropriate input while still loading real weights and running a forward pass.

Step 8 - Register Model

For public models, update earth2studio/models/dx/__init__.py alphabetically. Skip registration only when the user explicitly wants an internal or experimental file that should not be exported.

Step 9 - Documentation

For public models:

• Add to docs/modules/models_dx.rst alphabetically so API docs include the generated page.
• Add to docs/userguide/about/install.md if a model extra exists. Include model notes plus both pip install earth2studio[model-name] and uv add earth2studio --extra model-name instructions.
• Update CHANGELOG.md under ### Added.

Format and lint:

bash
make format && make lint && make license

Step 10 - Validation (if requested)

Follow references/validation-guide.md. Create uncommitted vanilla, Earth2Studio, comparison, and sanity-check scripts. Do not commit generated outputs, checkpoints, images, or local validation scripts.

For generative diagnostics, fix seeds and compare matching samples or report statistical/tolerance-based agreement when exact equality is impossible. Ask the user to visually inspect plots before proceeding.

Step 11 - PR (if requested)

Follow references/validation-guide.md and use:

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

Before creating the PR, verify dependency extras, all, install docs, API docs, changelog, tests, and validation artifacts are consistent. Do not include machine names, hostnames, absolute paths, cache paths, device inventory, or uploaded image links in PR text. Use plot placeholders for manual image upload.

Examples

Simple Diagnostic

text
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 Diagnostic

text
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.

Generative Diagnostic

text
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.

Troubleshooting

| Error | Solution | |-------|----------| | OptionalDependencyFailure | Install with uv sync --extra <model-extra> or fix the extra name | | Coordinate handshake fails | Check OrderedDict order and handshake_dim indices | | Wrong output shape | Verify output_coords lengths match returned tensor shape | | ModuleNotFoundError: pytest | Use uv run pytest, not bare pytest | | Package test fails on random input | Use a stable physically plausible input while still loading real weights |

Reminders

Do:

• Use uv run python and uv run pytest for all Python commands.
• Use @batch_coords() on output_coords.
• Use @torch.inference_mode() and @batch_func() on __call__.
• Keep batch as the first coordinate with np.empty(0) in input_coords.
• Validate coordinates with handshake_dim() and handshake_coords().
• Add sample in generative output_coords.
• Include the repo-standard SPDX/license header in every Python file.
• Use loguru.logger, never print(), inside earth2studio/.

Do not:

• Inherit from PrognosticMixin.
• Include lead_time coordinates.
• Create create_iterator.
• Create general base classes for a single wrapper without a clear multi-variant need.
• Commit API keys, credentials, validation scripts, plots, or generated outputs.
• Read from evals/targets/.

Self-Improvement

If this skill produces incorrect outputs, update it before continuing:

1. Identify the issue in the generated code or workflow.
2. Edit SKILL.md or the relevant file in references/ to fix the guidance.
3. Run focused validation for the changed skill files.
4. Commit the skill fix separately when working in a branch that expects commits.
5. Continue the model implementation with the corrected workflow.