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