Matplotlib : dessiner entre les sous-graphiques

La semaine dernière, je préparais un rapport d’analyse de données avec Jupyter, Pandas et Matplotlib (pour ne citer que quelques briques de ce fabuleux framework). Une des figures contenait deux sous-graphiques : le second étant un agrandissement d’une région du premier. Pour bien illustrer cela et, au passage, montrer au fanclub MATLAB à quel point ils sont restés “so 90s”, j’ai décidé de tracer une flèche entre les deux sous-graphiques.

Hélas, cette idée simple m’a conduit à trois heures de recherche intense sur Internet, entre essais, erreurs, et même quelques plantages. J’étais à deux doigts de m’incliner devant la technologie du siècle dernier… Mais pas question !

Artistes et transformations

Si vous voulez juste vous vanter du résultat final, passez cette section. Je vais juste expliquer vite-fait ce qui se passe sous le capot. Alors, figures, axes, abscisses, c’est quoi tout ça ? Je ne saurais trouver une meilleure explication que le tutoriel sur les “artistes” de Matplotlib, mais je vais essayer de résumer. Avec Matplotlib, vous dessinez sur un Canvas en utilisant un Renderer. Les objets capables d’utiliser un renderer pour dessiner sur le canvas s’appellent des “artistes”.

Il existe deux types d’artistes : les primitives et les conteneurs. Les primitives représentent les objets graphiques standards que nous voulons dessiner sur le canvas : Line2D, Rectangle, Text, AxesImage, etc. Les conteneurs sont des endroits pour les y mettre (Axis, Axes et Figure).

On en déduit qu’il nous faut créer une primitive, à savoir un Patch, et plus précisément dans notre cas un FancyArrowPatch, qu’on ajoutera à la figure. Facile.

Mais il y a un autre aspect à considérer. Comme nous ajoutons un patch à la figure, on peut s’attendre à ce qu’il soit dessiné en utilisant les coordonnées de la figure. Ce n’est pas très simple quand il s’agit de bien placer la flèche. Pour résoudre cela, il faut comprendre comment fonctionnent les transformations. Là encore, il existe un excellent tutoriel sur les transformations, et je ne saurais mieux expliquer. Lorsque vous traitez des coordonnées dans Matplotlib, il faut savoir à quel système de coordonnées elles se rapportent. Il y a 4 systèmes de coordonnées :

Le système de référence est display, et les objets de transformation décrivent comment transformer des coordonnées vers ce système de référence. Nous savons où positionner la flèche dans les données (c’est à dire ses coordonnées dans le système data), mais pas où cela se place sur la figure complète. En appliquant une transformation via l’objet de transformation des données, nous obtenons les coordonnées en pixels. Une étape. Ensuite, il suffit d’inverser la transformation en utilisant l’objet de transformation de la figure pour convertir depuis les coordonnées pixels au système de coordonnées de la figure.

Dernier point, il faut indiquer explicitement le système de coordonnées utilisé lors de la création du patch avec le paramètre transform.

Montre-moi le code

OK, sortez vos données, configurez le graphique ! Dans l’exemple ci-dessous, j’utilise un DataFrame Pandas, mais cela fonctionnerait tout aussi bien avec des tableaux NumPy.

import matplotlib.pyplot as plt
import matplotlib.patches

fig, axes = plt.subplots(2)
# Affiche la mesure complète pour comparaison sur l'axe 0
testdata.plot(ax=axes[0], x="mdfTime", y="psiDt")
# Place un petit rectangle pour marquer la zone agrandie
axes[0].add_patch(matplotlib.patches.Rectangle((200., -4.), 50., 6., transform=axes[0].transData, alpha=0.3, color="g"))
# Agrandissement de la mesure sur l'axe 1
testdata[int(200./0.02):int(250./0.02)].plot(ax=axes[1], x="mdfTime", y="psiDt")

Et enfin, la partie vraiment intéressante :

# Créer la flèche
# 1. Obtenir les objets de transformation pour les axes et la figure
ax0tr = axes[0].transData # Axe 0 -> Affichage
ax1tr = axes[1].transData # Axe 1 -> Affichage
figtr = fig.transFigure.inverted() # Affichage -> Figure
# 2. Transformer le point de départ de la flèche de l'axe 0 aux coordonnées de la figure
ptB = figtr.transform(ax0tr.transform((225., -10.)))
# 3. Transformer le point final de la flèche de l'axe 1 aux coordonnées de la figure
ptE = figtr.transform(ax1tr.transform((225., 1.)))
# 4. Créer le patch
arrow = matplotlib.patches.FancyArrowPatch(
    ptB, ptE, transform=fig.transFigure,  # Place la flèche dans le système de coord. de la figure
    fc="g", connectionstyle="arc3,rad=0.2", arrowstyle='simple', alpha=0.3,
    mutation_scale=40.
)
# 5. Ajouter le patch à la liste des objets à dessiner sur la figure
fig.patches.append(arrow)

Si vous travaillez avec des lignes, le processus est similaire, mais vous devez ajouter votre objet personnalisé à la propriété lines de la figure au lieu de patches. Jettez un œil à cette question sur StackOverflow.