Coverage for src/meshpy/cosserat_curve/cosserat_curve.py: 98%

1# The MIT License (MIT) 

2# 

3# Copyright (c) 2018-2025 MeshPy Authors 

4# 

5# Permission is hereby granted, free of charge, to any person obtaining a copy 

6# of this software and associated documentation files (the "Software"), to deal 

7# in the Software without restriction, including without limitation the rights 

8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

9# copies of the Software, and to permit persons to whom the Software is 

10# furnished to do so, subject to the following conditions: 

11# 

12# The above copyright notice and this permission notice shall be included in 

13# all copies or substantial portions of the Software. 

14# 

15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 

21# THE SOFTWARE. 

22"""Define a Cosserat curve object that can be used to describe warping of 

23curve-like objects.""" 

24 

25from typing import Tuple as _Tuple 

26 

27import numpy as _np 

28import pyvista as _pv 

29import quaternion as _quaternion 

30from numpy.typing import NDArray as _NDArray 

31from scipy import integrate as _integrate 

32from scipy import interpolate as _interpolate 

33from scipy import optimize as _optimize 

34 

35from meshpy.core.conf import mpy as _mpy 

36from meshpy.core.rotation import Rotation as _Rotation 

37from meshpy.core.rotation import rotate_coordinates as _rotate_coordinates 

38from meshpy.core.rotation import smallest_rotation as _smallest_rotation 

39 

40 

41def get_piecewise_linear_arc_length_along_points( 

42 coordinates: _np.ndarray, 

43) -> _np.ndarray: 

44 """Return the accumulated distance between the points. 

45 

46 Args 

47 ---- 

48 coordinates: 

49 Array containing the point coordinates 

50 """ 

51 

52 n_points = len(coordinates) 

53 point_distance = _np.linalg.norm(coordinates[1:] - coordinates[:-1], axis=1) 

54 point_arc_length = _np.zeros(n_points) 

55 for i in range(1, n_points): 

56 point_arc_length[i] = point_arc_length[i - 1] + point_distance[i - 1] 

57 return point_arc_length 

58 

59 

60def get_spline_interpolation( 

61 coordinates: _np.ndarray, point_arc_length: _np.ndarray 

62) -> _interpolate.BSpline: 

63 """Get a spline interpolation of the given points. 

64 

65 Args 

66 ---- 

67 coordinates: 

68 Array containing the point coordinates 

69 point_arc_length: 

70 Arc length for each coordinate 

71 

72 Return 

73 ---- 

74 centerline_interpolation: 

75 The spline interpolation object 

76 """ 

77 

78 # Interpolate coordinates along arc length 

79 centerline_interpolation = _interpolate.make_interp_spline( 

80 point_arc_length, coordinates 

81 ) 

82 return centerline_interpolation 

83 

84 

85def get_quaternions_along_curve( 

86 centerline: _interpolate.BSpline, point_arc_length: _np.ndarray 

87) -> _NDArray[_quaternion.quaternion]: 

88 """Get the quaternions along the curve based on smallest rotation mappings. 

89 

90 The initial rotation will be calculated based on the largest projection of the initial tangent 

91 onto the cartesian basis vectors. 

92 

93 Args 

94 ---- 

95 centerline: 

96 A function that returns the centerline position for a parameter coordinate t 

97 point_arc_length: 

98 Array of parameter coordinates for which the quaternions should be calculated 

99 """ 

100 

101 centerline_interpolation_derivative = centerline.derivative() 

102 

103 def basis(i): 

104 """Return the i-th Cartesian basis vector.""" 

105 basis = _np.zeros([3]) 

106 basis[i] = 1.0 

107 return basis 

108 

109 # Get the reference rotation 

110 t0 = centerline_interpolation_derivative(point_arc_length[0]) 

111 min_projection = _np.argmin(_np.abs([_np.dot(basis(i), t0) for i in range(3)])) 

112 last_rotation = _Rotation.from_basis(t0, basis(min_projection)) 

113 

114 # Get the rotation vectors along the curve. They are calculated with smallest rotation mappings. 

115 n_points = len(point_arc_length) 

116 quaternions = _np.zeros(n_points, dtype=_quaternion.quaternion) 

117 quaternions[0] = last_rotation.q 

118 for i in range(1, n_points): 

119 rotation = _smallest_rotation( 

120 last_rotation, 

121 centerline_interpolation_derivative(point_arc_length[i]), 

122 ) 

123 quaternions[i] = rotation.q 

