contour_tricks.py

import numpy as np
from matplotlib import pyplot as plt
from scipy.interpolate import griddata


def add_label(x, y, x0=None, y0=None, dx=0, dy=0, s='', **kwargs):
    "add label on the curve (x, y) at location specified by (x0, y0)"
    if x0 is not None:
        if x[-1] > x[0]:
            y0 = np.interp(x0, x, y)
        else:
            y0 = np.interp(x0, x[::-1], y[::-1])
    else:
        if y[-1] > y[0]:
            x0 = np.interp(y0, y, x)
        else:
            x0 = np.interp(y0, y[::-1], x[::-1])
    return plt.text(x0 + dx, y0 + dy, s=s, **kwargs)


def get_contour_line(x, y, z, levels):
    """
    Yields y(x) for constant z specified by levels.
    plot(x, ynew) is equivalent to contour(x, y, z, levels).
    it only works if y is a monotonic function of x, z

    x: (n,)
    y: (m,)
    z: (n, m)
    """
    xx, yy = np.meshgrid(x, y, indexing='ij')
    xx_zz = np.c_[xx.ravel(), z.ravel()]

    xnew, znew = np.meshgrid(x, levels, indexing='ij')
    ynew = griddata(xx_zz, yy.ravel(), (xnew, znew), method='cubic').T
    return ynew