def baseline_result(inputs: tf.Tensor,
equation: equations.Equation,
num_time_steps: int = 0,
accuracy_order: int = None) -> tf.Tensor:
"""Calculate derivatives and time-evolution using our baseline model.
Args:
inputs: float32 Tensor with dimensions [batch, x].
equation: equation being solved.
num_time_steps: integer number of time steps to integrate over.
accuracy_order: optional explicit accuracy order.
Returns:
Float32 Tensor with dimensions [batch, x, channel] with inferred space
derivatives, time derivative and the integrated solution.
"""
if accuracy_order is None:
equation = equation.to_exact()
elif type(equation) in equations.FLUX_EQUATION_TYPES:
equation = equation.to_conservative()
space_derivatives = baseline_space_derivatives(
inputs, equation, accuracy_order=accuracy_order)
time_derivative = apply_space_derivatives(space_derivatives, inputs,
equation)
if num_time_steps:
integrated_solution = baseline_time_evolution(inputs, num_time_steps,
equation)
else:
integrated_solution = None
return result_stack(space_derivatives, time_derivative,
integrated_solution)
def __init__(self,
equation: equations.Equation,
accuracy_order: Optional[int] = 1):
with tf.Graph().as_default():
self.t = tf.placeholder(tf.float32, shape=())
num_points = equation.grid.solution_num_points
self.inputs = tf.placeholder(tf.float32, shape=(num_points, ))
batched_inputs = self.inputs[tf.newaxis, :]
space_derivatives = model.baseline_space_derivatives(
batched_inputs, equation, accuracy_order=accuracy_order)
time_derivative = tf.squeeze(model.apply_space_derivatives(
space_derivatives, batched_inputs, equation),
axis=0)
self.value = equation.finalize_time_derivative(
self.t, time_derivative)
self._space_derivatives = {
k: tf.squeeze(space_derivatives[..., i], axis=0)
for i, k in enumerate(equation.DERIVATIVE_NAMES)
}
self.sess = tf.Session()
def integrate(equation: equations.Equation,
differentiator: Differentiator,
times: np.ndarray = _DEFAULT_TIMES,
warmup: float = 0,
integrate_method: str = 'RK23',
filter_interval: float = None,
filter_all_times: bool = False) -> xarray.Dataset:
"""Integrate an equation with possible warmup or periodic filtering."""
if filter_interval is not None:
warmup_odeint = functools.partial(
odeint_with_periodic_filtering,
filter_interval=filter_interval,
filter_order=max(equation.to_exact().DERIVATIVE_ORDERS))
else:
warmup_odeint = odeint
if warmup:
equation_exact = equation.to_exact()
diff_exact = exact_differentiator(equation_exact)
if filter_interval is not None:
warmup_times = np.arange(0, warmup + 1e-8, filter_interval)
else:
warmup_times = np.array([0, warmup])
y0_0 = equation_exact.initial_value()
solution_warmup, _ = warmup_odeint(y0_0,
diff_exact,
times=warmup_times,
method=integrate_method)
# use the sample after warmup to initialize later simulations
y0 = equation.grid.resample(solution_warmup[-1, :])
else:
y0 = equation.initial_value()
odeint_func = warmup_odeint if filter_all_times else odeint
solution, num_evals = odeint_func(y0,
differentiator,
times=warmup + times,
method=integrate_method)
results = xarray.Dataset(data_vars={'y': (('time', 'x'), solution)},
coords={
'time': warmup + times,
'x': equation.grid.solution_x,
'num_evals': num_evals
})
return results
def baseline_space_derivatives(inputs: tf.Tensor,
equation: equations.Equation,
accuracy_order: int = None) -> tf.Tensor:
"""Calculate spatial derivatives using a baseline metohd."""
assert_consistent_solution(equation, inputs)
spatial_derivatives_list = []
for derivative_name, derivative_order in zip(equation.DERIVATIVE_NAMES,
equation.DERIVATIVE_ORDERS):
if accuracy_order is None:
# use the best baseline method
assert equation.exact_type() is type(equation)
if equation.EXACT_METHOD is equations.ExactMethod.POLYNOMIAL:
grid = (0.5 + np.arange(-3, 3)) * equation.grid.solution_dx
method = FINITE_VOL if equation.CONSERVATIVE else FINITE_DIFF
derivative = polynomials.reconstruct(inputs, grid, method,
derivative_order)
elif equation.EXACT_METHOD is equations.ExactMethod.SPECTRAL:
derivative = duckarray.spectral_derivative(
inputs, derivative_order, equation.grid.period)
elif equation.EXACT_METHOD is equations.ExactMethod.WENO:
if derivative_name == 'u_minus':
derivative = duckarray.roll(weno.reconstruct_left(inputs),
1,
axis=-1)
elif derivative_name == 'u_plus':
derivative = duckarray.roll(weno.reconstruct_right(inputs),
1,
axis=-1)
else:
assert derivative_name == 'u_x'
grid = polynomials.regular_grid(
grid_offset=equation.GRID_OFFSET,
derivative_order=derivative_order,
accuracy_order=3,
dx=equation.grid.solution_dx)
method = FINITE_VOL if equation.CONSERVATIVE else FINITE_DIFF
derivative = polynomials.reconstruct(
inputs, grid, method, derivative_order)
else:
# explicit accuracy order provided
assert type(equation) not in equations.FLUX_EQUATION_TYPES
grid = polynomials.regular_grid(grid_offset=equation.GRID_OFFSET,
derivative_order=derivative_order,
accuracy_order=accuracy_order,
dx=equation.grid.solution_dx)
method = FINITE_VOL if equation.CONSERVATIVE else FINITE_DIFF
derivative = polynomials.reconstruct(inputs, grid, method,
derivative_order)
spatial_derivatives_list.append(derivative)
return tf.stack(spatial_derivatives_list, axis=-1)
def integrate_exact(
equation: equations.Equation,
times: np.ndarray = _DEFAULT_TIMES,
warmup: float = 0,
integrate_method: str = 'RK23',
filter_interval: float = None) -> xarray.Dataset:
"""Integrate only the exact model."""
equation = equation.to_exact()
differentiator = exact_differentiator(equation)
return integrate(equation, differentiator, times, warmup,
integrate_method=integrate_method,
filter_interval=filter_interval)
def __init__(self,
checkpoint_dir: str,
equation: equations.Equation,
hparams: tf.contrib.training.HParams):
with tf.Graph().as_default():
self.t = tf.placeholder(tf.float32, shape=())
num_points = equation.grid.solution_num_points
self.inputs = tf.placeholder(tf.float32, shape=(num_points,))
time_derivative = tf.squeeze(model.predict_time_derivative(
self.inputs[tf.newaxis, :], hparams), axis=0)
self.value = equation.finalize_time_derivative(self.t, time_derivative)
saver = tf.train.Saver()
self.sess = tf.Session()
saver.restore(self.sess, checkpoint_dir)
def exact_differentiator(equation: equations.Equation) -> Differentiator:
"""Return an "exact" differentiator for the given equation.
Args:
equation: equation for which to produce an "exact" differentiator.
Returns:
Differentiator to use for "exact" integration.
"""
if type(equation.to_exact()) is not type(equation):
raise TypeError('an exact equation must be provided')
if equation.BASELINE is equations.Baseline.POLYNOMIAL:
differentiator = PolynomialDifferentiator(equation,
accuracy_order=None)
elif equation.BASELINE is equations.Baseline.SPECTRAL:
differentiator = SpectralDifferentiator(equation)
else:
raise TypeError('unexpected equation: {}'.format(equation))
return differentiator
def apply_space_derivatives(derivatives: tf.Tensor, inputs: tf.Tensor,
equation: equations.Equation) -> tf.Tensor:
"""Combine spatial derivatives with input to calculate time derivatives.
Args:
derivatives: float32 tensor with dimensions [batch, x, derivative] giving
unnormalized spatial derivatives, e.g., as output from
predict_derivatives() or center_finite_differences().
inputs: float32 tensor with dimensions [batch, x].
equation: equation being solved.
Returns:
Float32 Tensor with diensions [batch, x] giving the time derivatives for
the given inputs and derivative model.
"""
derivatives_dict = {
k: derivatives[..., i]
for i, k in enumerate(equation.DERIVATIVE_NAMES)
}
return equation.equation_of_motion(inputs, derivatives_dict)