diff --git a/.coveragerc b/.coveragerc index ed02c8e7..52568494 100644 --- a/.coveragerc +++ b/.coveragerc @@ -4,6 +4,12 @@ exclude_lines = # Don't complain if non-runnable code isn't run: if 0: if __name__ == .__main__.: + ... + pass + def __repr__ + from + import + [run] branch = True diff --git a/.flake8 b/.flake8 index 34b2930b..7f925ebc 100644 --- a/.flake8 +++ b/.flake8 @@ -3,4 +3,8 @@ ignore = E203, E266, E501, W503, F403, F401, W191 max-line-length = 88 max-complexity = 18 select = B,C,E,F,W,T4,B9 -exclude = .git,__pycache__,elastica/rod.py \ No newline at end of file +exclude = + .git, + __pycache__, + docs/conf.py + tests diff --git a/Makefile b/Makefile index b45ffe6c..a1620765 100644 --- a/Makefile +++ b/Makefile @@ -12,6 +12,8 @@ flake8: @flake8 --version @flake8 elastica tests +test: + @python -m pytest all:black flake8 ci:black_check flake8 diff --git a/docs/api/rods.rst b/docs/api/rods.rst index c3248118..1b620efd 100644 --- a/docs/api/rods.rst +++ b/docs/api/rods.rst @@ -1,12 +1,12 @@ Rods -===== +==== .. automodule:: elastica.rod.rod_base :members: :exclude-members: __weakref__ Cosserat Rod -~~~~~~~~~~~~ +------------ +------------+-------------------+----------------------------------------+-----------------------------+ | | On Nodes (+1) | On Elements (n_elements) | On Voronoi (-1) | @@ -39,9 +39,22 @@ Cosserat Rod .. automodule:: elastica.rod.cosserat_rod :exclude-members: __weakref__, __init__, update_accelerations, zeroed_out_external_forces_and_torques, compute_internal_forces_and_torques :members: + :inherited-members: .. Constitutive Models .. ~~~~~~~~~~~~~~~~~~~ .. .. automodule:: elastica.rod.constitutive_model .. :members: .. :exclude-members: __weakref__ + + +Knot Theory (Mixin) +~~~~~~~~~~~~~~~~~~~ + +.. .. autoclass:: elastica.rod.knot_theory.KnotTheory + +.. .. autoclass:: elastica.rod.knot_theory.KnotTheoryCompatibleProtocol + +.. automodule:: elastica.rod.knot_theory + :exclude-members: __init__ + :members: diff --git a/docs/conf.py b/docs/conf.py index e6662bcf..964ad0a0 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -90,5 +90,8 @@ html_static_path = ['_static'] html_css_files = ['css/*', 'css/logo.css'] +# -- Options for autodoc --------------------------------------------------- +autodoc_member_order = 'bysource' + # -- Options for numpydoc --------------------------------------------------- numpydoc_show_class_members = False diff --git a/elastica/__init__.py b/elastica/__init__.py index b9a976d1..98bb897f 100644 --- a/elastica/__init__.py +++ b/elastica/__init__.py @@ -4,6 +4,7 @@ from collections import defaultdict from elastica.wrappers import * from elastica.rod.cosserat_rod import * +from elastica.rod.knot_theory import * from elastica.rigidbody import * from elastica.boundary_conditions import * from elastica.external_forces import * diff --git a/elastica/rod/__init__.py b/elastica/rod/__init__.py index 8d846bf6..08263853 100644 --- a/elastica/rod/__init__.py +++ b/elastica/rod/__init__.py @@ -1,5 +1,6 @@ __doc__ = """Rod classes and its data structures """ +from elastica.rod.knot_theory import * from elastica.rod.data_structures import * from elastica.rod.rod_base import RodBase diff --git a/elastica/rod/cosserat_rod.py b/elastica/rod/cosserat_rod.py index 4b635950..fa0eb12b 100644 --- a/elastica/rod/cosserat_rod.py +++ b/elastica/rod/cosserat_rod.py @@ -14,6 +14,7 @@ ) from elastica._rotations import _inv_rotate from elastica.rod.factory_function import allocate +from elastica.rod.knot_theory import KnotTheory from elastica._calculus import ( quadrature_kernel_for_block_structure, difference_kernel_for_block_structure, @@ -31,7 +32,7 @@ def _get_z_vector(): return np.array([0.0, 0.0, 1.0]).reshape(3, -1) -class CosseratRod(RodBase): +class CosseratRod(RodBase, KnotTheory): """ Cosserat Rod class. This is the preferred class for rods because it is derived from some of the essential base classes. diff --git a/elastica/rod/knot_theory.py b/elastica/rod/knot_theory.py new file mode 100644 index 00000000..f31d2abe --- /dev/null +++ b/elastica/rod/knot_theory.py @@ -0,0 +1,754 @@ +__doc__ = """ +This script is for computing the link-writhe-twist (LWT) of a rod using the method from Klenin & Langowski 2000 paper. +Algorithms are adapted from section S2 of Charles et. al. PRL 2019 paper. + +Following example cases includes computing LWT quantities to study the bifurcation: + +- `Example case (PlectonemesCase) `_ +- `Example case (SolenoidCase) `_ + +The details discussion is included in `N Charles et. al. PRL (2019) `_. +""" +__all__ = [ + "KnotTheoryCompatibleProtocol", + "KnotTheory", + "compute_twist", + "compute_link", + "compute_writhe", +] + +import sys + +if sys.version_info.minor >= 8: + # typing Protocol is introduced in python 3.8 + from typing import Protocol +elif sys.version_info.minor < 8: + # Protocol is implemented in typing_extensions for previous pythons + from typing_extensions import Protocol + +from typing import Union + +from numba import njit +import numpy as np + +from elastica.rod.rod_base import RodBase +from elastica._linalg import _batch_norm, _batch_dot, _batch_cross + + +class KnotTheoryCompatibleProtocol(Protocol): + """KnotTheoryCompatibleProtocol + + Required properties to use KnotTheory mixin + """ + + @property + def position_collection(self) -> np.ndarray: + ... + + @property + def director_collection(self) -> np.ndarray: + ... + + @property + def radius(self) -> np.ndarray: + ... + + @property + def base_length(self) -> np.ndarray: + ... + + +class KnotTheory: + """ + This mixin should be used in RodBase-derived class that satisfies KnotCompatibleProtocol. + The theory behind this module is based on the method from Klenin & Langowski 2000 paper. + + KnotTheory can be mixed with any rod-class based on RodBase:: + + class MyRod(RodBase, KnotTheory): + def __init__(self): + super().__init__() + rod = MyRod(...) + + total_twist = rod.compute_twist() + total_link = rod.compute_link() + + There are few alternative way of handling edge-condition in computing Link and Writhe. + Here, we provide three methods: "next_tangent", "end_to_end", and "net_tangent". + The default *type_of_additional_segment* is set to "next_tangent." + + ========================== ===================================== + type_of_additional_segment Description + ========================== ===================================== + next_tangent | Adds a two new point at the begining and end of the center line. + | Distance of these points are given in segment_length. + | Direction of these points are computed using the rod tangents at + | the begining and end. + end_to_end | Adds a two new point at the begining and end of the center line. + | Distance of these points are given in segment_length. + | Direction of these points are computed using the rod node end + | positions. + net_tangent | Adds a two new point at the begining and end of the center line. + | Distance of these points are given in segment_length. Direction of + | these points are point wise avarege of nodes at the first and + | second half of the rod. + ========================== ===================================== + + """ + + MIXIN_PROTOCOL = Union[RodBase, KnotTheoryCompatibleProtocol] + + def compute_twist(self: MIXIN_PROTOCOL): + """ + See :ref:`api/rods:Knot Theory (Mixin)` for the detail. + """ + total_twist, local_twist = compute_twist( + self.position_collection[None, ...], + self.director_collection[0][None, ...], + ) + return total_twist[0] + + def compute_writhe( + self: MIXIN_PROTOCOL, type_of_additional_segment: str = "next_tangent" + ): + """ + See :ref:`api/rods:Knot Theory (Mixin)` for the detail. + + Parameters + ---------- + type_of_additional_segment : str + Determines the method to compute new segments (elements) added to the rod. + Valid inputs are "next_tangent", "end_to_end", "net_tangent", otherwise program uses the center line. + """ + return compute_writhe( + self.position_collection[None, ...], + self.rest_lengths.sum(), + type_of_additional_segment, + )[0] + + def compute_link( + self: MIXIN_PROTOCOL, type_of_additional_segment: str = "next_tangent" + ): + """ + See :ref:`api/rods:Knot Theory (Mixin)` for the detail. + + Parameters + ---------- + type_of_additional_segment : str + Determines the method to compute new segments (elements) added to the rod. + Valid inputs are "next_tangent", "end_to_end", "net_tangent", otherwise program uses the center line. + """ + print(self.rest_lengths.sum()) + return compute_link( + self.position_collection[None, ...], + self.director_collection[0][None, ...], + self.radius[None, ...], + self.rest_lengths.sum(), + type_of_additional_segment, + )[0] + + +def compute_twist(center_line, normal_collection): + """ + Compute the twist of a rod, using center_line and normal collection. + + Methods used in this function is adapted from method 2a Klenin & Langowski 2000 paper. + + .. warning:: If center line is straight, although the normals of each element is pointing different direction computed twist will be zero. + + Typical runtime of this function is longer than simulation steps. While we provide a function to compute + topological quantities at every timesteps, **we highly recommend** to compute LWT during the post-processing + stage.:: + + import elastica + ... + normal_collection = director_collection[:,0,...] # shape of director (time, 3, 3, n_elems) + elastica.compute_twist( + center_line, # shape (time, 3, n_nodes) + normal_collection # shape (time, 3, n_elems) + ) + + Parameters + ---------- + center_line : numpy.ndarray + 3D (time, 3, n_nodes) array containing data with 'float' type. + Time history of rod node positions. + normal_collection : numpy.ndarray + 3D (time, 3, n_elems) array containing data with 'float' type. + Time history of rod elements normal direction. + + Returns + ------- + total_twist : numpy.ndarray + local_twist : numpy.ndarray + + """ + # fmt: off + # Format is turned off because I want the assertion message to display the line. + assert center_line.shape[2] == normal_collection.shape[2] + 1, \ + "Please check the shape (axis-2) of center_line(n_node=n_elems+1) or normal_collection(n_elems)." + assert center_line.shape[0] == normal_collection.shape[0], \ + "The number of timesteps (axis-0) must be equal" + assert center_line.shape[1] == normal_collection.shape[1] == 3, \ + "The dimension (axis-1) must be 3" + # fmt: on + + total_twist, local_twist = _compute_twist(center_line, normal_collection) + + return total_twist, local_twist + + +@njit(cache=True) +def _compute_twist(center_line, normal_collection): + """ + Parameters + ---------- + center_line : numpy.ndarray + 3D (time, 3, n_nodes) array containing data with 'float' type. + Time history of rod node positions. + normal_collection : numpy.ndarray + 3D (time, 3, n_elems) array containing data with 'float' type. + Time history of rod elements normal direction. + + Returns + ------- + total_twist : numpy.ndarray + local_twist : numpy.ndarray + """ + + timesize, _, blocksize = center_line.shape + + total_twist = np.zeros((timesize)) + local_twist = np.zeros((timesize, blocksize - 2)) # Only consider interior nodes + + # compute s vector + for k in range(timesize): + # s is a secondary vector field. + s = center_line[k, :, 1:] - center_line[k, :, :-1] + # Compute tangents + tangent = s / _batch_norm(s) + + # Compute the projection of normal collection (d1) on normal-binormal plane. + projection_of_normal_collection = ( + normal_collection[k, :, :] + - _batch_dot(tangent, normal_collection[k, :, :]) * tangent + ) + projection_of_normal_collection /= _batch_norm(projection_of_normal_collection) + + # Eq 27 in Klenin & Langowski 2000 + # p is defined on interior nodes + p = _batch_cross(s[:, :-1], s[:, 1:]) + p /= _batch_norm(p) + + # Compute the angle we need to turn d1 around s to get p + # sign part tells whether d1 must be rotated ccw(+) or cw(-) around s + alpha = np.sign( + _batch_dot( + _batch_cross(projection_of_normal_collection[:, :-1], p), s[:, :-1] + ) + ) * np.arccos(_batch_dot(projection_of_normal_collection[:, :-1], p)) + + gamma = np.sign( + _batch_dot( + _batch_cross(p, projection_of_normal_collection[:, 1:]), s[:, 1:] + ) + ) * np.arccos(_batch_dot(projection_of_normal_collection[:, 1:], p)) + + # An angle 1 is a full rotation, 0.5 is rotation by pi, 0.25 is pi/2 etc. + alpha /= 2 * np.pi + gamma /= 2 * np.pi + twist_temp = alpha + gamma + # Make sure twist is between (-1/2 to 1/2) as defined in pg 313 Klenin & Langowski 2000 + idx = np.where(twist_temp > 0.5)[0] + twist_temp[idx] -= 1 + idx = np.where(twist_temp < -0.5)[0] + twist_temp[idx] += 1 + + # Check if there is any Nan. Nan's appear when rod tangents are parallel to each other. + idx = np.where(np.isnan(twist_temp))[0] + twist_temp[idx] = 0.0 + + local_twist[k, :] = twist_temp + total_twist[k] = np.sum(twist_temp) + + return total_twist, local_twist + + +def compute_writhe(center_line, segment_length, type_of_additional_segment): + """ + This function computes the total writhe history of a rod. + + Equations used are from method 1a from Klenin & Langowski 2000 paper. + + Typical runtime of this function is longer than simulation steps. While we provide a function to compute + topological quantities at every timesteps, **we highly recommend** to compute LWT during the post-processing + stage.:: + + import elastica + ... + elastica.compute_writhe( + center_line, # shape (time, 3, n_nodes) + segment_length, + type_of_additional_segment="next_tangent" + ) + + Parameters + ---------- + center_line : numpy.ndarray + 3D (time, 3, n_nodes) array containing data with 'float' type. + Time history of rod node positions. + segment_length : float + Length of added segments. + type_of_additional_segment : str + Determines the method to compute new segments (elements) added to the rod. + Valid inputs are "next_tangent", "end_to_end", "net_tangent", otherwise program uses the center line. + + Returns + ------- + total_writhe : numpy.ndarray + + """ + # fmt: off + # Format is turned off because I want the assertion message to display the line. + assert center_line.shape[1] == 3, \ + "The dimension (axis-1) must be 3" + # fmt: on + + center_line_with_added_segments, _, _ = _compute_additional_segment( + center_line, segment_length, type_of_additional_segment + ) + + total_writhe = _compute_writhe(center_line_with_added_segments) + + return total_writhe + + +@njit(cache=True) +def _compute_writhe(center_line): + """ + Parameters + ---------- + center_line : numpy.ndarray + 3D (time, 3, n_nodes) array containing data with 'float' type. + Time history of rod node positions. + + Returns + ------- + total_writhe : numpy.ndarray + """ + + time, _, blocksize = center_line.shape + + omega_star = np.zeros((blocksize - 2, blocksize - 1)) + segment_writhe = np.zeros((blocksize - 2, blocksize - 1)) + total_writhe = np.zeros((time)) + + # Compute the writhe between each pair first. + for k in range(time): + for i in range(blocksize - 2): + for j in range(i + 1, blocksize - 1): + + point_one = center_line[k, :, i] + point_two = center_line[k, :, i + 1] + point_three = center_line[k, :, j] + point_four = center_line[k, :, j + 1] + + # Eq 15 in Klenin & Langowski 2000 + r12 = point_two - point_one + r34 = point_four - point_three + r14 = point_four - point_one + r13 = point_three - point_one + r23 = point_three - point_two + r24 = point_four - point_two + + n1 = np.cross(r13, r14) + n1 /= np.linalg.norm(n1) + n2 = np.cross(r14, r24) + n2 /= np.linalg.norm(n2) + n3 = np.cross(r24, r23) + n3 /= np.linalg.norm(n3) + n4 = np.cross(r23, r13) + n4 /= np.linalg.norm(n4) + + # Eq 16a in Klenin & Langowski 2000 + omega_star[i, j] = ( + np.arcsin(np.dot(n1, n2)) + + np.arcsin(np.dot(n2, n3)) + + np.arcsin(np.dot(n3, n4)) + + np.arcsin(np.dot(n4, n1)) + ) + + if np.isnan(omega_star[i, j]): + omega_star[i, j] = 0 + + # Eq 16b in Klenin & Langowski 2000 + segment_writhe[i, j] = ( + omega_star[i, j] + * np.sign(np.dot(np.cross(r34, r12), r13)) + / (4 * np.pi) + ) + + # Compute the total writhe + # Eq 13 in Klenin & Langowski 2000 + total_writhe[k] = 2 * np.sum(segment_writhe) + + return total_writhe + + +def compute_link( + center_line: np.ndarray, + normal_collection: np.ndarray, + radius: np.ndarray, + segment_length: float, + type_of_additional_segment: str, +): + """ + This function computes the total link history of a rod. + + Equations used are from method 1a from Klenin & Langowski 2000 paper. + + Typical runtime of this function is longer than simulation steps. While we provide a function to compute + topological quantities at every timesteps, **we highly recommend** to compute LWT during the post-processing + stage.:: + + import elastica + ... + normal_collection = director_collection[:,0,...] # shape of director (time, 3, 3, n_elems) + elastica.compute_link( + center_line, # shape (time, 3, n_nodes) + normal_collection, # shape (time 3, n_elems) + radius, # shape (time, n_elems) + segment_length, + type_of_additional_segment="next_tangent" + ) + + Parameters + ---------- + center_line : numpy.ndarray + 3D (time, 3, n_nodes) array containing data with 'float' type. + Time history of rod node positions. + normal_collection : numpy.ndarray + 3D (time, 3, n_elems) array containing data with 'float' type. + Time history of rod elements normal direction. + radius : numpy.ndarray + 2D (time, n_elems) array containing data with 'float' type. + Time history of rod element radius. + segment_length : float + Length of added segments. + type_of_additional_segment : str + Determines the method to compute new segments (elements) added to the rod. + Valid inputs are "next_tangent", "end_to_end", "net_tangent", otherwise program uses the center line. + + Returns + ------- + total_link : numpy.ndarray + + """ + # fmt: off + # Format is turned off because I want the assertion message to display the line. + assert center_line.shape[2] == normal_collection.shape[2] + 1 == radius.shape[1] + 1, \ + "Please check the shape (axis-2) of center_line(n_node=n_elems+1) or normal_collection(n_elems)." + assert center_line.shape[0] == normal_collection.shape[0] == radius.shape[0], \ + "The number of timesteps (axis-0) must be equal" + assert center_line.shape[1] == normal_collection.shape[1] == 3, \ + "The dimension (axis-1) for center_line and normal_collection must be 3" + # fmt: on + + # Compute auxiliary line + auxiliary_line = _compute_auxiliary_line(center_line, normal_collection, radius) + + # Add segments at the beginning and end of the rod center line and auxiliary line. + ( + center_line_with_added_segments, + beginning_direction, + end_direction, + ) = _compute_additional_segment( + center_line, segment_length, type_of_additional_segment + ) + auxiliary_line_with_added_segments = _compute_auxiliary_line_added_segments( + beginning_direction, end_direction, auxiliary_line, segment_length + ) + + """ + Total link of a rod is computed using the method 1a from Klenin & Langowski 2000 + """ + total_link = _compute_link( + center_line_with_added_segments, auxiliary_line_with_added_segments + ) + + return total_link + + +@njit(cache=True) +def _compute_auxiliary_line(center_line, normal_collection, radius): + """ + This function computes the auxiliary line using rod center line and normal collection. + + Parameters + ---------- + center_line : numpy.ndarray + 3D (time, 3, n_nodes) array containing data with 'float' type. + Time history of rod node positions. + normal_collection : numpy.ndarray + 3D (time, 3, n_elems) array containing data with 'float' type. + Time history of rod elements normal direction. + radius : numpy.ndarray + 2D (time, n_elems) array containing data with 'float' type. + Time history of rod element radius. + + Returns + ------- + auxillary_line : numpy.ndarray + + """ + time, _, blocksize = center_line.shape + auxiliary_line = np.zeros(center_line.shape) + projection_of_normal_collection = np.zeros((3, blocksize)) + radius_on_nodes = np.zeros((blocksize)) + + for i in range(time): + tangent = center_line[i, :, 1:] - center_line[i, :, :-1] + tangent /= _batch_norm(tangent) + # Compute the projection of normal collection (d1) on normal-binormal plane. + projection_of_normal_collection_temp = ( + normal_collection[i, :, :] + - _batch_dot(tangent, normal_collection[i, :, :]) * tangent + ) + projection_of_normal_collection_temp /= _batch_norm( + projection_of_normal_collection_temp + ) + + # First node have the same direction with second node. They share the same element. + # TODO: Instead of this maybe we should use the trapezoidal rule or averaging operator for normal and radius. + projection_of_normal_collection[:, 0] = projection_of_normal_collection_temp[ + :, 0 + ] + projection_of_normal_collection[:, 1:] = projection_of_normal_collection_temp[:] + radius_on_nodes[0] = radius[i, 0] + radius_on_nodes[1:] = radius[i, :] + + auxiliary_line[i, :, :] = ( + radius_on_nodes * projection_of_normal_collection + center_line[i, :, :] + ) + + return auxiliary_line + + +@njit(cache=True) +def _compute_link(center_line, auxiliary_line): + """ + + Parameters + ---------- + center_line : numpy.ndarray + 3D (time, 3, n_nodes) array containing data with 'float' type. + Time history of rod node positions. + auxiliary_line : numpy.ndarray + 3D (time, 3, n_nodes) array containing data with 'float' type. + Time history of auxiliary line. + + Returns + ------- + total_link : numpy.ndarray + + """ + timesize, _, blocksize_center_line = center_line.shape + blocksize_auxiliary_line = auxiliary_line.shape[-1] + + omega_star = np.zeros((blocksize_center_line - 1, blocksize_auxiliary_line - 1)) + segment_link = np.zeros((blocksize_center_line - 1, blocksize_auxiliary_line - 1)) + total_link = np.zeros((timesize)) + + # Compute the writhe between each pair first. + for k in range(timesize): + for i in range(blocksize_center_line - 1): + for j in range(blocksize_auxiliary_line - 1): + + point_one = center_line[k, :, i] + point_two = center_line[k, :, i + 1] + point_three = auxiliary_line[k, :, j] + point_four = auxiliary_line[k, :, j + 1] + + # Eq 15 in Klenin & Langowski 2000 + r12 = point_two - point_one + r34 = point_four - point_three + r14 = point_four - point_one + r13 = point_three - point_one + r23 = point_three - point_two + r24 = point_four - point_two + + n1 = np.cross(r13, r14) + n1 /= np.linalg.norm(n1) + n2 = np.cross(r14, r24) + n2 /= np.linalg.norm(n2) + n3 = np.cross(r24, r23) + n3 /= np.linalg.norm(n3) + n4 = np.cross(r23, r13) + n4 /= np.linalg.norm(n4) + + # Eq 16a in Klenin & Langowski 2000 + omega_star[i, j] = ( + np.arcsin(np.dot(n1, n2)) + + np.arcsin(np.dot(n2, n3)) + + np.arcsin(np.dot(n3, n4)) + + np.arcsin(np.dot(n4, n1)) + ) + + if np.isnan(omega_star[i, j]): + omega_star[i, j] = 0 + + # Eq 16b in Klenin & Langowski 2000 + segment_link[i, j] = ( + omega_star[i, j] + * np.sign(np.dot(np.cross(r34, r12), r13)) + / (4 * np.pi) + ) + + # Compute the total writhe + # Eq 6 in Klenin & Langowski 2000 + # Unlike the writhe, link computed using two curves so we do not multiply by 2 + total_link[k] = np.sum(segment_link) + + return total_link + + +@njit(cache=True) +def _compute_auxiliary_line_added_segments( + beginning_direction, end_direction, auxiliary_line, segment_length +): + """ + This code is for computing position of added segments to the auxiliary line. + + Parameters + ---------- + beginning_direction : numpy.ndarray + 2D (time, 3) array containing data with 'float' type. + Time history of center line tangent at the beginning. + end_direction : numpy.ndarray + 2D (time, 3) array containing data with 'float' type. + Time history of center line tangent at the end. + auxiliary_line : numpy.ndarray + 3D (time, 3, n_nodes) array containing data with 'float' type. + Time history of auxiliary line. + segment_length : float + Length of added segments. + + Returns + ------- + new_auxiliary_line : numpy.ndarray + + """ + timesize, _, blocksize = auxiliary_line.shape + + new_auxiliary_line = np.zeros((timesize, 3, blocksize + 2)) + + new_auxiliary_line[:, :, 1:-1] = auxiliary_line + + new_auxiliary_line[:, :, 0] = ( + auxiliary_line[:, :, 0] + beginning_direction * segment_length + ) + + new_auxiliary_line[:, :, -1] = ( + auxiliary_line[:, :, -1] + end_direction * segment_length + ) + + return new_auxiliary_line + + +@njit(cache=True) +def _compute_additional_segment( + center_line, segment_length, type_of_additional_segment +): + """ + This function adds two points at the end of center line. Distance from the center line is given by segment_length. + Direction from center line to the new point locations can be computed using 3 methods, which can be selected by + type. For more details section S2 of Charles et. al. PRL 2019 paper. + + next_tangent: + This function adds a two new point at the begining and end of the center line. Distance of these points are + given in segment_length. Direction of these points are computed using the rod tangents at the begining and end. + end_to_end: + This function adds a two new point at the begining and end of the center line. Distance of these points are + given in segment_length. Direction of these points are computed using the rod node end positions. + net_tangent: + This function adds a two new point at the begining and end of the center line. Distance of these points are + given in segment_length. Direction of these points are point wise avarege of nodes at the first and second half + of the rod. + + Parameters + ---------- + center_line : numpy.ndarray + 3D (time, 3, n_nodes) array containing data with 'float' type. + Time history of rod node positions. + segment_length : float + Length of added segments. + type_of_additional_segment : str + Determines the method to compute new segments (elements) added to the rod. + Valid inputs are "next_tangent", "end_to_end", "net_tangent". If None, returns the center line. + + Returns + ------- + center_line : numpy.ndarray + beginning_direction : numpy.ndarray + end_direction : numpy.ndarray + + """ + if type_of_additional_segment is None: + beginning_direction = np.zeros((center_line.shape[0], 3)) + end_direction = np.zeros((center_line.shape[0], 3)) + return center_line, beginning_direction, end_direction + + timesize, _, blocksize = center_line.shape + new_center_line = np.zeros( + (timesize, 3, blocksize + 2) + ) # +2 is for added two new points + beginning_direction = np.zeros((timesize, 3)) + end_direction = np.zeros((timesize, 3)) + + if type_of_additional_segment == "next_tangent": + for i in range(timesize): + # Direction of the additional point at the beginning of the rod + direction_of_rod_begin = center_line[i, :, 0] - center_line[i, :, 1] + direction_of_rod_begin /= np.linalg.norm(direction_of_rod_begin) + + # Direction of the additional point at the end of the rod + direction_of_rod_end = center_line[i, :, -1] - center_line[i, :, -2] + direction_of_rod_end /= np.linalg.norm(direction_of_rod_end) + elif type_of_additional_segment == "end_to_end": + for i in range(timesize): + # Direction of the additional point at the beginning of the rod + direction_of_rod_begin = center_line[i, :, 0] - center_line[i, :, -1] + direction_of_rod_begin /= np.linalg.norm(direction_of_rod_begin) + + # Direction of the additional point at the end of the rod + direction_of_rod_end = -direction_of_rod_begin + elif type_of_additional_segment == "net_tangent": + for i in range(timesize): + # Direction of the additional point at the beginning of the rod + n_nodes_begin = int(np.floor(blocksize / 2)) + average_begin = ( + np.sum(center_line[i, :, :n_nodes_begin], axis=1) / n_nodes_begin + ) + n_nodes_end = int(np.ceil(blocksize / 2)) + average_end = np.sum(center_line[i, :, n_nodes_end:], axis=1) / ( + blocksize - n_nodes_end + 1 + ) + direction_of_rod_begin = average_begin - average_end + direction_of_rod_begin /= np.linalg.norm(direction_of_rod_begin) + direction_of_rod_end = -direction_of_rod_begin + else: + raise NotImplementedError("unavailable type_of_additional_segment is given") + + # Compute new centerline and beginning/end direction + for i in range(timesize): + first_point = center_line[i, :, 0] + segment_length * direction_of_rod_begin + last_point = center_line[i, :, -1] + segment_length * direction_of_rod_end + + new_center_line[i, :, 1:-1] = center_line[i, :, :] + new_center_line[i, :, 0] = first_point + new_center_line[i, :, -1] = last_point + + beginning_direction[i, :] = direction_of_rod_begin + end_direction[i, :] = direction_of_rod_end + + return new_center_line, beginning_direction, end_direction diff --git a/examples/README.md b/examples/README.md index 0c3b67e4..e79241e4 100644 --- a/examples/README.md +++ b/examples/README.md @@ -48,11 +48,11 @@ Examples can serve as a starting template for customized usages. * __Features__: CosseratRod, ExternalContact * [RodSelfContact](./RodContactCase/RodSelfContact) * [PlectonemesCase](./RodContactCase/RodSelfContact/PlectonemesCase) - * __Purpose__: Demonstrates rod self contact with Plectoneme example. - * __Features__: CosseratRod, SelonoidsBC, SelfContact + * __Purpose__: Demonstrates rod self contact with Plectoneme example, and how to use link-writhe-twist after simulation completed. + * __Features__: CosseratRod, SelonoidsBC, SelfContact, Link-Writhe-Twist * [SolenoidsCase](./RodContactCase/RodSelfContact/SolenoidsCase) - * __Purpose__: Demonstrates rod self contact with Solenoid example. - * __Features__: CosseratRod, SelonoidsBC, SelfContact + * __Purpose__: Demonstrates rod self contact with Solenoid example, and how to use link-writhe-twist after simulation completed. + * __Features__: CosseratRod, SelonoidsBC, SelfContact, Link-Writhe-Twist ## Functional Examples diff --git a/examples/RodContactCase/RodSelfContact/PlectonemesCase/plectoneme_case.py b/examples/RodContactCase/RodSelfContact/PlectonemesCase/plectoneme_case.py index 54a13175..1f7e817b 100644 --- a/examples/RodContactCase/RodSelfContact/PlectonemesCase/plectoneme_case.py +++ b/examples/RodContactCase/RodSelfContact/PlectonemesCase/plectoneme_case.py @@ -5,6 +5,7 @@ from examples.RodContactCase.post_processing import ( plot_video_with_surface, plot_velocity, + plot_link_writhe_twist, ) @@ -45,7 +46,9 @@ class PlectonemesCase( direction = np.array([0.0, 1.0, 0]) normal = np.array([0.0, 0.0, 1.0]) -start = np.zeros(3,) +start = np.zeros( + 3, +) sherable_rod = CosseratRod.straight_rod( n_elem, @@ -70,9 +73,7 @@ class PlectonemesCase( class SelonoidsBC(ConstraintBase): - """ - - """ + """ """ def __init__( self, @@ -160,9 +161,7 @@ def constrain_rates(self, rod, time): # Add callback functions for plotting position of the rod later on class RodCallBack(CallBackBaseClass): - """ - - """ + """ """ def __init__(self, step_skip: int, callback_params: dict): CallBackBaseClass.__init__(self) @@ -187,6 +186,7 @@ def make_callback(self, system, time, current_step: int): + system.compute_shear_energy() ) self.callback_params["total_energy"].append(total_energy) + self.callback_params["directors"].append(system.director_collection.copy()) return @@ -194,7 +194,9 @@ def make_callback(self, system, time, current_step: int): post_processing_dict = defaultdict(list) # list which collected data will be append # set the diagnostics for rod and collect data plectonemes_sim.collect_diagnostics(sherable_rod).using( - RodCallBack, step_skip=step_skip, callback_params=post_processing_dict, + RodCallBack, + step_skip=step_skip, + callback_params=post_processing_dict, ) # finalize simulation @@ -217,3 +219,68 @@ def make_callback(self, system, time, current_step: int): y_limits=[-0.1, 1.1], z_limits=[-0.5, 0.5], ) + +# Compute topological quantities +time = np.array(post_processing_dict["time"]) +position_history = np.array(post_processing_dict["position"]) +radius_history = np.array(post_processing_dict["radius"]) +director_history = np.array(post_processing_dict["directors"]) + +# Compute twist density +theta = 2.0 * number_of_rotations * np.pi +angel_vel_scalar = theta / time_twist + +twist_time_interval_start_idx = np.where(time > time_start_twist)[0][0] +twist_time_interval_end_idx = np.where(time < (time_relax + time_twist))[0][-1] + +twist_density = ( + (time[twist_time_interval_start_idx:twist_time_interval_end_idx] - time_start_twist) + * angel_vel_scalar + * base_radius +) + +# Compute link-writhe-twist +normal_history = director_history[:, 0, :, :] +segment_length = 10 * base_length + +type_of_additional_segment = "next_tangent" + +total_twist, local_twist = compute_twist(position_history, normal_history) + +total_link = compute_link( + position_history, + normal_history, + radius_history, + segment_length, + type_of_additional_segment, +) + +total_writhe = compute_writhe( + position_history, segment_length, type_of_additional_segment +) + +# Plot link-writhe-twist +plot_link_writhe_twist( + twist_density, + total_twist[twist_time_interval_start_idx:twist_time_interval_end_idx], + total_writhe[twist_time_interval_start_idx:twist_time_interval_end_idx], + total_link[twist_time_interval_start_idx:twist_time_interval_end_idx], +) + +# Save simulation data +import os + +save_folder = os.path.join(os.getcwd(), "data") +os.makedirs(save_folder, exist_ok=True) +np.savez( + os.path.join(save_folder, "plectoneme_case_data.npz"), + time=time, + position_history=position_history, + radius_history=radius_history, + director_history=director_history, + base_length=np.array(base_length), + twist_density=twist_density, + total_twist=total_twist, + total_writhe=total_writhe, + total_link=total_link, +) diff --git a/examples/RodContactCase/RodSelfContact/SolenoidsCase/solenoid_case.py b/examples/RodContactCase/RodSelfContact/SolenoidsCase/solenoid_case.py index b65b8cd1..0ce2fdce 100644 --- a/examples/RodContactCase/RodSelfContact/SolenoidsCase/solenoid_case.py +++ b/examples/RodContactCase/RodSelfContact/SolenoidsCase/solenoid_case.py @@ -5,6 +5,7 @@ from examples.RodContactCase.post_processing import ( plot_video_with_surface, plot_velocity, + plot_link_writhe_twist, ) @@ -44,7 +45,9 @@ class SolenoidCase(BaseSystemCollection, Constraints, Connections, Forcing, Call direction = np.array([0.0, 1.0, 0]) normal = np.array([0.0, 0.0, 1.0]) -start = np.zeros(3,) +start = np.zeros( + 3, +) F_pulling_scalar = 300 @@ -70,9 +73,7 @@ class SolenoidCase(BaseSystemCollection, Constraints, Connections, Forcing, Call class SelonoidsBC(ConstraintBase): - """ - - """ + """ """ def __init__( self, @@ -157,7 +158,9 @@ def constrain_rates(self, rod, time): solenoid_sim.add_forcing_to(sherable_rod).using( EndpointForces, - np.zeros(3,), + np.zeros( + 3, + ), F_pulling_scalar * direction, ramp_up_time=time_start_twist - 1, ) @@ -167,9 +170,7 @@ def constrain_rates(self, rod, time): # Add callback functions for plotting position of the rod later on class RodCallBack(CallBackBaseClass): - """ - - """ + """ """ def __init__(self, step_skip: int, callback_params: dict): CallBackBaseClass.__init__(self) @@ -202,7 +203,9 @@ def make_callback(self, system, time, current_step: int): post_processing_dict = defaultdict(list) # list which collected data will be append # set the diagnostics for rod and collect data solenoid_sim.collect_diagnostics(sherable_rod).using( - RodCallBack, step_skip=step_skip, callback_params=post_processing_dict, + RodCallBack, + step_skip=step_skip, + callback_params=post_processing_dict, ) # finalize simulation @@ -226,12 +229,7 @@ def make_callback(self, system, time, current_step: int): z_limits=[-0.5, 0.5], ) -# Save data for post-processing and computing topological quantities -import os - -save_folder = os.path.join(os.getcwd(), "data") -os.makedirs(save_folder, exist_ok=True) - +# Compute topological quantities time = np.array(post_processing_dict["time"]) position_history = np.array(post_processing_dict["position"]) radius_history = np.array(post_processing_dict["radius"]) @@ -241,17 +239,57 @@ def make_callback(self, system, time, current_step: int): theta = 2.0 * number_of_rotations * np.pi angel_vel_scalar = theta / time_twist -twist_time_interval_start_idx = np.where(time>time_start_twist)[0][0] -twist_time_interval_end_idx = np.where(time<(time_relax + time_twist))[0][-1] +twist_time_interval_start_idx = np.where(time > time_start_twist)[0][0] +twist_time_interval_end_idx = np.where(time < (time_relax + time_twist))[0][-1] + +twist_density = ( + (time[twist_time_interval_start_idx:twist_time_interval_end_idx] - time_start_twist) + * angel_vel_scalar + * base_radius +) + +# Compute link-writhe-twist +normal_history = director_history[:, 0, :, :] +segment_length = 10 * base_length -twist_density = (time[twist_time_interval_start_idx:twist_time_interval_end_idx]-time_start_twist)*angel_vel_scalar * base_radius +type_of_additional_segment = "next_tangent" +total_twist, local_twist = compute_twist(position_history, normal_history) + +total_link = compute_link( + position_history, + normal_history, + radius_history, + segment_length, + type_of_additional_segment, +) + +total_writhe = compute_writhe( + position_history, segment_length, type_of_additional_segment +) + +# Plot link-writhe-twist +plot_link_writhe_twist( + twist_density, + total_twist[twist_time_interval_start_idx:twist_time_interval_end_idx], + total_writhe[twist_time_interval_start_idx:twist_time_interval_end_idx], + total_link[twist_time_interval_start_idx:twist_time_interval_end_idx], +) + +# Save simulation data +import os + +save_folder = os.path.join(os.getcwd(), "data") +os.makedirs(save_folder, exist_ok=True) np.savez( os.path.join(save_folder, "solenoid_case_data.npz"), time=time, position_history=position_history, radius_history=radius_history, director_history=director_history, - base_length = np.array(base_length), - twist_density = twist_density, + base_length=np.array(base_length), + twist_density=twist_density, + total_twist=total_twist, + total_writhe=total_writhe, + total_link=total_link, ) diff --git a/examples/RodContactCase/post_processing.py b/examples/RodContactCase/post_processing.py index 56432907..93c5f349 100644 --- a/examples/RodContactCase/post_processing.py +++ b/examples/RodContactCase/post_processing.py @@ -92,7 +92,7 @@ def plot_video_with_surface( ) if not inst_position.shape[1] == inst_radius.shape[0]: inst_position = 0.5 * ( - inst_position[..., 1:] + inst_position[..., :-1] + inst_position[..., 1:] + inst_position[..., :-1] ) rod_scatters[rod_idx]._offsets3d = ( @@ -154,7 +154,7 @@ def plot_video_with_surface( ) if not inst_position.shape[1] == inst_radius.shape[0]: inst_position = 0.5 * ( - inst_position[..., 1:] + inst_position[..., :-1] + inst_position[..., 1:] + inst_position[..., :-1] ) rod_lines[rod_idx].set_xdata(inst_position[0]) @@ -218,7 +218,7 @@ def plot_video_with_surface( ) if not inst_position.shape[1] == inst_radius.shape[0]: inst_position = 0.5 * ( - inst_position[..., 1:] + inst_position[..., :-1] + inst_position[..., 1:] + inst_position[..., :-1] ) rod_lines[rod_idx].set_xdata(inst_position[2]) @@ -283,7 +283,7 @@ def plot_video_with_surface( ) if not inst_position.shape[1] == inst_radius.shape[0]: inst_position = 0.5 * ( - inst_position[..., 1:] + inst_position[..., :-1] + inst_position[..., 1:] + inst_position[..., :-1] ) rod_lines[rod_idx].set_xdata(inst_position[0]) @@ -333,26 +333,38 @@ def plot_velocity( axs[0].set_ylabel("x velocity", fontsize=20) axs[1].plot( - time[:], avg_velocity_rod_one[:, 1], linewidth=3, + time[:], + avg_velocity_rod_one[:, 1], + linewidth=3, ) axs[1].plot( - time[:], avg_velocity_rod_two[:, 1], linewidth=3, + time[:], + avg_velocity_rod_two[:, 1], + linewidth=3, ) axs[1].set_ylabel("y velocity", fontsize=20) axs[2].plot( - time[:], avg_velocity_rod_one[:, 2], linewidth=3, + time[:], + avg_velocity_rod_one[:, 2], + linewidth=3, ) axs[2].plot( - time[:], avg_velocity_rod_two[:, 2], linewidth=3, + time[:], + avg_velocity_rod_two[:, 2], + linewidth=3, ) axs[2].set_ylabel("z velocity", fontsize=20) axs[3].semilogy( - time[:], total_energy_rod_one[:], linewidth=3, + time[:], + total_energy_rod_one[:], + linewidth=3, ) axs[3].semilogy( - time[:], total_energy_rod_two[:], linewidth=3, + time[:], + total_energy_rod_two[:], + linewidth=3, ) axs[3].semilogy( time[:], @@ -373,3 +385,35 @@ def plot_velocity( if SAVE_FIGURE: # fig.savefig(filename) plt.savefig(filename) + + +def plot_link_writhe_twist(twist_density, total_twist, total_writhe, total_link): + + plt.rcParams.update({"font.size": 22}) + fig = plt.figure(figsize=(10, 10), frameon=True, dpi=150) + + axs = [] + axs.append(plt.subplot2grid((1, 1), (0, 0))) + axs[0].plot( + twist_density, + total_twist, + label="twist", + ) + axs[0].plot( + twist_density, + total_writhe, + label="writhe", + ) + axs[0].plot( + twist_density, + total_link, + label="link", + ) + axs[0].set_xlabel("twist density", fontsize=20) + axs[0].set_ylabel("link-twist-writhe", fontsize=20) + axs[0].set_xlim(0, 2.0) + plt.tight_layout() + fig.align_ylabels() + fig.legend(prop={"size": 20}) + fig.savefig("link_twist_writhe.png") + plt.close(plt.gcf()) diff --git a/tests/test_boundary_conditions.py b/tests/test_boundary_conditions.py index d8143d78..7e3144c5 100644 --- a/tests/test_boundary_conditions.py +++ b/tests/test_boundary_conditions.py @@ -3,7 +3,7 @@ # System imports import numpy as np -from test_rod import MockTestRod +from test_rod.test_rods import MockTestRod from elastica.boundary_conditions import ( ConstraintBase, FreeBC, diff --git a/tests/test_interaction.py b/tests/test_interaction.py index eb010344..82ac01a8 100644 --- a/tests/test_interaction.py +++ b/tests/test_interaction.py @@ -12,7 +12,7 @@ SlenderBodyTheory, ) -from test_rod import MockTestRod +from test_rod.test_rods import MockTestRod class BaseRodClass(MockTestRod): diff --git a/tests/test_cosserat_rod.py b/tests/test_rod/test_cosserat_rod.py similarity index 100% rename from tests/test_cosserat_rod.py rename to tests/test_rod/test_cosserat_rod.py diff --git a/tests/test_rod/test_knot_theory.py b/tests/test_rod/test_knot_theory.py new file mode 100644 index 00000000..fca73b94 --- /dev/null +++ b/tests/test_rod/test_knot_theory.py @@ -0,0 +1,272 @@ +__doc__ = """ Knot Theory class testing """ + +import pytest +import numpy as np +from numpy.testing import assert_allclose + +from elastica.rod.data_structures import _bootstrap_from_data +from elastica.rod.data_structures import ( + _KinematicState, + _DynamicState, +) +from elastica.utils import MaxDimension + +from test_rods import MockTestRod + +from elastica.rod.rod_base import RodBase +from elastica.rod.knot_theory import ( + KnotTheoryCompatibleProtocol, + compute_twist, + compute_writhe, + compute_link, + _compute_additional_segment, +) + + +@pytest.fixture +def knot_theory(): + from elastica.rod import knot_theory + + return knot_theory + + +def test_knot_theory_protocol(): + # To clear the protocol test coverage + with pytest.raises(TypeError) as e_info: + protocol = KnotTheoryCompatibleProtocol() + assert "cannot be instantiated" in e_info + + +def test_knot_theory_mixin_methods(knot_theory): + class TestRodWithKnotTheory(MockTestRod, knot_theory.KnotTheory): + def __init__(self): + super().__init__() + self.radius = np.random.randn(MaxDimension.value(), self.n_elem) + + rod = TestRodWithKnotTheory() + assert hasattr( + rod, "MIXIN_PROTOCOL" + ), "Expected to mix-in variables: MIXIN_PROTOCOL" + assert hasattr( + rod, "compute_writhe" + ), "Expected to mix-in functionals into the rod class: compute_writhe" + assert hasattr( + rod, "compute_twist" + ), "Expected to mix-in functionals into the rod class: compute_twist" + assert hasattr( + rod, "compute_link" + ), "Expected to mix-in functionals into the rod class: compute_link" + + +def test_knot_theory_mixin_methods_with_no_radius(knot_theory): + class TestRodWithKnotTheoryWithoutRadius(MockTestRod, knot_theory.KnotTheory): + def __init__(self): + super().__init__() + + rod = TestRodWithKnotTheoryWithoutRadius() + with pytest.raises(AttributeError) as e_info: + rod.compute_writhe() + with pytest.raises(AttributeError) as e_info: + rod.compute_link() + + +@pytest.mark.parametrize( + "position_collection, director_collection, radius, segment_length, sol_total_twist, sol_total_writhe, sol_total_link", + # fmt: off + [ + ( + np.array([[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0]], # position_collection + dtype=np.float64).T, + np.array([[1, 0, 0], [0, 1, 1], [1, 1, 0], [0,1,0]], # director_collection + dtype=np.float64).T[None,...], + np.array([1, 2, 4, 2], dtype=np.float64), # radius + np.array([10.0]), # segment_length + 0.75, # solution total twist + -0.477268070084, # solution total writhe + -0.703465518706 + ), + ], + # solution total link + # fmt: on +) +def test_knot_theory_mixin_methods_arithmetic( + knot_theory, + position_collection, + director_collection, + radius, + segment_length, + sol_total_twist, + sol_total_writhe, + sol_total_link, +): + class TestRod(RodBase, knot_theory.KnotTheory): + def __init__( + self, position_collection, director_collection, radius, segment_length + ): + self.position_collection = position_collection + self.director_collection = director_collection + self.radius = radius + self.rest_lengths = segment_length + + test_rod = TestRod(position_collection, director_collection, radius, segment_length) + + twist = test_rod.compute_twist() + writhe = test_rod.compute_writhe() + link = test_rod.compute_link() + + assert np.isclose(twist, sol_total_twist) + assert np.isclose(writhe, sol_total_writhe) + assert np.isclose(link, sol_total_link) + + +def test_compute_twist_arithmetic(): + # fmt: off + center_line = np.array( + [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0]], + dtype=np.float64).T[None, ...] + normal_collection = np.array( + [[1, 0, 0], [0, 1, 1], [1, 1, 0], [0,1,0]], + dtype=np.float64).T[None, ...] + # fmt: on + a, b = compute_twist(center_line, normal_collection) + assert np.isclose(a[0], 0.75) + assert_allclose(b[0], np.array([0.25, 0.125, 0.375])) + + +@pytest.mark.parametrize( + "type_str, sol", + [ + ("next_tangent", -0.477268070084), + ("end_to_end", -0.37304522216388), + ("net_tangent", -0.26423311709925), + ], +) +def test_compute_writhe_arithmetic(type_str, sol): + # fmt: off + center_line = np.array( + [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0]], + dtype=np.float64).T[None, ...] + segment_length = 10.0 + # fmt: on + a = compute_writhe(center_line, segment_length, type_str) + assert np.isclose(a[0], sol) + + +@pytest.mark.parametrize( + "type_str, sol", + [ + ("next_tangent", -0.703465518706), + ("end_to_end", -0.4950786438825), + ("net_tangent", -0.321184858244), + ], +) +def test_compute_link_arithmetic(type_str, sol): + # fmt: off + center_line = np.array( + [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0]], + dtype=np.float64).T[None, ...] + normal_collection = np.array( + [[1, 0, 0], [0, 1, 1], [1, 1, 0], [0,1,0]], dtype=np.float64 + ).T[None, ...] + radius = np.array([1, 2, 4, 2], dtype=np.float64)[None,...] + segment_length = 10.0 + # fmt: on + a = compute_link(center_line, normal_collection, radius, segment_length, type_str) + assert np.isclose(a[0], sol) + + +@pytest.mark.parametrize("type_str", ["randomstr1", "nextnext_tangent", " "]) +def test_knot_theory_compute_additional_segment_integrity(type_str): + center_line = np.zeros([1, 3, 10]) + center_line[:, 2, :] = np.arange(10) + with pytest.raises(NotImplementedError) as e_info: + _compute_additional_segment(center_line, 10.0, type_str) + + +@pytest.mark.parametrize("n_elem", [2, 3, 8]) +@pytest.mark.parametrize("segment_length", [1.0, 10.0, 100.0]) +@pytest.mark.parametrize("type_str", ["next_tangent", "end_to_end", "net_tangent"]) +def test_knot_theory_compute_additional_segment_straight_case( + n_elem, segment_length, type_str +): + # If straight rod give, result should be same regardless of type + center_line = np.zeros([1, 3, n_elem]) + center_line[0, 2, :] = np.linspace(0, 5, n_elem) + ncl, bd, ed = _compute_additional_segment(center_line, segment_length, type_str) + assert_allclose(ncl[0, :, 0], np.array([0, 0, -segment_length])) + assert_allclose( + ncl[0, :, -1], np.array([0, 0, center_line[0, 2, -1] + segment_length]) + ) + assert_allclose(bd[0], np.array([0, 0, -1])) + assert_allclose(ed[0], np.array([0, 0, 1])) + + +def test_knot_theory_compute_additional_segment_next_tangent_case(): + center_line = np.array( + [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]], dtype=np.float64 + ).T[None, ...] + segment_length = 10 + ncl, bd, ed = _compute_additional_segment( + center_line, segment_length, "next_tangent" + ) + # fmt: off + assert_allclose(ncl[0], + np.array([[ 0., 0., 0., 0., 0., 0.], + [ 0., 0., 0., 1., 1., 1.], + [-10., 0., 1., 1., 0., -10.]])) + assert_allclose(bd[0], np.array([ 0., 0., -1.])) + assert_allclose(ed[0], np.array([ 0., 0., -1.])) + # fmt: on + + +def test_knot_theory_compute_additional_segment_end_to_end_case(): + center_line = np.array( + [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]], dtype=np.float64 + ).T[None, ...] + segment_length = 10 + ncl, bd, ed = _compute_additional_segment(center_line, segment_length, "end_to_end") + # fmt: off + assert_allclose(ncl[0], + np.array([[ 0., 0., 0., 0., 0., 0.], + [-10., 0., 0., 1., 1., 11.], + [ 0., 0., 1., 1., 0., 0.]])) + assert_allclose(bd[0], np.array([ 0., -1., 0.])) + assert_allclose(ed[0], np.array([-0., 1., -0.])) + # fmt: on + + +def test_knot_theory_compute_additional_segment_net_tangent_case(): + center_line = np.array( + [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]], dtype=np.float64 + ).T[None, ...] + segment_length = 10 + ncl, bd, ed = _compute_additional_segment( + center_line, segment_length, "net_tangent" + ) + # fmt: off + assert_allclose(ncl[0], + np.array([[ 0. , 0. , 0. , 0. , 0. , 0. ], + [-9.701425 , 0. , 0. , 1. , 1. , 10.701425 ], + [ 2.42535625, 0. , 1. , 1. , 0. , -2.42535625]])) + assert_allclose(bd[0], np.array([ 0. , -0.9701425 , 0.24253563])) + assert_allclose(ed[0], np.array([-0. , 0.9701425 , -0.24253563])) + # fmt: on + + +@pytest.mark.parametrize("timesteps", [1, 5, 10]) +@pytest.mark.parametrize("n_elem", [1, 3, 8]) +@pytest.mark.parametrize("segment_length", [1.0, 10.0, 100.0]) +def test_knot_theory_compute_additional_segment_none_case( + timesteps, n_elem, segment_length +): + center_line = np.random.random([timesteps, 3, n_elem]) + new_center_line, beginning_direction, end_direction = _compute_additional_segment( + center_line, segment_length, None + ) + + assert_allclose(new_center_line, center_line) + assert_allclose(beginning_direction, 0.0) + assert_allclose(end_direction, 0.0) + assert_allclose(new_center_line.shape, [timesteps, 3, n_elem]) + assert_allclose(beginning_direction.shape[0], timesteps) + assert_allclose(end_direction.shape[0], timesteps) diff --git a/tests/test_rod.py b/tests/test_rod/test_rods.py similarity index 97% rename from tests/test_rod.py rename to tests/test_rod/test_rods.py index 9577a3b8..bb1a8254 100644 --- a/tests/test_rod.py +++ b/tests/test_rod/test_rods.py @@ -17,14 +17,18 @@ class MockTestRod: def __init__(self): self.n_elem = 32 - self.position_collection = np.random.randn(MaxDimension.value(), self.n_elem) + self.position_collection = np.random.randn( + MaxDimension.value(), self.n_elem + 1 + ) self.director_collection = np.random.randn( MaxDimension.value(), MaxDimension.value(), self.n_elem ) - self.velocity_collection = np.random.randn(MaxDimension.value(), self.n_elem) + self.velocity_collection = np.random.randn( + MaxDimension.value(), self.n_elem + 1 + ) self.omega_collection = np.random.randn(MaxDimension.value(), self.n_elem) - self.mass = np.abs(np.random.randn(self.n_elem)) - self.external_forces = np.zeros(self.n_elem) + self.mass = np.abs(np.random.randn(self.n_elem + 1)) + self.external_forces = np.zeros(self.n_elem + 1) # Choosing 15 and 31 as nelems to reflect common expected