def test_swap_xy(self):
concentration = states.StateDefinition('concentration', (), (0, 0, 0),
(0, 0))
self.assertEqual(concentration.swap_xy(), concentration)
x_quantity = states.StateDefinition('quantity', (states.Dimension.X, ),
(1, 2, 3), (1, 0))
y_quantity = states.StateDefinition('quantity', (states.Dimension.Y, ),
(2, 1, 3), (0, 1))
self.assertEqual(x_quantity.swap_xy(), y_quantity)
self.assertEqual(y_quantity.swap_xy(), x_quantity)
def test_permutation(self):
definitions = {
'q_edge_x': states.StateDefinition('q', (), (0, 0, 0), (1, 0)),
'q_edge_y': states.StateDefinition('q', (), (0, 0, 0), (0, 1)),
}
transform = geometry.Permutation(definitions)
rs = np.random.RandomState(0)
inputs = {'q_edge_x': tf.convert_to_tensor(rs.randn(5, 5))}
result = transform.forward(inputs)
self.assertEqual(set(result), {'q_edge_y'})
np.testing.assert_array_equal(result['q_edge_y'].numpy().T,
inputs['q_edge_x'].numpy())
def test_face_average(self):
# construct a velocity field that covers every branch in the integral
vfield = velocity_fields.ConstantVelocityField(
x_wavenumbers=np.array([1, 0, 2, 0]),
y_wavenumbers=np.array([-1, 3, 0, 0]),
amplitudes=np.array([1.1, 1.2, 1.3, 1.4]),
phase_shifts=np.arange(4.0),
)
high_res_grid = grids.Grid.from_period(256, length=5)
low_res_grid = grids.Grid.from_period(32, length=5)
with self.subTest('point and face-average match at high resolution'):
kwargs = dict(t=0, grid=high_res_grid)
vx_point = vfield.get_velocity_x(face_average=False, **kwargs)
vy_point = vfield.get_velocity_y(face_average=False, **kwargs)
vx_average = vfield.get_velocity_x(face_average=True, **kwargs)
vy_average = vfield.get_velocity_y(face_average=True, **kwargs)
np.testing.assert_allclose(vx_point, vx_average, atol=1e-3)
np.testing.assert_allclose(vy_point, vy_average, atol=1e-3)
with self.subTest('average high-res and low-res average match'):
kwargs = dict(t=0, face_average=True)
vx_low = vfield.get_velocity_x(grid=low_res_grid, shift=(1, 0), **kwargs)
vy_low = vfield.get_velocity_y(grid=low_res_grid, shift=(0, 1), **kwargs)
vx_high = vfield.get_velocity_x(
grid=high_res_grid, shift=(1, 0), **kwargs)
vy_high = vfield.get_velocity_y(
grid=high_res_grid, shift=(0, 1), **kwargs)
vx_def = states.StateDefinition(
'v', (states.Dimension.X,), (0, 0, 0), offset=(1, 0))
vy_def = states.StateDefinition(
'v', (states.Dimension.Y,), (0, 0, 0), offset=(0, 1))
vx_high_average = tensor_ops.regrid(
vx_high, vx_def, high_res_grid, low_res_grid)
vy_high_average = tensor_ops.regrid(
vy_high, vy_def, high_res_grid, low_res_grid)
np.testing.assert_allclose(vx_low, vx_high_average)
np.testing.assert_allclose(vy_low, vy_high_average)
with self.subTest('divergence free'):
divergence = advection_equations.flux_to_time_derivative(
vx_low, vy_low, grid_step=1.0)
np.testing.assert_allclose(divergence, np.zeros_like(divergence),
atol=1e-8)
def test_symmetries_of_the_square(self):
definitions = {
'q': states.StateDefinition('q', (), (0, 0, 0), (0, 0)),
'q_x': states.StateDefinition('q', (), (0, 0, 0), (1, 0)),
'q_y': states.StateDefinition('q', (), (0, 0, 0), (0, 1)),
}
transforms = geometry.symmetries_of_the_square(definitions)
self.assertLen(transforms, 8)
rs = np.random.RandomState(0)
state = {
'q': tf.convert_to_tensor(rs.randn(5, 5)),
'q_x': tf.convert_to_tensor(rs.randn(5, 5)),
'q_y': tf.convert_to_tensor(rs.randn(5, 5)),
}
for transform in transforms:
with self.subTest(transform):
roundtripped = transform.inverse(transform.forward(state))
np.testing.assert_array_equal(state['q'].numpy(),
roundtripped['q'].numpy())
np.testing.assert_array_equal(state['q_x'].numpy(),
roundtripped['q_x'].numpy())
np.testing.assert_array_equal(state['q_y'].numpy(),
roundtripped['q_y'].numpy())
with self.subTest('uniqueness'):
rs = np.random.RandomState(0)
state = {
'q': tf.convert_to_tensor(rs.randn(5, 5)),
}
results_set = {
tuple(transform.forward(state)['q'].numpy().ravel().tolist())
for transform in transforms
}
self.assertLen(results_set, 8)
with self.subTest('permutation_counts'):
rs = np.random.RandomState(0)
state = {
'q': tf.convert_to_tensor(rs.randn(5, 5)),
'q_x': tf.convert_to_tensor(rs.randn(5, 5)),
}
counts = collections.defaultdict(int)
for transform in transforms:
result = transform.forward(state)
for k in result:
counts[k] += 1
self.assertEqual(counts, {'q': 8, 'q_x': 4, 'q_y': 4})
def test_reflection(self):
definitions = {
'q': states.StateDefinition('q', (), (0, 0, 0), (0, 0)),
'q_edge_x': states.StateDefinition('q', (), (0, 0, 0), (1, 0)),
'q_edge_y': states.StateDefinition('q', (), (0, 0, 0), (0, 1)),
'q_x': states.StateDefinition('q', (), (1, 0, 0), (0, 0)),
'q_xy': states.StateDefinition('q', (), (1, 1, 0), (0, 0)),
'q_xx': states.StateDefinition('q', (), (2, 0, 0), (0, 0)),
'x_velocity': states.StateDefinition('q', (X, ), (0, 0, 0),
(0, 0)),
}
transform = geometry.Reflection([X], definitions)
self.assertEqual(repr(transform), '<Reflection [X]>')
inputs = {
k: tf.convert_to_tensor(np.array([[0, 1, 2, 3]]).T)
for k in definitions
}
result = transform.forward(inputs)
np.testing.assert_array_equal(result['q'].numpy().T, [[3, 2, 1, 0]])
np.testing.assert_array_equal(result['q_edge_x'].numpy().T,
[[2, 1, 0, 3]])
np.testing.assert_array_equal(result['q_edge_y'].numpy().T,
[[3, 2, 1, 0]])
np.testing.assert_array_equal(result['q_x'].numpy().T,
[[-3, -2, -1, 0]])
np.testing.assert_array_equal(result['q_xy'].numpy().T,
[[-3, -2, -1, 0]])
np.testing.assert_array_equal(result['q_xx'].numpy().T, [[3, 2, 1, 0]])
np.testing.assert_array_equal(result['x_velocity'].numpy().T,
[[-3, -2, -1, 0]])
def test_regrid(self):
tensor = tf.convert_to_tensor([[1, 2, 3, 4], [5, 6, 7, 8]],
dtype=tf.float32)
source = grids.Grid(2, 4, step=1)
destination = grids.Grid(1, 2, step=2)
definition = states.StateDefinition('centered', (), (0, 0, 0), (0, 0))
actual = tensor_ops.regrid(tensor, definition, source, destination)
np.testing.assert_array_equal(actual, [[3.5, 5.5]])
definition = states.StateDefinition('x_offset', (), (0, 0, 0), (1, 0))
actual = tensor_ops.regrid(tensor, definition, source, destination)
np.testing.assert_array_equal(actual, [[5.5, 7.5]])
definition = states.StateDefinition('y_offset', (), (0, 0, 0), (0, 1))
actual = tensor_ops.regrid(tensor, definition, source, destination)
np.testing.assert_array_equal(actual, [[4, 6]])
definition = states.StateDefinition('xy_offset', (), (0, 0, 0), (1, 1))
actual = tensor_ops.regrid(tensor, definition, source, destination)
np.testing.assert_array_equal(actual, [[6, 8]])
def test_time_derivative(self):
initial_key = states.StateDefinition('key_a', (), (1, 2, 3), (4, 5))
expected = states.StateDefinition('key_a', (), (1, 2, 4), (4, 5))
self.assertEqual(expected, initial_key.time_derivative())
def test_with_prefix(self):
initial_key = states.StateDefinition('name', (), (1, 2, 3), (4, 5))
expected = states.StateDefinition('exact_name', (), (1, 2, 3), (4, 5))
result = initial_key.exact()
self.assertEqual(result, expected)