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] ######################################