Coverage for src/beamme/cosserat_curve/cosserat_curve.py: 98%
198 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-28 15:20 +0000
1# The MIT License (MIT)
2#
3# Copyright (c) 2018-2026 BeamMe 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 curve-like
23objects."""
25from pathlib import Path as _Path
26from xml.etree import ElementTree as _ET # nosec B405
28import numpy as _np
29import pyvista as _pv
30import quaternion as _quaternion
31from numpy.typing import NDArray as _NDArray
32from scipy import integrate as _integrate
33from scipy import interpolate as _interpolate
34from scipy import optimize as _optimize
36from beamme.core.conf import bme as _bme
37from beamme.core.rotation import Rotation as _Rotation
38from beamme.core.rotation import rotate_coordinates as _rotate_coordinates
39from beamme.core.rotation import smallest_rotation as _smallest_rotation
42def get_piecewise_linear_arc_length_along_points(
43 coordinates: _np.ndarray,
44) -> _np.ndarray:
45 """Return the accumulated distance between the points.
47 Args
48 ----
49 coordinates:
50 Array containing the point coordinates
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
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.
65 Args
66 ----
67 coordinates:
68 Array containing the point coordinates
69 point_arc_length:
70 Arc length for each coordinate
72 Return
73 ----
74 centerline_interpolation:
75 The spline interpolation object
76 """
77 # Interpolate coordinates along arc length
78 # Note: The numeric evaluation of the spline interpolation can depend on the
79 # operating system, thus introducing slight numerical differences (~1e-12).
80 centerline_interpolation = _interpolate.make_interp_spline(
81 point_arc_length, coordinates
82 )
83 return centerline_interpolation
86def get_quaternions_along_curve(
87 centerline: _interpolate.BSpline, point_arc_length: _np.ndarray
88) -> _NDArray[_quaternion.quaternion]:
89 """Get the quaternions along the curve based on smallest rotation mappings.
91 The initial rotation will be calculated based on the largest projection of the initial tangent
92 onto the cartesian basis vectors.
94 Args
95 ----
96 centerline:
97 A function that returns the centerline position for a parameter coordinate t
98 point_arc_length:
99 Array of parameter coordinates for which the quaternions should be calculated
100 """
101 centerline_interpolation_derivative = centerline.derivative()
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
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))
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
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 n_points = len(coordinates)
136 relative_distances = _np.zeros(n_points - 1)
137 relative_distances_rotation = _np.zeros(n_points - 1, dtype=_quaternion.quaternion)
138 relative_rotations = _np.zeros(n_points - 1, dtype=_quaternion.quaternion)
140 for i_segment in range(n_points - 1):
141 relative_distance = coordinates[i_segment + 1] - coordinates[i_segment]
142 relative_distance_local = _quaternion.rotate_vectors(
143 quaternions[i_segment].conjugate(), relative_distance
144 )
145 relative_distances[i_segment] = _np.linalg.norm(relative_distance_local)
147 smallest_relative_rotation_onto_distance = _smallest_rotation(
148 _Rotation(),
149 relative_distance_local,
150 )
151 relative_distances_rotation[i_segment] = (
152 smallest_relative_rotation_onto_distance.get_numpy_quaternion()
153 )
155 relative_rotations[i_segment] = (
156 quaternions[i_segment].conjugate() * quaternions[i_segment + 1]
157 )
159 return relative_distances, relative_distances_rotation, relative_rotations
162class CosseratCurve(object):
163 """Represent a Cosserat curve in space."""
165 def __init__(
166 self,
167 point_coordinates: _np.ndarray,
168 *,
169 starting_triad_guess: _Rotation | None = None,
170 ):
171 """Initialize the Cosserat curve based on points in 3D space.
173 Args:
174 point_coordinates: Array containing the point coordinates
175 starting_triad_guess: Optional initial guess for the starting triad.
176 If provided, this introduces a constant twist angle along the curve.
177 The twist angle is computed between:
178 - The given starting guess triad, and
179 - The automatically calculated triad, rotated onto the first basis vector
180 of the starting guess triad using the smallest rotation.
181 """
182 self.coordinates = point_coordinates.copy()
183 self.n_points = len(self.coordinates)
185 # Interpolate coordinates along piece wise linear arc length
186 point_arc_length_piecewise_linear = (
187 get_piecewise_linear_arc_length_along_points(self.coordinates)
188 )
189 centerline_interpolation_piecewise_linear = get_spline_interpolation(
190 self.coordinates, point_arc_length_piecewise_linear
191 )
192 centerline_interpolation_piecewise_linear_p = (
193 centerline_interpolation_piecewise_linear.derivative(1)
194 )
196 def ds(t):
197 """Arc length along interpolated spline."""
198 return _np.linalg.norm(centerline_interpolation_piecewise_linear_p(t))
200 # Integrate the arc length along the interpolated centerline, this will result
201 # in a more accurate centerline arc length
202 self.point_arc_length = _np.zeros(self.n_points)
203 for i in range(len(point_arc_length_piecewise_linear) - 1):
204 self.point_arc_length[i + 1] = (
205 self.point_arc_length[i]
206 + _integrate.quad(
207 ds,
208 point_arc_length_piecewise_linear[i],
209 point_arc_length_piecewise_linear[i + 1],
210 )[0]
211 )
213 # Set the interpolation of the (positional) centerline
214 self.set_centerline_interpolation()
216 # Get the quaternions along the centerline based on smallest rotation mappings
217 self.quaternions = get_quaternions_along_curve(
218 self.centerline_interpolation, self.point_arc_length
219 )
221 # Get the relative quantities used to warp the curve
222 (
223 self.relative_distances,
224 self.relative_distances_rotation,
225 self.relative_rotations,
226 ) = get_relative_distance_and_rotations(self.coordinates, self.quaternions)
228 # Check if we have to apply a twist for the rotations
229 if starting_triad_guess is not None:
230 first_rotation = _Rotation.from_quaternion(self.quaternions[0])
231 starting_triad_e1 = starting_triad_guess * [1, 0, 0]
232 if _np.dot(first_rotation * [1, 0, 0], starting_triad_e1) < 0.5:
233 raise ValueError(
234 "The angle between the first basis vectors of the guess triad you"
235 " provided and the automatically calculated one is too large,"
236 " please check your input data."
237 )
238 smallest_rotation_to_guess_tangent = _smallest_rotation(
239 first_rotation, starting_triad_e1
240 )
241 relative_rotation = (
242 smallest_rotation_to_guess_tangent.inv() * starting_triad_guess
243 )
244 psi = relative_rotation.get_rotation_vector()
245 if _np.linalg.norm(psi[1:]) > _bme.eps_quaternion:
246 raise ValueError(
247 "The twist angle can not be extracted as the relative rotation is not plane!"
248 )
249 twist_angle = psi[0]
250 self.twist(twist_angle)
252 def set_centerline_interpolation(self):
253 """Set the interpolation of the centerline based on the coordinates and arc
254 length stored in this object."""
255 self.centerline_interpolation = get_spline_interpolation(
256 self.coordinates, self.point_arc_length
257 )
259 def translate(self, vector):
260 """Translate the curve by the given vector."""
261 self.coordinates += vector
262 self.set_centerline_interpolation()
264 def rotate(self, rotation: _Rotation, *, origin=None):
265 """Rotate the curve and the quaternions."""
266 self.quaternions = rotation.get_numpy_quaternion() * self.quaternions
267 self.coordinates = _rotate_coordinates(
268 self.coordinates, rotation, origin=origin
269 )
270 self.set_centerline_interpolation()
272 def twist(self, twist_angle: float) -> None:
273 """Apply a constant twist rotation along the Cosserat curve.
275 Args:
276 twist_angle: The rotation angle (in radiants).
277 """
278 material_twist_rotation = _Rotation(
279 [1, 0, 0], twist_angle
280 ).get_numpy_quaternion()
282 self.quaternions = self.quaternions * material_twist_rotation
283 self.relative_distances_rotation = (
284 material_twist_rotation.conjugate()
285 * self.relative_distances_rotation
286 * material_twist_rotation
287 )
288 self.relative_rotations = (
289 material_twist_rotation.conjugate()
290 * self.relative_rotations
291 * material_twist_rotation
292 )
294 def get_centerline_position_and_rotation(
295 self, arc_length: float, **kwargs
296 ) -> tuple[_np.ndarray, _NDArray[_quaternion.quaternion]]:
297 """Return the position and rotation at a given centerline arc length."""
298 pos, rot = self.get_centerline_positions_and_rotations([arc_length], **kwargs)
299 return pos[0], rot[0]
301 def get_centerline_positions_and_rotations(
302 self, points_on_arc_length, *, factor=1.0
303 ) -> tuple[_np.ndarray, _NDArray[_quaternion.quaternion]]:
304 """Return the position and rotation at given centerline arc lengths.
306 If the points are outside of the valid interval, a linear extrapolation will be
307 performed for the displacements and the rotations will be held constant.
309 This function also allows to scale the curvature along the curve, allowing for a
310 "natural" unwrapping of general curves in 3D. We achieve this by scaling the
311 "final" curvature along the beam and then evaluating the curve that follows this
312 curvature (this would actually require to solve an ODE, but we avoid this by
313 using a piecewise constant approximation).
315 Args
316 ----
317 points_on_arc_length: list(float)
318 A sorted list with the arc lengths along the curve centerline
319 factor: float
320 Factor to scale the curvature along the curve.
321 factor == 1
322 Use the default positions and the triads obtained via a smallest rotation mapping
323 0 <factor < 1
324 Integrate (piecewise constant as evaluated with get_relative_distance_and_rotations)
325 the scaled curvature of the curve to obtain a intuitive wrapping. (factor=0 gives
326 a straight line)
327 """
328 # Get the points that are within the arc length of the given curve.
329 points_on_arc_length = _np.asarray(points_on_arc_length)
330 points_in_bounds = _np.logical_and(
331 points_on_arc_length > self.point_arc_length[0],
332 points_on_arc_length < self.point_arc_length[-1],
333 )
334 index_in_bound = _np.where(points_in_bounds == True)[0]
335 index_out_of_bound = _np.where(points_in_bounds == False)[0]
336 points_on_arc_length_in_bound = [
337 self.point_arc_length[0],
338 *points_on_arc_length[index_in_bound],
339 self.point_arc_length[-1],
340 ]
342 if factor < (1.0 - _bme.eps_quaternion):
343 coordinates = _np.zeros_like(self.coordinates)
344 quaternions = _np.zeros_like(self.quaternions)
345 coordinates[0] = self.coordinates[0]
346 quaternions[0] = self.quaternions[0]
347 for i_segment in range(self.n_points - 1):
348 relative_distance_rotation = _quaternion.slerp_evaluate(
349 _quaternion.quaternion(1),
350 self.relative_distances_rotation[i_segment],
351 factor,
352 )
353 # In the initial configuration (factor=0) we get a straight curve, so we need
354 # to use the arc length here. In the final configuration (factor=1) we want to
355 # exactly recover the input points, so we need the piecewise linear distance.
356 # Between them, we interpolate.
357 relative_distance = (factor * self.relative_distances[i_segment]) + (
358 1.0 - factor
359 ) * (
360 self.point_arc_length[i_segment + 1]
361 - self.point_arc_length[i_segment]
362 )
363 coordinates[i_segment + 1] = (
364 _quaternion.rotate_vectors(
365 quaternions[i_segment] * relative_distance_rotation,
366 [relative_distance, 0, 0],
367 )
368 + coordinates[i_segment]
369 )
370 quaternions[i_segment + 1] = quaternions[
371 i_segment
372 ] * _quaternion.slerp_evaluate(
373 _quaternion.quaternion(1),
374 self.relative_rotations[i_segment],
375 factor,
376 )
377 arc_length_spline_interpolation = get_spline_interpolation(
378 coordinates, self.point_arc_length
379 )
380 else:
381 coordinates = self.coordinates
382 quaternions = self.quaternions
383 arc_length_spline_interpolation = self.centerline_interpolation
385 sol_r = _np.zeros([len(points_on_arc_length_in_bound), 3])
386 sol_q = _np.zeros(
387 len(points_on_arc_length_in_bound), dtype=_quaternion.quaternion
388 )
389 for i_point, centerline_arc_length in enumerate(points_on_arc_length_in_bound):
390 if (
391 centerline_arc_length >= self.point_arc_length[0]
392 and centerline_arc_length <= self.point_arc_length[-1]
393 ):
394 for i in range(1, self.n_points):
395 centerline_index = i - 1
396 if self.point_arc_length[i] > centerline_arc_length:
397 break
399 # Get the two rotation vectors and arc length values
400 arc_lengths = self.point_arc_length[
401 centerline_index : centerline_index + 2
402 ]
403 q1 = quaternions[centerline_index]
404 q2 = quaternions[centerline_index + 1]
406 # Linear interpolate the arc length
407 xi = (centerline_arc_length - arc_lengths[0]) / (
408 arc_lengths[1] - arc_lengths[0]
409 )
411 # Perform a spline interpolation for the positions and a slerp
412 # interpolation for the rotations
413 sol_r[i_point] = arc_length_spline_interpolation(centerline_arc_length)
414 sol_q[i_point] = _quaternion.slerp_evaluate(q1, q2, xi)
415 else:
416 raise ValueError("Centerline value out of bounds")
418 # Set the already computed results in the final data structures
419 sol_r_final = _np.zeros([len(points_on_arc_length), 3])
420 sol_q_final = _np.zeros(len(points_on_arc_length), dtype=_quaternion.quaternion)
421 if len(index_in_bound) > 0:
422 sol_r_final[index_in_bound] = sol_r[index_in_bound - index_in_bound[0] + 1]
423 sol_q_final[index_in_bound] = sol_q[index_in_bound - index_in_bound[0] + 1]
425 # Perform the extrapolation at both ends of the curve
426 for i in index_out_of_bound:
427 arc_length = points_on_arc_length[i]
428 if arc_length <= self.point_arc_length[0]:
429 index = 0
430 elif arc_length >= self.point_arc_length[-1]:
431 index = -1
432 else:
433 raise ValueError("Should not happen")
435 length = arc_length - self.point_arc_length[index]
436 r = sol_r[index]
437 q = sol_q[index]
438 sol_r_final[i] = r + _Rotation.from_quaternion(q) * [length, 0, 0]
439 sol_q_final[i] = q
441 return sol_r_final, sol_q_final
443 def project_point(self, p, t0=None) -> float:
444 """Project a point to the curve, return the parameter coordinate for the
445 projection point."""
446 centerline_interpolation_p = self.centerline_interpolation.derivative(1)
447 centerline_interpolation_pp = self.centerline_interpolation.derivative(2)
449 def f(t):
450 """Function to find the root of."""
451 r = self.centerline_interpolation(t)
452 rp = centerline_interpolation_p(t)
453 return _np.dot(r - p, rp)
455 def fp(t):
456 """Derivative of the Function to find the root of."""
457 r = self.centerline_interpolation(t)
458 rp = centerline_interpolation_p(t)
459 rpp = centerline_interpolation_pp(t)
460 return _np.dot(rp, rp) + _np.dot(r - p, rpp)
462 if t0 is None:
463 t0 = 0.0
465 return _optimize.newton(f, t0, fprime=fp)
467 def get_pyvista_polyline(self, *, factor: float = 1.0) -> _pv.PolyData:
468 """Create a pyvista representation of the curve with the evaluated triad basis
469 vectors.
471 Args:
472 factor: Factor to scale the curvature along the curve (see
473 `get_centerline_positions_and_rotations` for details).
475 Returns:
476 A pyvista PolyData object representing the curve.
477 """
478 positions, rotations = self.get_centerline_positions_and_rotations(
479 self.point_arc_length, factor=factor
480 )
482 poly_line = _pv.PolyData()
483 poly_line.points = positions
484 cell = _np.arange(0, self.n_points, dtype=int)
485 cell = _np.insert(cell, 0, self.n_points)
486 poly_line.lines = cell
488 rotation_matrices = _quaternion.as_rotation_matrix(rotations)
489 for i_dir in range(3):
490 poly_line.point_data.set_array(
491 rotation_matrices[:, :, i_dir], f"base_vector_{i_dir + 1}"
492 )
494 return poly_line
496 def write_vtk(self, path) -> None:
497 """Save a vtk representation of the curve."""
498 self.get_pyvista_polyline().save(path)
500 def write_pvd_series(
501 self,
502 pvd_path: _Path | str,
503 *,
504 factors: list[float] | None = None,
505 n_steps: int | None = None,
506 binary: bool = True,
507 ) -> None:
508 """Save a pvd series representing the curve at different states.
510 Args:
511 pvd_path: Path where to save the pvd file.
512 factors: List of factors to scale the curvature along the curve. Mutually exclusive with 'n_steps'.
513 n_steps: Number of steps to create a uniform series of factors. Mutually exclusive with 'factors'.
514 binary: If True, save the vtk files in binary format.
515 """
516 pvd_path = _Path(pvd_path)
517 if pvd_path.suffix != ".pvd":
518 raise ValueError(
519 f"The output path must have a .pvd suffix, got {pvd_path.suffix}"
520 )
522 if factors is not None and n_steps is not None:
523 raise ValueError(
524 "The keyword arguments 'factors' and 'n_steps' are mutually exclusive."
525 )
526 if factors is None and n_steps is None:
527 raise ValueError(
528 "One of the keyword arguments 'factors' or 'n_steps' must be provided."
529 )
530 if factors is None:
531 factors = _np.linspace(0.0, 1.0, num=n_steps)
533 pvd_file = _ET.Element("VTKFile", type="Collection", version="0.1")
534 collection = _ET.SubElement(pvd_file, "Collection")
535 width = max(1, len(str(len(factors) - 1)))
536 for i_step, factor in enumerate(factors):
537 # TODO: Check if we can use vtp here instead of vtu. Currently this does
538 # not work with how we compare files in testing. Since vtu and vtp are
539 # basically the same in this case, this solution is fine at the moment.
540 factor_file = pvd_path.parent / f"{pvd_path.stem}.{i_step:0{width}d}.vtu"
541 _pv.UnstructuredGrid(self.get_pyvista_polyline(factor=factor)).save(
542 factor_file, binary=binary
543 )
544 _ET.SubElement(
545 collection,
546 "DataSet",
547 timestep=str(factor),
548 group="",
549 part="0",
550 file=str(factor_file.relative_to(pvd_path.parent)),
551 )
553 tree = _ET.ElementTree(pvd_file)
554 _ET.indent(tree, space=" ", level=0)
555 tree.write(pvd_path, encoding="utf-8", xml_declaration=True)