def test_integration_in_constant_velocity_field(self, equation, atol):
grid = grids.Grid.from_period(100, length=100)
model = models.FiniteDifferenceModel(equation, grid)
initial_concentration = equations.symmetrized_gaussian(
grid, 50, 50, gaussian_width=20)
self.assertGreaterEqual(initial_concentration.min(), 0.0)
self.assertLessEqual(initial_concentration.max(), 1.0 + 1e-7)
initial_state = {
'concentration': initial_concentration.astype(np.float32),
'x_velocity': np.ones(grid.shape, np.float32),
'y_velocity': np.zeros(grid.shape, np.float32)
}
steps = round(1 / equation.cfl_safety_factor) * np.arange(10)
integrated = integrate.integrate_steps(model, initial_state, steps)
actual = integrated['concentration'].numpy()
expected = np.stack(
[np.roll(initial_concentration, i, axis=0) for i in range(10)])
np.testing.assert_allclose(actual, expected, atol=atol)
if equation.METHOD is polynomials.Method.FINITE_VOLUME:
np.testing.assert_allclose(actual[-1].sum(),
actual[0].sum(),
rtol=1e-7)
else:
np.testing.assert_allclose(actual[-1].sum(),
actual[0].sum(),
rtol=1e-3)
if equation.MONOTONIC:
self.assertGreaterEqual(actual[-1].min(), 0.0)
self.assertLessEqual(actual[-1].max(), 1.0)
def integrate_for_initial_conditions(
self, state: Dict[str, tf.Tensor]) -> List[Dict[str, tf.Tensor]]:
"""Integrate from random seeds for initial conditions."""
assert tf.executing_eagerly()
integrated = integrate.integrate_steps(
self.model, state, self.initial_condition_steps, axis=0)
return unstack_dict(integrated, len(self.initial_condition_steps))
def postprocess(self, state: Dict[str, tf.Tensor]) -> StateGridTensors:
"""Resample states to low resolution."""
result = []
key_defs = self.equation.key_definitions
# exact solution
coarse_state = self.equation.regrid(state, self.simulation_grid,
self.output_grid)
for k, v in coarse_state.items():
result.append(((key_defs[k].exact(), self.output_grid), v))
# baseline solution
initial_coarse_state = nest.map_structure(operator.itemgetter(0),
coarse_state)
integrated_baseline = integrate.integrate_steps(
self.coarse_model, initial_coarse_state,
np.arange(self.example_num_time_steps))
for k, v in integrated_baseline.items():
result.append(((key_defs[k].baseline(), self.output_grid), v))
return merge(result)
def integrate_each_example(
self, state: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
"""Integrate each initial condition to produce examples for training."""
assert tf.executing_eagerly()
steps = self.time_step_ratio * np.arange(self.example_num_time_steps)
return integrate.integrate_steps(self.model, state, steps, axis=0)