Aurora adapter Python API¶

This backend-specific reference is generated from the sibling geoxplain-aurora-adapter checkout used to build the site. These APIs target Microsoft Aurora; they are separate from the model-agnostic GeoXplain Python API. The package's root exports these target, compute, result, overlay, and listener APIs.

Targets¶

geoxplain_aurora_adapter.schema.spec.Target ¶

Namespace for TargetSpec factory methods.

Usage::

import geoxplain_aurora_adapter as ax

# Single grid point
target = ax.Target.point(var="q", level=850, lat=46.2, lon=8.8,
                         timestamp="2024-03-20T00:00:00Z")

# Box of default size (2.0° lat × 3.0° lon) centered at (lat, lon)
target = ax.Target.box(var="q", level=850, lat=46.25, lon=8.75,
                       timestamp="2024-03-20T00:00:00Z")

# Box of custom size
target = ax.Target.box(var="q", level=850, lat=46.25, lon=8.75,
                       size=(1.5, 2.5),
                       timestamp="2024-03-20T00:00:00Z")

geoxplain_aurora_adapter.schema.spec.TargetSpec `dataclass` ¶

TargetSpec(
    var: str,
    level: Optional[int],
    mode: str,
    timestamp: str,
    lat: Optional[float] = None,
    lon: Optional[float] = None,
    size: Optional[tuple[float, float]] = None,
)

Fully-resolved specification of the XIA attribution target.

Fields

var: Variable name in the model output (e.g. "q", "t", "zwd"). level: Pressure level in hPa (e.g. 850). None for surface-only vars. mode: Spatial selection mode: "point" or "box". timestamp: ISO-8601 string for the second (t1) input timestep, e.g. "2024-03-20T00:00:00Z". This is what you pass when constructing a target, and it is preserved unchanged as the frame's displayed timestamp (frame.target.timestamp == frame.timestamp == t1). The explained prediction is the 6 h-ahead step t2 = t1 + lead, recorded in the frame's lead_hours metadata rather than by shifting the timestamp. lat/lon: Point coordinates (mode="point") or box center (mode="box"). Longitudes accepted in either -180..180 or 0..360 convention. size: Box full extent in degrees (dlat, dlon) (mode="box" only). The actual bounds are [lat-dlat/2, lat+dlat/2] × [lon-dlon/2, lon+dlon/2].

point `classmethod` ¶

point(
    *,
    var: str,
    level: Optional[int],
    lat: float,
    lon: float,
    timestamp: str,
) -> "TargetSpec"

Single-grid-point target.

box `classmethod` ¶

box(
    *,
    var: str,
    level: Optional[int],
    lat: float,
    lon: float,
    timestamp: str,
    size: tuple[float, float] = DEFAULT_BOX_SIZE,
) -> "TargetSpec"

Box-mean target centered at (lat, lon) with extent size.

box_bounds ¶

box_bounds() -> tuple[float, float, float, float]

Return (south, north, west, east) for mode="box".

to_dict ¶

to_dict() -> dict

from_dict `classmethod` ¶

from_dict(d: dict) -> 'TargetSpec'

as_widget_dict ¶

as_widget_dict() -> dict

Return the target in the model-agnostic GeoXplain viewer format.

Compute methods¶

geoxplain_aurora_adapter.api.methods.run_saliency ¶

run_saliency(
    target: TargetSpec,
    input: list[str],
    *,
    remote: Optional[str] = None,
    checkpoint_path: Optional[str] = None,
    timeframes: int = 1,
    step_hours: int = 6,
    levels: Optional[int | list[int]] = None,
    **options,
) -> XiaResult

Compute vanilla gradient saliency attribution.

Parameters:

Name	Type	Description	Default
`target`	`TargetSpec`		required
`input`	`list[str]`		required
`remote`	`Optional[str]`	(e.g. ``"http://gpu01:8765"`` or ``"http://localhost:8765"`` for an SSH-tunnelled cluster). If ``None``, runs in-process on the local GPU.	`None`
`checkpoint_path`	`Optional[str]`		`None`
`timeframes`	`int`	`greater than 1 return one multi-frame XiaResult.`	`1`
`step_hours`	`int`		`6`
`levels`	`Optional[int \| list[int]]`	``[925, 850, 700]``. ``None`` (the default) returns every level in :data:`AURORA_LEVELS`. Only affects atmospheric input variables; surface variables are unaffected. The single backward pass computes gradients for all levels regardless, so this filters the output rather than reducing GPU cost.	`None`
`**options`			`{}`

