def __init__(self,
func,
y0,
rtol,
atol,
implicit=True,
first_step=None,
max_order=_MAX_ORDER,
safety=0.9,
ifactor=10.0,
dfactor=0.2,
**unused_kwargs):
_handle_unused_kwargs(self, unused_kwargs)
del unused_kwargs
self.func = func
self.y0 = y0
self.rtol = rtol if _is_iterable(rtol) else [rtol] * len(y0)
self.atol = atol if _is_iterable(atol) else [atol] * len(y0)
self.implicit = implicit
self.first_step = first_step
self.max_order = int(max(_MIN_ORDER, min(max_order, _MAX_ORDER)))
self.safety = _convert_to_tensor(safety,
dtype=tf.float64,
device=y0[0].device)
self.ifactor = _convert_to_tensor(ifactor,
dtype=tf.float64,
device=y0[0].device)
self.dfactor = _convert_to_tensor(dfactor,
dtype=tf.float64,
device=y0[0].device)
def __init__(self,
func,
y0,
rtol,
atol,
first_step=None,
safety=0.9,
ifactor=10.0,
dfactor=0.2,
max_num_steps=2**31 - 1,
**unused_kwargs):
_handle_unused_kwargs(self, unused_kwargs)
del unused_kwargs
self.func = func
self.y0 = y0
self.rtol = rtol if _is_iterable(rtol) else [rtol] * len(y0)
self.atol = atol if _is_iterable(atol) else [atol] * len(y0)
self.first_step = first_step
self.safety = _convert_to_tensor(safety,
dtype=tf.float64,
device=y0[0].device)
self.ifactor = _convert_to_tensor(ifactor,
dtype=tf.float64,
device=y0[0].device)
self.dfactor = _convert_to_tensor(dfactor,
dtype=tf.float64,
device=y0[0].device)
self.max_num_steps = _convert_to_tensor(max_num_steps,
dtype=tf.int32,
device=y0[0].device)
def _interp_evaluate(coefficients, t0, t1, t):
"""Evaluate polynomial interpolation at the given time point.
Args:
coefficients: list of Tensor coefficients as created by `interp_fit`.
t0: scalar float64 Tensor giving the start of the interval.
t1: scalar float64 Tensor giving the end of the interval.
t: scalar float64 Tensor giving the desired interpolation point.
Returns:
Polynomial interpolation of the coefficients at time `t`.
"""
dtype = coefficients[0][0].dtype
device = coefficients[0][0].device
t0 = _convert_to_tensor(t0, dtype=dtype, device=device)
t1 = _convert_to_tensor(t1, dtype=dtype, device=device)
t = _convert_to_tensor(t, dtype=dtype, device=device)
assert (t0 <= t) & (
t <=
t1), 'invalid interpolation, fails `t0 <= t <= t1`: {}, {}, {}'.format(
t0, t, t1)
x = tf.cast(((t - t0) / (t1 - t0)), dtype)
x = move_to_device(x, device)
xs = [move_to_device(tf.convert_to_tensor(1, dtype=dtype), device), x]
for _ in range(2, len(coefficients)):
xs.append(xs[-1] * x)
return tuple(
_dot_product(coefficients_, reversed(xs))
for coefficients_ in zip(*coefficients))
def before_integrate(self, t):
if self.first_step is None:
first_step = _convert_to_tensor(_select_initial_step(self.func, t[0], self.y0, 4, self.rtol, self.atol),
device=t.device)
else:
first_step = _convert_to_tensor(0.01, dtype=t.dtype, device=t.device)
self.rk_state = _RungeKuttaState(
self.y0,
cast_double(self.func(t[0], self.y0)), t[0], t[0], first_step,
tuple(map(lambda x: [x] * 7, self.y0))
)
def before_integrate(self, t):
if self.first_step is None:
first_step = _convert_to_tensor(_select_initial_step(self.func, t[0], self.y0, 4, self.rtol, self.atol),
device=t.device)
else:
first_step = _convert_to_tensor(self.first_step, dtype=t.dtype, device=t.device)
self.rk_state = _RungeKuttaState(
self.y0,
self.func(t[0], self.y0), t[0], t[0], first_step,
[self.y0] * 7
)
def before_integrate(self, t):
f0 = self.func(tf.cast(t[0], self.y0[0].dtype), self.y0)
if self.first_step is None:
first_step = _select_initial_step(self.func, t[0], self.y0, 1, self.rtol[0], self.atol[0], f0=f0)
first_step = move_to_device(tf.cast(first_step, t.dtype), t.device)
else:
first_step = _convert_to_tensor(self.first_step, dtype=t.dtype, device=t.device)
self.rk_state = _RungeKuttaState(self.y0, f0, t[0], t[0], first_step, interp_coeff=[self.y0] * 5)
def advance(self, final_t):
final_t = _convert_to_tensor(final_t, device=self.vcabm_state.prev_t[0].device)
while final_t > self.vcabm_state.prev_t[0]:
# print("VCABM State T = ", final_t.numpy(), self.vcabm_state.y_n)
self.vcabm_state = self._adaptive_adams_step(self.vcabm_state, final_t)
assert tf.equal(final_t, self.vcabm_state.prev_t[0])
return self.vcabm_state.y_n
def _runge_kutta_step(func, y0, f0, t0, dt, tableau):
"""Take an arbitrary Runge-Kutta step and estimate error.
Args:
func: Function to evaluate like `func(t, y)` to compute the time derivative
of `y`.
y0: Tensor initial value for the state.
f0: Tensor initial value for the derivative, computed from `func(t0, y0)`.
t0: float64 scalar Tensor giving the initial time.
dt: float64 scalar Tensor giving the size of the desired time step.
tableau: optional _ButcherTableau describing how to take the Runge-Kutta
step.
name: optional name for the operation.
Returns:
Tuple `(y1, f1, y1_error, k)` giving the estimated function value after
the Runge-Kutta step at `t1 = t0 + dt`, the derivative of the state at `t1`,
estimated error at `t1`, and a list of Runge-Kutta coefficients `k` used for
calculating these terms.
"""
y0 = cast_double(y0)
f0 = cast_double(f0)
dtype = y0[0].dtype
device = y0[0].device
t0 = _convert_to_tensor(t0, dtype=dtype, device=device)
dt = _convert_to_tensor(dt, dtype=dtype, device=device)
k = tuple(map(lambda x: [x], f0))
for alpha_i, beta_i in zip(tableau.alpha, tableau.beta):
ti = t0 + alpha_i * dt
yi = tuple(y0_ + _scaled_dot_product(dt, cast_double(beta_i), k_)
for y0_, k_ in zip(y0, k))
tuple(k_.append(cast_double(f_)) for k_, f_ in zip(k, func(ti, yi)))
if not (tableau.c_sol[-1] == 0 and tableau.c_sol[:-1] == tableau.beta[-1]):
# This property (true for Dormand-Prince) lets us save a few FLOPs.
yi = tuple(y0_ + _scaled_dot_product(dt, tableau.c_sol, k_)
for y0_, k_ in zip(y0, k))
y1 = yi
f1 = tuple(k_[-1] for k_ in k)
y1_error = tuple(_scaled_dot_product(dt, tableau.c_error, k_) for k_ in k)
return (y1, f1, y1_error, k)
def _optimal_step_size(last_step, mean_error_ratio, safety=0.9, ifactor=10.0, dfactor=0.2, order=5):
"""Calculate the optimal size for the next Runge-Kutta step."""
if mean_error_ratio == 0:
return last_step * ifactor
if mean_error_ratio < 1:
dfactor = _convert_to_tensor(1., dtype=tf.float64, device=mean_error_ratio.device)
error_ratio = tf.cast(mean_error_ratio, last_step.dtype)
exponent = tf.convert_to_tensor(1. / order, dtype=last_step.dtype)
factor = tf.maximum(1. / ifactor, tf.minimum((error_ratio ** exponent) / safety, 1. / dfactor))
return last_step / factor
