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)
