def _integrate_backward(self, q, y_of_t, normA, normB, i0, k_ab4, dt_ab4):
"""
Use AB4 to integrate backward in time, starting at index i0.
k_ab4 is [dydt(i0 + 3), dydt(i0 + 2), dydt(i0 + 1)]
dt_ab4 is [t(i0 + 3) - t(i0 + 2), t(i0 + 2) - t(i0 + 1), t(i0 + 1) - t(i0)]
"""
if i0 > len(self.t) - 7:
raise Exception("i0 must be <= len(self.t) - 7")
# Setup AB4
k1, k2, k3 = k_ab4
dt1, dt2, dt3 = dt_ab4
# Setup dt array, removing the half steps
dt_array = np.append(2 * self.diff_t[:6:2], self.diff_t[6:])
for i_output in range(i0)[::-1]:
node_index = i_output + 4
if i_output < 2:
node_index = 2 + 2*i_output
dt4 = dt_array[i_output]
k4 = self.get_time_deriv_from_index(node_index, q, y_of_t[i_output+1])
ynext = y_of_t[i_output+1] - _utils.ab4_dy(k1, k2, k3, k4, dt1, dt2, dt3, dt4)
y_of_t[i_output] = _utils.normalize_y(ynext, normA, normB)
# Setup for next iteration
k1, k2, k3 = k2, k3, k4
dt1, dt2, dt3 = dt2, dt3, dt4
return y_of_t
def _integrate_forward(self, q, y_of_t, normA, normB, i0, k_ab4, dt_ab4):
"""
Use AB4 to integrate forward in time, starting at index i0.
i0 refers to the index of y_of_t, which should be the latest index at which
we already have the solution; typically i0=3 after three steps of RK4.
k_ab4 is [dydt(i0 - 3), dydt(i0 - 2), dydt(i0 - 1)]
dt_ab4 is [t(i0 - 2) - t(i0 - 3), t(i0 - 1) - t(i0 - 2), t(i0) - t(i0 - 1)]
where for both k_ab4 and dt_ab4 the indices correspond to y_of_t nodes and skip
fractional nodes.
"""
if i0 < 3:
raise Exception("i0 must be at least 3!")
# Setup AB4
k1, k2, k3 = k_ab4
dt1, dt2, dt3 = dt_ab4
# Run AB4
for i, dt4 in enumerate(self.diff_t[i0+3:]): #i0+3 due to 3 half time steps
i_output = i0+i
k4 = self.get_time_deriv_from_index(i_output+3, q, y_of_t[i_output])
ynext = y_of_t[i_output] + _utils.ab4_dy(k1, k2, k3, k4, dt1, dt2, dt3, dt4)
y_of_t[i_output+1] = _utils.normalize_y(ynext, normA, normB)
# Setup for next iteration
k1, k2, k3 = k2, k3, k4
dt1, dt2, dt3 = dt2, dt3, dt4
return y_of_t
def _initialize(self, q, chiA0, chiB0, init_quat, init_orbphase, t_ref, normA, normB):
"""
Initializes an array of data with the initial conditions.
If t_ref does not correspond to a time node, takes one small time step to
the nearest time node.
"""
# data is [q0, qx, qy, qz, orbphase, chiAx, chiAy, chiAz, chiBx, chiBy, chiBz]
# We do three steps of RK4, so we have 3 fewer timesteps in the output
# compared to self.t
data = np.zeros((self.L-3, 11))
y0 = np.append(np.array([1., 0., 0., 0., init_orbphase]), np.append(chiA0, chiB0))
if init_quat is not None:
y0[:4] = init_quat
if t_ref is None:
data[0, :] = y0
i0 = 0
else:
# Step to the closest time node using forward Euler
times = np.append(self.t[:6:2], self.t[6:])
i0 = np.argmin(abs(times - t_ref))
t0 = times[i0]
dydt0 = self.get_time_deriv(t_ref, q, y0)
y_node = y0 + (t0 - t_ref) * dydt0
y_node = _utils.normalize_y(y_node, normA, normB)
data[i0, :] = y_node
return data, i0
def _initial_RK4(self, q, y_of_t, normA, normB):
"""This is used to initialize the AB4 system when t_ref=t_0 (default)"""
# Three steps of RK4
k_ab4 = []
dt_ab4 = []
for i, dt in enumerate(self.diff_t[:6:2]):
k1 = self.get_time_deriv_from_index(2*i, q, y_of_t[i])
k_ab4.append(k1)
dt_ab4.append(2*dt)
k2 = self.get_time_deriv_from_index(2*i+1, q, y_of_t[i] + dt*k1)
k3 = self.get_time_deriv_from_index(2*i+1, q, y_of_t[i] + dt*k2)
k4 = self.get_time_deriv_from_index(2*i+2, q, y_of_t[i] + 2*dt*k3)
ynext = y_of_t[i] + (dt/3.)*(k1 + 2*k2 + 2*k3 + k4)
y_of_t[i+1] = _utils.normalize_y(ynext, normA, normB)
return k_ab4, dt_ab4, y_of_t
def _one_forward_RK4_step(self, q, y_of_t, normA, normB, i0):
"""Steps forward one step using RK4"""
# i0 is on the y_of_t grid, which has 3 fewer samples than the self.t grid
i_t = i0 + 3
if i0 < 3:
i_t = i0*2
t1 = self.t[i_t]
t2 = self.t[i_t + 1]
if i0 < 3:
t2 = self.t[i_t + 2]
half_dt = 0.5*(t2 - t1)
k1 = self.get_time_deriv(t1, q, y_of_t[i0])
k2 = self.get_time_deriv(t1 + half_dt, q, y_of_t[i0] + half_dt*k1)
k3 = self.get_time_deriv(t1 + half_dt, q, y_of_t[i0] + half_dt*k2)
k4 = self.get_time_deriv(t2, q, y_of_t[i0] + 2*half_dt*k3)
ynext = y_of_t[i0] + (half_dt/3.)*(k1 + 2*k2 + 2*k3 + k4)
y_of_t[i0+1] = _utils.normalize_y(ynext, normA, normB)
return y_of_t, k1