Returns:

Type	Description
`XiaResult`	One frame, or a multi-frame bundle when `timeframes > 1`.

geoxplain_aurora_adapter.api.methods.run_ig ¶

run_ig(
    target: TargetSpec,
    input: list[str],
    *,
    remote: Optional[str] = None,
    checkpoint_path: Optional[str] = None,
    timeframes: int = 1,
    step_hours: int = 6,
    n_steps: int = 32,
    baseline_sigma_deg: float = 2.5,
    levels: Optional[int | list[int]] = None,
    **options,
) -> XiaResult

Compute Integrated Gradients attribution.

Parameters:

Name	Type	Description	Default
`target`	`TargetSpec`		required
`input`	`list[str]`		required
`remote`	`Optional[str]`		`None`
`checkpoint_path`	`Optional[str]`		`None`
`timeframes`	`int`		`1`
`step_hours`	`int`		`6`
`n_steps`	`int`		`32`
`baseline_sigma_deg`	`float`		`2.5`
`levels`	`Optional[int \| list[int]]`	``[925, 850, 700]``. ``None`` (the default) returns every level in :data:`AURORA_LEVELS`. Only affects atmospheric input variables; surface variables are unaffected. Each integration step's backward pass computes gradients for all levels regardless, so this filters the output rather than reducing GPU cost.	`None`

Returns:

Type	Description
`XiaResult`	One frame, or a multi-frame bundle when `timeframes > 1`.

geoxplain_aurora_adapter.api.methods.run_rise ¶

run_rise(
    target: TargetSpec,
    input: list[str],
    *,
    remote: Optional[str] = None,
    checkpoint_path: Optional[str] = None,
    timeframes: int = 1,
    step_hours: int = 6,
    n_masks: int = 1200,
    cells_h: int = 400,
    cells_w: int = 800,
    seed: int = 42,
    baseline_sigma_deg: float = 2.5,
    **options,
) -> XiaResult

Compute RISE (Randomized Input Sampling for Explanation) attribution.

Parameters:

Name	Type	Default
`target`	`TargetSpec`	required
`input`	`list[str]`	required
`remote`	`Optional[str]`	`None`
`checkpoint_path`	`Optional[str]`	`None`
`timeframes`	`int`	`1`
`step_hours`	`int`	`6`
`n_masks`	`int`	`1200`
`cells_h`	`int`	`400`
`seed`	`int`	`42`
`baseline_sigma_deg`	`float`	`2.5`

Returns:

Type	Description
`XiaResult`	One frame, or a multi-frame bundle when `timeframes > 1`.

geoxplain_aurora_adapter.api.methods.run_vit_cx ¶

run_vit_cx(
    target: TargetSpec,
    input: list[str],
    *,
    remote: Optional[str] = None,
    checkpoint_path: Optional[str] = None,
    timeframes: int = 1,
    step_hours: int = 6,
    hook_stage: int = 1,
    n_clusters: int = 4096,
    smooth_sigma: float | tuple | None = 0,
    baseline_sigma_deg: float = 2.5,
    **options,
) -> XiaResult

Compute ViT-CX (cluster-based causal attribution) attribution.

Parameters:

Name	Type	Description	Default
`target`	`TargetSpec`		required
`input`	`list[str]`		required
`remote`	`Optional[str]`		`None`
`checkpoint_path`	`Optional[str]`		`None`
`timeframes`	`int`		`1`
`step_hours`	`int`		`6`
`hook_stage`	`int`		`1`
`n_clusters`	`int`	`per variable. Bounds run cost regardless of input.`	`4096`
`smooth_sigma`	`float \| tuple \| None`	Default ``0`` disables smoothing and keeps the raw cluster map. A non-zero scalar applies the same sigma to both axes; a ``(lat, lon)`` pair sets them independently.	`0`
`baseline_sigma_deg`	`float`		`2.5`

Returns:

Type	Description
`XiaResult`	One frame, or a multi-frame bundle when `timeframes > 1`.

geoxplain_aurora_adapter.api.methods.run_rollout ¶

