def test_functional_gradient(self): def f(x: math.Tensor, y: math.Tensor): assert isinstance(x, math.Tensor) assert isinstance(y, math.Tensor) pred = x loss = math.l2_loss(pred - y) return loss, pred for backend in BACKENDS: if backend.supports(Backend.functional_gradient): with backend: x_data = math.tensor(2.) y_data = math.tensor(1.) dx, = math.functional_gradient(f, wrt=[0], get_output=False)(x_data, y_data) math.assert_close(dx, 1, msg=backend.name) dx, dy = math.functional_gradient(f, [0, 1], get_output=False)(x_data, y_data) math.assert_close(dx, 1, msg=backend.name) math.assert_close(dy, -1, msg=backend.name) (loss, pred), (dx, dy) = math.functional_gradient( f, [0, 1], get_output=True)(x_data, y_data) math.assert_close(loss, 0.5, msg=backend.name) math.assert_close(pred, x_data, msg=backend.name) math.assert_close(dx, 1, msg=backend.name) math.assert_close(dy, -1, msg=backend.name)
def test_batch_independence(self): def simulate(centers): world = World() fluid = world.add(Fluid(Domain(x=5, y=4, boundaries=CLOSED, bounds=Box(0, [40, 32])), buoyancy_factor=0.1, batch_size=centers.shape[0]), physics=IncompressibleFlow()) world.add(Inflow(Sphere(center=centers, radius=3), rate=0.2)) world.add( Fan(Sphere(center=centers, radius=5), acceleration=[1.0, 0])) world.step(dt=1.5) world.step(dt=1.5) world.step(dt=1.5) assert not math.close(fluid.density.values, 0) print() return fluid.density.values.batch[0], fluid.velocity.values.batch[ 0] d1, v1 = simulate(math.tensor([[5, 16], [5, 4]], names='batch,vector')) d2, v2 = simulate(math.tensor([[5, 16], [5, 16]], names='batch,vector')) math.assert_close(d1, d2) math.assert_close(v1, v2)
def test_np_speed_sum(self): print() np1, np2 = rnpv(64), rnpv(256) t1 = math.tensor(np1, batch('batch'), spatial('x, y'), channel('vector')) t2 = math.tensor(np2, batch('batch'), spatial('x, y'), channel('vector')) _assert_equally_fast(lambda: np.sum(np1), lambda: math.sum(t1, dim=t1.shape), n=10000) _assert_equally_fast(lambda: np.sum(np2), lambda: math.sum(t2, dim=t1.shape), n=10000)
def points(self, points: Tensor or Number or tuple or list, values: Tensor or Number = None, radius: Tensor or float or int or None = None, extrapolation: math.Extrapolation = math.extrapolation.ZERO, color: str or Tensor or tuple or list or None = None) -> PointCloud: """ Create a `phi.field.PointCloud` from the given `points`. The created field has no channel dimensions and all points carry the value `1`. Args: points: point locations in physical units values: (optional) values of the particles, defaults to 1. radius: (optional) size of the particles extrapolation: (optional) extrapolation to use, defaults to extrapolation.ZERO color: (optional) color used when plotting the points Returns: `phi.field.PointCloud` object """ extrapolation = extrapolation if isinstance(extrapolation, math.Extrapolation) else self.boundaries[extrapolation] if radius is None: radius = math.mean(self.bounds.size) * 0.005 # --- Parse points: tuple / list --- if isinstance(points, (tuple, list)): if len(points) == 0: # no points points = math.zeros(instance(points=0), channel(vector=1)) elif isinstance(points[0], Number): # single point points = math.tensor([points], instance('points'), channel('vector')) else: points = math.tensor(points, instance('points'), channel('vector')) elements = Sphere(points, radius) if values is None: values = math.tensor(1.) return PointCloud(elements, values, extrapolation, add_overlapping=False, bounds=self.bounds, color=color)
def test_multi_dim_tensor_from_numpy(self): v = math.tensor(np.ones([1, 4, 3, 2]), batch('batch'), spatial('x,y'), channel('vector')) self.assertEqual((1, 4, 3, 2), v.shape.sizes) v = math.tensor(np.ones([10, 4, 3, 2]), batch('batch'), spatial('x,y'), channel('vector')) self.assertEqual((10, 4, 3, 2), v.shape.sizes)
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_extrapolate_valid_3D_3x3x3_2(self): valid = tensor([[[0, 0, 0], [0, 0, 0], [0, 0, 0]], [[0, 0, 1], [0, 0, 0], [1, 0, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 0]]], spatial('x, y, z')) values = tensor([[[0, 0, 0], [0, 0, 0], [0, 0, 0]], [[1, 0, 4], [0, 0, 0], [2, 0, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 0]]], spatial('x, y, z')) expected_valid = math.ones(spatial(x=3, y=3, z=3)) expected_values = tensor([[[3, 4, 4], [2, 3, 4], [2, 2, 3]], [[3, 4, 4], [2, 3, 4], [2, 2, 3]], [[3, 4, 4], [2, 3, 4], [2, 2, 3]]], spatial('x, y, z')) extrapolated_values, extrapolated_valid = math.extrapolate_valid_values(values, valid, 2) math.assert_close(extrapolated_values, expected_values) math.assert_close(extrapolated_valid, expected_valid)
def test_cos(self): for backend in BACKENDS: with backend: math.assert_close(math.cos(math.zeros(spatial(x=4))), 1, abs_tolerance=1e-6, msg=backend.name) math.assert_close(math.cos(math.tensor(math.PI / 2)), 0, abs_tolerance=1e-6, msg=backend.name) math.assert_close(math.cos(math.tensor(math.PI)), -1, abs_tolerance=1e-6, msg=backend.name) math.assert_close(math.cos(math.tensor(math.PI * 3 / 2)), 0, abs_tolerance=1e-6, msg=backend.name)
def test_hessian(self): def f(x, y): return math.l1_loss(x**2 * y), x, y eval_hessian = math.hessian(f, wrt=[ 0, ], get_output=True, get_gradient=True, dim_suffixes=('1', '2')) for backend in BACKENDS: if backend.supports(Backend.hessian): with backend: x = math.tensor([(0.01, 1, 2)], channel('vector', 'v')) y = math.tensor([1., 2.], batch('batch')) (L, x, y), g, H, = eval_hessian(x, y) math.assert_close(L, (5.0001, 10.0002), msg=backend.name) math.assert_close(g.batch[0].vector[0], (0.02, 2, 4), msg=backend.name) math.assert_close(g.batch[1].vector[0], (0.04, 4, 8), msg=backend.name) math.assert_close(2, H.v1[0].v2[0].batch[0], H.v1[1].v2[1].batch[0], H.v1[2].v2[2].batch[0], msg=backend.name) math.assert_close(4, H.v1[0].v2[0].batch[1], H.v1[1].v2[1].batch[1], H.v1[2].v2[2].batch[1], msg=backend.name)
def test_closest_grid_values_1d(self): grid = math.tensor([0, 1, 2, 3], names='x') coords = math.tensor([[0.1], [1.9], [0.5], [3.1]], names='x,vector') closest = math.closest_grid_values(grid, coords, extrapolation.ZERO) math.assert_close( closest, math.tensor([(0, 1), (1, 2), (0, 1), (3, 0)], names='x,closest_x'))
def test_cumulative_sum(self): t = math.tensor([(0, 1, 2, 3), (-1, -2, -3, -4)], spatial('y,x')) for backend in BACKENDS: with backend: t_ = math.tensor(t) x_ = math.cumulative_sum(t_, 'x') math.assert_close(x_, [(0, 1, 3, 6), (-1, -3, -6, -10)], msg=backend.name) y_ = math.cumulative_sum(t_, t.shape[0]) math.assert_close(y_, [(0, 1, 2, 3), (-1, -1, -1, -1)], msg=backend.name)
def test_np_speed_op2(self): print() np1, np2 = rnpv(64), rnpv(64) t1 = math.tensor(np1, batch('batch'), spatial('x, y'), channel('vector')) t2 = math.tensor(np2, batch('batch'), spatial('x, y'), channel('vector')) _assert_equally_fast(lambda: np1 + np2, lambda: t1 + t2, n=10000) np1, np2 = rnpv(256), rnpv(256) t1 = math.tensor(np1, batch('batch'), spatial('x, y'), channel('vector')) t2 = math.tensor(np2, batch('batch'), spatial('x, y'), channel('vector')) _assert_equally_fast(lambda: np1 + np2, lambda: t1 + t2, n=1000)
def test_tensor_from_constant(self): for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND): with backend: for const in (1, 1.5, True, 1+1j): tens = math.tensor(const, convert=False) self.assertEqual(math.NUMPY_BACKEND, math.choose_backend(tens)) math.assert_close(tens, const) tens = math.tensor(const) self.assertEqual(backend, math.choose_backend(tens), f'{const} was not converted to the specified backend') math.assert_close(tens, const)
def test_tensor_from_tuple_of_numbers(self): data_tuple = (1, 2, 3) for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND): with backend: tens = math.tensor(data_tuple, convert=False) self.assertEqual(math.NUMPY_BACKEND, math.choose_backend(tens)) math.assert_close(tens, data_tuple) tens = math.tensor(data_tuple) self.assertEqual(backend, math.choose_backend(tens)) math.assert_close(tens, data_tuple)
def test_tensor_from_tuple_of_tensor_like(self): native = ((1, 2, 3), math.zeros(vector=3)) for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND): with backend: tens = math.tensor(native, names=['stack', 'vector'], convert=False) self.assertEqual(math.NUMPY_BACKEND, math.choose_backend(tens)) self.assertEqual(shape(stack=2, vector=3), tens.shape) tens = math.tensor(native, names=['stack', 'vector']) self.assertEqual(backend, math.choose_backend(tens)) self.assertEqual(shape(stack=2, vector=3), tens.shape)
def test_tensor_from_native(self): for creation_backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND): native = creation_backend.ones((4,)) for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND): with backend: tens = math.tensor(native, convert=False) self.assertEqual(creation_backend, math.choose_backend(tens)) math.assert_close(tens, native) tens = math.tensor(native) self.assertEqual(backend, math.choose_backend(tens), f'Conversion failed from {creation_backend} to {backend}') math.assert_close(tens, native)
def test_boolean_mask_batched(self): for backend in BACKENDS: with backend: x = math.expand(math.range(spatial('x'), 4), batch(batch=2)) * math.tensor([1, -1]) mask = math.tensor([[True, False, True, False], [False, True, False, False]], batch('batch'), spatial('x')) selected = math.boolean_mask(x, 'x', mask) expected_0 = math.tensor([(0, -0), (2, -2)], spatial('x'), channel('vector')) expected_1 = math.tensor([(1, -1)], spatial('x'), channel('vector')) math.assert_close(selected.batch[0], expected_0, msg=backend.name) math.assert_close(selected.batch[1], expected_1, msg=backend.name) math.assert_close(selected, x.x[mask], msg=backend.name)
def test_extrapolate_valid(self): valid = tensor([[0, 0, 0], [0, 1, 1], [1, 0, 0]], 'x, y') values = tensor([[1, 0, 0], [0, 4, 0], [2, 0, 0]], 'x, y') expected_valid = tensor([[0, 1, 1], [1, 1, 1], [1, 1, 1]], 'x, y') expected_values = tensor([[1, 4, 0], [3, 4, 0], [2, 3, 0]], 'x, y') new_values, new_valid = math.extrapolate_valid_values(values, valid, 1) self.assertTrue(new_values == expected_values) self.assertTrue(new_valid == expected_valid)
def test_grid_sample_gradient_1d(self): for backend in BACKENDS: if backend.supports(Backend.gradients): with backend: grid = math.tensor([0., 1, 2, 3], spatial('x')) coords = math.tensor([0.5, 1.5], instance('points')) with math.record_gradients(grid, coords): sampled = math.grid_sample(grid, coords, extrapolation.ZERO) loss = math.mean(math.l2_loss(sampled)) / 2 grad_grid, grad_coords = math.gradients(loss, grid, coords) math.assert_close(grad_grid, math.tensor([0.125, 0.5, 0.375, 0], spatial('x')), msg=backend) math.assert_close(grad_coords, math.tensor([0.25, 0.75], instance('points')), msg=backend)
def test__periodic_2d_arakawa_poisson_bracket(self): """test _periodic_2d_arakawa_poisson_bracket implementation""" with math.precision(64): # Define parameters to test test_params = { 'grid_size': [(4, 4), (32, 32)], 'dx': [0.1, 1], 'gen_func': [ lambda grid_size: np.random.rand(*grid_size).reshape( grid_size) ] } # Generate test cases as the product test_cases = [ dict(zip(test_params, v)) for v in product(*test_params.values()) ] for params in test_cases: grid_size = params['grid_size'] d1 = params['gen_func'](grid_size) d2 = params['gen_func'](grid_size) dx = params['dx'] padding = extrapolation.PERIODIC ref = self.arakawa_reference_implementation( np.pad(d1.copy(), 1, mode='wrap'), np.pad(d2.copy(), 1, mode='wrap'), dx)[1:-1, 1:-1] d1_tensor = field.CenteredGrid( values=math.tensor(d1, names=['x', 'y']), bounds=geom.Box([0, 0], list(grid_size)), extrapolation=padding) d2_tensor = field.CenteredGrid( values=math.tensor(d2, names=['x', 'y']), bounds=geom.Box([0, 0], list(grid_size)), extrapolation=padding) val = math._nd._periodic_2d_arakawa_poisson_bracket( d1_tensor.values, d2_tensor.values, dx) try: math.assert_close(ref, val, rel_tolerance=1e-14, abs_tolerance=1e-14) except BaseException as e: abs_error = math.abs(val - ref) max_abs_error = math.max(abs_error) max_rel_error = math.max(math.abs(abs_error / ref)) variation_str = "\n".join([ f"max_absolute_error: {max_abs_error}", f"max_relative_error: {max_rel_error}", ]) print(ref) print(val) raise AssertionError(e, params, variation_str)
def _op2(self, other, operator) -> 'SampledField': if isinstance(other, Field): other_values = reduce_sample(other, self._elements) values = operator(self._values, other_values) extrapolation_ = operator(self._extrapolation, other.extrapolation) return self.with_values(values).with_extrapolation(extrapolation_) else: if isinstance(other, (tuple, list)) and len(other) == self.spatial_rank: other = math.tensor(other, self.points.shape['vector']) else: other = math.tensor(other) values = operator(self._values, other) return self.with_values(values)
def test_tensor_from_tensor(self): ref = math.batch_stack([math.zeros(x=5), math.zeros(x=4)], 'stack') for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND): with backend: tens = math.tensor(ref, convert=False) self.assertEqual(math.NUMPY_BACKEND, math.choose_backend(tens)) self.assertEqual(2, tens.shape.get_size('stack')) self.assertEqual(('stack', 'x'), tens.shape.names) tens = math.tensor(ref) self.assertEqual(backend, math.choose_backend(tens)) self.assertEqual(backend, math.choose_backend(tens.stack[0])) self.assertEqual(backend, math.choose_backend(tens.stack[1])) tens = math.tensor(ref, names=('n1', 'n2')) self.assertEqual(backend, math.choose_backend(tens))
def test_tensor_from_native(self): for creation_backend in BACKENDS: native = creation_backend.ones((4, )) for backend in BACKENDS: with backend: tens = math.tensor(native, convert=False) self.assertEqual(creation_backend, tens.default_backend) math.assert_close(tens, native) tens = math.tensor(native) self.assertEqual( backend, tens.default_backend, f'Conversion failed from {creation_backend} to {backend}' ) math.assert_close(tens, native)
def test_tensor_from_tensor(self): ref = math.stack([math.zeros(spatial(x=5)), math.zeros(spatial(x=4))], batch('stack')) for backend in BACKENDS: with backend: tens = math.tensor(ref, convert=False) self.assertEqual(math.NUMPY, math.choose_backend(tens)) self.assertEqual(2, tens.shape.get_size('stack')) self.assertEqual(('stack', 'x'), tens.shape.names) tens = math.tensor(ref) self.assertEqual(backend, math.choose_backend(tens)) self.assertEqual(backend, math.choose_backend(tens.stack[0])) self.assertEqual(backend, math.choose_backend(tens.stack[1])) tens = math.tensor(ref, batch('n1', 'n2')) self.assertEqual(backend, math.choose_backend(tens))
def test_slice_by_item_name(self): t = math.tensor(spatial(x=4, y=3)) math.assert_close(t.dims['x'], 4) math.assert_close(t.dims['y'], 3) math.assert_close(t.dims['y,x'], (3, 4)) math.assert_close(t.dims[('y', 'x')], (3, 4)) math.assert_close(t.dims[spatial('x,y')], (4, 3))
def test_boolean_mask_dim_missing(self): for backend in BACKENDS: with backend: x = math.random_uniform(spatial(x=2)) mask = math.tensor([True, False, True, True], batch('selection')) selected = math.boolean_mask(x, 'selection', mask) self.assertEqual(set(spatial(x=2) & batch(selection=3)), set(selected.shape), msg=backend.name)
def test_grid_sample(self): for backend in BACKENDS: with backend: grid = math.sum(math.meshgrid(x=[1, 2, 3], y=[0, 3]), 'vector') # 1 2 3 | 4 5 6 coords = math.tensor([(0, 0), (0.5, 0), (0, 0.5), (-2, -1)], instance('list'), channel('vector')) interp = math.grid_sample(grid, coords, extrapolation.ZERO) math.assert_close(interp, [1, 1.5, 2.5, 0], msg=backend.name)
def read_single_field(file: str, convert_to_backend=True) -> SampledField: stored = np.load(file, allow_pickle=True) ftype = stored['field_type'] implemented_types = ('CenteredGrid', 'StaggeredGrid') if ftype in implemented_types: data = stored['data'] dim_item_names = stored.get('dim_item_names', (None, ) * len(data.shape)) data = NativeTensor( data, math.Shape(data.shape, tuple(stored['dim_names']), tuple(stored['dim_types']), tuple(dim_item_names))) if convert_to_backend: data = math.tensor(data, convert=convert_to_backend) bounds_item_names = stored.get('bounds_item_names', (None, ) * len(stored['lower'] + stored['upper'])) lower = math.wrap( stored['lower'], math.channel(vector=tuple(bounds_item_names)) ) if stored['lower'].ndim > 0 else math.wrap(stored['lower']) upper = math.wrap(stored['upper'], math.channel(vector=tuple(bounds_item_names))) extrapolation = math.extrapolation.from_dict( stored['extrapolation'][()]) if ftype == 'CenteredGrid': return CenteredGrid(data, bounds=geom.Box(lower, upper), extrapolation=extrapolation) elif ftype == 'StaggeredGrid': data_ = unstack_staggered_tensor(data, extrapolation) return StaggeredGrid(data_, bounds=geom.Box(lower, upper), extrapolation=extrapolation) raise NotImplementedError(f"{ftype} not implemented ({implemented_types})")
def test_extrapolate_valid_3x3(self): valid = tensor([[0, 0, 0], [0, 0, 1], [1, 0, 0]], spatial('x, y')) values = tensor([[1, 0, 2], [0, 0, 4], [2, 0, 0]], spatial('x, y')) expected_valid = tensor([[0, 1, 1], [1, 1, 1], [1, 1, 1]], spatial('x, y')) expected_values = tensor([[1, 4, 4], [2, 3, 4], [2, 3, 4]], spatial('x, y')) extrapolated_values, extrapolated_valid = math.extrapolate_valid_values(values, valid) math.assert_close(extrapolated_values, expected_values) math.assert_close(extrapolated_valid, expected_valid)
def native_call(f, *inputs, channels_last=None, channel_dim='vector', extrapolation=None) -> SampledField or Tensor: """ Similar to `phi.math.native_call()`. Args: f: Function to be called on native tensors of `inputs.values`. The function output must have the same dimension layout as the inputs and the batch size must be identical. *inputs: `SampledField` or `phi.Tensor` instances. extrapolation: (Optional) Extrapolation of the output field. If `None`, uses the extrapolation of the first input field. Returns: `SampledField` matching the first `SampledField` in `inputs`. """ input_tensors = [ i.values if isinstance(i, SampledField) else tensor(i) for i in inputs ] values = math.native_call(f, *input_tensors, channels_last=channels_last, channel_dim=channel_dim) for i in inputs: if isinstance(i, SampledField): result = i.with_values(values=values) if extrapolation is not None: result = result.with_extrapolation(extrapolation) return result else: raise AssertionError("At least one input must be a SampledField.")