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)
