def test_extrapolate_valid_3D_3x3x3_2(self):
valid = tensor([[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]],
[[0, 0, 1],
[0, 0, 0],
[1, 0, 0]],
[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]]], spatial('x, y, z'))
values = tensor([[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]],
[[1, 0, 4],
[0, 0, 0],
[2, 0, 0]],
[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]]], spatial('x, y, z'))
expected_valid = math.ones(spatial(x=3, y=3, z=3))
expected_values = tensor([[[3, 4, 4],
[2, 3, 4],
[2, 2, 3]],
[[3, 4, 4],
[2, 3, 4],
[2, 2, 3]],
[[3, 4, 4],
[2, 3, 4],
[2, 2, 3]]], spatial('x, y, z'))
extrapolated_values, extrapolated_valid = math.extrapolate_valid_values(values, valid, 2)
math.assert_close(extrapolated_values, expected_values)
math.assert_close(extrapolated_valid, expected_valid)
def test_extrapolate_valid_3x3_sanity(self):
values = tensor([[0, 0, 0],
[0, 1, 0],
[0, 0, 0]], spatial('x, y'))
valid = values
extrapolated_values, extrapolated_valid = math.extrapolate_valid_values(values, valid)
expected_values = math.ones(spatial(x=3, y=3))
expected_valid = extrapolated_values
math.assert_close(extrapolated_values, expected_values)
math.assert_close(extrapolated_valid, expected_valid)
def test_extrapolate_valid(self):
valid = tensor([[0, 0, 0], [0, 1, 1], [1, 0, 0]], 'x, y')
values = tensor([[1, 0, 0], [0, 4, 0], [2, 0, 0]], 'x, y')
expected_valid = tensor([[0, 1, 1], [1, 1, 1], [1, 1, 1]], 'x, y')
expected_values = tensor([[1, 4, 0], [3, 4, 0], [2, 3, 0]], 'x, y')
new_values, new_valid = math.extrapolate_valid_values(values, valid, 1)
self.assertTrue(new_values == expected_values)
self.assertTrue(new_valid == expected_valid)
def test_extrapolate_valid_3x3(self):
valid = tensor([[0, 0, 0],
[0, 0, 1],
[1, 0, 0]], spatial('x, y'))
values = tensor([[1, 0, 2],
[0, 0, 4],
[2, 0, 0]], spatial('x, y'))
expected_valid = tensor([[0, 1, 1],
[1, 1, 1],
[1, 1, 1]], spatial('x, y'))
expected_values = tensor([[1, 4, 4],
[2, 3, 4],
[2, 3, 4]], spatial('x, y'))
extrapolated_values, extrapolated_valid = math.extrapolate_valid_values(values, valid)
math.assert_close(extrapolated_values, expected_values)
math.assert_close(extrapolated_valid, expected_valid)
def extrapolate_valid(grid: GridType,
valid: GridType,
distance_cells=1) -> tuple:
"""
Extrapolates values of `grid` which are marked by nonzero values in `valid` using `phi.math.extrapolate_valid_values().
If `values` is a StaggeredGrid, its components get extrapolated independently.
Args:
grid: Grid holding the values for extrapolation
valid: Grid (same type as `values`) marking the positions for extrapolation with nonzero values
distance_cells: Number of extrapolation steps
Returns:
grid: Grid with extrapolated values.
valid: binary Grid marking all valid values after extrapolation.
"""
assert isinstance(
valid, type(grid)), 'Type of valid Grid must match type of grid.'
if isinstance(grid, CenteredGrid):
new_values, new_valid = extrapolate_valid_values(
grid.values, valid.values, distance_cells)
return grid.with_values(new_values), valid.with_values(new_valid)
elif isinstance(grid, StaggeredGrid):
new_values = []
new_valid = []
for cgrid, cvalid in zip(unstack(grid, 'vector'),
unstack(valid, 'vector')):
new_tensor, new_mask = extrapolate_valid(
cgrid, valid=cvalid, distance_cells=distance_cells)
new_values.append(new_tensor.values)
new_valid.append(new_mask.values)
return grid.with_values(math.stack(new_values,
channel(grid))), valid.with_values(
math.stack(
new_valid, channel(grid)))
else:
raise NotImplementedError()