def test_measurement_doesnt_write_if_not_allowed(_test_df_just_names,
_test_dir):
meas = OpticsMeasurement(_test_dir)
meas.allow_write = False
meas.beta_x = (_test_df_just_names, "_free")
assert not os.path.isfile(
os.path.join(CURRENT_DIR, "_test", "getbetax_free.out"))
def test_measurement_writes_double_plane(_test_df_just_names, _test_dir):
meas = OpticsMeasurement(_test_dir)
meas.allow_write = True
meas.beta_x = (_test_df_just_names, "_free")
meas.beta["y"] = (_test_df_just_names, "")
assert os.path.isfile(
os.path.join(CURRENT_DIR, "_test", "getbetax_free.out"))
assert os.path.isfile(os.path.join(CURRENT_DIR, "_test", "getbetay.out"))
def test_measurement_writes_double_plane(_test_df_just_names, _test_dir):
meas = OpticsMeasurement(_test_dir)
meas.allow_write = True
meas.beta_x = (_test_df_just_names, "_free")
meas.beta["y"] = (_test_df_just_names, "")
assert os.path.isfile(
os.path.join(CURRENT_DIR, "_test", "getbetax_free.out"))
assert os.path.isfile(
os.path.join(CURRENT_DIR, "_test", "getbetay.out"))
def test_measurement_maybe_doesnt_call_on_error(_meas_dir_empty):
meas = OpticsMeasurement(_meas_dir_empty)
def assert_false(args):
assert False
meas.maybe_call.beta_x(assert_false, "whatever")
def test_measurement_empty_dir(_meas_dir_empty):
meas = OpticsMeasurement(_meas_dir_empty)
with pytest.raises(IOError):
meas.beta_x
with pytest.raises(IOError):
meas.beta["x"]
with pytest.raises(IOError):
meas.coupling
def test_measurement_maybe_calls_on_success(_meas_dir):
meas = OpticsMeasurement(_meas_dir)
def raise_to_be_catched(tfs_file):
tfs_file.BETX # Check that the columns can be accessed
raise _KnownError()
with pytest.raises(_KnownError):
meas.maybe_call.beta_x(raise_to_be_catched)
def test_measurement_maybe_more_args(_meas_dir):
meas = OpticsMeasurement(_meas_dir)
def test_funct(tfs_file, number, kwnumber=None):
tfs_file.BETY # Check that the columns can be accessed
return number, kwnumber
rnumber, rkwnumber = meas.maybe_call.beta["y"](test_funct, 5, 7)
assert rnumber == 5
assert rkwnumber == 7
def getdiff(meas_path=None, beta_file_name="getbeta"):
""" Calculates the differences between measurement, corrected and uncorrected model.
After running madx and creating the model with (twiss_cor) and without (twiss_no)
corrections, run this functions to get tfs-files with the differences between measurements
and models.
Args:
meas_path (str): Path to the measurement folder.
Needs to contain twiss_cor.dat and twiss_no.dat.
"""
if meas_path is None:
meas_path = sys.argv[1]
LOG.debug("Started 'getdiff' for measurment dir '{:s}'".format(meas_path))
if not isdir(meas_path):
raise IOError("No valid measurement directory:" + meas_path)
corrected_model_path = join(meas_path, TWISS_CORRECTED)
uncorrected_model_path = join(meas_path, TWISS_NOT_CORRECTED)
meas = OpticsMeasurement(meas_path)
twiss_cor = read_tfs(corrected_model_path).set_index('NAME')
twiss_no = read_tfs(uncorrected_model_path).set_index('NAME')
coup_cor = TwissOptics(twiss_cor,
quick_init=True).get_coupling(method='cmatrix')
coup_no = TwissOptics(twiss_no,
quick_init=True).get_coupling(method='cmatrix')
model = pd.merge(twiss_cor,
twiss_no,
how='inner',
left_index=True,
right_index=True,
suffixes=('_c', '_n'))
coupling_model = pd.merge(coup_cor,
coup_no,
how='inner',
left_index=True,
right_index=True,
suffixes=('_c', '_n'))
#coupling_model['NAME'] = coupling_model.index.values
for plane in ['x', 'y']:
_write_betabeat_diff_file(meas_path, meas, model, plane,
beta_file_name)
_write_phase_diff_file(meas_path, meas, model, plane)
_write_disp_diff_file(meas_path, meas, model, plane)
_write_closed_orbit_diff_file(meas_path, meas, model, plane)
_write_coupling_diff_file(meas_path, meas, coupling_model)
_write_norm_disp_diff_file(meas_path, meas, model)
_write_chromatic_coupling_files(meas_path, corrected_model_path)
_write_betastar_diff_file(meas_path, meas, twiss_cor, twiss_no)
LOG.debug("Finished 'getdiff'.")
def _get_measurment_data(keys, meas_dir, beta_file_name, w_dict):
""" Retruns a dictionary full of get_llm data """
measurement = {}
filtered_keys = [k for k in keys if w_dict[k] != 0]
getllm_data = OpticsMeasurement(meas_dir)
for key in filtered_keys:
if key == "MUX":
measurement['MUX'] = getllm_data.phase_x
elif key == 'MUY':
measurement['MUY'] = getllm_data.phase_y
elif key == "DX":
measurement['DX'] = getllm_data.disp_x
elif key == "DY":
measurement['DY'] = getllm_data.disp_y
elif key == "NDX":
measurement['NDX'] = getllm_data.norm_disp
elif key in ('F1001R', 'F1001I', 'F1010R', 'F1010I'):
measurement[key] = getllm_data.coupling
elif key == "Q":
measurement["Q"] = pd.DataFrame(
{
# Just fractional tunes:
'VALUE':
np.remainder(
[getllm_data.phase_x['Q1'], getllm_data.phase_y['Q2']],
[1, 1]),
# TODO measured errors not in the file
'ERROR':
np.array([0.001, 0.001])
},
index=['Q1', 'Q2'])
else:
# a beta key
if beta_file_name == "getbeta":
if key in ("BBX", "BETX"):
measurement[key] = getllm_data.beta_x
elif key in ("BBY", "BETY"):
measurement[key] = getllm_data.beta_y
elif beta_file_name == "getampbeta":
if key in ("BBX", "BETX"):
measurement[key] = getllm_data.amp_beta_x
elif key in ("BBY", "BETY"):
measurement[key] = getllm_data.amp_beta_y
elif beta_file_name == "getkmodbeta":
if key in ("BBX", "BETX"):
measurement[key] = getllm_data.kmod_beta_x
elif key in ("BBY", "BETY"):
measurement[key] = getllm_data.kmod_beta_y
return filtered_keys, measurement
def segment_by_segment(accel_cls, options):
"""
TODO
"""
segments = _parse_segments(options.segments)
elements = _parse_elements(options.elements)
if not segments and not elements:
raise SbsDefinitionError("No segments or elements provided in the input.")
if _there_are_duplicated_names(segments, elements):
raise SbsDefinitionError("Duplicated names in segments and elements.")
model = tfs_pandas.read_tfs(options.model).set_index("NAME", drop=False)
meas = OpticsMeasurement(options.measurement)
elem_segments = [Segment.init_from_element(name) for name in elements]
for segment in elem_segments + segments:
propagable = run_for_segment(accel_cls, segment, model, meas,
options.optics, options.output)
write_beatings(segment, propagable, options.output)
def test_measurement_load_norm_dips(_meas_dir):
meas = OpticsMeasurement(_meas_dir)
meas.norm_disp.NDX
def test_measurement_load_no_plane(_meas_dir):
meas = OpticsMeasurement(_meas_dir)
meas.coupling.F1001W
def test_measurement_load_normal(_meas_dir):
meas = OpticsMeasurement(_meas_dir)
meas.disp_x.DX
meas.disp_y.DY
def test_measurement_load_free2(_meas_dir):
meas = OpticsMeasurement(_meas_dir)
meas.phase_x.PHASEX
meas.phase_y.PHASEY
def test_measurement_load_free(_meas_dir):
meas = OpticsMeasurement(_meas_dir)
meas.beta["x"].BETX
meas.beta_y.BETY
def test_measurement_doesnt_write_if_not_allowed(_test_df_just_names, _test_dir):
meas = OpticsMeasurement(_test_dir)
meas.allow_write = False
meas.beta_x = (_test_df_just_names, "_free")
assert not os.path.isfile(
os.path.join(CURRENT_DIR, "_test", "getbetax_free.out"))