def test_initialization_with_cuds(self):
# given
points = [
[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
bonds = [[0, 1], [0, 3], [1, 3, 2]]
point_temperature = [10., 20., 30., 40.]
bond_temperature = [60., 80., 190., 5.]
reference = Particles('test')
# add particles
particle_iter = (Particle(coordinates=point,
data=DataContainer(TEMPERATURE=temp))
for temp, point in zip(point_temperature, points))
point_uids = reference.add(particle_iter)
# add bonds
bond_iter = (Bond(particles=[point_uids[index] for index in indices],
data=DataContainer(TEMPERATURE=temp))
for temp, indices in zip(bond_temperature, bonds))
bond_uids = reference.add(bond_iter)
# when
container = VTKParticles.from_particles(reference)
# then
number_of_particles = sum(1 for _ in container.iter(
item_type=CUBA.PARTICLE))
self.assertEqual(number_of_particles, len(point_uids))
for expected in reference.iter(item_type=CUBA.PARTICLE):
self.assertEqual(container.get(expected.uid), expected)
number_of_bonds = sum(1 for _ in container.iter(item_type=CUBA.BOND))
self.assertEqual(number_of_bonds, len(bond_uids))
for expected in reference.iter(item_type=CUBA.BOND):
self.assertEqual(container.get(expected.uid), expected)
class TestParticlesSource(unittest.TestCase):
tested_class = CUDSSource
def setUp(self):
self.points = [
[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
self.bonds = [[0, 1], [0, 3], [1, 3, 2]]
self.point_temperature = [10., 20., 30., 40.]
self.bond_temperature = [60., 80., 190., 5.]
self.container = Particles('test')
# add particles
def particle_iter():
for temp, point in zip(self.point_temperature, self.points):
yield Particle(coordinates=point,
data=DataContainer(TEMPERATURE=temp))
self.point_uids = self.container.add(particle_iter())
# add bonds
def bond_iter():
for temp, indices in zip(self.bond_temperature, self.bonds):
yield Bond(particles=[self.point_uids[index]
for index in indices],
data=DataContainer(TEMPERATURE=temp))
self.bond_uids = self.container.add(bond_iter())
# for testing save/load visualization
self.temp_dir = tempfile.mkdtemp()
def tearDown(self):
shutil.rmtree(self.temp_dir)
def test_source_from_vtk_particles(self):
# given
container = VTKParticles('test')
container.add(self.container.iter(item_type=CUBA.PARTICLE))
container.add(self.container.iter(item_type=CUBA.BOND))
# when
source = self.tested_class(cuds=container)
# then
self.assertIs(source.data, container.data_set)
self.assertIs(source.cuds, container)
def test_particles(self):
# given
container = self.container
# when
source = self.tested_class(cuds=container)
# then
points = source.data.points.to_array()
dataset = source.data
vtk_cuds = source._vtk_cuds
number_of_particles = len(self.points)
self.assertEqual(len(points), number_of_particles)
self.assertEqual(len(vtk_cuds.particle2index), number_of_particles)
self.assertEqual(dataset.point_data.number_of_arrays, 1)
temperature = dataset.point_data.get_array('TEMPERATURE')
for key, index in vtk_cuds.particle2index.iteritems():
point = container.get(key)
assert_array_equal(points[index], point.coordinates)
self.assertEqual(temperature[index], point.data[CUBA.TEMPERATURE])
def test_bonds(self):
# given
container = self.container
# when
source = self.tested_class(cuds=container)
# then
dataset = source.data
vtk_cuds = source._vtk_cuds
bonds = [
bond for bond in cell_array_slicer(dataset.lines.to_array())]
number_of_bonds = len(self.bonds)
self.assertEqual(len(bonds), number_of_bonds)
self.assertEqual(len(vtk_cuds.bond2index), number_of_bonds)
self.assertEqual(dataset.cell_data.number_of_arrays, 1)
temperature = dataset.cell_data.get_array('TEMPERATURE')
for key, index in vtk_cuds.bond2index.iteritems():
bond = container.get(key)
particles = [
vtk_cuds.particle2index[uid] for uid in bond.particles]
self.assertEqual(bonds[index], particles)
self.assertEqual(temperature[index], bond.data[CUBA.TEMPERATURE])
def test_particles_source_name(self):
# given
particles = Particles(name='my_particles')
# when
source = self.tested_class(cuds=particles)
# then
self.assertEqual(source.name, 'my_particles (CUDS Particles)')
def check_save_load_visualization(self, engine):
# set up the visualization
container = self.container
source = self.tested_class(cuds=container)
engine.add_source(source)
# save the visualization
saved_viz_file = os.path.join(self.temp_dir, 'test_saved_viz.mv2')
engine.save_visualization(saved_viz_file)
engine.close_scene(engine.current_scene)
# restore the visualization
engine.load_visualization(saved_viz_file)
# then
source_in_scene = engine.current_scene.children[0]
points = source_in_scene.data.points.to_array()
dataset = source_in_scene.data
number_of_particles = len(self.points)
# data is restored
self.assertEqual(len(points), number_of_particles)
self.assertEqual(dataset.point_data.number_of_arrays, 1)
# But cuds and vtk_cuds are not available
self.assertIsNone(source_in_scene._vtk_cuds)
self.assertIsNone(source_in_scene._cuds)
@unittest.skipIf(is_mayavi_older("4.4.4"),
"Mayavi < 4.4.4 has problem with load_visualization")
def test_save_load_visualization_with_mlab(self):
# test mlab.get_engine
engine = mlab.get_engine()
try:
self.check_save_load_visualization(engine)
finally:
mlab.clf()
mlab.close(all=True)
def test_save_load_visualization_with_null_engine(self):
self.check_save_load_visualization(NullEngine())
# according to a Maxwell-Boltzmann distribution. In this example
# we define a uniform random distribution, the MD algorithm will
# in any case result in a MB one quickly.
p.data[CUBA.VELOCITY] = [
numpy.random.uniform(-0.5, 0.5),
numpy.random.uniform(-0.5, 0.5),
numpy.random.uniform(-0.5, 0.5)
]
pc.add([p])
input_particles.append(p)
# Calculate the velocity of center of mass and reset it to zero to
# avoid drift of the whole system.
v_cm = [0, 0, 0]
for p in pc.iter(item_type=CUBA.PARTICLE):
v_cm[0] += p.data[CUBA.VELOCITY][0]
v_cm[1] += p.data[CUBA.VELOCITY][1]
v_cm[2] += p.data[CUBA.VELOCITY][2]
number_of_points = pc.count_of(CUBA.PARTICLE)
v_cm[0] /= number_of_points
v_cm[1] /= number_of_points
v_cm[2] /= number_of_points
for p in pc.iter(item_type=CUBA.PARTICLE):
p.data[CUBA.VELOCITY][0] -= v_cm[0]
p.data[CUBA.VELOCITY][1] -= v_cm[1]
p.data[CUBA.VELOCITY][2] -= v_cm[2]
######################################
