def test_unstack(self):
a = math.random_uniform(batch(b=10), spatial(x=4, y=3), channel(vector=2))
u = math.unstack(a, 'vector')
self.assertIsInstance(u, tuple)
self.assertEqual(len(u), 2)
math.assert_close(u, math.unstack(a, channel))
# Multiple dimensions
u = math.unstack(a, 'x,y')
self.assertIsInstance(u, tuple)
self.assertEqual(len(u), 12)
math.assert_close(u, math.unstack(a, spatial))
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 multi_advect(self, fields, interpolation="LINEAR", dt=1):
assert isinstance(
fields,
(list, tuple)), "first parameter must be either a tuple or list"
inputs_lists = []
coords_lists = []
value_generators = []
for field in fields:
if isinstance(field, StaggeredGrid):
i, c, v = self._mac_block_advection(field.staggered, dt)
else:
i, c, v = self._centered_block_advection(field, dt)
inputs_lists.append(i)
coords_lists.append(c)
value_generators.append(v)
inputs = math.concat(sum(inputs_lists, []), 0)
coords = math.concat(sum(coords_lists, []), 0)
all_advected = math.resample(inputs,
coords,
interpolation=interpolation,
boundary="REPLICATE")
all_advected = math.reshape(all_advected, [self.spatial_rank, -1] +
list(all_advected.shape[1:]))
all_advected = math.unstack(all_advected)
results = []
abs_i = 0
for i in range(len(inputs_lists)):
n = len(inputs_lists[0])
assigned_advected = all_advected[abs_i:abs_i + n]
results.append(value_generators[i](assigned_advected))
abs_i += n
return results
def unstack_staggered_tensor(tensor):
tensors = math.unstack(tensor, -1)
result = []
for i, dim in enumerate(math.spatial_dimensions(tensor)):
slices = [slice(None, -1) if d != dim else slice(None) for d in math.spatial_dimensions(tensor)]
result.append(math.expand_dims(tensors[i][tuple([slice(None)]+slices)], -1))
return result
def __dataop__(self, other, linear_if_scalar, data_operator):
if isinstance(other, StaggeredGrid):
assert self.compatible(
other), 'Fields are not compatible: %s and %s' % (self, other)
data = [
data_operator(c1, c2) for c1, c2 in zip(self.data, other.data)
]
flags = propagate_flags_operation(self.flags + other.flags, False,
self.rank, self.component_count)
elif math.ndims(other) > 0 and math.staticshape(other)[-1] > 1:
other_components = math.unstack(math.as_tensor(other),
axis=-1,
keepdims=True)
data = [
data_operator(c1, c2)
for c1, c2 in zip(self.data, other_components)
]
flags = propagate_flags_operation(self.flags, False, self.rank,
self.component_count)
else:
data = [data_operator(c1, other) for c1 in self.data]
flags = propagate_flags_operation(self.flags, linear_if_scalar,
self.rank, self.component_count)
return self.copied_with(data=np.array(data, dtype=np.object),
flags=flags)
def unstack(self):
flags = propagate_flags_children(self.flags, self.rank, 1)
return [
ConstantField(c,
'%s[%d]' % (self.name, i),
flags=flags,
batch_size=self._batch_size)
for i, c in enumerate(math.unstack(self.data, -1))
]
def unstack(self):
flags = propagate_flags_children(self.flags, self.rank, 1)
return [
CenteredGrid(math.expand_dims(component),
box=self.box,
name='%s[...,%d]' % (self.name, i),
flags=flags,
batch_size=self._batch_size)
for i, component in enumerate(math.unstack(self.data, -1))
]
def unstack(self):
flags = propagate_flags_children(self.flags, self.rank, 1)
return [
GeometryMask(self.geometries,
c,
'%s[%d]' % (self.name, i),
flags,
batch_size=self._batch_size)
for i, c in enumerate(math.unstack(self.data, -1))
]
def unstack(self):
flags = propagate_flags_children(self.flags, self.rank, 1)
components = math.unstack(self.data, axis=-1, keepdims=True)
return [
CenteredGrid(component,
box=self.box,
flags=flags,
batch_size=self._batch_size)
for i, component in enumerate(components)
]
def unstack(self):
unstacked = {}
for arg in self.function_args:
if isinstance(arg, Field):
unstacked[arg] = arg.unstack()
elif math.is_tensor(arg) and math.ndims(arg) > 0:
unstacked[arg] = math.unstack(arg, axis=-1, keepdims=True)
else:
unstacked[arg] = [arg] * self.component_count
assert len(unstacked[arg]) == self.component_count
result = [_SymbolicOpField(self.function, [unstacked[arg][i] for arg in self.function_args]) for i in range(self.component_count)]
return result
def staggered_tensor(self) -> Tensor:
"""
Stacks all component grids into a single uniform `phi.math.Tensor`.
The individual components are padded to a common (larger) shape before being stacked.
The shape of the returned tensor is exactly one cell larger than the grid `resolution` in every spatial dimension.
Returns:
Uniform `phi.math.Tensor`.
"""
padded = []
for dim, component in zip(self.resolution.names,
math.unstack(self.values, 'vector')):
widths = {d: (0, 1) for d in self.resolution.names}
lo_valid, up_valid = self.extrapolation.valid_outer_faces(dim)
widths[dim] = (int(not lo_valid), int(not up_valid))
padded.append(math.pad(component, widths, mode=self.extrapolation))
result = math.stack(padded, channel('vector'))
assert result.shape.is_uniform
return result
def global_to_child(self, location):
""" Inverse transform """
delta = location - self.center
if math.staticshape(location)[-1] == 2:
angle = -math.batch_align(self.angle, 0, location)
sin = math.sin(angle)
cos = math.cos(angle)
y, x = math.unstack(delta, axis=-1)
if GLOBAL_AXIS_ORDER.is_x_first:
x, y = y, x
rot_x = cos * x - sin * y
rot_y = sin * x + cos * y
rotated = math.stack([rot_y, rot_x], axis=-1)
elif math.staticshape(location)[-1] == 3:
angle = math.batch_align(self.angle, 1, location)
raise NotImplementedError(
'not yet implemented') # ToDo apply angle
else:
raise NotImplementedError('Rotation only supported in 2D and 3D')
final = rotated + self.center
return final
def sample_at(self, points):
points_rank = math.spatial_rank(points)
src_rank = math.spatial_rank(self.location)
# --- Expand shapes to format (batch_size, points_dims..., src_dims..., channels) ---
points = math.expand_dims(points, axis=-2, number=src_rank)
src_points = math.expand_dims(self.location,
axis=-2,
number=points_rank)
src_strength = math.expand_dims(self.strength, axis=-1)
src_strength = math.batch_align(src_strength, 0, self.location)
src_strength = math.expand_dims(src_strength,
axis=-1,
number=points_rank)
src_axes = tuple(range(-2, -2 - src_rank, -1))
# --- Compute distances and falloff ---
distances = points - src_points
if self.falloff is not None:
raise NotImplementedError()
# distances_squared = math.sum(distances ** 2, axis=-1, keepdims=True)
# unit_distances = distances / math.sqrt(distances_squared)
# strength = src_strength * math.exp(-distances_squared)
else:
strength = src_strength
# --- Compute velocities ---
if math.staticshape(points)[-1] == 2: # Curl in 2D
dist_1, dist_2 = math.unstack(distances, axis=-1)
if GLOBAL_AXIS_ORDER.is_x_first:
velocity = strength * math.stack([-dist_2, dist_1], axis=-1)
else:
velocity = strength * math.stack([dist_2, -dist_1], axis=-1)
elif math.staticshape(points)[-1] == 3: # Curl in 3D
raise NotImplementedError('not yet implemented')
else:
raise AssertionError(
'Vector product not available in > 3 dimensions')
velocity = math.sum(velocity, axis=src_axes)
return velocity
def vector_field(field2d, settings):
assert isinstance(field2d, (CenteredGrid, StaggeredGrid))
if isinstance(field2d, StaggeredGrid):
field2d = field2d.at_centers()
assert isinstance(field2d, CenteredGrid)
assert field2d.rank == 2
batch = settings.get('batch', 0)
batch = min(batch, field2d.data.shape[0])
arrow_origin = settings.get('arrow_origin', 'tip')
assert arrow_origin in ('base', 'center', 'tip')
max_resolution = settings.get('max_arrow_resolution', 40)
max_arrows = settings.get('max_arrows', 300)
draw_full_arrows = settings.get('draw_full_arrows', False)
y, x = math.unstack(field2d.points.data[0, ..., -2:], axis=-1)
data_y, data_x = math.unstack(field2d.data[batch, ...], -1)[-2:]
while np.prod(x.shape) > max_resolution**2:
y = y[::2, ::2]
x = x[::2, ::2]
data_y = data_y[::2, ::2]
data_x = data_x[::2, ::2]
y = y.flatten()
x = x.flatten()
data_y = data_y.flatten()
data_x = data_x.flatten()
if max_arrows is not None and len(x) > max_arrows:
length = np.sqrt(data_y**2 + data_x**2)
keep_indices = np.argsort(length)[-max_arrows:]
# size = np.max(field2d.box.size)
# threshold = size * negligible_threshold
# keep_condition = (np.abs(data_x) > threshold) | (np.abs(data_y) > threshold)
# keep_indices = np.where(keep_condition)
y = y[keep_indices]
x = x[keep_indices]
data_y = data_y[keep_indices]
data_x = data_x[keep_indices]
if arrow_origin == 'tip':
x -= data_x
y -= data_y
elif arrow_origin == 'center':
x -= 0.5 * data_x
y -= 0.5 * data_y
if draw_full_arrows:
result = plotly_figures.create_quiver(x, y, data_x, data_y,
scale=1.0) # 7 points per arrow
result.update_xaxes(
range=[field2d.box.get_lower(1),
field2d.box.get_upper(1)])
result.update_yaxes(
range=[field2d.box.get_lower(0),
field2d.box.get_upper(0)])
return result
else:
lines_y = np.stack([y, y + data_y, [None] * len(x)],
-1).flatten() # 3 points per arrow
lines_x = np.stack([x, x + data_x, [None] * len(x)], -1).flatten()
return {
'data': [{
'mode': 'lines',
'x': lines_x,
'y': lines_y,
'type': 'scatter',
}],
'layout': {
'xaxis': {
'range':
[field2d.box.get_lower(1),
field2d.box.get_upper(1)]
},
'yaxis': {
'range':
[field2d.box.get_lower(0),
field2d.box.get_upper(0)]
},
}
}
