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 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 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 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