run_rollout(
    target: TargetSpec,
    input: list[str],
    *,
    method: str = "saliency",
    timeframes: int,
    remote: Optional[str] = None,
    checkpoint_path: Optional[str] = None,
    n_steps: int = 32,
    baseline_sigma_deg: float = 2.5,
    timeout_s: float = 1800.0,
    poll_interval_s: float = 2.0,
    poll_max_s: float = 30.0,
) -> XiaResult

Compute autoregressive rollout XIA in fixed six-hour Aurora frames.

Parameters:

Name	Type	Description	Default
`target`	`TargetSpec`	Scalar output target for the first frame.	required
`input`	`list[str]`	Input variable names to attribute.	required
`method`	`str`	`"saliency"` or `"ig"`. RISE and ViT-CX are recognized method identifiers but are not implemented for rollout.	`'saliency'`
`timeframes`	`int`	Required positive number of autoregressive frames.	required
`remote`	`Optional[str]`	Listener URL, or `None` for in-process execution.	`None`
`checkpoint_path`	`Optional[str]`	Override the default checkpoint in local mode only.	`None`
`n_steps`	`int`	Integrated Gradients options, ignored for Saliency.	`32`
`baseline_sigma_deg`	`int`	Integrated Gradients options, ignored for Saliency.	`32`
`timeout_s`	`float`	Remote-client timeout and polling controls.	`1800.0`
`poll_interval_s`	`float`	Remote-client timeout and polling controls.	`1800.0`
`poll_max_s`	`float`	Remote-client timeout and polling controls.	`1800.0`

Returns:

Type	Description
`XiaResult`	One bundle containing all rollout frames.

Raises:

Type	Description
`ValueError`	If `method` is unknown or `timeframes` is not positive.
`NotImplementedError`	If rollout is requested with RISE or ViT-CX.

Weather overlays¶

geoxplain_aurora_adapter.api.methods.pull_overlay ¶

pull_overlay(
    variable: str,
    dates: str | list[str] | tuple[str, ...] | None = None,
    *,
    level: Optional[int] = None,
    remote: Optional[str] = None,
    name: Optional[str] = None,
    unit: Optional[str] = None,
    colormap: Optional[str] = None,
    visible: bool = True,
    overlay_time: str = "input",
    step_hours: int = 6,
    timeout_s: float = 1800.0,
    poll_interval_s: float = 2.0,
    poll_max_s: float = 30.0,
) -> OverlayResult

Pull ERA5 fields as timestamped viewer overlays.

dates accepts ISO timestamps, dates, ranges, or None to reuse the XIA frame timestamps recorded this session. The displayed frame keeps its requested timestamp (Aurora's most-recent input step t1), so overlay_time chooses the field time relative to it: "input" (0 h, the default — the frame's own time t1), "prior" (-6 h, the earlier input step t0), or "predicted" (+6 h, the forecast valid time t2). It also sets a time_label annotation on the result for the non-default choices ("prior" → "Aurora input step t0", "predicted" → "Forecast valid time t2"), which the viewer shows next to the offset. Remote calls use the listener at remote; local calls read from the configured dataset.

Attribution results¶

geoxplain_aurora_adapter.schema.result.XiaFrame `dataclass` ¶

XiaFrame(
    target: TargetSpec,
    timestamp: str,
    attributions: dict[str, dict[str, ndarray]],
    diverging: bool,
    meta: dict = dict(),
)

A single time step of an :class:XiaResult bundle.

Attributes:

Name	Type	Description
`target`	`TargetSpec`	The scalar that was explained (what variable, where, when).
`timestamp`	`str`	ISO-8601 string of the requested time — Aurora's most-recent input step t1 — which is what the viewer displays and equals `target.timestamp`. The explained prediction is the 6 h-ahead step t2 = t1 + lead; `meta["lead_hours"]` records the lead, and `meta["input_timestamp"]` echoes t1 for convenience.
`attributions`	`dict[str, dict[str, ndarray]]`	`attributions[wrt_var][level_key]` is a `(721, 1440)` float32 array recording the sensitivity of the target to the input field `wrt_var` at `level_key`. Atmospheric level keys have the form `"z-{N}"` (higher `N` = higher in the tool); surface variables use `"sfc"`.
`diverging`	`bool`	Whether the attribution maps contain significant negative values. Auto-detected from the sign distribution; controls whether the visualization widget uses a diverging colormap.
`meta`	`dict`	Optional per-frame diagnostic metadata: `"target_score"` (scalar value being explained, saliency/IG only), `"runtime_s"`, etc.

