import warnings warnings.filterwarnings("ignore") import sys import argparse import xarray as xr import numpy as np import matplotlib.pyplot as plt import matplotlib.colors as mcolors import cartopy.crs as ccrs import cartopy.feature as cfeature import seaborn as sns import gradio as gr import torch from matplotlib.patches import Rectangle from matplotlib.gridspec import GridSpec # Configure matplotlib rcParams for consistent font sizes plt.rcParams.update({ 'axes.titlesize': 14, # Title font size for plots 'axes.labelsize': 14, # Axis label font size 'xtick.labelsize': 14, # X-axis tick label font size 'ytick.labelsize': 14, # Y-axis tick label font size 'legend.fontsize': 14, # Legend font size 'figure.titlesize': 16, # Figure title font size 'font.size': 14, # Base font size }) # Define specific font sizes for different plot types PLOT_TITLE_SIZE = plt.rcParams['axes.titlesize'] # 14 IMPORTANCE_TITLE_SIZE = 12 # Smaller for importance plots # Parse command line arguments before Gradio processes them def parse_args(): parser = argparse.ArgumentParser(description='TeleViT XAI Viewer', add_help=False) parser.add_argument('mode', nargs='?', choices=['local'], help='Use local data instead of remote') # Parse known args to avoid conflicts with Gradio args, unknown = parser.parse_known_args() # Remove our custom arguments from sys.argv so Gradio doesn't see them if args.mode == 'local': sys.argv = [sys.argv[0]] + unknown return args.mode == 'local' # Parse arguments early USE_LOCAL_DATA = parse_args() # --- Constants and configuration (copied/kept consistent with main.py) --- COMMON_VARS = { "input_global_shape": [14, 1, 360, 180], "input_local_shape": [14, 1, 80, 80], "input_oci_shape": [10, 10], "patch_global_shape": [14, 1, 30, 30], "patch_local_shape": [14, 1, 16, 16], "patch_oci_shape": [1, 1], "local_vars": ['lst_day', 'mslp', 'ndvi', 'pop_dens', 'ssrd', 'sst', 'swvl1', 't2m_mean', 'tp', 'vpd'], "global_vars": ['lst_day', 'mslp', 'ndvi', 'pop_dens', 'ssrd', 'sst', 'swvl1', 't2m_mean', 'tp', 'vpd'], "positional_vars": ['cos_lat', 'sin_lat', 'cos_lon', 'sin_lon'], "oci_vars": ['oci_censo', 'oci_ea', 'oci_epo', 'oci_ao', 'oci_nao', 'oci_nina34_anom', 'oci_pdo', 'oci_pna', 'oci_soi', 'oci_wp'], "var_2_label": {'oci_censo' : 'Bivariate ENSO', 'oci_ea' : 'Eastern Atlantic', 'oci_epo' : 'East Pacific', 'oci_ao' : 'Arctic', 'oci_nao' : 'North Atlantic', 'oci_nina34_anom' : 'NINA 3.4', 'oci_pdo' : 'Pacific Decadal', 'oci_pna' : 'Pacific/North Amer.', 'oci_soi' : 'Southern', 'oci_wp' : 'Western Pacific', 'argmax_var': 'Most Important Variable', "lst_day": "Land Surface Temperature", "mslp": "Mean Sea Level Pressure", "ndvi": "Normalized Difference Vegetation Index", "pop_dens": "Population Density", "ssrd": "Surface Solar Radiation Downwards", "sst": "Sea Surface Temperature", "swvl1": "Soil Wetness Layer 1", "t2m_mean": "2m Temperature", "tp": "Total Precipitation", "vpd": "Vapour Pressure Deficit" }, 'multiplier': {'local': 80*80, 'global': 360*180, 'oci': 10} } # Set DATA_PATH based on argument if USE_LOCAL_DATA: DATA_PATH = './data' else: DATA_PATH = 'https://huggingface.co/datasets/iprapas/televit-xai-data/resolve/main' LEAD_TIMES = [0, 1, 2, 4, 8, 16] TIME_RANGE = (0, 45) # Mapping between preset places and (latitude_idx, longitude_idx) PLACE_TO_IDX: dict[str, tuple[int, int]] = { 'north_america': (2, 3), 'mediterranean': (2, 9), 'north_africa': (3, 10), 'amazon': (4, 5), 'south_africa': (5, 10), 'australia': (5, 15), 'southeast_asia': (4, 14), } # Lazily initialized datasets _ds_dict: dict | None = None def _load_datasets_once() -> dict: global _ds_dict if _ds_dict is not None: return _ds_dict path_dict: dict[str, str] = {} path_dict['attentions'] = f'{DATA_PATH}/attentions-televit_ig-2019.zip' path_dict['predictions'] = f'{DATA_PATH}/predictions-televit_ig-2019.zip' path_dict['seasfire'] = f'{DATA_PATH}/seasfire_2019_v0.4.zip' path_dict['seasfire_1deg'] = f'{DATA_PATH}/seasfire_2019_1deg_v0.4.zip' for input_type in ['local', 'global', 'oci']: path_dict[f'xai-{input_type}'] = f'{DATA_PATH}/intgrad-televit_ig-{input_type}-2019.zip' ds_dict: dict[str, xr.Dataset] = {} for key, value in path_dict.items(): # Open remote zipped zarr stores ds_dict[key] = xr.open_zarr(f'zip://::{value}', consolidated=True) _ds_dict = ds_dict return _ds_dict def get_available_vars(dataset_type: str) -> list[str]: return ['argmax_var'] + COMMON_VARS[f"{dataset_type}_vars"] def get_code_to_label_map(dataset_type: str) -> dict[str, str]: codes = get_available_vars(dataset_type) return {code: COMMON_VARS['var_2_label'].get(code, code) for code in codes} def get_label_to_code_map(dataset_type: str) -> dict[str, str]: m = get_code_to_label_map(dataset_type) return {label: code for code, label in m.items()} def get_var_labels(dataset_type: str) -> list[str]: return list(get_code_to_label_map(dataset_type).values()) def render_plot(var: str, time: int, lead_time: int, aggregation: str, dataset_type: str): ds_dict = _load_datasets_once() agg = None if (aggregation is None or str(aggregation).lower() == 'none') else aggregation # Select the dataset and the requested variable slice ds_key = f"xai-{dataset_type}" ds = ds_dict[ds_key].sel(lead_time=lead_time)[var] if agg == 'median': data = ds.median(dim="time") time_label = "median" else: data = ds.isel(time=int(time)) # format time for title tval = str(data.time.values) time_label = tval.split("T")[0] # Prepare colormap / normalization if var == 'argmax_var': labels = COMMON_VARS[f"{dataset_type}_vars"] num_classes = len(labels) palette = sns.color_palette("muted", n_colors=num_classes) cmap = mcolors.ListedColormap(palette) boundaries = np.arange(num_classes + 1) - 0.5 norm = mcolors.BoundaryNorm(boundaries=boundaries, ncolors=num_classes) else: cmap = 'viridis' norm = mcolors.Normalize() # Projection and figure include_secondary = (dataset_type == 'oci' and var != 'argmax_var') proj = ccrs.Robinson() if include_secondary: fig = plt.figure(figsize=(12, 9)) ax = fig.add_subplot(2, 1, 1, projection=proj) else: fig, ax = plt.subplots(figsize=(12, 6), subplot_kw={'projection': proj}) ax.add_feature(cfeature.LAND, facecolor='lightgray', zorder=0) ax.add_feature(cfeature.COASTLINE, linewidth=0.6) ax.set_global() ax.gridlines(draw_labels=False, linestyle='--', linewidth=0.5, alpha=0.7) # Plot im = ax.pcolormesh(data['longitude'], data['latitude'], data*COMMON_VARS['multiplier'][dataset_type] if var != 'argmax_var' else data, transform=ccrs.PlateCarree(), cmap=cmap, norm=norm, alpha=0.9) # Colorbar if var == 'argmax_var': labels = COMMON_VARS[f"{dataset_type}_vars"] num_classes = len(labels) cbar = plt.colorbar(im, ax=ax, orientation='vertical', ticks=range(num_classes), pad=0.05, shrink=0.7) cbar.ax.set_yticklabels([COMMON_VARS["var_2_label"].get(x, x) for x in labels]) cbar.set_label("Variables") else: cbar = plt.colorbar(im, ax=ax, orientation='vertical', pad=0.05, shrink=0.7) cbar.set_label(COMMON_VARS["var_2_label"].get(var, var)) cbar.ax.tick_params(left=False, right=False, labelleft=False, labelbottom=False) # Title if var == 'argmax_var': ax.set_title(f"Most important {dataset_type} variable - time={time_label} - lead time={lead_time}x8d", fontsize=IMPORTANCE_TITLE_SIZE) else: ax.set_title(f"{COMMON_VARS['var_2_label'].get(var, var)} importance - time={time_label} - lead time={lead_time}x8d", fontsize=IMPORTANCE_TITLE_SIZE) # Secondary subplot for OCI raw variable when applicable if include_secondary: try: ax2 = fig.add_subplot(2, 1, 2) ds_dict['seasfire'][var].plot(ax=ax2) ax2.set_title(f"{COMMON_VARS['var_2_label'].get(var, var)} (seasfire)") # Mark selected time when applicable if agg is None: try: selected_time = data['time'].values ax2.axvline(selected_time, color='red', linestyle='--', linewidth=1.5) except Exception: pass except Exception: ax2 = fig.add_subplot(2, 1, 2) ax2.axis('off') ax2.text(0.5, 0.5, 'Unable to plot seasfire variable', ha='center', va='center') fig.tight_layout() return fig # --- Gradio UI --- def _update_var_choices(dataset_type: str): labels = get_var_labels(dataset_type) default_label = COMMON_VARS['var_2_label'].get('argmax_var', 'argmax_var') return gr.update(choices=labels, value=default_label) def _update_time_interactive(aggregation: str): agg = None if (aggregation is None or str(aggregation).lower() == 'none') else aggregation return gr.update(interactive=(agg is None)) def _predict(dataset_type: str, var_label: str, lead_time: int, aggregation: str, time: int): try: # Map human-friendly label back to code code = get_label_to_code_map(dataset_type).get(var_label, var_label) fig = render_plot(var=code, time=time, lead_time=int(lead_time), aggregation=aggregation, dataset_type=dataset_type) return fig except Exception as e: raise gr.Error(str(e)) def _predict_fixed(dataset_type: str): def _fn(var_label: str, lead_time: int, aggregation: str, time: int): return _predict(dataset_type, var_label, lead_time, aggregation, time) return _fn def _build_tab_for(dataset_type: str): with gr.Tab(dataset_type.capitalize() + " importance per patch"): gr.Markdown(f""" Visualize importance maps for the {dataset_type} inputs. - Each map shows the summed integrated gradients per local patch for the selected variable. - "Most Important Variable" colors each patch by the variable with the maximum summed attribution. - Select a lead time (8-day steps) and a time index; or use median to aggregate across time. """) var = gr.Dropdown(label="Variable", choices=get_var_labels(dataset_type), value=COMMON_VARS['var_2_label'].get('argmax_var', 'argmax_var')) lead_time = gr.Dropdown(label="Lead Time (x8d)", choices=LEAD_TIMES, value=0) aggregation = gr.Dropdown(label="Aggregation", choices=["none", "median"], value="none") time = gr.Slider(label="Time Index", minimum=TIME_RANGE[0], maximum=TIME_RANGE[1], step=1, value=0, interactive=True) submit = gr.Button("Render") out = gr.Plot(label="Map") aggregation.change(_update_time_interactive, inputs=aggregation, outputs=time) submit.click(_predict_fixed(dataset_type), inputs=[var, lead_time, aggregation, time], outputs=out) def _get_attentions(ds_attn): input_global_shape = [10, 1, 180, 360] input_local_shape = [10, 1, 80, 80] input_oci_shape = [10, 10] patch_global_shape = [10, 1, 30, 30] patch_local_shape = [10, 1, 16, 16] patch_oci_shape = [1, 1] token_global_shape = [x//y for x,y in zip(input_global_shape, patch_global_shape)] # [1,1,6,12] actually for 180x360 num_global_tokens = np.prod(token_global_shape) token_local_shape = [x//y for x,y in zip(input_local_shape, patch_local_shape)] # [1,1,5,5] num_local_tokens = np.prod(token_local_shape) token_oci_shape = [x//y for x,y in zip(input_oci_shape, patch_oci_shape)] # [10,10] num_oci_tokens = np.prod(token_oci_shape) attns = ds_attn attns_t = attns local_to_others_attn = attns_t[:num_local_tokens].mean(dim=0) local_to_local_attn = local_to_others_attn[:num_local_tokens].reshape(*token_local_shape).mean(dim=(0,1)) start_idx = num_local_tokens end_idx = start_idx + num_global_tokens local_to_global_attn = local_to_others_attn[start_idx:end_idx].reshape(*token_global_shape).mean(dim=(0,1)) start_idx = end_idx end_idx = start_idx + num_oci_tokens local_to_oci_attn = local_to_others_attn[start_idx:end_idx].reshape(*token_oci_shape) return local_to_local_attn, local_to_global_attn, local_to_oci_attn def _predict_global_comparison(lead_time: int, time: int): ds_dict = _load_datasets_once() ds = ds_dict['seasfire'] if 'seasfire' in ds_dict else None preds_ds = ds_dict['predictions'] if 'predictions' in ds_dict else None plots = [] # Create Target plot (global) fig_target, ax_target = plt.subplots(figsize=(15, 8), subplot_kw={'projection': ccrs.PlateCarree()}) if ds is not None: target_data = ds.isel(time=time) # Add cartopy features ax_target.add_feature(cfeature.COASTLINE, linewidth=0.5) ax_target.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3) # Plot target data - using the mask as specified: (gwis_ba > 0).where(gwis_ba > 0) mask = target_data['gwis_ba'] mask_filtered = np.log(mask + 1).where(mask > 0) im_target = ax_target.pcolormesh(target_data['longitude'], target_data['latitude'], mask_filtered, transform=ccrs.PlateCarree(), alpha=0.9, cmap='Reds') # Set global extent ax_target.set_global() ax_target.gridlines(draw_labels=False, linestyle='--', alpha=0.5, xlocs=np.arange(-180, 181, 30), ylocs=np.arange(-90, 91, 30)) # Format time for title tval = str(target_data.time.values) time_label = tval.split("T")[0] ax_target.set_title(f'Global Targets (Burned Area > 0) - time={time_label}') cbar_target = fig_target.colorbar(im_target, ax=ax_target, fraction=0.046, pad=0.04) cbar_target.set_label('Fire Occurrence') else: ax_target.axis('off') ax_target.text(0.5, 0.5, 'Target dataset not available', ha='center', va='center', transform=ax_target.transAxes) fig_target.tight_layout() plots.append(fig_target) # Create Prediction plot (global) fig_pred, ax_pred = plt.subplots(figsize=(15, 8), subplot_kw={'projection': ccrs.PlateCarree()}) try: if preds_ds is not None: preds = preds_ds[f'predictions_{lead_time}'].isel(time=time) # Add cartopy features ax_pred.add_feature(cfeature.COASTLINE, linewidth=0.5) ax_pred.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3) # Plot predictions im_pred = ax_pred.pcolormesh(preds['longitude'], preds['latitude'], np.ma.masked_where(preds.values < 0.01, preds.values), transform=ccrs.PlateCarree(), alpha=0.9, cmap='Spectral_r', vmin=0, vmax=1) # Set global extent ax_pred.set_global() ax_pred.gridlines(draw_labels=False, linestyle='--', alpha=0.5, xlocs=np.arange(-180, 181, 30), ylocs=np.arange(-90, 91, 30)) # Format time for title tval = str(preds.time.values) time_label = tval.split("T")[0] ax_pred.set_title(f'Global Predictions - time={time_label} - lead time={lead_time}x8d') cbar_pred = fig_pred.colorbar(im_pred, ax=ax_pred, fraction=0.046, pad=0.04) cbar_pred.set_label('Prediction Confidence') else: ax_pred.axis('off') ax_pred.text(0.5, 0.5, 'Predictions dataset not available', ha='center', va='center', transform=ax_pred.transAxes) except Exception as e: ax_pred.axis('off') ax_pred.text(0.5, 0.5, f'Unable to plot predictions: {str(e)}', ha='center', va='center', transform=ax_pred.transAxes) fig_pred.tight_layout() plots.append(fig_pred) return plots def _predict_attentions(place: str, latitude_idx: int, longitude_idx: int, lead_time: int, time: int): ds_dict = _load_datasets_once() # Resolve place presets if place and place in PLACE_TO_IDX: latitude_idx, longitude_idx = PLACE_TO_IDX[place] ds_attns_tmp = ds_dict[f'attentions'][f'attentions_{lead_time}'].isel(time=time, latitude=latitude_idx, longitude=longitude_idx) longitude = ds_attns_tmp.longitude latitude = ds_attns_tmp.latitude local_to_local, local_to_global_attn, local_to_oci_attn = _get_attentions(torch.from_numpy(ds_attns_tmp.to_numpy())) ds = ds_dict['seasfire'] if 'seasfire' in ds_dict else None ds_1deg = ds_dict['seasfire_1deg'] if 'seasfire_1deg' in ds_dict else None preds_ds = ds_dict['predictions'] if 'predictions' in ds_dict else None plots = [] # Create Target plot fig_target, ax_target = plt.subplots(figsize=(10, 8), subplot_kw={'projection': ccrs.PlateCarree()}) if ds is not None: target_data = ds.sel(longitude=slice(longitude-10, longitude+10), latitude=slice(latitude+10, latitude-10)).isel(time=time) import numpy as np import matplotlib.colors as mcolors # Create binary mask (1 where burned area, 0 elsewhere) mask = (target_data['gwis_ba'] > 0).astype(int) # Create categorical colormap: index 0 = transparent, index 1 = red cmap = mcolors.ListedColormap([(1, 0, 0, 0), # RGBA for transparent (1, 0, 0, 1)]) # RGBA for red bounds = [-0.5, 0.5, 1.5] norm = mcolors.BoundaryNorm(bounds, cmap.N) # Add cartopy features ax_target.add_feature(cfeature.COASTLINE, linewidth=0.5) ax_target.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3) # Plot binary mask im_target = ax_target.pcolormesh( target_data['longitude'], target_data['latitude'], mask, transform=ccrs.PlateCarree(), cmap=cmap, norm=norm ) # Set extent based on data ax_target.set_extent( [target_data['longitude'].min(), target_data['longitude'].max(), target_data['latitude'].min(), target_data['latitude'].max()], ccrs.PlateCarree() ) ax_target.gridlines(draw_labels=False, linestyle='--', alpha=0.5, xlocs=np.arange(-180, 181, 4), ylocs=np.arange(-90, 91, 4)) ax_target.set_title('Burned Area Mask') fig_target.tight_layout() plots.append(fig_target) # Create Prediction plot fig_pred, ax_pred = plt.subplots(figsize=(10, 8), subplot_kw={'projection': ccrs.PlateCarree()}) try: if preds_ds is not None: preds = preds_ds[f'predictions_{lead_time}'].isel(time=time) preds_sel = preds.sel(longitude=slice(longitude-10, longitude+10), latitude=slice(latitude+10, latitude-10)) # Add cartopy features ax_pred.add_feature(cfeature.COASTLINE, linewidth=0.5) ax_pred.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3) # Plot predictions im_pred = ax_pred.pcolormesh(preds_sel['longitude'], preds_sel['latitude'], preds_sel.values, transform=ccrs.PlateCarree(), alpha=0.9, cmap='Spectral_r') # Set extent and styling ax_pred.set_extent([preds_sel['longitude'].min(), preds_sel['longitude'].max(), preds_sel['latitude'].min(), preds_sel['latitude'].max()], ccrs.PlateCarree()) ax_pred.gridlines(draw_labels=False, linestyle='--', alpha=0.5, xlocs=np.arange(-180, 181, 4), ylocs=np.arange(-90, 91, 4)) ax_pred.set_title('Prediction') cbar_pred = fig_pred.colorbar(im_pred, ax=ax_pred, fraction=0.046, pad=0.04) cbar_pred.set_label('Confidence') else: ax_pred.axis('off') ax_pred.text(0.5, 0.5, 'Predictions dataset not available', ha='center', va='center', transform=ax_pred.transAxes) except Exception: ax_pred.axis('off') ax_pred.text(0.5, 0.5, 'Unable to plot predictions', ha='center', va='center', transform=ax_pred.transAxes) fig_pred.tight_layout() plots.append(fig_pred) # Create Local to Local Attention plot fig_local, ax_local = plt.subplots(figsize=(10, 8), subplot_kw={'projection': ccrs.PlateCarree()}) # Add cartopy features ax_local.add_feature(cfeature.COASTLINE, linewidth=0.5) ax_local.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3) if ds is not None: target_data = ds.sel(longitude=slice(longitude-10, longitude+10), latitude=slice(latitude+10, latitude-10)).isel(time=time) # Create coordinate grids for the attention data lon_coords = target_data['longitude'] lat_coords = target_data['latitude'] # Interpolate attention data to match the spatial resolution attention_interp = torch.nn.functional.interpolate( local_to_local.unsqueeze(0).unsqueeze(0), size=(len(lat_coords), len(lon_coords)), mode='bilinear' ).squeeze().detach().numpy() # Plot attention data im_local = ax_local.pcolormesh(lon_coords, lat_coords, attention_interp, transform=ccrs.PlateCarree(), cmap='viridis', alpha=0.9) # Set extent and styling ax_local.set_extent([lon_coords.min(), lon_coords.max(), lat_coords.min(), lat_coords.max()], ccrs.PlateCarree()) else: # Fallback: create simple coordinate grid around the selected location lon_range = np.linspace(longitude-10, longitude+10, 20) lat_range = np.linspace(latitude-10, latitude+10, 20) attention_interp = torch.nn.functional.interpolate( local_to_local.unsqueeze(0).unsqueeze(0), size=(20, 20), mode='bilinear' ).squeeze().detach().numpy() im_local = ax_local.pcolormesh(lon_range, lat_range, attention_interp, transform=ccrs.PlateCarree(), cmap='viridis', alpha=0.9) ax_local.set_extent([lon_range.min(), lon_range.max(), lat_range.min(), lat_range.max()], ccrs.PlateCarree()) ax_local.gridlines(draw_labels=False, linestyle='--', alpha=0.5, xlocs=np.arange(-180, 181, 4), ylocs=np.arange(-90, 91, 4)) ax_local.set_title('Local to Local Attention') cbar_local = fig_local.colorbar(im_local, ax=ax_local, fraction=0.046, pad=0.04) cbar_local.set_label('Attention weight') fig_local.tight_layout() plots.append(fig_local) # Create Local to OCI Attention plot fig_oci, ax_oci = plt.subplots(figsize=(10, 8)) im_oci = ax_oci.imshow(local_to_oci_attn, cmap='viridis', alpha=0.9) ax_oci.set_title('Local to OCI Attention') ax_oci.set_yticks(np.arange(10)) ax_oci.set_yticklabels([COMMON_VARS["var_2_label"].get(x, x) for x in COMMON_VARS["oci_vars"]]) ax_oci.set_xticks(np.arange(10)) ax_oci.set_xticklabels(np.arange(10, 0, -1)) ax_oci.set_ylabel('OCI Variables') ax_oci.set_xlabel('Months Before Prediction') cbar_oci = fig_oci.colorbar(im_oci, ax=ax_oci, fraction=0.046, pad=0.04) cbar_oci.set_label('Attention weight') fig_oci.tight_layout() plots.append(fig_oci) # Create Local to Global Attention plot fig_global, ax_global = plt.subplots(figsize=(15, 8), subplot_kw={'projection': ccrs.PlateCarree()}) # Add cartopy features ax_global.add_feature(cfeature.COASTLINE, linewidth=0.5) ax_global.add_feature(cfeature.LAND, facecolor='lightgray', alpha=0.3) if ds_1deg is not None: # Create coordinate grids for the global attention data global_lons = ds_1deg['longitude'] global_lats = ds_1deg['latitude'] # Interpolate attention data to match the global grid attention_global_interp = torch.nn.functional.interpolate( local_to_global_attn.unsqueeze(0).unsqueeze(0), size=(len(global_lats), len(global_lons)), mode='bilinear' ).squeeze().detach().numpy() # Plot global attention im_global = ax_global.pcolormesh(global_lons, global_lats, attention_global_interp, transform=ccrs.PlateCarree(), cmap='viridis', alpha=0.9) # Draw rectangle for the corresponding local patch on the global grid try: lat_min = latitude - 10 lat_max = latitude + 10 lon_min = longitude - 10 lon_max = longitude + 10 lat_min_v = float(np.asarray(getattr(lat_min, 'values', lat_min))) lat_max_v = float(np.asarray(getattr(lat_max, 'values', lat_max))) lon_min_v = float(np.asarray(getattr(lon_min, 'values', lon_min))) lon_max_v = float(np.asarray(getattr(lon_max, 'values', lon_max))) # Draw rectangle in geographic coordinates from matplotlib.patches import Rectangle as GeoRectangle rect = GeoRectangle((lon_min_v, lat_min_v), lon_max_v - lon_min_v, lat_max_v - lat_min_v, linewidth=2.5, edgecolor='cyan', facecolor='none', alpha=0.9, transform=ccrs.PlateCarree()) ax_global.add_patch(rect) except Exception: pass else: # Fallback: create global coordinate grid global_lons = np.linspace(-180, 180, 360) global_lats = np.linspace(-90, 90, 180) attention_global_interp = torch.nn.functional.interpolate( local_to_global_attn.unsqueeze(0).unsqueeze(0), size=(180, 360), mode='bilinear' ).squeeze().detach().numpy() im_global = ax_global.pcolormesh(global_lons, global_lats, attention_global_interp, transform=ccrs.PlateCarree(), cmap='viridis', alpha=0.9) # Set global extent and styling ax_global.set_global() ax_global.gridlines(draw_labels=False, linestyle='--', alpha=0.5, xlocs=np.arange(-180, 181, 30), ylocs=np.arange(-90, 91, 30)) ax_global.set_title('Local to Global Attention') cbar_global = fig_global.colorbar(im_global, ax=ax_global, fraction=0.046, pad=0.04) cbar_global.set_label('Attention weight') fig_global.tight_layout() plots.append(fig_global) return plots def _build_attentions_tab(): with gr.Tab("Attentions"): gr.Markdown( """ Explore attention maps at a specific location and time. - Local→Local shows how tokens within the local patch attend to each other (upsampled to the 80×80 input grid). - Local→Global shows attention from local tokens to global tokens (upsampled to the 180×360 grid). - Local→OCI shows attention from local tokens to the 10×10 monthly climate indices grid. - The attention overlays are drawn on top of their originating input grids (land/LSM or target), to provide spatial context. """ ) place = gr.Dropdown(label="Place", choices=[""] + list(PLACE_TO_IDX.keys()), value="mediterranean") latitude_idx = gr.Slider(label="Latitude Index", minimum=0, maximum=9, step=1, value=2) longitude_idx = gr.Slider(label="Longitude Index", minimum=0, maximum=18, step=1, value=9) lead_time = gr.Dropdown(label="Lead Time (x8d)", choices=LEAD_TIMES, value=0) time = gr.Slider(label="Time Index", minimum=TIME_RANGE[0], maximum=TIME_RANGE[1], step=1, value=0) # Hidden state to indicate when lat/lon updates originate from a place selection syncing_from_place = gr.State(False) submit = gr.Button("Render") target_plot = gr.Plot(label="Target") pred_plot = gr.Plot(label="Prediction") local_plot = gr.Plot(label="Local to Local Attention") oci_plot = gr.Plot(label="Local to OCI Attention") global_plot = gr.Plot(label="Local to Global Attention") outputs = [target_plot, pred_plot, local_plot, oci_plot, global_plot] # When place changes, update lat/lon and set syncing flag so slider change handlers don't clear place def _on_place_change(selected_place: str, _flag: bool): if selected_place and selected_place in PLACE_TO_IDX: lat, lon = PLACE_TO_IDX[selected_place] return gr.update(value=lat), gr.update(value=lon), gr.update(value=True) return gr.update(), gr.update(), gr.update(value=False) place.change(_on_place_change, inputs=[place, syncing_from_place], outputs=[latitude_idx, longitude_idx, syncing_from_place]) # When user manually changes lat/lon, update place if coordinates match a preset, otherwise clear it def _on_index_change(lat: int, lon: int, flag: bool): if flag: # Reset the flag, keep the current place return gr.update(), gr.update(value=False) # Manual change: check if the coordinates match any preset place for place_name, (place_lat, place_lon) in PLACE_TO_IDX.items(): if lat == place_lat and lon == place_lon: # Coordinates match a preset place, update to that place return gr.update(value=place_name), gr.update(value=False) # Coordinates don't match any preset, clear place return gr.update(value=""), gr.update(value=False) latitude_idx.change(_on_index_change, inputs=[latitude_idx, longitude_idx, syncing_from_place], outputs=[place, syncing_from_place]) longitude_idx.change(_on_index_change, inputs=[latitude_idx, longitude_idx, syncing_from_place], outputs=[place, syncing_from_place]) submit.click(_predict_attentions, inputs=[place, latitude_idx, longitude_idx, lead_time, time], outputs=outputs) def _build_global_predictions_tab(): with gr.Tab("Global Predictions vs Targets"): gr.Markdown( """ Visualize global predictions vs targets for wildfire occurrence. - **Targets**: Shows fire occurrence (burned area > 0) from the GWIS dataset. - **Predictions**: Shows model predictions for the selected lead time. - Select lead time (8-day steps) and time index to compare model predictions with actual fire occurrence. - Both plots are shown on a global scale to understand spatial patterns of fire activity and model performance. """ ) lead_time = gr.Dropdown(label="Lead Time (x8d)", choices=LEAD_TIMES, value=0) time = gr.Slider(label="Time Index", minimum=TIME_RANGE[0], maximum=TIME_RANGE[1], step=1, value=0) submit = gr.Button("Render") target_plot = gr.Plot(label="Global Targets") pred_plot = gr.Plot(label="Global Predictions") outputs = [target_plot, pred_plot] submit.click(_predict_global_comparison, inputs=[lead_time, time], outputs=outputs) def _warmup(): # Load all datasets once on first page load _load_datasets_once() return None def build_interface() -> gr.Blocks: data_source_note = f"**Data Source:** {'Local data (`./data`)' if USE_LOCAL_DATA else 'Remote data (HuggingFace)'}" loading_note = "Note: The first render may take longer as datasets are loaded from " + ("local storage." if USE_LOCAL_DATA else "remote storage.") with gr.Blocks(title="TeleViT XAI Viewer") as demo: gr.Markdown(f""" ### TeleViT XAI Viewer {data_source_note} Select dataset type tab, then choose variable, time, lead time, and aggregation to visualize the corresponding map. {loading_note} Scientific context: - For local, global, and oci inputs, the maps show the sum of integrated gradients within each local patch for the selected variable. This shows the importance of the variable in the local patch. Choosing "Most Important Variable" renders, for each patch, the variable with the highest summed attribution (argmax over variables). - Lead time indexes the prediction horizon in steps of 8 days; time selects the date of the forecast. "median" aggregates attributions across time before rendering. - Attention maps visualize transformer attention from local tokens to: Local (within-patch), Global (global tokens), and OCI (10×10 monthly climate indices). These attentions are upsampled to their native input grids (80×80 local; 180×360 global; 10×10 OCI) and overlaid on the corresponding backgrounds to provide spatial context. In the Attentions tab you can choose a preset place or indices, plus lead time and time, to render these views. """) with gr.Tabs(): _build_attentions_tab() _build_tab_for("local") _build_tab_for("global") _build_tab_for("oci") _build_global_predictions_tab() # Trigger dataset loading on first page load (must be inside Blocks context) demo.load(_warmup) return demo demo = build_interface() def main(): # Print data source information if USE_LOCAL_DATA: print("Using local data from:", DATA_PATH) else: print("Using remote data from:", DATA_PATH) # Ensure datasets are lazily loaded on first use demo.queue() demo.launch(server_name="0.0.0.0", server_port=7860) if __name__ == "__main__": main()