GeoXplain Aurora Adapter

On-the-fly XIA attribution for the Microsoft Aurora weather model, with one Python API that can run locally, through a GPU listener, or through a SLURM-backed listener.

Supports adapted versions of Saliency, Integrated Gradients, RISE and Vit-CX.

import geoxplain_aurora_adapter as ax

result = ax.run_saliency(
    target=ax.Target.point(
        var="q",
        level=850,
        lat=46.2,
        lon=8.8,
        timestamp="2024-03-20T00:00:00Z",
    ),
    input=["t", "q", "z"],
    # remote=None                    # local GPU process
    # remote="http://gpu01:8765"     # GPU listener
    # remote="http://localhost:8765" # tunnel to login-node listener
)

result.save("ticino_q850_saliency.xia.npz")

XiaResult is a self-describing .xia.npz bundle that can be passed directly to the GeoXplain Aurora visualization widget.

Recommended Setup

Install the base package first. A plain install is enough to run the CLI setup guide.

pip install geoxplain-aurora-adapter
geoxplain-aurora-adapter

From a source checkout, use python -m pip install -e . instead.

geoxplain-aurora-adapter setup asks for the preferred deployment mode, writes the config needed, and prints what to install where.

Shortcut setup commands:

geoxplain-aurora-adapter setup --client
geoxplain-aurora-adapter setup --local
geoxplain-aurora-adapter setup --gpu-listener
geoxplain-aurora-adapter setup --login-node

The saved mode is only a preference. Listener runs can still override it with geoxplain-aurora-adapter listen --mode ...

More detailed installation notes are in docs/installation.md.

Start The Listener

For gpu-listener or sbatch installs, start the listener with:

geoxplain-aurora-adapter listen

CLI flags can override config for a single run:

geoxplain-aurora-adapter listen --yes \
    --mode sbatch-oneshot \
    --account PROJECT_ACCOUNT \
    --partition GPU_PARTITION \
    --time 00:30:00

To discard the saved listener config and rerun setup:

geoxplain-aurora-adapter setup --reset

Network binding and access

The listener's HTTP API is unauthenticated. It therefore binds 127.0.0.1 (loopback) by default and is meant to be reached over an SSH tunnel:

ssh -L 8765:localhost:8765 <login-node>   # then remote="http://localhost:8765"

Only bind a public address (--host 0.0.0.0) on a trusted/firewalled network; the listener prints a warning when you do.

Deployment Modes

The Python call site stays the same. Set remote="http://..." to delegate work to a listener.

API Sketch

target = ax.Target.box(
    var="q",
    level=850,
    lat=46.25,
    lon=8.75,
    size=(1.5, 2.5),
    timestamp="2020-04-20T12:00:00Z",
)

result = ax.run_saliency(target=target, input=["t", "q", "z"])
result = ax.run_ig(target=target, input=["t", "q", "z"], n_steps=32)
result = ax.run_rise(target=target, input=["t", "q", "z"], n_masks=200)
result = ax.run_vit_cx(target=target, input=["t", "q", "z"], n_clusters=256)

Batch timeframes return one multi-frame XiaResult:

result = ax.run_saliency(
    target=target,
    input=["t", "q", "z"],
    timeframes=6,
    step_hours=6,
    remote="http://localhost:8765",
)

Weather overlays use the same local/remote dispatch path:

overlay = ax.pull_overlay(
    "q",
    "2024-04-20",
    level=850,
    remote="http://localhost:8765",
    name="Specific Humidity 850 hPa",
    unit="kg/kg",
)

Omit dates to infer them from the explanations run this session. With overlay_time, this can be shifted by a fixed amount.

Result Format

Results are saved as .xia.npz archives. Each XiaResult contains one or more frames, and each frame carries:

• target metadata
• timestamp
• attribution maps keyed by input variable and vertical layer
• per-frame metadata such as target score or runtime

result.save("case.xia.npz")
restored = ax.XiaResult.load("case.xia.npz")

Transport

Remote execution uses the same FastAPI/msgpack protocol for GPU listeners and SLURM-backed listeners:

More Details

• Installation and setup
• Remote execution and retention
• Weather overlays