def test_fourier_laplace_2d_periodic(self):
"""test for convergence of the laplace operator"""
test_params = {
'size': [16, 32, 40],
'L': [1, 2,
3], # NOTE: Cannot test with less than 1 full wavelength
}
test_cases = [
dict(zip(test_params, v)) for v in product(*test_params.values())
]
for params in test_cases:
vec = math.meshgrid(x=params['size'], y=params['size'])
sine_field = math.prod(
math.sin(2 * PI * params['L'] * vec / params['size'] + 1),
'vector')
sin_lap_ref = -2 * (
2 * PI * params['L'] / params['size']
)**2 * sine_field # leading 2 from from x-y cross terms
sin_lap = math.fourier_laplace(sine_field, 1)
try:
math.assert_close(sin_lap,
sin_lap_ref,
rel_tolerance=0,
abs_tolerance=1e-5)
except BaseException as e:
abs_error = math.abs(sin_lap - sin_lap_ref)
max_abs_error = math.max(abs_error)
max_rel_error = math.max(math.abs(abs_error / sin_lap_ref))
variation_str = "\n".join([
f"max_absolute_error: {max_abs_error}",
f"max_relative_error: {max_rel_error}",
])
print(f"{variation_str}\n{params}")
raise AssertionError(e, f"{variation_str}\n{params}")
def test__periodic_2d_arakawa_poisson_bracket(self):
"""test _periodic_2d_arakawa_poisson_bracket implementation"""
with math.precision(64):
# Define parameters to test
test_params = {
'grid_size': [(4, 4), (32, 32)],
'dx': [0.1, 1],
'gen_func': [
lambda grid_size: np.random.rand(*grid_size).reshape(
grid_size)
]
}
# Generate test cases as the product
test_cases = [
dict(zip(test_params, v))
for v in product(*test_params.values())
]
for params in test_cases:
grid_size = params['grid_size']
d1 = params['gen_func'](grid_size)
d2 = params['gen_func'](grid_size)
dx = params['dx']
padding = extrapolation.PERIODIC
ref = self.arakawa_reference_implementation(
np.pad(d1.copy(), 1, mode='wrap'),
np.pad(d2.copy(), 1, mode='wrap'), dx)[1:-1, 1:-1]
d1_tensor = field.CenteredGrid(
values=math.tensor(d1, names=['x', 'y']),
bounds=geom.Box([0, 0], list(grid_size)),
extrapolation=padding)
d2_tensor = field.CenteredGrid(
values=math.tensor(d2, names=['x', 'y']),
bounds=geom.Box([0, 0], list(grid_size)),
extrapolation=padding)
val = math._nd._periodic_2d_arakawa_poisson_bracket(
d1_tensor.values, d2_tensor.values, dx)
try:
math.assert_close(ref,
val,
rel_tolerance=1e-14,
abs_tolerance=1e-14)
except BaseException as e:
abs_error = math.abs(val - ref)
max_abs_error = math.max(abs_error)
max_rel_error = math.max(math.abs(abs_error / ref))
variation_str = "\n".join([
f"max_absolute_error: {max_abs_error}",
f"max_relative_error: {max_rel_error}",
])
print(ref)
print(val)
raise AssertionError(e, params, variation_str)
def l1_loss(tensor, batch_norm=True, reduce_batches=True):
if isinstance(tensor, StaggeredGrid):
tensor = tensor.staggered
if reduce_batches:
total_loss = math.sum(math.abs(tensor))
else:
total_loss = math.sum(math.abs(tensor),
axis=list(range(1, len(tensor.shape))))
if batch_norm and reduce_batches:
batch_size = math.shape(tensor)[0]
return total_loss / math.to_float(batch_size)
else:
return total_loss
def extrapolation_helper(elements, t_shift, v_field, mask):
shift = math.ceil(math.max(
math.abs(elements.center - points.center))) - t_shift
t_shift += shift
v_field, mask = extrapolate_valid(v_field, mask, int(shift))
v_field *= accessible
return v_field, mask, t_shift
def push(self,
positions: Tensor,
outward: bool = True,
shift_amount: float = 0) -> Tensor:
loc_to_center = positions - self.center
sgn_dist_from_surface = math.abs(loc_to_center) - self.half_size
if outward:
# --- get negative distances (particles are inside) towards the nearest boundary and add shift_amount ---
distances_of_interest = (sgn_dist_from_surface == math.max(
sgn_dist_from_surface, 'vector')) & (sgn_dist_from_surface < 0)
shift = distances_of_interest * (sgn_dist_from_surface -
shift_amount)
else:
shift = (sgn_dist_from_surface + shift_amount) * (
sgn_dist_from_surface > 0
) # get positive distances (particles are outside) and add shift_amount
shift = math.where(
math.abs(shift) > math.abs(loc_to_center),
math.abs(loc_to_center),
shift) # ensure inward shift ends at center
return positions + math.where(loc_to_center < 0, 1, -1) * shift
def approximate_signed_distance(self, location):
"""
Computes the signed L-infinity norm (manhattan distance) from the location to the nearest side of the box.
For an outside location `l` with the closest surface point `s`, the distance is `max(abs(l - s))`.
For inside locations it is `-max(abs(l - s))`.
:param location: float tensor of shape (batch_size, ..., rank)
:return: float tensor of shape (*location.shape[:-1], 1).
"""
lower, upper = math.batch_align([self.lower, self.upper], 1, location)
center = 0.5 * (lower + upper)
extent = upper - lower
distance = math.abs(location - center) - extent * 0.5
return math.max(distance, axis=-1, keepdims=True)
def plot_solves():
"""
While `plot_solves()` is active, certain performance optimizations and algorithm implementations may be disabled.
"""
from . import math
import pylab
cycle = pylab.rcParams['axes.prop_cycle'].by_key()['color']
with math.SolveTape(record_trajectories=True) as solves:
try:
yield solves
finally:
for i, result in enumerate(solves):
assert isinstance(result, math.SolveInfo)
from phi.math._tensors import disassemble_tree
_, (residual, ) = disassemble_tree(result.residual)
residual_mse = math.mean(math.sqrt(math.sum(residual**2)),
residual.shape.without('trajectory'))
residual_mse_max = math.max(
math.sqrt(math.sum(residual**2)),
residual.shape.without('trajectory'))
# residual_mean = math.mean(math.abs(residual), residual.shape.without('trajectory'))
residual_max = math.max(math.abs(residual),
residual.shape.without('trajectory'))
pylab.plot(residual_mse.numpy(),
label=f"{i}: {result.method}",
color=cycle[i % len(cycle)])
pylab.plot(residual_max.numpy(),
'--',
alpha=0.2,
color=cycle[i % len(cycle)])
pylab.plot(residual_mse_max.numpy(),
alpha=0.2,
color=cycle[i % len(cycle)])
print(
f"Solve {i}: {result.method} ({1000 * result.solve_time:.1f} ms)\n"
f"\t{result.solve}\n"
f"\t{result.msg}\n"
f"\tConverged: {result.converged}\n"
f"\tDiverged: {result.diverged}\n"
f"\tIterations: {result.iterations}\n"
f"\tFunction evaulations: {result.function_evaluations.trajectory[-1]}"
)
pylab.yscale('log')
pylab.ylabel("Residual: MSE / max / individual max")
pylab.xlabel("Iteration")
pylab.title(f"Solve Convergence")
pylab.legend(loc='upper right')
pylab.savefig(f"pressure-solvers-FP32.png")
pylab.show()
def approximate_signed_distance(self, location):
"""
Computes the signed L-infinity norm (manhattan distance) from the location to the nearest side of the box.
For an outside location `l` with the closest surface point `s`, the distance is `max(abs(l - s))`.
For inside locations it is `-max(abs(l - s))`.
Args:
location: float tensor of shape (batch_size, ..., rank)
Returns:
float tensor of shape (*location.shape[:-1], 1).
"""
center = 0.5 * (self.lower + self.upper)
extent = self.upper - self.lower
distance = math.abs(location - center) - extent * 0.5
return math.max(distance, 'vector')
def extrapolate(input_field, valid_mask, voxel_distance=10):
"""
Create a signed distance field for the grid, where negative signs are fluid cells and positive signs are empty cells. The fluid surface is located at the points where the interpolated value is zero. Then extrapolate the input field into the air cells.
:param domain: Domain that can create new Fields
:param input_field: Field to be extrapolated
:param valid_mask: One dimensional binary mask indicating where fluid is present
:param voxel_distance: Optional maximal distance (in number of grid cells) where signed distance should still be calculated / how far should be extrapolated.
:return: ext_field: a new Field with extrapolated values, s_distance: tensor containing signed distance field, depending only on the valid_mask
"""
ext_data = input_field.data
dx = input_field.dx
if isinstance(input_field, StaggeredGrid):
ext_data = input_field.staggered_tensor()
valid_mask = math.pad(valid_mask, [[0, 0]] +
[[0, 1]] * input_field.rank + [[0, 0]],
"constant")
dims = range(input_field.rank)
# Larger than voxel_distance to be safe. It could start extrapolating velocities from outside voxel_distance into the field.
signs = -1 * (2 * valid_mask - 1)
s_distance = 2.0 * (voxel_distance + 1) * signs
surface_mask = create_surface_mask(valid_mask)
# surface_mask == 1 doesn't output a tensor, just a scalar, but >= works.
# Initialize the voxel_distance with 0 at the surface
# Previously initialized with -0.5*dx, i.e. the cell is completely full (center is 0.5*dx inside the fluid surface). For stability and looks this was changed to 0 * dx, i.e. the cell is only half full. This way small changes to the SDF won't directly change neighbouring empty cells to fluid cells.
s_distance = math.where((surface_mask >= 1),
-0.0 * math.ones_like(s_distance), s_distance)
directions = np.array(
list(itertools.product(*np.tile((-1, 0, 1), (len(dims), 1)))))
# First make a move in every positive direction (StaggeredGrid velocities there are correct, we want to extrapolate these)
if isinstance(input_field, StaggeredGrid):
for d in directions:
if (d <= 0).all():
continue
# Shift the field in direction d, compare new distances to old ones.
d_slice = tuple([(slice(1, None) if d[i] == -1 else
slice(0, -1) if d[i] == 1 else slice(None))
for i in dims])
d_field = math.pad(
ext_data, [[0, 0]] +
[([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0])
for i in dims] + [[0, 0]], "symmetric")
d_field = d_field[(slice(None), ) + d_slice + (slice(None), )]
d_dist = math.pad(
s_distance, [[0, 0]] +
[([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0])
for i in dims] + [[0, 0]], "symmetric")
d_dist = d_dist[(slice(None), ) + d_slice + (slice(None), )]
d_dist += np.sqrt((dx * d).dot(dx * d)) * signs
if (d.dot(d) == 1) and (d >= 0).all():
# Pure axis direction (1,0,0), (0,1,0), (0,0,1)
updates = (math.abs(d_dist) <
math.abs(s_distance)) & (surface_mask <= 0)
updates_velocity = updates & (signs > 0)
ext_data = math.where(
math.concat([(math.zeros_like(updates_velocity)
if d[i] == 1 else updates_velocity)
for i in dims],
axis=-1), d_field, ext_data)
s_distance = math.where(updates, d_dist, s_distance)
else:
# Mixed axis direction (1,1,0), (1,1,-1), etc.
continue
for _ in range(voxel_distance):
buffered_distance = 1.0 * s_distance # Create a copy of current voxel_distance. This should not be necessary...
for d in directions:
if (d == 0).all():
continue
# Shift the field in direction d, compare new distances to old ones.
d_slice = tuple([(slice(1, None) if d[i] == -1 else
slice(0, -1) if d[i] == 1 else slice(None))
for i in dims])
d_field = math.pad(
ext_data, [[0, 0]] +
[([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0])
for i in dims] + [[0, 0]], "symmetric")
d_field = d_field[(slice(None), ) + d_slice + (slice(None), )]
d_dist = math.pad(
s_distance, [[0, 0]] +
[([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0])
for i in dims] + [[0, 0]], "symmetric")
d_dist = d_dist[(slice(None), ) + d_slice + (slice(None), )]
d_dist += np.sqrt((dx * d).dot(dx * d)) * signs
# We only want to update velocity that is outside of fluid
updates = (math.abs(d_dist) <
math.abs(buffered_distance)) & (surface_mask <= 0)
updates_velocity = updates & (signs > 0)
ext_data = math.where(
math.concat([updates_velocity] * math.spatial_rank(ext_data),
axis=-1), d_field, ext_data)
buffered_distance = math.where(updates, d_dist, buffered_distance)
s_distance = buffered_distance
# Cut off inaccurate values
distance_limit = -voxel_distance * (2 * valid_mask - 1)
s_distance = math.where(
math.abs(s_distance) < voxel_distance, s_distance, distance_limit)
if isinstance(input_field, StaggeredGrid):
ext_field = input_field.with_data(ext_data)
stagger_slice = tuple([slice(0, -1) for i in dims])
s_distance = s_distance[(slice(None), ) + stagger_slice +
(slice(None), )]
else:
ext_field = input_field.copied_with(data=ext_data)
return ext_field, s_distance
def abs(self):
return StaggeredGrid(math.abs(self.staggered))
def abs(self):
return self.with_data(math.abs(self.data))
