geoxplain-aurora-adapter 1.0.0

On-the-fly XAI computation with remote backends for the Aurora weather model

GeoXplain Aurora Adapter

GeoXplain	GeoXplain Aurora Adapter	Documentation	Live demo
Core toolkit	Current repository	User guide and API	Hosted viewer

Cite us

Will be published shortly.

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.

Profile	Use when	Install
client	This machine only submits requests to a listener.	[client]
local	Notebooks run directly inside a GPU allocation.	[gpu]
gpu-listener	This machine runs an HTTP listener inside an existing GPU allocation.	[gpu,server]
login-node	A login-node listener submits oneshot jobs or keeps a persistent GPU worker warm.	login node: [server]; GPU worker: [gpu,server,client]

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

Mode	Where it runs	When to use
local	Notebook process on a GPU node	You already have a GPU allocation.
gpu-listener	HTTP listener inside a GPU allocation	Clients can reach the GPU node directly or through a tunnel.
sbatch-oneshot	Login-node listener plus one SLURM job per request	Default sbatch-backed mode; no warm worker needed.
sbatch-persistent	Login-node listener plus one long-lived GPU worker	Faster repeated calls after model warmup.

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:

Endpoint	Purpose
POST /run	Submit one target.
POST /run_batch	Submit multiple timeframes.
GET /jobs/{job_id}	Poll status, ETA, progress, and log tail.
GET /jobs/{job_id}/result	Fetch the packed XiaResult.
GET /health	Inspect backend mode, queue depth, and resolved config.

More Details

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