124 last_rotation = rotation 

125 return quaternions 

126 

127 

128def get_relative_distance_and_rotations( 

129 coordinates: _np.ndarray, quaternions: _NDArray[_quaternion.quaternion] 

130) -> _Tuple[ 

131 _np.ndarray, _NDArray[_quaternion.quaternion], _NDArray[_quaternion.quaternion] 

132]: 

133 """Get relative distances and rotations that can be used to evaluate 

134 "intermediate" states of the Cosserat curve.""" 

135 

136 n_points = len(coordinates) 

137 relative_distances = _np.zeros(n_points - 1) 

138 relative_distances_rotation = _np.zeros(n_points - 1, dtype=_quaternion.quaternion) 

139 relative_rotations = _np.zeros(n_points - 1, dtype=_quaternion.quaternion) 

140 

141 for i_segment in range(n_points - 1): 

142 relative_distance = coordinates[i_segment + 1] - coordinates[i_segment] 

143 relative_distance_local = _quaternion.rotate_vectors( 

144 quaternions[i_segment].conjugate(), relative_distance 

145 ) 

146 relative_distances[i_segment] = _np.linalg.norm(relative_distance_local) 

147 

148 smallest_relative_rotation_onto_distance = _smallest_rotation( 

149 _Rotation(), 

150 relative_distance_local, 

151 ) 

152 relative_distances_rotation[i_segment] = _quaternion.from_float_array( 

153 smallest_relative_rotation_onto_distance.q 

154 ) 

155 

156 relative_rotations[i_segment] = ( 

157 quaternions[i_segment].conjugate() * quaternions[i_segment + 1] 

158 ) 

159 

160 return relative_distances, relative_distances_rotation, relative_rotations 

161 

162 

163class CosseratCurve(object): 

164 """Represent a Cosserat curve in space.""" 

165 

166 def __init__(self, point_coordinates: _np.ndarray): 

167 """Initialize the Cosserat curve based on points in 3D space. 

168 

169 Args 

170 ---- 

171 point_coordinates: 

172 Array containing the point coordinates 

173 """ 

174 

175 self.coordinates = point_coordinates.copy() 

176 self.n_points = len(self.coordinates) 

177 

178 # Interpolate coordinates along piece wise linear arc length 

179 point_arc_length_piecewise_linear = ( 

180 get_piecewise_linear_arc_length_along_points(self.coordinates) 

181 ) 

182 centerline_interpolation_piecewise_linear = get_spline_interpolation( 

183 self.coordinates, point_arc_length_piecewise_linear 

184 ) 

185 centerline_interpolation_piecewise_linear_p = ( 

186 centerline_interpolation_piecewise_linear.derivative(1) 

187 ) 

188 

189 def ds(t): 

190 """Arc length along interpolated spline.""" 

191 return _np.linalg.norm(centerline_interpolation_piecewise_linear_p(t)) 

192 

193 # Integrate the arc length along the interpolated centerline, this will result 

194 # in a more accurate centerline arc length 

195 self.point_arc_length = _np.zeros(self.n_points) 

196 for i in range(len(point_arc_length_piecewise_linear) - 1): 

197 self.point_arc_length[i + 1] = ( 

198 self.point_arc_length[i] 

199 + _integrate.quad( 

200 ds, 

201 point_arc_length_piecewise_linear[i], 

202 point_arc_length_piecewise_linear[i + 1], 

203 )[0] 

204 ) 

205 

206 # Set the interpolation of the (positional) centerline 

207 self.set_centerline_interpolation() 

208 

209 # Get the quaternions along the centerline based on smallest rotation mappings 

210 self.quaternions = get_quaternions_along_curve( 

211 self.centerline_interpolation, self.point_arc_length 

212 ) 

213 

214 # Get the relative quantities used to warp the curve 

215 ( 

216 self.relative_distances, 

217 self.relative_distances_rotation, 

218 self.relative_rotations, 

219 ) = get_relative_distance_and_rotations(self.coordinates, self.quaternions) 

220 

221 def set_centerline_interpolation(self): 

222 """Set the interpolation of the centerline based on the coordinates and 

223 arc length stored in this object.""" 

224 self.centerline_interpolation = get_spline_interpolation( 

225 self.coordinates, self.point_arc_length 

226 ) 

227 

228 def translate(self, vector): 

229 """Translate the curve by the given vector.""" 

230 

231 self.coordinates += vector 

232 self.set_centerline_interpolation() 

233 

