def test_magnet_setup(self):
"""Test magnet_setup logic"""
# Check that this works in the absence of fields
main_window = envelope_tool.MainWindow()
for button in ["&Cancel", "&Okay"]:
main_window.window.get_frame("&Magnet Setup", "button").Clicked()
magnet_setup = main_window.magnet_setup
magnet_setup.window.get_frame(button, "button").Clicked()
main_window.window.get_frame("E&xit", "button").Clicked()
# now with fields
sys.argv += ["--configuration_file", TEST_DIR + "test_config.py"]
main_window = envelope_tool.MainWindow()
main_window.window.get_frame("&Magnet Setup", "button").Clicked()
magnet_setup = main_window.magnet_setup
# Check field data is represented (manipulator worked)
test_data = {"Q1": -1., "D2": -0.9, "Q2": 1.}
for test, data in test_data.iteritems():
text = magnet_setup.window.get_text_entry(test, type(1.))
self.assertAlmostEqual(text, data)
# Check cancel action
magnet_setup.window.set_text_entry("Q1", 2.)
magnet_setup.window.get_frame("&Cancel", "button").Clicked()
main_window.window.get_frame("&Magnet Setup", "button").Clicked()
magnet_setup = main_window.magnet_setup
q1 = magnet_setup.window.get_text_entry("Q1", type(1.))
self.assertAlmostEqual(q1, test_data["Q1"])
# Check okay action
magnet_setup.window.set_text_entry("Q1", 2.)
magnet_setup.window.get_frame("&Okay", "button").Clicked()
main_window.window.get_frame("&Magnet Setup", "button").Clicked()
magnet_setup = main_window.magnet_setup
q1 = magnet_setup.window.get_text_entry("Q1", type(1.))
self.assertAlmostEqual(q1, 2.)
sys.argv = sys.argv[0:-2]
def test_plot_setup(self):
"""Test plot_setup logic"""
main_window = envelope_tool.MainWindow()
main_window.window.get_frame("&Plot Setup", "button").Clicked()
plot_setup = main_window.plot_setup
# check select action
first_var = plot_setup.window.get_frame("first_var_0", "drop_down")
self.assertEqual(first_var.GetNumberOfEntries(), 1)
plot_setup.window.get_frame("variable_type_0", "drop_down").Select(1)
self.assertGreater(first_var.GetNumberOfEntries(), 1)
first_var.Select(1)
# check cancel action
test_options = main_window.plot_setup_options
plot_setup.window.get_frame("&Cancel", "button").Clicked()
self.assertEqual(main_window.plot_setup, None)
self.assertEqual(main_window.plot_setup_options, test_options)
# Check okay action
main_window.window.get_frame("&Plot Setup", "button").Clicked()
plot_setup = main_window.plot_setup
plot_setup.window.get_frame("variable_type_0", "drop_down").Select(1)
plot_setup.window.get_frame("first_var_0", "drop_down").Select(1)
plot_setup.window.get_frame("plot_apertures",
"check_button").SetState(0)
test_options = [{
"variable_type": 1,
"first_var": 1,
"plot_apertures": False
}]
plot_setup.window.get_frame("&Okay", "button").Clicked()
self.assertEqual(main_window.plot_setup, None)
self.assertEqual(main_window.plot_setup_options, test_options)
def setUp(self):
envelope_tool.set_share_dirs()
self.red_dict = {
'pid': 13,
'px': 1.,
'py': 2.,
'pz': 3.,
'x': -1,
'y': -2,
'z': -3
}
self.hit_red = Hit.new_from_dict(self.red_dict)
self.hit_full = Hit.new_from_dict(
{
'pid': 13,
'px': 1.,
'py': 2.,
'pz': 3.,
'x': -1,
'y': -2,
'z': -3,
'mass': common.pdg_pid_to_mass[13],
'charge': common.pdg_pid_to_charge[13]
}, 'energy')
self.main_window = envelope_tool.MainWindow()
self.main_window.window.get_frame("&Beam Setup", "button").Clicked()
self.beam_setup = self.main_window.beam_setup
self.cov = CovarianceMatrix()
for i in range(1, 7):
for j in range(1, 7):
self.cov.set_element(i, j, i * j)
self.cov.set_element(i, i, i * i + 50.)
def test_lattice_get_set_beam(self):
"""Test lattice - get set beam"""
sys.argv += ["--configuration_file", TEST_DIR + "test_config.py"]
_lattice = envelope_tool.MainWindow().lattice
self.assertGreater(len(_lattice.ref_list), 1)
self.assertGreater(len(_lattice.ellipse_list), 1)
hit_ref, ellipse_ref = _lattice.get_beam()
hit_ref['pid'] = 2212
ellipse_ref.set_element(1, 1, 999)
_lattice.set_beam(hit_ref, ellipse_ref)
hit, ellipse = _lattice.get_beam()
self.assertEqual(hit_ref['pid'], hit['pid'])
self._ell_equal(ellipse, ellipse_ref)
self.assertEqual(len(_lattice.ref_list), 1)
self.assertEqual(len(_lattice.ellipse_list), 1)
def test_lattice_run_lattice_ellipses(self):
"""Test lattice - run lattice"""
# check that ellipse propagates okay
# MICE SFoFo lattice nearly periodic with beta=431.
sys.argv += ["--configuration_file", TEST_DIR + "test_config_2.py"]
_lattice = envelope_tool.MainWindow().lattice
self.assertGreater(len(_lattice.ref_list), 1)
self.assertGreater(len(_lattice.ellipse_list), 1)
ref, ellipse = _lattice.get_beam()
ellipse = maus_cpp.covariance_matrix.create_from_penn_parameters(
ref["mass"], ref["p"], 1.0, 431., 10., 10.)
_lattice.set_beam(ref, ellipse)
_lattice.run_lattice()
last_ellipse = _lattice.ellipse_list[-1]
ratio_sxx = ellipse.get_element(3, 3) / last_ellipse.get_element(3, 3)
self.assertLess(abs(ratio_sxx - 1.), 5e-2)
ratio_spp = ellipse.get_element(4, 4) / last_ellipse.get_element(4, 4)
self.assertLess(abs(ratio_spp - 1.), 5e-2)
def test_lattice_get_set_fields(self):
"""Test lattice - get set fields"""
sys.argv += ["--configuration_file", TEST_DIR + "test_config.py"]
main_window = envelope_tool.MainWindow()
# get_set_fields
field_list = main_window.lattice.get_field_list()
test_q1 = {
"field_name": "Q1",
"field_type": "Multipole",
"scale_factor": -1.,
"aperture": {
"x": 172.,
"y": 172.,
"z": 550.
},
"outer": {
"x": 250.,
"y": 250.,
"z": 550.
},
"position": {
"x": 1600.,
"y": 0.,
"z": 0.
},
"rotation": {
"x": 0.,
"y": 45. / 180. * math.pi,
"z": 0.
}
}
self.assertEqual(len(field_list), 3)
sorted_names = sorted([field["field_name"] for field in field_list])
self.assertEqual(sorted_names, ["D2", "Q1", "Q2"])
self._field_list_equal(field_list, test_q1)
for field in field_list:
field["scale_factor"] = 2.
main_window.lattice.set_fields(field_list)
field_list_2 = main_window.lattice.get_field_list()
test_q1["scale_factor"] = 2.
self._field_list_equal(field_list_2, test_q1)
def test_envelope_tool(self):
"""Test buttons in envelope_tool window work okay"""
main_window = envelope_tool.MainWindow()
self.assertEqual(main_window.beam_setup, None)
beam_button = main_window.window.get_frame("&Beam Setup", "button")
beam_button.Clicked()
self.assertNotEqual(main_window.beam_setup, None)
self.assertEqual(main_window.magnet_setup, None)
magnet_button = main_window.window.get_frame("&Magnet Setup", "button")
magnet_button.Clicked()
self.assertNotEqual(main_window.magnet_setup, None)
self.assertEqual(main_window.plot_setup, None)
magnet_button = main_window.window.get_frame("&Plot Setup", "button")
magnet_button.Clicked()
self.assertNotEqual(main_window.plot_setup, None)
exit_button = main_window.window.get_frame("E&xit", "button")
exit_button.Clicked()
time.sleep(1)
self.assertEqual(main_window.window.main_frame, None)
def test_lattice_run_lattice_ref(self):
"""Test lattice - run lattice"""
# try changing the field configuration and reference particle; check
# that the reference particle comes out okay
sys.argv += ["--configuration_file", TEST_DIR + "test_config.py"]
_lattice = envelope_tool.MainWindow().lattice
self.assertGreater(len(_lattice.ref_list), 1)
self.assertGreater(len(_lattice.ellipse_list), 1)
last_hit = _lattice.ref_list[-1]
field_list = _lattice.get_field_list()
for item in field_list:
item["scale_factor"] *= 2.
_lattice.set_fields(field_list)
hit_ref, ellipse_ref = _lattice.get_beam()
hit_ref['p'] = hit_ref['p'] * 2.
hit_ref.mass_shell_condition('energy')
_lattice.set_beam(hit_ref, ellipse_ref)
_lattice.run_lattice()
test_hit = _lattice.ref_list[-1]
for pos in ['x', 'y', 'z']:
self.assertAlmostEqual(last_hit[pos], test_hit[pos], 3)
for mom in ['px', 'py', 'pz']:
self.assertAlmostEqual(last_hit[mom] * 2., test_hit[mom], 1)