def test_linear_solve_matrix_tape(self): y = CenteredGrid(1, extrapolation.ZERO, x=3) * (1, 2) x0 = CenteredGrid(0, extrapolation.ZERO, x=3) for method in ['CG', 'CG-adaptive', 'auto']: solve = math.Solve(method, 0, 1e-3, x0=x0, max_iterations=100) with math.SolveTape() as solves: x = field.solve_linear(math.jit_compile_linear(field.laplace), y, solve) math.assert_close(x.values, [[-1.5, -2, -1.5], [-3, -4, -3]], abs_tolerance=1e-3) assert len(solves) == 1 assert solves[0] == solves[solve] math.assert_close(solves[solve].residual.values, 0, abs_tolerance=1e-3) assert math.close(solves[solve].iterations, 2) or math.close( solves[solve].iterations, -1) with math.SolveTape(record_trajectories=True) as solves: x = field.solve_linear(math.jit_compile_linear(field.laplace), y, solve) math.assert_close(x.values, [[-1.5, -2, -1.5], [-3, -4, -3]], abs_tolerance=1e-3) assert solves[solve].x.trajectory.size == 3 math.assert_close(solves[solve].residual.trajectory[-1].values, 0, abs_tolerance=1e-3)
def simulate(centers): world = World() fluid = world.add(Fluid(Domain(x=5, y=4, 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]
def _sample(self, geometry: Geometry) -> Tensor: if geometry == self.bounds: return math.mean(self._values, self._resolution) if isinstance(geometry, GeometryStack): sampled = [self.sample(g) for g in geometry.geometries] return math.stack(sampled, geometry.stack_dim) if isinstance(geometry, GridCell): if self.elements == geometry: return self.values elif math.close(self.dx, geometry.size): fast_resampled = self._shift_resample(geometry.resolution, geometry.bounds) if fast_resampled is not NotImplemented: return fast_resampled points = geometry.center local_points = self.box.global_to_local(points) * self.resolution - 0.5 return math.grid_sample(self.values, local_points, self.extrapolation)
def sample_in(self, geometry: Geometry, reduce_channels=()) -> Tensor: if reduce_channels: assert len(reduce_channels) == 1 geometries = geometry.unstack(reduce_channels[0]) components = self.vector.unstack(len(geometries)) sampled = [c.sample_in(g) for c, g in zip(components, geometries)] return math.channel_stack(sampled, 'vector') if isinstance(geometry, GeometryStack): sampled = [self.sample_in(g) for g in geometry.geometries] return math.batch_stack(sampled, geometry.stack_dim_name) if isinstance(geometry, GridCell): if self.elements == geometry: return self.values elif math.close(self.dx, geometry.size): fast_resampled = self._shift_resample(geometry.resolution, geometry.bounds) if fast_resampled is not NotImplemented: return fast_resampled return self.sample_at(geometry.center, reduce_channels)
def global_to_local(self, global_position: Tensor) -> Tensor: if math.close(self.lower, 0): return global_position / self.size else: return (global_position - self.lower) / self.size
def __eq__(self, other): return isinstance( other, AbstractBox ) and self._lower.shape == other.lower.shape and math.close( self._lower, other.lower)
def test_close_different_shapes(self): a = math.ones(channel(vector='x,y')) b = math.wrap(3) self.assertFalse(math.close(a, b)) self.assertFalse(math.close(cached(a), b)) math.assert_close(a + 2, b)
def __eq__(self, other): return isinstance(other, BaseBox)\ and set(self.shape) == set(other.shape)\ and math.close(self._center, other.center)\ and math.close(self._half_size, other.half_size)
def __eq__(self, other): return isinstance(other, BaseBox)\ and set(self.shape) == set(other.shape)\ and self.size.shape.get_size('vector') == other.size.shape.get_size('vector')\ and math.close(self._lower, other.lower)\ and math.close(self._upper, other.upper)