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(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_constant_diffusion(self):
DOMAIN = Domain(x=4, y=3, boundaries=PERIODIC)
grid = DOMAIN.scalar_grid(1)
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_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_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_make_incompressible_batched(self, grid_type):
DOMAIN = Domain(x=16, y=16, boundaries=CLOSED, bounds=Box[0:100, 0:100])
smoke = DOMAIN.scalar_grid(Sphere(center=(math.random_uniform(batch=2) * 100, 10), radius=5))
velocity = DOMAIN.vector_grid(0, grid_type)
for _ in range(2):
velocity += smoke * (0, 0.1) >> velocity
velocity, pressure, _, _ = fluid.make_incompressible(velocity, DOMAIN)
math.assert_close(divergence(velocity).values, 0, abs_tolerance=2e-5)
return velocity.values
def test_runge_kutta_4(self):
domain = Domain(x=4, y=3)
points = domain.distribute_points(domain.bounds, points_per_cell=2)
v = domain.vector_grid(Noise(vector=2))
field.assert_close(points, advect.runge_kutta_4(points, v, 0),
advect.runge_kutta_4(points, v * 0, 0))
sv = domain.staggered_grid(Noise())
field.assert_close(points, advect.runge_kutta_4(points, sv, 0),
advect.runge_kutta_4(points, sv * 0, 0))
def test_domain_grid_from_constant(self):
domain = Domain(x=4, y=3)
# int / float
grid = domain.scalar_grid(1)
math.assert_close(grid.values, 1)
# complex
grid = domain.scalar_grid(1 + 1j)
math.assert_close(grid.values, 1 + 1j)
# NumPy
grid = domain.scalar_grid(numpy.array(1))
math.assert_close(grid.values, 1)
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_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_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_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_make_incompressible_gradients_equal_tf_torch(self):
DOMAIN = Domain(x=16, y=16, boundaries=OPEN, bounds=Box[0:100, 0:100]) # TODO CLOSED solve fails because div is not subtracted from dx
velocity0 = DOMAIN.staggered_grid(Noise(vector=2))
grads = []
for backend in [TF_BACKEND, TORCH_BACKEND]:
with backend:
velocity = param = velocity0.with_(values=math.tensor(velocity0.values))
with math.record_gradients(param.values):
solve = math.LinearSolve()
velocity, _, _, _ = fluid.make_incompressible(velocity, DOMAIN, solve_params=solve)
loss = field.l2_loss(velocity)
assert math.isfinite(loss)
grad = math.gradients(loss, param.values)
assert math.all(math.isfinite(grad))
grads.append(grad)
math.assert_close(*grads, abs_tolerance=1e-5)
def test_effects(self):
world = World()
fluid = world.add(Fluid(Domain(x=16, y=16)),
physics=IncompressibleFlow())
fan = world.add(Fan(Sphere((10, 8), 5), [-1, 0]))
obstacle = world.add(Obstacle(Box[0:1, 0:1]))
world.step(dt=1)
world.step(dt=0.5)
assert fluid.age == fan.age == obstacle.age == 1.5
def test_gradient_function(self):
def f(x: StaggeredGrid, y: CenteredGrid):
pred = x + (y >> x)
loss = field.l2_loss(pred)
return loss
domain = Domain(x=4, y=3)
x = domain.staggered_grid(1)
y = domain.vector_grid(1)
with torch.TORCH_BACKEND:
dx, = field.gradient_function(f)(x, y)
self.assertIsInstance(dx, StaggeredGrid)
loss, dx, dy = field.gradient_function(f, (0, 1),
get_output=True)(x, y)
self.assertIsInstance(loss, math.Tensor)
self.assertIsInstance(dx, StaggeredGrid)
self.assertIsInstance(dy, CenteredGrid)
def test_smoke_plume(self):
world = World()
world.batch_size = 3
fluid = world.add(Fluid(Domain(x=16, y=16)),
physics=IncompressibleFlow())
inflow = world.add(Inflow(Sphere((8, 8), radius=4)))
world.step()
world.step(fluid)
self.assertAlmostEqual(fluid.age, 2.0)
self.assertAlmostEqual(inflow.age, 1.0)
def test_gradient_function(self):
def f(x: StaggeredGrid, y: CenteredGrid):
pred = x + (y >> x)
loss = field.l2_loss(pred)
return loss
domain = Domain(x=4, y=3)
x = domain.staggered_grid(1)
y = domain.vector_grid(1)
for backend in BACKENDS:
if backend.supports(Backend.gradients):
with backend:
dx, = field.functional_gradient(f)(x, y)
self.assertIsInstance(dx, StaggeredGrid)
loss, dx, dy = field.functional_gradient(f, (0, 1), get_output=True)(x, y)
self.assertIsInstance(loss, math.Tensor)
self.assertIsInstance(dx, StaggeredGrid)
self.assertIsInstance(dy, CenteredGrid)
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_domain_grid_from_constant(self):
domain = Domain(x=4, y=3)
# int / float
grid = domain.scalar_grid(1)
math.assert_close(grid.values, 1)
# complex
grid = domain.grid(1 + 1j)
math.assert_close(grid.values, 1 + 1j)
# NumPy
grid = domain.grid(numpy.array(1))
math.assert_close(grid.values, 1)
# PyTorch
grid = domain.grid(torch.tensor(1, dtype=torch.float32))
math.assert_close(grid.values, 1)
# TensorFlow
grid = domain.grid(tf.constant(1, tf.float32))
math.assert_close(grid.values, 1)
def simulate(centers):
world = World()
fluid = world.add(Fluid(Domain(x=5,
y=4,
boundaries=CLOSED,
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 test_properties_dict(self):
world = World()
world.add(Fluid(Domain(x=16, y=16)), physics=IncompressibleFlow())
world.add(Inflow(Sphere((8, 8), radius=4)))
world.add(Fan(Sphere((10, 8), 5), [-1, 0]))
struct.properties_dict(world.state)
def test_domain_grid_memory_allocation(self):
domain = Domain(x=10000, y=10000, z=10000, w=10000)
grid = domain.scalar_grid()
self.assertEqual((10000, ) * 4, grid.shape.sizes)
sgrid = domain.staggered_grid()
self.assertEqual((10000, 10000, 10000, 10000, 4), sgrid.shape.sizes)
def test_varying_boundaries(self):
fluid = Fluid(Domain(x=16, y=16, boundaries=[(CLOSED, OPEN), CLOSED]))
IncompressibleFlow().step(fluid)
def test_sqrt(self):
grid = Domain(x=4, y=3).staggered_grid(2)
field.assert_close(field.sqrt(grid), numpy.sqrt(2))
def test_direct_fluid(self):
fluid = Fluid(Domain(x=16, y=16))
fluid2 = IncompressibleFlow().step(fluid)
assert fluid2.age == 1.0
assert fluid.age == 0.0
assert fluid2.name == fluid.name
def test_exp(self):
grid = Domain(x=4, y=3).staggered_grid(0)
field.assert_close(field.exp(grid), 1)
def test_demo_vs_numpy(self):
dt = 0.1
# Physical parameters
L = 4 * 2 * np.pi # 2 * np.pi / 0.15 # 4*2*np.pi#/0.15
c1 = 5 # 0.01
# Numerical Parameters
grid_pts = 64
nu = 0 # 1e-8
N = 0 # 3
arakawa_coeff = 0
kappa_coeff = 0
# Derived Parameters
dx = L / grid_pts
k0 = 2 * np.pi / L
# Get input data
rnd_noise = np.random.rand(grid_pts * grid_pts).reshape(
grid_pts, grid_pts)
sine = get_2d_sine((grid_pts, grid_pts), L)
init_values = sine / 1000
density_coeff = 1
omega_coeff = -1 / 2
phi_coeff = -1 / 2
x = grid_pts
y = grid_pts
params = dict(c1=c1,
nu=nu,
N=N,
arakawa_coeff=arakawa_coeff,
kappa_coeff=kappa_coeff)
# NumPy reference
hw = HW(
c1=c1,
nu=nu,
N=N,
arakawa_coeff=arakawa_coeff,
kappa_coeff=kappa_coeff,
debug=False,
)
hw_state_numpy = Namespace(
density=init_values * density_coeff,
omega=init_values * omega_coeff,
phi=init_values * phi_coeff,
age=0,
dx=dx,
)
# Phi version
domain = Domain(x=x,
y=y,
boundaries=PERIODIC,
bounds=Box[0:L, 0:L])
hw_state_phi = Namespace(
density=domain.grid(
math.tensor(init_values * density_coeff, names=["x",
"y"])),
omega=domain.grid(
math.tensor(init_values * omega_coeff, names=["x", "y"])),
phi=domain.grid(
math.tensor(init_values * phi_coeff, names=["x", "y"])),
# domain=domain,
age=0,
dx=dx,
)
from functools import partial
get_phi = partial(get_domain_phi, domain=domain)
rk4_step2 = partial(rk4_step, params=params, get_phi=get_phi)
app = App("Hasegawa-Wakatani", dt=dt)
app.set_state(hw_state_phi,
step_function=rk4_step2,
show=["density", "omega", "phi"])
app.prepare()
# Run
def compare(iterable):
for k in iterable[0].keys():
compare = []
for state in iterable:
if isinstance(state[k], field._grid.Grid):
val = state[k].values.numpy(order="x,y,z")[0]
else:
val = state[k]
compare.append(val)
assert len(compare) == 2
print(f" {k:<7}:" +
f" {np.max(np.abs(compare[0]-compare[1])):.7f}" +
f" {np.array_equal(*compare)}" +
f" {np.max(np.abs(compare[0]-compare[1])):.2g}")
return True
np_times = []
phi_times = []
for i in range(0, STEPS + 1):
print(f"step {i:>3} {i*dt:>12.7f}")
# Numpy
t0 = time.time()
hw_state_numpy = hw.rk4_step(hw_state_numpy, dt=dt)
np_time = time.time() - t0
np_times.append(np_time)
# Phi
t0 = time.time()
app.step()
phi_time = time.time() - t0
phi_times.append(phi_time)
hw_state_phi = app.state
compare([hw_state_numpy, hw_state_phi])
# if i % 100 == 0:
# plot({'numpy': hw_state_numpy.density,
# 'phi': hw_state_phi.density.values.numpy(order='zyx')[0],
# 'diff': hw_state_numpy.density - hw_state_phi.density.values.numpy(order='zyx')[0]},
# title=f"step: {i}")
# assert np.allclose(hw_state_numpy.density, hw_state_phi.density.values.numpy(order='zyx')[0])
print(f"Comparison | NumPy | PhiFlow")
def get_str(name, func, vals):
return " | ".join([
f"{name:<10}",
f"{func(vals[0]):<10.4f}",
f"{func(vals[1]):<10.4f}",
])
print(get_str("Mean (s)", np.mean, (np_times, phi_times)))
print(get_str("Median (s)", np.median, (np_times, phi_times)))
print(get_str("Min (s)", np.min, (np_times, phi_times)))
print(get_str("Max (s)", np.max, (np_times, phi_times)))
def test_domain_grid_from_field(self):
large_grid = Domain(x=4, y=3).grid(Noise())
small_grid = Domain(x=3, y=2).grid(large_grid)
math.assert_close(large_grid.values.x[:-1].y[:-1], small_grid.values)
def test_domain_grid_from_tensor(self):
domain = Domain(x=4, y=3)
grid = domain.vector_grid(Noise(vector=2))
grid2 = domain.vector_grid(grid.values)
math.assert_close(grid.values, grid2.values)
def test_domain_grid_from_function(self):
grid = Domain(x=4, y=3).scalar_grid(lambda x: math.sum(x**2, 'vector'))
math.assert_close(grid.values.x[0].y[0], 0.5)
self.assertEqual(grid.shape.volume, 12)
grid = Domain(x=4, y=3).scalar_grid(lambda x: 1)
math.assert_close(grid.values, 1)
def test_custom_spatial_dims(self):
domain = Domain(a=4, b=3)
grid = domain.scalar_grid(1)
self.assertEqual(math.shape(a=4, b=3), grid.shape)
grid = domain.staggered_grid(1)
self.assertEqual(math.shape(a=4, b=3, vector=2), grid.shape)
