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

Workspace

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

Reference Files

Load these files on demand during the matching workflow:

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:

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.