Search code examples
pythonnumpyscipyresampling

Resample trajectory to have equal euclidean distance in each sample

Let's say we have a list of x,y points:

x = [0, 0, 0]
y = [0, 10, 100]

The Euclidean distance between points is now [10, 90].
I'm looking for a function that accepts x, y and the sample_rate, and could output equal distance points. e.g.:

x = [0, 0, 0]
y = [0, 10, 100]

resample_distance = 10
resampler(x, y, resample_distance)
# Outputs:
# [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
# [0, 0, 0, 0, 0,0, 0, 0, 0, 0, 0]

Using an similiar answer, we can achieve almost correct values, but that's not accurate:

resample_trajectory_same_distance(data[0], data[1], 10)
# Output:
# [ 0.        , 10.27027027, 20.81081081, 31.08108108, 41.62162162, 51.89189189, 62.43243243, 72.7027027 , 83.24324324, 93.78378378]
# [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]

Using any 3rd party libs like numpy, scipy, etc is OK.

Solution

  • I implemented next solution.

    All functions for efficiency are compiled by Numba compiler/optimizer supporting Just-in-Time/Ahead-of-Time technologies. Numba converts all marked by @numba.njit decorator functions to pure C/C++ code automatically on the fly whenever Python code is started, then C++ code is compiled to machine code. No interaction with Python is done in such functions, only low-level fast structures are used inside. Hence Numba usually is able to increase speed of almost any code by factor of 50x-200x times on average, very fast! Such Python compiled code usually achieves speed very near to speed of same algorithms implemented by hand in pure C/C++. To use Numba one just needs to install just two next python packages: python -m pip install numpy numba.

    Next steps are done in my code:

    1. Input function is represented by two 1D arrays x and y.
    2. Input function (points) is then Approximated (or called differently as Interpolated) by one of two types of piece-wise functions - 1) Linear 2) Cubic Spline.
    3. Linear piece-wise approximation function just connects given array of points (x0, y0) so that pair of two consecutive points (x1, y1) and (x2, y2) are connected by straight line (segment).
    4. Cubic Spline is more advanced way of smoothly approximating function, it connects all points (x0, y0) so that each pair of points (x1, y1) and (x2, y2) is connected by cubic segment represented by cubic polynomial so that it goes through these two end-points plus first and second derivative of adjacent segments within common point are equal, these all ensures that function looks very smooth and nice, and is very popular in computer graphics for doing natural/realistic approximation/visualization of functions/surfaces.
    5. Then using very fast Binary Search Algorithm I'm searching one-by-one for points on this Interpolation function, points such that Euclidean Distance between each two consecutive points equals exactly to desired (provided to algorithm) value d.
    6. All above is just computational part. The rest of steps does visualizing part, drawing plots using matplotlib library. Detailed description of plots goes after code right before plots.

    In order to use this implemented Euclidean Equal-Distance Resampling algorithm you need just to import my next script-module and do xr, yr = module_name.resample_euclid_equidist(x, y, dist) where input and output x and y are both 1D numpy arrays with floating point numbers, this will return input points resampled at dist euclidean distance. More examples of usage are located in my code's test() function. Very first run is quite slow (can take around 15 seconds), this run is just a compilation run, all my code is automatically precompiled to C/C++ and then machine code, next runs are very fast, especially resampling function itself just takes some milliseconds. Also to use just computational part of code you need to install python -m pip install numpy numba, and to run whole of my code including tests and visualization (just run my script) you need to install python -m pip install numpy numba matplotlib just one time.

    Try it online!

    # Needs:
#     For computation: python -m pip install numpy numba
#     For testing:     python -m pip install matplotlib

if __name__ == '__main__':
    print('Compiling...', flush = True)
    
import numba, numpy as np

# Linear Approximator related functions

# Spline value calculating function, given params and "x"
@numba.njit(cache = True, fastmath = True, inline = 'always')
def func_linear(x, ix, x0, y0, k):
    return (x - x0[ix]) * k[ix] + y0[ix]
    
# Compute piece-wise linear function for "x" out of sorted "x0" points
@numba.njit([f'f{ii}[:](f{ii}[:], f{ii}[:], f{ii}[:], f{ii}[:], f{ii}[:], f{ii}[:])' for ii in (4, 8)],
    cache = True, fastmath = True, inline = 'always')
def piece_wise_linear(x, x0, y0, k, dummy0, dummy1):
    xsh = x.shape
    x = x.ravel()
    ix = np.searchsorted(x0[1 : -1], x)
    y = func_linear(x, ix, x0, y0, k)
    y = y.reshape(xsh)
    return y
    
# Spline Approximator related functions
    
