def _interp_fit_adaptive_heun(y0, y1, k, dt, tableau=_ADAPTIVE_HEUN_TABLEAU):
"""Fit an interpolating polynomial to the results of a Runge-Kutta step."""
dt = tf.cast(dt, y0[0].dtype)
y_mid = tuple(y0_ + _scaled_dot_product(dt, AH_C_MID, k_) for y0_, k_ in zip(y0, k))
f0 = tuple(k_[0] for k_ in k)
f1 = tuple(k_[-1] for k_ in k)
return _interp_fit(y0, y1, y_mid, f0, f1, dt)
def _interp_fit_bosh3(y0, y1, k, dt):
"""Fit an interpolating polynomial to the results of a Runge-Kutta step."""
dt = tf.cast(dt, y0[0].dtype)
y_mid = tuple(y0_ + _scaled_dot_product(dt, BS_C_MID, k_) for y0_, k_ in zip(y0, k))
f0 = tuple(k_[0] for k_ in k)
f1 = tuple(k_[-1] for k_ in k)
return _interp_fit(y0, y1, y_mid, f0, f1, dt)
def _interp_fit_dopri8(y0, y1, k, dt, tableau=_DOPRI8_TABLEAU):
"""Fit an interpolating polynomial to the results of a Runge-Kutta step."""
y_mid = tuple(y0_ + _scaled_dot_product(dt, c_mid, k_)
for y0_, k_ in zip(y0, k))
f0 = tuple(k_[0] for k_ in k)
f1 = tuple(k_[-1] for k_ in k)
return _interp_fit(y0, y1, y_mid, f0, f1, dt)
def _interp_fit_dopri5(y0,
y1,
k,
dt,
tableau=_DORMAND_PRINCE_SHAMPINE_TABLEAU):
"""Fit an interpolating polynomial to the results of a Runge-Kutta step."""
dt = tf.cast(dt, y0[0].dtype)
y_mid = tuple(y0_ + _scaled_dot_product(dt, DPS_C_MID, k_)
for y0_, k_ in zip(y0, k))
f0 = tuple(k_[0] for k_ in k)
f1 = tuple(k_[-1] for k_ in k)
return _interp_fit(y0, y1, y_mid, f0, f1, dt)