<jupyter_start><jupyter_text>Plot features maps  Positions<NAME> (20, 50)<jupyter_code>normalized_attention = [
    (atten_map-atten_map.mean(dim=(-2, -1), keepdim=True))/atten_map.std(dim=(-2, -1), keepdim=True)
    for atten_map in attentions
]
cmap = plt.get_cmap('coolwarm')
norm = mpl_c.BoundaryNorm(np.linspace(-2.5, 2.5, 41), cmap.N)

gs = mpl_gs.GridSpec(nrows=2, ncols=1)

figsize = [s*0.48 for s in plt.rcParams['figure.figsize']]

fig = plt.figure(figsize=figsize, dpi=300)

ax_feat_0 = fig.add_subplot(gs[0], projection=ccrs.PlateCarree())
feat_plt, lons = add_cyclic_point(normalized_attention[0][56].numpy(), coord=data_module.ds_train.dataset.ifs['longitude'])
ax_feat_0.pcolormesh(lons, data_module.ds_train.dataset.ifs['latitude'], feat_plt, cmap=cmap, norm=norm)
ax_feat_0.text(-0.05, 0.5, '(1,56)', fontsize=10, va='center', ha='center',
               rotation='vertical', transform=ax_feat_0.transAxes)
ax_feat_0.coastlines(linewidth=0.3)

ax_feat_1 = fig.add_subplot(gs[1], projection=ccrs.PlateCarree())
feat_plt, lons = add_cyclic_point(normalized_attention[4][28], coord=data_module.ds_train.dataset.ifs['longitude'])
ax_feat_1.pcolormesh(lons, data_module.ds_train.dataset.ifs['latitude'], feat_plt, cmap=cmap, norm=norm)
ax_feat_1.text(-0.01, 0.5, '(5,28)', fontsize=10, va='center', ha='right',
               rotation='vertical', transform=ax_feat_1.transAxes)
ax_feat_1.coastlines(linewidth=0.3)
fig.savefig('../paper/figures/fig_attention.png', dpi=300, bbox_inches='tight')<jupyter_output>findfont: Font family ['serif'] not found. Falling back to DejaVu Sans.
{'fraction_non_alphanumeric': 0.1708542713567839, 'fraction_numerical': 0.042085427135678394, 'mean_word_length': 6.695652173913044, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '541450', 'license': 'MIT', 'n_tokens_mistral': 685, 'n_tokens_neox': 644, 'n_words': 89, 'path': 'notebooks/fig_attention_overview.ipynb', 'repo_name': 'tobifinn/ensemble_transformer', 'size': '1602'}