feat: add simulator.close() method for collective end-call for callbacks

Author: skim0119

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_CHECKING, Type, Generator, Any, overload
+from typing import TYPE_CHECKING, Type, Generator, Any, overload, Callable
 from typing import final
 from elastica.typing import (
     SystemType,
@@ -72,6 +72,9 @@ def __init__(self) -> None:
         self._feature_group_callback: OperatorGroupFIFO[
             OperatorCallbackType, ModuleProtocol
         ] = OperatorGroupFIFO()
+        self._feature_group_on_close: OperatorGroupFIFO[Callable, ModuleProtocol] = (
+            OperatorGroupFIFO()
+        )
         self._feature_group_finalize: list[OperatorFinalizeType] = []
         # We need to initialize our mixin classes
         super().__init__()
@@ -282,6 +285,15 @@ def apply_callbacks(self, time: np.float64, current_step: int) -> None:
         for func in self._feature_group_callback:
             func(time=time, current_step=current_step)
 
+    @final
+    def close(self) -> None:
+        """
+        Call close functions for all features.
+        Features are registered in _feature_group_on_close.
+        """
+        for func in self._feature_group_on_close:
+            func()
+
 
 if TYPE_CHECKING:
     from .protocol import SystemCollectionProtocol

elastica/modules/callbacks.py

@@ -57,6 +57,7 @@ def collect_diagnostics(
         _callback: ModuleProtocol = _CallBack(sys_idx)
         self._callback_list.append(_callback)
         self._feature_group_callback.append_id(_callback)
+        self._feature_group_on_close.append_id(_callback)
 
         return _callback
 
@@ -70,6 +71,9 @@ def _finalize_callback(self: SystemCollectionWithCallbackProtocol) -> None:
                 callback_instance.make_callback, system=self[sys_id]
             )
             self._feature_group_callback.add_operators(callback, [callback_operator])
+            self._feature_group_on_close.add_operators(
+                callback, [callback_instance.on_close]
+            )
 
         self._callback_list.clear()
         del self._callback_list

elastica/modules/memory_block.py

@@ -51,25 +51,11 @@ def construct_memory_block_structures(
             temp_list_for_rigid_body_systems_idx.append(system_idx)
 
         elif isinstance(sys_to_be_added, SurfaceBase):
+            # TODO: Surface type is passive system
             pass
-            # surface_system = cast(SurfaceType, sys_to_be_added)
-            # raise NotImplementedError(
-            #     "Surfaces are not yet implemented in memory block construction."
-            # )
 
         else:
-            raise TypeError(
-                "{0}\n"
-                "is not a system passing validity\n"
-                "checks for constructing block structure. If you are sure that\n"
-                "{0}\n"
-                "satisfies all criteria for being a system, please add\n"
-                "it here with correct memory block implementation.\n"
-                "The allowed types are\n"
-                "{1} {2} {3}".format(
-                    sys_to_be_added.__class__, RodBase, RigidBodyBase, SurfaceBase
-                )
-            )
+            continue  # No error:: any typechecking should be finished by BaseSystemCollection._check_type
 
     if temp_list_for_cosserat_rod_systems:
         _memory_blocks.append(

elastica/modules/protocol.py

@@ -1,4 +1,4 @@
-from typing import Protocol, Generator, TypeVar, Any, Type, overload, Iterator
+from typing import Protocol, Generator, TypeVar, Any, Type, overload, Iterator, Callable
 from typing import TYPE_CHECKING
 from typing_extensions import Self  # python 3.11: from typing import Self
 
@@ -80,6 +80,10 @@ def apply_callbacks(self, time: np.float64, current_step: int) -> None: ...
 
     def finalize(self) -> None: ...
 
+    _feature_group_on_close: "OperatorGroupFIFO[Callable, ModuleProtocol]"
+
+    def close(self) -> None: ...
+
 
 # Mixin Protocols (Used to type Self)
 class ConnectedSystemCollectionProtocol(SystemCollectionProtocol, Protocol):

elastica/timestepper/protocol.py

@@ -25,7 +25,10 @@ def n_stages(self) -> int: ...
     def step_methods(self) -> SteppersOperatorsType: ...
 
     def step(
-        self, SystemCollection: SystemCollectionType, time: np.float64, dt: np.float64
+        self,
+        SystemCollection: SystemCollectionType,
+        time: np.float64 | float,
+        dt: np.float64 | float,
     ) -> np.float64: ...
 
     def step_single_instance(

elastica/timestepper/symplectic_steppers.py

@@ -66,11 +66,15 @@ def n_stages(self: SymplecticStepperProtocol) -> int:
     def step(
         self: SymplecticStepperProtocol,
         SystemCollection: SystemCollectionType,
-        time: np.float64,
-        dt: np.float64,
+        time: np.float64 | float,
+        dt: np.float64 | float,
     ) -> np.float64:
         return SymplecticStepperMixin.do_step(
-            self, self.steps_and_prefactors, SystemCollection, time, dt
+            self,
+            self.steps_and_prefactors,
+            SystemCollection,
+            np.float64(time),
+            np.float64(dt),
         )
 
     # TODO: Merge with .step method in the future.