def setUp(self): magnet1 = MagnetModel(1.0, 0.0, 1.5) magnet2 = MagnetModel(-1.0, 1.0, 1.5) magnet3 = MagnetModel(-1.0, -1.0, 1.5) magnet4 = MagnetModel(0.0, 0.0, 0.5) self.magnets = [magnet1, magnet2, magnet3, magnet4] self.plane_distance = 0.25 self.friction = 0.3 self.gravity_pullback = 0.5 self.pendulum = PendulumModel(self.friction, self.gravity_pullback, self.plane_distance) self.pendulum.magnets = self.magnets self.time_step = 0.01 self.integrator = EulerIntegrator(self.time_step) self.r = 255 self.g = 0 self.b = 0 self.image_generator = BasicImageGenerator(self.r, self.g, self.b) self.size = 5 self.resolution = 640 self.basins_generator = BasinsGenerator(self.size, self.resolution) self.basins_generator.pendulum_model = self.pendulum self.basins_generator.integrator = self.integrator self.basins_generator.image_generator = self.image_generator
class TestBasinsGenerator(unittest.TestCase): def setUp(self): magnet1 = MagnetModel(1.0, 0.0, 1.5) magnet2 = MagnetModel(-1.0, 1.0, 1.5) magnet3 = MagnetModel(-1.0, -1.0, 1.5) magnet4 = MagnetModel(0.0, 0.0, 0.5) self.magnets = [magnet1, magnet2, magnet3, magnet4] self.plane_distance = 0.25 self.friction = 0.3 self.gravity_pullback = 0.5 self.pendulum = PendulumModel(self.friction, self.gravity_pullback, self.plane_distance) self.pendulum.magnets = self.magnets self.time_step = 0.01 self.integrator = EulerIntegrator(self.time_step) self.r = 255 self.g = 0 self.b = 0 self.image_generator = BasicImageGenerator(self.r, self.g, self.b) self.size = 5 self.resolution = 640 self.basins_generator = BasinsGenerator(self.size, self.resolution) self.basins_generator.pendulum_model = self.pendulum self.basins_generator.integrator = self.integrator self.basins_generator.image_generator = self.image_generator def test_basins_initalization(self): self.assertEqual(self.basins_generator.size, self.size) self.assertEqual(self.basins_generator.resolution, self.resolution) self.assertEqual(self.basins_generator.cuda_device_number, 0) self.basins_generator.pendulum_model = self.pendulum self.basins_generator.integrator = self.integrator self.basins_generator.image_generator = self.image_generator self.assertEqual(self.basins_generator.pendulum_model, self.pendulum) self.assertEqual(self.basins_generator.integrator, self.integrator) self.assertEqual(self.basins_generator.image_generator, self.image_generator) self.assertEqual(self.basins_generator.pendulum_model.magnets, self.magnets) def test_basins_gpu_calculation(self): vel_vect = [0, 0] sim_time = 5 delta = 0.2 file_name = "test_image" self.basins_generator.calculate_basins(vel_vect, sim_time, delta) self.basins_generator.draw_basins(file_name) self.assertEqual(len(self.basins_generator.result_data), 640)