import numpy as np
from scipy.optimize import leastsq
from scipy.spatial import cKDTree


def invert_cmap_leastsq(pix, l, colors):
    """
    Given a sequence of pixels, convert each to an equivalent index in the
    color sequence l, colors

    Performs a least squares minimization of the error to find the index.
    color(f) is interpolated.  Error is quadrature sum of color erros.
    """
    pix = pix[..., :3]  # only want rgb
    shape = pix.shape

    ans = np.zeros(list(shape)[:-1], dtype=float)

    def residuals(f, RGB):
        return [np.interp(f, l, colors[:, i]) - RGB[i] for i in range(3)]

    guess = 0.5
    it = np.nditer(pix[..., 0], flags=['multi_index'])
    while not it.finished:
        index = it.multi_index
        # use the answer from the previous iteration to guess the next answer
        guess, stuff = leastsq(residuals, guess, args=(pix[index],))
        ans[index] = guess
        it.iternext()

    return ans


def invert_cmap_argmin(pix, l, colors):
    """
    Given a sequence of pixels, convert each to an equivalent index in the
    color sequence l, colors

    Finds the error between the target pixes and all colors in the cmap.
    Picks the color with the smallest error.
    Error is quadrature sum of color erros.
    """
    shape = pix.shape

    ans = np.zeros(list(shape)[:-1], dtype=float)

    it = np.nditer(pix[..., 0], flags=['multi_index'])
    while not it.finished:
        index = it.multi_index

        err = colors - pix[index]
        err = np.sum(err * err, axis=1)
        ans[index] = l[err.argmin()]
        it.iternext()

    return ans


def invert_cmap_kdtree(pix, l, colors):
    """
    Given a sequence of pixels, convert each to an equivalent index in the
    color sequence l, colors

    Uses a scipy.spatial.cKDTree to find the color with closest coordinates in
    RGB space.
    """
    kd = cKDTree(colors)
    ni, nj, nc = pix.shape
    pix = pix.reshape((ni * nj, nc))
    d, i = kd.query(pix)
    i = i.reshape((ni, nj))
    return l[i]


def order_corners(corners):
    """
    bottom-left, bottom-right, top-right, top-left
    min(x**2+y**2), argmin(y), argmax(x), left
    """
    corners = list(corners)
    assert len(corners) == 4
    assert len(corners[0]) == 2
    ordered = []

    radii = [x ** 2 + y ** 2 for x, y in corners]
    i = np.argmin(radii)
    ordered.append(corners.pop(i))

    i = np.argmin([y for x, y in corners])
    ordered.append(corners.pop(i))

    i = np.argmax([x for x, y in corners])
    ordered.append(corners.pop(i))

    ordered += corners

    return ordered


def get_corner_grid(corners, ni, nj):
    """
    Get uniform grid between corners

    Parameters
    ----------
    corners : list  - of x,y coordinates
    ni : int - number of points in first index
    nj : int - number of points in second index

    Returns
    -------
    x : ndarray, shpae=(ni, nj) - grid of first dimension coordinates
    y : ndarray, shpae=(ni, nj) - grid of second dimension coordinates
    """
    bl, br, tr, tl = order_corners(corners)
    x = np.zeros((ni, nj))

    left_x = _midspace(tl[0], bl[0], ni)
    right_x = _midspace(tr[0], br[0], ni)

    bot_y = _midspace(bl[1], br[1], nj)
    top_y = _midspace(tl[1], tr[1], nj)

    di = 1. / ni
    dj = 1. / nj
    ispan = np.arange(di / 2, 1., di)
    jspan = np.arange(dj / 2, 1., dj)

    x = np.empty((ni, nj))
    x[:, :] = jspan
    x *= (right_x - left_x)[:, None]
    x += left_x[:, None]

    y = np.empty((ni, nj))
    y[:, :] = ispan[:, None]
    y *= (top_y - bot_y)
    y += bot_y
    return x, y


def _midspace(start, end, n):
    """
    Cut region start/end into n regions and return the midpoint of each bin
    """
    x = np.linspace(start, end, n + 1)
    x = x[:-1] + np.diff(x)
    return x

@@ -27,1593 +27,30 @@
 ipy.
-optimize import leastsq%0Afrom scipy.spatial import cKDTree%0A%0A%0Adef invert_cmap_leastsq(pix, l, colors):%0A    %22%22%22%0A    Given a sequence of pixels, convert each to an equivalent index in the%0A    color sequence l, colors%0A%0A    Performs a least squares minimization of the error to find the index.%0A    color(f) is interpolated.  Error is quadrature sum of color erros.%0A    %22%22%22%0A    pix = pix%5B..., :3%5D  # only want rgb%0A    shape = pix.shape%0A%0A    ans = np.zeros(list(shape)%5B:-1%5D, dtype=float)%0A%0A    def residuals(f, RGB):%0A        return %5Bnp.interp(f, l, colors%5B:, i%5D) - RGB%5Bi%5D for i in range(3)%5D%0A%0A    guess = 0.5%0A    it = np.nditer(pix%5B..., 0%5D, flags=%5B'multi_index'%5D)%0A    while not it.finished:%0A        index = it.multi_index%0A        # use the answer from the previous iteration to guess the next answer%0A        guess, stuff = leastsq(residuals, guess, args=(pix%5Bindex%5D,))%0A        ans%5Bindex%5D = guess%0A        it.iternext()%0A%0A    return ans%0A%0A%0Adef invert_cmap_argmin(pix, l, colors):%0A    %22%22%22%0A    Given a sequence of pixels, convert each to an equivalent index in the%0A    color sequence l, colors%0A%0A    Finds the error between the target pixes and all colors in the cmap.%0A    Picks the color with the smallest error.%0A    Error is quadrature sum of color erros.%0A    %22%22%22%0A    shape = pix.shape%0A%0A    ans = np.zeros(list(shape)%5B:-1%5D, dtype=float)%0A%0A    it = np.nditer(pix%5B..., 0%5D, flags=%5B'multi_index'%5D)%0A    while not it.finished:%0A        index = it.multi_index%0A%0A        err = colors - pix%5Bindex%5D%0A        err = np.sum(err * err, axis=1)%0A        ans%5Bindex%5D = l%5Berr.argmin()%5D%0A        it.iternext()%0A%0A    return ans
+spatial import cKDTree
 %0A%0A%0Ad

remove slow invert_cmap_* variants