def gravity_tensor(gravity, rank): if isinstance(gravity, Gravity): gravity = gravity.gravity if math.is_scalar(gravity): return math.to_float( math.expand_dims([gravity] + [0] * (rank - 1), 0, rank + 1)) else: assert math.staticshape(gravity)[-1] == rank return math.to_float( math.expand_dims(gravity, 0, rank + 2 - len(math.staticshape(gravity))))
def sample_at(self, points): if self.mode == 'EXP': envelope = math.exp(-0.5 * math.sum( (points - self.center)**2, axis=-1, keepdims=True) / self.size**2) envelope = math.to_float(envelope) return envelope * self.factor elif self.mode == 'RECT': conf = np.zeros(points.shape) conf[:, self.center[0] - self.size:self.center[0] + self.size, self.center[1] - self.size:self.center[1] + self.size, :] = np.ones( conf[:, self.center[0] - self.size:self.center[0] + self.size, self.center[1] - self.size:self.center[1] + self.size, :].shape) return conf[:, :, :, :-1] * self.factor elif self.mode == 'RANDOM': conf = np.random.random_sample(points.shape) conf[:, 0, :, :] *= 0 conf[:, -1, :, :] *= 0 conf[:, :, 0, :] *= 0 conf[:, :, -1, :] *= 0 return conf[:, :, :, :-1] * self.factor
def grid_sample(self, resolution, size, batch_size=1, channels=None): channels = channels or self.channels or len(size) shape = (batch_size, ) + tuple(resolution) + (channels, ) rndj = math.to_complex( self.math.random_normal(shape)) + 1j * math.to_complex( self.math.random_normal(shape)) # Note: there is no complex32 k = math.fftfreq( resolution) * resolution / size * self.scale # in physical units k = math.sum(k**2, axis=-1, keepdims=True) lowest_frequency = 0.1 weight_mask = 1 / (1 + math.exp( (lowest_frequency - k) * 1e3)) # High pass filter # --- Compute 1/k --- k[(0, ) * len(k.shape)] = np.inf inv_k = 1 / k inv_k[(0, ) * len(k.shape)] = 0 # --- Compute result --- fft = rndj * inv_k**self.smoothness * weight_mask array = math.real(math.ifft(fft)) array /= math.std(array, axis=tuple(range(1, math.ndims(array))), keepdims=True) array -= math.mean(array, axis=tuple(range(1, math.ndims(array))), keepdims=True) array = math.to_float(array) return array
def __call__(self, *args, **kwargs): if len(args) == 2 and len(kwargs) == 0: lower, upper = args return AABox(lower, upper) elif len(args) == 1 and 'size' in kwargs: center, = args size = kwargs['size'] return Cuboid(center, 0.5 * math.to_float(size)) elif 'size' in kwargs and 'center' in kwargs: center, size = kwargs['center'], kwargs['size'] return Cuboid(center, 0.5 * math.to_float(size)) elif len(args) == 1 and len(kwargs) == 0: return AABox(0, args[0]) else: raise ValueError('Cannot create box from args=%s, kwargs=%s' % (args, kwargs))
def distribute_points(density, particles_per_cell=1, distribution='uniform'): """ Distribute points according to the distribution specified in density. :param density: binary tensor :param particles_per_cell: integer :param distribution: 'uniform' or 'center' :return: tensor of shape (batch_size, point_count, rank) """ assert distribution in ('center', 'uniform') index_array = [] batch_size = math.staticshape(density)[0] if math.staticshape(density)[0] is not None else 1 for batch in range(batch_size): indices = math.where(density[batch, ..., 0] > 0) indices = math.to_float(indices) temp = [] for _ in range(particles_per_cell): if distribution == 'center': temp.append(indices + 0.5) elif distribution == 'uniform': temp.append(indices + math.random_uniform(math.shape(indices))) index_array.append(math.concat(temp, axis=0)) try: index_array = math.stack(index_array) return index_array except ValueError: raise ValueError("all arrays in the batch must have the same number of active cells.")
def data(self, data): assert math.is_tensor(data), data rank = math.ndims(data) assert rank <= 3 if rank < 3: assert rank in (0, 1) # Scalar field / vector field data = math.expand_dims(data, 0, 3 - rank) return math.to_float(data)
def value_at(self, location): 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) bool_inside = distance_squared <= radius**2 return math.to_float(bool_inside)
def gravity_tensor(gravity, rank): if isinstance(gravity, Gravity): gravity = gravity.gravity if math.is_scalar(gravity): gravity = gravity * GLOBAL_AXIS_ORDER.up_vector(rank) assert math.staticshape(gravity)[-1] == rank return math.to_float( math.expand_dims(gravity, 0, rank + 2 - len(math.staticshape(gravity))))
def test_cast(self): for backend in BACKENDS: with backend: x = math.random_uniform(dtype=DType(float, 64)) self.assertEqual(DType(float, 32), math.to_float(x).dtype, msg=backend.name) self.assertEqual(DType(complex, 64), math.to_complex(x).dtype, msg=backend.name) with math.precision(64): self.assertEqual(DType(float, 64), math.to_float(x).dtype, msg=backend.name) self.assertEqual(DType(complex, 128), math.to_complex(x).dtype, msg=backend.name) self.assertEqual(DType(int, 64), math.to_int64(x).dtype, msg=backend.name) self.assertEqual(DType(int, 32), math.to_int32(x).dtype, msg=backend.name) self.assertEqual(DType(float, 16), math.cast(x, DType(float, 16)).dtype, msg=backend.name) self.assertEqual(DType(complex, 128), math.cast(x, DType(complex, 128)).dtype, msg=backend.name) try: math.cast(x, DType(float, 3)) self.fail(msg=backend.name) except KeyError: pass
def _padded_resample(self, points): data = self.padded([[1, 1]] * self.rank).data local_points = self.box.global_to_local(points) local_points = local_points * math.to_float( self.resolution) + 0.5 # depends on amount of padding resampled = math.resample(data, local_points, boundary='replicate', interpolation=self.interpolation) return resampled
def sample_at(self, points, collapse_dimensions=True): if not isinstance(self.extrapolation, six.string_types): return self._padded_resample(points) local_points = self.box.global_to_local(points) local_points = math.mul(local_points, math.to_float( self.resolution)) - 0.5 resampled = math.resample(self.data, local_points, boundary=_pad_mode(self.extrapolation), interpolation=self.interpolation) return resampled
def sample_at(self, points): local_points = self.box.global_to_local(points) local_points = math.mul(local_points, math.to_float( self.resolution)) - 0.5 resampled = math.resample(self.data, local_points, boundary=_pad_mode(self.extrapolation), interpolation=self.interpolation, constant_values=_pad_value( self.extrapolation_value)) return resampled
def test_batched_forced_burgers_2d(self): world = World(batch_size=3) burgers = world.add(Domain([4, 4]).centered_grid(0, batch_size=world.batch_size, name='velocity')) k = math.to_float(numpy.random.uniform(3, 6, [world.batch_size, 2])) amplitude = numpy.random.uniform(-0.5, 0.5, [world.batch_size]) force = SinPotential(k, phase_offset=numpy.random.uniform(0, 2 * numpy.pi, [world.batch_size]), data=amplitude) physics = ForcingPhysics(numpy.random.uniform(-0.4, 0.4, [world.batch_size])) effect = FieldEffect(force, ['velocity']) world.add(effect, physics=physics) burgers.step() burgers.step()
def sample_at(self, points, collapse_dimensions=True): if not isinstance(self.extrapolation, six.string_types): return self._padded_resample(points) local_points = self.box.global_to_local(points) local_points = local_points * math.to_float(self.resolution) - 0.5 if self.extrapolation == 'periodic': data = math.pad(self.data, [[0, 0]] + [[0, 1]] * self.rank + [[0, 0]], mode='wrap') local_points = local_points % math.to_float( math.staticshape(self.data)[1:-1]) resampled = math.resample(data, local_points, interpolation=self.interpolation) else: boundary = 'replicate' if self.extrapolation == 'boundary' else 'zero' resampled = math.resample(self.data, local_points, boundary=boundary, interpolation=self.interpolation) return resampled
def buoyancy(density, gravity, buoyancy_factor): """ Computes the buoyancy force proportional to the density. :param density: CenteredGrid :param gravity: vector or float :param buoyancy_factor: float :return: StaggeredGrid for the domain of the density """ if isinstance(gravity, (int, float)): gravity = math.to_float(math.as_tensor([gravity] + ([0] * (density.rank - 1)))) result = StaggeredGrid.from_scalar(density, -gravity * buoyancy_factor) return result
def getpoints(box, resolution): idx_zyx = np.meshgrid(*[ np.linspace(0.5 / dim, 1 - 0.5 / dim, dim) for dim in resolution ], indexing="ij") local_coords = math.to_float( math.expand_dims(math.stack(idx_zyx, axis=-1), 0)) points = box.local_to_global(local_coords) return CenteredGrid(points, box, name='grid_centers(%s, %s)' % (box, resolution), flags=[SAMPLE_POINTS])
def general_sample_at(self, points, reduce): local_points = self.box.global_to_local(points) local_points = math.mul(local_points, math.to_float( self.resolution)) - 0.5 result = general_grid_sample_nd( self.data, local_points, boundary=_pad_mode(self.extrapolation), constant_values=_pad_value(self.extrapolation_value), math=math.choose_backend([self.data, points]), reduce=reduce) return result
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 l_n_loss(tensor, n, batch_norm=True, reduce_batches=True): if isinstance(tensor, StaggeredGrid): tensor = tensor.staggered if reduce_batches: total_loss = math.sum(tensor**n) / n else: total_loss = math.sum(tensor**n, axis=list(range(1, len(tensor.shape)))) / n if batch_norm: batch_size = math.shape(tensor)[0] return total_loss / math.to_float(batch_size) else: return total_loss
def sample_at(self, points, collapse_dimensions=True): if not isinstance(self.extrapolation, six.string_types): return self._padded_resample(points) local_points = self.box.global_to_local(points) local_points = math.mul(local_points, math.to_float( self.resolution)) - 0.5 boundary = { 'periodic': 'circular', 'boundary': 'replicate', 'constant': 'constant' }[self.extrapolation] resampled = math.resample(self.data, local_points, boundary=boundary, interpolation=self.interpolation) return resampled
def buoyancy(density, gravity, buoyancy_factor): """ Computes the buoyancy force proportional to the density. :param density: CenteredGrid :param gravity: vector or float :param buoyancy_factor: float :return: StaggeredGrid for the domain of the density """ warnings.warn( 'buoyancy() is deprecated. Use (density * -gravity * buoyancy_factor).at(target_grid) instead.', DeprecationWarning) if isinstance(gravity, (int, float)): gravity = math.to_float( math.as_tensor([gravity] + ([0] * (density.rank - 1)))) result = StaggeredGrid.from_scalar(density, -gravity * buoyancy_factor) return result
def grid_sample(self, resolution: math.Shape, size, shape: math.Shape = None): shape = (self._shape if shape is None else shape).combined(resolution) rndj = math.to_complex(random_normal(shape)) + 1j * math.to_complex(random_normal(shape)) # Note: there is no complex32 with math.NUMPY_BACKEND: k = math.fftfreq(resolution) * resolution / size * self.scale # in physical units k = math.vec_squared(k) lowest_frequency = 0.1 weight_mask = 1 / (1 + math.exp((lowest_frequency - k) * 1e3)) # High pass filter # --- Compute 1/k --- k.native()[(0,) * len(k.shape)] = np.inf inv_k = 1 / k inv_k.native()[(0,) * len(k.shape)] = 0 # --- Compute result --- fft = rndj * inv_k ** self.smoothness * weight_mask array = math.real(math.ifft(fft)) array /= math.std(array, dim=array.shape.non_batch) array -= math.mean(array, dim=array.shape.non_batch) array = math.to_float(array) return array
def _distribute_points(mask: math.Tensor, points_per_cell: int = 1, center: bool = False) -> math.Tensor: """ Generates points (either uniformly distributed or at the cell centers) according to the given tensor mask. Args: mask: Tensor with nonzero values at the indices where particles should get generated. points_per_cell: Number of particles to generate at each marked index center: Set points to cell centers. If False, points will be distributed using a uniform distribution within each cell. Returns: A tensor containing the positions of the generated points. """ indices = math.to_float(math.nonzero(mask, list_dim=instance('points'))) temp = [] for _ in range(points_per_cell): if center: temp.append(indices + 0.5) else: temp.append(indices + (math.random_uniform(indices.shape))) return math.concat(temp, dim=instance('points'))
def _read_npy_files(items): data = [read_zipped_array(item.decode())[0, ...] for item in items] data = [math.to_float(array) for array in data] return data
def half_size(self, half_size): return math.to_float(half_size)
def center(self, center): return math.to_float(center)
def upper(self, upper): return math.to_float(upper)
def lower(self, lower): return math.to_float(lower)
def location(self, loc): loc = math.to_float(loc) assert math.staticshape(loc)[-1] in (2, 3) if math.ndims(loc) < 2: loc = math.expand_dims(loc, axis=0, number=2 - math.ndims(loc)) return loc
def strength(self, strength): return math.to_float(strength)