How to create a matplotlib axis tight to data?

I am struggling in creating a curved axis that would follow a contour plot. Supposing I have a contour plot like in the figure:

How could I force the axis to follow the contour, like in the following?

I found this link with something like it in end of the page but I can't find out how to do it.

Here is the code I made to create the figure.

import matplotlib.pyplot as plt
from numpy import linspace, meshgrid, zeros

n = 100
x = linspace(-1, 1, n)
y = linspace(-1, 1, n)
X, Y = meshgrid(x, y)
Z = zeros(X.shape)
for i in range(n):
    X[i, :] = linspace(-1. + i / n * 0.2, 1. - i / n * 0.2, n)
    Y[:, i] = linspace(-x[i] ** 2 - 2, x[i] ** 2 + 2, n)
    for j in range(n):
        Z[i, j] = x[i] ** 2 + y[j] ** 2

fig, ax = plt.subplots()

ax.contour(X, Y, Z, 10, linewidths=0.5, colors='k')
c = ax.contourf(X, Y, Z, 100, cmap=plt.cm.jet)
ax.set_xlabel("x (m)")
ax.set_xlabel("y (m)")
plt.show()

Thanks for the help. :)