-
Notifications
You must be signed in to change notification settings - Fork 109
/
symplectic_steppers.py
322 lines (254 loc) · 10.2 KB
/
symplectic_steppers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
__doc__ = """Symplectic time steppers and concepts for integrating the kinematic and dynamic equations of rod-like objects. """
import numpy as np
# from elastica._elastica_numba._timestepper._symplectic_steppers import (
# SymplecticStepperTag,
# PositionVerlet,
# PEFRL,
# )
# from elastica.timestepper._stepper_interface import (
# _TimeStepper,
# _LinearExponentialIntegratorMixin,
# )
from elastica.rod.data_structures import (
overload_operator_kinematic_numba,
overload_operator_dynamic_numba,
)
"""
Developer Note
--------------
For the reasons why we define Mixin classes here, the developer
is referred to the same section on `explicit_steppers.py`.
"""
class _SystemInstanceStepper:
@staticmethod
def do_step(
TimeStepper, _steps_and_prefactors, System, time: np.float64, dt: np.float64
):
for (kin_prefactor, kin_step, dyn_step) in _steps_and_prefactors[:-1]:
kin_step(TimeStepper, System, time, dt)
time += kin_prefactor(TimeStepper, dt)
System.update_internal_forces_and_torques(time)
dyn_step(TimeStepper, System, time, dt)
# Peel the last kinematic step and prefactor alone
last_kin_prefactor = _steps_and_prefactors[-1][0]
last_kin_step = _steps_and_prefactors[-1][1]
last_kin_step(TimeStepper, System, time, dt)
return time + last_kin_prefactor(TimeStepper, dt)
class _SystemCollectionStepper:
"""
Symplectic stepper collection class
"""
@staticmethod
def do_step(
TimeStepper,
_steps_and_prefactors,
SystemCollection,
time: np.float64,
dt: np.float64,
):
"""
Function for doing symplectic stepper over the user defined rods (system).
Parameters
----------
SystemCollection: rod object
time: float
dt: float
Returns
-------
"""
for (kin_prefactor, kin_step, dyn_step) in _steps_and_prefactors[:-1]:
for system in SystemCollection._memory_blocks:
kin_step(TimeStepper, system, time, dt)
time += kin_prefactor(TimeStepper, dt)
# Constrain only values
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)
# system.update_internal_forces_and_torques()
# Add external forces, controls etc.
SystemCollection.synchronize(time)
for system in SystemCollection._memory_blocks:
dyn_step(TimeStepper, system, time, dt)
# Constrain only rates
SystemCollection.constrain_rates(time)
# Peel the last kinematic step and prefactor alone
last_kin_prefactor = _steps_and_prefactors[-1][0]
last_kin_step = _steps_and_prefactors[-1][1]
for system in SystemCollection._memory_blocks:
last_kin_step(TimeStepper, system, time, dt)
time += last_kin_prefactor(TimeStepper, dt)
SystemCollection.constrain_values(time)
# Call back function, will call the user defined call back functions and store data
SystemCollection.apply_callbacks(time, int(time / dt))
# Zero out the external forces and torques
for system in SystemCollection._memory_blocks:
system.reset_external_forces_and_torques(time)
return time
class SymplecticStepperMethods:
def __init__(self, timestepper_instance):
take_methods_from = timestepper_instance
# Let the total number of steps for the Symplectic method
# be (2*n + 1) (for time-symmetry). What we do is collect
# the first n + 1 entries down in _steps and _prefac below, and then
# reverse and append it to itself.
self._steps = [
v
for (k, v) in take_methods_from.__class__.__dict__.items()
if k.endswith("step")
]
# Prefac here is necessary because the linear-exponential integrator
# needs only the prefactor and not the dt.
self._prefactors = [
v
for (k, v) in take_methods_from.__class__.__dict__.items()
if k.endswith("prefactor")
]
# # We are getting function named as _update_internal_forces_torques from dictionary,
# # it turns a list.
# self._update_internal_forces_torques = [
# v
# for (k, v) in take_methods_from.__class__.__dict__.items()
# if k.endswith("forces_torques")
# ]
def mirror(in_list):
"""Mirrors an input list ignoring the last element
If steps = [A, B, C]
then this call makes it [A, B, C, B, A]
Parameters
----------
in_list : input list to be mirrored, modified in-place
Returns
-------
"""
# syntax is very ugly
if len(in_list) > 1:
in_list.extend(in_list[-2::-1])
elif in_list:
in_list.append(in_list[0])
mirror(self._steps)
mirror(self._prefactors)
assert (
len(self._steps) == 2 * len(self._prefactors) - 1
), "Size mismatch in the number of steps and prefactors provided for a Symplectic Stepper!"
self._kinematic_steps = self._steps[::2]
self._dynamic_steps = self._steps[1::2]
# Avoid this check for MockClasses
if len(self._kinematic_steps) > 0:
assert (
len(self._kinematic_steps) == len(self._dynamic_steps) + 1
), "Size mismatch in the number of kinematic and dynamic steps provided for a Symplectic Stepper!"
from itertools import zip_longest
def NoOp(*args):
pass
self._steps_and_prefactors = tuple(
zip_longest(
self._prefactors,
self._kinematic_steps,
self._dynamic_steps,
fillvalue=NoOp,
)
)
def step_methods(self):
return self._steps_and_prefactors
@property
def n_stages(self):
return len(self._steps_and_prefactors)
class SymplecticStepperTag:
def __init__(self):
pass
class PositionVerlet:
"""
Position Verlet symplectic time stepper class, which
includes methods for second-order position Verlet.
"""
Tag = SymplecticStepperTag()
def __init__(self):
pass
def _first_prefactor(self, dt):
return 0.5 * dt
def _first_kinematic_step(self, System, time: np.float64, dt: np.float64):
prefac = self._first_prefactor(dt)
overload_operator_kinematic_numba(
System.n_nodes,
prefac,
System.kinematic_states.position_collection,
System.kinematic_states.director_collection,
System.velocity_collection,
System.omega_collection,
)
def _first_dynamic_step(self, System, time: np.float64, dt: np.float64):
overload_operator_dynamic_numba(
System.dynamic_states.rate_collection,
System.dynamic_rates(time, dt),
)
class PEFRL:
"""
Position Extended Forest-Ruth Like Algorithm of
I.M. Omelyan, I.M. Mryglod and R. Folk, Computer Physics Communications 146, 188 (2002),
http://arxiv.org/abs/cond-mat/0110585
"""
# xi and chi are confusing, but be careful!
ξ = np.float64(0.1786178958448091e0) # ξ
λ = -np.float64(0.2123418310626054e0) # λ
χ = -np.float64(0.6626458266981849e-1) # χ
# Pre-calculate other coefficients
lambda_dash_coeff = 0.5 * (1.0 - 2.0 * λ)
xi_chi_dash_coeff = 1.0 - 2.0 * (ξ + χ)
Tag = SymplecticStepperTag()
def __init__(self):
pass
def _first_kinematic_prefactor(self, dt):
return self.ξ * dt
def _first_kinematic_step(self, System, time: np.float64, dt: np.float64):
prefac = self._first_kinematic_prefactor(dt)
overload_operator_kinematic_numba(
System.n_nodes,
prefac,
System.kinematic_states.position_collection,
System.kinematic_states.director_collection,
System.velocity_collection,
System.omega_collection,
)
# System.kinematic_states += prefac * System.kinematic_rates(time, prefac)
def _first_dynamic_step(self, System, time: np.float64, dt: np.float64):
prefac = self.lambda_dash_coeff * dt
overload_operator_dynamic_numba(
System.dynamic_states.rate_collection,
System.dynamic_rates(time, prefac),
)
# System.dynamic_states += prefac * System.dynamic_rates(time, prefac)
def _second_kinematic_prefactor(self, dt):
return self.χ * dt
def _second_kinematic_step(self, System, time: np.float64, dt: np.float64):
prefac = self._second_kinematic_prefactor(dt)
overload_operator_kinematic_numba(
System.n_nodes,
prefac,
System.kinematic_states.position_collection,
System.kinematic_states.director_collection,
System.velocity_collection,
System.omega_collection,
)
# System.kinematic_states += prefac * System.kinematic_rates(time, prefac)
def _second_dynamic_step(self, System, time: np.float64, dt: np.float64):
prefac = self.λ * dt
overload_operator_dynamic_numba(
System.dynamic_states.rate_collection,
System.dynamic_rates(time, prefac),
)
# System.dynamic_states += prefac * System.dynamic_rates(time, prefac)
def _third_kinematic_prefactor(self, dt):
return self.xi_chi_dash_coeff * dt
def _third_kinematic_step(self, System, time: np.float64, dt: np.float64):
prefac = self._third_kinematic_prefactor(dt)
# Need to fill in
overload_operator_kinematic_numba(
System.n_nodes,
prefac,
System.kinematic_states.position_collection,
System.kinematic_states.director_collection,
System.velocity_collection,
System.omega_collection,
)
# System.kinematic_states += prefac * System.kinematic_rates(time, prefac)