def test_state_collection(self):
fluid = Fluid(Domain([1, 1]))
fluid2 = Fluid(Domain([2, 2]))
c1 = StateCollection([fluid])
assert c1.fluid is fluid
assert fluid in c1
assert c1[fluid] is fluid
assert isinstance(repr(c1), six.string_types)
assert len(c1) == len(c1.shape) == len(c1.staticshape) == len(c1.dtype)
assert c1.shape.fluid.density.data == (1, 1, 1, 1)
self.assertIsInstance(c1.dtype.fluid.density.data, numpy.dtype)
c2 = StateCollection()
assert len(c2) == 0
c2 = c2.state_added(fluid)
assert c2 == c1
assert hash(c2) == hash(c1)
c3 = c2.state_replaced(fluid2)
assert c3 != c2
assert c3.fluid is fluid2
c4 = c3.state_removed(fluid2)
assert len(c4) == 0
c5 = struct.map(lambda x: x, c1)
assert isinstance(c5, StateCollection)
assert c5 == c1
def test_convenience_initializers(self):
domain = Domain(64)
numpy.testing.assert_equal(domain.resolution, [64])
numpy.testing.assert_equal(domain.box.size, [64])
domain = Domain(64, box=10)
numpy.testing.assert_equal(domain.resolution, [64])
numpy.testing.assert_equal(domain.box.size, [10])
def test_boundary_definitions(self):
domain = Domain([128, 128, 16], (OPEN, OPEN, OPEN))
self.assertEqual(domain.boundaries, OPEN)
domain = Domain([64, 32], boundaries=[(OPEN, OPEN), (OPEN, OPEN)])
self.assertEqual(domain.boundaries, OPEN)
try:
Domain([64, 32], None)
self.fail()
except AssertionError:
pass
def test_names(self):
c = CollectiveState()
self.assertEqual(c.states, {})
c = c.state_added(Fluid(Domain([64])))
try:
c = c.state_added(Fluid(Domain([80])))
self.fail()
except AssertionError:
pass
c = c.state_replaced(Fluid(Domain([80])))
numpy.testing.assert_equal(c.fluid.density.data.shape, [1, 80, 1])
world = World(add_default_objects=True)
assert world.gravity.state is world.state.gravity
def test_tf_worldgraph(self):
world = World()
fluid = world.add(Fluid(Domain([16, 16])))
tf_bake_graph(world, Session(Scene.create('data', copy_calling_script=False)))
world.step()
self.assertIsInstance(fluid.state, Fluid)
self.assertIsInstance(fluid.state.density.data, numpy.ndarray)
def test_properties_dict(self):
world = World()
world.add(Fluid(Domain([16, 16])), physics=IncompressibleFlow())
world.add(Inflow(Sphere((8, 8), radius=4)))
# world.add(ConstantDensity(box[0:2, 6:10], 1.0))
world.add(Fan(Sphere((10, 8), 5), [-1, 0]))
struct.properties_dict(world.state)
def test_gradient_batch_independence(self):
session = Session(None) # Used to run the TensorFlow graph
world = World()
fluid = world.add(Fluid(Domain([40, 32], boundaries=CLOSED), buoyancy_factor=0.1, batch_size=2), physics=IncompressibleFlow())
world.add(Inflow(Sphere(center=numpy.array([[5, 4], [5, 8]]), radius=3), rate=0.2))
fluid.velocity = variable(fluid.velocity) # create TensorFlow variable
# fluid.velocity *= 0
initial_state = fluid.state # Remember the state at t=0 for later visualization
session.initialize_variables()
for frame in range(3):
world.step(dt=1.5)
target = session.run(fluid.density).data[0, ...]
loss = tf.nn.l2_loss(fluid.density.data[1, ...] - target)
self_loss = tf.nn.l2_loss(fluid.density.data[0, ...] - target)
# loss = self_loss
optim = tf.train.GradientDescentOptimizer(learning_rate=0.2).minimize(loss)
session.initialize_variables()
for optim_step in range(3):
_, loss_value, sl_value = session.run([optim, loss, self_loss])
staggered_velocity = session.run(initial_state.velocity).staggered_tensor()
numpy.testing.assert_equal(staggered_velocity[0, ...], 0)
assert numpy.all(~numpy.isnan(staggered_velocity))
def test_direct_fluid(self):
fluid = Fluid(Domain([16, 16]))
assert fluid.default_physics() == INCOMPRESSIBLE_FLOW
fluid2 = INCOMPRESSIBLE_FLOW.step(fluid)
assert fluid2.age == 1.0
assert fluid.age == 0.0
assert fluid2.name == fluid.name
def generate_test_structs():
return [
manta.centered_grid(numpy.zeros([1, 4, 1])), [('Item', )], {
'A': 'Entry A',
'Vel': manta.staggered_grid(numpy.zeros([1, 5, 5, 2]))
},
CollectiveState((Fluid(Domain([4])), ))
]
def test_effects(self):
world = World()
fluid = world.add(Fluid(Domain([16, 16])))
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_fluid_initializers(self):
def typetest(fluid):
self.assertIsInstance(fluid, Fluid)
self.assertIsInstance(fluid.velocity, StaggeredGrid)
numpy.testing.assert_equal(fluid.density.data.shape, [1, 4, 4, 1])
numpy.testing.assert_equal(fluid.velocity.resolution, [4, 4])
numpy.testing.assert_equal(fluid.velocity.data[0].resolution,
[5, 4])
typetest(Fluid(Domain([4, 4]), density=0.0, velocity=0.0))
typetest(Fluid(Domain([4, 4]), density=1.0, velocity=1.0))
typetest(Fluid(Domain([4, 4]), density=0, velocity=math.zeros))
typetest(
Fluid(Domain([4, 4]),
density=lambda s: math.randn(s),
velocity=lambda s: math.randn(s)))
typetest(
Fluid(Domain([4, 4]),
density=numpy.zeros([1, 4, 4, 1]),
velocity=numpy.zeros([1, 5, 5, 2])))
typetest(
Fluid(Domain([4, 4]),
density=numpy.zeros([1, 4, 4, 1]),
velocity=numpy.zeros([1, 5, 5, 2])))
typetest(Fluid(Domain([4, 4])))
def test_simpleplume(self):
world = World()
world.batch_size = 3
fluid = world.add(Fluid(Domain([16, 16])))
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_complex_step(self):
q = QuantumWave(Domain([4, 4]))
q = q.copied_with(amplitude=WavePacket([2, 2], 1.0, [0.5, 0]))
pot = StepPotential(box[0:1, 0:1], 1.0)
SCHROEDINGER.step(q,
1.0,
potentials=[pot],
obstacles=[Obstacle(box[3:4, 0:1])])
numpy.testing.assert_equal(q.amplitude.data.shape, [1, 4, 4, 1])
def test_batched_forced_burgers_2d(self):
world = World(batch_size=3)
burgers = world.add(Domain([4, 4]).centered_grid(0, batch_size=world.batch_size, name='velocity'))
k = math.to_float(numpy.random.uniform(3, 6, [world.batch_size, 2]))
amplitude = numpy.random.uniform(-0.5, 0.5, [world.batch_size])
force = SinPotential(k, phase_offset=numpy.random.uniform(0, 2 * numpy.pi, [world.batch_size]), data=amplitude)
physics = ForcingPhysics(numpy.random.uniform(-0.4, 0.4, [world.batch_size]))
effect = FieldEffect(force, ['velocity'])
world.add(effect, physics=physics)
burgers.step()
burgers.step()
def test_copy(self):
with struct.unsafe():
fluid = Fluid(Domain([4]), density='Density', velocity='Velocity')
v = fluid.copied_with(velocity='V2')
self.assertEqual(v.velocity, 'V2')
self.assertEqual(v.density, 'Density')
try:
fluid.copied_with(velocity='D2')
self.fail()
except AssertionError:
pass
def simulate(centers):
world = World()
fluid = world.add(Fluid(Domain([5, 4], boundaries=CLOSED, box=AABox(0, [40, 32])),
buoyancy_factor=0.1,
batch_size=centers.shape[0]),
physics=IncompressibleFlow(pressure_solver=SparseCG(max_iterations=3)))
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)
print()
return fluid.density.data[0, ...], fluid.velocity.unstack()[0].data[0, ...], fluid.velocity.unstack()[1].data[0, ...]
def test_precision_16(self):
try:
math.set_precision(16)
fluid = Fluid(Domain([16, 16]), density=math.maximum(0, Noise()))
self.assertEqual(fluid.density.data.dtype, numpy.float16)
self.assertEqual(fluid.velocity.unstack()[0].data.dtype,
numpy.float16)
fluid = IncompressibleFlow().step(fluid, dt=1.0)
self.assertEqual(fluid.density.data.dtype, numpy.float16)
self.assertEqual(fluid.velocity.unstack()[0].data.dtype,
numpy.float16)
finally:
math.set_precision(32) # Reset environment
def test_fluid_tf(self):
world = World()
fluid = Fluid(Domain([16, 16]))
world.add(fluid)
world.add(Inflow(Sphere((8, 8), radius=4)))
world.add(Obstacle(box[4:16, 0:8]))
fluid_in = fluid.copied_with(density=placeholder, velocity=placeholder)
fluid_out = world.step(fluid_in)
self.assertIsInstance(fluid_out, Fluid)
session = Session(Scene.create('data', copy_calling_script=False))
fluid = session.run(fluid_out, {fluid_in: fluid})
fluid = session.run(fluid_out, {fluid_in: fluid})
self.assertIsInstance(fluid, Fluid)
def test_read_write_struct(self):
for scene in Scene.list('data'):
scene.remove()
state = Fluid(Domain([4, 4]))
scene = Scene.create('data')
scene.write(state, frame=0)
self.assertTrue(isfile(scene.subpath('density_000000.npz')))
self.assertTrue(isfile(scene.subpath('velocity_000000.npz')))
loaded_state = scene.read(state, frame=0)
self.assertIsInstance(loaded_state, Fluid)
self.assertIsInstance(loaded_state.velocity, StaggeredGrid)
self.assertIsInstance(loaded_state.density, CenteredGrid)
_differences = struct.compare([loaded_state.density, state.density])
self.assertEqual(loaded_state.density, state.density)
print_differences(loaded_state.velocity.data, state.velocity.data)
np.testing.assert_equal(loaded_state.velocity.data[0].data,
state.velocity.data[0].data)
scene.write(np.ones([1, 4, 4, 1]) * 2, frame=1)
self.assertTrue(isfile(scene.subpath('unnamed_000001.npz')))
self.assertEqual(scene.read(None, frame=1)[0, 0, 0, 0], 2)
scene.write([np.ones([1, 4, 4, 1])], ['Ones'], frame=2)
self.assertTrue(isfile(scene.subpath('Ones_000002.npz')))
mystruct = [{
'Two': np.ones([1, 4, 4, 1]) * 2,
'Three': np.ones([1, 4, 4, 1]) * 3
}]
scene.write(mystruct, frame=3)
self.assertTrue(isfile(scene.subpath('0_Three_000003.npz')))
self.assertTrue(isfile(scene.subpath('0_Two_000003.npz')))
loaded_struct = scene.read(mystruct, frame=3)
self.assertIsInstance(loaded_struct, list)
np.testing.assert_equal(mystruct[0]['Two'][0, 0, 0, 0],
loaded_struct[0]['Two'][0, 0, 0, 0])
scene.remove()
def test_simple_step(self):
q = QuantumWave(Domain([4, 5]))
q = q.copied_with(amplitude=WavePacket([2, 2], 1.0, [0.5, 0]))
q = SCHROEDINGER.step(q, 1.0)
numpy.testing.assert_equal(q.amplitude.data.shape, [1, 4, 5, 1])
def test_new_grids(self):
fluid = Fluid(Domain([16, 16]), batch_size=3)
centered_ones = fluid.centered_grid('f', 1)
numpy.testing.assert_equal(centered_ones.data, 1)
staggered_ones = fluid.staggered_grid('v', 1)
numpy.testing.assert_equal(staggered_ones.data[0].data, 1)
def test_staggered_curl2d(self):
domain = Domain([32, 32])
pot = CenteredGrid.sample(Noise(), domain)
vel = staggered_curl_2d(pot)
div = vel.divergence()
np.testing.assert_almost_equal(div.data, 0, decimal=5)
def test_effects(self):
world = World()
world.add(Fluid(Domain([16, 16])))
world.add(Fan(Sphere((10, 8), 5), [-1, 0]))
world.step()
world.step()
def test_varying_boundaries(self):
fluid = Fluid(Domain([16, 16], boundaries=[(CLOSED, OPEN), CLOSED]))
INCOMPRESSIBLE_FLOW.step(fluid)