Typing: rods

elastica/__init__.py

@@ -1,7 +1,5 @@
 from collections import defaultdict
 from elastica.rod.knot_theory import (
-    KnotTheory,
-    KnotTheoryCompatibleProtocol,
     compute_link,
     compute_twist,
     compute_writhe,

elastica/memory_block/memory_block_rod.py

@@ -1,7 +1,7 @@
 __doc__ = """Create block-structure class for collection of Cosserat rod systems."""
 import numpy as np
 from typing import Literal, Callable
-from elastica.typing import SystemIdxType
+from elastica.typing import SystemIdxType, RodType
 from elastica.memory_block.memory_block_rod_base import (
     make_block_memory_metadata,
     make_block_memory_periodic_boundary_metadata,
@@ -12,7 +12,6 @@
     CosseratRod,
     _compute_sigma_kappa_for_blockstructure,
 )
-from elastica.rod.rod_base import RodBase
 from elastica._synchronize_periodic_boundary import (
     _synchronize_periodic_boundary_of_vector_collection,
     _synchronize_periodic_boundary_of_scalar_collection,
@@ -31,7 +30,7 @@ class MemoryBlockCosseratRod(CosseratRod, _RodSymplecticStepperMixin):
     """
 
     def __init__(
-        self, systems: list[RodBase], system_idx_list: list[SystemIdxType]
+        self, systems: list[RodType], system_idx_list: list[SystemIdxType]
     ) -> None:
 
         # separate straight and ring rods
@@ -201,7 +200,7 @@ def __init__(
         # Initialize the mixin class for symplectic time-stepper.
         _RodSymplecticStepperMixin.__init__(self)
 
-    def _allocate_block_variables_in_nodes(self, systems: list[RodBase]) -> None:
+    def _allocate_block_variables_in_nodes(self, systems: list[RodType]) -> None:
         """
         This function takes system collection and allocates the variables on
         node for block-structure and references allocated variables back to the
@@ -251,7 +250,7 @@ def _allocate_block_variables_in_nodes(self, systems: list[RodBase]) -> None:
             value_type="vector",
         )
 
-    def _allocate_block_variables_in_elements(self, systems: list[RodBase]) -> None:
+    def _allocate_block_variables_in_elements(self, systems: list[RodType]) -> None:
         """
         This function takes system collection and allocates the variables on
         elements for block-structure and references allocated variables back to the
@@ -342,7 +341,7 @@ def _allocate_block_variables_in_elements(self, systems: list[RodBase]) -> None:
             value_type="tensor",
         )
 
-    def _allocate_blocks_variables_in_voronoi(self, systems: list[RodBase]) -> None:
+    def _allocate_blocks_variables_in_voronoi(self, systems: list[RodType]) -> None:
         """
         This function takes system collection and allocates the variables on
         voronoi for block-structure and references allocated variables back to the
@@ -410,7 +409,7 @@ def _allocate_blocks_variables_in_voronoi(self, systems: list[RodBase]) -> None:
         )
 
     def _allocate_blocks_variables_for_symplectic_stepper(
-        self, systems: list[RodBase]
+        self, systems: list[RodType]
     ) -> None:
         """
         This function takes system collection and allocates the variables used by symplectic
@@ -470,7 +469,7 @@ def _allocate_blocks_variables_for_symplectic_stepper(
     def _map_system_properties_to_block_memory(
         self,
         mapping_dict: dict,
-        systems: list[RodBase],
+        systems: list[RodType],
         block_memory: np.ndarray,
         domain_type: Literal["node", "element", "voronoi"],
         value_type: Literal["scalar", "vector", "tensor"],
@@ -485,7 +484,7 @@ def _map_system_properties_to_block_memory(
         ----------
         mapping_dict: dict
             Dictionary with attribute names as keys and block row index as values.
-        systems: list[RodBase]
+        systems: list[RodType]
             A sequence containing Cosserat rod objects to map from.
         block_memory: ndarray
             Memory block that, at the end of the method execution, contains all designated

elastica/memory_block/protocol.py

@@ -0,0 +1,12 @@
+from typing import Protocol
+
+from elastica.rod.protocol import CosseratRodProtocol
+from elastica.systems.protocol import SymplecticSystemProtocol
+
+
+class BlockCosseratRodProtocol(CosseratRodProtocol, SymplecticSystemProtocol, Protocol):
+    pass
+
+
+# class BlockRigidBodyProtocol(RigidBodyProtocol, SymplecticSystemProtocol, Protocol):
+#     pass

elastica/modules/base_system.py

@@ -20,9 +20,9 @@
 
 from collections.abc import MutableSequence
 
-from elastica.rod import RodBase
-from elastica.rigidbody import RigidBodyBase
-from elastica.surface import SurfaceBase
+from elastica.rod.rod_base import RodBase
+from elastica.rigidbody.rigid_body import RigidBodyBase
+from elastica.surface.surface_base import SurfaceBase
 
 from .memory_block import construct_memory_block_structures
 from .operator_group import OperatorGroupFIFO

elastica/modules/memory_block.py

@@ -5,10 +5,11 @@
 
 from elastica.typing import SystemType, SystemIdxType
 
-from elastica.rod import RodBase
-from elastica.rigidbody import RigidBodyBase
-from elastica.surface import SurfaceBase
-from elastica.memory_block import MemoryBlockCosseratRod, MemoryBlockRigidBody
+from elastica.rod.rod_base import RodBase
+from elastica.rigidbody.rigid_body import RigidBodyBase
+from elastica.surface.surface_base import SurfaceBase
+from elastica.memory_block.memory_block_rod import MemoryBlockCosseratRod
+from elastica.memory_block.memory_block_rigid_body import MemoryBlockRigidBody
 
 
 def construct_memory_block_structures(systems: list[SystemType]) -> list[SystemType]:

elastica/rigidbody/rigid_body.py

@@ -1,5 +1,7 @@
 __doc__ = """"""
 
+from typing import Type
+
 import numpy as np
 from abc import ABC
 from elastica._linalg import _batch_matvec, _batch_cross
@@ -18,9 +20,8 @@ class RigidBodyBase(ABC):
 
     """
 
-    REQUISITE_MODULES = []
-
-
+    REQUISITE_MODULES: list[Type] = []
+
     def __init__(self) -> None:
 
         self.position_collection: f_arr_t

elastica/rod/__init__.py

@@ -1,12 +1,4 @@
 __doc__ = """Rod classes and its data structures """
 
 
-from elastica.rod.knot_theory import KnotTheory
 from elastica.rod.rod_base import RodBase
-from elastica.rod.data_structures import (
-    _RodSymplecticStepperMixin,
-    _State,
-    _DerivativeState,
-    _KinematicState,
-    _DynamicState,
-)