# Solves linear system given by Tridiagonal Matrix
# Helper for calculating cubic splines
@numba.njit(cache = True, fastmath = True, inline = 'always')
def tri_diag_solve(A, B, C, F):
    n = B.size
    assert A.ndim == B.ndim == C.ndim == F.ndim == 1 and (
        A.size == B.size == C.size == F.size == n
    ) #, (A.shape, B.shape, C.shape, F.shape)
    Bs, Fs = np.zeros_like(B), np.zeros_like(F)
    Bs[0], Fs[0] = B[0], F[0]
    for i in range(1, n):
        Bs[i] = B[i] - A[i] / Bs[i - 1] * C[i - 1]
        Fs[i] = F[i] - A[i] / Bs[i - 1] * Fs[i - 1]
    x = np.zeros_like(B)
    x[-1] = Fs[-1] / Bs[-1]
    for i in range(n - 2, -1, -1):
        x[i] = (Fs[i] - C[i] * x[i + 1]) / Bs[i]
    return x
    
# Calculate cubic spline params
@numba.njit(cache = True, fastmath = True, inline = 'always')
def calc_spline_params(x, y):
    a = y
    h = np.diff(x)
    c = np.concatenate((np.zeros((1,), dtype = y.dtype),
        np.append(tri_diag_solve(h[:-1], (h[:-1] + h[1:]) * 2, h[1:],
        ((a[2:] - a[1:-1]) / h[1:] - (a[1:-1] - a[:-2]) / h[:-1]) * 3), 0)))
    d = np.diff(c) / (3 * h)
    b = (a[1:] - a[:-1]) / h + (2 * c[1:] + c[:-1]) / 3 * h
    return a[1:], b, c[1:], d
    
# Spline value calculating function, given params and "x"
@numba.njit(cache = True, fastmath = True, inline = 'always')
def func_spline(x, ix, x0, a, b, c, d):
    dx = x - x0[1:][ix]
    return a[ix] + (b[ix] + (c[ix] + d[ix] * dx) * dx) * dx
    
# Compute piece-wise spline function for "x" out of sorted "x0" points
@numba.njit([f'f{ii}[:](f{ii}[:], f{ii}[:], f{ii}[:], f{ii}[:], f{ii}[:], f{ii}[:])' for ii in (4, 8)],
    cache = True, fastmath = True, inline = 'always')
def piece_wise_spline(x, x0, a, b, c, d):
    xsh = x.shape
    x = x.ravel()
    ix = np.searchsorted(x0[1 : -1], x)
    y = func_spline(x, ix, x0, a, b, c, d)
    y = y.reshape(xsh)
    return y
    
# Appximates function given by (x0, y0) by piece-wise spline or linear
def approx_func(x0, y0, t = 'spline'): # t is spline/linear
    assert x0.ndim == 1 and y0.ndim == 1 and x0.size == y0.size#, (x0.shape, y0.shape)
    n = x0.size - 1
    if t == 'linear':
        k = np.diff(y0) / np.diff(x0)
        return piece_wise_linear, (x0, y0, k, np.zeros((0,), dtype = y0.dtype), np.zeros((0,), dtype = y0.dtype))
    elif t == 'spline':
        a, b, c, d = calc_spline_params(x0, y0)
        return piece_wise_spline, (x0, a, b, c, d)
    else:
        assert False, t

# Main function that computes Euclidian Equi-Distant points based on approximation function
@numba.njit(
    [f'f{ii}[:, :](f{ii}[:], f{ii}[:], f{ii}, b1, b1, f{ii}, f{ii}, f{ii}[:], f{ii}[:], f{ii}[:], f{ii}[:], f{ii}[:])' for ii in (4, 8)],
    cache = True, fastmath = True)
