Address #148 by adding ConfigurableFixedConstraint boundary condition class

elastica/boundary_conditions.py

@@ -5,6 +5,7 @@
     "FreeRod",  # Deprecated: remove v0.3.0
     "OneEndFixedBC",
     "OneEndFixedRod",  # Deprecated: remove v0.3.0
+    "ConfigurableFixedConstraint",
     "FixedConstraint",
     "HelicalBucklingBC",
 ]
@@ -19,7 +20,8 @@
 import numba
 from numba import njit
 
-from elastica._rotations import _get_rotation_matrix
+from elastica._linalg import _batch_matvec
+from elastica._rotations import _get_rotation_matrix, _inv_rotate
 from elastica.rod import RodBase
 from elastica.rigidbody import RigidBodyBase
 
@@ -258,28 +260,31 @@ class OneEndFixedRod(OneEndFixedBC):
     )
 
 
-class FixedConstraint(ConstraintBase):
+class ConfigurableFixedConstraint(ConstraintBase):
     """
-    This boundary condition class fixes the specified node or orientations.
+    This boundary condition class fixes the specified node / link with 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 z-axis in inertial frame:
 
     >>> simulator.constrain(rod).using(
-    ...    FixedConstraint,
-    ...    constrained_position_idx=(0,1,-2,-1),
-    ...    constrained_director_idx=(0,-1)
+    ...    ConfigurableFixedConstraint,
+    ...    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 x-y plane and allow all rotational DoF:
 
     >>> simulator.constrain(rod).using(
-    ...    FixedConstraint,
-    ...    constrained_position_idx=(10)
+    ...    ConfigurableFixedConstraint,
+    ...    constrained_position_idx=(-1,),
+    ...    translational_constraint_selector=np.array([True, True, False]),
     ... )
     """
 
@@ -294,6 +299,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: np.array = np.array([True, True, True])
+            np.array of type bool indicating which translational Degree of Freedom (DoF) to constrain.
+            If entry is True, the DOF will be constrained.
+        rotational_constraint_selector: np.array = np.array([True, True, True])
+            np.array of type bool indicating which translational Degree of Freedom (DoF) to constrain.
+            If entry is True, the DOF will be constrained.
         """
         super().__init__(**kwargs)
         pos, dir = [], []
@@ -311,6 +322,29 @@ def __init__(self, *fixed_data, **kwargs):
         self.fixed_positions = np.array(pos)
         self.fixed_directors = np.array(dir)
 
+        translational_constraint_selector = kwargs.get(
+            "translational_constraint_selector", np.array([True, True, True])
+        )
+        rotational_constraint_selector = kwargs.get(
+            "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, rod: Union[Type[RodBase], Type[RigidBodyBase]], time: float
     ) -> None:
@@ -319,12 +353,14 @@ def constrain_values(
                 rod.position_collection,
                 self.fixed_positions,
                 self.constrained_position_idx,
+                self.translational_constraint_selector,
             )
         if self.constrained_director_idx.size:
             self.nb_constraint_rotational_values(
                 rod.director_collection,
                 self.fixed_directors,
                 self.constrained_director_idx,
+                self.rotational_constraint_selector,
             )
 
     def constrain_rates(
@@ -334,64 +370,108 @@ def constrain_rates(
             self.nb_constrain_translational_rates(
                 rod.velocity_collection,
                 self.constrained_position_idx,
+                self.translational_constraint_selector,
             )
         if self.constrained_director_idx.size:
             self.nb_constrain_rotational_rates(
+                rod.director_collection,
                 rod.omega_collection,
                 self.constrained_director_idx,
+                self.rotational_constraint_selector,
             )
 
     @staticmethod
     @njit(cache=True)
-    def nb_constraint_rotational_values(
-        director_collection, fixed_director_collection, indices
+    def nb_constrain_translational_values(
+        position_collection, fixed_position_collection, indices, constraint_selector
     ) -> None:
         """
         Computes constrain values in numba njit decorator
 
         Parameters
         ----------
-        director_collection : numpy.ndarray
-            3D (dim, dim, blocksize) array containing data with `float` type.
-        fixed_director_collection : numpy.ndarray
-            3D (dim, dim, blocksize) array containing data with `float` type.
+        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]
-            director_collection[..., k] = fixed_director_collection[i, ...]
+            # add the old position values using the inverse constraint selector (e.g. DoF)
+            new_position_values = (1 - constraint_selector) * position_collection[
+                ..., k
+            ]
+            # add the fixed position values using the constraint selector (e.g. constraint dimensions)
+            new_position_values += (
+                constraint_selector * fixed_position_collection[i, ...]
+            )
+            position_collection[..., k] = new_position_values
 
     @staticmethod
