diff --git a/docs/api/constraints.rst b/docs/api/constraints.rst index 40a4b2df..32308e94 100644 --- a/docs/api/constraints.rst +++ b/docs/api/constraints.rst @@ -18,6 +18,7 @@ Constraints are equivalent to displacement boundary condition. ConstraintBase FreeBC OneEndFixedBC + GeneralConstraint FixedConstraint HelicalBucklingBC FreeRod @@ -31,6 +32,7 @@ Constraint / Boundary Condition Rod Rigid Body =============================== ==== =========== FreeBC ✅ ✅ OneEndFixedBC ✅ ✅ +GeneralConstraint ✅ ✅ FixedConstraint ✅ ✅ HelicalBucklingBC ✅ ❌ =============================== ==== =========== @@ -62,6 +64,9 @@ Built-in Constraints .. autoclass:: OneEndFixedBC :special-members: __init__ +.. autoclass:: GeneralConstraint + :special-members: __init__ + .. autoclass:: FixedConstraint :special-members: __init__ diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 9c903eb1..4e11205f 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -5,6 +5,7 @@ "FreeRod", # Deprecated: remove v0.3.0 "OneEndFixedBC", "OneEndFixedRod", # Deprecated: remove v0.3.0 + "GeneralConstraint", "FixedConstraint", "HelicalBucklingBC", ] @@ -19,9 +20,11 @@ import numba from numba import njit +from elastica._linalg import _batch_matvec, _batch_matrix_transpose from elastica._rotations import _get_rotation_matrix from elastica.rod import RodBase from elastica.rigidbody import RigidBodyBase +from elastica.typing import SystemType class ConstraintBase(ABC): @@ -258,32 +261,41 @@ class OneEndFixedRod(OneEndFixedBC): ) -class FixedConstraint(ConstraintBase): +class GeneralConstraint(ConstraintBase): """ - This boundary condition class fixes the specified node or orientations. + This boundary condition class allows the specified node/link to have a configurable constraint. Index can be passed to fix either or both the position or the director. Constraining position is equivalent to setting 0 translational DOF. Constraining director is equivalent to setting 0 rotational DOF. Examples -------- - How to fix two ends of the rod: + How to fix all translational and rotational dof except allowing twisting around the z-axis in an inertial frame: >>> simulator.constrain(rod).using( - ... FixedConstraint, - ... constrained_position_idx=(0,1,-2,-1), - ... constrained_director_idx=(0,-1) + ... GeneralConstraint, + ... constrained_position_idx=(0,), + ... constrained_director_idx=(0,), + ... translational_constraint_selector=np.array([True, True, True]), + ... rotational_constraint_selector=np.array([True, True, False]), ... ) - How to pin the middle of the rod (10th node), without constraining the rotational DOF. + How to allow the end of the rod to move in the XY plane and allow all rotational dof: >>> simulator.constrain(rod).using( - ... FixedConstraint, - ... constrained_position_idx=(10) + ... GeneralConstraint, + ... constrained_position_idx=(-1,), + ... translational_constraint_selector=np.array([True, True, False]), ... ) """ - def __init__(self, *fixed_data, **kwargs): + def __init__( + self, + *fixed_data, + translational_constraint_selector: Optional[np.ndarray] = None, + rotational_constraint_selector: Optional[np.array] = None, + **kwargs, + ): """ Initialization of the constraint. Any parameter passed to 'using' will be available in kwargs. @@ -294,6 +306,12 @@ def __init__(self, *fixed_data, **kwargs): Tuple of position-indices that will be constrained constrained_director_idx : tuple Tuple of director-indices that will be constrained + translational_constraint_selector: Optional[np.ndarray] + np.array of type bool indicating which translational degrees of freedom (dof) to constrain. + If entry is True, the corresponding dof will be constrained. If None, we constrain all dofs. + rotational_constraint_selector: Optional[np.ndarray] + np.array of type bool indicating which translational degrees of freedom (dof) to constrain. + If entry is True, the corresponding dof will be constrained. """ super().__init__(**kwargs) pos, dir = [], [] @@ -317,6 +335,204 @@ def __init__(self, *fixed_data, **kwargs): # transpose from (blocksize, dim, dim) to (dim, dim, blocksize) self.fixed_directors = np.array(dir).transpose((1, 2, 0)) + if translational_constraint_selector is None: + translational_constraint_selector = np.array([True, True, True]) + if rotational_constraint_selector is None: + rotational_constraint_selector = np.array([True, True, True]) + # properly validate the user-provided constraint selectors + assert ( + type(translational_constraint_selector) == np.ndarray + and translational_constraint_selector.dtype == bool + and translational_constraint_selector.shape == (3,) + ), "Translational constraint selector must be a 1D boolean array of length 3." + assert ( + type(rotational_constraint_selector) == np.ndarray + and rotational_constraint_selector.dtype == bool + and rotational_constraint_selector.shape == (3,) + ), "Rotational constraint selector must be a 1D boolean array of length 3." + # cast booleans to int + self.translational_constraint_selector = ( + translational_constraint_selector.astype(int) + ) + self.rotational_constraint_selector = rotational_constraint_selector.astype(int) + + def constrain_values(self, system: SystemType, time: float) -> None: + if self.constrained_position_idx.size: + self.nb_constrain_translational_values( + system.position_collection, + self.fixed_positions, + self.constrained_position_idx, + self.translational_constraint_selector, + ) + + def constrain_rates(self, system: SystemType, time: float) -> None: + if self.constrained_position_idx.size: + self.nb_constrain_translational_rates( + system.velocity_collection, + self.constrained_position_idx, + self.translational_constraint_selector, + ) + if self.constrained_director_idx.size: + self.nb_constrain_rotational_rates( + system.director_collection, + system.omega_collection, + self.constrained_director_idx, + self.rotational_constraint_selector, + ) + + @staticmethod + @njit(cache=True) + def nb_constrain_translational_values( + position_collection, fixed_position_collection, indices, constraint_selector + ) -> None: + """ + Computes constrain values in numba njit decorator + + Parameters + ---------- + position_collection : numpy.ndarray + 2D (dim, blocksize) array containing data with `float` type. + fixed_position_collection : numpy.ndarray + 2D (dim, blocksize) array containing data with `float` type. + indices : numpy.ndarray + 1D array containing the index of constraining nodes + constraint_selector: numpy.ndarray + 1D array of type int and size (3,) indicating which translational Degrees of Freedom (DoF) to constrain. + Entries are integers in {0, 1} (e.g. a binary values of either 0 or 1). + If entry is 1, the concerning DoF will be constrained, otherwise it will be free for translation. + Selector shall be specified in the inertial frame + """ + block_size = indices.size + for i in range(block_size): + k = indices[i] + # First term: add the old position values using the inverse constraint selector (e.g. DoF) + # Second term: add the fixed position values using the constraint selector (e.g. constraint dimensions) + position_collection[..., k] = ( + 1 - constraint_selector + ) * position_collection[ + ..., k + ] + constraint_selector * fixed_position_collection[ + ..., i + ] + + @staticmethod + @njit(cache=True) + def nb_constrain_translational_rates( + velocity_collection, indices, constraint_selector + ) -> None: + """ + Compute constrain rates in numba njit decorator + + Parameters + ---------- + velocity_collection : numpy.ndarray + 2D (dim, blocksize) array containing data with `float` type. + indices : numpy.ndarray + 1D array containing the index of constraining nodes + constraint_selector: numpy.ndarray + 1D array of type int and size (3,) indicating which translational Degrees of Freedom (DoF) to constrain. + Entries are integers in {0, 1} (e.g. a binary values of either 0 or 1). + If entry is 1, the concerning DoF will be constrained, otherwise it will be free for translation. + Selector shall be specified in the inertial frame + """ + + block_size = indices.size + for i in range(block_size): + k = indices[i] + # set the dofs to 0 where the constraint_selector mask is active + velocity_collection[..., k] = ( + 1 - constraint_selector + ) * velocity_collection[..., k] + + @staticmethod + @njit(cache=True) + def nb_constrain_rotational_rates( + director_collection, omega_collection, indices, constraint_selector + ) -> None: + """ + Compute constrain rates in numba njit decorator + + Parameters + ---------- + director_collection : numpy.ndarray + 2D (dim, blocksize) array containing data with `float` type. + omega_collection : numpy.ndarray + 2D (dim, blocksize) array containing data with `float` type. + indices : numpy.ndarray + 1D array containing the index of constraining nodes + constraint_selector: numpy.ndarray + 1D array of type int and size (3,) indicating which rotational Degrees of Freedom (DoF) to constrain. + Entries are integers in {0, 1} (e.g. a binary values of either 0 or 1). + If an entry is 1, the rotation around the respective axis will be constrained, + otherwise the system can freely rotate around the axis. + The selector shall be specified in the lab frame + """ + directors = director_collection[..., indices] + + # rotate angular velocities to lab frame + omega_collection_lab_frame = _batch_matvec( + _batch_matrix_transpose(directors), omega_collection[..., indices] + ) + + # apply constraint selector to angular velocities in lab frame + omega_collection_not_constrained = ( + 1 - np.expand_dims(constraint_selector, 1) + ) * omega_collection_lab_frame + + # rotate angular velocities vector back to local frame and apply to omega_collection + omega_collection[..., indices] = _batch_matvec( + directors, omega_collection_not_constrained + ) + + +class FixedConstraint(GeneralConstraint): + """ + This boundary condition class fixes the specified node or orientations. + Index can be passed to fix either or both the position or the director. + Constraining position is equivalent to setting 0 translational DOF. + Constraining director is equivalent to setting 0 rotational DOF. + + Examples + -------- + How to fix two ends of the rod: + + >>> simulator.constrain(rod).using( + ... FixedConstraint, + ... constrained_position_idx=(0,-1), + ... constrained_director_idx=(0,-1) + ... ) + + How to pin the middle of the rod (10th node), without constraining the rotational DOF. + + >>> simulator.constrain(rod).using( + ... FixedConstraint, + ... constrained_position_idx=(10) + ... ) + + See Also + --------- + GeneralConstraint: Generalized constraint with configurable DOF. + """ + + def __init__(self, *args, **kwargs): + """ + + Initialization of the constraint. Any parameter passed to 'using' will be available in kwargs. + + Parameters + ---------- + constrained_position_idx : tuple + Tuple of position-indices that will be constrained + constrained_director_idx : tuple + Tuple of director-indices that will be constrained + """ + super().__init__( + *args, + translational_constraint_selector=np.array([True, True, True]), + rotational_constraint_selector=np.array([True, True, True]), + **kwargs, + ) + def constrain_values( self, rod: Union[Type[RodBase], Type[RigidBodyBase]], time: float ) -> None: @@ -354,7 +570,6 @@ def nb_constraint_rotational_values( ) -> None: """ Computes constrain values in numba njit decorator - Parameters ---------- director_collection : numpy.ndarray @@ -363,7 +578,6 @@ def nb_constraint_rotational_values( 3D (dim, dim, blocksize) array containing data with `float` type. indices : numpy.ndarray 1D array containing the index of constraining nodes - """ block_size = indices.size for i in range(block_size): @@ -377,7 +591,6 @@ def nb_constrain_translational_values( ) -> None: """ Computes constrain values in numba njit decorator - Parameters ---------- position_collection : numpy.ndarray @@ -386,7 +599,6 @@ def nb_constrain_translational_values( 2D (dim, blocksize) array containing data with `float` type. indices : numpy.ndarray 1D array containing the index of constraining nodes - """ block_size = indices.size for i in range(block_size): @@ -398,7 +610,6 @@ def nb_constrain_translational_values( def nb_constrain_translational_rates(velocity_collection, indices) -> None: """ Compute constrain rates in numba njit decorator - Parameters ---------- velocity_collection : numpy.ndarray @@ -419,7 +630,6 @@ def nb_constrain_translational_rates(velocity_collection, indices) -> None: def nb_constrain_rotational_rates(omega_collection, indices) -> None: """ Compute constrain rates in numba njit decorator - Parameters ---------- omega_collection : numpy.ndarray diff --git a/examples/BoundaryConditionsCases/bc_cases_postprocessing.py b/examples/BoundaryConditionsCases/bc_cases_postprocessing.py new file mode 100644 index 00000000..71531105 --- /dev/null +++ b/examples/BoundaryConditionsCases/bc_cases_postprocessing.py @@ -0,0 +1,149 @@ +import numpy as np +from matplotlib import pyplot as plt +from matplotlib.colors import to_rgb +from mpl_toolkits import mplot3d +from scipy.spatial.transform import Rotation + +from elastica.rigidbody import Cylinder +from elastica._linalg import _batch_matvec + + +def plot_position( + plot_params_rod1: dict, + filename="tip_position.png", + SAVE_FIGURE=False, +): + position_of_rod1 = np.array(plot_params_rod1["position"]) + + fig = plt.figure(figsize=(10, 10), frameon=True, dpi=150) + ax = fig.add_subplot(111) + + ax.grid(b=True, which="minor", color="k", linestyle="--") + ax.grid(b=True, which="major", color="k", linestyle="-") + ax.plot(position_of_rod1[:, 0, -1], position_of_rod1[:, 1, -1], "r-", label="rod1") + + fig.legend(prop={"size": 20}) + plt.xlabel("x") + plt.ylabel("y") + + plt.show() + + if SAVE_FIGURE: + fig.savefig(filename) + + +def plot_orientation(title, time, directors): + quat = [] + for t in range(len(time)): + quat_t = Rotation.from_matrix(directors[t].T).as_quat() + quat.append(quat_t) + quat = np.array(quat) + + plt.figure(num=title) + plt.plot(time, quat[:, 0], label="x") + plt.plot(time, quat[:, 1], label="y") + plt.plot(time, quat[:, 2], label="z") + plt.plot(time, quat[:, 3], label="w") + plt.title(title) + plt.legend() + plt.xlabel("Time [s]") + plt.ylabel("Quaternion") + plt.show() + + +def plot_video( + plot_params_rod1: dict, + video_name="video.mp4", + fps=100, +): # (time step, x/y/z, node) + import matplotlib.animation as manimation + + time = plot_params_rod1["time"] + position_of_rod1 = np.array(plot_params_rod1["position"]) + + print("plot video") + FFMpegWriter = manimation.writers["ffmpeg"] + metadata = dict(title="Movie Test", artist="Matplotlib", comment="Movie support!") + writer = FFMpegWriter(fps=fps, metadata=metadata) + fig = plt.figure(figsize=(10, 8), frameon=True, dpi=150) + with writer.saving(fig, video_name, 100): + for time in range(1, len(time)): + fig.clf() + ax = plt.axes(projection="3d") # fig.add_subplot(111) + ax.grid(b=True, which="minor", color="k", linestyle="--") + ax.grid(b=True, which="major", color="k", linestyle="-") + ax.plot( + position_of_rod1[time, 0], + position_of_rod1[time, 1], + position_of_rod1[time, 2], + "or", + label="rod1", + ) + + ax.set_xlim(-0.25, 0.25) + ax.set_ylim(-0.25, 0.25) + ax.set_zlim(0, 0.61) + ax.set_xlabel("x") + ax.set_ylabel("y") + ax.set_zlabel("z") + writer.grab_frame() + + +def plot_video_xy( + plot_params_rod1: dict, + video_name="video_xy.mp4", + fps=100, +): # (time step, x/y/z, node) + import matplotlib.animation as manimation + + time = plot_params_rod1["time"] + position_of_rod1 = np.array(plot_params_rod1["position"]) + + print("plot video xy") + FFMpegWriter = manimation.writers["ffmpeg"] + metadata = dict(title="Movie Test", artist="Matplotlib", comment="Movie support!") + writer = FFMpegWriter(fps=fps, metadata=metadata) + fig = plt.figure() + plt.axis("equal") + with writer.saving(fig, video_name, 100): + for time in range(1, len(time)): + fig.clf() + plt.plot( + position_of_rod1[time, 0], position_of_rod1[time, 1], "or", label="rod1" + ) + + plt.xlim([-0.25, 0.25]) + plt.ylim([-0.25, 0.25]) + plt.xlabel("x") + plt.ylabel("y") + writer.grab_frame() + + +def plot_video_xz( + plot_params_rod1: dict, + video_name="video_xz.mp4", + fps=100, +): # (time step, x/y/z, node) + import matplotlib.animation as manimation + + time = plot_params_rod1["time"] + position_of_rod1 = np.array(plot_params_rod1["position"]) + + print("plot video xz") + FFMpegWriter = manimation.writers["ffmpeg"] + metadata = dict(title="Movie Test", artist="Matplotlib", comment="Movie support!") + writer = FFMpegWriter(fps=fps, metadata=metadata) + fig = plt.figure() + plt.axis("equal") + with writer.saving(fig, video_name, 100): + for time in range(1, len(time)): + fig.clf() + plt.plot( + position_of_rod1[time, 0], position_of_rod1[time, 2], "or", label="rod1" + ) + + plt.xlim([-0.25, 0.25]) + plt.ylim([0, 0.61]) + plt.xlabel("x") + plt.ylabel("z") + writer.grab_frame() diff --git a/examples/BoundaryConditionsCases/general_constraint_allow_yaw.py b/examples/BoundaryConditionsCases/general_constraint_allow_yaw.py new file mode 100644 index 00000000..50088323 --- /dev/null +++ b/examples/BoundaryConditionsCases/general_constraint_allow_yaw.py @@ -0,0 +1,140 @@ +__doc__ = """Fixed joint example, for detailed explanation refer to Zhang et. al. Nature Comm. methods section.""" + +import matplotlib.pyplot as plt +import numpy as np +import sys + +# FIXME without appending sys.path make it more generic +sys.path.append("../../") +from elastica import * +from examples.BoundaryConditionsCases.bc_cases_postprocessing import ( + plot_position, + plot_orientation, + plot_video, + plot_video_xy, + plot_video_xz, +) + + +class GeneralConstraintSimulator( + BaseSystemCollection, Constraints, Connections, Forcing, Damping, CallBacks +): + pass + + +general_constraint_sim = GeneralConstraintSimulator() + +# setting up test params +n_elem = 10 +direction = np.array([0.0, 0.0, 1.0]) +normal = np.array([0.0, 1.0, 0.0]) +base_length = 0.2 +base_radius = 0.007 +base_area = np.pi * base_radius ** 2 +density = 1750 +E = 3e7 +poisson_ratio = 0.5 +shear_modulus = E / (poisson_ratio + 1.0) + +# setting up timestepper and video +final_time = 10 +dl = base_length / n_elem +dt = 1e-5 +total_steps = int(final_time / dt) +fps = 100 # frames per second of the video +diagnostic_step_skip = 1 / (fps * dt) + +start_rod_1 = np.zeros((3,)) +start_rod_2 = start_rod_1 + direction * base_length + +# Create rod 1 +rod1 = CosseratRod.straight_rod( + n_elem, + start_rod_1, + direction, + normal, + base_length, + base_radius, + density, + 0.0, # internal damping constant, deprecated in v0.3.0 + E, + shear_modulus=shear_modulus, +) +general_constraint_sim.append(rod1) + +# Apply boundary conditions to rod1: only allow yaw +general_constraint_sim.constrain(rod1).using( + GeneralConstraint, + constrained_position_idx=(0,), + constrained_director_idx=(0,), + translational_constraint_selector=np.array([True, True, True]), + rotational_constraint_selector=np.array([True, True, False]), +) + +# add forces to endpoint +general_constraint_sim.add_forcing_to(rod1).using( + EndpointForces, + start_force=np.zeros((3,)), + end_force=1e-0 * np.ones((3,)), + ramp_up_time=0.5, +) +# add uniform torsion torque +general_constraint_sim.add_forcing_to(rod1).using( + UniformTorques, torque=1e-3, direction=np.array([0, 0, 1]) +) + +# add damping +damping_constant = 0.4 +general_constraint_sim.dampen(rod1).using( + ExponentialDamper, + damping_constant=damping_constant, + time_step=dt, +) + + +pp_list_rod1 = defaultdict(list) + + +general_constraint_sim.collect_diagnostics(rod1).using( + MyCallBack, step_skip=diagnostic_step_skip, callback_params=pp_list_rod1 +) + +general_constraint_sim.finalize() +timestepper = PositionVerlet() + +print("Total steps", total_steps) +integrate(timestepper, general_constraint_sim, final_time, total_steps) + + +plot_orientation( + "Orientation of first element of rod 1", + pp_list_rod1["time"], + np.array(pp_list_rod1["directors"])[..., 0], +) + +PLOT_FIGURE = True +SAVE_FIGURE = True +PLOT_VIDEO = True + +# plotting results +if PLOT_FIGURE: + filename = "configurable_fixed_constraint_example_last_node_pos_xy.png" + plot_position(pp_list_rod1, filename, SAVE_FIGURE) + +if PLOT_VIDEO: + filename = "configurable_fixed_constraint_example" + plot_video( + pp_list_rod1, + video_name=filename + ".mp4", + fps=fps, + ) + plot_video_xy( + pp_list_rod1, + video_name=filename + "_xy.mp4", + fps=fps, + ) + plot_video_xz( + pp_list_rod1, + video_name=filename + "_xz.mp4", + fps=fps, + ) diff --git a/examples/README.md b/examples/README.md index 5211b129..06fb8f82 100644 --- a/examples/README.md +++ b/examples/README.md @@ -66,6 +66,9 @@ Examples can serve as a starting template for customized usages. * [SolenoidsCase](./RodContactCase/RodSelfContact/SolenoidsCase) * __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 +* [BoundaryConditionsCases](./BoundaryConditionsCases) + * __Purpose__: Demonstrate the usage of boundary conditions for constraining the movement of the system. + * __Features__: GeneralConstraint, CosseratRod ## Functional Examples diff --git a/tests/test_boundary_conditions.py b/tests/test_boundary_conditions.py index af289e77..9d327256 100644 --- a/tests/test_boundary_conditions.py +++ b/tests/test_boundary_conditions.py @@ -3,23 +3,27 @@ # System imports import numpy as np -from tests.test_rod.test_rods import MockTestRod from elastica.boundary_conditions import ( ConstraintBase, FreeBC, OneEndFixedBC, FixedConstraint, + GeneralConstraint, HelicalBucklingBC, ) -from numpy.testing import assert_allclose +from elastica._linalg import _batch_matvec from elastica.utils import Tolerance +from numpy.testing import assert_allclose import pytest from pytest import main +from scipy.spatial.transform import Rotation +from tests.test_rod.test_rods import MockTestRod test_built_in_boundary_condition_impls = [ FreeBC, OneEndFixedBC, FixedConstraint, + GeneralConstraint, HelicalBucklingBC, ] @@ -167,18 +171,18 @@ def test_one_end_fixed_bc(): @pytest.mark.parametrize("n_director_constraint", [0, 2, 6, 9]) def test_fixed_constraint(seed, n_position_constraint, n_director_constraint): rng = np.random.default_rng(seed) - N = 20 - test_rod = MockTestRod() + N = test_rod.n_elem + start_position_collection = rng.random((n_position_constraint, 3)) start_director_collection = rng.random((n_director_constraint, 3, 3)) - fixed_rod = FixedConstraint( + fixed_constrained = FixedConstraint( *start_position_collection, *start_director_collection, _system=test_rod ) pos_indices = rng.choice(N, size=n_position_constraint, replace=False) dir_indices = rng.choice(N, size=n_director_constraint, replace=False) - fixed_rod._constrained_position_idx = pos_indices.copy() - fixed_rod._constrained_director_idx = dir_indices.copy() + fixed_constrained._constrained_position_idx = pos_indices.copy() + fixed_constrained._constrained_director_idx = dir_indices.copy() test_position_collection = rng.random((3, N)) test_rod.position_collection = ( @@ -188,7 +192,7 @@ def test_fixed_constraint(seed, n_position_constraint, n_director_constraint): test_rod.director_collection = ( test_director_collection.copy() ) # We need copy of the list not a reference to this array - fixed_rod.constrain_values(test_rod, time=0) + fixed_constrained.constrain_values(test_rod, time=0) test_position_collection[..., pos_indices] = start_position_collection.transpose( (1, 0) ) @@ -210,7 +214,7 @@ def test_fixed_constraint(seed, n_position_constraint, n_director_constraint): test_rod.omega_collection = ( test_omega_collection.copy() ) # We need copy of the list not a reference to this array - fixed_rod.constrain_rates(test_rod, time=0) + fixed_constrained.constrain_rates(test_rod, time=0) test_velocity_collection[..., pos_indices] = 0.0 test_omega_collection[..., dir_indices] = 0.0 assert_allclose( @@ -221,6 +225,162 @@ def test_fixed_constraint(seed, n_position_constraint, n_director_constraint): ) +@pytest.mark.parametrize("seed", [1, 10, 100]) +@pytest.mark.parametrize("num_translational_constraint", [0, 1, 2]) +@pytest.mark.parametrize("num_rotational_constraint", [0, 1, 2]) +@pytest.mark.parametrize( + "constraint_selector", + [ + np.array([False, False, False]), + np.array([True, False, False]), + np.array([False, True, False]), + np.array([False, False, True]), + np.array([True, True, False]), + np.array([True, False, True]), + np.array([False, True, True]), + np.array([True, True, True]), + ], +) +def test_general_constraint( + seed, + num_translational_constraint, + num_rotational_constraint, + constraint_selector, +): + rng = np.random.default_rng(seed) + test_rod = MockTestRod() + + start_position_collection = rng.random((num_translational_constraint, 3)) + start_director_collection = rng.random((num_rotational_constraint, 3, 3)) + translational_constraint_selector = constraint_selector + rotational_constraint_selector = constraint_selector + general_constraint = GeneralConstraint( + *start_position_collection, + *start_director_collection, + translational_constraint_selector=translational_constraint_selector, + rotational_constraint_selector=rotational_constraint_selector, + _system=test_rod, + ) + pos_indices = rng.choice( + test_rod.n_elem, size=num_translational_constraint, replace=False + ) + dir_indices = rng.choice( + test_rod.n_elem, size=num_rotational_constraint, replace=False + ) + general_constraint._constrained_position_idx = pos_indices.copy() + general_constraint._constrained_director_idx = dir_indices.copy() + + test_position_collection = rng.random((3, test_rod.n_elem)) + test_rod.position_collection = ( + test_position_collection.copy() + ) # We need copy of the list not a reference to this array + general_constraint.constrain_values(test_rod, time=0) + test_position_collection[..., pos_indices] = start_position_collection.transpose( + (1, 0) + ) + + assert_allclose( + test_position_collection[translational_constraint_selector, :], + test_rod.position_collection[translational_constraint_selector, :], + atol=Tolerance.atol(), + ) + + test_velocity_collection = rng.random((3, test_rod.n_elem)) + test_rod.velocity_collection = ( + test_velocity_collection.copy() + ) # We need copy of the list not a reference to this array + test_director_collection = rng.random((3, 3, test_rod.n_elem)) + test_rod.director_collection = ( + test_director_collection.copy() + ) # We need copy of the list not a reference to this array + test_omega_collection = rng.random((3, test_rod.n_elem)) + test_rod.omega_collection = ( + test_omega_collection.copy() + ) # We need copy of the list not a reference to this array + general_constraint.constrain_rates(test_rod, time=0) + + # test `nb_constrain_translational_rates` + if translational_constraint_selector[0]: + test_velocity_collection[0, pos_indices] = 0.0 + if translational_constraint_selector[1]: + test_velocity_collection[1, pos_indices] = 0.0 + if translational_constraint_selector[2]: + test_velocity_collection[2, pos_indices] = 0.0 + assert_allclose( + test_velocity_collection, test_rod.velocity_collection, atol=Tolerance.atol() + ) + + # test `nb_constrain_rotational_rates` for directors not equal to identity matrix + # rotate angular velocities to inertial frame + omega_collection_lab_frame = _batch_matvec( + test_director_collection[ + ..., + ].transpose(1, 0, 2), + test_omega_collection, + ) + # apply constraint selector to angular velocities in inertial frame + omega_collection_not_constrained = omega_collection_lab_frame.copy() + if rotational_constraint_selector[0]: + omega_collection_not_constrained[0, dir_indices] = 0.0 + if rotational_constraint_selector[1]: + omega_collection_not_constrained[1, dir_indices] = 0.0 + if rotational_constraint_selector[2]: + omega_collection_not_constrained[2, dir_indices] = 0.0 + # rotate angular velocities vector back to local frame and apply to omega_collection + test_omega_collection[..., dir_indices] = _batch_matvec( + test_director_collection, omega_collection_not_constrained + )[..., dir_indices] + assert_allclose( + test_omega_collection, test_rod.omega_collection, atol=Tolerance.atol() + ) + + # test `nb_constrain_rotational_rates` for directors equal to identity matrix + test_director_collection = ( + np.eye(3).reshape(3, 3, 1).repeat(test_rod.n_elem, axis=2) + ) + test_rod.director_collection = ( + test_director_collection.copy() + ) # We need copy of the list not a reference to this array + test_omega_collection = rng.random((3, test_rod.n_elem)) + test_rod.omega_collection = ( + test_omega_collection.copy() + ) # We need copy of the list not a reference to this array + general_constraint.constrain_rates(test_rod, time=0) + if rotational_constraint_selector[0]: + test_omega_collection[0, dir_indices] = 0.0 + if rotational_constraint_selector[1]: + test_omega_collection[1, dir_indices] = 0.0 + if rotational_constraint_selector[2]: + test_omega_collection[2, dir_indices] = 0.0 + assert_allclose( + test_omega_collection, test_rod.omega_collection, atol=Tolerance.atol() + ) + + # test `nb_constrain_rotational_rates` for directors equal to 90 degrees rotation around z-axis + rot_mat_90deg_yaw = Rotation.from_euler("z", np.pi / 2).as_matrix() + test_director_collection = rot_mat_90deg_yaw.reshape((3, 3, 1)).repeat( + test_rod.n_elem, axis=2 + ) + test_rod.director_collection = ( + test_director_collection.copy() + ) # We need copy of the list not a reference to this array + test_omega_collection = rng.random((3, test_rod.n_elem)) + test_rod.omega_collection = ( + test_omega_collection.copy() + ) # We need copy of the list not a reference to this array + general_constraint.constrain_rates(test_rod, time=0) + # because of the 90 degree rotation around z-axis, the x and y selectors are switched + if rotational_constraint_selector[1]: + test_omega_collection[0, dir_indices] = 0.0 + if rotational_constraint_selector[0]: + test_omega_collection[1, dir_indices] = 0.0 + if rotational_constraint_selector[2]: + test_omega_collection[2, dir_indices] = 0.0 + assert_allclose( + test_omega_collection, test_rod.omega_collection, atol=Tolerance.atol() + ) + + def test_helical_buckling_bc(): twisting_time = 500.0