def _resample_inner(x, y, d, is_spline, strict, aerr, rerr, a0, a1, a2, a3, a4):
    rs, r = 0, np.zeros((1 << 10, 2), dtype = y.dtype)
    t0 = np.zeros((1,), dtype = y.dtype)
    i, x0, y0 = 0, x[0], y[0]
    #print(i, x0, y0, np.sin(x0))
    while True:
        if rs >= r.size:
            r = np.concatenate((r, np.zeros(r.shape, dtype = r.dtype))) # Grow array
        r[rs, 0] = x0
        r[rs, 1] = y0
        rs += 1
        if i + 1 >= x.size:
            break
        ie = min(i + 1 + np.searchsorted(x[i + 1:], x0 + d), x.size - 1)
        for ie in range(i + 1 if strict else ie, ie + 1):
            xl = max(x0, x[ie - 1 if strict else i])
            xr = max(x0, x[ie])
            # Do binary search to find next point
            for ii in range(1000):
                if xr - xl <= aerr:
                    break # Already very small delta X interval
                xm = (xl + xr) / 2
                t0[0] = xm
                if is_spline:
                    ym = piece_wise_spline(t0, a0, a1, a2, a3, a4)[0]
                else:
                    ym = piece_wise_linear(t0, a0, a1, a2, a3, a4)[0]
                
                # Compute Euclidian distance
                dx_, dy_ = xm - x0, ym - y0
                dm = np.sqrt(dx_ * dx_ + dy_ * dy_)

                if -rerr <= dm / d - 1. <= rerr:
                    break # We got d with enough precision
                if dm >= d:
                    xr = xm
                else:
                    xl = xm
            else:
                assert False # To many iterations
            if -rerr <= dm / d - 1. <= rerr:
                break # Next point found
        else:
            break # No next point found, we're finished
        i = np.searchsorted(x, xm) - 1
        #print('_0', i, x0, y0, np.sin(x0), dist(x0, xm, y0, ym), dist(x0, xm, np.sin(x0), np.sin(xm)))
        x0, y0 = xm, ym
        #print('_1', i, x0, y0, np.sin(x0), dm)
    return r[:rs]
    
# Resamples (x, y) points using given approximation function type
# so that euclidian distance between each resampled points equals to "d".
# If strict = True then strictly closest (by X) next point at distance "d"
# is chosen, which can imply more computations, when strict = False then
# any found point with distance "d" is taken as next.
def resample_euclid_equidist(
    x, y, d, *,
    aerr = 2 ** -21, rerr = 2 ** -9, approx = 'spline',
    return_approx = False, strict = True,
):
    assert d > 0, d
    dtype = np.dtype(y.dtype).type
    x, y, d, aerr, rerr = [dtype(e) for e in [x, y, d, aerr, rerr]]
    ixs = np.argsort(x)
    x, y = x[ixs], y[ixs]
    f, fargs = approx_func(x, y, approx)
    r = _resample_inner(x, y, d, approx == 'spline', strict, aerr, rerr, *fargs)
    return (r[:, 0], r[:, 1]) + ((), (lambda x: f(x, *fargs),))[return_approx]

def test():
    import matplotlib.pyplot as plt, numpy as np, time
    np.random.seed(0)
    # Input
    n = 50
    x = np.sort(np.random.uniform(0., 10 * np.pi, (n,)))
    y = np.sin(x) * 5 + np.sin(1 + 2.5 * x) * 3 + np.sin(2 + 0.5 * x) * 2
    # Visualize
    for isl, sl in enumerate(['spline', 'linear']):
        # Compute resampled points
        for i in range(3):
            tb = time.time()
            xa, ya, fa = resample_euclid_equidist(x, y, 1.5, approx = sl, return_approx = True)
            print(sl, 'try', i, 'run time', round(time.time() - tb, 4), 'sec', flush = True)
        # Compute spline/linear approx points
        fax = np.linspace(x[0], x[-1], 1000)
        fay = fa(fax)
        # Plotting
        plt.rcParams['figure.figsize'] = (9.6, 5.4)
        for ci, (cx, cy, fn) in enumerate([
            (x, y, 'original'), (fax, fay, f'approx_{sl}'), (xa, ya, 'euclid_euqidist'),
        ]):
            p, = plt.plot(cx, cy)
            p.set_label(fn)
            if ci >= 2:
                plt.scatter(cx, cy, marker = '.', color = p.get_color())
                if False:
                    # Show distances
                    def dist(x0, x1, y0, y1):
                        # Compute Euclidian distance
                        dx, dy = x1 - x0, y1 - y0
                        return np.sqrt(dx * dx + dy * dy)
                    for i in range(cx.size - 1):
                        plt.annotate(
                            round(dist(cx[i], cx[i + 1], cy[i], cy[i + 1]), 2),
                            (cx[i], cy[i]), fontsize = 'xx-small',
                        )
        plt.gca().set_aspect('equal', adjustable = 'box')
        plt.legend()
        plt.show()
        plt.clf()

if __name__ == '__main__':
    test()

    Below are resulting plots. As an example function is taken y = np.sin(x) * 5 + np.sin(1 + 2.5 * x) * 3 + np.sin(2 + 0.5 * x) * 2 sampled at 50 uniformly random points for 0 <= x <= 10 * pi range. Plots: blue is original function, connected with straight line points, orange is approximation function (spline or linear) this is just interpolating function, it is drawn as hundreds of points that's why looks smooth, green is resulting Euclidian-Equal-Distance points, exactly what was task about, euclidian length of each segment between two green small circles is equal exactly to desired distance d. First screen represents approximation by piece-wise cubic spline. Second screen represents approximation for by piece-wise linear function for exactly same input points.

    Spline:

    spline

    Linear:

    linear