1#!/usr/bin/env python
2
3# Author: Andrew Jewett (jewett.aij at g mail)
4# License: MIT License  (See LICENSE.md)
5# Copyright (c) 2017, California Institute of Technology
6# All rights reserved.
7
8g_program_name = __file__.split('/')[-1]  # = 'interpolate_curve.py'
9g_version_str = '0.3.1'
10g_date_str = '2020-4-11'
11
12g_usage_str = """
13Usage:
14
15   """ + g_program_name + """ Ndesired [scale] [alpha] < coords_orig.raw > coords.raw
16
17Example:
18
19   """ + g_program_name + """ 30117 3.0 0.5 < coords_orig.raw > coords.raw
20
21"""
22
23
24import sys
25from math import *
26import numpy as np
27import bisect
28
29
30
31## Tri Diagonal Matrix Algorithm(a.k.a Thomas algorithm) solver stolen from:
32# https://gist.github.com/cbellei/8ab3ab8551b8dfc8b081c518ccd9ada9
33# (also taken from https://gist.github.com/ofan666/1875903)
34
35import numpy as np
36
37## Tri Diagonal Matrix Algorithm(a.k.a Thomas algorithm) solver
38def TDMAsolver(a, b, c, d):
39    '''
40    TDMA solver, a b c d can be NumPy array type or Python list type.
41    Solves a (non-cyclic) system of equations of the form:
42
43      a_i*x_{i-1} + b_i*x_i + c_i*x_{i+1} = d_i
44      where a_1=0, and c_n=0
45
46    refer to http://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
47    and to http://www.cfd-online.com/Wiki/Tridiagonal_matrix_algorithm_-_TDMA_(Thomas_algorithm)
48    '''
49    nf = len(d) # number of equations
50    ac, bc, cc, dc = map(np.array, (a, b, c, d)) # copy arrays
51    for it in range(1, nf):
52        mc = ac[it-1]/bc[it-1]
53        bc[it] = bc[it] - mc*cc[it-1]
54        dc[it] = dc[it] - mc*dc[it-1]
55
56    xc = bc
57    xc[-1] = dc[-1]/bc[-1]
58
59    for il in range(nf-2, -1, -1):
60        xc[il] = (dc[il]-cc[il]*xc[il+1])/bc[il]
61
62    return xc
63
64
65
66def CalcNaturalCubicSplineCoeffs(r, spline_exponent_alpha=0.5):
67    N = len(r)
68    assert(N >= 4)
69    D = len(r[0])
70    tcontrol = np.zeros(N)
71
72    # e_d2rdt2[d][i] is the second derivative of the spline at the ith
73    #               control point (and in the dth direction)
74    #
75    # Once we have figured out e_d2rdt2[d][i], we can calculate the spline
76    # anywhere using:
77    #   SplineEval(t, h_i,
78    #              e_{i+1} / (6*h_i),
79    #              e_i / (6*h_{i+1}),
80    #              r_{i+1}/h_i - e_{i+1}*h_i,
81    #              r_i/h_i - e_i*h_i)
82
83    e_d2rdt2 = np.zeros((D,N))
84
85    # We want to solve this system of equations for e_1, e_2, ..., e_{n-2}:
86    #    (note: e_i is shorthand for e_d2rdt2[d][i])
87    #
88    # h_{i-1}*e_{i-1} + u_i*e_i + h_{i+1} * e_{i+1}  =  v_i
89    #   where h_i, u_i and v_i are shorthand for:
90    #
91    # h_i = change in "time" parameter for the ith interval, which is usually:
92    #     = |r_{i+1} - r_i |^alpha    (alpha varies depending on settings)
93    #   and
94    # u_i = 2 * (h_{i-1} + h_i )
95    # v_i = 6 * (b_i     - b_{i-1})
96    # b_i =  (r_i - r_{i-1}) / h_i
97    #
98    # ...subject to the constraints that the curve is straight at the endpoints:
99    # e_0 = 0       <-- first control point  (indexing begins at 0)
100    # e_{n-1} = 0   <-- this is the last control point  (indexing begins at 0)
101
102    h_dt = np.zeros(N-1)   # h_dt[i] is the i'th time interval
103                         # in the parameterization
104    b_drdt = np.zeros((N-1,D))  # b_drdt is a discrete version of the derivative
105
106    t_total=0.0
107    for i in range(0, N-1):
108        tcontrol[i] = t_total
109        sqr_distance_i_ip1 = 0.0
110        for d in range(0, D):
111            sqr_distance_i_ip1 += (r[i+1][d] - r[i][d])**2
112        h_dt[i] = sqr_distance_i_ip1**(0.5*spline_exponent_alpha)
113        for d in range(0, D):
114            b_drdt[i][d] = (r[i+1][d] - r[i][d]) / h_dt[i]
115        t_total += h_dt[i]
116    tcontrol[N-1] = t_total
117
118    a_coeff = np.zeros(N)
119    b_coeff = np.zeros(N)
120    c_coeff = np.zeros(N)
121    d_coeff = np.zeros((D,N))
122
123    for i in range(1, N-1):
124        # h_dt[i] is the difference in "time" in the parametric curve
125        # between pairs of control points.  If spline_exponenent_alpha is 0
126        # then the time interval between control points is uniform.
127        # ("spline_exponent_alpha" is 0.5 for centripital Catmull-Rom splines.)
128        a_coeff[i]    =      h_dt[i-1]
129        b_coeff[i]    = 2.0*(h_dt[i-1] + h_dt[i])
130        c_coeff[i]    =      h_dt[i]
131        for d in range(0, D):
132            d_coeff[d][i] = 6.0*(b_drdt[i][d] - b_drdt[i-1][d])
133
134    a_coeff[0] = 0.0
135    b_coeff[0] = 1.0
136    c_coeff[0] = 0.0
137    for d in range(0, D):
138        d_coeff[d][0] = 0.0
139    a_coeff[N-1] = 0.0
140    b_coeff[N-1] = 1.0
141    c_coeff[N-1] = 0.0
142    for d in range(0, D):
143        d_coeff[d][N-1] = 0.0
144
145    for d in range(0, D):
146        e_d2rdt2[d] = TDMAsolver(a_coeff, b_coeff, c_coeff, d_coeff[d])
147
148    # alternately, if that fails, try the matrix inverter that comes with numpy:
149    #M = np.zeros((N,N))
150    #for i in range(0,N):
151    #    if i-1>=0:
152    #        M[i][i-1] = a_coeff[i]
153    #    M[i][i]   = b_coeff[i]
154    #    if i+1 < N:
155    #        M[i][i+1] = c_coeff[i]
156    #for d in range(0, D):
157    #    e_d2rdt2[d] = np.linalg.solve(M, d_coeff[d])
158
159    e_d2rdt2 = e_d2rdt2.transpose()
160    c3a = np.zeros((N, D))
161    c3b = np.zeros((N, D))
162    c1a = np.zeros((N, D))
163    c1b = np.zeros((N, D))
164    # faster to precompute these coefficients in advance:
165    for i in range(0, N-1):
166        # c3a = e_{i+1} / (6*h_i)
167        # c3b =     e_i / (6*h_{i+1})
168        # c1a = r_{i+1}/h_i - e_{i+1}*h_i
169        # c1b =  r_i / h_i - e_i*h_i)
170        c3a[i] = e_d2rdt2[i+1] / (6*h_dt[i])
171        c3b[i] = e_d2rdt2[i]   / (6*h_dt[i])
172        c1a[i] = r[i+1]/h_dt[i]  -  e_d2rdt2[i+1]*h_dt[i]/6.0
173        c1b[i] = r[i]/h_dt[i]    -  e_d2rdt2[i]*h_dt[i]/6.0
174
175    # Return these spline coefficients to the caller.
176    # We can use these to quickly evaluate the spline repetatively later on
177
178    return c3a, c3b, c1a, c1b, tcontrol
179
180
181
182
183
184def SplineEval(t, t_interval, c3a, c3b, c1a, c1b):
185    ta = t
186    tb = t_interval - t
187    return (c3a*ta*ta + c1a)*ta + (c3b*tb*tb + c1b)*tb
188
189
190
191def SplineEvalD1(t, t_interval, c3a, c3b, c1a, c1b):
192    ta = t
193    tb = t_interval - t
194    return c3a*ta*ta*3.0 + c1a - (c3b*tb*tb*3.0 + c1b)
195
196
197
198def SplineEvalD2(t, t_interval, c3a, c3b, c1a, c1b):
199    ta = t
200    tb = t_interval - t
201    return c3a*ta*6.0 + c3b*tb*6.0
202
203
204
205def SplineInterpEval(t, c3a, c3b, c1a, c1b, tcontrol):
206    i = bisect.bisect(tcontrol, t) - 1
207    if i < 0: i = 0
208    if i >= len(tcontrol)-1: i = len(tcontrol)-2
209    return SplineEval(t-tcontrol[i], tcontrol[i+1]-tcontrol[i],
210                      c3a[i], c3b[i], c1a[i], c1b[i])
211
212
213def SplineInterpEvalD1(t, c3a, c3b, c1a, c1b, tcontrol):
214    i = bisect.bisect(tcontrol, t) - 1
215    if i < 0: i = 0
216    if i >= len(tcontrol)-1: i = len(tcontrol)-2
217    return SplineEvalD1(t-tcontrol[i], tcontrol[i+1]-tcontrol[i],
218                        c3a[i], c3b[i], c1a[i], c1b[i])
219
220
221
222def SplineInterpEvalD2(t, c3a, c3b, c1a, c1b, tcontrol):
223    i = bisect.bisect(tcontrol, t) - 1
224    if i < 0: i = 0
225    if i >= len(tcontrol)-1: i = len(tcontrol)-2
226    return SplineEvalD2(t-tcontrol[i], tcontrol[i+1]-tcontrol[i],
227                        c3a[i], c3b[i], c1a[i], c1b[i])
228
229
230
231def SplineCurvature2D(t, t_interval, c3a, c3b, c1a, c1b):
232    # first derivatives
233    x1,y1 = SplineEvalD1(t, t_interval, c3a, c3b, c1a, c1b)
234    # second derivatives
235    x2,y2 = SplineEvalD2(t, t_interval, c3a, c3b, c1a, c1b)
236    denom = pow((x1*x1 + y1*y1), 1.5)
237    curvature = (x1*y2 - x2*y1) / denom
238    return curvature
239
240
241
242def SplineInterpCurvature2D(t, c3a, c3b, c1a, c1b, tcontrol):
243    i = bisect.bisect(tcontrol, t) - 1
244    if i < 0: i = 0
245    if i >= len(tcontrol)-1: i = len(tcontrol)-2
246    return SplineCurvature2D(t-tcontrol[i], tcontrol[i+1]-tcontrol[i],
247                             c3a[i], c3b[i], c1a[i], c1b[i])
248
249
250
251def ResampleCurve(x_orig, num_points, alpha=0.5):
252    """
253       Given a list (or numpy array) of points in n-dimensional space that lie
254    along some curve, this function returns a new list of "num_points" points
255    which are uniformly distributed along the original curve.  This is done
256    using cubic spline interpolation (which is tuned by the "alpha" parameter).
257
258    https://en.wikipedia.org/wiki/Cubic_Hermite_spline#Catmull%E2%80%93Rom_spline
259
260       Here "uniformly" means distributed at even intervals in the spline-
261    parameter-space, not in physical distance. (If the original points were
262    not uniformly distributed along the curve, then new points won't be either.)
263
264       This function has not been optimized for speed.
265    """
266
267    assert(len(x_orig) >= 4)
268    num_dimensions = len(x_orig[0])
269    x_new = np.zeros((num_points, num_dimensions))
270
271    c3a, c3b, c1a, c1b, tcontrol = \
272        CalcNaturalCubicSplineCoeffs(x_orig, alpha)
273    tmin = 0.0
274    tmax = tcontrol[-1]
275    for i in range(0, num_points):
276        t = tmin + i*((tmax - tmin) / (num_points-1))
277        x_new[i] = SplineInterpEval(t, c3a, c3b, c1a, c1b, tcontrol)
278
279    return x_new
280
281
282class InputError(Exception):
283    """ A generic exception object containing a string for error reporting.
284        (Raising this exception implies that the caller has provided
285         a faulty input file or argument.)
286
287    """
288
289    def __init__(self, err_msg):
290        self.err_msg = err_msg
291
292    def __str__(self):
293        return self.err_msg
294
295    def __repr__(self):
296        return str(self)
297
298
299def main():
300
301    try:
302
303        #######  Main Code Below: #######
304        sys.stderr.write(g_program_name+' v'+g_version_str+' '+g_date_str+'\n')
305
306        spline_exponent_alpha = 0.5
307        use_linear_interpolation = False
308
309        # Parse the argument list:
310        if len(sys.argv) <= 1:
311            raise InputError("Missing arguments\n"+g_usage_str+"\n")
312
313
314        n_new = int(sys.argv[1])
315
316        if len(sys.argv) > 2:
317            scale = float(sys.argv[2])
318        else:
319            scale = 1.0
320
321        if len(sys.argv) > 3:
322            spline_exponent_alpha = float(sys.argv[3])
323
324        coords_orig = []
325
326        lines = sys.stdin.readlines()
327
328        for line in lines:
329            tokens = line.split()
330            if (len(tokens) > 0):
331                coords_orig.append(list(map(float, tokens)))
332                g_dim = len(tokens)
333
334        n_orig = len(coords_orig)
335
336        x_orig = np.array(coords_orig)
337
338        if n_orig < 2:
339            raise InputError("Input file contains less than two lines of coordinates.")
340
341        if n_orig < 4:
342            use_linear_interpolation = True
343
344        if n_new < 2:
345            raise InputError("Output file will contain less than two lines of coordinates.")
346
347        #coords_new = [[0.0 for d in range(0, g_dim)] for i in range(0, n_new)]
348        x_new = np.zeros((n_new, g_dim))
349
350        if use_linear_interpolation:
351            for i_new in range(0, n_new):
352                I_orig = (i_new) * (float(n_orig-1) / float(n_new-1))
353                i_orig = int(floor(I_orig))
354                i_remainder = I_orig - i_orig
355
356                if (i_new < n_new-1):
357                    for d in range(0, g_dim):
358                        x_new[i_new][d] = (x_orig[i_orig][d]
359                                           +
360                                           i_remainder*(x_orig[i_orig+1][d]-
361                                                        x_orig[i_orig][d]))
362                else:
363                    for d in range(0, g_dim):
364                        x_new[i_new][d] = x_orig[n_orig-1][d]
365
366        else:
367            x_new = ResampleCurve(x_orig, n_new, spline_exponent_alpha)
368
369        # print the coordates
370        for i in range(0, n_new):
371            for d in range(0, g_dim-1):
372                sys.stdout.write(str(x_new[i][d]*scale) + ' ')
373            sys.stdout.write(str(x_new[i][g_dim-1]*scale) + "\n")
374
375    except (ValueError, InputError) as err:
376        sys.stderr.write('\n' + 'Error:\n\n' + str(err) + '\n')
377        sys.exit(1)
378
379
380if __name__ == '__main__':
381    main()
382