234 def rotate(self, rotation: _Rotation, *, origin=None): 

235 """Rotate the curve and the quaternions.""" 

236 

237 self.quaternions = _quaternion.from_float_array(rotation.q) * self.quaternions 

238 self.coordinates = _rotate_coordinates( 

239 self.coordinates, rotation, origin=origin 

240 ) 

241 self.set_centerline_interpolation() 

242 

243 def get_centerline_position_and_rotation( 

244 self, arc_length: float, **kwargs 

245 ) -> _Tuple[_np.ndarray, _NDArray[_quaternion.quaternion]]: 

246 """Return the position and rotation at a given centerline arc 

247 length.""" 

248 pos, rot = self.get_centerline_positions_and_rotations([arc_length]) 

249 return pos[0], rot[0] 

250 

251 def get_centerline_positions_and_rotations( 

252 self, points_on_arc_length, *, factor=1.0 

253 ) -> _Tuple[_np.ndarray, _NDArray[_quaternion.quaternion]]: 

254 """Return the position and rotation at given centerline arc lengths. 

255 

256 If the points are outside of the valid interval, a linear extrapolation will be 

257 performed for the displacements and the rotations will be held constant. 

258 

259 Args 

260 ---- 

261 points_on_arc_length: list(float) 

262 A sorted list with the arc lengths along the curve centerline 

263 factor: float 

264 Factor to scale the curvature along the curve. 

265 factor == 1 

266 Use the default positions and the triads obtained via a smallest rotation mapping 

267 factor < 1 

268 Integrate (piecewise constant as evaluated with get_relative_distance_and_rotations) 

269 the scaled curvature of the curve to obtain a intuitive wrapping 

270 """ 

271 

272 # Get the points that are within the arc length of the given curve. 

273 points_on_arc_length = _np.asarray(points_on_arc_length) 

274 points_in_bounds = _np.logical_and( 

275 points_on_arc_length > self.point_arc_length[0], 

276 points_on_arc_length < self.point_arc_length[-1], 

277 ) 

278 index_in_bound = _np.where(points_in_bounds == True)[0] 

279 index_out_of_bound = _np.where(points_in_bounds == False)[0] 

280 points_on_arc_length_in_bound = [ 

281 self.point_arc_length[0], 

282 *points_on_arc_length[index_in_bound], 

283 self.point_arc_length[-1], 

284 ] 

285 

286 if factor < (1.0 - _mpy.eps_quaternion): 

287 coordinates = _np.zeros_like(self.coordinates) 

288 quaternions = _np.zeros_like(self.quaternions) 

289 coordinates[0] = self.coordinates[0] 

290 quaternions[0] = self.quaternions[0] 

291 for i_segment in range(self.n_points - 1): 

292 relative_distance_rotation = _quaternion.slerp_evaluate( 

293 _quaternion.quaternion(1), 

294 self.relative_distances_rotation[i_segment], 

295 factor, 

296 ) 

297 # In the initial configuration (factor=0) we get a straight curve, so we need 

298 # to use the arc length here. In the final configuration (factor=1) we want to 

299 # exactly recover the input points, so we need the piecewise linear distance. 

300 # Between them, we interpolate. 

301 relative_distance = (factor * self.relative_distances[i_segment]) + ( 

302 1.0 - factor 

303 ) * ( 

304 self.point_arc_length[i_segment + 1] 

305 - self.point_arc_length[i_segment] 

306 ) 

307 coordinates[i_segment + 1] = ( 

308 _quaternion.rotate_vectors( 

309 quaternions[i_segment] * relative_distance_rotation, 

310 [relative_distance, 0, 0], 

311 ) 

312 + coordinates[i_segment] 

313 ) 

314 quaternions[i_segment + 1] = quaternions[ 

315 i_segment 

316 ] * _quaternion.slerp_evaluate( 

317 _quaternion.quaternion(1), 

318 self.relative_rotations[i_segment], 

319 factor, 

320 ) 

321 else: 

322 coordinates = self.coordinates 

323 quaternions = self.quaternions 

324 

325 sol_r = _np.zeros([len(points_on_arc_length_in_bound), 3]) 

326 sol_q = _np.zeros( 

327 len(points_on_arc_length_in_bound), dtype=_quaternion.quaternion 

328 ) 

329 for i_point, centerline_arc_length in enumerate(points_on_arc_length_in_bound): 

330 if ( 

331 centerline_arc_length >= self.point_arc_length[0] 

332 and centerline_arc_length <= self.point_arc_length[-1] 

333 ): 