as_widget_dict ¶

as_widget_dict() -> dict

Return this frame's target as a geoxplain-compatible dict.

The visualization side's importer calls as_widget_dict() on each frame (not on frame.target) to place the target point/box. When an in-memory :class:XiaResult is handed straight to the widget — the live/remote path, with no .xia.npz round-trip — the frame object is this class, so it must expose the method itself; otherwise the target is silently dropped. Delegates to :meth:TargetSpec.as_widget_dict.

geoxplain_aurora_adapter.schema.result.XiaResult `dataclass` ¶

XiaResult(
    method: str,
    frames: list[XiaFrame],
    layer_labels: dict[str, str] = dict(),
    meta: dict = dict(),
    method_label: str = "",
)

Self-describing XIA attribution bundle (one method, one or more frames).

Attributes:

Name	Type	Description
`method`	`str`	XIA method id: `"saliency"`, `"ig"`, `"rise"`, or `"vit_cx"`. A stable machine identifier (the compute/dispatch layers branch on it).
`method_label`	`str`	Human-readable method name for display, e.g. `"Integrated Gradients"`. Empty when unknown; consumers fall back to `method` in that case.
`frames`	`list[XiaFrame]`	One :class:`XiaFrame` per time step. Single-frame results are the common case; use :meth:`single` to build one ergonomically.
`layer_labels`	`dict[str, str]`	Optional `{level_key: display_name}` map shared across frames, e.g. `{"z-2": "850 hPa", "sfc": "Surface"}`. Absent keys fall back to the bare number (`"z-2"` → `"2"`) or `"Surface"` for `"sfc"`.
`meta`	`dict`	Bundle-level diagnostic metadata: `"checkpoint_hash"`, `"host"`, `"slurm_job_id"`, etc.

single `classmethod` ¶

single(
    method: str,
    target: TargetSpec,
    timestamp: str,
    attributions: dict[str, dict[str, ndarray]],
    diverging: bool,
    meta: dict | None = None,
    *,
    layer_labels: dict[str, str] | None = None,
    frame_meta: dict | None = None,
    method_label: str = "",
) -> "XiaResult"

Build a one-frame bundle.

meta becomes the bundle-level metadata; frame_meta (if given) is attached to the single frame. For the common case where there is only one frame, passing diagnostics via meta is fine.

save ¶

save(path: str | PathLike) -> None

Save to a .xia.npz archive.

If path does not end with ".xia.npz", the suffix is appended automatically.

load `classmethod` ¶

load(path: str | PathLike) -> 'XiaResult'

Load from a .xia.npz archive (format_version == 2).

to_msgpack ¶

to_msgpack() -> bytes

Serialize to msgpack bytes for HTTP transport.

Arrays are embedded as raw float32 LE bytes + shape info. The result is designed to be deserialised by XiaResult.from_msgpack.

from_msgpack `classmethod` ¶

from_msgpack(data: bytes) -> 'XiaResult'

Deserialize from msgpack bytes produced by XiaResult.to_msgpack.

summary ¶

summary() -> str

One-line human-readable summary.

Overlay results¶

geoxplain_aurora_adapter.schema.overlay.OverlayFrame `dataclass` ¶

OverlayFrame(timestamp: str, data: ndarray)

One timestamped raw weather-field overlay frame.

geoxplain_aurora_adapter.schema.overlay.OverlayResult `dataclass` ¶

OverlayResult(
    variable: str,
    level: Optional[int],
    frames: list[OverlayFrame],
    label: str,
    unit: str = "",
    colormap: str = "viridis",
    visible: bool = True,
    overlay_offset_hours: int = 0,
    time_label: Optional[str] = None,
    lat: Optional[ndarray] = None,
    lon: Optional[ndarray] = None,
    meta: dict = dict(),
)

Self-contained weather-field overlay bundle.

frames are raw ERA5/Aurora-grid arrays. The visualization package applies its existing preprocessing pipeline when the overlay is added to a widget.

timestamps `property` ¶

timestamps: list[str]

arrays ¶

arrays() -> ndarray

save ¶

save(path: str | PathLike) -> None

Save to a .overlay.npz archive.

load `classmethod` ¶

