Typing: contact

elastica/__init__.py

@@ -36,10 +36,8 @@
 )
 from elastica.joint import (
     FreeJoint,
-    ExternalContact,
     FixedJoint,
     HingeJoint,
-    SelfContact,
 )
 from elastica.contact_forces import (
     NoContact,

elastica/_synchronize_periodic_boundary.py

@@ -7,7 +7,7 @@
 import numpy as np
 from numpy.typing import NDArray
 from elastica.boundary_conditions import ConstraintBase
-from elastica.typing import SystemType
+from elastica.typing import RodType
 
 
 @njit(cache=True)  # type: ignore
@@ -92,15 +92,15 @@ class _ConstrainPeriodicBoundaries(ConstraintBase):
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
 
-    def constrain_values(self, rod: SystemType, time: np.floating) -> None:
+    def constrain_values(self, rod: RodType, time: np.floating) -> None:
         _synchronize_periodic_boundary_of_vector_collection(
             rod.position_collection, rod.periodic_boundary_nodes_idx
         )
         _synchronize_periodic_boundary_of_matrix_collection(
             rod.director_collection, rod.periodic_boundary_elems_idx
         )
 
-    def constrain_rates(self, rod: SystemType, time: np.floating) -> None:
+    def constrain_rates(self, rod: RodType, time: np.floating) -> None:
         _synchronize_periodic_boundary_of_vector_collection(
             rod.velocity_collection, rod.periodic_boundary_nodes_idx
         )

elastica/boundary_conditions.py

@@ -1,7 +1,7 @@
 __doc__ = """ Built-in boundary condition implementationss """
 
 import warnings
-from typing import Any, Optional
+from typing import Any, Optional, TypeVar, Generic
 
 import numpy as np
 from numpy.typing import NDArray
@@ -12,10 +12,13 @@
 
 from elastica._linalg import _batch_matvec, _batch_matrix_transpose
 from elastica._rotations import _get_rotation_matrix
-from elastica.typing import SystemType, RodType
+from elastica.typing import SystemType, RodType, RigidBodyType
 
 
-class ConstraintBase(ABC):
+S = TypeVar("S")
+
+
+class ConstraintBase(ABC, Generic[S]):
     """Base class for constraint and displacement boundary condition implementation.
 
     Notes
@@ -31,7 +34,7 @@ class ConstraintBase(ABC):
 
     """
 
-    _system: SystemType
+    _system: S
     _constrained_position_idx: np.ndarray
     _constrained_director_idx: np.ndarray
 
@@ -51,24 +54,24 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             )
 
     @property
-    def system(self) -> SystemType:
+    def system(self) -> S:
         """get system (rod or rigid body) reference"""
         return self._system
 
     @property
-    def constrained_position_idx(self) -> Optional[np.ndarray]:
+    def constrained_position_idx(self) -> np.ndarray:
         """get position-indices passed to "using" """
         # TODO: This should be immutable somehow
         return self._constrained_position_idx
 
     @property
-    def constrained_director_idx(self) -> Optional[np.ndarray]:
+    def constrained_director_idx(self) -> np.ndarray:
         """get director-indices passed to "using" """
         # TODO: This should be immutable somehow
         return self._constrained_director_idx
 
     @abstractmethod
-    def constrain_values(self, system: SystemType, time: np.floating) -> None:
+    def constrain_values(self, system: S, time: np.floating) -> None:
         # TODO: In the future, we can remove rod and use self.system
         """
         Constrain values (position and/or directors) of a rod object.
@@ -83,7 +86,7 @@ def constrain_values(self, system: SystemType, time: np.floating) -> None:
         pass
 
     @abstractmethod
-    def constrain_rates(self, system: SystemType, time: np.floating) -> None:
+    def constrain_rates(self, system: S, time: np.floating) -> None:
         # TODO: In the future, we can remove rod and use self.system
         """
         Constrain rates (velocity and/or omega) of a rod object.
@@ -107,11 +110,15 @@ class FreeBC(ConstraintBase):
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
 
-    def constrain_values(self, system: SystemType, time: np.floating) -> None:
+    def constrain_values(
+        self, system: "RodType | RigidBodyType", time: np.floating
+    ) -> None:
         """In FreeBC, this routine simply passes."""
         pass
 
-    def constrain_rates(self, system: SystemType, time: np.floating) -> None:
+    def constrain_rates(
+        self, system: "RodType | RigidBodyType", time: np.floating
+    ) -> None:
         """In FreeBC, this routine simply passes."""
         pass
 
