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'))))