wip: testing functional operator listing

GazzolaLab · Jun 30, 2024 · 271c0fc · 271c0fc
1 parent cddfa60
commit 271c0fc
Showing 1 changed file with 5 additions and 18 deletions.
diff --git 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):