@@ -165,7 +172,9 @@ def __init__(
         self.fixed_position_collection = np.array(fixed_position)
         self.fixed_directors_collection = np.array(fixed_directors)
 
-    def constrain_values(self, system: SystemType, time: np.floating) -> None:
+    def constrain_values(
+        self, system: "RodType | RigidBodyType", time: np.floating
+    ) -> None:
         # system.position_collection[..., 0] = self.fixed_position
         # system.director_collection[..., 0] = self.fixed_directors
         self.compute_constrain_values(
@@ -175,7 +184,9 @@ def constrain_values(self, system: SystemType, time: np.floating) -> None:
             self.fixed_directors_collection,
         )
 
-    def constrain_rates(self, system: SystemType, time: np.floating) -> None:
+    def constrain_rates(
+        self, system: "RodType | RigidBodyType", time: np.floating
+    ) -> None:
         # system.velocity_collection[..., 0] = 0.0
         # system.omega_collection[..., 0] = 0.0
         self.compute_constrain_rates(
@@ -340,7 +351,9 @@ def __init__(
         )
         self.rotational_constraint_selector = rotational_constraint_selector.astype(int)
 
-    def constrain_values(self, system: SystemType, time: np.floating) -> None:
+    def constrain_values(
+        self, system: "RodType | RigidBodyType", time: np.floating
+    ) -> None:
         if self.constrained_position_idx.size:
             self.nb_constrain_translational_values(
                 system.position_collection,
@@ -349,7 +362,9 @@ def constrain_values(self, system: SystemType, time: np.floating) -> None:
                 self.translational_constraint_selector,
             )
 
-    def constrain_rates(self, system: SystemType, time: np.floating) -> None:
+    def constrain_rates(
+        self, system: "RodType | RigidBodyType", time: np.floating
+    ) -> None:
         if self.constrained_position_idx.size:
             self.nb_constrain_translational_rates(
                 system.velocity_collection,
@@ -525,7 +540,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             **kwargs,
         )
 
-    def constrain_values(self, system: SystemType, time: np.floating) -> None:
+    def constrain_values(
+        self, system: "RodType | RigidBodyType", time: np.floating
+    ) -> None:
         if self.constrained_position_idx.size:
             self.nb_constrain_translational_values(
                 system.position_collection,
@@ -539,7 +556,9 @@ def constrain_values(self, system: SystemType, time: np.floating) -> None:
                 self.constrained_director_idx,
             )
 
-    def constrain_rates(self, system: SystemType, time: np.floating) -> None:
+    def constrain_rates(
+        self, system: "RodType | RigidBodyType", time: np.floating
+    ) -> None:
         if self.constrained_position_idx.size:
             self.nb_constrain_translational_rates(
                 system.velocity_collection,
@@ -743,15 +762,19 @@ def __init__(
             @ director_end
         )  # rotation_matrix wants vectors 3,1
 
-    def constrain_values(self, rod: SystemType, time: np.floating) -> None:
+    def constrain_values(
+        self, rod: "RodType | RigidBodyType", time: np.floating
+    ) -> None:
         if time > self.twisting_time:
             rod.position_collection[..., 0] = self.final_start_position
             rod.position_collection[..., -1] = self.final_end_position
 
             rod.director_collection[..., 0] = self.final_start_directors
             rod.director_collection[..., -1] = self.final_end_directors
 
-    def constrain_rates(self, rod: SystemType, time: np.floating) -> None:
+    def constrain_rates(
+        self, rod: "RodType | RigidBodyType", time: np.floating
+    ) -> None:
         if time > self.twisting_time:
             rod.velocity_collection[..., 0] = 0.0
             rod.omega_collection[..., 0] = 0.0

elastica/callback_functions.py

@@ -1,19 +1,21 @@
 __doc__ = """ Module contains callback classes to save simulation data for rod-like objects """
-from elastica.typing import SystemType
+from typing import Any, Optional, TypeVar, Generic
+from elastica.typing import RodType, RigidBodyType, SystemType
 
 import os
 import sys
 import numpy as np
+from numpy.typing import NDArray
 import logging
-from typing import Any, Optional
 
 
 from collections import defaultdict
 
-from elastica.typing import RodType, SystemType
 
+T = TypeVar("T")
 
-class CallBackBaseClass:
+
+class CallBackBaseClass(Generic[T]):
     """
     This is the base class for callbacks for rod-like objects.
 
@@ -30,9 +32,7 @@ def __init__(self) -> None:
         """
         pass
 
-    def make_callback(
-        self, system: SystemType, time: np.floating, current_step: int
-    ) -> None:
+    def make_callback(self, system: T, time: np.floating, current_step: int) -> None:
         """
         This method is called every time step. Users can define
         which parameters are called back and recorded. Also users
@@ -81,7 +81,7 @@ def __init__(self, step_skip: int, callback_params: dict) -> None:
         self.callback_params = callback_params
 
     def make_callback(
-        self, system: SystemType, time: np.floating, current_step: int
+        self, system: "RodType | RigidBodyType", time: np.floating, current_step: int
     ) -> None:
 
         if current_step % self.sample_every == 0:
@@ -176,7 +176,9 @@ def __init__(
         self.file_save_interval = file_save_interval
 
         # Data collector
-        self.buffer = defaultdict(list)
+        self.buffer: dict[str, list[NDArray[np.floating] | np.floating | int]] = (
+            defaultdict(list)
+        )
         self.buffer_size = 0
 
         # Module
@@ -199,7 +201,7 @@ def __init__(
             self._ext = "pkl"
 
     def make_callback(
-        self, system: SystemType, time: np.floating, current_step: int
+        self, system: "RodType | RigidBodyType", time: np.floating, current_step: int
     ) -> None:
         """