Matplotlib - твинкс с make_axes_locatable

Question

Я пытаюсь использовать твинкс с цветной панелью, созданной с помощью make_axes loctable. В другом вопросе я видел, что это может быть сложно, потому что axes_grid1 позиционирует оси только во время рисования ( здесь). Но я не мог адаптировать решение к своей проблеме.

Вот мой код:

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

def func_colorbar(mappable,cax=None,size="5%",pad=0.05,position='right',extend='neither'):
    ax = mappable.axes
    fig = ax.figure
    if position=='right':
        orientation='vertical'
    else:
        orientation='horizontal'
    if cax is None:
        divider = make_axes_locatable(ax)
        cax = divider.append_axes(position, size=size, pad=pad)
        if position=='top':
            cax.xaxis.tick_top()
    return fig.colorbar(mappable, cax=cax,extend=extend,orientation=orientation)

x=np.linspace(0,1,100)
y=np.linspace(0,1,100)
z=x*y

fig,axe=plt.subplots(1,1,figsize=(10,10))
im=axe.scatter(x,y,c=z)

ax2=axe.twinx()
ax2.plot(x,x**2)

func_colorbar(im)

plt.show()

Добавление ax2 портит макет. Ожидаемый результат можно получить, просто заменив func_colorbar (im) на fig.colorbar(im), но мне любопытно узнать, как решить проблему, если я хочу использовать axes_grid1 (что иногда может быть полезно). Ниже вы найдете то, что я пробовал, после ответа на этот вопрос. Однако похоже, что это не работает...

x=np.linspace(0,1,100)
y=np.linspace(0,1,100)
z=x*y

fig,axe=plt.subplots(1,1,figsize=(10,10))
im=axe.scatter(x,y,c=z)

ax2=axe.twinx()
ax2.plot(x,x**2)

func_colorbar(im)

def resize(event):
    axpos = axe.get_position()
    newpos = [axpos.x0, axpos.y0,  axpos.width, axpos.height] 
    ax2.set_position(newpos)

cid = fig.canvas.mpl_connect('draw_event', resize)
cid2 = fig.canvas.mpl_connect('resize_event', resize)

plt.show()

Спасибо за помощь!