From 3b4b73544a99105e9dc3551ac4a29434ee92eb03 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 20 Jan 2025 17:07:03 -0600 Subject: [PATCH 1/4] add different order symplectic steppers --- elastica/__init__.py | 8 +- elastica/timestepper/__init__.py | 11 +- elastica/timestepper/protocol.py | 19 +- elastica/timestepper/symplectic_steppers.py | 251 +++++++++++++++++++- 4 files changed, 270 insertions(+), 19 deletions(-) diff --git a/elastica/__init__.py b/elastica/__init__.py index 21dfddb72..f7095cfe4 100644 --- a/elastica/__init__.py +++ b/elastica/__init__.py @@ -75,11 +75,15 @@ from elastica.utils import isqrt from elastica.timestepper import ( integrate, + extend_stepper_interface, PositionVerlet, + VelocityVerlet, + ThirdOrderSymplectic, + FourthOrderSymplectic, + SemiImplicitEuler, PEFRL, - RungeKutta4, EulerForward, - extend_stepper_interface, + RungeKutta4, ) from elastica.memory_block.memory_block_rigid_body import MemoryBlockRigidBody from elastica.memory_block.memory_block_rod import MemoryBlockCosseratRod diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 25f883411..b0853ed80 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -8,7 +8,14 @@ from elastica.systems import is_system_a_collection -from .symplectic_steppers import PositionVerlet, PEFRL +from .symplectic_steppers import ( + PositionVerlet, + SemiImplicitEuler, + VelocityVerlet, + ThirdOrderSymplectic, + FourthOrderSymplectic, + PEFRL, +) from .explicit_steppers import RungeKutta4, EulerForward from .protocol import StepperProtocol, SymplecticStepperProtocol @@ -16,7 +23,7 @@ # Deprecated: Remove in the future version # Many script still uses this method to control timestep. Keep it for backward compatibility def extend_stepper_interface( - stepper: StepperProtocol, system_collection: SystemCollectionType + stepper: SymplecticStepperProtocol, system_collection: SystemCollectionType ) -> Tuple[ Callable[ [StepperProtocol, SystemCollectionType, np.float64, np.float64], np.float64 diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index 1a64c7254..6900843ec 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -19,10 +19,7 @@ class StepperProtocol(Protocol): def __init__(self) -> None: ... - @property - def n_stages(self) -> int: ... - - def step_methods(self) -> SteppersOperatorsType: ... + def build_step_methods(self) -> SteppersOperatorsType: ... def step( self, SystemCollection: SystemCollectionType, time: np.float64, dt: np.float64 @@ -34,16 +31,20 @@ def step_single_instance( class SymplecticStepperProtocol(StepperProtocol, Protocol): - """symplectic stepper protocol.""" + """Symplectic stepper protocol.""" def get_steps(self) -> list[OperatorType]: ... def get_prefactors(self) -> list[OperatorType]: ... - -class MemoryProtocol(Protocol): - @property - def initial_state(self) -> bool: ... + @staticmethod + def do_step( + TimeStepper: "StepperProtocol", + steps_and_prefactors: SteppersOperatorsType, + SystemCollection: SystemCollectionType, + time: np.float64, + dt: np.float64, + ) -> np.float64: ... class ExplicitStepperProtocol(StepperProtocol, Protocol): diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 3937e337b..7771f8b7f 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -33,9 +33,11 @@ class SymplecticStepperMixin: def __init__(self: SymplecticStepperProtocol): - self.steps_and_prefactors: Final[SteppersOperatorsType] = self.step_methods() + self.steps_and_prefactors: Final[SteppersOperatorsType] = ( + self.build_step_methods() + ) - def step_methods(self: SymplecticStepperProtocol) -> SteppersOperatorsType: + def build_step_methods(self: SymplecticStepperProtocol) -> SteppersOperatorsType: # Let the total number of steps for the Symplectic method # be (2*n + 1) (for time-symmetry). _steps: list[OperatorType] = self.get_steps() @@ -62,10 +64,6 @@ def no_operation(*args: Any) -> None: ) ) - @property - def n_stages(self: SymplecticStepperProtocol) -> int: - return len(self.steps_and_prefactors) - def step( self: SymplecticStepperProtocol, SystemCollection: SystemCollectionType, @@ -202,6 +200,247 @@ def _first_dynamic_step( ) +class VelocityVerlet(SymplecticStepperMixin): + """ + Velocity Verlet symplectic time stepper class. + """ + + def get_steps(self) -> list[OperatorType]: + return [ + self._no_operation, + self._first_dynamic_step, + self._first_kinematic_step, + self._second_dynamic_step, + self._no_operation, + ] + + def get_prefactors(self) -> list[OperatorType]: + return [ + self._first_prefactor, + self._second_prefactor, + self._third_prefactor, + ] + + def _no_operation(*args: Any) -> None: + pass + + def _first_prefactor(self, dt: np.float64) -> np.float64: + return 0.0 + + def _second_prefactor(self, dt: np.float64) -> np.float64: + return 0.0 + + def _third_prefactor(self, dt: np.float64) -> np.float64: + return 1.0 * dt + + def _first_kinematic_step( + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 + ) -> None: + overload_operator_kinematic_numba( + System.n_nodes, + dt, + System.kinematic_states.position_collection, + System.kinematic_states.director_collection, + System.velocity_collection, + System.omega_collection, + ) + + def _first_dynamic_step( + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 + ) -> None: + overload_operator_dynamic_numba( + System.dynamic_states.rate_collection, + System.dynamic_rates(time, dt / 2.0), + ) + + def _second_dynamic_step( + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 + ) -> None: + overload_operator_dynamic_numba( + System.dynamic_states.rate_collection, + System.dynamic_rates(time, dt / 2.0), + ) + + +class SemiImplicitEuler(SymplecticStepperMixin): + """ + Semi-Implicit (Explicit) symplectic time stepper class. + Slightly (30%) faster than the Position Verlet with reduced + convergence. + First order symplectic euler method. + """ + + def get_steps(self) -> list[OperatorType]: + return [ + self._no_operation, + self._dynamic_step, + self._kinematic_step, + ] + + def get_prefactors(self) -> list[OperatorType]: + return [ + self._first_prefactor, + self._second_prefactor, + ] + + def _no_operation(*args: Any) -> None: + pass + + def _first_prefactor(self, dt: np.float64) -> np.float64: + return 0.0 + + def _second_prefactor(self, dt: np.float64) -> np.float64: + return 1.0 * dt + + def _kinematic_step( + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 + ) -> None: + overload_operator_kinematic_numba( + System.n_nodes, + dt, + System.kinematic_states.position_collection, + System.kinematic_states.director_collection, + System.velocity_collection, + System.omega_collection, + ) + + def _dynamic_step( + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 + ) -> None: + overload_operator_dynamic_numba( + System.dynamic_states.rate_collection, + System.dynamic_rates(time, dt), + ) + + +class ThirdOrderSymplectic(SymplecticStepperMixin): + """ + Third order symplectic time stepper class based on Ruth 1983. + """ + + def get_steps(self) -> list[OperatorType]: + return [ + self._kinematic_step_with_prefactor(2.0 / 3.0), + self._dynamic_step_with_prefactor(7.0 / 24.0), + self._kinematic_step_with_prefactor(-2.0 / 3.0), + self._dynamic_step_with_prefactor(3.0 / 4.0), + self._kinematic_step_with_prefactor(1.0), + self._dynamic_step_with_prefactor(-1.0 / 24.0), + self._no_operation, + ] + + def get_prefactors(self) -> list[OperatorType]: + return [ + self._prefactor(0.0), + self._prefactor(0.0), + self._prefactor(0.0), + self._prefactor(1.0), + ] + + def _no_operation(*args: Any) -> None: + pass + + def _prefactor(self, factor: np.float64) -> np.float64: + def func(dt: np.float64) -> np.float64: + return dt * factor + + return func + + def _kinematic_step_with_prefactor(self, factor: np.float64) -> OperatorType: + def func( + System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 + ) -> None: + overload_operator_kinematic_numba( + System.n_nodes, + factor * dt, + System.kinematic_states.position_collection, + System.kinematic_states.director_collection, + System.velocity_collection, + System.omega_collection, + ) + + return func + + def _dynamic_step_with_prefactor(self, factor: np.float64) -> OperatorType: + def func( + System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 + ) -> None: + overload_operator_dynamic_numba( + System.dynamic_states.rate_collection, + System.dynamic_rates(time, factor * dt), + ) + + return func + + +class FourthOrderSymplectic(SymplecticStepperMixin): + """ + Fourth order symplectic time stepper class based on Ruth 1983. + """ + + def get_steps(self) -> list[OperatorType]: + c1 = c4 = 1.0 / (2.0 * (2.0 - 2.0 ** (1.0 / 3.0))) + c2 = c3 = (1.0 - 2.0 ** (1.0 / 3.0)) / (2.0 * (2.0 - 2.0 ** (1.0 / 3.0))) + d1 = d3 = 1.0 / (2.0 - 2.0 ** (1.0 / 3.0)) + d2 = -(2.0 ** (1.0 / 3.0)) / (2.0 - 2.0 ** (1.0 / 3.0)) + d4 = 0.0 + return [ + self._kinematic_step_with_prefactor(c4), + self._dynamic_step_with_prefactor(d4), + self._kinematic_step_with_prefactor(c3), + self._dynamic_step_with_prefactor(d3), + self._kinematic_step_with_prefactor(c2), + self._dynamic_step_with_prefactor(d2), + self._kinematic_step_with_prefactor(c1), + self._dynamic_step_with_prefactor(d1), + self._no_operation, + ] + + def get_prefactors(self) -> list[OperatorType]: + return [ + self._prefactor(0.0), + self._prefactor(0.0), + self._prefactor(0.0), + self._prefactor(0.0), + self._prefactor(1.0), + ] + + def _no_operation(*args: Any) -> None: + pass + + def _prefactor(self, factor: np.float64) -> np.float64: + def func(dt: np.float64) -> np.float64: + return dt * factor + + return func + + def _kinematic_step_with_prefactor(self, factor: np.float64) -> OperatorType: + def func( + System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 + ) -> None: + overload_operator_kinematic_numba( + System.n_nodes, + factor * dt, + System.kinematic_states.position_collection, + System.kinematic_states.director_collection, + System.velocity_collection, + System.omega_collection, + ) + + return func + + def _dynamic_step_with_prefactor(self, factor: np.float64) -> OperatorType: + def func( + System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 + ) -> None: + overload_operator_dynamic_numba( + System.dynamic_states.rate_collection, + System.dynamic_rates(time, factor * dt), + ) + + return func + + class PEFRL(SymplecticStepperMixin): """ Position Extended Forest-Ruth Like Algorithm of From 3964e4ead742ca81ccdb8952178e5ef7b981b15b Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 20 Jan 2025 17:08:58 -0600 Subject: [PATCH 2/4] modify butterfly case to test multiple timestepper --- examples/ButterflyCase/butterfly.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/examples/ButterflyCase/butterfly.py b/examples/ButterflyCase/butterfly.py index d8ffd0b82..792ae74f8 100644 --- a/examples/ButterflyCase/butterfly.py +++ b/examples/ButterflyCase/butterfly.py @@ -1,7 +1,7 @@ import numpy as np from matplotlib import pyplot as plt from matplotlib.colors import to_rgb - +from tqdm import tqdm import elastica as ea from elastica.utils import MaxDimension @@ -118,12 +118,19 @@ def make_callback(self, system, time, current_step: int): butterfly_sim.finalize() timestepper = ea.PositionVerlet() -# timestepper = PEFRL() +# timestepper = ea.SemiImplicitEuler() +# timestepper = ea.VelocityVerlet() +# timestepper = ea.ThirdOrderSymplectic() +# timestepper = ea.FourthOrderSymplectic() +# timestepper = ea.PEFRL() +# timestepper = ea.EulerForward() +# timestepper = ea.RungeKutta4() dt = 0.01 * dl total_steps = int(final_time / dt) -print("Total steps", total_steps) -ea.integrate(timestepper, butterfly_sim, final_time, total_steps) +time = 0.0 +for step in tqdm(range(total_steps)): + time = timestepper.step(butterfly_sim, time, dt) if PLOT_FIGURE: # Plot the histories From 390895c5fedc7a5def3451ef25676de27333fd7f Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 20 Jan 2025 17:09:40 -0600 Subject: [PATCH 3/4] restore explicit steppers: EulerForward and RK4 --- elastica/timestepper/explicit_steppers.py | 394 ++++++++-------------- elastica/timestepper/protocol.py | 64 +--- 2 files changed, 137 insertions(+), 321 deletions(-) diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 5fb1531f7..780e90498 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -3,7 +3,7 @@ from typing import Any import numpy as np -from copy import copy +from copy import copy, deepcopy from elastica.typing import ( SystemType, @@ -13,80 +13,11 @@ StateType, ) from elastica.systems.protocol import ExplicitSystemProtocol -from .protocol import ExplicitStepperProtocol, MemoryProtocol - - -""" -Developer Note --------------- -## Motivation for choosing _Mixin classes below - -The constraint/problem is that we do not know what -`System` we are integrating apriori. For a single -standalone `System` (which defines a `__call__` -operator and has its own states), we should just -step it like a single system. - -Instead if we get a `SystemCollection` made up of -bunch of smaller systems (like Cosserat Rods), we now -need to loop over all these smaller systems and perform -state updates there. Not only that we may also need -to communicate between such smaller systems. - -One way to solve this issue is to give the integrator -two methods: - -- `do_step`, which does the time-stepping for only a -`System` -- `do_system_step` which does the time-stepping for -a `SystemCollection` - -The problem with this approach is that -1. We have more methods than we actually use -(indeed we can only integrate either a `System` or -a `SystemCollection` but not both) -2. From an interface point of view, its ugly and not -graceful (at least IMO). - -The second approach is what I have chosen here, -which is to create two mixin classes : one for -integrating `System` and one for integrating -`SystemCollection`. And then depending upon the runtime -type of the object to be integrated, we can dynamically -mixin the required class. - -This approach overcomes the disadvantages of the -previous approach (as there's only one `do_step` method -associated with a Stepper at any given point of time), -at the expense of being a tad bit harder to understand -(which this documentation will hopefully fix). In essence, -we "smartly" use a mixin class to define the necessary -`do_step` method, which the `integrate` function then uses. -""" - - -class EulerForwardMemory: - def __init__(self, initial_state: StateType) -> None: - self.initial_state = initial_state - - -class RungeKutta4Memory: - """ - Stores all states of Rk within the time-stepper. Works as long as the states - are all one big numpy array, made possible by carefully using views. - - Convenience wrapper around Stateless that provides memory - """ - - def __init__( - self, - initial_state: StateType, - ) -> None: - self.initial_state = initial_state - self.k_1 = initial_state - self.k_2 = initial_state - self.k_3 = initial_state - self.k_4 = initial_state +from elastica.rod.data_structures import ( + overload_operator_kinematic_numba, + overload_operator_dynamic_numba, +) +from .protocol import ExplicitStepperProtocol class ExplicitStepperMixin: @@ -94,21 +25,45 @@ class ExplicitStepperMixin: Can also be used as a mixin with optional cls argument below """ - def __init__(self: ExplicitStepperProtocol): - self.steps_and_prefactors = self.step_methods() + def system_inplace_update( + self: ExplicitStepperProtocol, + system1: ExplicitSystemProtocol, + system2: ExplicitSystemProtocol, + time: np.float64, + dt: np.float64, + ): + """ + y_n+1 = y_n + prefac * f(y_n, t_n) + """ + overload_operator_kinematic_numba( + system1.n_nodes, + dt, + system1.kinematic_states.position_collection, + system1.kinematic_states.director_collection, + system2.velocity_collection, + system2.omega_collection, + ) + + overload_operator_dynamic_numba( + system1.dynamic_states.rate_collection, + system2.dynamic_rates(time, dt), + ) + + def system_rate_update( + self, SystemCollection: SystemCollectionType, time: np.float64 + ): + # Constrain + SystemCollection.constrain_values(time) - def step_methods(self: ExplicitStepperProtocol) -> SteppersOperatorsType: - stages = self.get_stages() - updates = self.get_updates() + # We need internal forces and torques because they are used by interaction module. + for system in SystemCollection.block_systems(): + system.compute_internal_forces_and_torques(time) - assert len(stages) == len( - updates - ), "Number of stages and updates should be equal to one another" - return tuple(zip(stages, updates)) + # Add external forces, controls etc. + SystemCollection.synchronize(time) - @property - def n_stages(self: ExplicitStepperProtocol) -> int: - return len(self.steps_and_prefactors) + # Constrain only rates + SystemCollection.constrain_rates(time) def step( self: ExplicitStepperProtocol, @@ -116,49 +71,19 @@ def step( time: np.float64, dt: np.float64, ) -> np.float64: - if isinstance( - self, EulerForward - ): # TODO: Cleanup - use depedency injection instead - Memory = EulerForwardMemory - elif isinstance(self, RungeKutta4): - Memory = RungeKutta4Memory # type: ignore[assignment] - else: - raise NotImplementedError(f"Memory class not defined for {self}") - memory_collection = tuple( - [Memory(initial_state=system.state) for system in SystemCollection] - ) - return ExplicitStepperMixin.do_step(self, self.steps_and_prefactors, SystemCollection, memory_collection, time, dt) # type: ignore[attr-defined] + self.stage(SystemCollection, time, dt) - @staticmethod - def do_step( - TimeStepper: ExplicitStepperProtocol, - steps_and_prefactors: SteppersOperatorsType, - SystemCollection: SystemCollectionType, - MemoryCollection: Any, # TODO - time: np.float64, - dt: np.float64, - ) -> np.float64: - for stage, update in steps_and_prefactors: - SystemCollection.synchronize(time) - for system, memory in zip(SystemCollection[:-1], MemoryCollection[:-1]): - stage(system, memory, time, dt) - _ = update(system, memory, time, dt) + # Timestep update + next_time = time + dt - stage(SystemCollection[-1], MemoryCollection[-1], time, dt) - time = update(SystemCollection[-1], MemoryCollection[-1], time, dt) - return time + # Call back function, will call the user defined call back functions and store data + SystemCollection.apply_callbacks(next_time, round(next_time / dt)) - def step_single_instance( - self: ExplicitStepperProtocol, - System: SystemType, - Memory: MemoryProtocol, - time: np.float64, - dt: np.float64, - ) -> np.float64: - for stage, update in self.steps_and_prefactors: - stage(System, Memory, time, dt) - time = update(System, Memory, time, dt) - return time + # Zero out the external forces and torques + for system in SystemCollection.block_systems(): + system.zeroed_out_external_forces_and_torques(next_time) + + return next_time class EulerForward(ExplicitStepperMixin): @@ -166,152 +91,97 @@ class EulerForward(ExplicitStepperMixin): Classical Euler Forward stepper. Stateless, coordinates operations only. """ - def get_stages(self) -> list[OperatorType]: - return [self._first_stage] - - def get_updates(self) -> list[OperatorType]: - return [self._first_update] - - def _first_stage( + def stage( self, - System: ExplicitSystemProtocol, - Memory: EulerForwardMemory, + SystemCollection: SystemCollectionType, time: np.float64, dt: np.float64, ) -> None: - pass - - def _first_update( - self, - System: ExplicitSystemProtocol, - Memory: EulerForwardMemory, - time: np.float64, - dt: np.float64, - ) -> np.float64: - System.state += dt * System(time, dt) # type: ignore[arg-type] - return time + dt + self.system_rate_update(SystemCollection, time) + for system in SystemCollection.block_systems(): + self.system_inplace_update(system, system, time, dt) class RungeKutta4(ExplicitStepperMixin): """ - Stateless runge-kutta4. coordinates operations only, memory needs - to be externally managed and allocated. + Stateless runge-kutta4. coordinates operations only. """ - def get_stages(self) -> list[OperatorType]: - return [ - self._first_stage, - self._second_stage, - self._third_stage, - self._fourth_stage, - ] - - def get_updates(self) -> list[OperatorType]: - return [ - self._first_update, - self._second_update, - self._third_update, - self._fourth_update, - ] - - # These methods should be static, but because we need to enable automatic - # discovery in ExplicitStepper, these are bound to the RungeKutta4 class - # For automatic discovery, the order of declaring stages here is very important - def _first_stage( - self, - System: ExplicitSystemProtocol, - Memory: RungeKutta4Memory, - time: np.float64, - dt: np.float64, - ) -> None: - Memory.initial_state = copy(System.state) - Memory.k_1 = dt * System(time, dt) # type: ignore[operator, assignment] - - def _first_update( - self, - System: ExplicitSystemProtocol, - Memory: RungeKutta4Memory, - time: np.float64, - dt: np.float64, - ) -> np.float64: - # prepare for next stage - System.state = Memory.initial_state + 0.5 * Memory.k_1 # type: ignore[operator] - return time + 0.5 * dt - - def _second_stage( + def stage( self, - System: ExplicitSystemProtocol, - Memory: RungeKutta4Memory, - time: np.float64, - dt: np.float64, - ) -> None: - Memory.k_2 = dt * System(time, dt) # type: ignore[operator, assignment] - - def _second_update( - self, - System: ExplicitSystemProtocol, - Memory: RungeKutta4Memory, - time: np.float64, - dt: np.float64, - ) -> np.float64: - # prepare for next stage - System.state = Memory.initial_state + 0.5 * Memory.k_2 # type: ignore[operator] - return time - - def _third_stage( - self, - System: ExplicitSystemProtocol, - Memory: RungeKutta4Memory, - time: np.float64, - dt: np.float64, - ) -> None: - Memory.k_3 = dt * System(time, dt) # type: ignore[operator, assignment] - - def _third_update( - self, - System: ExplicitSystemProtocol, - Memory: RungeKutta4Memory, - time: np.float64, - dt: np.float64, - ) -> np.float64: - # prepare for next stage - System.state = Memory.initial_state + Memory.k_3 # type: ignore[operator] - return time + 0.5 * dt - - def _fourth_stage( - self, - System: ExplicitSystemProtocol, - Memory: RungeKutta4Memory, + SystemCollection: SystemCollectionType, time: np.float64, dt: np.float64, ) -> None: - Memory.k_4 = dt * System(time, dt) # type: ignore[operator, assignment] - - def _fourth_update( - self, - System: ExplicitSystemProtocol, - Memory: RungeKutta4Memory, - time: np.float64, - dt: np.float64, - ) -> np.float64: - # prepare for next stage - System.state = ( - Memory.initial_state - + (Memory.k_1 + 2.0 * Memory.k_2 + 2.0 * Memory.k_3 + Memory.k_4) / 6.0 # type: ignore[operator] - ) - return time - - -# class ExplicitLinearExponentialIntegrator( -# _LinearExponentialIntegratorMixin, ExplicitStepper -# ): -# def __init__(self): -# _LinearExponentialIntegratorMixin.__init__(self) -# ExplicitStepper.__init__(self, _LinearExponentialIntegratorMixin) -# -# -# class StatefulLinearExponentialIntegrator(_StatefulStepper): -# def __init__(self): -# super(StatefulLinearExponentialIntegrator, self).__init__() -# self.stepper = ExplicitLinearExponentialIntegrator() -# self.linear_operator = None + """ + In-place update of system states with Runge-Kutta 4 integration + """ + k1_system = SystemCollection # Alias + k2_system = deepcopy(SystemCollection) + k3_system = deepcopy(SystemCollection) + k4_system = deepcopy(SystemCollection) + + # First stage + self.system_rate_update(k1_system, time) + + # Second stage + for system1, system2 in zip( + k2_system.block_systems(), k1_system.block_systems() + ): + self.system_inplace_update(system1, system2, time, dt / 2.0) + self.system_rate_update(k2_system, time + dt / 2.0) + + # Third stage + for system1, system2 in zip( + k3_system.block_systems(), k2_system.block_systems() + ): + self.system_inplace_update(system1, system2, time, dt / 2.0) + self.system_rate_update(k3_system, time + dt / 2.0) + + # Fourth stage + for system1, system2 in zip( + k4_system.block_systems(), k3_system.block_systems() + ): + self.system_inplace_update(system1, system2, time, dt) + self.system_rate_update(k3_system, time + dt) + + # Combine stages + for system1, k1, k2, k3, k4 in zip( + SystemCollection.block_systems(), + k1_system.block_systems(), + k2_system.block_systems(), + k3_system.block_systems(), + k4_system.block_systems(), + ): + velocity_update = ( + k1.velocity_collection + + 2 * k2.velocity_collection + + 2 * k3.velocity_collection + + k4.velocity_collection + ) / 6 + omega_update = ( + k1.omega_collection + + 2 * k2.omega_collection + + 2 * k3.omega_collection + + k4.omega_collection + ) / 6 + dynamic_rates_update = ( + k1.dynamic_rates(time, dt / 6.0) + + k2.dynamic_rates(time, dt / 3.0) + + k3.dynamic_rates(time, dt / 3.0) + + k4.dynamic_rates(time, dt / 6.0) + ) # Time is dummy. + + overload_operator_kinematic_numba( + system1.n_nodes, + dt, + system1.kinematic_states.position_collection, + system1.kinematic_states.director_collection, + velocity_update, + omega_update, + ) + + overload_operator_dynamic_numba( + system1.dynamic_states.rate_collection, + dynamic_rates_update, + ) diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index 6900843ec..6d41bbbbb 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -48,62 +48,8 @@ def do_step( class ExplicitStepperProtocol(StepperProtocol, Protocol): - """symplectic stepper protocol.""" - - def get_stages(self) -> list[OperatorType]: ... - - def get_updates(self) -> list[OperatorType]: ... - - -# class _LinearExponentialIntegratorMixin: -# """ -# Linear Exponential integrator mixin wrapper. -# """ -# -# def __init__(self): -# pass -# -# def _do_stage(self, System, Memory, time, dt): -# # TODO : Make more general, system should not be calculating what the state -# # transition matrix directly is, but rather it should just give -# Memory.linear_operator = System.get_linear_state_transition_operator(time, dt) -# -# def _do_update(self, System, Memory, time, dt): -# # FIXME What's the right formula when doing update? -# # System.linearly_evolving_state = _batch_matmul( -# # System.linearly_evolving_state, -# # Memory.linear_operator -# # ) -# System.linearly_evolving_state = np.einsum( -# "ijk,ljk->ilk", System.linearly_evolving_state, Memory.linear_operator -# ) -# return time + dt -# -# def _first_prefactor(self, dt): -# """Prefactor call to satisfy interface of SymplecticStepper. Should never -# be used in actual code. -# -# Parameters -# ---------- -# dt : the time step of simulation -# -# Raises -# ------ -# RuntimeError -# """ -# raise RuntimeError( -# "Symplectic prefactor of LinearExponentialIntegrator should not be called!" -# ) -# -# # Code repeat! -# # Easy to avoid, but keep for performance. -# def _do_one_step(self, System, time, prefac): -# System.linearly_evolving_state = np.einsum( -# "ijk,ljk->ilk", -# System.linearly_evolving_state, -# System.get_linear_state_transition_operator(time, prefac), -# ) -# return ( -# time # TODO fix hack that treats time separately here. Shuold be time + dt -# ) -# # return time + dt + """Explicit stepper protocol.""" + + def stage( + self, System: SystemType, time: np.float64, dt: np.float64 + ) -> np.float64: ... From 3eddef6c2de0108107228531f5a380f78cf89993 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 20 Jan 2025 17:10:58 -0600 Subject: [PATCH 4/4] add TODO for future improvement --- elastica/timestepper/explicit_steppers.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 780e90498..c3903dca8 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -1,9 +1,7 @@ __doc__ = """Explicit timesteppers and concepts""" -from typing import Any - import numpy as np -from copy import copy, deepcopy +from copy import deepcopy from elastica.typing import ( SystemType, @@ -116,6 +114,10 @@ def stage( """ In-place update of system states with Runge-Kutta 4 integration """ + + # TODO: Try to avoid deepcopying the whole system collection + # Related to #33: save/restart module + # Maybe search for low-storage RK scheme k1_system = SystemCollection # Alias k2_system = deepcopy(SystemCollection) k3_system = deepcopy(SystemCollection)