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)] }, } }