334 for i in range(1, self.n_points): 

335 centerline_index = i - 1 

336 if self.point_arc_length[i] > centerline_arc_length: 

337 break 

338 

339 # Get the two rotation vectors and arc length values 

340 arc_lengths = self.point_arc_length[ 

341 centerline_index : centerline_index + 2 

342 ] 

343 q1 = quaternions[centerline_index] 

344 q2 = quaternions[centerline_index + 1] 

345 

346 # Linear interpolate the arc length 

347 xi = (centerline_arc_length - arc_lengths[0]) / ( 

348 arc_lengths[1] - arc_lengths[0] 

349 ) 

350 

351 # Perform a spline interpolation for the positions and a slerp 

352 # interpolation for the rotations 

353 sol_r[i_point] = get_spline_interpolation( 

354 coordinates, self.point_arc_length 

355 )(centerline_arc_length) 

356 sol_q[i_point] = _quaternion.as_float_array( 

357 _quaternion.slerp_evaluate(q1, q2, xi) 

358 ) 

359 else: 

360 raise ValueError("Centerline value out of bounds") 

361 

362 # Set the already computed results in the final data structures 

363 sol_r_final = _np.zeros([len(points_on_arc_length), 3]) 

364 sol_q_final = _np.zeros(len(points_on_arc_length), dtype=_quaternion.quaternion) 

365 if len(index_in_bound) > 0: 

366 sol_r_final[index_in_bound] = sol_r[index_in_bound - index_in_bound[0] + 1] 

367 sol_q_final[index_in_bound] = sol_q[index_in_bound - index_in_bound[0] + 1] 

368 

369 # Perform the extrapolation at both ends of the curve 

370 for i in index_out_of_bound: 

371 arc_length = points_on_arc_length[i] 

372 if arc_length <= self.point_arc_length[0]: 

373 index = 0 

374 elif arc_length >= self.point_arc_length[-1]: 

375 index = -1 

376 else: 

377 raise ValueError("Should not happen") 

378 

379 length = arc_length - self.point_arc_length[index] 

380 r = sol_r[index] 

381 q = sol_q[index] 

382 sol_r_final[i] = r + _Rotation.from_quaternion( 

383 _quaternion.as_float_array(q) 

384 ) * [length, 0, 0] 

385 sol_q_final[i] = q 

386 

387 return sol_r_final, sol_q_final 

388 

389 def project_point(self, p, t0=None) -> float: 

390 """Project a point to the curve, return the parameter coordinate for 

391 the projection point.""" 

392 

393 centerline_interpolation_p = self.centerline_interpolation.derivative(1) 

394 centerline_interpolation_pp = self.centerline_interpolation.derivative(2) 

395 

396 def f(t): 

397 """Function to find the root of.""" 

398 r = self.centerline_interpolation(t) 

399 rp = centerline_interpolation_p(t) 

400 return _np.dot(r - p, rp) 

401 

402 def fp(t): 

403 """Derivative of the Function to find the root of.""" 

404 r = self.centerline_interpolation(t) 

405 rp = centerline_interpolation_p(t) 

406 rpp = centerline_interpolation_pp(t) 

407 return _np.dot(rp, rp) + _np.dot(r - p, rpp) 

408 

409 if t0 is None: 

410 t0 = 0.0 

411 

412 return _optimize.newton(f, t0, fprime=fp) 

413 

414 def get_pyvista_polyline(self) -> _pv.PolyData: 

415 """Create a pyvista (vtk) representation of the curve with the 

416 evaluated triad basis vectors.""" 

417 

418 poly_line = _pv.PolyData() 

419 poly_line.points = self.coordinates 

420 cell = _np.arange(0, self.n_points, dtype=int) 

421 cell = _np.insert(cell, 0, self.n_points) 

422 poly_line.lines = cell 

423 

424 rotation_matrices = _np.zeros((len(self.quaternions), 3, 3)) 

425 for i_quaternion, q in enumerate(self.quaternions): 

426 R = _quaternion.as_rotation_matrix(q) 

427 rotation_matrices[i_quaternion] = R 

428 

429 for i_dir in range(3): 

430 poly_line.point_data.set_array( 

431 rotation_matrices[:, :, i_dir], f"base_vector_{i_dir + 1}" 

432 ) 

433 

434 return poly_line 

435 

436 def write_vtk(self, path) -> None: 

437 """Save a vtk representation of the curve.""" 

438 self.get_pyvista_polyline().save(path)