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_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 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_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_consistency_implicit(self):
DOMAIN = Domain(x=200, boundaries=PERIODIC)
DIFFUSIVITY = 0.5
grid = DOMAIN.scalar_grid((1,) * 100 + (0,) * 100)
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_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 test_implicit_stability(self):
DOMAIN = Domain(x=6, boundaries=PERIODIC)
DIFFUSIVITY = 10
grid = DOMAIN.scalar_grid((1,) * 3 + (0,) * 3)
try:
implicit = diffuse.implicit(grid, DIFFUSIVITY, 1, order=10)
print(implicit.values)
field.assert_close(0 <= implicit <= 1.0001, True)
except AssertionError as err:
print(err)
pass # solve did not converge
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_trace_function(self):
def f(x: StaggeredGrid, y: CenteredGrid):
return x + (y >> x)
ft = field.trace_function(f)
domain = Domain(x=4, y=3)
x = domain.staggered_grid(1)
y = domain.vector_grid(1)
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_batched_files(self):
DOMAIN = Domain(x=32, y=32, boundaries=CLOSED)
smoke = DOMAIN.scalar_grid(1) * math.random_uniform(count=2, config=3)
vel = DOMAIN.staggered_grid(2) * math.random_uniform(count=2, vel=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_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_write_read_batch_duplicate(self):
DOMAIN = Domain(x=32, y=32, boundaries=CLOSED)
smoke = DOMAIN.scalar_grid(1) * math.random_uniform(count=2)
vel = DOMAIN.staggered_grid(2) * math.random_uniform(count=2)
# write
scene = Scene.create(DIR, more=2)
scene.write({'smoke': smoke, 'vel': vel})
# read batch
smoke_ = scene.read('smoke')
vel_ = scene.read('vel')
self.assertEqual(4, smoke_.shape.batch.volume)
self.assertEqual(4, vel_.shape.batch.volume)
field.assert_close(smoke, smoke_)
field.assert_close(vel, vel_)
scene.remove()
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_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_advection(adv):
domain = Domain(x=4, y=3)
s = domain.scalar_grid(Noise())
v = domain.vector_grid(Noise(vector=2))
field.assert_close(s, adv(s, v, 0), adv(s, v * 0, 1))
sv = domain.staggered_grid(Noise())
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_equality_1d_periodic(self):
DOMAIN = Domain(x=200, boundaries=PERIODIC)
DIFFUSIVITY = 0.5
grid = DOMAIN.scalar_grid((1,) * 100 + (0,) * 100)
explicit = diffuse.explicit(grid, DIFFUSIVITY, 1, substeps=1000)
implicit = diffuse.implicit(grid, DIFFUSIVITY, 1, order=10)
fourier = diffuse.fourier(grid, DIFFUSIVITY, 1)
field.assert_close(explicit, implicit, rel_tolerance=0, abs_tolerance=0.01)
field.assert_close(explicit, implicit, fourier, rel_tolerance=0, abs_tolerance=0.1)
# print(f"{explicit.values[:6]} Explicit")
# print(f"{implicit.values[:6]} Implicit")
# print(f"{fourier.values[:6]} Fourier")
# print()
back_explicit = diffuse.explicit(explicit, DIFFUSIVITY, -1, substeps=1000)
back_implicit = diffuse.implicit(implicit, DIFFUSIVITY, -1, order=10)
back_fourier = diffuse.fourier(fourier, DIFFUSIVITY, -1)
# print(f"{back_explicit.values[:6]} Explicit")
# print(f"{back_implicit.values[:6]} Implicit")
# print(f"{back_fourier.values[:6]} Fourier")
field.assert_close(grid, back_explicit, back_implicit, back_fourier, rel_tolerance=0, abs_tolerance=0.1)
def test_equality_1d_periodic(self):
DIFFUSIVITY = 0.5
grid = CenteredGrid((1, ) * 100 + (0, ) * 100,
extrapolation.PERIODIC,
x=200)
explicit = diffuse.explicit(grid, DIFFUSIVITY, 1, substeps=1000)
implicit = diffuse.implicit(grid, DIFFUSIVITY, 1, order=10)
fourier = diffuse.fourier(grid, DIFFUSIVITY, 1)
field.assert_close(explicit,
implicit,
rel_tolerance=0,
abs_tolerance=0.01)
field.assert_close(explicit,
implicit,
fourier,
rel_tolerance=0,
abs_tolerance=0.1)
# print(f"{explicit.values[:6]} Explicit")
# print(f"{implicit.values[:6]} Implicit")
# print(f"{fourier.values[:6]} Fourier")
# print()
back_explicit = diffuse.explicit(explicit,
DIFFUSIVITY,
-1,
substeps=1000)
back_implicit = diffuse.implicit(implicit, DIFFUSIVITY, -1, order=10)
back_fourier = diffuse.fourier(fourier, DIFFUSIVITY, -1)
# print(f"{back_explicit.values[:6]} Explicit")
# print(f"{back_implicit.values[:6]} Implicit")
# print(f"{back_fourier.values[:6]} Fourier")
field.assert_close(grid,
back_explicit,
back_implicit,
back_fourier,
rel_tolerance=0,
abs_tolerance=0.1)
def test_round_(self):
grid = StaggeredGrid(1.7, x=4, y=3)
field.assert_close(field.round(grid), 2)
def test_isfinite(self):
grid = StaggeredGrid(1, x=4, y=3)
field.assert_close(field.isfinite(grid), True)
def test_write_read(self):
DOMAIN = Domain(x=32, y=32, boundaries=CLOSED)
smoke = DOMAIN.scalar_grid(1)
vel = DOMAIN.staggered_grid(2)
# write
scene = Scene.create(DIR)
scene.write(smoke=smoke, vel=vel)
self.assertEqual(1, len(scene.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_real(self):
grid = StaggeredGrid(1, x=4, y=3)
field.assert_close(field.real(grid), grid)
def test_imag(self):
grid = StaggeredGrid(1, x=4, y=3)
field.assert_close(field.imag(grid), 0)
def test_sqrt(self):
grid = StaggeredGrid(2, x=4, y=3)
field.assert_close(field.sqrt(grid), numpy.sqrt(2))
def test_sign(self):
grid = StaggeredGrid(0.5, x=4, y=3)
field.assert_close(field.sign(grid), 1)
def test_abs(self):
grid = StaggeredGrid(-1, x=4, y=3)
field.assert_close(field.abs(grid), abs(grid), 1)
def test_exp(self):
grid = Domain(x=4, y=3).staggered_grid(0)
field.assert_close(field.exp(grid), 1)
def test_cos(self):
grid = StaggeredGrid(0, x=4, y=3)
field.assert_close(field.cos(grid), 1)
def test_ceil(self):
grid = StaggeredGrid(1.1, x=4, y=3)
field.assert_close(field.ceil(grid), 2)
def test_floor(self):
grid = StaggeredGrid(2.8, x=4, y=3)
field.assert_close(field.floor(grid), 2)