-    @njit(cache=True)
-    def nb_constrain_translational_values(
-        position_collection, fixed_position_collection, indices
+    # @njit(cache=True)
+    def nb_constraint_rotational_values(
+        director_collection, fixed_director_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.
+        director_collection : numpy.ndarray
+            3D (dim, dim, blocksize) array containing data with `float` type.
+        fixed_director_collection : numpy.ndarray
+            3D (dim, 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.
+            Selector shall be specified in the inertial frame
         """
         block_size = indices.size
         for i in range(block_size):
             k = indices[i]
-            position_collection[0, k] = fixed_position_collection[i, 0]
-            position_collection[1, k] = fixed_position_collection[i, 1]
-            position_collection[2, k] = fixed_position_collection[i, 2]
+
+            # Rotation matrix from fixed director (e.g. saved at the first time-step) to current director
+            # C_{fixed to actual} = C_{fixed to inertial} @ C_{inertial to actual}
+            dev_rot = fixed_director_collection[i, ...] @ director_collection[..., k].T
+
+            from scipy.spatial.transform import Rotation
+
+            # XYZ Euler angles for C_{fixed to actual}
+            euler_angle = Rotation.from_matrix(dev_rot).as_euler("xyz")
+
+            # We re-set the Euler angles for constrained rotation axes to zero
+            allowed_euler_angle = (1 - constraint_selector) * euler_angle
+            # Transform allowed euler angles back to rotation matrix C_{fixed to allowed)
+            allowed_rot = Rotation.from_euler("xyz", allowed_euler_angle).as_matrix()
+
+            # Transform allowed rotation matrix to C_{inertial to allowed}
+            # This describes rotation from inertial frame to desired frame (e.g. containing allowed rotations)
+            # C_{inertial to allowed} = C_{inertial to fixed} @ C_{fixed to allowed}
+            allowed_directors = fixed_director_collection[i, ...].T @ allowed_rot
+
+            # write C_{allowed to inertial} to director_collection
+            director_collection[..., k] = allowed_directors.T
+
+            # old implementation without DoF
+            # director_collection[..., k] = fixed_director_collection[i, ...]
 
     @staticmethod
     @njit(cache=True)
-    def nb_constrain_translational_rates(velocity_collection, indices) -> None:
+    def nb_constrain_translational_rates(
+        velocity_collection, indices, constraint_selector
+    ) -> None:
         """
         Compute constrain rates in numba njit decorator
 
@@ -401,35 +481,100 @@ def nb_constrain_translational_rates(velocity_collection, indices) -> None:
             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]
-            velocity_collection[0, k] = 0.0
-            velocity_collection[1, k] = 0.0
-            velocity_collection[2, k] = 0.0
+            velocity_collection[..., k] = (
+                1 - constraint_selector
+            ) * velocity_collection[..., k]
 
     @staticmethod
     @njit(cache=True)
-    def nb_constrain_rotational_rates(omega_collection, indices) -> None:
+    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.
+            Selector shall be specified in the inertial frame
         """
+        directors = director_collection[..., indices]
 
-        block_size = indices.size
-        for i in range(block_size):
-            k = indices[i]
-            omega_collection[0, k] = 0.0
-            omega_collection[1, k] = 0.0
-            omega_collection[2, k] = 0.0
+        # rotate angular velocities to inertial frame
+        omegas_inertial = _batch_matvec(
+            directors.transpose(1, 0, 2), omega_collection[..., indices]
+        )
+
+        # apply constraint selector to angular velocities in inertial frame
+        omegas_allowed = (1 - np.expand_dims(constraint_selector, 1)) * omegas_inertial
+
+        # rotate angular velocities vector back to local frame and apply to omega_collection
+        omega_collection[..., indices] = _batch_matvec(directors, omegas_allowed)
+
+
+class FixedConstraint(ConfigurableFixedConstraint):
+    """
+    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,-2,-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)
+    ... )
+    """
+
+    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,
+        )
 
 
 class HelicalBucklingBC(ConstraintBase):