def grid_sample(self,
resolution: math.Shape,
size,
shape: math.Shape = None):
shape = (self._shape if shape is None else shape) & resolution
for dim in channel(self._shape):
if dim.item_names[0] is None:
warnings.warn(
f"Please provide item names for Noise dim {dim} using {dim}='x,y,z'",
FutureWarning)
shape &= channel(**{dim.name: resolution.names})
rndj = math.to_complex(random_normal(shape)) + 1j * math.to_complex(
random_normal(shape)) # Note: there is no complex32
with math.NUMPY:
k = math.fftfreq(resolution) * resolution / math.tensor(
size) * math.tensor(self.scale) # in physical units
k = math.vec_squared(k)
lowest_frequency = 0.1
weight_mask = math.to_float(k > lowest_frequency)
# --- 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 test_tensor_as_field(self):
t = math.random_normal(spatial(x=4, y=3), channel(vector='x,y'))
grid = field.tensor_as_field(t)
self.assertIsInstance(grid, CenteredGrid)
math.assert_close(grid.dx, 1)
math.assert_close(grid.points.x[0].y[0], 0)
t = math.random_normal(instance(points=5), channel(vector='x,y'))
points = field.tensor_as_field(t)
self.assertTrue((points.elements.bounding_radius() > 0).all)
self.assertIsInstance(points, PointCloud)
def test_op2_incompatible_item_names(self):
t1 = math.random_normal(channel(vector='x,y,z'))
t2 = math.random_normal(channel(vector='r,g,b'))
self.assertEqual(('r', 'g', 'b'), t2.vector.item_names)
try:
t1 + t2
self.fail("Tensors with incompatible item names cannot be added")
except math.IncompatibleShapes:
pass
t1 + t1
t2_ = t2 + math.random_normal(channel(vector=3))
self.assertEqual(('r', 'g', 'b'), t2_.vector.item_names)
t2_ = math.random_normal(channel(vector=3)) + t2
self.assertEqual(('r', 'g', 'b'), t2_.vector.item_names)
def test_zeros_nonuniform(self):
nonuniform = shape_stack('stack', BATCH_DIM, shape(time=1, x=3, y=3),
shape(x=3, y=4), shape())
self.assertEqual(math.zeros(nonuniform).shape, nonuniform)
self.assertEqual(math.ones(nonuniform).shape, nonuniform)
self.assertEqual(math.random_normal(nonuniform).shape, nonuniform)
self.assertEqual(math.random_uniform(nonuniform).shape, nonuniform)
def test_linear_operator(self):
GLOBAL_AXIS_ORDER.x_last()
direct = math.random_normal(batch=3, x=4, y=3) # , vector=2
op = lin_placeholder(direct)
def linear_function(val):
val = -val
val *= 2
val = math.pad(val, {'x': (2, 0), 'y': (0, 1)}, extrapolation.PERIODIC)
val = val.x[:-2].y[1:] + val.x[2:].y[:-1]
val = math.pad(val, {'x': (0, 0), 'y': (0, 1)}, extrapolation.ZERO)
val = math.pad(val, {'x': (2, 2), 'y': (0, 1)}, extrapolation.BOUNDARY)
# sl = sl.vector[0]
return val
val = val.x[1:4].y[:2]
return math.sum([val, sl], axis=0) - sl
functions = [
linear_function,
lambda val: math.gradient(val, difference='forward', padding=extrapolation.ZERO, dims='x').gradient[0],
lambda val: math.gradient(val, difference='backward', padding=extrapolation.PERIODIC, dims='x').gradient[0],
lambda val: math.gradient(val, difference='central', padding=extrapolation.BOUNDARY, dims='x').gradient[0],
]
for f in functions:
direct_result = f(direct)
# print(direct_result.batch[0], 'Direct result')
op_result = f(op)
# print(op_result.build_sparse_coordinate_matrix().todense())
self.assertIsInstance(op_result, ShiftLinOp)
op_result = NativeTensor(op_result.native(), op_result.shape)
# print(op_result.batch[0], 'Placeholder result')
math.assert_close(direct_result, op_result)
def test_zeros_nonuniform(self):
nonuniform = shape_stack(batch('stack'),
batch(time=1) & spatial(x=3, y=3),
spatial(x=3, y=4), channel())
self.assertEqual(math.zeros(nonuniform).shape, nonuniform)
self.assertEqual(math.ones(nonuniform).shape, nonuniform)
self.assertEqual(math.random_normal(nonuniform).shape, nonuniform)
self.assertEqual(math.random_uniform(nonuniform).shape, nonuniform)
def test_fft_dims(self):
for backend in BACKENDS:
with backend:
x = math.random_normal(batch(x=8, y=6, z=4))
k3 = math.fft(x, 'x,y,z')
k = x
for dim in 'xyz':
k = math.fft(k, dim)
math.assert_close(k, k3, abs_tolerance=1e-5, msg=backend.name)
def test_ifft(self):
dimensions = 'xyz'
for backend in BACKENDS:
with backend:
for d in range(1, len(dimensions) + 1):
x = math.random_normal(spatial(**{dim: 6 for dim in dimensions[:d]})) + math.tensor((0, 1), batch('batch'))
k = math.fft(x)
x_ = math.ifft(k)
math.assert_close(x, x_, abs_tolerance=1e-5, msg=backend.name)
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 test_convert_point_cloud(self):
loc = math.random_uniform(instance(points=4), channel(vector=2))
val = math.random_normal(instance(points=4), channel(vector=2))
points = PointCloud(Sphere(loc, radius=1), val)
for backend in BACKENDS:
converted = field.convert(points, backend)
self.assertEqual(converted.values.default_backend, backend)
self.assertEqual(converted.elements.center.default_backend,
backend)
self.assertEqual(converted.elements.radius.default_backend,
backend)
def test_repr(self):
print("--- Eager ---")
print(repr(math.zeros(batch(b=10))))
print(repr(math.zeros(batch(b=10)) > 0))
print(repr(math.ones(channel(vector=3))))
print(repr(math.ones(channel(vector=3), dtype=DType(int, 64))))
print(repr(math.ones(channel(vector=3), dtype=DType(float, 64))))
print(repr(math.ones(batch(vector=3))))
print(repr(math.random_normal(batch(b=10))))
print(
repr(
math.random_normal(batch(b=10), dtype=DType(float, 64)) *
1e-6))
def tracable(x):
print(x)
return x
print("--- Placeholders ---")
for backend in BACKENDS:
if backend.supports(Backend.jit_compile):
with backend:
math.jit_compile(tracable)(math.ones(channel(vector=3)))
def test_grid_sample_backend_equality_2d_batched(self):
grid = math.random_normal(mybatch=10, y=10, x=7)
coords = math.random_uniform(mybatch=10, x=3, y=2) * (12, 9)
grid_ = math.tensor(grid.native('mybatch,x,y'), 'mybatch,x,y')
coords_ = coords.vector.flip()
for extrap in (extrapolation.ZERO, extrapolation.ONE,
extrapolation.BOUNDARY, extrapolation.PERIODIC):
sampled = []
for backend in BACKENDS:
with backend:
grid, coords, grid_, coords_ = math.tensors(
grid, coords, grid_, coords_)
sampled.append(math.grid_sample(grid, coords, extrap))
sampled.append(math.grid_sample(grid_, coords_, extrap))
math.assert_close(*sampled, abs_tolerance=1e-5)
def test_jit_compile_linear(self):
math.GLOBAL_AXIS_ORDER.x_last()
x = math.random_normal(batch(batch=3)
& spatial(x=4, y=3)) # , vector=2
def linear_function(val):
val = -val
val *= 2
val = math.pad(val, {
'x': (2, 0),
'y': (0, 1)
}, math.extrapolation.PERIODIC)
val = val.x[:-2].y[1:] + val.x[2:].y[:-1]
val = math.pad(val, {
'x': (0, 0),
'y': (0, 1)
}, math.extrapolation.ZERO)
val = math.pad(val, {
'x': (2, 2),
'y': (0, 1)
}, math.extrapolation.BOUNDARY)
return math.sum([val, val], dim='0') - val
functions = [
linear_function,
lambda val: math.spatial_gradient(val,
difference='forward',
padding=math.extrapolation.ZERO,
dims='x').gradient[0],
lambda val: math.spatial_gradient(val,
difference='backward',
padding=math.extrapolation.
PERIODIC,
dims='x').gradient[0],
lambda val: math.spatial_gradient(val,
difference='central',
padding=math.extrapolation.
BOUNDARY,
dims='x').gradient[0],
]
for f in functions:
direct_result = f(x)
jit_f = math.jit_compile_linear(f)
jit_result = jit_f(x)
math.assert_close(direct_result, jit_result)
def test_split_dimension(self):
grid = math.random_normal(batch(batch=10) & spatial(x=4, y=3) & channel(vector=2))
points = math.pack_dims(grid, grid.shape.spatial, instance('points'))
split = points.points.split(grid.shape.spatial)
self.assertEqual(grid.shape, split.shape)
math.assert_close(grid, split)
def test_serialize_tensor(self):
t = math.random_normal(batch(batch=10), spatial(x=4, y=3),
channel(vector=2))
math.assert_close(t, math.from_dict(math.to_dict(t)))
def test_join_dimensions(self):
grid = math.random_normal(batch=10, x=4, y=3, vector=2)
points = math.join_dimensions(grid, grid.shape.spatial, 'points')
self.assertEqual(('batch', 'points', 'vector'), points.shape.names)
self.assertEqual(grid.shape.volume, points.shape.volume)
self.assertEqual(grid.shape.non_spatial, points.shape.non_spatial)
def test_join_dimensions(self):
grid = math.random_normal(batch(batch=10) & spatial(x=4, y=3) & channel(vector=2))
points = math.pack_dims(grid, grid.shape.spatial, instance('points'))
self.assertEqual(('batch', 'points', 'vector'), points.shape.names)
self.assertEqual(grid.shape.volume, points.shape.volume)
self.assertEqual(grid.shape.non_spatial, points.shape.non_instance)
def test_split_dimension(self):
grid = math.random_normal(batch=10, x=4, y=3, vector=2)
points = math.join_dimensions(grid, grid.shape.spatial, 'points')
split = points.points.split(grid.shape.spatial)
self.assertEqual(grid.shape, split.shape)
math.assert_close(grid, split)
def test_plot_point_cloud_3d_points(self):
self._test_plot(
PointCloud(
math.random_normal(instance(points=5),
channel(vector='x,y,z'))))
