def test_staggered_grid_with_extrapolation(self): grid = StaggeredGrid(Noise(vector='x,y'), extrapolation.BOUNDARY, x=20, y=10) grid_0 = grid.with_extrapolation(extrapolation.ZERO) self.assertEqual(grid.resolution, grid_0.resolution) grid_ = grid_0.with_extrapolation(extrapolation.BOUNDARY) self.assertEqual(grid.resolution, grid_.resolution) math.assert_close(grid_.values.vector['x'].x[0], 0) math.assert_close(grid_.values.vector['x'].x[-1], 0) math.assert_close(grid_.values.vector['y'].y[0], 0) math.assert_close(grid_.values.vector['y'].y[-1], 0)
def test_slice_staggered_grid_along_batch(self): v = StaggeredGrid(Noise(batch(batch=10)), x=10, y=20) b1 = v[{'batch': 1}] b2 = v.batch[1] b3 = field.unstack(v, 'batch')[1] self.assertIsInstance(b1, StaggeredGrid) field.assert_close(b1, b2, b3)
def grid(self, value: Field or Tensor or Number or Geometry or callable = 0., type: type = CenteredGrid, extrapolation: math.Extrapolation = None) -> CenteredGrid or StaggeredGrid: """ *Deprecated* due to inconsistent extrapolation selection. Use `scalar_grid()` or `vector_grid()` instead. Creates a grid matching the resolution and bounds of the domain. The grid is created from the given `value` which must be one of the following: * Number (int, float, complex or zero-dimensional tensor): all grid values will be equal to `value`. This has a near-zero memory footprint. * Field: the given value is resampled to the grid cells of this Domain. * Tensor with spatial dimensions matching the domain resolution: grid values will equal `value`. * Geometry: grid values are determined from the volume overlap between grid cells and geometry. Non-overlapping = 0, fully enclosed grid cell = 1. * function(location: Tensor) returning one of the above. Args: value: constant, Field, Tensor or function specifying the grid values type: type of Grid to create, must be either CenteredGrid or StaggeredGrid extrapolation: (optional) grid extrapolation, defaults to Domain.boundaries['scalar_extrapolation'] Returns: Grid of specified type """ warnings.warn("Domain.grid is deprecated. Use scalar_grid or vector_grid instead.", DeprecationWarning) extrapolation = extrapolation or self.boundaries['scalar_extrapolation'] if type is CenteredGrid: return CenteredGrid.sample(value, self.resolution, self.bounds, extrapolation) elif type is StaggeredGrid: return StaggeredGrid.sample(value, self.resolution, self.bounds, extrapolation) else: raise ValueError('Unknown grid type: %s' % type)
def _test_advection(adv): s = CenteredGrid(Noise(), x=4, y=3) v = CenteredGrid(Noise(vector=2), x=4, y=3) field.assert_close(s, adv(s, v, 0), adv(s, v * 0, 1)) sv = StaggeredGrid(Noise(), x=4, y=3) field.assert_close(s, adv(s, sv, 0), adv(s, sv * 0, 1)) field.assert_close(sv, adv(sv, sv, 0), adv(sv, sv * 0, 1))
def test_slice_staggered_grid_along_vector(self): v = StaggeredGrid(Noise(batch(batch=10)), x=10, y=20) x1 = v[{'vector': 0}] x2 = v.vector[0] x3 = v.vector['x'] x4 = field.unstack(v, 'vector')[0] self.assertIsInstance(x1, CenteredGrid) field.assert_close(x1, x2, x3, x4)
def test_grid_constant_extrapolation(self): grid = CenteredGrid(math.random_uniform(spatial(x=50, y=10)), 0., Box[0:1, 0:1]) self.assertEqual(grid.extrapolation, extrapolation.ZERO) grid = CenteredGrid(0, 0, Box[0:1, 0:1], x=50, y=10) self.assertEqual(grid.extrapolation, extrapolation.ZERO) grid = StaggeredGrid(0, 0, Box[0:1, 0:1], x=50, y=10) self.assertEqual(grid.extrapolation, extrapolation.ZERO)
def test_runge_kutta_4(self): domain = Domain(x=4, y=3) points = domain.distribute_points(domain.bounds, points_per_cell=2) v = CenteredGrid(Noise(vector=2), x=4, y=3) field.assert_close(points, advect.runge_kutta_4(points, v, 0), advect.runge_kutta_4(points, v * 0, 0)) sv = StaggeredGrid(Noise(), x=4, y=3) field.assert_close(points, advect.runge_kutta_4(points, sv, 0), advect.runge_kutta_4(points, sv * 0, 0))
def test_make_incompressible_gradients_equal_tf_torch(self): velocity0 = StaggeredGrid(Noise(), ZERO, x=16, y=16, bounds=Box[0:100, 0:100]) grads = [] for backend in BACKENDS: if backend.supports(Backend.record_gradients): with backend: velocity = param = velocity0.with_values( math.tensor(velocity0.values)) with math.record_gradients(param.values): velocity, _ = fluid.make_incompressible(velocity) loss = field.l2_loss(velocity) assert math.isfinite(loss).all grad = math.gradients(loss, param.values) assert math.isfinite(grad).all grads.append(grad) math.assert_close(*grads, abs_tolerance=1e-5)
def test_diffuse_staggered_batched(self): for diffusivity in [1, 0.5, math.wrap([1., 0.5], batch('batch'))]: grid = StaggeredGrid(Noise(batch(batch=2), vector=2), extrapolation.PERIODIC, x=4, y=3) diffuse.explicit(grid, diffusivity, 1, substeps=10) diffuse.implicit(grid, diffusivity, 1, order=2) diffuse.fourier(grid, diffusivity, 1) grid = StaggeredGrid(Noise(batch(batch=2), vector=2), extrapolation.ZERO, x=4, y=3) diffuse.explicit(grid, diffusivity, 1, substeps=10) # diffuse.implicit(grid, diffusivity, 1, order=2) # not yet supported grid = StaggeredGrid(Noise(batch(batch=2), vector=2), extrapolation.BOUNDARY, x=4, y=3) diffuse.explicit(grid, diffusivity, 1, substeps=10)
def test_trace_function(self): def f(x: StaggeredGrid, y: CenteredGrid): return x + (y @ x) ft = field.jit_compile(f) x = StaggeredGrid(1, x=4, y=3) y = CenteredGrid((1, 1), x=4, y=3) res_f = f(x, y) res_ft = ft(x, y) self.assertEqual(res_f.shape, res_ft.shape) field.assert_close(res_f, res_ft)
def test_write_read_batch_matching(self): smoke = CenteredGrid(1, extrapolation.BOUNDARY, x=32, y=32) * math.random_uniform(batch(count=2)) vel = StaggeredGrid(2, 0, x=32, y=32) * math.random_uniform( batch(count=2)) # write scene = Scene.create(DIR, count=2) scene.write({'smoke': smoke, 'vel': vel}) # read batch smoke_ = scene.read('smoke') vel_ = scene.read('vel') field.assert_close(smoke, smoke_) field.assert_close(vel, vel_) scene.remove()
def test_staggered_grid_sizes_by_extrapolation(self): s = spatial(x=20, y=10) for initializer in [0, Noise(vector=2), (0, 1), Sphere((0, 0), 1)]: g_const = StaggeredGrid(initializer, extrapolation.ZERO, resolution=s) self.assertEqual(g_const.shape, spatial(x=20, y=10) & channel(vector=2)) self.assertEqual(g_const.values.vector[0].shape, spatial(x=19, y=10)) g_periodic = StaggeredGrid(initializer, extrapolation.PERIODIC, resolution=s) self.assertEqual(g_periodic.shape, spatial(x=20, y=10) & channel(vector=2)) self.assertEqual(g_periodic.values.vector[0].shape, spatial(x=20, y=10)) g_boundary = StaggeredGrid(initializer, extrapolation.BOUNDARY, resolution=s) self.assertEqual(g_boundary.shape, spatial(x=20, y=10) & channel(vector=2)) self.assertEqual(g_boundary.values.vector[0].shape, spatial(x=21, y=10))
def vector_grid( self, value: Field or Tensor or Number or Geometry or callable = 0., type: type = CenteredGrid, extrapolation: math.Extrapolation = None ) -> CenteredGrid or StaggeredGrid: """ Creates a vector grid matching the resolution and bounds of the domain. The grid is created from the given `value` which must be one of the following: * Number (int, float, complex or zero-dimensional tensor): all grid values will be equal to `value`. This has a near-zero memory footprint. * Field: the given value is resampled to the grid cells of this Domain. * Tensor with spatial dimensions matcing the domain resolution: grid values will equal `value`. * Geometry: grid values are determined from the volume overlap between grid cells and geometry. Non-overlapping = 0, fully enclosed grid cell = 1. * function(location: Tensor) returning one of the above. The returned grid will have a vector dimension with size equal to the rank of the domain. Args: value: constant, Field, Tensor or function specifying the grid values type: class of Grid to create, must be either CenteredGrid or StaggeredGrid extrapolation: (optional) grid extrapolation, defaults to Domain.boundaries['vector_extrapolation'] Returns: Grid of specified type """ extrapolation = extrapolation or self.boundaries['vector_extrapolation'] if type is CenteredGrid: grid = CenteredGrid.sample(value, self.resolution, self.bounds, extrapolation) if grid.shape.channel.rank == 0: grid = grid.with_(values=math.expand_channel( grid.values, 'vector', dim_size=self.rank)) else: assert grid.shape.channel.sizes[0] == self.rank return grid elif type is StaggeredGrid: return StaggeredGrid.sample(value, self.resolution, self.bounds, extrapolation) else: raise ValueError('Unknown grid type: %s' % type)
def test_gradient_function(self): def f(x: StaggeredGrid, y: CenteredGrid): pred = x + (y @ x) loss = field.l2_loss(pred) return loss grad = field.functional_gradient(f, get_output=False) fgrad = field.functional_gradient(f, (0, 1), get_output=True) x = StaggeredGrid(1, x=4, y=3) y = CenteredGrid((1, 1), x=4, y=3) for backend in BACKENDS: if backend.supports(Backend.gradients): with backend: dx, = grad(x, y) self.assertIsInstance(dx, StaggeredGrid) loss, (dx, dy) = fgrad(x, y) self.assertIsInstance(loss, math.Tensor) self.assertIsInstance(dx, StaggeredGrid) self.assertIsInstance(dy, CenteredGrid)
def test_write_read(self): smoke = CenteredGrid(1, extrapolation.BOUNDARY, x=32, y=32) vel = StaggeredGrid(2, 0, x=32, y=32) # write scene = Scene.create(DIR) scene.write(smoke=smoke, vel=vel) self.assertEqual(1, len(scene.frames)) self.assertEqual(1, len(scene.complete_frames)) self.assertEqual(2, len(scene.fieldnames)) # read single smoke_ = scene.read('smoke') vel_ = scene.read('vel') field.assert_close(smoke, smoke_) field.assert_close(vel, vel_) self.assertEqual(smoke.extrapolation, smoke_.extrapolation) self.assertEqual(vel.extrapolation, vel_.extrapolation) # read multiple smoke__, vel__ = scene.read(['smoke', 'vel']) # deprecated field.assert_close(smoke, smoke__) field.assert_close(vel, vel__) smoke__, vel__ = scene.read('smoke', 'vel') field.assert_close(smoke, smoke__) field.assert_close(vel, vel__) scene.remove()
def test_isfinite(self): grid = StaggeredGrid(1, x=4, y=3) field.assert_close(field.isfinite(grid), True)
def test_abs(self): grid = StaggeredGrid(-1, x=4, y=3) field.assert_close(field.abs(grid), abs(grid), 1)
def test_imag(self): grid = StaggeredGrid(1, x=4, y=3) field.assert_close(field.imag(grid), 0)
def test_real(self): grid = StaggeredGrid(1, x=4, y=3) field.assert_close(field.real(grid), grid)
def test_zero_staggered_grid(self): for data in [(0, 0), 0, Noise(), lambda x: x]: grid = StaggeredGrid(data, 0, x=4, y=3) self.assertEqual(('x', 'y'), grid.values.vector.item_names) self.assertEqual(('x', 'y'), grid.dx.vector.item_names)
def test_iter_dim(self): slices = tuple(StaggeredGrid(0, x=4, y=3).vector) self.assertEqual(2, len(slices)) self.assertFalse(slices[0].shape.non_spatial) self.assertEqual(('x', 'y'), slices[0].bounds.size.vector.item_names)
def test_sample_at(self): field = AngularVelocity([0, 0]) self.assertEqual(channel(vector=2), field.shape.channel) field @ CenteredGrid(0, x=4, y=3) field @ StaggeredGrid(0, x=4, y=3)
def test_create_grid_int_resolution(self): g = CenteredGrid(0, 0, Box(x=4, y=3), resolution=10) self.assertEqual(g.shape, spatial(x=10, y=10)) g = StaggeredGrid(0, 0, Box(x=4, y=3), resolution=10) self.assertEqual(spatial(g), spatial(x=10, y=10))
def test_cos(self): grid = StaggeredGrid(0, x=4, y=3) field.assert_close(field.cos(grid), 1)
def test_curl_2d_staggered_to_centered(self): velocity = StaggeredGrid((2, 0), extrapolation.BOUNDARY, x=2, y=2) curl = field.curl(velocity) self.assertEqual(spatial(x=3, y=3), curl.resolution)
def test_ceil(self): grid = StaggeredGrid(1.1, x=4, y=3) field.assert_close(field.ceil(grid), 2)
def test_round_(self): grid = StaggeredGrid(1.7, x=4, y=3) field.assert_close(field.round(grid), 2)
def test_downsample_staggered_2d(self): grid = StaggeredGrid(1, extrapolation.BOUNDARY, x=32, y=40) downsampled = field.downsample2x(grid) self.assertEqual(set(spatial(x=16, y=20) & channel(vector=2)), set(downsampled.shape))
def test_infinite_cylinder_to_grid(self): cylinder = geom.infinite_cylinder(x=2, y=1.5, radius=.8, inf_dim='z') StaggeredGrid(cylinder, 0, x=4, y=3, z=2)
def test_sign(self): grid = StaggeredGrid(0.5, x=4, y=3) field.assert_close(field.sign(grid), 1)