interp_along_curve.py

interp_along_curve.py

import numpy as np
import matplotlib as mpl

def line(p0, p1, n):
    p0 = np.asarray(p0, dtype=np.float32)
    p1 = np.asarray(p1, dtype=np.float32)
    n = int(n)
    dp = p1 - p0
    l = np.linspace(0.0, 1.0, n)
    dim = len(p0)
    coords = np.zeros([dim, n])
    for d in range(dim):
        coords[d, :] = p0[d] + dp[d] * l
    ds = np.sqrt(np.dot(dp, dp))
    distance = np.linspace(0.0, ds, n)
    return coords, distance

 
def spline(knots, n, dtype=np.float32, **kwargs):
    """
    Args:
        knots : Array-like, of size 3*N, where N is the number of knots.
            For example, knots = [[x0, x1, x2, ...], [y0, y1, y2, ...],
            [z0, z1, z2, ...]].
        n : int. Number of points to be generated along the spline.
        **kwargs : Keyword arguments used by splpre.
    Returns:
        coords : ndarray of size 3*n.
        s : 1d ndarray of accumulated distances at points along the spline.
    """
    import scipy.interpolate as interpolate
    knots = np.asarray(knots, dtype=dtype)
    kmax = len(knots[0]) - 1
    k = min(kmax, kwargs.pop("k", kmax))
    tck, u = interpolate.splprep(knots, k=k, **kwargs)
    u = np.linspace(0, 1, n + 1, endpoint=True)
    coords = np.vstack(interpolate.splev(u, tck))
    x, y, z = coords
    dx = coords[:, 1:] - coords[:, :-1]
    ds = np.zeros_like(coords[0])
    ds[1:] = np.sqrt(np.sum(dx**2, axis=0))
    s = np.cumsum(ds)
    return coords, s

def interp_cartesian(data,
                     coords,
                     coords_t,
                     method='map_coordinates',
                     **kwargs):
    """
    Args:
        data : ndarray.
        coords : A list of coordinates that data lives on. For example, if data is 2d,
            corrds can be [x, y], where x and y are 1d arrays of coordinates.
        coords_t : ndim * npts array.
        method : Method of interpolation. Must be one of 'nearest', 'map_coordinates'.
		kwargs : For map_coordinates method, 'kind' will be used by interp1d to
		    interpolate true coordinates to normalized image coordinates (e.g., 0-1).
			Other args will be used by map_coordinates, for example, order = 0-5,
			mode = 'constant', 'nearest', 'reflect' or 'wrap' for points outside the
			boundaries of true coordinates.
    Returns:
        data_t : 1d array of size npts of interpolated data.
    """
    dim = data.ndim
    if method == 'nearest':
        idx = []
        for d in range(dim)[::-1]:
            x = coords[d]
            xt = coords_t[d]
            nx, = x.shape
            idx.append(np.clip(np.searchsorted(x, xt, **kwargs), 0, nx))
        if type(data) == np.ndarray:
            data_t = data[idx]
        else:  # e.g. h5py dataset does not fully support numpy ndarray's fancy indexing
            idx = np.array(idx)
            npts = coords_t.shape[1]
            pts = []
            for ipt in range(npts):
                pts.append(data[tuple(idx[:, ipt])])
            data_t = np.asarray(pts)
    elif method == 'map_coordinates':
        import scipy.interpolate as interpolate
        import scipy.ndimage as ndimage
        kind = kwargs.pop('kind', 'linear')
        xti = []
        for d in range(dim)[::-1]:
            x = coords[d]
            xi = np.arange(len(coords[d]))  # image coordinates, i.e., 0 to n-1
            xt = coords_t[d]
            xti.append(interpolate.interp1d(x, xi, kind=kind)(xt))
        data_t = ndimage.map_coordinates(data, np.vstack(xti), **kwargs)
    else:
        raise KeyError("method must be one of 'nearest', 'map_coordinates'")
    return data_t