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_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 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_plot_multiple(self):
grid = CenteredGrid(Noise(batch(b=2)), 0, Box[0:1, 0:1], x=50, y=10)
grid2 = CenteredGrid(grid, 0, Box[0:2, 0:1], x=20, y=50)
points = wrap([(.2, .4), (.9, .8)], instance('points'),
channel('vector'))
cloud = PointCloud(Sphere(points, radius=0.1), bounds=Box(0, [1, 1]))
titles = math.wrap([['b=0', 'b=0', 'points'], ['b=1', 'b=1', '']],
spatial('rows,cols'))
self._test_plot(grid, grid2, cloud, row_dims='b', title=titles)
def test_linear_solve_matrix_batched(
self): # TODO also test batched matrix
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)
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)
def test_solve_linear_matrix_dirichlet(self):
for backend in BACKENDS:
with backend:
y = CenteredGrid(1, extrapolation.ONE, x=3)
x0 = CenteredGrid(0, extrapolation.ONE, x=3)
solve = math.Solve('CG', 0, 1e-3, x0=x0, max_iterations=100)
x_ref = field.solve_linear(field.laplace, y, solve)
x_jit = field.solve_linear(
math.jit_compile_linear(field.laplace), y, solve)
math.assert_close(x_ref.values,
x_jit.values, [-0.5, -1, -0.5],
abs_tolerance=1e-3,
msg=backend)
def test_plot_vector_3d_batched(self):
sphere = CenteredGrid(SoftGeometryMask(
Sphere(x=.5, y=.5, z=.5, radius=.4)),
x=10,
y=10,
z=10,
bounds=Box(x=1, y=1, z=1)) * (.1, 0, 0)
cylinder = CenteredGrid(geom.infinite_cylinder(
x=16, y=16, inf_dim='z', radius=10),
x=32,
y=32,
z=32) * (0, 0, .1)
self._test_plot(sphere, cylinder)
def test_plot_scalar_2d_batch(self):
self._test_plot(
CenteredGrid(Noise(batch(b=2)),
0,
x=64,
y=8,
bounds=Box(0, [1, 1])))
self._test_plot(CenteredGrid(Noise(batch(b=2)),
0,
x=64,
y=8,
bounds=Box(0, [1, 1])),
row_dims='b',
size=(2, 4))
def test_integrate_all(self):
grid = CenteredGrid(field.Noise(vector=2),
extrapolation.ZERO,
x=10,
y=10,
bounds=Box[0:10, 0:10])
math.assert_close(field.integrate(grid, grid.bounds),
math.sum(grid.values, 'x,y'))
grid = CenteredGrid(field.Noise(vector=2),
extrapolation.ZERO,
x=10,
y=10,
bounds=Box[0:1, 0:1])
math.assert_close(field.integrate(grid, grid.bounds),
math.sum(grid.values, 'x,y') / 100)
def test_constant_diffusion(self):
grid = CenteredGrid(1, extrapolation.PERIODIC, x=4, y=3)
explicit = diffuse.explicit(grid, 1, 1, substeps=10)
implicit = diffuse.implicit(grid, 1, 1, order=2)
fourier = diffuse.fourier(grid, 1, 1)
math.assert_close(grid.values, explicit.values, implicit.values,
fourier.values)
def test_plot_vector_2d_batch(self):
self._test_plot(
CenteredGrid(Noise(batch(b=2), vector=2),
extrapolation.ZERO,
bounds=Box[0:1, 0:1],
x=10,
y=10))
def test_plot_vector_grid_2d(self):
self._test_plot(
CenteredGrid(Noise(vector=2),
extrapolation.ZERO,
x=64,
y=8,
bounds=Box(0, [1, 1])) * 0.1)
def test_solve_linear_matrix(self):
for backend in BACKENDS:
with backend:
y = CenteredGrid(1, extrapolation.ZERO, x=3)
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)
x = field.solve_linear(
math.jit_compile_linear(field.laplace), y, solve)
math.assert_close(x.values, [-1.5, -2, -1.5],
abs_tolerance=1e-3,
msg=backend)
def dash_graph_plot(data, settings: dict) -> dict:
if data is None:
return EMPTY_FIGURE
try:
if isinstance(data, np.ndarray):
data = math.wrap(data)
if isinstance(data, math.Tensor):
data = CenteredGrid(data, Box(0, math.wrap(data.shape, 'vector')))
if isinstance(data, (CenteredGrid, StaggeredGrid)):
component = settings.get('component', 'x')
if data.spatial_rank == 1:
return plot(data, settings)
if data.spatial_rank == 2:
if component == 'vec2' and data.shape.channel.volume >= 2:
return vector_field(data, settings)
else:
return heatmap(data, settings)
if data.spatial_rank == 3:
if component == 'vec2' and data.shape.channel.volum >= 2:
return vector_field(slice_2d(data, settings), settings)
else:
return heatmap(slice_2d(data, settings), settings)
if isinstance(data, PointCloud):
return cloud_plot(data, settings)
warnings.warn('No figure recipe for data %s' % data)
except BaseException as err:
print(f"Error during plotting: {err}")
return EMPTY_FIGURE
def test_zero_vector_grid(self):
for data in [(0, 0),
Noise(vector='x,y'),
Noise(vector=2), lambda x: x]:
grid = CenteredGrid(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 heatmap(grid: Grid, cols: int, rows: int, title: str):
inner_cols = cols - 10
inner_rows = rows - 2
grid @= CenteredGrid(0,
extrapolation.ZERO,
x=inner_cols,
y=inner_rows,
bounds=grid.bounds)
data = grid.values.numpy('y,x')
min_, max_ = numpy.min(data), numpy.max(data)
col_indices = (data - min_) / (max_ - min_) * len(FILLED)
col_indices = numpy.clip(
numpy.round(col_indices).astype(numpy.int8), 0,
len(FILLED) - 1)
lines = []
# lines.append(" " + "_" * inner_cols + " ")
title = title[:inner_cols]
padded_title = " " * ((inner_cols - len(title)) // 2) + title + " " * (
(inner_cols - len(title) + 1) // 2)
lines.append(" " + underline(padded_title) + "\033[0m ")
for index_row in col_indices[::-1]:
line = [FILLED[col_index] for col_index in index_row]
lines.append(" |" + "".join(line) + "|")
lines[-1] = lines[-1][:3] + underline(
lines[-1][3:inner_cols + 3]) + lines[-1][inner_cols + 3:]
return lines
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_slice_centered_grid(self):
g = CenteredGrid(Noise(batch(batch=10), channel(vector=2)), x=10, y=20)
s1 = g[{'vector': 0, 'batch': 1, 'x': 1}]
s2 = g.vector[0].batch[1].x[1]
self.assertIsInstance(s1, CenteredGrid)
self.assertEqual(s1.bounds, Box[1:2, 0:20])
field.assert_close(s1, s2)
def _test_make_incompressible(self, grid_type: type,
extrapolation: math.Extrapolation,
**batch_dims):
result = None
for i, backend in enumerate(BACKENDS):
with backend:
smoke = CenteredGrid(Sphere(
center=(math.random_uniform(batch(**batch_dims)) * 100,
10),
radius=5),
extrapolation,
x=16,
y=20,
bounds=Box[0:100, 0:100])
velocity = grid_type(0,
extrapolation,
x=16,
y=20,
bounds=Box[0:100, 0:100])
for _ in range(2):
velocity += smoke * (0, 0.1) @ velocity
velocity, _ = fluid.make_incompressible(velocity)
math.assert_close(divergence(velocity).values,
0,
abs_tolerance=2e-5)
if result is None:
result = velocity
else:
field.assert_close(
result,
abs_tolerance=1e-5,
msg=
f"Simulation with {backend} does not match {BACKENDS[:i]}"
)
def test_diffuse_centered_batched(self):
grid = CenteredGrid(Noise(batch=2, vector=2),
extrapolation.PERIODIC,
x=4,
y=3)
diffuse.explicit(grid, 1, 1, substeps=10)
diffuse.implicit(grid, 1, 1, order=2)
diffuse.fourier(grid, 1, 1)
def test_consistency_implicit(self):
DIFFUSIVITY = 0.5
grid = CenteredGrid((1, ) * 100 + (0, ) * 100,
extrapolation.PERIODIC,
x=200)
for extrap in (extrapolation.ZERO, extrapolation.BOUNDARY,
extrapolation.PERIODIC):
grid = grid.with_extrapolation(extrap)
implicit = diffuse.implicit(grid, DIFFUSIVITY, 1, order=10)
back_implicit = diffuse.implicit(implicit,
DIFFUSIVITY,
-1,
order=10)
field.assert_close(grid,
back_implicit,
rel_tolerance=0,
abs_tolerance=0.1)
def test_plot_multi_1d(self):
self._test_plot(
CenteredGrid(
lambda x: math.stack({
'sin': math.sin(x),
'cos': math.cos(x)
}, channel('curves')),
x=100,
bounds=Box(x=2 * math.pi)))
def test_linear_solve_matrix_jit(self):
@math.jit_compile
def solve(y, method):
solve = math.Solve(method, 0, 1e-3, x0=x0, max_iterations=100)
return field.solve_linear(math.jit_compile_linear(field.laplace),
y, solve)
for backend in BACKENDS:
with backend:
x0 = CenteredGrid(0, extrapolation.ZERO, x=3)
for method in ['CG']:
x = solve(CenteredGrid(0, extrapolation.ZERO, x=3),
method=method)
math.assert_close(x.values, 0, abs_tolerance=1e-3)
x = solve(CenteredGrid(1, extrapolation.ZERO, x=3),
method=method)
math.assert_close(x.values, [-1.5, -2, -1.5],
abs_tolerance=1e-3)
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_implicit_stability(self):
DIFFUSIVITY = 10
grid = CenteredGrid((1, ) * 3 + (0, ) * 3, extrapolation.PERIODIC, x=6)
try:
implicit = diffuse.implicit(grid, DIFFUSIVITY, 1, order=10)
print(implicit.values)
field.assert_close(0 <= implicit <= 1.0001, True)
except NotConverged as err:
print(err)
pass # solve_linear did not converge
def slice_2d(field3d, settings):
if isinstance(field3d, np.ndarray):
field3d = CenteredGrid(field3d)
if isinstance(field3d, StaggeredGrid):
component = settings.get('component', 'length')
if component in ('z', 'y', 'x'):
field3d = field3d.unstack()[{'z': physics_config.z, 'y': physics_config.y, 'x': physics_config.x}[component] % 3]
else:
field3d = field3d.at_centers()
assert isinstance(field3d, CenteredGrid) and field3d.spatial_rank == 3
depth = settings.get('depth', 0)
projection = settings.get('projection', FRONT)
removed_axis = {FRONT: physics_config.y, RIGHT: physics_config.x, TOP: physics_config.z}[projection] % 3
data = field3d.values[(slice(None),) + tuple([min(depth, field3d.resolution[i]) if i == removed_axis else slice(None) for i in range(3)]) + (slice(None),)]
if projection == RIGHT and not physics_config.is_x_first:
data = np.transpose(data, axes=(0, 2, 1, 3))
return CenteredGrid(data, box=field3d.box.without_axis(removed_axis))
def test_plot_point_cloud_2d_large(self):
spheres = PointCloud(
Sphere(wrap([(2, 4), (9, 8), (7, 8)], instance('points'),
channel('vector')),
radius=1))
cells = PointCloud(
geom.pack_dims(
CenteredGrid(0, 0, x=3, y=3, bounds=Box[4:6, 2:4]).elements,
'x,y', instance('points')))
cloud = field.stack([spheres, cells], instance('stack'))
self._test_plot(cloud)
def test_overlay(self):
grid = CenteredGrid(Noise(),
extrapolation.ZERO,
x=64,
y=8,
bounds=Box(0, [1, 1]))
points = wrap([(.2, .4), (.9, .8)], instance('points'),
channel('vector'))
cloud = PointCloud(Sphere(points, radius=.1))
self._test_plot(overlay(grid, grid * (0.1, 0.02), cloud),
title='Overlay')
def test_plot_point_cloud_2d(self):
spheres = PointCloud(
Sphere(wrap([(.2, .4), (.9, .8), (.7, .8)], instance('points'),
channel('vector')),
radius=.1))
cells = PointCloud(
geom.pack_dims(
CenteredGrid(0, 0, x=3, y=3,
bounds=Box[.4:.6, .2:.4]).elements, 'x,y',
instance('points')))
cloud = field.stack([spheres, cells], instance('stack'))
self._test_plot(cloud)
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)