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

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:

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

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

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:

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