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
+    "ConfigurableConstraint",
     "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 ConfigurableConstraint(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)
+    ...    ConfigurableConstraint,
+    ...    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)
+    ...    ConfigurableConstraint,
+    ...    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 degrees of freedom (dof) to constrain.
+            If entry is True, the corresponding dof will be constrained.
+        rotational_constraint_selector: np.array = np.array([True, True, True])
+            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 +328,29 @@ 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))
 
+        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:
@@ -325,12 +359,7 @@ def constrain_values(
                 rod.position_collection,
                 self.fixed_positions,
                 self.constrained_position_idx,
-            )
-        if self.constrained_director_idx.size:
-            self.nb_constraint_rotational_values(
-                rod.director_collection,
-                self.fixed_directors,
-                self.constrained_director_idx,
+                self.translational_constraint_selector,
             )
 
     def constrain_rates(
@@ -340,40 +369,20 @@ 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
-    ) -> 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.
-        indices : numpy.ndarray
-            1D array containing the index of constraining nodes
-
-        """
-        block_size = indices.size
-        for i in range(block_size):
-            k = indices[i]
-            director_collection[..., k] = fixed_director_collection[..., i]
-
     @staticmethod
     @njit(cache=True)
     def nb_constrain_translational_values(
-        position_collection, fixed_position_collection, indices
+        position_collection, fixed_position_collection, indices, constraint_selector
     ) -> None:
         """
         Computes constrain values in numba njit decorator
@@ -386,16 +395,30 @@ 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
-
+        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]
-            position_collection[..., k] = fixed_position_collection[..., 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) -> None:
+    def nb_constrain_translational_rates(
+        velocity_collection, indices, constraint_selector
+    ) -> None:
         """
         Compute constrain rates in numba njit decorator
 
@@ -405,35 +428,105 @@ 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
+            # 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(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.
+            The selector shall be specified in the lab 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 lab frame
+        omega_collection_lab_frame = _batch_matvec(
+            directors.transpose(1, 0, 2), 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(ConfigurableConstraint):
+    """
+    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)
+    ... )
+    """
+
+    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):