diff --git a/elastica/__init__.py b/elastica/__init__.py index 7cdfc7f39..f705615a4 100644 --- a/elastica/__init__.py +++ b/elastica/__init__.py @@ -36,10 +36,8 @@ ) from elastica.joint import ( FreeJoint, - ExternalContact, FixedJoint, HingeJoint, - SelfContact, ) from elastica.contact_forces import ( NoContact, diff --git a/elastica/_synchronize_periodic_boundary.py b/elastica/_synchronize_periodic_boundary.py index 7d4c4882b..49784f6c6 100644 --- a/elastica/_synchronize_periodic_boundary.py +++ b/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,7 +92,7 @@ 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 ) @@ -100,7 +100,7 @@ def constrain_values(self, rod: SystemType, time: np.floating) -> None: 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 ) diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 3a1b99642..6b25b8653 100644 --- a/elastica/boundary_conditions.py +++ b/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,7 +762,9 @@ 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 @@ -751,7 +772,9 @@ def constrain_values(self, rod: SystemType, time: np.floating) -> None: 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 diff --git a/elastica/callback_functions.py b/elastica/callback_functions.py index dac062080..8e89683de 100644 --- a/elastica/callback_functions.py +++ b/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: """ diff --git a/elastica/collision/AABBCollection.py b/elastica/collision/AABBCollection.py index 22b716065..9385dab40 100644 --- a/elastica/collision/AABBCollection.py +++ b/elastica/collision/AABBCollection.py @@ -1,17 +1,20 @@ """ Axis Aligned Bounding Boxes for coarse collision detection """ +from typing_extensions import Self + import numpy as np +from numpy.typing import NDArray from elastica.utils import MaxDimension class AABBCollection: def __init__( self, - elemental_position_collection, - dimension_collection, + elemental_position_collection: NDArray[np.floating], + dimension_collection: NDArray[np.floating], elements_per_aabb: int, - ): + ) -> None: """ Doesn't differentiate tangent direction from the rest : potentially harmful as maybe you don't need to expand to radius amount in tangential direction @@ -36,7 +39,9 @@ def __init__( self.update(elemental_position_collection, dimension_collection) @classmethod - def make_from_aabb(cls, aabb_collection, scale_factor=4): + def make_from_aabb( + cls, aabb_collection: list["AABBCollection"], scale_factor: int = 4 + ) -> Self: # Make position collection and dimension collection arrays from aabb_collection # Wasted effort, but only once during construction n_aabb_from_lower_level = len(aabb_collection) @@ -59,7 +64,7 @@ def make_from_aabb(cls, aabb_collection, scale_factor=4): return cls(elemental_position_collection, dimension_collection, scale_factor) - def _update(self, aabb_collection): + def _update(self, aabb_collection: list["AABBCollection"]) -> None: # Updates internal state from another aabb """ @@ -78,7 +83,11 @@ def _update(self, aabb_collection): temp = np.array([aabb.aabb[..., 1, 0] for aabb in aabb_collection]) self.aabb[..., 1, 0] = np.amax(temp, axis=0) - def update(self, elemental_position_collection, dimension_collection): + def update( + self, + elemental_position_collection: NDArray[np.floating], + dimension_collection: NDArray[np.floating], + ) -> None: # Initialize the boxes for i in range(self.n_aabb): start = i * self.elements_per_aabb @@ -91,7 +100,7 @@ def update(self, elemental_position_collection, dimension_collection): ) + np.amax(dimension_collection[..., start:stop], axis=1) -def find_nearest_integer_square_root(x: int): +def find_nearest_integer_square_root(x: int) -> int: from math import sqrt return round(sqrt(x)) @@ -101,8 +110,11 @@ class AABBHierarchy: """Simple hierarchy for handling cylinder collisions alone, meant for a rod""" def __init__( - self, position_collection, dimension_collection, avg_n_dofs_in_final_level - ): + self, + position_collection: NDArray[np.floating], + dimension_collection: NDArray[np.floating], + avg_n_dofs_in_final_level: int, + ) -> None: """ scaling is always set to 4, so that theres' 1 major AABBCollection, then scaling_factor smaller AABBs, then scaling factor even smaller AABBs (which cover the elements @@ -121,10 +133,10 @@ def __init__( ) # nearest power of 4 that is less than the number - n_levels_bound_below = np.int( + n_levels_bound_below = int( np.floor(0.5 * np.log2(potential_n_aabbs_in_final_level)) ) - n_levels_bound_above = np.int( + n_levels_bound_above = int( np.ceil(0.5 * np.log2(potential_n_aabbs_in_final_level)) ) # Check which is the closest and use that as the number of levels @@ -214,11 +226,15 @@ def __init__( # Add one for the middle level # self.aabb.append(AABBCollection(position_collection, dimension_collection, self.n_aabbs_in_first_level)) - def n_aabbs_at_level(self, i: int): + def n_aabbs_at_level(self, i: int) -> int: assert i < self.n_levels return 4 ** (i) - def update(self, position_collection, dimension_collection): + def update( + self, + position_collection: NDArray[np.floating], + dimension_collection: NDArray[np.floating], + ) -> None: # Update bottom level first, the first level entries n_aabbs_in_final_level = self.n_aabbs_at_level(self.n_levels - 1) stop = 0 @@ -261,5 +277,8 @@ def update(self, position_collection, dimension_collection): count_elapsed_n_aabbs += n_aabbs_in_next_level -def are_aabb_intersecting(first_aabb_collection, second_aabb_collection): +def are_aabb_intersecting( + first_aabb_collection: NDArray[np.floating], + second_aabb_collection: NDArray[np.floating], +) -> bool: return True diff --git a/elastica/contact_forces.py b/elastica/contact_forces.py index 9e9f5a2de..22adcb3e9 100644 --- a/elastica/contact_forces.py +++ b/elastica/contact_forces.py @@ -1,10 +1,13 @@ __doc__ = """ Numba implementation module containing contact between rods and rigid bodies and other rods rigid bodies or surfaces.""" -from typing import Optional -from elastica.typing import RodType, SystemType, AllowedContactType -from elastica.rod import RodBase -from elastica.rigidbody import Cylinder, Sphere -from elastica.surface import Plane +from typing import TypeVar, Generic, Type +from elastica.typing import RodType, SystemType, SurfaceType + +from elastica.rod.rod_base import RodBase +from elastica.rigidbody.cylinder import Cylinder +from elastica.rigidbody.sphere import Sphere +from elastica.surface.plane import Plane +from elastica.surface.surface_base import SurfaceBase from elastica.contact_utils import ( _prune_using_aabbs_rod_cylinder, _prune_using_aabbs_rod_rod, @@ -23,7 +26,11 @@ from numpy.typing import NDArray -class NoContact: +S1 = TypeVar("S1") # TODO: Find bound +S2 = TypeVar("S2") + + +class NoContact(Generic[S1, S2]): """ This is the base class for contact applied between rod-like objects and allowed contact objects. @@ -38,51 +45,48 @@ def __init__(self) -> None: """ NoContact class does not need any input parameters. """ + pass + + @property + def _allowed_system_one(self) -> list[Type]: + # Modify this list to include the allowed system types for contact + return [RodBase] + + @property + def _allowed_system_two(self) -> list[Type]: + # Modify this list to include the allowed system types for contact + return [RodBase] def _check_systems_validity( self, - system_one: SystemType, - system_two: AllowedContactType, + system_one: S1, + system_two: S2, ) -> None: """ - This checks the contact order between a SystemType object and an AllowedContactType object, the order should follow: Rod, Rigid body, Surface. - In NoContact class, this just checks if system_two is a rod then system_one must be a rod. + Here, we check the allowed system types for contact. + For derived classes, this method can be overridden to enforce specific system types + for contact model. + """ + common_check_systems_validity(system_one, self._allowed_system_one) + common_check_systems_validity(system_two, self._allowed_system_two) - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if issubclass(system_two.__class__, RodBase): - if not issubclass(system_one.__class__, RodBase): - raise TypeError( - "Systems provided to the contact class have incorrect order. \n" - " First system is {0} and second system is {1}. \n" - " If the first system is a rod, the second system can be a rod, rigid body or surface. \n" - " If the first system is a rigid body, the second system can be a rigid body or surface.".format( - system_one.__class__, system_two.__class__ - ) - ) + common_check_systems_identity(system_one, system_two) def apply_contact( self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> Optional[tuple[NDArray[np.floating], NDArray[np.intp]]]: + system_one: S1, + system_two: S2, + ) -> None: """ - Apply contact forces and torques between SystemType object and AllowedContactType object. + Apply contact forces and torques between two system object.. In NoContact class, this routine simply passes. Parameters ---------- - system_one : SystemType - Rod or rigid-body object - system_two : AllowedContactType - Rod, rigid-body, or surface object + system_one + system_two """ pass @@ -116,49 +120,14 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: self.k = k self.nu = nu - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodRodContact class both systems must be distinct rods. - - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, RodBase - ): - raise TypeError( - "Systems provided to the contact class have incorrect order. \n" - " First system is {0} and second system is {1}. \n" - " Both systems must be distinct rods".format( - system_one.__class__, system_two.__class__ - ) - ) - if system_one == system_two: - raise TypeError( - "First rod is identical to second rod. \n" - "Rods must be distinct for RodRodConact. \n" - "If you want self contact, use RodSelfContact instead" - ) - def apply_contact(self, system_one: RodType, system_two: RodType) -> None: """ Apply contact forces and torques between RodType object and RodType object. Parameters ---------- - system_one: object - Rod object. - system_two: object - Rod object. + system_one: RodType + system_two: RodType """ # First, check for a global AABB bounding box, and see whether that @@ -229,8 +198,8 @@ def __init__( self, k: np.floating, nu: np.floating, - velocity_damping_coefficient: np.floating = 0.0, - friction_coefficient: np.floating = 0.0, + velocity_damping_coefficient: float = 0.0, + friction_coefficient: float = 0.0, ) -> None: """ @@ -252,34 +221,12 @@ def __init__( self.velocity_damping_coefficient = velocity_damping_coefficient self.friction_coefficient = friction_coefficient - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodCylinderContact class first_system should be a rod and second_system should be a cylinder. + @property + def _allowed_system_two(self) -> list[Type]: + # Modify this list to include the allowed system types for contact + return [Cylinder] - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, Cylinder - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a cylinder".format( - system_one.__class__, system_two.__class__ - ) - ) - - def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: + def apply_contact(self, system_one: RodType, system_two: Cylinder) -> None: # First, check for a global AABB bounding box, and see whether that # intersects if _prune_using_aabbs_rod_cylinder( @@ -288,8 +235,8 @@ def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: system_one.lengths, system_two.position_collection, system_two.director_collection, - system_two.radius[0], - system_two.length[0], + system_two.radius, + system_two.length, ): return @@ -356,33 +303,15 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: def _check_systems_validity( self, - system_one: SystemType, - system_two: AllowedContactType, + system_one: RodType, + system_two: RodType, ) -> None: """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodSelfContact class first_system and second_system should be the same rod. - - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType + Overriding the base class method to check if the two systems are identical. """ - if ( - not issubclass(system_one.__class__, RodBase) - or not issubclass(system_two.__class__, RodBase) - or system_one != system_two - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system and second system should be the same rod \n" - " If you want rod rod contact, use RodRodContact instead".format( - system_one.__class__, system_two.__class__ - ) - ) + common_check_systems_validity(system_one, self._allowed_system_one) + common_check_systems_validity(system_two, self._allowed_system_two) + common_check_systems_different(system_one, system_two) def apply_contact(self, system_one: RodType, system_two: RodType) -> None: """ @@ -390,10 +319,8 @@ def apply_contact(self, system_one: RodType, system_two: RodType) -> None: Parameters ---------- - system_one: object - Rod object. - system_two: object - Rod object. + system_one: RodType + system_two: RodType """ _calculate_contact_forces_self_rod( @@ -439,8 +366,8 @@ def __init__( self, k: np.floating, nu: np.floating, - velocity_damping_coefficient: np.floating = 0.0, - friction_coefficient: np.floating = 0.0, + velocity_damping_coefficient: float = 0.0, + friction_coefficient: float = 0.0, ) -> None: """ Parameters @@ -459,44 +386,20 @@ def __init__( self.k = k self.nu = nu self.velocity_damping_coefficient = velocity_damping_coefficient - self.friction_coefficient = friction_coefficient + self.friction_coefficient = np.float64(friction_coefficient) - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodSphereContact class first_system should be a rod and second_system should be a sphere. - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, Sphere - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a sphere".format( - system_one.__class__, system_two.__class__ - ) - ) + @property + def _allowed_system_two(self) -> list[Type]: + return [Sphere] - def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: + def apply_contact(self, system_one: RodType, system_two: Sphere) -> None: """ Apply contact forces and torques between RodType object and Sphere object. Parameters ---------- - system_one: object - Rod object. - system_two: object - Sphere object. + system_one: RodType + system_two: Sphere """ # First, check for a global AABB bounding box, and see whether that @@ -507,7 +410,7 @@ def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: system_one.lengths, system_two.position_collection, system_two.director_collection, - system_two.radius[0], + system_two.radius, ): return @@ -578,33 +481,11 @@ def __init__( self.nu = nu self.surface_tol = 1e-4 - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodPlaneContact class first_system should be a rod and second_system should be a plane. - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, Plane - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane".format( - system_one.__class__, system_two.__class__ - ) - ) + @property + def _allowed_system_two(self) -> list[Type]: + return [SurfaceBase] - def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: + def apply_contact(self, system_one: RodType, system_two: SurfaceType) -> None: """ Apply contact forces and torques between RodType object and Plane object. @@ -692,42 +573,18 @@ def __init__( self.kinetic_mu_sideways, ) = kinetic_mu_array - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodSphereContact class first_system should be a rod and second_system should be a plane. - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, Plane - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane".format( - system_one.__class__, system_two.__class__ - ) - ) + @property + def _allowed_system_two(self) -> list[Type]: + return [SurfaceBase] - def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: + def apply_contact(self, system_one: RodType, system_two: SurfaceType) -> None: """ Apply contact forces and torques between RodType object and Plane object with anisotropic friction. Parameters ---------- - system_one: object - Rod object. - system_two: object - Plane object. + system_one: RodType + system_two: SurfaceType """ @@ -794,35 +651,15 @@ def __init__( self.nu = nu self.surface_tol = 1e-4 - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodPlaneContact class first_system should be a cylinder and second_system should be a plane. - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, Cylinder) or not issubclass( - system_two.__class__, Plane - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a cylinder, second should be a plane".format( - system_one.__class__, system_two.__class__ - ) - ) + @property + def _allowed_system_one(self) -> list[Type]: + return [Cylinder] - def apply_contact( - self, system_one: Cylinder, system_two: SystemType - ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: + @property + def _allowed_system_two(self) -> list[Type]: + return [SurfaceBase] + + def apply_contact(self, system_one: Cylinder, system_two: SurfaceType) -> None: """ This function computes the plane force response on the cylinder, in the case of contact. Contact model given in Eqn 4.8 Gazzola et. al. RSoS 2018 paper @@ -830,13 +667,11 @@ def apply_contact( Parameters ---------- - system_one: object - Cylinder object. - system_two: object - Plane object. + system_one: Cylinder + system_two: SurfaceBase """ - return _calculate_contact_forces_cylinder_plane( + _calculate_contact_forces_cylinder_plane( system_two.origin, system_two.normal, self.surface_tol, @@ -847,3 +682,49 @@ def apply_contact( system_one.velocity_collection, system_one.external_forces, ) + + +def common_check_systems_identity( + system_one: S1, + system_two: S2, +) -> None: + """ + This checks if two objects are identical. + + Raises + ------ + TypeError + If two objects are identical. + """ + if system_one == system_two: + raise TypeError( + "First system is identical to second system. Systems must be distinct for contact." + ) + + +def common_check_systems_different( + system_one: S1, + system_two: S2, +) -> None: + """ + This checks if two objects are identical. + + Raises + ------ + TypeError + If two objects are not identical. + """ + if system_one != system_two: + raise TypeError("First system must be identical to the second system.") + + +def common_check_systems_validity( + system: S1 | S2, allowed_system: list[Type[S1] | Type[S2]] +) -> None: + # Check validity + if not isinstance(system, tuple(allowed_system)): + system_name = system.__class__.__name__ + allowed_system_names = [candidate.__name__ for candidate in allowed_system] + raise TypeError( + f"System provided ({system_name}) must be derived from {allowed_system_names}." + ) diff --git a/elastica/contact_utils.py b/elastica/contact_utils.py index 657b6dd00..a67793a07 100644 --- a/elastica/contact_utils.py +++ b/elastica/contact_utils.py @@ -7,37 +7,34 @@ from elastica._linalg import ( _batch_norm, ) -from typing import Literal, Sequence, TypeVar +from typing import Literal, Sequence -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _dot_product(a: Sequence[np.floating], b: Sequence[np.floating]) -> np.floating: - sum: np.floating = 0.0 + total: np.float64 = np.float64(0.0) for i in range(3): - sum += a[i] * b[i] - return sum + total += a[i] * b[i] + return total -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _norm(a: Sequence[np.floating]) -> float: return sqrt(_dot_product(a, a)) -_SupportsCompareT = TypeVar("_SupportsCompareT") - - -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _clip(x: np.floating, low: np.floating, high: np.floating) -> np.floating: return max(low, min(x, high)) # Can this be made more efficient than 2 comp, 1 or? -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _out_of_bounds(x: np.floating, low: np.floating, high: np.floating) -> bool: return (x < low) or (x > high) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _find_min_dist( x1: NDArray[np.floating], e1: NDArray[np.floating], @@ -108,7 +105,7 @@ def _find_min_dist( return x2 + s * e2 - x1 - t * e1, x2 + s * e2, x1 - t * e1 -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _aabbs_not_intersecting( aabb_one: NDArray[np.floating], aabb_two: NDArray[np.floating] ) -> Literal[1, 0]: @@ -123,7 +120,7 @@ def _aabbs_not_intersecting( return 0 -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _prune_using_aabbs_rod_cylinder( rod_one_position_collection: NDArray[np.floating], rod_one_radius_collection: NDArray[np.floating], @@ -165,7 +162,7 @@ def _prune_using_aabbs_rod_cylinder( return _aabbs_not_intersecting(aabb_cylinder, aabb_rod) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _prune_using_aabbs_rod_rod( rod_one_position_collection: NDArray[np.floating], rod_one_radius_collection: NDArray[np.floating], @@ -203,7 +200,7 @@ def _prune_using_aabbs_rod_rod( return _aabbs_not_intersecting(aabb_rod_two, aabb_rod_one) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _prune_using_aabbs_rod_sphere( rod_one_position_collection: NDArray[np.floating], rod_one_radius_collection: NDArray[np.floating], @@ -242,7 +239,7 @@ def _prune_using_aabbs_rod_sphere( return _aabbs_not_intersecting(aabb_sphere, aabb_rod) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _find_slipping_elements( velocity_slip: NDArray[np.floating], velocity_threshold: np.floating ) -> NDArray[np.floating]: @@ -285,7 +282,7 @@ def _find_slipping_elements( return slip_function -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _node_to_element_mass_or_force(input: NDArray[np.floating]) -> NDArray[np.floating]: """ This function converts the mass/forces on rod nodes to @@ -323,7 +320,7 @@ def _node_to_element_mass_or_force(input: NDArray[np.floating]) -> NDArray[np.fl return output -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _elements_to_nodes_inplace( vector_in_element_frame: NDArray[np.floating], vector_in_node_frame: NDArray[np.floating], @@ -349,7 +346,7 @@ def _elements_to_nodes_inplace( vector_in_node_frame[i, k + 1] += 0.5 * vector_in_element_frame[i, k] -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _node_to_element_position( node_position_collection: NDArray[np.floating], ) -> NDArray[np.floating]: @@ -397,7 +394,7 @@ def _node_to_element_position( return element_position_collection -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _node_to_element_velocity( mass: NDArray[np.floating], node_velocity_collection: NDArray[np.floating] ) -> NDArray[np.floating]: diff --git a/elastica/dissipation.py b/elastica/dissipation.py index dfe5373af..2cc532918 100644 --- a/elastica/dissipation.py +++ b/elastica/dissipation.py @@ -5,7 +5,7 @@ """ from abc import ABC, abstractmethod -from typing import Any +from typing import Any, Generic, TypeVar from elastica.typing import RodType, SystemType @@ -15,7 +15,10 @@ from numpy.typing import NDArray -class DamperBase(ABC): +T = TypeVar("T") + + +class DamperBase(Generic[T], ABC): """Base class for damping module implementations. Notes @@ -25,11 +28,11 @@ class DamperBase(ABC): Attributes ---------- - system : SystemType (RodBase or RigidBodyBase) + system : RodBase """ - _system: SystemType + _system: T # TODO typing can be made better def __init__(self, *args: Any, **kwargs: Any) -> None: @@ -43,7 +46,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ) @property - def system(self) -> SystemType: + def system(self) -> T: """ get system (rod or rigid body) reference @@ -55,7 +58,7 @@ def system(self) -> SystemType: return self._system @abstractmethod - def dampen_rates(self, system: SystemType, time: np.floating) -> None: + def dampen_rates(self, system: T, time: np.floating) -> None: # TODO: In the future, we can remove rod and use self.system """ Dampen rates (velocity and/or omega) of a rod object. @@ -148,7 +151,7 @@ def __init__( * np.diagonal(self._system.inv_mass_second_moment_of_inertia).T ) - def dampen_rates(self, rod: SystemType, time: np.floating) -> None: + def dampen_rates(self, rod: RodType, time: np.floating) -> None: rod.velocity_collection[:] = ( rod.velocity_collection * self.translational_damping_coefficient ) diff --git a/elastica/experimental/interaction.py b/elastica/experimental/interaction.py index 776f0da7e..2a328934a 100644 --- a/elastica/experimental/interaction.py +++ b/elastica/experimental/interaction.py @@ -3,9 +3,9 @@ import numpy as np from elastica.external_forces import NoForces +from elastica.contact_utils import _find_slipping_elements +from elastica._contact_functions import _calculate_contact_forces_cylinder_plane from elastica.interaction import ( - find_slipping_elements, - apply_normal_force_numba_rigid_body, InteractionPlaneRigidBody, ) @@ -93,7 +93,7 @@ def anisotropic_friction_numba_rigid_body( ( plane_response_force_mag, no_contact_point_idx, - ) = apply_normal_force_numba_rigid_body( + ) = _calculate_contact_forces_cylinder_plane( plane_origin, plane_normal, surface_tol, @@ -122,7 +122,7 @@ def anisotropic_friction_numba_rigid_body( + kinetic_mu_backward * (1 - velocity_sign_along_axial_direction) ) # Call slip function to check if elements slipping or not - slip_function_along_axial_direction = find_slipping_elements( + slip_function_along_axial_direction = _find_slipping_elements( velocity_along_axial_direction, slip_velocity_tol ) kinetic_friction_force_along_axial_direction = -( @@ -151,7 +151,7 @@ def anisotropic_friction_numba_rigid_body( + kinetic_mu_backward * (1 - velocity_sign_along_binormal_direction) ) # Call slip function to check if elements slipping or not - slip_function_along_binormal_direction = find_slipping_elements( + slip_function_along_binormal_direction = _find_slipping_elements( velocity_along_binormal_direction, slip_velocity_tol ) kinetic_friction_force_along_binormal_direction = -( diff --git a/elastica/external_forces.py b/elastica/external_forces.py index 00cbe60d0..ae85a3569 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -1,19 +1,23 @@ __doc__ = """ Numba implementation module for boundary condition implementations that apply external forces to the system.""" +from typing import TypeVar, Generic import numpy as np from numpy.typing import NDArray from elastica._linalg import _batch_matvec -from elastica.typing import SystemType, RodType +from elastica.typing import SystemType, RodType, RigidBodyType from elastica.utils import _bspline from numba import njit from elastica._linalg import _batch_product_i_k_to_ik -class NoForces: +S = TypeVar("S") + + +class NoForces(Generic[S]): """ This is the base class for external forcing boundary conditions applied to rod-like objects. @@ -30,7 +34,7 @@ def __init__(self) -> None: """ pass - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces(self, system: S, time: float = 0.0) -> None: """Apply forces to a rod-like object. In NoForces class, this routine simply passes. @@ -45,7 +49,7 @@ def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: """ pass - def apply_torques(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_torques(self, system: S, time: float = 0.0) -> None: """Apply torques to a rod-like object. In NoForces class, this routine simply passes. @@ -86,7 +90,9 @@ def __init__( super(GravityForces, self).__init__() self.acc_gravity = acc_gravity - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces( + self, system: "RodType | RigidBodyType", time: float = 0.0 + ) -> None: self.compute_gravity_forces( self.acc_gravity, system.mass, system.external_forces ) @@ -134,7 +140,7 @@ def __init__( self, start_force: NDArray[np.floating], end_force: NDArray[np.floating], - ramp_up_time: np.floating, + ramp_up_time: float, ) -> None: """ @@ -156,7 +162,9 @@ def __init__( assert ramp_up_time > 0.0 self.ramp_up_time = ramp_up_time - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces( + self, system: "RodType | RigidBodyType", time: float = 0.0 + ) -> None: self.compute_end_point_forces( system.external_forces, self.start_force, @@ -171,7 +179,7 @@ def compute_end_point_forces( external_forces: NDArray[np.floating], start_force: NDArray[np.floating], end_force: NDArray[np.floating], - time: np.floating, + time: float, ramp_up_time: np.floating, ) -> None: """ @@ -191,9 +199,9 @@ def compute_end_point_forces( Applied forces are ramped up until ramp up time. """ - factor: np.floating = min(1.0, time / ramp_up_time) - external_forces[..., 0] += start_force * factor - external_forces[..., -1] += end_force * factor + factor = min(1.0, time / ramp_up_time) + external_forces[..., 0] += start_force * factor # type: ignore[operator] + external_forces[..., -1] += end_force * factor # type: ignore[operator] class UniformTorques(NoForces): @@ -225,7 +233,9 @@ def __init__( super(UniformTorques, self).__init__() self.torque = torque * direction - def apply_torques(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_torques( + self, system: "RodType | RigidBodyType", time: float = 0.0 + ) -> None: n_elems = system.n_elems torque_on_one_element = ( _batch_product_i_k_to_ik(self.torque, np.ones((n_elems))) / n_elems @@ -263,7 +273,7 @@ def __init__( super(UniformForces, self).__init__() self.force = (force * direction).reshape(3, 1) - def apply_forces(self, rod: SystemType, time: np.floating = 0.0) -> None: + def apply_forces(self, rod: "RodType | RigidBodyType", time: float = 0.0) -> None: force_on_one_element = self.force / rod.n_elems rod.external_forces += force_on_one_element @@ -327,7 +337,7 @@ def __init__( Phase shift of traveling wave. direction: numpy.ndarray 1D (dim) array containing data with 'float' type. Muscle torque direction. - ramp_up_time: float + ramp_up_time: np.floating Applied muscle torques are ramped up until ramp up time. with_spline: boolean Option to use beta-spline. @@ -360,7 +370,7 @@ def __init__( else: self.my_spline = np.full_like(self.s, fill_value=1.0) - def apply_torques(self, rod: SystemType, time: np.floating = 0.0) -> None: + def apply_torques(self, rod: "RodType | RigidBodyType", time: float = 0.0) -> None: self.compute_muscle_torques( time, self.my_spline, @@ -377,7 +387,7 @@ def apply_torques(self, rod: SystemType, time: np.floating = 0.0) -> None: @staticmethod @njit(cache=True) # type: ignore def compute_muscle_torques( - time: np.floating, + time: float, my_spline: NDArray[np.floating], s: np.floating, angular_frequency: np.floating, @@ -389,7 +399,7 @@ def compute_muscle_torques( external_torques: NDArray[np.floating], ) -> None: # Ramp up the muscle torque - factor: np.floating = min(1.0, time / ramp_up_time) + factor = np.float64(min(np.float64(1.0), time / ramp_up_time)) # From the node 1 to node nelem-1 # Magnitude of the torque. Am = beta(s) * sin(2pi*t/T + 2pi*s/lambda + phi) # There is an inconsistency with paper and Elastica cpp implementation. In paper sign in @@ -493,7 +503,7 @@ def __init__( self, start_force_mag: np.floating, end_force_mag: np.floating, - ramp_up_time: np.floating = 0.0, + ramp_up_time: float = 0.0, tangent_direction: NDArray[np.floating] = np.array([0, 0, 1]), normal_direction: NDArray[np.floating] = np.array([0, 1, 0]), ) -> None: @@ -526,7 +536,9 @@ def __init__( assert ramp_up_time >= 0.0 self.ramp_up_time = ramp_up_time - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces( + self, system: "RodType | RigidBodyType", time: float = 0.0 + ) -> None: if time < self.ramp_up_time: # When time smaller than ramp up time apply the force in normal direction diff --git a/elastica/interaction.py b/elastica/interaction.py index f0b9a9b28..24ecb92e7 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -1,7 +1,6 @@ __doc__ = """ Numba implementation module containing interactions between a rod and its environment.""" -from typing import Any, NoReturn import numpy as np from elastica.external_forces import NoForces from numba import njit @@ -17,48 +16,12 @@ from numpy.typing import NDArray -from elastica.typing import SystemType - - -def find_slipping_elements(velocity_slip: Any, velocity_threshold: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._find_slipping_elements()\n" - "instead for finding slipping elements." - ) - - -def node_to_element_mass_or_force(input: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._node_to_element_mass_or_force()\n" - "instead for converting the mass/forces on rod nodes to elements." - ) - - -def nodes_to_elements(input: Any) -> NoReturn: - # Remove the function beyond v0.4.0 - raise NotImplementedError( - "This function is removed in v0.3.1. Please use\n" - "elastica.interaction.node_to_element_mass_or_force()\n" - "instead for node-to-element interpolation of mass/forces." - ) - - -@njit(cache=True) # type: ignore -def elements_to_nodes_inplace( - vector_in_element_frame: Any, vector_in_node_frame: Any -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._elements_to_nodes_inplace()\n" - "instead for updating nodal forces using the forces computed on elements." - ) +from elastica.typing import SystemType, RodType, RigidBodyType # base class for interaction # only applies normal force no friction -class InteractionPlane: +class InteractionPlane(NoForces): """ The interaction plane class computes the plane reaction force on a rod-like object. For more details regarding the contact module refer to @@ -109,9 +72,7 @@ def __init__( self.plane_normal = plane_normal.reshape(3) self.surface_tol = 1e-4 - def apply_normal_force( - self, system: SystemType - ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: + def apply_forces(self, system: RodType, time: float = 0.0) -> None: """ In the case of contact with the plane, this function computes the plane reaction force on the element. @@ -144,33 +105,13 @@ def apply_normal_force( ) -def apply_normal_force_numba( - plane_origin: Any, - plane_normal: Any, - surface_tol: Any, - k: Any, - nu: Any, - radius: Any, - mass: Any, - position_collection: Any, - velocity_collection: Any, - internal_forces: Any, - external_forces: Any, -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. For rod plane contact please use: \n" - "elastica._contact_functions._calculate_contact_forces_rod_plane() \n" - "For detail, refer to issue #113." - ) - - # class for anisotropic frictional plane # NOTE: friction coefficients are passed as arrays in the order # mu_forward : mu_backward : mu_sideways # head is at x[0] and forward means head to tail # same convention for kinetic and static # mu named as to which direction it opposes -class AnisotropicFrictionalPlane(NoForces, InteractionPlane): +class AnisotropicFrictionalPlane(InteractionPlane): """ This anisotropic friction plane class is for computing anisotropic friction forces on rods. @@ -247,12 +188,14 @@ def __init__( # kinetic and static friction should separate functions # for now putting them together to figure out common variables - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces( + self, system: "RodType | RigidBodyType", time: float = 0.0 + ) -> None: """ Call numba implementation to apply friction forces Parameters ---------- - system + system : RodType | RigidBodyType time """ @@ -283,38 +226,6 @@ def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: ) -def anisotropic_friction( - plane_origin: Any, - plane_normal: Any, - surface_tol: Any, - slip_velocity_tol: Any, - k: Any, - nu: Any, - kinetic_mu_forward: Any, - kinetic_mu_backward: Any, - kinetic_mu_sideways: Any, - static_mu_forward: Any, - static_mu_backward: Any, - static_mu_sideways: Any, - radius: Any, - mass: Any, - tangents: Any, - position_collection: Any, - director_collection: Any, - velocity_collection: Any, - omega_collection: Any, - internal_forces: Any, - external_forces: Any, - internal_torques: Any, - external_torques: Any, -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. For anisotropic_friction please use: \n" - "elastica._contact_functions._calculate_contact_forces_rod_plane_with_anisotropic_friction() \n" - "For detail, refer to issue #113." - ) - - # Slender body module @njit(cache=True) # type: ignore def sum_over_elements(input: NDArray[np.floating]) -> np.floating: @@ -349,39 +260,13 @@ def sum_over_elements(input: NDArray[np.floating]) -> np.floating: This version: 513 ns ± 24.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each) """ - output: np.floating = 0.0 + output: np.floating = np.float64(0.0) for i in range(input.shape[0]): output += input[i] return output -def node_to_element_position(node_position_collection: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. For node-to-element_position() interpolation please use: \n" - "elastica.contact_utils._node_to_element_position() for rod position \n" - "For detail, refer to issue #113." - ) - - -def node_to_element_velocity(mass: Any, node_velocity_collection: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. For node-to-element_velocity() interpolation please use: \n" - "elastica.contact_utils._node_to_element_velocity() for rod velocity. \n" - "For detail, refer to issue #113." - ) - - -def node_to_element_pos_or_vel(vector_in_node_frame: Any) -> NoReturn: - # Remove the function beyond v0.4.0 - raise NotImplementedError( - "This function is removed in v0.3.0. For node-to-element interpolation please use: \n" - "elastica.contact_utils._node_to_element_position() for rod position \n" - "elastica.contact_utils._node_to_element_velocity() for rod velocity. \n" - "For detail, refer to issue #80." - ) - - @njit(cache=True) # type: ignore def slender_body_forces( tangents: NDArray[np.floating], @@ -512,7 +397,7 @@ def __init__(self, dynamic_viscosity: np.floating) -> None: super(SlenderBodyTheory, self).__init__() self.dynamic_viscosity = dynamic_viscosity - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces(self, system: RodType, time: float = 0.0) -> None: """ This function applies hydrodynamic forces on body using the slender body theory given in @@ -537,7 +422,7 @@ def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: # base class for interaction # only applies normal force no friction -class InteractionPlaneRigidBody: +class InteractionPlaneRigidBody(NoForces): def __init__( self, k: np.floating, @@ -551,9 +436,7 @@ def __init__( self.plane_normal = plane_normal.reshape(3) self.surface_tol = 1e-4 - def apply_normal_force( - self, system: SystemType - ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: + def apply_forces(self, system: RigidBodyType, time: float = 0.0) -> None: """ This function computes the plane force response on the rigid body, in the case of contact. Contact model given in Eqn 4.8 Gazzola et. al. RSoS 2018 paper @@ -566,7 +449,7 @@ def apply_normal_force( ------- magnitude of the plane response """ - return _calculate_contact_forces_cylinder_plane( + _calculate_contact_forces_cylinder_plane( self.plane_origin, self.plane_normal, self.surface_tol, @@ -577,23 +460,3 @@ def apply_normal_force( system.velocity_collection, system.external_forces, ) - - -@njit(cache=True) # type: ignore -def apply_normal_force_numba_rigid_body( - plane_origin: Any, - plane_normal: Any, - surface_tol: Any, - k: Any, - nu: Any, - length: Any, - position_collection: Any, - velocity_collection: Any, - external_forces: Any, -) -> NoReturn: - - raise NotImplementedError( - "This function is removed in v0.3.2. For cylinder plane contact please use: \n" - "elastica._contact_functions._calculate_contact_forces_cylinder_plane() \n" - "For detail, refer to issue #113." - ) diff --git a/elastica/joint.py b/elastica/joint.py index eb3fd5d18..812a754c9 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -1,11 +1,12 @@ __doc__ = """ Module containing joint classes to connect multiple rods together. """ +__all__ = ["FreeJoint", "HingeJoint", "FixedJoint", "get_relative_rotation_two_systems"] from typing import Optional from elastica._rotations import _inv_rotate -from elastica.typing import SystemType, RodType, ConnectionIndex +from elastica.typing import SystemType, RodType, ConnectionIndex, RigidBodyType + import numpy as np -import logging from numpy.typing import NDArray @@ -47,9 +48,9 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: def apply_forces( self, - system_one: SystemType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: SystemType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: """ @@ -57,11 +58,11 @@ def apply_forces( Parameters ---------- - system_one : SystemType + system_one : RodType | RigidBodyType Rod or rigid-body object index_one : ConnectionIndex Index of first rod for joint. - system_two : SystemType + system_two : RodType | RigidBodyType Rod or rigid-body object index_two : ConnectionIndex Index of second rod for joint. @@ -90,9 +91,9 @@ def apply_forces( def apply_torques( self, - system_one: SystemType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: SystemType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: """ @@ -102,11 +103,11 @@ def apply_torques( Parameters ---------- - system_one : SystemType + system_one : RodType | RigidBodyType Rod or rigid-body object index_one : ConnectionIndex Index of first rod for joint. - system_two : SystemType + system_two : RodType | RigidBodyType Rod or rigid-body object index_two : ConnectionIndex Index of second rod for joint. @@ -171,18 +172,18 @@ def __init__( # Apply force is same as free joint def apply_forces( self, - system_one: SystemType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: SystemType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: return super().apply_forces(system_one, index_one, system_two, index_two) def apply_torques( self, - system_one: SystemType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: SystemType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: # current tangent direction of the `index_two` element of system two @@ -238,7 +239,7 @@ def __init__( k: np.floating, nu: np.floating, kt: np.floating, - nut: np.floating = 0.0, + nut: np.floating = np.float64(0.0), rest_rotation_matrix: Optional[NDArray[np.floating]] = None, ) -> None: """ @@ -278,18 +279,18 @@ def __init__( # Apply force is same as free joint def apply_forces( self, - system_one: SystemType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: SystemType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: return super().apply_forces(system_one, index_one, system_two, index_two) def apply_torques( self, - system_one: SystemType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: SystemType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: # collect directors of systems one and two @@ -335,10 +336,10 @@ def apply_torques( def get_relative_rotation_two_systems( - system_one: SystemType, - index_one: int, - system_two: SystemType, - index_two: int, + system_one: "RodType | RigidBodyType", + index_one: ConnectionIndex, + system_two: "RodType | RigidBodyType", + index_two: ConnectionIndex, ) -> NDArray[np.floating]: """ Compute the relative rotation matrix C_12 between system one and system two at the specified elements. @@ -366,13 +367,13 @@ def get_relative_rotation_two_systems( Parameters ---------- - system_one : SystemType + system_one : RodType | RigidBodyType Rod or rigid-body object - index_one : int + index_one : ConnectionIndex Index of first rod for joint. - system_two : SystemType + system_two : RodType | RigidBodyType Rod or rigid-body object - index_two : int + index_two : ConnectionIndex Index of second rod for joint. Returns @@ -384,230 +385,3 @@ def get_relative_rotation_two_systems( system_one.director_collection[..., index_one] @ system_two.director_collection[..., index_two].T ) - - -class ExternalContact(FreeJoint): - """ - This class is for applying contact forces between rod-cylinder and rod-rod. - If you are want to apply contact forces between rod and cylinder, first system is always rod and second system - is always cylinder. - In addition to the contact forces, user can define apply friction forces between rod and cylinder that - are in contact. For details on friction model refer to this [1]_. - TODO: Currently friction force is between rod-cylinder, in future implement friction forces between rod-rod. - - Notes - ----- - The `velocity_damping_coefficient` is set to a high value (e.g. 1e4) to minimize slip and simulate stiction - (static friction), while friction_coefficient corresponds to the Coulombic friction coefficient. - - Examples - -------- - How to define contact between rod and cylinder. - - >>> simulator.connect(rod, cylinder).using( - ... ExternalContact, - ... k=1e4, - ... nu=10, - ... velocity_damping_coefficient=10, - ... kinetic_friction_coefficient=10, - ... ) - - How to define contact between rod and rod. - - >>> simulator.connect(rod, rod).using( - ... ExternalContact, - ... k=1e4, - ... nu=10, - ... ) - - .. [1] Preclik T., Popa Constantin., Rude U., Regularizing a Time-Stepping Method for Rigid Multibody Dynamics, Multibody Dynamics 2011, ECCOMAS. URL: https://www10.cs.fau.de/publications/papers/2011/Preclik_Multibody_Ext_Abstr.pdf - """ - - # Dev note: - # Most of the cylinder-cylinder contact SHOULD be implemented - # as given in this `paper `, - # but the elastica-cpp kernels are implemented. - # This is maybe to speed-up the kernel, but it's - # potentially dangerous as it does not deal with "end" conditions - # correctly. - - def __init__( - self, - k: np.floating, - nu: np.floating, - velocity_damping_coefficient: np.floating = 0, - friction_coefficient: np.floating = 0, - ) -> None: - """ - - Parameters - ---------- - k : float - Contact spring constant. - nu : float - Contact damping constant. - velocity_damping_coefficient : float - Velocity damping coefficient between rigid-body and rod contact is used to apply friction force in the - slip direction. - friction_coefficient : float - For Coulombic friction coefficient for rigid-body and rod contact. - """ - super().__init__(k, nu) - self.velocity_damping_coefficient = velocity_damping_coefficient - self.friction_coefficient = friction_coefficient - log = logging.getLogger(self.__class__.__name__) - log.warning( - # Remove warning and add error if ExternalContact is used in v0.3.3 - # Remove the option to use ExternalContact, beyond v0.3.3 - "The option to use the ExternalContact joint for the rod-rod and rod-cylinder contact is now deprecated.\n" - "Instead, for rod-rod contact or rod-cylinder contact,use RodRodContact or RodCylinderContact from the add-on Contact mixin class.\n" - "For reference see the classes elastica.contact_forces.RodRodContact() and elastica.contact_forces.RodCylinderContact().\n" - "For usage check examples/RigidbodyCases/RodRigidBodyContact/rod_cylinder_contact.py and examples/RodContactCase/RodRodContact/rod_rod_contact_parallel_validation.py.\n" - " The option to use the ExternalContact joint for the rod-rod and rod-cylinder will be removed in the future (v0.3.3).\n" - ) - - def apply_forces( - self, - rod_one: SystemType, - index_one: ConnectionIndex, - rod_two: SystemType, - index_two: ConnectionIndex, - ) -> None: - # del index_one, index_two - from elastica.contact_utils import ( - _prune_using_aabbs_rod_cylinder, - _prune_using_aabbs_rod_rod, - ) - from elastica._contact_functions import ( - _calculate_contact_forces_rod_cylinder, - _calculate_contact_forces_rod_rod, - ) - - # TODO: raise error during the initialization if rod one is rigid body. - - # If rod two has one element, then it is rigid body. - if rod_two.n_elems == 1: - cylinder_two = rod_two - # First, check for a global AABB bounding box, and see whether that - # intersects - if _prune_using_aabbs_rod_cylinder( - rod_one.position_collection, - rod_one.radius, - rod_one.lengths, - cylinder_two.position_collection, - cylinder_two.director_collection, - cylinder_two.radius[0], - cylinder_two.length[0], - ): - return - - x_cyl = ( - cylinder_two.position_collection[..., 0] - - 0.5 * cylinder_two.length * cylinder_two.director_collection[2, :, 0] - ) - - rod_element_position = 0.5 * ( - rod_one.position_collection[..., 1:] - + rod_one.position_collection[..., :-1] - ) - _calculate_contact_forces_rod_cylinder( - rod_element_position, - rod_one.lengths * rod_one.tangents, - cylinder_two.position_collection[..., 0], - x_cyl, - cylinder_two.length * cylinder_two.director_collection[2, :, 0], - rod_one.radius + cylinder_two.radius, - rod_one.lengths + cylinder_two.length, - rod_one.internal_forces, - rod_one.external_forces, - cylinder_two.external_forces, - cylinder_two.external_torques, - cylinder_two.director_collection[:, :, 0], - rod_one.velocity_collection, - cylinder_two.velocity_collection, - self.k, - self.nu, - self.velocity_damping_coefficient, - self.friction_coefficient, - ) - - else: - # First, check for a global AABB bounding box, and see whether that - # intersects - - if _prune_using_aabbs_rod_rod( - rod_one.position_collection, - rod_one.radius, - rod_one.lengths, - rod_two.position_collection, - rod_two.radius, - rod_two.lengths, - ): - return - - _calculate_contact_forces_rod_rod( - rod_one.position_collection[ - ..., :-1 - ], # Discount last node, we want element start position - rod_one.radius, - rod_one.lengths, - rod_one.tangents, - rod_one.velocity_collection, - rod_one.internal_forces, - rod_one.external_forces, - rod_two.position_collection[ - ..., :-1 - ], # Discount last node, we want element start position - rod_two.radius, - rod_two.lengths, - rod_two.tangents, - rod_two.velocity_collection, - rod_two.internal_forces, - rod_two.external_forces, - self.k, - self.nu, - ) - - -class SelfContact(FreeJoint): - """ - This class is modeling self contact of rod. - - """ - - def __init__(self, k: np.floating, nu: np.floating) -> None: - super().__init__(k, nu) - log = logging.getLogger(self.__class__.__name__) - log.warning( - # Remove warning and add error if SelfContact is used in v0.3.3 - # Remove the option to use SelfContact, beyond v0.3.3 - "The option to use the SelfContact joint for the rod self contact is now deprecated.\n" - "Instead, for rod self contact use RodSelfContact from the add-on Contact mixin class.\n" - "For reference see the class elastica.contact_forces.RodSelfContact(), and for usage check examples/RodContactCase/RodSelfContact/solenoids.py.\n" - "The option to use the SelfContact joint for the rod self contact will be removed in the future (v0.3.3).\n" - ) - - def apply_forces( - self, - rod_one: SystemType, - index_one: ConnectionIndex, - rod_two: SystemType, - index_two: ConnectionIndex, - ) -> None: - # del index_one, index_two - from elastica._contact_functions import ( - _calculate_contact_forces_self_rod, - ) - - _calculate_contact_forces_self_rod( - rod_one.position_collection[ - ..., :-1 - ], # Discount last node, we want element start position - rod_one.radius, - rod_one.lengths, - rod_one.tangents, - rod_one.velocity_collection, - rod_one.external_forces, - self.k, - self.nu, - ) diff --git a/elastica/memory_block/memory_block_rigid_body.py b/elastica/memory_block/memory_block_rigid_body.py index 0ec9d4256..38ebdb8f1 100644 --- a/elastica/memory_block/memory_block_rigid_body.py +++ b/elastica/memory_block/memory_block_rigid_body.py @@ -1,7 +1,7 @@ __doc__ = """Create block-structure class for collection of rigid body systems.""" from typing import Literal import numpy as np -from elastica.typing import SystemType, SystemIdxType +from elastica.typing import SystemIdxType, RigidBodyType from elastica.rigidbody import RigidBodyBase from elastica.rigidbody.data_structures import _RigidRodSymplecticStepperMixin @@ -9,11 +9,11 @@ class MemoryBlockRigidBody(RigidBodyBase, _RigidRodSymplecticStepperMixin): def __init__( - self, systems: list[RigidBodyBase], system_idx_list: list[SystemIdxType] + self, systems: list[RigidBodyType], system_idx_list: list[SystemIdxType] ) -> None: - self.n_bodies = len(systems) - self.n_elems = self.n_bodies + self.n_systems = len(systems) + self.n_elems = self.n_systems self.n_nodes = self.n_elems self.system_idx_list = np.array(system_idx_list, dtype=np.int64) @@ -26,7 +26,7 @@ def __init__( # Initialize the mixin class for symplectic time-stepper. _RigidRodSymplecticStepperMixin.__init__(self) - def _allocate_block_variables_scalars(self, systems: list[RigidBodyBase]) -> None: + def _allocate_block_variables_scalars(self, systems: list[RigidBodyType]) -> None: """ This function takes system collection and allocates the variables for block-structure and references allocated variables back to the systems. @@ -61,7 +61,7 @@ def _allocate_block_variables_scalars(self, systems: list[RigidBodyBase]) -> Non value_type="scalar", ) - def _allocate_block_variables_vectors(self, systems: list[RigidBodyBase]) -> None: + def _allocate_block_variables_vectors(self, systems: list[RigidBodyType]) -> None: """ This function takes system collection and allocates the vector variables for block-structure and references allocated vector variables back to the systems. @@ -97,7 +97,7 @@ def _allocate_block_variables_vectors(self, systems: list[RigidBodyBase]) -> Non value_type="vector", ) - def _allocate_block_variables_matrix(self, systems: list[RigidBodyBase]) -> None: + def _allocate_block_variables_matrix(self, systems: list[RigidBodyType]) -> None: """ This function takes system collection and allocates the matrix variables for block-structure and references allocated matrix variables back to the systems. @@ -134,7 +134,7 @@ def _allocate_block_variables_matrix(self, systems: list[RigidBodyBase]) -> None ) def _allocate_block_variables_for_symplectic_stepper( - self, systems: list[RigidBodyBase] + self, systems: list[RigidBodyType] ) -> None: """ This function takes system collection and allocates the variables used by symplectic @@ -181,7 +181,7 @@ def _allocate_block_variables_for_symplectic_stepper( def _map_system_properties_to_block_memory( self, mapping_dict: dict, - systems: list[RigidBodyBase], + systems: list[RigidBodyType], block_memory: np.ndarray, value_type: Literal["scalar", "vector", "tensor"], ) -> None: @@ -191,7 +191,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[RigidBodyBase] + systems: list[RigidBodyType] A sequence containing rigid body objects to map from. block_memory: ndarray Memory block that, at the end of the method execution, contains all designated diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index 4a1e8bd9f..6cc07a086 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -2,10 +2,6 @@ import numpy as np from typing import Literal, Callable 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, -) from elastica.rod.data_structures import _RodSymplecticStepperMixin from elastica.reset_functions_for_block_structure import _reset_scalar_ghost from elastica.rod.cosserat_rod import ( @@ -18,6 +14,11 @@ _synchronize_periodic_boundary_of_matrix_collection, ) +from .utils import ( + make_block_memory_metadata, + make_block_memory_periodic_boundary_metadata, +) + class MemoryBlockCosseratRod(CosseratRod, _RodSymplecticStepperMixin): """ @@ -32,6 +33,7 @@ class MemoryBlockCosseratRod(CosseratRod, _RodSymplecticStepperMixin): def __init__( self, systems: list[RodType], system_idx_list: list[SystemIdxType] ) -> None: + self.n_systems = len(systems) # separate straight and ring rods system_straight_rod = [] @@ -61,8 +63,8 @@ def __init__( [x.n_elems for x in system_ring_rod], dtype=np.int64 ) - n_straight_rods = len(system_straight_rod) - n_ring_rods = len(system_ring_rod) + n_straight_rods: int = len(system_straight_rod) + n_ring_rods: int = len(system_ring_rod) # self.n_elems_in_rods = np.array([x.n_elems for x in systems], dtype=np.int) self.n_elems_in_rods = np.hstack((n_elems_straight_rods, n_elems_ring_rods + 2)) @@ -132,13 +134,13 @@ def __init__( # we place first straight rods, then ring rods. # TODO: in future consider a better implementation for packing problem. # +1 is because we want to start from next idx, where periodic boundary starts - self.periodic_boundary_nodes_idx += ( + self.periodic_boundary_nodes_idx += ( # type: ignore self.ghost_nodes_idx[n_straight_rods - 1] + 1 ) - self.periodic_boundary_elems_idx += ( + self.periodic_boundary_elems_idx += ( # type: ignore self.ghost_elems_idx[1::2][n_straight_rods - 1] + 1 ) - self.periodic_boundary_voronoi_idx += ( + self.periodic_boundary_voronoi_idx += ( # type: ignore self.ghost_voronoi_idx[2::3][n_straight_rods - 1] + 1 ) diff --git a/elastica/memory_block/memory_block_rod_base.py b/elastica/memory_block/memory_block_rod_base.py index d8b91ff26..87fcdd6d4 100644 --- a/elastica/memory_block/memory_block_rod_base.py +++ b/elastica/memory_block/memory_block_rod_base.py @@ -1,169 +1,8 @@ -__doc__ = """Create block-structure class for collection of Cosserat rod systems.""" +__doc__ = """Deprecated module. Use memory_blocks.utils instead.""" import numpy as np import numpy.typing as npt - -def make_block_memory_metadata( - n_elems_in_rods: npt.NDArray[np.integer], -) -> tuple[ - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], -]: - """ - This function, takes number of elements of each rod as a numpy array and computes, - ghost nodes, elements and voronoi element indexes and numbers and returns it. - - Parameters - ---------- - n_elems_in_rods: npt.NDArray - An integer array containing the number of elements in each of the n rod. - - Returns - ------- - n_elems_with_ghosts: int64 - Total number of elements with ghost elements included. There are two ghost elements - between each pair of two rods adjacent in memory block. - ghost_nodes_idx: ndarray - An integer array of length n - 1 containing the indices of ghost nodes in memory block. - ghost_elements_idx: npt.NDArray - An integer array of length 2 * (n - 1) containing the indices of ghost elements in memory block. - ghost_voronoi_idx: npt.NDArray - An integer array of length 2 * (n - 1) containing the indices of ghost Voronoi nodes in memory block. - """ - - n_nodes_in_rods = n_elems_in_rods + 1 - n_rods = n_elems_in_rods.shape[0] - - # Gap between two rods have one ghost node - # n_nodes_with_ghosts = np.sum(n_nodes_in_rods) + (n_rods - 1) - # Gap between two rods have two ghost elements : comes out to n_nodes_with_ghosts - 1 - n_elems_with_ghosts = np.sum(n_elems_in_rods) + 2 * (n_rods - 1) - # Gap between two rods have three ghost voronois : comes out to n_nodes_with_ghosts - 2 - # n_voronoi_with_ghosts = np.sum(n_voronois_in_rods) + 3 * (n_rods - 1) - - ghost_nodes_idx = np.cumsum(n_nodes_in_rods[:-1], dtype=np.int64) - # Add [0, 1, 2, ... n_rods-2] to the ghost_nodes idx to accommodate miscounting - ghost_nodes_idx += np.arange(0, n_rods - 1, dtype=np.int64) - - ghost_elems_idx = np.zeros((2 * (n_rods - 1),), dtype=np.int64) - ghost_elems_idx[::2] = ghost_nodes_idx - 1 - ghost_elems_idx[1::2] = ghost_nodes_idx.copy() - - ghost_voronoi_idx = np.zeros((3 * (n_rods - 1),), dtype=np.int64) - ghost_voronoi_idx[::3] = ghost_nodes_idx - 2 - ghost_voronoi_idx[1::3] = ghost_nodes_idx - 1 - ghost_voronoi_idx[2::3] = ghost_nodes_idx.copy() - - return n_elems_with_ghosts, ghost_nodes_idx, ghost_elems_idx, ghost_voronoi_idx - - -def make_block_memory_periodic_boundary_metadata( - n_elems_in_rods: npt.NDArray[np.integer], -) -> tuple[ - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], -]: - """ - This function, takes the number of elements of ring rods and computes the periodic boundary node, - element and voronoi index. - - Parameters - ---------- - n_elems_in_rods : npt.NDArray - 1D (n_ring_rods,) array containing data with 'float' type. Elements of this array contains total number of - elements of one rod, including periodic boundary elements. - - Returns - ------- - n_elems - - periodic_boundary_node : npt.NDArray - 2D (2, n_periodic_boundary_nodes) array containing data with 'float' type. Vector containing periodic boundary - elements index. First dimension is the periodic boundary index, second dimension is the referenced cell index. - - periodic_boundary_elems_idx : npt.NDArray - 2D (2, n_periodic_boundary_elems) array containing data with 'float' type. Vector containing periodic boundary - nodes index. First dimension is the periodic boundary index, second dimension is the referenced cell index. - - periodic_boundary_voronoi_idx : npt.NDArray - 2D (2, n_periodic_boundary_voronoi) array containing data with 'float' type. Vector containing periodic boundary - voronoi index. First dimension is the periodic boundary index, second dimension is the referenced cell index. - - """ - - n_elem: npt.NDArray[np.integer] = n_elems_in_rods.copy() - n_rods = n_elems_in_rods.shape[0] - - periodic_boundary_node_idx = np.zeros((2, 3 * n_rods), dtype=np.int64) - # count ghost nodes, first rod does not have a ghost node at the start, so exclude first rod. - periodic_boundary_node_idx[0, 0::3][1:] = 1 - # This is for the first periodic node at the end - periodic_boundary_node_idx[0, 1::3] = 1 + n_elem - # This is for the second periodic node at the end - periodic_boundary_node_idx[0, 2::3] = 1 - periodic_boundary_node_idx[0, :] = np.cumsum(periodic_boundary_node_idx[0, :]) - # Add [0, 1, 2, ..., n_rods] to the periodic boundary nodes to accommodate miscounting - periodic_boundary_node_idx[0, :] += np.repeat( - np.arange(0, n_rods, dtype=np.int64), 3 - ) - # Now fill the reference node idx, to copy and correct periodic boundary nodes - # First fill with the reference node idx of the first periodic node. This is the last node of the actual rod - # (without ghost and periodic nodes). - periodic_boundary_node_idx[1, 0::3] = periodic_boundary_node_idx[0, 1::3] - 1 - # Second fill with the reference node idx of the second periodic node. This is the first node of the actual rod - # (without ghost and periodic nodes). - periodic_boundary_node_idx[1, 1::3] = periodic_boundary_node_idx[0, 0::3] + 1 - # Third fill with the reference node idx of the third periodic node. This is the second node of the actual rod - # (without ghost and periodic nodes). - periodic_boundary_node_idx[1, 2::3] = periodic_boundary_node_idx[0, 0::3] + 2 - - periodic_boundary_elems_idx = np.zeros((2, 2 * n_rods), dtype=np.int64) - # count ghost elems, first rod does not have a ghost elem at the start, so exclude first rod. - periodic_boundary_elems_idx[0, 0::2][1:] = 2 - # This is for the first periodic elem at the end - periodic_boundary_elems_idx[0, 1::2] = 1 + n_elem - periodic_boundary_elems_idx[0, :] = np.cumsum(periodic_boundary_elems_idx[0, :]) - # Add [0, 1, 2, ..., n_rods] to the periodic boundary elems to accommodate miscounting - periodic_boundary_elems_idx[0, :] += np.repeat( - np.arange(0, n_rods, dtype=np.int64), 2 - ) - # Now fill the reference element idx, to copy and correct periodic boundary elements - # First fill with the reference element idx of the first periodic element. This is the last element of the actual - # rod - # (without ghost and periodic elements). - periodic_boundary_elems_idx[1, 0::2] = periodic_boundary_elems_idx[0, 1::2] - 1 - # Second fill with the reference element idx of the second periodic element. This is the first element of the actual - # rod - # (without ghost and periodic elements). - periodic_boundary_elems_idx[1, 1::2] = periodic_boundary_elems_idx[0, 0::2] + 1 - - periodic_boundary_voronoi_idx = np.zeros((2, n_rods), dtype=np.int64) - # count ghost voronoi, first rod does not have a ghost voronoi at the start, so exclude first rod. - periodic_boundary_voronoi_idx[0, 0::1][1:] = 3 - # This is for the first periodic voronoi at the end - periodic_boundary_voronoi_idx[0, 1:] += n_elem[:-1] - periodic_boundary_voronoi_idx[0, :] = np.cumsum(periodic_boundary_voronoi_idx[0, :]) - # Add [0, 1, 2, ..., n_rods] to the periodic boundary voronoi to accommodate miscounting - periodic_boundary_voronoi_idx[0, :] += np.repeat( - np.arange(0, n_rods, dtype=np.int64), 1 - ) - # Now fill the reference voronoi idx, to copy and correct periodic boundary voronoi - # Fill with the reference voronoi idx of the periodic voronoi. This is the last voronoi of the actual rod - # (without ghost and periodic voronoi). - periodic_boundary_voronoi_idx[1, :] = ( - periodic_boundary_voronoi_idx[0, :] + n_elem[:] - ) - - # Increase the n_elem in rods by 2 because we are adding two periodic boundary elements - n_elem = n_elem + 2 - - return ( - n_elem, - periodic_boundary_node_idx, - periodic_boundary_elems_idx, - periodic_boundary_voronoi_idx, - ) +from .utils import ( + make_block_memory_metadata, + make_block_memory_periodic_boundary_metadata, +) diff --git a/elastica/memory_block/protocol.py b/elastica/memory_block/protocol.py index 5cf002462..279f01580 100644 --- a/elastica/memory_block/protocol.py +++ b/elastica/memory_block/protocol.py @@ -1,12 +1,22 @@ from typing import Protocol +from elastica.systems.protocol import SystemProtocol + +import numpy as np from elastica.rod.protocol import CosseratRodProtocol +from elastica.rigidbody.protocol import RigidBodyProtocol from elastica.systems.protocol import SymplecticSystemProtocol +class BlockSystemProtocol(SystemProtocol, Protocol): + @property + def n_systems(self) -> int: + """Number of systems in the block.""" + + class BlockCosseratRodProtocol(CosseratRodProtocol, SymplecticSystemProtocol, Protocol): pass -# class BlockRigidBodyProtocol(RigidBodyProtocol, SymplecticSystemProtocol, Protocol): -# pass +class BlockRigidBodyProtocol(RigidBodyProtocol, SymplecticSystemProtocol, Protocol): + pass diff --git a/elastica/memory_block/utils.py b/elastica/memory_block/utils.py new file mode 100644 index 000000000..c7219e405 --- /dev/null +++ b/elastica/memory_block/utils.py @@ -0,0 +1,169 @@ +__doc__ = """Create block-structure class for collection of Cosserat rod systems.""" +import numpy as np +import numpy.typing as npt + + +def make_block_memory_metadata( + n_elems_in_rods: npt.NDArray[np.integer], +) -> tuple[ + int, + npt.NDArray[np.integer], + npt.NDArray[np.integer], + npt.NDArray[np.integer], +]: + """ + This function, takes number of elements of each rod as a numpy array and computes, + ghost nodes, elements and voronoi element indexes and numbers and returns it. + + Parameters + ---------- + n_elems_in_rods: npt.NDArray + An integer array containing the number of elements in each of the n rod. + + Returns + ------- + n_elems_with_ghosts: int64 + Total number of elements with ghost elements included. There are two ghost elements + between each pair of two rods adjacent in memory block. + ghost_nodes_idx: ndarray + An integer array of length n - 1 containing the indices of ghost nodes in memory block. + ghost_elements_idx: npt.NDArray + An integer array of length 2 * (n - 1) containing the indices of ghost elements in memory block. + ghost_voronoi_idx: npt.NDArray + An integer array of length 2 * (n - 1) containing the indices of ghost Voronoi nodes in memory block. + """ + + n_nodes_in_rods = n_elems_in_rods + 1 + n_rods = n_elems_in_rods.shape[0] + + # Gap between two rods have one ghost node + # n_nodes_with_ghosts = np.sum(n_nodes_in_rods) + (n_rods - 1) + # Gap between two rods have two ghost elements : comes out to n_nodes_with_ghosts - 1 + n_elems_with_ghosts = np.sum(n_elems_in_rods) + 2 * (n_rods - 1) + # Gap between two rods have three ghost voronois : comes out to n_nodes_with_ghosts - 2 + # n_voronoi_with_ghosts = np.sum(n_voronois_in_rods) + 3 * (n_rods - 1) + + ghost_nodes_idx = np.cumsum(n_nodes_in_rods[:-1], dtype=np.int64) + # Add [0, 1, 2, ... n_rods-2] to the ghost_nodes idx to accommodate miscounting + ghost_nodes_idx += np.arange(0, n_rods - 1, dtype=np.int64) + + ghost_elems_idx = np.zeros((2 * (n_rods - 1),), dtype=np.int64) + ghost_elems_idx[::2] = ghost_nodes_idx - 1 + ghost_elems_idx[1::2] = ghost_nodes_idx.copy() + + ghost_voronoi_idx = np.zeros((3 * (n_rods - 1),), dtype=np.int64) + ghost_voronoi_idx[::3] = ghost_nodes_idx - 2 + ghost_voronoi_idx[1::3] = ghost_nodes_idx - 1 + ghost_voronoi_idx[2::3] = ghost_nodes_idx.copy() + + return n_elems_with_ghosts, ghost_nodes_idx, ghost_elems_idx, ghost_voronoi_idx + + +def make_block_memory_periodic_boundary_metadata( + n_elems_in_rods: npt.NDArray[np.integer], +) -> tuple[ + npt.NDArray[np.integer], + npt.NDArray[np.integer], + npt.NDArray[np.integer], + npt.NDArray[np.integer], +]: + """ + This function, takes the number of elements of ring rods and computes the periodic boundary node, + element and voronoi index. + + Parameters + ---------- + n_elems_in_rods : npt.NDArray + 1D (n_ring_rods,) array containing data with 'int' type. Elements of this array contains total number of + elements of one rod, including periodic boundary elements. + + Returns + ------- + n_elems + + periodic_boundary_node : npt.NDArray + 2D (2, n_periodic_boundary_nodes) array containing data with 'int' type. Vector containing periodic boundary + elements index. First dimension is the periodic boundary index, second dimension is the referenced cell index. + + periodic_boundary_elems_idx : npt.NDArray + 2D (2, n_periodic_boundary_elems) array containing data with 'int' type. Vector containing periodic boundary + nodes index. First dimension is the periodic boundary index, second dimension is the referenced cell index. + + periodic_boundary_voronoi_idx : npt.NDArray + 2D (2, n_periodic_boundary_voronoi) array containing data with 'int' type. Vector containing periodic boundary + voronoi index. First dimension is the periodic boundary index, second dimension is the referenced cell index. + + """ + + n_elem: npt.NDArray[np.integer] = n_elems_in_rods.copy() + n_rods = n_elems_in_rods.shape[0] + + periodic_boundary_node_idx = np.zeros((2, 3 * n_rods), dtype=np.int64) + # count ghost nodes, first rod does not have a ghost node at the start, so exclude first rod. + periodic_boundary_node_idx[0, 0::3][1:] = 1 + # This is for the first periodic node at the end + periodic_boundary_node_idx[0, 1::3] = 1 + n_elem + # This is for the second periodic node at the end + periodic_boundary_node_idx[0, 2::3] = 1 + periodic_boundary_node_idx[0, :] = np.cumsum(periodic_boundary_node_idx[0, :]) + # Add [0, 1, 2, ..., n_rods] to the periodic boundary nodes to accommodate miscounting + periodic_boundary_node_idx[0, :] += np.repeat( + np.arange(0, n_rods, dtype=np.int64), 3 + ) + # Now fill the reference node idx, to copy and correct periodic boundary nodes + # First fill with the reference node idx of the first periodic node. This is the last node of the actual rod + # (without ghost and periodic nodes). + periodic_boundary_node_idx[1, 0::3] = periodic_boundary_node_idx[0, 1::3] - 1 + # Second fill with the reference node idx of the second periodic node. This is the first node of the actual rod + # (without ghost and periodic nodes). + periodic_boundary_node_idx[1, 1::3] = periodic_boundary_node_idx[0, 0::3] + 1 + # Third fill with the reference node idx of the third periodic node. This is the second node of the actual rod + # (without ghost and periodic nodes). + periodic_boundary_node_idx[1, 2::3] = periodic_boundary_node_idx[0, 0::3] + 2 + + periodic_boundary_elems_idx = np.zeros((2, 2 * n_rods), dtype=np.int64) + # count ghost elems, first rod does not have a ghost elem at the start, so exclude first rod. + periodic_boundary_elems_idx[0, 0::2][1:] = 2 + # This is for the first periodic elem at the end + periodic_boundary_elems_idx[0, 1::2] = 1 + n_elem + periodic_boundary_elems_idx[0, :] = np.cumsum(periodic_boundary_elems_idx[0, :]) + # Add [0, 1, 2, ..., n_rods] to the periodic boundary elems to accommodate miscounting + periodic_boundary_elems_idx[0, :] += np.repeat( + np.arange(0, n_rods, dtype=np.int64), 2 + ) + # Now fill the reference element idx, to copy and correct periodic boundary elements + # First fill with the reference element idx of the first periodic element. This is the last element of the actual + # rod + # (without ghost and periodic elements). + periodic_boundary_elems_idx[1, 0::2] = periodic_boundary_elems_idx[0, 1::2] - 1 + # Second fill with the reference element idx of the second periodic element. This is the first element of the actual + # rod + # (without ghost and periodic elements). + periodic_boundary_elems_idx[1, 1::2] = periodic_boundary_elems_idx[0, 0::2] + 1 + + periodic_boundary_voronoi_idx = np.zeros((2, n_rods), dtype=np.int64) + # count ghost voronoi, first rod does not have a ghost voronoi at the start, so exclude first rod. + periodic_boundary_voronoi_idx[0, 0::1][1:] = 3 + # This is for the first periodic voronoi at the end + periodic_boundary_voronoi_idx[0, 1:] += n_elem[:-1] + periodic_boundary_voronoi_idx[0, :] = np.cumsum(periodic_boundary_voronoi_idx[0, :]) + # Add [0, 1, 2, ..., n_rods] to the periodic boundary voronoi to accommodate miscounting + periodic_boundary_voronoi_idx[0, :] += np.repeat( + np.arange(0, n_rods, dtype=np.int64), 1 + ) + # Now fill the reference voronoi idx, to copy and correct periodic boundary voronoi + # Fill with the reference voronoi idx of the periodic voronoi. This is the last voronoi of the actual rod + # (without ghost and periodic voronoi). + periodic_boundary_voronoi_idx[1, :] = ( + periodic_boundary_voronoi_idx[0, :] + n_elem[:] + ) + + # Increase the n_elem in rods by 2 because we are adding two periodic boundary elements + n_elem = n_elem + 2 + + return ( + n_elem, + periodic_boundary_node_idx, + periodic_boundary_elems_idx, + periodic_boundary_voronoi_idx, + ) diff --git a/elastica/mesh/protocol.py b/elastica/mesh/protocol.py new file mode 100644 index 000000000..766651fd5 --- /dev/null +++ b/elastica/mesh/protocol.py @@ -0,0 +1,10 @@ +from typing import Protocol + +import numpy as np +from numpy.typing import NDArray + + +class MeshProtocol(Protocol): + faces: NDArray[np.floating] + face_centers: NDArray[np.floating] + face_normals: NDArray[np.floating] diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index fd999665f..c31581d66 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -5,7 +5,7 @@ Basic coordinating for multiple, smaller systems that have an independently integrable interface (i.e. works with symplectic or explicit routines `timestepper.py`.) """ -from typing import Type, Generator, Iterable, Any +from typing import Type, Generator, Iterable, Any, overload from typing import final from elastica.typing import ( SystemType, @@ -13,7 +13,7 @@ OperatorType, OperatorCallbackType, OperatorFinalizeType, - AllowedContactType, + BlockType, ) import numpy as np @@ -69,7 +69,7 @@ def __init__(self) -> None: # List of systems to be integrated self._systems: list[SystemType] = [] - self._memory_blocks: list[SystemType] = [] + self.__final_systems: list[SystemType] = [] # Flag Finalize: Finalizing twice will cause an error, # but the error message is very misleading @@ -101,6 +101,12 @@ def _check_type(self, sys_to_be_added: Any) -> bool: def __len__(self) -> int: return len(self._systems) + @overload + def __getitem__(self, idx: int, /) -> SystemType: ... + + @overload + def __getitem__(self, idx: slice, /) -> list[SystemType]: ... + def __getitem__(self, idx, /): # type: ignore return self._systems[idx] @@ -159,9 +165,9 @@ def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: return sys_idx @final - def blocks(self) -> Generator[SystemType, None, None]: + def systems(self) -> Generator[SystemType, None, None]: # assert self._finalize_flag, "The simulator is not finalized." - for block in self._memory_blocks: + for block in self.__final_systems: yield block @final @@ -177,10 +183,10 @@ def finalize(self) -> None: self._finalize_flag = True # construct memory block - self._memory_blocks = construct_memory_block_structures(self._systems) - for block in self.blocks(): - # append the memory block to the simulation as a system. Memory block is the final system in the simulation. - self.append(block) + self.__final_systems = construct_memory_block_structures(self._systems) + # TODO: try to remove the _systems list for memory optimization + # self._systems.clear() + # del self._systems # Recurrent call finalize functions for all components. for finalize in self._feature_group_finalize: diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index a9ffb19d8..71b15834d 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -10,8 +10,9 @@ from elastica.typing import ( SystemIdxType, OperatorFinalizeType, - SystemType, ConnectionIndex, + RodType, + RigidBodyType, ) import numpy as np import functools @@ -39,8 +40,8 @@ def __init__(self: SystemCollectionProtocol) -> None: def connect( self: SystemCollectionProtocol, - first_rod: SystemType, - second_rod: SystemType, + first_rod: "RodType | RigidBodyType", + second_rod: "RodType | RigidBodyType", first_connect_idx: ConnectionIndex = None, second_connect_idx: ConnectionIndex = None, ) -> ModuleProtocol: @@ -51,9 +52,9 @@ def connect( Parameters ---------- - first_rod : SystemType + first_rod : RodType | RigidBodyType Rod-like object - second_rod : SystemType + second_rod : RodType | RigidBodyType Rod-like object first_connect_idx : Optional[int] Index of first rod for joint. @@ -64,12 +65,11 @@ def connect( ------- """ + # For each system identified, get max dofs sys_idx_first = self._get_sys_idx_if_valid(first_rod) sys_idx_second = self._get_sys_idx_if_valid(second_rod) - - # For each system identified, get max dofs - sys_dofs_first = self._systems[sys_idx_first].n_elems - sys_dofs_second = self._systems[sys_idx_second].n_elems + sys_dofs_first = first_rod.n_elems + sys_dofs_second = second_rod.n_elems # Create _Connect object, cache it and return to user _connect: ModuleProtocol = _Connect( @@ -90,9 +90,9 @@ def _finalize_connections(self: SystemCollectionProtocol) -> None: def apply_forces_and_torques( time: np.floating, connect_instance: FreeJoint, - system_one: SystemType, + system_one: "RodType | RigidBodyType", first_connect_idx: ConnectionIndex, - system_two: SystemType, + system_two: "RodType | RigidBodyType", second_connect_idx: ConnectionIndex, ) -> None: connect_instance.apply_forces( diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 9c70e640e..bfa452124 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -11,7 +11,7 @@ from elastica.boundary_conditions import ConstraintBase -from elastica.typing import SystemType, SystemIdxType, ConstrainingIndex +from elastica.typing import SystemIdxType, ConstrainingIndex, RigidBodyType, RodType from .protocol import SystemCollectionProtocol, ModuleProtocol @@ -34,7 +34,9 @@ def __init__(self: SystemCollectionProtocol) -> None: self._feature_group_constrain_rates.append(self._constrain_rates) self._feature_group_finalize.append(self._finalize_constraints) - def constrain(self: SystemCollectionProtocol, system: SystemType) -> ModuleProtocol: + def constrain( + self: SystemCollectionProtocol, system: "RodType | RigidBodyType" + ) -> ModuleProtocol: """ This method enforces a displacement boundary conditions to the relevant user-defined system or rod-like object. You must input the system or rod-like @@ -65,7 +67,7 @@ def _finalize_constraints(self: SystemCollectionProtocol) -> None: """ from elastica._synchronize_periodic_boundary import _ConstrainPeriodicBoundaries - for block in self.blocks(): + for block in self.systems(): # append the memory block to the simulation as a system. Memory block is the final system in the simulation. if hasattr(block, "ring_rod_flag"): # Apply the constrain to synchronize the periodic boundaries of the memory rod. Find the memory block @@ -178,7 +180,7 @@ def using( def id(self) -> SystemIdxType: return self._sys_idx - def instantiate(self, system: SystemType) -> ConstraintBase: + def instantiate(self, system: "RodType | RigidBodyType") -> ConstraintBase: """Constructs a constraint after checks""" if not hasattr(self, "_bc_cls"): raise RuntimeError( diff --git a/elastica/modules/contact.py b/elastica/modules/contact.py index 924e7a4a9..69cb11db0 100644 --- a/elastica/modules/contact.py +++ b/elastica/modules/contact.py @@ -11,7 +11,6 @@ SystemIdxType, OperatorFinalizeType, SystemType, - AllowedContactType, ) from .protocol import SystemCollectionProtocol, ModuleProtocol @@ -48,7 +47,7 @@ def __init__(self: SystemCollectionProtocol) -> None: def detect_contact_between( self: SystemCollectionProtocol, first_system: SystemType, - second_system: AllowedContactType, + second_system: SystemType, ) -> ModuleProtocol: """ This method adds contact detection between two objects using the selected contact class. @@ -57,9 +56,7 @@ def detect_contact_between( Parameters ---------- first_system : SystemType - Rod or rigid body object - second_system : AllowedContactType - Rod, rigid body or surface object + second_system : SystemType Returns ------- diff --git a/elastica/modules/damping.py b/elastica/modules/damping.py index 4fa0cf055..a3a2615e5 100644 --- a/elastica/modules/damping.py +++ b/elastica/modules/damping.py @@ -15,7 +15,7 @@ import numpy as np from elastica.dissipation import DamperBase -from elastica.typing import SystemType, SystemIdxType +from elastica.typing import RodType, SystemType, SystemIdxType from .protocol import SystemCollectionProtocol, ModuleProtocol @@ -37,7 +37,7 @@ def __init__(self: SystemCollectionProtocol) -> None: self._feature_group_constrain_rates.append(self._dampen_rates) self._feature_group_finalize.append(self._finalize_dampers) - def dampen(self: SystemCollectionProtocol, system: SystemType) -> ModuleProtocol: + def dampen(self: SystemCollectionProtocol, system: RodType) -> ModuleProtocol: """ This method applies damping on relevant user-defined system or rod-like object. You must input the system or rod-like diff --git a/elastica/modules/memory_block.py b/elastica/modules/memory_block.py index f13760d88..3f69b57e0 100644 --- a/elastica/modules/memory_block.py +++ b/elastica/modules/memory_block.py @@ -2,8 +2,14 @@ This function is a module to construct memory blocks for different types of systems, such as Cosserat Rods, Rigid Body etc. """ - -from elastica.typing import SystemType, SystemIdxType +from typing import cast +from elastica.typing import ( + RodType, + RigidBodyType, + SystemType, + SystemIdxType, + BlockSystemType, +) from elastica.rod.rod_base import RodBase from elastica.rigidbody.rigid_body import RigidBodyBase @@ -22,20 +28,22 @@ def construct_memory_block_structures(systems: list[SystemType]) -> list[SystemT ------- """ - _memory_blocks: list[SystemType] = [] - temp_list_for_cosserat_rod_systems: list[RodBase] = [] - temp_list_for_rigid_body_systems: list[RigidBodyBase] = [] + _memory_blocks: list[BlockSystemType] = [] + temp_list_for_cosserat_rod_systems: list[RodType] = [] + temp_list_for_rigid_body_systems: list[RigidBodyType] = [] temp_list_for_cosserat_rod_systems_idx: list[SystemIdxType] = [] temp_list_for_rigid_body_systems_idx: list[SystemIdxType] = [] for system_idx, sys_to_be_added in enumerate(systems): if isinstance(sys_to_be_added, RodBase): - temp_list_for_cosserat_rod_systems.append(sys_to_be_added) + rod_system = cast(RodType, sys_to_be_added) + temp_list_for_cosserat_rod_systems.append(rod_system) temp_list_for_cosserat_rod_systems_idx.append(system_idx) elif isinstance(sys_to_be_added, RigidBodyBase): - temp_list_for_rigid_body_systems.append(sys_to_be_added) + rigid_body_system = cast(RigidBodyType, sys_to_be_added) + temp_list_for_rigid_body_systems.append(rigid_body_system) temp_list_for_rigid_body_systems_idx.append(system_idx) elif isinstance(sys_to_be_added, SurfaceBase): @@ -72,4 +80,4 @@ def construct_memory_block_structures(systems: list[SystemType]) -> list[SystemT ) ) - return _memory_blocks + return list(_memory_blocks) diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 5ce5ba7cf..b65419158 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,14 +1,16 @@ -from typing import Protocol, Generator, TypeVar, Any, Type +from typing import Protocol, Generator, TypeVar, Any, Type, overload from typing_extensions import Self # 3.11: from typing import Self + from abc import abstractmethod + from elastica.typing import ( SystemIdxType, OperatorType, OperatorCallbackType, OperatorFinalizeType, SystemType, - AllowedContactType, ConnectionIndex, + BlockType, ) from elastica.joint import FreeJoint from elastica.callback_functions import CallBackBaseClass @@ -34,6 +36,14 @@ def id(self) -> Any: ... class SystemCollectionProtocol(Protocol): _systems: list[SystemType] + def __len__(self) -> int: ... + + @overload + def __getitem__(self, i: slice) -> list[SystemType]: ... + @overload + def __getitem__(self, i: int) -> SystemType: ... + def __getitem__(self, i: slice | int) -> "list[SystemType] | SystemType": ... + @property def _feature_group_synchronize(self) -> OperatorGroupFIFO: ... @@ -57,11 +67,9 @@ def apply_callbacks(self, time: np.floating, current_step: int) -> None: ... @property def _feature_group_finalize(self) -> list[OperatorFinalizeType]: ... - def blocks(self) -> Generator[SystemType, None, None]: ... + def systems(self) -> Generator[SystemType, None, None]: ... - def _get_sys_idx_if_valid( - self, sys_to_be_added: SystemType | AllowedContactType - ) -> SystemIdxType: ... + def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ... # Connection API _finalize_connections: OperatorFinalizeType @@ -123,7 +131,7 @@ def add_forcing_to(self, system: SystemType) -> ModuleProtocol: @abstractmethod def detect_contact_between( - self, first_system: SystemType, second_system: AllowedContactType + self, first_system: SystemType, second_system: SystemType ) -> ModuleProtocol: raise NotImplementedError diff --git a/elastica/restart.py b/elastica/restart.py index 643e0ed03..cd31d520f 100644 --- a/elastica/restart.py +++ b/elastica/restart.py @@ -30,7 +30,7 @@ def all_equal(iterable: Iterable[Any]) -> bool: def save_state( simulator: Iterable, directory: str = "", - time: np.floating = 0.0, + time: np.floating = np.float64(0.0), verbose: bool = False, ) -> None: """ diff --git a/elastica/rigidbody/cylinder.py b/elastica/rigidbody/cylinder.py index 079ce7d53..053dfdee3 100644 --- a/elastica/rigidbody/cylinder.py +++ b/elastica/rigidbody/cylinder.py @@ -72,8 +72,8 @@ def _check_lower_bound( dim: int_t = MaxDimension.value() # This is for a rigid body cylinder - self.volume = np.pi * base_radius * base_radius * base_length - self.mass = np.array([self.volume * self.density], dtype=np.float64) + self.volume = np.float64(np.pi * base_radius * base_radius * base_length) + self.mass = np.float64(self.volume * self.density) # Second moment of inertia area = np.pi * base_radius * base_radius diff --git a/elastica/rigidbody/data_structures.py b/elastica/rigidbody/data_structures.py index d6edb4f22..95944f759 100644 --- a/elastica/rigidbody/data_structures.py +++ b/elastica/rigidbody/data_structures.py @@ -1,7 +1,8 @@ __doc__ = "Data structure wrapper for rod components" from elastica.rod.data_structures import _RodSymplecticStepperMixin -from typing import Any + +pass """ # FIXME : Explicit Stepper doesn't work as States lose the @@ -40,14 +41,5 @@ def __call__(self, time, *args, **kwargs): return self.__deriv_state """ - -class _RigidRodSymplecticStepperMixin(_RodSymplecticStepperMixin): - def __init__(self) -> None: - super(_RigidRodSymplecticStepperMixin, self).__init__() - # Expose rate returning functions in the interface - # to be used by the time-stepping algorithm - # dynamic rates needs to call update_accelerations and henc - # is another function - - def update_internal_forces_and_torques(self, *args: Any, **kwargs: Any) -> None: - pass +# TODO: Temporary solution as the structure for RigidBody is similar to Rod +_RigidRodSymplecticStepperMixin = _RodSymplecticStepperMixin diff --git a/elastica/rigidbody/mesh_rigid_body.py b/elastica/rigidbody/mesh_rigid_body.py index 4a026f9a8..f72925b01 100644 --- a/elastica/rigidbody/mesh_rigid_body.py +++ b/elastica/rigidbody/mesh_rigid_body.py @@ -1,5 +1,8 @@ __doc__ = """rigid body class based on mesh""" +from numpy.typing import NDArray +from elastica.typing import MeshType + import numpy as np import numba from elastica._linalg import _batch_cross, _batch_norm @@ -10,12 +13,12 @@ class MeshRigidBody(RigidBodyBase): def __init__( self, - mesh, - center_of_mass, - mass_second_moment_of_inertia, - density, - volume, - ): + mesh: MeshType, + center_of_mass: NDArray[np.floating], + mass_second_moment_of_inertia: NDArray[np.floating], + density: np.floating, + volume: np.floating, + ) -> None: """ Mesh rigid body initializer. @@ -37,11 +40,11 @@ def __init__( """ # rigid body does not have elements it only have one node. We are setting n_elems to # zero for only make code to work. _bootstrap_from_data requires n_elems to be defined - self.n_elems = 1 # center_mass + self.n_elems: int = 1 # center_mass self.density = density self.volume = volume - self.mass = np.array([self.volume * self.density]) + self.mass = np.float64(self.volume * self.density) self.mass_second_moment_of_inertia = mass_second_moment_of_inertia.reshape( MaxDimension.value(), MaxDimension.value(), 1 ) @@ -112,7 +115,7 @@ def __init__( MaxDimension.value(), 1 ) - def update_faces(self): + def update_faces(self) -> None: _update_faces( self.director_collection, self.face_centers, @@ -128,7 +131,7 @@ def update_faces(self): ) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _update_faces( director_collection, face_centers, diff --git a/elastica/rigidbody/protocol.py b/elastica/rigidbody/protocol.py new file mode 100644 index 000000000..22bc049d9 --- /dev/null +++ b/elastica/rigidbody/protocol.py @@ -0,0 +1,18 @@ +from typing import Protocol + +import numpy as np +from numpy.typing import NDArray + +from elastica.systems.protocol import SystemProtocol, SlenderBodyGeometryProtocol + + +class RigidBodyProtocol(SystemProtocol, SlenderBodyGeometryProtocol, Protocol): + + mass: np.floating + volume: np.floating + length: np.floating + tangents: NDArray[np.floating] + radius: np.floating + + mass_second_moment_of_inertia: NDArray[np.floating] + inv_mass_second_moment_of_inertia: NDArray[np.floating] diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index 7ce554730..900d9978d 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -34,7 +34,14 @@ def __init__(self) -> None: self.external_forces: f_arr_t self.external_torques: f_arr_t - self.mass: f_arr_t + self.internal_forces: f_arr_t + self.internal_torques: f_arr_t + + self.mass: np.floating + self.volume: np.floating + self.length: np.floating + self.tangents: f_arr_t + self.radius: np.floating self.mass_second_moment_of_inertia: f_arr_t self.inv_mass_second_moment_of_inertia: f_arr_t @@ -66,6 +73,12 @@ def zeroed_out_external_forces_and_torques(self, time: float_t) -> None: self.external_forces *= 0.0 self.external_torques *= 0.0 + def compute_internal_forces_and_torques(self, time: np.floating) -> None: + """ + For rigid body, there is no internal forces and torques + """ + pass + def compute_position_center_of_mass(self) -> f_arr_t: """ Return positional center of mass diff --git a/elastica/rigidbody/sphere.py b/elastica/rigidbody/sphere.py index 92c68ac96..675571a55 100644 --- a/elastica/rigidbody/sphere.py +++ b/elastica/rigidbody/sphere.py @@ -37,10 +37,10 @@ def __init__(self, center: f_arr_t, base_radius: float_t, density: float_t) -> N self.radius = base_radius self.density = density - self.length = 2 * base_radius + self.length = np.float64(2 * base_radius) # This is for a rigid body cylinder - self.volume = 4.0 / 3.0 * np.pi * base_radius**3 - self.mass = np.array([self.volume * self.density], dtype=np.float64) + self.volume = np.float64(4.0 / 3.0 * np.pi * base_radius**3) + self.mass = np.float64(self.volume * self.density) normal = np.array([1.0, 0.0, 0.0], dtype=np.float64).reshape(dim, 1) tangents = np.array([0.0, 0.0, 1.0], dtype=np.float64).reshape(dim, 1) binormal = _batch_cross(tangents, normal) diff --git a/elastica/rod/data_structures.py b/elastica/rod/data_structures.py index dcbbee20c..46588785b 100644 --- a/elastica/rod/data_structures.py +++ b/elastica/rod/data_structures.py @@ -9,9 +9,9 @@ from elastica._linalg import _batch_matmul if TYPE_CHECKING: - from elastica.memory_block.protocol import BlockCosseratRodProtocol + from elastica.systems.protocol import SymplecticSystemProtocol else: - BlockCosseratRodProtocol = "BlockCosseratRodProtocol" + SymplecticSystemProtocol = "SymplecticSystemProtocol" # FIXME : Explicit Stepper doesn't work as States lose the # views they initially had when working with a timestepper. @@ -50,7 +50,7 @@ class _RodSymplecticStepperMixin: - def __init__(self: BlockCosseratRodProtocol) -> None: + def __init__(self: SymplecticSystemProtocol) -> None: self.kinematic_states = _KinematicState( self.position_collection, self.director_collection ) @@ -67,24 +67,14 @@ def __init__(self: BlockCosseratRodProtocol) -> None: # is another function self.kinematic_rates = self.dynamic_states.kinematic_rates - def update_internal_forces_and_torques( - self: BlockCosseratRodProtocol, time: np.floating - ) -> None: - self.compute_internal_forces_and_torques(time) - def dynamic_rates( - self: BlockCosseratRodProtocol, + self: SymplecticSystemProtocol, time: np.floating, prefac: np.floating, ) -> NDArray[np.floating]: self.update_accelerations(time) return self.dynamic_states.dynamic_rates(time, prefac) - def reset_external_forces_and_torques( - self: BlockCosseratRodProtocol, time: np.floating - ) -> None: - self.zeroed_out_external_forces_and_torques(time) - def _bootstrap_from_data( stepper_type: str, diff --git a/elastica/rod/knot_theory.py b/elastica/rod/knot_theory.py index 404131e35..ac440d140 100644 --- a/elastica/rod/knot_theory.py +++ b/elastica/rod/knot_theory.py @@ -10,11 +10,9 @@ The details discussion is included in `N Charles et. al. PRL (2019) `_. """ -pass - -from numba import njit import numpy as np from numpy.typing import NDArray +from numba import njit from elastica.rod.rod_base import RodBase from elastica._linalg import _batch_norm, _batch_dot, _batch_cross diff --git a/elastica/rod/protocol.py b/elastica/rod/protocol.py index ba83bb225..62d33fe57 100644 --- a/elastica/rod/protocol.py +++ b/elastica/rod/protocol.py @@ -3,9 +3,7 @@ import numpy as np from numpy.typing import NDArray -pass - -from elastica.systems.protocol import SystemProtocol +from elastica.systems.protocol import SystemProtocol, SlenderBodyGeometryProtocol class _RodEnergy(Protocol): @@ -14,7 +12,9 @@ def compute_bending_energy(self) -> NDArray[np.float64]: ... def compute_shear_energy(self) -> NDArray[np.float64]: ... -class CosseratRodProtocol(SystemProtocol, _RodEnergy, Protocol): +class CosseratRodProtocol( + SystemProtocol, SlenderBodyGeometryProtocol, _RodEnergy, Protocol +): mass: NDArray[np.floating] volume: NDArray[np.floating] @@ -44,3 +44,6 @@ class CosseratRodProtocol(SystemProtocol, _RodEnergy, Protocol): ghost_elems_idx: NDArray[np.integer] ring_rod_flag: bool + periodic_boundary_nodes_idx: NDArray[np.integer] + periodic_boundary_elems_idx: NDArray[np.integer] + periodic_boundary_voronoi_idx: NDArray[np.integer] diff --git a/elastica/rod/rod_base.py b/elastica/rod/rod_base.py index 76e262f15..11dc9211f 100644 --- a/elastica/rod/rod_base.py +++ b/elastica/rod/rod_base.py @@ -32,3 +32,7 @@ def __init__(self) -> None: self.ghost_voronoi_idx: NDArray[np.integer] self.ghost_elems_idx: NDArray[np.integer] + + self.periodic_boundary_nodes_idx: NDArray[np.integer] + self.periodic_boundary_elems_idx: NDArray[np.integer] + self.periodic_boundary_voronoi_idx: NDArray[np.integer] diff --git a/elastica/surface/plane.py b/elastica/surface/plane.py index 82edd1a64..7f17537f0 100644 --- a/elastica/surface/plane.py +++ b/elastica/surface/plane.py @@ -23,7 +23,7 @@ def __init__(self, plane_origin: np.ndarray, plane_normal: np.ndarray): assert np.allclose( np.linalg.norm(plane_normal), 1, - atol=Tolerance.atol(), + atol=float(Tolerance.atol()), ), "plane normal is not a unit vector" self.normal = np.asarray(plane_normal).reshape(3) self.origin = np.asarray(plane_origin).reshape(3, 1) diff --git a/elastica/surface/surface_base.py b/elastica/surface/surface_base.py index 421cb1304..c5c7a2ff7 100644 --- a/elastica/surface/surface_base.py +++ b/elastica/surface/surface_base.py @@ -1,6 +1,9 @@ __doc__ = """Base class for surfaces""" from typing import Type +import numpy as np +from numpy.typing import NDArray + class SurfaceBase: """ @@ -18,4 +21,5 @@ def __init__(self) -> None: """ SurfaceBase does not take any arguments. """ - pass + self.normal: NDArray[np.floating] # (3,) + self.origin: NDArray[np.floating] # (3, 1) diff --git a/elastica/systems/__init__.py b/elastica/systems/__init__.py index 246629c39..b61b905ed 100644 --- a/elastica/systems/__init__.py +++ b/elastica/systems/__init__.py @@ -1,11 +1,9 @@ -from typing import Iterator, TypeVar, Generic, Type -from elastica.timestepper.protocol import ExplicitStepperProtocol -from elastica.typing import SystemCollectionType +from typing import Type -from copy import copy +from elastica.typing import SystemType, SystemCollectionType -def is_system_a_collection(system: object) -> bool: +def is_system_a_collection(system: "SystemType | SystemCollectionType") -> bool: # Check if system is a "collection" of smaller systems # by checking for the [] method """ @@ -40,87 +38,4 @@ def is_system_a_collection(system: object) -> bool: from elastica.modules import BaseSystemCollection __sys_get_item = getattr(system, "__getitem__", None) - return issubclass(system.__class__, BaseSystemCollection) or callable( - __sys_get_item - ) - - -# FIXME: Move memory related functions to separate module or as part of the timestepper - - -# TODO: Use MemoryProtocol -def make_memory_for_explicit_stepper( - stepper: ExplicitStepperProtocol, system: SystemCollectionType -) -> "MemoryCollection": - # TODO Automated logic (class creation, memory management logic) agnostic of stepper details (RK, AB etc.) - - from elastica.timestepper.explicit_steppers import ( - RungeKutta4, - EulerForward, - ) - - # is_this_system_a_collection = is_system_a_collection(system) - - memory_cls: Type - if RungeKutta4 in stepper.__class__.mro(): - # Bad way of doing it, introduces tight coupling - # this should rather be taken from the class itself - class MemoryRungeKutta4: - def __init__(self) -> None: - self.initial_state = None - self.k_1 = None - self.k_2 = None - self.k_3 = None - self.k_4 = None - - memory_cls = MemoryRungeKutta4 - elif EulerForward in stepper.__class__.mro(): - - class MemoryEulerForward: - def __init__(self) -> None: - self.initial_state = None - self.k = None - - memory_cls = MemoryEulerForward - else: - raise NotImplementedError("Making memory for other types not supported") - - return MemoryCollection(memory_cls(), len(system)) - - -M = TypeVar("M", bound="MemoryCollection") - - -class MemoryCollection(Generic[M]): - """Slots of memories for timestepper in a cohesive unit. - - A `MemoryCollection` object is meant to be used in conjunction - with a `SystemCollection`, where each independent `System` to - be integrated has its own `Memory`. - - Example - ------- - - A RK4 integrator needs to store k_1, k_2, k_3, k_4 (intermediate - results from four stages) for each `System`. The restriction for - having a memory slot arises because the `Systems` are usually - not independent of one another and may need communication after - every stage. - """ - - def __init__(self, memory: M, n_memory_slots: int): - super(MemoryCollection, self).__init__() - - self.__memories: list[M] = [] - for _ in range(n_memory_slots - 1): - self.__memories.append(copy(memory)) - self.__memories.append(memory) - - def __getitem__(self, idx: int) -> M: - return self.__memories[idx] - - def __len__(self) -> int: - return len(self.__memories) - - def __iter__(self) -> Iterator[M]: - return self.__memories.__iter__() + return isinstance(system, BaseSystemCollection) or callable(__sys_get_item) diff --git a/elastica/systems/analytical.py b/elastica/systems/analytical.py index b03a917aa..317b3f99e 100644 --- a/elastica/systems/analytical.py +++ b/elastica/systems/analytical.py @@ -162,10 +162,10 @@ def energy(st): current_energy = energy(self._state) return current_energy, anal_energy - def update_internal_forces_and_torques(self, time): + def compute_internal_forces_and_torques(self, time): pass - def reset_external_forces_and_torques(self, time): + def zeroed_out_external_forces_and_torques(self, time): pass @@ -300,7 +300,9 @@ class CollectiveSystem: def __init__(self): self._memory_blocks = [] - self.systems = self._memory_blocks + + def systems(self): + return self._memory_blocks def __getitem__(self, idx): return self._memory_blocks[idx] @@ -344,8 +346,8 @@ def __init__(self): super( ScalarExponentialDampedHarmonicOscillatorCollectiveSystem, self ).__init__() - self.systems.append(ScalarExponentialDecaySystem()) - self.systems.append(DampedSimpleHarmonicOscillatorSystem()) + self._memory_blocks.append(ScalarExponentialDecaySystem()) + self._memory_blocks.append(DampedSimpleHarmonicOscillatorSystem()) def make_simple_system_with_positions_directors( diff --git a/elastica/systems/memory.py b/elastica/systems/memory.py new file mode 100644 index 000000000..c669be9b8 --- /dev/null +++ b/elastica/systems/memory.py @@ -0,0 +1,84 @@ +from typing import Iterator, TypeVar, Generic, Type +from elastica.timestepper.protocol import ExplicitStepperProtocol +from elastica.typing import SystemCollectionType + +from copy import copy + + +# FIXME: Move memory related functions to separate module or as part of the timestepper +# TODO: Use MemoryProtocol +def make_memory_for_explicit_stepper( + stepper: ExplicitStepperProtocol, system: SystemCollectionType +) -> "MemoryCollection": + # TODO Automated logic (class creation, memory management logic) agnostic of stepper details (RK, AB etc.) + + from elastica.timestepper.explicit_steppers import ( + RungeKutta4, + EulerForward, + ) + + # is_this_system_a_collection = is_system_a_collection(system) + + memory_cls: Type + if RungeKutta4 in stepper.__class__.mro(): + # Bad way of doing it, introduces tight coupling + # this should rather be taken from the class itself + class MemoryRungeKutta4: + def __init__(self) -> None: + self.initial_state = None + self.k_1 = None + self.k_2 = None + self.k_3 = None + self.k_4 = None + + memory_cls = MemoryRungeKutta4 + elif EulerForward in stepper.__class__.mro(): + + class MemoryEulerForward: + def __init__(self) -> None: + self.initial_state = None + self.k = None + + memory_cls = MemoryEulerForward + else: + raise NotImplementedError("Making memory for other types not supported") + + return MemoryCollection(memory_cls(), len(system)) + + +M = TypeVar("M", bound="MemoryCollection") + + +class MemoryCollection(Generic[M]): + """Slots of memories for timestepper in a cohesive unit. + + A `MemoryCollection` object is meant to be used in conjunction + with a `SystemCollection`, where each independent `System` to + be integrated has its own `Memory`. + + Example + ------- + + A RK4 integrator needs to store k_1, k_2, k_3, k_4 (intermediate + results from four stages) for each `System`. The restriction for + having a memory slot arises because the `Systems` are usually + not independent of one another and may need communication after + every stage. + """ + + def __init__(self, memory: M, n_memory_slots: int): + super(MemoryCollection, self).__init__() + + self.__memories: list[M] = [] + for _ in range(n_memory_slots - 1): + self.__memories.append(copy(memory)) + self.__memories.append(memory) + + def __getitem__(self, idx: int) -> M: + return self.__memories[idx] + + def __len__(self) -> int: + return len(self.__memories) + + def __iter__(self) -> Iterator[M]: + return self.__memories.__iter__() diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index d49129743..ef4cb7a6a 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -21,13 +21,21 @@ def compute_velocity_center_of_mass(self) -> NDArray[np.float64]: ... def compute_position_center_of_mass(self) -> NDArray[np.float64]: ... -class SystemProtocol(_SystemWithEnergy, _SystemWithCenterOfMass, Protocol): +class SystemProtocol(Protocol): """ Protocol for all elastica system """ REQUISITE_MODULES: list[Type] + def compute_internal_forces_and_torques(self, time: np.floating) -> None: ... + + def update_accelerations(self, time: np.floating) -> None: ... + + def zeroed_out_external_forces_and_torques(self, time: np.floating) -> None: ... + + +class SlenderBodyGeometryProtocol(Protocol): @property def n_nodes(self) -> int: ... @@ -35,29 +43,21 @@ def n_nodes(self) -> int: ... def n_elems(self) -> int: ... position_collection: NDArray[np.floating] - velocity_collection: NDArray[np.floating] - acceleration_collection: NDArray[np.floating] - director_collection: NDArray[np.floating] omega_collection: NDArray[np.floating] alpha_collection: NDArray[np.floating] - - internal_forces: NDArray[np.floating] - internal_torques: NDArray[np.floating] + director_collection: NDArray[np.floating] external_forces: NDArray[np.floating] external_torques: NDArray[np.floating] - def compute_internal_forces_and_torques(self, time: np.floating) -> None: ... - - def update_accelerations(self, time: np.floating) -> None: ... - - def zeroed_out_external_forces_and_torques(self, time: np.floating) -> None: ... + internal_forces: NDArray[np.floating] + internal_torques: NDArray[np.floating] -class SymplecticSystemProtocol(SystemProtocol, Protocol): +class SymplecticSystemProtocol(SystemProtocol, SlenderBodyGeometryProtocol, Protocol): """ Protocol for system with symplectic state variables """ @@ -80,7 +80,7 @@ def dynamic_rates( ) -> NDArray[np.floating]: ... -class ExplicitSystemProtocol(SystemProtocol, Protocol): +class ExplicitSystemProtocol(SystemProtocol, SlenderBodyGeometryProtocol, Protocol): # TODO: Temporarily made to handle explicit stepper. # Need to be refactored as the explicit stepper is further developed. def __call__(self, time: np.floating, dt: np.floating) -> np.floating: ... diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index aeb4cab03..5a5877f67 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -188,7 +188,7 @@ def _first_update( time: np.floating, dt: np.floating, ) -> np.floating: - System.state += dt * System(time, dt) + System.state += dt * System(time, dt) # type: ignore[arg-type] return time + dt diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 84a87ce20..5d7c0ca37 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -95,7 +95,7 @@ def do_step( """ for kin_prefactor, kin_step, dyn_step in steps_and_prefactors[:-1]: - for system in SystemCollection._memory_blocks: + for system in SystemCollection.systems(): kin_step(system, time, dt) time += kin_prefactor(dt) @@ -104,14 +104,14 @@ def do_step( SystemCollection.constrain_values(time) # We need internal forces and torques because they are used by interaction module. - for system in SystemCollection._memory_blocks: - system.update_internal_forces_and_torques(time) + for system in SystemCollection.systems(): + system.compute_internal_forces_and_torques(time) # system.update_internal_forces_and_torques() # Add external forces, controls etc. SystemCollection.synchronize(time) - for system in SystemCollection._memory_blocks: + for system in SystemCollection.systems(): dyn_step(system, time, dt) # Constrain only rates @@ -121,7 +121,7 @@ def do_step( last_kin_prefactor = steps_and_prefactors[-1][0] last_kin_step = steps_and_prefactors[-1][1] - for system in SystemCollection._memory_blocks: + for system in SystemCollection.systems(): last_kin_step(system, time, dt) time += last_kin_prefactor(dt) SystemCollection.constrain_values(time) @@ -130,14 +130,14 @@ def do_step( SystemCollection.apply_callbacks(time, round(time / dt)) # Zero out the external forces and torques - for system in SystemCollection._memory_blocks: - system.reset_external_forces_and_torques(time) + for system in SystemCollection.systems(): + system.zeroed_out_external_forces_and_torques(time) return time def step_single_instance( self: SymplecticStepperProtocol, - System: SystemType, + System: SymplecticSystemProtocol, time: np.floating, dt: np.floating, ) -> np.floating: @@ -145,7 +145,7 @@ def step_single_instance( for kin_prefactor, kin_step, dyn_step in self.steps_and_prefactors[:-1]: kin_step(System, time, dt) time += kin_prefactor(dt) - System.update_internal_forces_and_torques(time) + System.compute_internal_forces_and_torques(time) dyn_step(System, time, dt) # Peel the last kinematic step and prefactor alone diff --git a/elastica/typing.py b/elastica/typing.py index 177c7ea4c..d13e927dc 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -4,7 +4,7 @@ """ from typing import TYPE_CHECKING -from typing import Type, Callable, Any, ParamSpec +from typing import Callable, Any, ParamSpec from typing import TypeAlias import numpy as np @@ -14,21 +14,27 @@ # Used for type hinting without circular imports # NEVER BACK-IMPORT ANY ELASTICA MODULES HERE from .rod.protocol import CosseratRodProtocol - from .rigidbody import RigidBodyBase - from .surface import SurfaceBase + from .rigidbody.protocol import RigidBodyProtocol + from .surface.surface_base import SurfaceBase from .modules.base_system import BaseSystemCollection from .modules.protocol import SystemCollectionProtocol from .rod.data_structures import _State as State - from .systems.protocol import SymplecticSystemProtocol, ExplicitSystemProtocol + from .systems.protocol import ( + SystemProtocol, + SymplecticSystemProtocol, + ExplicitSystemProtocol, + ) from .timestepper.protocol import ( StepperProtocol, SymplecticStepperProtocol, MemoryProtocol, ) - from memory_block.protocol import ( + from .memory_block.protocol import ( BlockCosseratRodProtocol, - ) # , BlockRigidBodyProtocol + BlockRigidBodyProtocol, + BlockSystemProtocol, + ) # Modules Base Classes from .boundary_conditions import FreeBC @@ -38,8 +44,11 @@ from .external_forces import NoForces from .joint import FreeJoint + from .mesh.protocol import MeshProtocol -SystemType: TypeAlias = "SymplecticSystemProtocol" # | ExplicitSystemProtocol + +SystemType: TypeAlias = "SystemProtocol" +BlockSystemType: TypeAlias = "BlockSystemProtocol" SystemIdxType: TypeAlias = int # ModuleObjectTypes: TypeAlias = ( @@ -57,12 +66,10 @@ RodType: TypeAlias = "CosseratRodProtocol" RigidBodyType: TypeAlias = "RigidBodyProtocol" +SurfaceType: TypeAlias = "SurfaceBase" SystemCollectionType: TypeAlias = "SystemCollectionProtocol" -AllowedContactType: TypeAlias = ( - SystemType | Type["SurfaceBase"] -) # FIXME: SurfaceBase needs to be treated differently -BlockType: TypeAlias = "BlockCosseratRodProtocol" # | "BlockRigidBodyProtocol" +BlockType: TypeAlias = "BlockCosseratRodProtocol | BlockRigidBodyProtocol" # Indexing types # TODO: Maybe just use slice?? @@ -73,6 +80,7 @@ # Operators in elastica.modules OperatorParam = ParamSpec("OperatorParam") -OperatorType: TypeAlias = Callable[OperatorParam, None] OperatorCallbackType: TypeAlias = Callable[..., None] OperatorFinalizeType: TypeAlias = Callable[..., None] + +MeshType: TypeAlias = "MeshProtocol" diff --git a/elastica/utils.py b/elastica/utils.py index d31cf3cdc..429ae63ab 100644 --- a/elastica/utils.py +++ b/elastica/utils.py @@ -80,7 +80,7 @@ def atol() -> float: ------- atol : library-wide set absolute tolerance for kernels """ - return finfo(float64).eps * 1e4 + return float(finfo(float64).eps * 1e4) @staticmethod def rtol() -> float: @@ -91,7 +91,7 @@ def rtol() -> float: ------- tol : library-wide set relative tolerance for kernels """ - return finfo(float64).eps * 1e11 + return float(finfo(float64).eps * 1e11) def perm_parity(lst: list[int]) -> int: @@ -178,8 +178,8 @@ def extend_instance(obj: Any, cls: Any) -> None: obj.__class__ = type(base_cls_name, (cls, base_cls), {}) -def _bspline( - t_coeff: NDArray, l_centerline: np.floating = 1.0 +def _bspline( # type: ignore[no-any-unimported] + t_coeff: NDArray, l_centerline: np.floating = np.float64(1.0) ) -> tuple[BSpline, NDArray, NDArray]: """Generates a bspline object that plots the spline interpolant for any vector x. Optionally takes in a centerline length, set to 1.0 by @@ -208,7 +208,7 @@ def _bspline( return __bspline_impl__(control_pts, t_coeff, degree) -def __bspline_impl__( +def __bspline_impl__( # type: ignore[no-any-unimported] x_pts: NDArray, t_c: NDArray, degree: int ) -> tuple[BSpline, NDArray, NDArray]: """""" diff --git a/elastica/wrappers.py b/elastica/wrappers.py deleted file mode 100644 index fa3990f80..000000000 --- a/elastica/wrappers.py +++ /dev/null @@ -1,23 +0,0 @@ -import warnings - -__all__ = [ - "BaseSystemCollection", - "Connections", - "Constraints", - "Forcing", - "CallBacks", - "Damping", -] - -from elastica.modules.base_system import BaseSystemCollection -from elastica.modules.connections import Connections -from elastica.modules.constraints import Constraints -from elastica.modules.forcing import Forcing -from elastica.modules.callbacks import CallBacks -from elastica.modules.damping import Damping - - -warnings.warn( - "elastica.wrappers is refactored to elastica.modules in version 0.3.0.", - DeprecationWarning, -) diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py index 68dea81dd..c27ffd37c 100755 --- a/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py +++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py @@ -22,10 +22,11 @@ _node_to_element_mass_or_force, ) from numba import njit -from elastica.rod import RodBase +from elastica.rod.rod_base import RodBase from elastica.surface import Plane +from elastica.surface.surface_base import SurfaceBase from elastica.contact_forces import NoContact -from elastica.typing import RodType, SystemType, AllowedContactType +from elastica.typing import RodType, SystemType @njit(cache=True) @@ -286,31 +287,10 @@ def __init__( self.kinetic_mu_sideways, ) = kinetic_mu_array - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodSphereContact class first_system should be a rod and second_system should be a plane. - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, Plane - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane".format( - system_one.__class__, system_two.__class__ - ) - ) + @property + def _allowed_system_two(self) -> list[Type]: + # Modify this list to include the allowed system types for contact + return [SurfaceBase] def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: """ diff --git a/pyproject.toml b/pyproject.toml index 71b03f98d..2628c4f9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -145,6 +145,7 @@ warn_unused_ignores = false exclude = [ "elastica/systems/analytical.py", "elastica/experimental/*", + "elastica/mesh/mesh_initializer.py", ] [tool.coverage.report] diff --git a/tests/test_callback_functions.py b/tests/test_callback_functions.py index 8c6362b36..0e5f3e21e 100644 --- a/tests/test_callback_functions.py +++ b/tests/test_callback_functions.py @@ -104,14 +104,14 @@ def test_my_call_back_base_class(self, n_elems): class TestExportCallBackClass: @pytest.mark.parametrize("method", ["0", 1, "numba", "test", "some string", None]) - def test_export_call_back_unavailable_save_methods(self, method): + def test_export_call_back_unavailable_save_methods(self, tmp_path, method): with pytest.raises(AssertionError) as excinfo: - callback = ExportCallBack(1, "rod", "tempdir", method) + callback = ExportCallBack(1, "rod", tmp_path.as_posix(), method) @pytest.mark.parametrize("method", ExportCallBack.AVAILABLE_METHOD) - def test_export_call_back_available_save_methods(self, method): + def test_export_call_back_available_save_methods(self, tmp_path, method): try: - callback = ExportCallBack(1, "rod", "tempdir", method) + callback = ExportCallBack(1, "rod", tmp_path.as_posix(), method) except Error: pytest.fail( f"Could not create callback module with available method {method}" @@ -236,7 +236,7 @@ def test_export_call_back_clear_test(self): assert os.path.exists(saved_path_name), "File is not saved." @pytest.mark.parametrize("n_elems", [2, 4, 16]) - def test_export_call_back_class_tempfile_option(self, n_elems): + def test_export_call_back_class_tempfile_option(self, tmp_path, n_elems): """ This test case is for testing ExportCallBack function, saving into temporary files. """ @@ -256,7 +256,7 @@ def test_export_call_back_class_tempfile_option(self, n_elems): } callback = ExportCallBack( - step_skip, "rod", "tempdir", "tempfile", file_save_interval=10 + step_skip, "rod", tmp_path.as_posix(), "tempfile", file_save_interval=10 ) for i in range(10): callback.make_callback(mock_rod, time[i], current_step[i]) diff --git a/tests/test_contact_classes.py b/tests/test_contact_classes.py index 94b41a457..509d8e32d 100644 --- a/tests/test_contact_classes.py +++ b/tests/test_contact_classes.py @@ -52,8 +52,8 @@ def mock_cylinder_init(self): self.director_collection = np.array( [[[1.0], [0.0], [0.0]], [[0.0], [1.0], [0.0]], [[0.0], [0.0], [1.0]]] ) - self.radius = np.array([1.0]) - self.length = np.array([2.0]) + self.radius = 1.0 + self.length = 2.0 self.external_forces = np.array([[0.0], [0.0], [0.0]]) self.external_torques = np.array([[0.0], [0.0], [0.0]]) self.velocity_collection = np.array([[0.0], [0.0], [0.0]]) @@ -69,7 +69,7 @@ def mock_sphere_init(self): self.director_collection = np.array( [[[1.0], [0.0], [0.0]], [[0.0], [1.0], [0.0]], [[0.0], [0.0], [1.0]]] ) - self.radius = np.array([1.0]) + self.radius = 1.0 self.velocity_collection = np.array([[0.0], [0.0], [0.0]]) self.external_forces = np.array([[0.0], [0.0], [0.0]]) self.external_torques = np.array([[0.0], [0.0], [0.0]]) @@ -104,43 +104,41 @@ def test_check_systems_validity_with_invalid_systems( "Testing Rod Cylinder Contact wrapper with incorrect type for second argument" with pytest.raises(TypeError) as excinfo: rod_cylinder_contact._check_systems_validity(mock_rod, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a cylinder" - ).format(mock_rod.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['Cylinder']." == str( + excinfo.value + ) - "Testing Rod Cylinder Contact wrapper with incorrect type for first argument" with pytest.raises(TypeError) as excinfo: rod_cylinder_contact._check_systems_validity(mock_list, mock_rod) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a cylinder" - ).format(mock_list.__class__, mock_rod.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Rod Cylinder Contact wrapper with incorrect order" with pytest.raises(TypeError) as excinfo: rod_cylinder_contact._check_systems_validity(mock_cylinder, mock_rod) - print(excinfo.value) assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a cylinder" - ).format(mock_cylinder.__class__, mock_rod.__class__) == str(excinfo.value) + "System provided (MockCylinder) must be derived from ['RodBase']." + == str(excinfo.value) + ) + + with pytest.raises(TypeError) as excinfo: + rod_cylinder_contact._check_systems_validity(mock_list, mock_cylinder) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) def test_contact_rod_cylinder_with_collision_with_k_without_nu_and_friction( self, ): - "Testing Rod Cylinder Contact wrapper with Collision with analytical verified values" + # Testing Rod Cylinder Contact wrapper with Collision with analytical verified values mock_rod = MockRod() mock_cylinder = MockCylinder() rod_cylinder_contact = RodCylinderContact(k=1.0, nu=0.0) rod_cylinder_contact.apply_contact(mock_rod, mock_cylinder) - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_cylinder()'""" + # Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_cylinder()' assert_allclose( mock_rod.external_forces, np.array([[0.166666, 0.333333, 0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), @@ -293,32 +291,27 @@ def test_check_systems_validity_with_invalid_systems( mock_list = [1, 2, 3] rod_rod_contact = RodRodContact(k=1.0, nu=0.0) - "Testing Rod Rod Contact wrapper with incorrect type for second argument" + # Testing Rod Rod Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: rod_rod_contact._check_systems_validity(mock_rod_one, mock_list) - assert ( - "Systems provided to the contact class have incorrect order. \n" - " First system is {0} and second system is {1}. \n" - " Both systems must be distinct rods" - ).format(mock_rod_one.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Rod Rod Contact wrapper with incorrect type for first argument" + # Testing Rod Rod Contact wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: rod_rod_contact._check_systems_validity(mock_list, mock_rod_one) - assert ( - "Systems provided to the contact class have incorrect order. \n" - " First system is {0} and second system is {1}. \n" - " Both systems must be distinct rods" - ).format(mock_list.__class__, mock_rod_one.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Rod Rod Contact wrapper with same rod for both arguments" + # Testing Rod Rod Contact wrapper with same rod for both arguments with pytest.raises(TypeError) as excinfo: rod_rod_contact._check_systems_validity(mock_rod_one, mock_rod_one) assert ( - "First rod is identical to second rod. \n" - "Rods must be distinct for RodRodConact. \n" - "If you want self contact, use RodSelfContact instead" - ) == str(excinfo.value) + "First system is identical to second system. Systems must be distinct for contact." + == str(excinfo.value) + ) def test_contact_with_two_rods_with_collision_with_k_without_nu(self): @@ -439,35 +432,26 @@ def test_check_systems_validity_with_invalid_systems( mock_list = [1, 2, 3] self_contact = RodSelfContact(k=1.0, nu=0.0) - "Testing Self Contact wrapper with incorrect type for second argument" + # Testing Self Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: self_contact._check_systems_validity(mock_rod_one, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system and second system should be the same rod \n" - " If you want rod rod contact, use RodRodContact instead" - ).format(mock_rod_one.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Self Contact wrapper with incorrect type for first argument" + # Testing Self Contact wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: self_contact._check_systems_validity(mock_list, mock_rod_one) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system and second system should be the same rod \n" - " If you want rod rod contact, use RodRodContact instead" - ).format(mock_list.__class__, mock_rod_one.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Self Contact wrapper with different rods" + # Testing Self Contact wrapper with different rods with pytest.raises(TypeError) as excinfo: self_contact._check_systems_validity(mock_rod_one, mock_rod_two) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system and second system should be the same rod \n" - " If you want rod rod contact, use RodRodContact instead" - ).format(mock_rod_one.__class__, mock_rod_two.__class__) == str(excinfo.value) + assert "First system must be identical to the second system." == str( + excinfo.value + ) def test_self_contact_with_rod_self_collision(self): @@ -530,46 +514,39 @@ def test_check_systems_validity_with_invalid_systems( mock_sphere = MockSphere() rod_sphere_contact = RodSphereContact(k=1.0, nu=0.0) - "Testing Rod Sphere Contact wrapper with incorrect type for second argument" + # Testing Rod Sphere Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: rod_sphere_contact._check_systems_validity(mock_rod, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a sphere" - ).format(mock_rod.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['Sphere']." == str( + excinfo.value + ) - "Testing Rod Sphere Contact wrapper with incorrect type for first argument" + # Testing Rod Sphere Contact wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: rod_sphere_contact._check_systems_validity(mock_list, mock_rod) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a sphere" - ).format(mock_list.__class__, mock_rod.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Rod Sphere Contact wrapper with incorrect order" + # Testing Rod Sphere Contact wrapper with incorrect order with pytest.raises(TypeError) as excinfo: rod_sphere_contact._check_systems_validity(mock_sphere, mock_rod) - print(excinfo.value) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a sphere" - ).format(mock_sphere.__class__, mock_rod.__class__) == str(excinfo.value) + assert "System provided (MockSphere) must be derived from ['RodBase']." == str( + excinfo.value + ) def test_contact_rod_sphere_with_collision_with_k_without_nu_and_friction( self, ): - "Testing Rod Sphere Contact wrapper with Collision with analytical verified values" + # "Testing Rod Sphere Contact wrapper with Collision with analytical verified values mock_rod = MockRod() mock_sphere = MockSphere() rod_sphere_contact = RodSphereContact(k=1.0, nu=0.0) rod_sphere_contact.apply_contact(mock_rod, mock_sphere) - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_sphere_with_k_without_nu_and_friction()'""" + # Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_sphere_with_k_without_nu_and_friction()' assert_allclose( mock_sphere.external_forces, np.array([[-0.5], [0], [0]]), atol=1e-6 ) @@ -630,23 +607,19 @@ def test_check_systems_validity_with_invalid_systems( mock_list = [1, 2, 3] rod_plane_contact = RodPlaneContact(k=1.0, nu=0.0) - "Testing Rod Plane Contact wrapper with incorrect type for second argument" + # Testing Rod Plane Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: rod_plane_contact._check_systems_validity(mock_rod, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane" - ).format(mock_rod.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['SurfaceBase']." == str( + excinfo.value + ) - "Testing Rod Plane Contact wrapper with incorrect type for first argument" + # Testing Rod Plane Contact wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: rod_plane_contact._check_systems_validity(mock_list, mock_plane) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane" - ).format(mock_list.__class__, mock_plane.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) def test_rod_plane_contact_without_contact(self): """ @@ -905,23 +878,19 @@ def test_check_systems_validity_with_invalid_systems( np.array([0.0, 0.0, 0.0]), # forward, backward, sideways ) - "Testing Rod Plane Contact wrapper with incorrect type for second argument" + # Testing Rod Plane Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: rod_plane_contact._check_systems_validity(mock_rod, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane" - ).format(mock_rod.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['SurfaceBase']." == str( + excinfo.value + ) - "Testing Rod Plane wrapper with incorrect type for first argument" + # Testing Rod Plane wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: rod_plane_contact._check_systems_validity(mock_list, mock_plane) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane" - ).format(mock_list.__class__, mock_plane.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) @pytest.mark.parametrize("velocity", [-1.0, -3.0, 1.0, 5.0, 2.0]) def test_axial_kinetic_friction(self, velocity): @@ -1212,7 +1181,7 @@ def initializer( # create plane plane = MockPlane() - plane.origin = np.array([0.0, -cylinder.radius[0] + shift, 0.0]).reshape(3, 1) + plane.origin = np.array([0.0, -cylinder.radius + shift, 0.0]).reshape(3, 1) plane.normal = plane_normal.reshape( 3, ) @@ -1232,23 +1201,19 @@ def test_check_systems_validity_with_invalid_systems( mock_list = [1, 2, 3] cylinder_plane_contact = CylinderPlaneContact(k=1.0, nu=0.0) - "Testing Cylinder Plane Contact wrapper with incorrect type for second argument" + # Testing Cylinder Plane Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: cylinder_plane_contact._check_systems_validity(mock_cylinder, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a cylinder, second should be a plane" - ).format(mock_cylinder.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['SurfaceBase']." == str( + excinfo.value + ) - "Testing Cylinder Plane wrapper with incorrect type for first argument" + # Testing Cylinder Plane wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: cylinder_plane_contact._check_systems_validity(mock_list, mock_plane) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a cylinder, second should be a plane" - ).format(mock_list.__class__, mock_plane.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['Cylinder']." == str( + excinfo.value + ) def test_cylinder_plane_contact_without_contact(self): """ diff --git a/tests/test_interaction.py b/tests/test_interaction.py index 8d2b0e13c..9d30b8067 100644 --- a/tests/test_interaction.py +++ b/tests/test_interaction.py @@ -6,13 +6,8 @@ from elastica.utils import Tolerance, MaxDimension from elastica.interaction import ( InteractionPlane, - find_slipping_elements, AnisotropicFrictionalPlane, - node_to_element_mass_or_force, SlenderBodyTheory, - nodes_to_elements, - elements_to_nodes_inplace, - apply_normal_force_numba_rigid_body, ) from elastica.contact_utils import ( _node_to_element_mass_or_force, @@ -105,7 +100,7 @@ def test_interaction_without_contact(self, n_elem): [rod, interaction_plane, external_forces] = self.initializer(n_elem, shift) - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) correct_forces = external_forces # since no contact assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -127,7 +122,7 @@ def test_interaction_plane_without_k_and_nu(self, n_elem): [rod, interaction_plane, external_forces] = self.initializer(n_elem) - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) correct_forces = np.zeros((3, n_elem + 1)) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -158,7 +153,7 @@ def test_interaction_plane_with_k_without_nu(self, n_elem, k_w): correct_forces[..., 0] *= 0.5 correct_forces[..., -1] *= 0.5 - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -194,7 +189,7 @@ def test_interaction_plane_without_k_with_nu(self, n_elem, nu_w): correct_forces[..., 0] *= 0.5 correct_forces[..., -1] *= 0.5 - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -224,7 +219,7 @@ def test_interaction_when_rod_is_under_plane(self, n_elem): n_elem, shift=offset_of_plane_with_respect_to_rod, plane_normal=plane_normal ) - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) correct_forces = np.zeros((3, n_elem + 1)) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -269,7 +264,7 @@ def test_interaction_when_rod_is_under_plane_with_k_without_nu(self, n_elem, k_w correct_forces[..., 0] *= 0.5 correct_forces[..., -1] *= 0.5 - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -317,58 +312,11 @@ def test_interaction_when_rod_is_under_plane_without_k_with_nu(self, n_elem, nu_ correct_forces[..., 0] *= 0.5 correct_forces[..., -1] *= 0.5 - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) -class TestAuxiliaryFunctions: - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_linear_interpolation_slip_error_message(self, n_elem): - velocity_threshold = 1.0 - - # if slip velocity larger than threshold - velocity_slip = np.repeat( - np.array([0.0, 0.0, 2.0]).reshape(3, 1), n_elem, axis=1 - ) - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._find_slipping_elements()\n" - "instead for finding slipping elements." - ) - with pytest.raises(NotImplementedError) as error_info: - slip_function = find_slipping_elements(velocity_slip, velocity_threshold) - assert error_info.value.args[0] == error_message - - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_node_to_element_mass_or_force_error_message(self, n_elem): - random_vector = np.random.rand(3).reshape(3, 1) - input = np.repeat(random_vector, n_elem + 1, axis=1) - input[..., 0] *= 0.5 - input[..., -1] *= 0.5 - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._node_to_element_mass_or_force()\n" - "instead for converting the mass/forces on rod nodes to elements." - ) - with pytest.raises(NotImplementedError) as error_info: - output = node_to_element_mass_or_force(input) - assert error_info.value.args[0] == error_message - - @pytest.mark.parametrize("n_elem", [2, 10]) - def test_not_impl_error_for_nodes_to_elements(self, n_elem): - random_vector = np.random.rand(3).reshape(3, 1) - input = np.repeat(random_vector, n_elem + 1, axis=1) - error_message = ( - "This function is removed in v0.3.1. Please use\n" - "elastica.interaction.node_to_element_mass_or_force()\n" - "instead for node-to-element interpolation of mass/forces." - ) - with pytest.raises(NotImplementedError) as error_info: - nodes_to_elements(input) - assert error_info.value.args[0] == error_message - - class TestAnisotropicFriction: def initializer( self, @@ -739,140 +687,33 @@ def test_static_rolling_friction_total_torque_larger_than_static_friction_force( # Slender Body Theory Unit Tests +from elastica.interaction import ( + sum_over_elements, +) -try: - from elastica.interaction import ( - sum_over_elements, - node_to_element_position, - node_to_element_velocity, - node_to_element_pos_or_vel, - ) - - # These functions are used in the case if Numba is available - class TestAuxiliaryFunctionsForSlenderBodyTheory: - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_sum_over_elements(self, n_elem): - """ - This function test sum over elements function with - respect to default python function .sum(). We write - this function because with numba we can get the sum - faster. - Parameters - ---------- - n_elem - - Returns - ------- - - """ - - input_variable = np.random.rand(n_elem) - correct_output = input_variable.sum() - output = sum_over_elements(input_variable) - assert_allclose(correct_output, output, atol=Tolerance.atol()) - - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_node_to_element_position_error_message(self, n_elem): - """ - This function tests node_to_element_position function. We are - converting node positions to element positions. Here also - we are using numba to speed up the process. - - Parameters - ---------- - n_elem - - Returns - ------- - - """ - random = np.random.rand() # Adding some random numbers - input_position = random * np.ones((3, n_elem + 1)) - correct_output = random * np.ones((3, n_elem)) - - error_message = ( - "This function is removed in v0.3.2. For node-to-element_position() interpolation please use: \n" - "elastica.contact_utils._node_to_element_position() for rod position \n" - "For detail, refer to issue #113." - ) - with pytest.raises(NotImplementedError) as error_info: - output = node_to_element_position(input_position) - assert error_info.value.args[0] == error_message - - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_node_to_element_velocity_error_message(self, n_elem): - """ - This function tests node_to_element_velocity function. We are - converting node velocities to element velocities. Here also - we are using numba to speed up the process. - - Parameters - ---------- - n_elem - - Returns - ------- - - """ - random = np.random.rand() # Adding some random numbers - input_velocity = random * np.ones((3, n_elem + 1)) - input_mass = 2.0 * random * np.ones(n_elem + 1) - correct_output = random * np.ones((3, n_elem)) - - error_message = ( - "This function is removed in v0.3.2. For node-to-element_velocity() interpolation please use: \n" - "elastica.contact_utils._node_to_element_velocity() for rod velocity. \n" - "For detail, refer to issue #113." - ) - with pytest.raises(NotImplementedError) as error_info: - output = node_to_element_velocity( - mass=input_mass, node_velocity_collection=input_velocity - ) - assert error_info.value.args[0] == error_message - - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_not_impl_error_for_node_to_element_pos_or_vel(self, n_elem): - random = np.random.rand() # Adding some random numbers - input_velocity = random * np.ones((3, n_elem + 1)) - error_message = ( - "This function is removed in v0.3.0. For node-to-element interpolation please use: \n" - "elastica.contact_utils._node_to_element_position() for rod position \n" - "elastica.contact_utils._node_to_element_velocity() for rod velocity. \n" - "For detail, refer to issue #80." - ) - with pytest.raises(NotImplementedError) as error_info: - node_to_element_pos_or_vel(input_velocity) - assert error_info.value.args[0] == error_message - - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_elements_to_nodes_inplace_error_message(self, n_elem): - """ - This function tests _elements_to_nodes_inplace. We are - converting node velocities to element velocities. Here also - we are using numba to speed up the process. - - Parameters - ---------- - n_elem - - Returns - ------- - - """ - random = np.random.rand() # Adding some random numbers - vector_in_element_frame = random * np.ones((3, n_elem)) - vector_in_node_frame = np.zeros((3, n_elem + 1)) - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._elements_to_nodes_inplace()\n" - "instead for updating nodal forces using the forces computed on elements." - ) - with pytest.raises(NotImplementedError) as error_info: - elements_to_nodes_inplace(vector_in_element_frame, vector_in_node_frame) - assert error_info.value.args[0] == error_message -except ImportError: - pass +# These functions are used in the case if Numba is available +class TestAuxiliaryFunctionsForSlenderBodyTheory: + @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) + def test_sum_over_elements(self, n_elem): + """ + This function test sum over elements function with + respect to default python function .sum(). We write + this function because with numba we can get the sum + faster. + Parameters + ---------- + n_elem + + Returns + ------- + + """ + + input_variable = np.random.rand(n_elem) + correct_output = input_variable.sum() + output = sum_over_elements(input_variable) + assert_allclose(correct_output, output, atol=Tolerance.atol()) class TestSlenderBody: @@ -1011,42 +852,3 @@ def test_slender_body_matrix_product_only_xy(self, n_elem): slender_body_theory.apply_forces(rod) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) - - -@pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) -def test_apply_normal_force_numba_rigid_body_error_message(n_elem): - """ - This function _elements_to_nodes_inplace. We are - converting node velocities to element velocities. Here also - we are using numba to speed up the process. - - Parameters - ---------- - n_elem - - Returns - ------- - - """ - - position_collection = np.zeros((3, n_elem + 1)) - position_collection[0, :] = np.linspace(0, 1.0, n_elem + 1) - - error_message = ( - "This function is removed in v0.3.2. For cylinder plane contact please use: \n" - "elastica._contact_functions._calculate_contact_forces_cylinder_plane() \n" - "For detail, refer to issue #113." - ) - with pytest.raises(NotImplementedError) as error_info: - apply_normal_force_numba_rigid_body( - plane_origin=np.array([0.0, 0.0, 0.0]), - plane_normal=np.array([0.0, 0.0, 1.0]), - surface_tol=1e-4, - k=1.0, - nu=1.0, - length=1.0, - position_collection=position_collection, - velocity_collection=np.zeros((3, n_elem + 1)), - external_forces=np.zeros((3, n_elem + 1)), - ) - assert error_info.value.args[0] == error_message diff --git a/tests/test_joint.py b/tests/test_joint.py index 40cbd7c27..6d5165914 100644 --- a/tests/test_joint.py +++ b/tests/test_joint.py @@ -13,7 +13,6 @@ import numpy as np import pytest from scipy.spatial.transform import Rotation -from elastica.joint import ExternalContact, SelfContact # TODO: change tests and made them independent of rod, at least assigin hardcoded values for forces and torques @@ -412,329 +411,3 @@ def mock_rigid_body_init(self): self.external_forces = np.array([[0.0], [0.0], [0.0]]) self.external_torques = np.array([[0.0], [0.0], [0.0]]) self.velocity_collection = np.array([[0.0], [0.0], [0.0]]) - - -MockRod = type("MockRod", (RodBase,), {"__init__": mock_rod_init}) - -MockRigidBody = type( - "MockRigidBody", (RigidBodyBase,), {"__init__": mock_rigid_body_init} -) - - -class TestExternalContact: - def test_external_contact_rod_rigid_body_with_collision_with_k_without_nu_and_friction( - self, - ): - - "Testing External Contact wrapper with Collision with analytical verified values" - - mock_rod = MockRod() - mock_rigid_body = MockRigidBody() - ext_contact = ExternalContact(k=1.0, nu=0.0) - ext_contact.apply_forces(mock_rod, 0, mock_rigid_body, 1) - - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_claculate_contact_forces_rod_rigid_body()'""" - assert_allclose( - mock_rod.external_forces, - np.array([[0.166666, 0.333333, 0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), - atol=1e-6, - ) - - assert_allclose( - mock_rigid_body.external_forces, np.array([[-0.5], [0.0], [0.0]]), atol=1e-6 - ) - - assert_allclose( - mock_rigid_body.external_torques, np.array([[0.0], [0.0], [0.0]]), atol=1e-6 - ) - - def test_external_contact_rod_rigid_body_with_collision_with_nu_without_k_and_friction( - self, - ): - - "Testing External Contact wrapper with Collision with analytical verified values" - - mock_rod = MockRod() - "Moving rod towards the cylinder with a velocity of -1 in x-axis" - mock_rod.velocity_collection = np.array([[-1, 0, 0], [-1, 0, 0], [-1, 0, 0]]) - mock_rigid_body = MockRigidBody() - "Moving cylinder towards the rod with a velocity of 1 in x-axis" - mock_rigid_body.velocity_collection = np.array([[1], [0], [0]]) - ext_contact = ExternalContact(k=0.0, nu=1.0) - ext_contact.apply_forces(mock_rod, 0, mock_rigid_body, 1) - - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_claculate_contact_forces_rod_rigid_body()'""" - assert_allclose( - mock_rod.external_forces, - np.array([[0.5, 1, 0], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - - assert_allclose( - mock_rigid_body.external_forces, np.array([[-1.5], [0], [0]]), atol=1e-6 - ) - - assert_allclose( - mock_rigid_body.external_torques, np.array([[0.0], [0.0], [0.0]]), atol=1e-6 - ) - - def test_external_contact_rod_rigid_body_with_collision_with_k_and_nu_without_friction( - self, - ): - - "Testing External Contact wrapper with Collision with analytical verified values" - - mock_rod = MockRod() - "Moving rod towards the cylinder with a velocity of -1 in x-axis" - mock_rod.velocity_collection = np.array([[-1, 0, 0], [-1, 0, 0], [-1, 0, 0]]) - mock_rigid_body = MockRigidBody() - "Moving cylinder towards the rod with a velocity of 1 in x-axis" - mock_rigid_body.velocity_collection = np.array([[1], [0], [0]]) - ext_contact = ExternalContact(k=1.0, nu=1.0) - ext_contact.apply_forces(mock_rod, 0, mock_rigid_body, 1) - - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_claculate_contact_forces_rod_rigid_body()'""" - assert_allclose( - mock_rod.external_forces, - np.array([[0.666666, 1.333333, 0], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - - assert_allclose( - mock_rigid_body.external_forces, np.array([[-2], [0], [0]]), atol=1e-6 - ) - - assert_allclose( - mock_rigid_body.external_torques, np.array([[0.0], [0.0], [0.0]]), atol=1e-6 - ) - - def test_external_contact_rod_rigid_body_with_collision_with_k_and_nu_and_friction( - self, - ): - - "Testing External Contact wrapper with Collision with analytical verified values" - - mock_rod = MockRod() - "Moving rod towards the cylinder with a velocity of -1 in x-axis" - mock_rod.velocity_collection = np.array([[-1, 0, 0], [-1, 0, 0], [-1, 0, 0]]) - mock_rigid_body = MockRigidBody() - "Moving cylinder towards the rod with a velocity of 1 in x-axis" - mock_rigid_body.velocity_collection = np.array([[1], [0], [0]]) - ext_contact = ExternalContact( - k=1.0, nu=1.0, velocity_damping_coefficient=0.1, friction_coefficient=0.1 - ) - ext_contact.apply_forces(mock_rod, 0, mock_rigid_body, 1) - - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_claculate_contact_forces_rod_rigid_body()'""" - assert_allclose( - mock_rod.external_forces, - np.array( - [ - [0.666666, 1.333333, 0], - [0.033333, 0.066666, 0], - [0.033333, 0.066666, 0], - ] - ), - atol=1e-6, - ) - - assert_allclose( - mock_rigid_body.external_forces, np.array([[-2], [-0.1], [-0.1]]), atol=1e-6 - ) - - assert_allclose( - mock_rigid_body.external_torques, np.array([[0.0], [0.0], [0.0]]), atol=1e-6 - ) - - def test_external_contact_rod_rigid_body_without_collision(self): - - "Testing External Contact wrapper without Collision with analytical verified values" - - mock_rod = MockRod() - mock_rigid_body = MockRigidBody() - ext_contact = ExternalContact(k=1.0, nu=0.5) - - """Setting rigid body position such that there is no collision""" - mock_rigid_body.position_collection = np.array([[400], [500], [600]]) - mock_rod_external_forces_before_execution = mock_rod.external_forces.copy() - mock_rigid_body_external_forces_before_execution = ( - mock_rigid_body.external_forces.copy() - ) - mock_rigid_body_external_torques_before_execution = ( - mock_rigid_body.external_torques.copy() - ) - ext_contact.apply_forces(mock_rod, 0, mock_rigid_body, 1) - - assert_allclose( - mock_rod.external_forces, mock_rod_external_forces_before_execution - ) - assert_allclose( - mock_rigid_body.external_forces, - mock_rigid_body_external_forces_before_execution, - ) - assert_allclose( - mock_rigid_body.external_torques, - mock_rigid_body_external_torques_before_execution, - ) - - def test_external_contact_with_two_rods_with_collision_with_k_without_nu(self): - - "Testing External Contact wrapper with two rods with analytical verified values" - "Test values have been copied from 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_rod()'" - - mock_rod_one = MockRod() - mock_rod_two = MockRod() - mock_rod_two.position_collection = np.array([[4, 5, 6], [0, 0, 0], [0, 0, 0]]) - ext_contact = ExternalContact(k=1.0, nu=0.0) - ext_contact.apply_forces(mock_rod_one, 0, mock_rod_two, 0) - - assert_allclose( - mock_rod_one.external_forces, - np.array([[0, -0.666666, -0.333333], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - assert_allclose( - mock_rod_two.external_forces, - np.array([[0.333333, 0.666666, 0], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - - def test_external_contact_with_two_rods_with_collision_without_k_with_nu(self): - - "Testing External Contact wrapper with two rods with analytical verified values" - "Test values have been copied from 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_rod()'" - - mock_rod_one = MockRod() - mock_rod_two = MockRod() - - """Moving the rods towards each other with a velocity of 1 along the x-axis.""" - mock_rod_one.velocity_collection = np.array([[1, 0, 0], [1, 0, 0], [1, 0, 0]]) - mock_rod_two.velocity_collection = np.array( - [[-1, 0, 0], [-1, 0, 0], [-1, 0, 0]] - ) - mock_rod_two.position_collection = np.array([[4, 5, 6], [0, 0, 0], [0, 0, 0]]) - ext_contact = ExternalContact(k=0.0, nu=1.0) - ext_contact.apply_forces(mock_rod_one, 0, mock_rod_two, 0) - - assert_allclose( - mock_rod_one.external_forces, - np.array( - [[0, -0.333333, -0.166666], [0, 0, 0], [0, 0, 0]], - ), - atol=1e-6, - ) - assert_allclose( - mock_rod_two.external_forces, - np.array([[0.166666, 0.333333, 0], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - - def test_external_contact_with_two_rods_with_collision_with_k_and_nu(self): - - "Testing External Contact wrapper with two rods with analytical verified values" - "Test values have been copied from 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_rod()'" - - mock_rod_one = MockRod() - mock_rod_two = MockRod() - - """Moving the rods towards each other with a velocity of 1 along the x-axis.""" - mock_rod_one.velocity_collection = np.array([[1, 0, 0], [1, 0, 0], [1, 0, 0]]) - mock_rod_two.velocity_collection = np.array( - [[-1, 0, 0], [-1, 0, 0], [-1, 0, 0]] - ) - mock_rod_two.position_collection = np.array([[4, 5, 6], [0, 0, 0], [0, 0, 0]]) - ext_contact = ExternalContact(k=1.0, nu=1.0) - ext_contact.apply_forces(mock_rod_one, 0, mock_rod_two, 0) - - assert_allclose( - mock_rod_one.external_forces, - np.array( - [[0, -1, -0.5], [0, 0, 0], [0, 0, 0]], - ), - atol=1e-6, - ) - assert_allclose( - mock_rod_two.external_forces, - np.array([[0.5, 1, 0], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - - def test_external_contact_with_two_rods_without_collision(self): - - "Testing External Contact wrapper with two rods with analytical verified values" - - mock_rod_one = MockRod() - mock_rod_two = MockRod() - - "Setting rod two position such that there is no collision" - mock_rod_two.position_collection = np.array( - [[100, 101, 102], [0, 0, 0], [0, 0, 0]] - ) - ext_contact = ExternalContact(k=1.0, nu=1.0) - mock_rod_one_external_forces_before_execution = ( - mock_rod_one.external_forces.copy() - ) - mock_rod_two_external_forces_before_execution = ( - mock_rod_two.external_forces.copy() - ) - ext_contact.apply_forces(mock_rod_one, 0, mock_rod_two, 0) - - assert_allclose( - mock_rod_one.external_forces, mock_rod_one_external_forces_before_execution - ) - assert_allclose( - mock_rod_two.external_forces, mock_rod_two_external_forces_before_execution - ) - - -class TestSelfContact: - def test_self_contact_with_rod_self_collision(self): - - "Testing Self Contact wrapper rod self collision with analytical verified values" - - mock_rod = MockRod() - - "Test values have been copied from 'test_contact_specific_functions.py/test_calculate_contact_forces_self_rod()'" - mock_rod.n_elems = 3 - mock_rod.position_collection = np.array( - [[1, 4, 4, 1], [0, 0, 1, 1], [0, 0, 0, 0]] - ) - mock_rod.radius = np.array([1, 1, 1]) - mock_rod.lengths = np.array([3, 1, 3]) - mock_rod.tangents = np.array( - [[1.0, 0.0, -1.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]] - ) - mock_rod.velocity_collection = np.array( - [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]] - ) - mock_rod.internal_forces = np.array( - [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]] - ) - mock_rod.external_forces = np.array( - [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]] - ) - sel_contact = SelfContact(k=1.0, nu=0.0) - sel_contact.apply_forces(mock_rod, 0, mock_rod, 0) - - assert_allclose( - mock_rod.external_forces, - np.array( - [[0, 0, 0, 0], [-0.333333, -0.666666, 0.666666, 0.333333], [0, 0, 0, 0]] - ), - atol=1e-6, - ) - - def test_self_contact_with_rod_no_self_collision(self): - - "Testing Self Contact wrapper rod no self collision with analytical verified values" - - mock_rod = MockRod() - - "the initially set rod does not have self collision" - mock_rod_external_forces_before_execution = mock_rod.external_forces.copy() - sel_contact = SelfContact(k=1.0, nu=1.0) - sel_contact.apply_forces(mock_rod, 0, mock_rod, 0) - - assert_allclose( - mock_rod.external_forces, mock_rod_external_forces_before_execution - ) diff --git a/tests/test_math/test_timestepper.py b/tests/test_math/test_timestepper.py index bf4b2fa2c..2afa6de65 100644 --- a/tests/test_math/test_timestepper.py +++ b/tests/test_math/test_timestepper.py @@ -245,7 +245,7 @@ def test_explicit_steppers(self, explicit_stepper): # Before stepping, let's extend the interface of the stepper # while providing memory slots - from elastica.systems import make_memory_for_explicit_stepper + from elastica.systems.memory import make_memory_for_explicit_stepper memory_collection = make_memory_for_explicit_stepper(stepper, collective_system) from elastica.timestepper import extend_stepper_interface diff --git a/tests/test_modules/test_connections.py b/tests/test_modules/test_connections.py index 334ece41b..53ed22f46 100644 --- a/tests/test_modules/test_connections.py +++ b/tests/test_modules/test_connections.py @@ -313,15 +313,15 @@ def mock_init(self, *args, **kwargs): ) # Constrain any and all systems - system_collection_with_connections.connect(0, 1).using( - MockConnect, 2, 42 - ) # index based connect + # system_collection_with_connections.connect(0, 1).using( + # MockConnect, 2, 42 + # ) # index based connect system_collection_with_connections.connect(mock_rod_one, mock_rod_two).using( MockConnect, 2, 3 ) # system based connect - system_collection_with_connections.connect(0, mock_rod_one).using( - MockConnect, 1, 2 - ) # index/system based connect + # system_collection_with_connections.connect(0, mock_rod_one).using( + # MockConnect, 1, 2 + # ) # index/system based connect return system_collection_with_connections, MockConnect diff --git a/tests/test_modules/test_contact.py b/tests/test_modules/test_contact.py index 82c41f696..5b9ffe75b 100644 --- a/tests/test_modules/test_contact.py +++ b/tests/test_modules/test_contact.py @@ -240,9 +240,17 @@ def mock_init(self, *args, **kwargs): pass # in place class - MockContact = type( - "MockContact", (self.NoContact, object), {"__init__": mock_init} - ) + class MockContact(self.NoContact): + def __init__(self, *args, **kwargs): + pass + + @property + def _allowed_system_one(self): + return [TestContactMixin.MockRod] + + @property + def _allowed_system_two(self): + return [TestContactMixin.MockRod] # Constrain any and all systems system_collection_with_contacts.detect_contact_between(0, 1).using( @@ -283,9 +291,17 @@ def mock_init(self, *args, **kwargs): pass # in place class - MockContact = type( - "MockContact", (self.NoContact, object), {"__init__": mock_init} - ) + class MockContact(self.NoContact): + def __init__(self, *args, **kwargs): + pass + + @property + def _allowed_system_one(self): + return [TestContactMixin.MockRod] + + @property + def _allowed_system_two(self): + return [TestContactMixin.MockRigidBody] # incorrect order contact system_collection_with_contacts.detect_contact_between( @@ -302,17 +318,12 @@ def test_contact_check_order(self, load_contact_objects_with_incorrect_order): contact_cls, ) = load_contact_objects_with_incorrect_order - mock_rod = self.MockRod(2, 3, 4, 5) - mock_rigid_body = self.MockRigidBody(5.0, 5.0) - with pytest.raises(TypeError) as excinfo: system_collection_with_contacts._finalize_contact() assert ( - "Systems provided to the contact class have incorrect order. \n" - " First system is {0} and second system is {1}. \n" - " If the first system is a rod, the second system can be a rod, rigid body or surface. \n" - " If the first system is a rigid body, the second system can be a rigid body or surface." - ).format(mock_rigid_body.__class__, mock_rod.__class__) == str(excinfo.value) + "System provided (MockRigidBody) must be derived from ['MockRod']" + in str(excinfo.value) + ) @pytest.fixture def load_system_with_rods_in_contact(self, load_system_with_contacts): diff --git a/tests/test_modules/test_memory_block_rigid_body.py b/tests/test_modules/test_memory_block_rigid_body.py index 7eb4882e6..07b11d041 100644 --- a/tests/test_modules/test_memory_block_rigid_body.py +++ b/tests/test_modules/test_memory_block_rigid_body.py @@ -81,7 +81,7 @@ def test_memory_block_rigid_body(n_bodies): memory_block = MemoryBlockRigidBody(systems, system_idx_list) - assert memory_block.n_bodies == n_bodies + assert memory_block.n_systems == n_bodies assert memory_block.n_elems == n_bodies assert memory_block.n_nodes == n_bodies diff --git a/tests/test_modules/test_memory_block_base.py b/tests/test_modules/test_memory_block_utils.py similarity index 98% rename from tests/test_modules/test_memory_block_base.py rename to tests/test_modules/test_memory_block_utils.py index 9dde39234..580addabd 100644 --- a/tests/test_modules/test_memory_block_base.py +++ b/tests/test_modules/test_memory_block_utils.py @@ -3,7 +3,7 @@ import pytest import numpy as np from numpy.testing import assert_array_equal -from elastica.memory_block.memory_block_rod_base import ( +from elastica.memory_block.utils import ( make_block_memory_metadata, make_block_memory_periodic_boundary_metadata, ) diff --git a/tests/test_restart.py b/tests/test_restart.py index 0f814b21a..f19cf0377 100644 --- a/tests/test_restart.py +++ b/tests/test_restart.py @@ -48,13 +48,13 @@ def load_collection(self): return sc, rod_list - def test_restart_save_load(self, load_collection): + def test_restart_save_load(self, tmp_path, load_collection): simulator_class, rod_list = load_collection # Finalize simulator simulator_class.finalize() - directory = "restart_test_data/" + directory = (tmp_path / "restart_test_data").as_posix() time = np.random.rand() # save state @@ -79,7 +79,7 @@ def test_restart_save_load(self, load_collection): assert_allclose(test_value, correct_value) - def run_sim(self, final_time, load_from_restart, save_data_restart): + def run_sim(self, final_time, load_from_restart, save_data_restart, tmp_path): class BaseSimulatorClass( BaseSystemCollection, Constraints, Forcing, Connections, CallBacks ): @@ -118,7 +118,7 @@ class BaseSimulatorClass( # Finalize simulator simulator_class.finalize() - directory = "restart_test_data/" + directory = (tmp_path / "restart_test_data").as_posix() time_step = 1e-4 total_steps = int(final_time / time_step) @@ -149,22 +149,31 @@ class BaseSimulatorClass( return recorded_list @pytest.mark.parametrize("final_time", [0.2, 1.0]) - def test_save_restart_run_sim(self, final_time): + def test_save_restart_run_sim(self, tmp_path, final_time): # First half of simulation _ = self.run_sim( - final_time / 2, load_from_restart=False, save_data_restart=True + final_time / 2, + load_from_restart=False, + save_data_restart=True, + tmp_path=tmp_path, ) # Second half of simulation recorded_list = self.run_sim( - final_time / 2, load_from_restart=True, save_data_restart=False + final_time / 2, + load_from_restart=True, + save_data_restart=False, + tmp_path=tmp_path, ) recorded_list_second_half = recorded_list.copy() # Full simulation recorded_list = self.run_sim( - final_time, load_from_restart=False, save_data_restart=False + final_time, + load_from_restart=False, + save_data_restart=False, + tmp_path=tmp_path, ) recorded_list_full_sim = recorded_list.copy() @@ -197,13 +206,13 @@ def load_collection(self): return sc, cylinder_list - def test_restart_save_load(self, load_collection): + def test_restart_save_load(self, tmp_path, load_collection): simulator_class, cylinder_list = load_collection # Finalize simulator simulator_class.finalize() - directory = "restart_test_data/" + directory = (tmp_path / "restart_test_data").as_posix() time = np.random.rand() # save state diff --git a/tests/test_rigid_body/test_rigid_body_data_structures.py b/tests/test_rigid_body/test_rigid_body_data_structures.py index b482202b8..0ca7715a0 100644 --- a/tests/test_rigid_body/test_rigid_body_data_structures.py +++ b/tests/test_rigid_body/test_rigid_body_data_structures.py @@ -64,7 +64,7 @@ def __init__(self, start_position, start_director): # Givees position, director etc. super(SimpleSystemWithPositionsDirectors, self).__init__() - def _compute_internal_forces_and_torques(self, time): + def compute_internal_forces_and_torques(self, time): pass def update_accelerations(self, time):