load(path: str | PathLike) -> 'OverlayResult'

Load from a .overlay.npz archive.

to_msgpack ¶

to_msgpack() -> bytes

Serialize to msgpack bytes for HTTP transport.

from_msgpack `classmethod` ¶

from_msgpack(data: bytes) -> 'OverlayResult'

Deserialize bytes produced by :meth:to_msgpack.

summary ¶

summary() -> str

Listener¶

geoxplain_aurora_adapter.api.listener.listen_for_request ¶

listen_for_request(
    port: int = 8765,
    host: str = "127.0.0.1",
    persistent: bool = False,
    prewarm: bool = False,
    memory_retention: Optional[str] = None,
    result_retention: Optional[str] = None,
    **sbatch_kwargs,
) -> None

Start the geoxplain_aurora_adapter HTTP listener server.

Mode is auto-detected:

+------------------+------------------+---------------------------------+ | GPU visible? | sbatch on PATH? | Resolved mode | +==================+==================+=================================+ | yes | yes | sbatch-oneshot (or sbatch-persistent if persistent) | | no | yes | sbatch-oneshot (or sbatch-persistent if persistent) | | yes | no | gpu-listener | | no | no | Error | +------------------+------------------+---------------------------------+

Parameters:

Name	Type	Description	Default
`port`	`int`		`8765`
`host`	`str`	HTTP API is unauthenticated, so reach it via an SSH tunnel. Pass ``"0.0.0.0"`` only on a trusted/firewalled network).	`'127.0.0.1'`
`persistent`	`bool`	`default sbatch-oneshot (one sbatch per request).`	`False`
`prewarm`	`bool`	`rather than waiting for the first request.`	`False`
`memory_retention`	`Optional[str]`	listener's memory (default ``"1h"``; ``"never"`` disables).	`None`
`result_retention`	`Optional[str]`	are kept (default ``"never"``).	`None`
`**sbatch_kwargs`		Also accepted as CLI flags by ``geoxplain-aurora-adapter listen``.	`{}`

Aurora adapter Python API¶

Targets¶

geoxplain_aurora_adapter.schema.spec.Target ¶

geoxplain_aurora_adapter.schema.spec.TargetSpec dataclass ¶

point classmethod ¶

box classmethod ¶

box_bounds ¶

to_dict ¶

from_dict classmethod ¶

as_widget_dict ¶

Compute methods¶

geoxplain_aurora_adapter.api.methods.run_saliency ¶

geoxplain_aurora_adapter.api.methods.run_ig ¶

geoxplain_aurora_adapter.api.methods.run_rise ¶

geoxplain_aurora_adapter.api.methods.run_vit_cx ¶

geoxplain_aurora_adapter.api.methods.run_rollout ¶

Weather overlays¶

geoxplain_aurora_adapter.api.methods.pull_overlay ¶

Attribution results¶

geoxplain_aurora_adapter.schema.result.XiaFrame dataclass ¶

as_widget_dict ¶

geoxplain_aurora_adapter.schema.result.XiaResult dataclass ¶

single classmethod ¶

save ¶

load classmethod ¶

to_msgpack ¶

from_msgpack classmethod ¶

summary ¶

Overlay results¶

geoxplain_aurora_adapter.schema.overlay.OverlayFrame dataclass ¶

geoxplain_aurora_adapter.schema.overlay.OverlayResult dataclass ¶

timestamps property ¶

arrays ¶

save ¶

load classmethod ¶

to_msgpack ¶

from_msgpack classmethod ¶

summary ¶

Listener¶

geoxplain_aurora_adapter.api.listener.listen_for_request ¶

geoxplain_aurora_adapter.schema.spec.TargetSpec `dataclass` ¶

point `classmethod` ¶

box `classmethod` ¶

from_dict `classmethod` ¶

geoxplain_aurora_adapter.schema.result.XiaFrame `dataclass` ¶

geoxplain_aurora_adapter.schema.result.XiaResult `dataclass` ¶

single `classmethod` ¶

load `classmethod` ¶

from_msgpack `classmethod` ¶

geoxplain_aurora_adapter.schema.overlay.OverlayFrame `dataclass` ¶

geoxplain_aurora_adapter.schema.overlay.OverlayResult `dataclass` ¶

timestamps `property` ¶

load `classmethod` ¶

from_msgpack `classmethod` ¶