diff --git a/elastica/modules/contact.py b/elastica/modules/contact.py index 76491664..ed942de2 100644 --- a/elastica/modules/contact.py +++ b/elastica/modules/contact.py @@ -9,14 +9,13 @@ from typing_extensions import Self from elastica.typing import ( SystemIdxType, - OperatorFinalizeType, + OperatorType, StaticSystemType, SystemType, ) from .protocol import SystemCollectionProtocol, ModuleProtocol import logging -import functools import numpy as np @@ -79,17 +78,6 @@ def _finalize_contact(self: SystemCollectionProtocol) -> None: # (first_rod_idx, second_rod_idx) # to apply the contacts to - def apply_contact( - time: np.float64, - contact_instance: NoContact, - first_sys_idx: SystemIdxType, - second_sys_idx: SystemIdxType, - ) -> None: - contact_instance.apply_contact( - system_one=self[first_sys_idx], - system_two=self[second_sys_idx], - ) - for contact in self._contacts: first_sys_idx, second_sys_idx = contact.id() contact_instance = contact.instantiate() @@ -98,12 +86,11 @@ def apply_contact( self[first_sys_idx], self[second_sys_idx], ) - func = functools.partial( - apply_contact, - contact_instance=contact_instance, - first_sys_idx=first_sys_idx, - second_sys_idx=second_sys_idx, + func: OperatorType = lambda time: contact_instance.apply_contact( + system_one=self[first_sys_idx], + system_two=self[second_sys_idx], ) + self._feature_group_synchronize.add_operators(contact, [func]) if not self._feature_group_synchronize.is_last(contact):