from geoxplain import GeoXplainWidget
from geoxplain.xia_result import load_xia_result

w = GeoXplainWidget(height=640)

zurich_attribution = load_xia_result('zurich.xia.npz')
w.add_attribution(zurich_attribution)

w

from geoxplain import GeoXplainWidget

w_manual = GeoXplainWidget(height=640)

# Map each pressure level (z-index) to a saved .npy attribution grid.
# These were extracted from the first rollout frame (specific humidity, q)
# of zurich_rollout_cut.xia.npz.
grids = {
    'z-1': 'zurich_rollout_q_925hPa.npy',  # orders by z-index automatically
    'z-2': 'zurich_rollout_q_850hPa.npy'
}

# Optional, give layers names
layer_labels = {
    'z-1': '925 hPa',
    'z-2': '850 hPa'
}

w_manual.add_attribution(
    grids,
    method='Integrated Gradients',
    timestamp='2024-07-15T18:00:00Z',
    target={'south': 46.88, 'west': 7.74, 'north': 47.88, 'east': 9.34},
    layer_labels=layer_labels,
)

w_manual

import numpy as np
from geoxplain import GeoXplainWidget

# Build a synthetic global grid (lat 90 -> -90, lon 0 -> 360)
H, W = 721, 1440
lat_vals = np.linspace(90, -90, H)
lon_vals = np.linspace(0, 360, W, endpoint=False)
lon_grid, lat_grid = np.meshgrid(lon_vals, lat_vals)

# A diverging Gaussian blob centred over Europe
sigma = 5.0
blob_pos = np.exp(-((lat_grid - 47) ** 2 + (lon_grid - 20) ** 2) / (2 * sigma ** 2))
blob_neg = -0.6 * np.exp(-((lat_grid - 52) ** 2 + (lon_grid - 30) ** 2) / (2 * sigma ** 2))
synthetic = (blob_pos + blob_neg).astype(np.float32)

w_blob = GeoXplainWidget(height=640)

# Attribution straight from in-memory arrays, one grid per level
w_blob.add_attribution(
    {
        'z-2': synthetic,
        'z-3': synthetic * 0.8,
        'z-5': synthetic * 0.5,
    },
    method='Synthetic Blob',
    layer_labels={'z-2': '850 hPa', 'z-3': '700 hPa', 'z-5': '500 hPa'},
)

w_blob