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)