def create_surface_mask(liquid_mask):
"""
Computes inner contours of the liquid_mask.
A cell i is flagged 1 if liquid_mask[i] = 1 and it has a non-liquid neighbour.
:param liquid_mask: binary tensor
:return: tensor
"""
# When we create inner contour, we don't want the fluid-wall boundaries to show up as surface, so we should pad with symmetric edge values.
mask = math.pad(liquid_mask, [[0, 0]] +
[[1, 1]] * math.spatial_rank(liquid_mask) + [[0, 0]],
"constant")
dims = range(math.spatial_rank(mask))
bcs = math.zeros_like(liquid_mask)
# Move in every possible direction to assure corners are properly set.
directions = np.array(
list(itertools.product(*np.tile((-1, 0, 1), (len(dims), 1)))))
for d in directions:
d_slice = tuple([(slice(2, None) if d[i] == -1 else
slice(0, -2) if d[i] == 1 else slice(1, -1))
for i in dims])
center_slice = tuple([slice(1, -1) for _ in dims])
# Create inner contour of particles
bc_d = math.maximum(mask[(slice(None),) + d_slice + (slice(None),)],
mask[(slice(None),) + center_slice + (slice(None),)]) - \
mask[(slice(None),) + d_slice + (slice(None),)]
bcs = math.maximum(bcs, bc_d)
return bcs
def _shift_resample(self,
resolution: Shape,
bounds: Box,
threshold=1e-5,
max_padding=20):
assert math.all_available(
bounds.lower, bounds.upper
), "Shift resampling requires 'bounds' to be available."
lower = math.to_int32(
math.ceil(
math.maximum(0, self.box.lower - bounds.lower) / self.dx -
threshold))
upper = math.to_int32(
math.ceil(
math.maximum(0, bounds.upper - self.box.upper) / self.dx -
threshold))
total_padding = (math.sum(lower) + math.sum(upper)).numpy()
if total_padding > max_padding:
return NotImplemented
elif total_padding > 0:
from phi.field import pad
padded = pad(
self, {
dim: (int(lower[i]), int(upper[i]))
for i, dim in enumerate(self.shape.spatial.names)
})
grid_box, grid_resolution, grid_values = padded.box, padded.resolution, padded.values
else:
grid_box, grid_resolution, grid_values = self.box, self.resolution, self.values
origin_in_local = grid_box.global_to_local(
bounds.lower) * grid_resolution
data = math.sample_subgrid(grid_values, origin_in_local, resolution)
return data
def _shift_resample(self, resolution, box, threshold=1e-5, max_padding=20):
lower = math.to_int(
math.ceil(
math.maximum(0, self.box.lower - box.lower) / self.dx -
threshold))
upper = math.to_int(
math.ceil(
math.maximum(0, box.upper - self.box.upper) / self.dx -
threshold))
total_padding = math.sum(lower) + math.sum(upper)
if total_padding > max_padding:
return NotImplemented
elif total_padding > 0:
from phi.field import pad
padded = pad(
self, {
dim: (int(lower[i]), int(upper[i]))
for i, dim in enumerate(self.shape.spatial.names)
})
grid_box, grid_resolution, grid_values = padded.box, padded.resolution, padded.values
else:
grid_box, grid_resolution, grid_values = self.box, self.resolution, self.values
origin_in_local = grid_box.global_to_local(box.lower) * grid_resolution
data = math.sample_subgrid(grid_values, origin_in_local, resolution)
return data
def _stagger_sample(self, box, resolution):
"""
Samples this field on a staggered grid.
In addition to sampling, extrapolates the field using an occupancy mask generated from the points.
:param box: physical dimensions of the grid
:param resolution: grid resolution
:return: StaggeredGrid
"""
resolution = np.array(resolution)
valid_indices = math.to_int(math.floor(self.sample_points))
valid_indices = math.minimum(math.maximum(0, valid_indices), resolution - 1)
# Correct format for math.scatter
valid_indices = batch_indices(valid_indices)
active_mask = math.scatter(self.sample_points, valid_indices, 1, math.concat([[valid_indices.shape[0]], resolution, [1]], axis=-1), duplicates_handling='any')
mask = math.pad(active_mask, [[0, 0]] + [[1, 1]] * self.rank + [[0, 0]], "constant")
if isinstance(self.data, (int, float, np.ndarray)):
values = math.zeros_like(self.sample_points) + self.data
else:
values = self.data
result = []
ones_1d = math.unstack(math.ones_like(values), axis=-1)[0]
staggered_shape = [i + 1 for i in resolution]
dx = box.size / resolution
dims = range(len(resolution))
for d in dims:
staggered_offset = math.stack([(0.5 * dx[i] * ones_1d if i == d else 0.0 * ones_1d) for i in dims], axis=-1)
indices = math.to_int(math.floor(self.sample_points + staggered_offset))
valid_indices = math.maximum(0, math.minimum(indices, resolution))
valid_indices = batch_indices(valid_indices)
values_d = math.expand_dims(math.unstack(values, axis=-1)[d], axis=-1)
result.append(math.scatter(self.sample_points, valid_indices, values_d, [indices.shape[0]] + staggered_shape + [1], duplicates_handling=self.mode))
d_slice = tuple([(slice(0, -2) if i == d else slice(1,-1)) for i in dims])
u_slice = tuple([(slice(2, None) if i == d else slice(1,-1)) for i in dims])
active_mask = math.minimum(mask[(slice(None),) + d_slice + (slice(None),)], active_mask)
active_mask = math.minimum(mask[(slice(None),) + u_slice + (slice(None),)], active_mask)
staggered_tensor_prep = unstack_staggered_tensor(math.concat(result, axis=-1))
grid_values = StaggeredGrid(staggered_tensor_prep)
# Fix values at boundary of liquids (using StaggeredGrid these might not receive a value, so we replace it with a value inside the liquid)
grid_values, _ = extrapolate(grid_values, active_mask, voxel_distance=2)
return grid_values
def cell_index(self, global_position):
local_position = self.box.global_to_local(
global_position) * self.resolution
position = math.to_int(local_position - 0.5)
position = math.maximum(0, position)
position = math.minimum(position, self.resolution - 1)
return position
def _grid_sample(self, box, resolution):
"""
Samples this field on a regular grid.
:param box: physical dimensions of the grid
:param resolution: grid resolution
:return: CenteredGrid
"""
sample_indices_nd = math.to_int(
math.round(box.global_to_local(self.sample_points) * resolution))
sample_indices_nd = math.minimum(
math.maximum(0, sample_indices_nd), resolution - 1
) # Snap outside points to edges, otherwise scatter raises an error
# Correct format for math.scatter
valid_indices = _batch_indices(sample_indices_nd)
shape = (math.shape(
self.data)[0], ) + tuple(resolution) + (self.data.shape[-1], )
scattered = math.scatter(self.sample_points,
valid_indices,
self.data,
shape,
duplicates_handling=self.mode)
return CenteredGrid(data=scattered,
box=box,
extrapolation='constant',
name=self.name + '_centered')
def _grid_sample(self, box, resolution):
"""
Samples this field on a regular grid.
:param box: physical dimensions of the grid
:param resolution: grid resolution
:return: CenteredGrid
"""
valid_indices = math.to_int(math.floor(self.sample_points))
valid_indices = math.minimum(math.maximum(0, valid_indices), resolution - 1)
# Correct format for math.scatter
valid_indices = batch_indices(valid_indices)
scattered = math.scatter(self.sample_points, valid_indices, self.data, math.concat([[valid_indices.shape[0]], resolution, [1]], axis=-1), duplicates_handling=self.mode)
return CenteredGrid(data=scattered, box=box, extrapolation='constant', name=self.name+'_centered')
def test_precision_16(self):
try:
math.set_precision(16)
fluid = Fluid(Domain([16, 16]), density=math.maximum(0, Noise()))
self.assertEqual(fluid.density.data.dtype, numpy.float16)
self.assertEqual(fluid.velocity.unstack()[0].data.dtype,
numpy.float16)
fluid = IncompressibleFlow().step(fluid, dt=1.0)
self.assertEqual(fluid.density.data.dtype, numpy.float16)
self.assertEqual(fluid.velocity.unstack()[0].data.dtype,
numpy.float16)
finally:
math.set_precision(32) # Reset environment
def maximum(f1: Field or Geometry or float, f2: Field or Geometry or float):
"""
Element-wise maximum.
One of the given fields needs to be an instance of `SampledField` and the the result will be sampled at the corresponding points.
If both are `SampledFields` but have different points, `f1` takes priority.
Args:
f1: `Field` or `Geometry` or constant.
f2: `Field` or `Geometry` or constant.
Returns:
`SampledField`
"""
f1, f2 = _auto_resample(f1, f2)
return f1.with_values(math.maximum(f1.values, f2.values))
def approximate_signed_distance(self, location):
"""
Computes the exact distance from location to the closest point on the sphere.
Very close to the sphere center, the distance takes a constant value.
:param location: float tensor of shape (batch_size, ..., rank)
:return: float tensor of shape (*location.shape[:-1], 1).
"""
center = math.batch_align(self.center, 1, location)
radius = math.batch_align(self.radius, 0, location)
distance_squared = math.sum((location - center)**2,
axis=-1,
keepdims=True)
distance_squared = math.maximum(
distance_squared,
radius * 1e-2) # Prevent infinite gradient at sphere center
distance = math.sqrt(distance_squared)
return distance - radius
def approximate_signed_distance(self, location):
"""
Computes the exact distance from location to the closest point on the sphere.
Very close to the sphere center, the distance takes a constant value.
Args:
location: float tensor of shape (batch_size, ..., rank)
Returns:
float tensor of shape (*location.shape[:-1], 1).
"""
distance_squared = math.vec_squared(location - self.center)
distance_squared = math.maximum(
distance_squared,
self.radius * 1e-2) # Prevent infinite gradient at sphere center
distance = math.sqrt(distance_squared)
return distance - self.radius
def _required_paddings_transposed(box, dx, target):
lower = math.to_int(
math.ceil(math.maximum(0, box.lower - target.lower) / dx))
upper = math.to_int(
math.ceil(math.maximum(0, target.upper - box.upper) / dx))
return [lower, upper]
def test_maximum(self):
v = math.ones(x=4, y=3, vector=2)
math.assert_close(math.maximum(0, v), 1)
math.assert_close(math.maximum(0, -v), 0)
def soft_sqrt(self):
return StaggeredGrid(math.sqrt(math.maximum(self.staggered, 1e-20)))
def test_maximum(self):
v = math.ones(spatial(x=4, y=3) & channel(vector=2))
math.assert_close(math.maximum(0, v), 1)
math.assert_close(math.maximum(0, -v), 0)
