def _calc_tune_response(self):
"""Response vectors for Tune."""
LOG.debug("Calculate Tune Response Matrix")
with timeit(lambda t: LOG.debug(f" Time needed: {t} s")):
tw = self._twiss
k1_el = self._elements_in["K1L"]
if len(k1_el) > 0:
dtune = dict.fromkeys(PLANES)
dtune["X"] = 1 / (
4 * np.pi) * tw.loc[k1_el, [f"{BETA}X"]].transpose()
dtune["X"].index = ["DQX"]
dtune["Y"] = -1 / (
4 * np.pi) * tw.loc[k1_el, [f"{BETA}Y"]].transpose()
dtune["Y"].index = ["DQY"]
else:
LOG.debug(
" No 'K1L' variables found. Tune Response will be empty.")
dtune = {
"X": pd.DataFrame(None, index=["DQX"]),
"Y": pd.DataFrame(None, index=["DQY"])
}
return dtune
def convert_old_directory_to_new(opt: DotDict) -> None:
"""
Looks in the provided directory for expected ``BetaBeat.src`` output files, converts it to the output
format used by ``omc3`` and write them to the new location.
Args:
opt (DotDict): The entrypoint parameters parsed from the command line.
"""
inputdir_path = Path(opt.inputdir)
outputdir_path = Path(opt.outputdir)
suffix = opt.suffix
LOGGER.info(
f"Converting old BetaBeat.src outputs at '{inputdir_path.absolute()}' to omc3 format"
)
with contexts.timeit(lambda spanned: LOGGER.info(
f"Total time for conversion: {spanned}s")):
for plane in PLANES:
convert_old_beta_from_amplitude(inputdir_path, outputdir_path,
suffix, plane)
convert_old_beta_from_phase(inputdir_path, outputdir_path, suffix,
plane)
convert_old_phase(inputdir_path, outputdir_path, suffix, plane)
convert_old_total_phase(inputdir_path, outputdir_path, suffix,
plane)
# TODO phase vs phasetot inconsistent naming NAME S , first and second BPMs swapped locations
convert_old_closed_orbit(inputdir_path, outputdir_path, plane)
convert_old_dispersion(inputdir_path, outputdir_path, plane)
convert_old_coupling(inputdir_path, outputdir_path, suffix)
convert_old_normalised_dispersion(inputdir_path, outputdir_path, "X")
def _calc_beta_response(self):
"""Response Matrix for delta beta."""
LOG.debug("Calculate Beta Response Matrix")
with timeit(lambda t: LOG.debug(f" Time needed: {t} s")):
tw = self._twiss
adv = self._phase_advances
el_out = self._elements_out
k1_el = self._elements_in["K1L"]
dbeta = dict.fromkeys(PLANES)
for plane in PLANES:
col_beta = f"{BETA}{plane}"
q = tw.Q1 if plane == "X" else tw.Q2
coeff_sign = -1 if plane == "X" else 1
pi2tau = 2 * np.pi * tau(adv[plane].loc[k1_el, el_out], q)
dbeta[plane] = pd.DataFrame(
tw.loc[el_out, col_beta].to_numpy()[None, :] *
tw.loc[k1_el, col_beta].to_numpy()[:, None] *
np.cos(2 * pi2tau.to_numpy()) *
(coeff_sign / (2 * np.sin(2 * np.pi * q))),
index=k1_el,
columns=el_out,
).transpose()
return dbeta
def _closed_orbit_analysis(bpm_data, model, bpm_res):
lin_frame = pd.DataFrame(index=bpm_data.index.to_numpy(),
data=OrderedDict([("NAME", bpm_data.index.to_numpy()),
("S", np.arange(bpm_data.index.size) if model is None
else model.loc[bpm_data.index])]))
lin_frame['BPM_RES'] = 0.0 if bpm_res is None else bpm_res.loc[lin_frame.index]
with timeit(lambda spanned: LOGGER.debug(f"Time for orbit_analysis: {spanned}")):
lin_frame = _get_orbit_data(lin_frame, bpm_data)
return lin_frame, bpm_data.subtract(bpm_data.mean(axis=1), axis=0)
def _measure_optics(lins, optics_opt):
if len(lins) == 0:
lins = optics_opt.files
inputs = measure_optics.InputFiles(lins, optics_opt)
iotools.create_dirs(optics_opt.outputdir)
calibrations = measure_optics.copy_calibration_files(
optics_opt.outputdir, optics_opt.calibrationdir)
inputs.calibrate(calibrations)
with timeit(lambda spanned: LOGGER.info(
f"Total time for optics measurements: {spanned}")):
measure_optics.measure_optics(inputs, optics_opt)
def run_per_bunch(tbt_data, harpy_input):
"""
Cleans data, analyses frequencies and searches resonances
Args:
tbt_data: single bunch TbtData
harpy_input: Analysis settings
Returns:
Dictionary of TfsDataFrames per plane
"""
model = None if harpy_input.model is None else tfs.read(harpy_input.model, index="NAME").loc[:, 'S']
bpm_datas, usvs, lins, bad_bpms = {}, {}, {}, {}
output_file_path = _get_output_path_without_suffix(harpy_input.outputdir, harpy_input.files)
for plane in PLANES:
bpm_data = _get_cut_tbt_matrix(tbt_data, harpy_input.turns, plane)
bpm_data = _scale_to_meters(bpm_data, harpy_input.unit)
bpm_data, usvs[plane], bad_bpms[plane], bpm_res = clean.clean(harpy_input, bpm_data, model)
lins[plane], bpm_datas[plane] = _closed_orbit_analysis(bpm_data, model, bpm_res)
tune_estimates = harpy_input.tunes if harpy_input.autotunes is None else frequency.estimate_tunes(
harpy_input, usvs if harpy_input.clean else
dict(X=clean.svd_decomposition(bpm_datas["X"], harpy_input.sing_val),
Y=clean.svd_decomposition(bpm_datas["Y"], harpy_input.sing_val)))
spectra = {}
for plane in PLANES:
with timeit(lambda spanned: LOGGER.debug(f"Time for harmonic_analysis: {spanned}")):
harpy_results, spectra[plane], bad_bpms_summaries = frequency.harpy_per_plane(
harpy_input, bpm_datas[plane], usvs[plane], tune_estimates, plane)
if "bpm_summary" in harpy_input.to_write:
bad_bpms[plane].extend(bad_bpms_summaries)
_write_bad_bpms(output_file_path, plane, bad_bpms[plane])
if "spectra" in harpy_input.to_write or "full_spectra" in harpy_input.to_write:
_write_spectrum(output_file_path, plane, spectra[plane])
lins[plane] = lins[plane].loc[harpy_results.index].join(harpy_results)
if harpy_input.is_free_kick:
lins[plane] = kicker.phase_correction(bpm_datas[plane], lins[plane], plane)
measured_tunes = [lins["X"]["TUNEX"].mean(), lins["Y"]["TUNEY"].mean(),
lins["X"]["TUNEZ"].mean() if tune_estimates[2] > 0 else 0]
for plane in PLANES:
lins[plane] = lins[plane].join(frequency.find_resonances(
measured_tunes, bpm_datas[plane].shape[1], plane, spectra[plane]))
lins[plane] = _add_calculated_phase_errors(lins[plane])
lins[plane] = _sync_phase(lins[plane], plane)
lins[plane] = _rescale_amps_to_main_line_and_compute_noise(lins[plane], plane)
lins[plane] = lins[plane].sort_values('S', axis=0, ascending=True)
lins[plane] = tfs.TfsDataFrame(lins[plane], headers=_compute_headers(lins[plane], tbt_data.date))
if "lin" in harpy_input.to_write:
_write_lin_tfs(output_file_path, plane, lins[plane])
return lins
def _calc_phase_response(self):
"""Response Matrix for delta DPhi.
This calculation could also be achieved by applying np.cumsum to the DataFrames of
_calc_phase_adv_response() (tested!), but _calc_phase_response() is about 4x faster.
"""
LOG.debug("Calculate Phase Response Matrix")
with timeit(lambda t: LOG.debug(f" Time needed: {t} s")):
tw = self._twiss
adv = self._phase_advances
k1_el = self._elements_in["K1L"]
el_out = self._elements_out
if len(k1_el) > 0:
dmu = dict.fromkeys(PLANES)
pi = pd.DataFrame(
tw["S"].to_numpy()[:, None] <
tw["S"].to_numpy()[None, :], # pi(i,j) = s(i) < s(j)
index=tw.index,
columns=tw.index,
dtype=int)
pi_term = pi.loc[k1_el, el_out].to_numpy()
for plane in PLANES:
col_beta = f"{BETA}{plane}"
q = tw.Q1 if plane == "X" else tw.Q2
coeff_sign = 1 if plane == "X" else -1
pi2tau = 2 * np.pi * tau(
adv[plane].loc[k1_el, [DUMMY_ID] + el_out], q)
brackets = 2 * pi_term + (
(np.sin(2 * pi2tau.loc[:, el_out].to_numpy()) - np.sin(
2 * pi2tau.loc[:, DUMMY_ID].to_numpy()[:, None])) /
np.sin(2 * np.pi * q))
dmu[plane] = pd.DataFrame(
tw.loc[k1_el, col_beta].to_numpy()[:, None] *
brackets * (coeff_sign / (8 * np.pi)),
index=k1_el,
columns=el_out,
).transpose()
else:
LOG.debug(
" No 'K1L' variables found. Phase Response will be empty."
)
dmu = {
"X": pd.DataFrame(None, index=el_out),
"Y": pd.DataFrame(None, index=el_out)
}
return dmu
def _run_harpy(harpy_options):
iotools.create_dirs(harpy_options.outputdir)
with timeit(
lambda spanned: LOGGER.info(f"Total time for Harpy: {spanned}")):
lins = []
all_options = _replicate_harpy_options_per_file(harpy_options)
tbt_datas = [(tbt.read_tbt(option.files,
datatype=option.tbt_datatype), option)
for option in all_options]
for tbt_data, option in tbt_datas:
lins.extend([
handler.run_per_bunch(bunch_data, bunch_options)
for bunch_data, bunch_options in _multibunch(tbt_data, option)
])
return lins
def clean(harpy_input, bpm_data, model):
"""
Cleans BPM TbT matrix: removes BPMs not present in the model and based on specified cuts.
Also cleans the noise using singular value decomposition.
Args:
harpy_input: The input object containing the analysis settings
bpm_data: DataFrame of BPM TbT matrix indexed by BPM names
model: model containing BPMs longitudinal locations indexed by BPM names
Returns:
Clean BPM matrix, its decomposition, bad BPMs summary and estimated BPM resolutions
"""
bpm_data, bpms_not_in_model = _get_only_model_bpms(bpm_data, model)
if bpm_data.empty:
raise AssertionError("Check BPMs names! None of the BPMs were found in the model!")
if not harpy_input.clean:
return bpm_data, None, bpms_not_in_model, None
with timeit(lambda spanned: LOGGER.debug(f"Time for filtering: {spanned}")):
bpm_data, bad_bpms_clean = _cut_cleaning(harpy_input, bpm_data, model)
with timeit(lambda spanned: LOGGER.debug(f"Time for SVD clean: {spanned}")):
bpm_data, bpm_res, bad_bpms_svd, usv = _svd_clean(bpm_data, harpy_input)
all_bad_bpms = bpms_not_in_model + bad_bpms_clean + bad_bpms_svd
return bpm_data, usv, all_bad_bpms, bpm_res
def test_measure_optics(tmp_path, input_data, lin_slice, compensation,
coupling_method, range_of_bpms, three_bpm_method,
second_order_disp):
data = input_data["free" if compensation == 'none' else "driven"]
lins, optics_opt = data['lins'], data['optics_opt']
optics_opt.update(
outputdir=tmp_path,
compensation=compensation,
coupling_method=coupling_method,
range_of_bpms=range_of_bpms,
three_bpm_method=three_bpm_method,
second_order_disp=second_order_disp,
chromatic_beating=lin_slice == slice(None, 7),
)
inputs = measure_optics.InputFiles(lins[lin_slice], optics_opt)
with timeit(lambda spanned: LOG.debug(
f"\nTotal time for optics measurements: {spanned}")):
measure_optics.measure_optics(inputs, optics_opt)
evaluate_accuracy(optics_opt.outputdir, LIMITS)
def create_response(
accel_inst: Accelerator,
vars_categories: Sequence[str],
optics_params: List[str],
) -> dict:
""" Wrapper to create response via TwissResponse """
LOG.debug("Creating response via TwissResponse.")
varmap_path = check_varmap_file(accel_inst, vars_categories)
with timeit(
lambda t: LOG.debug(f"Total time getting TwissResponse: {t} s")):
tr = TwissResponse(accel_inst, vars_categories, varmap_path)
response: dict = tr.get_response_for(optics_params)
if not any([resp.size for resp in response.values()]):
raise ValueError(
"Responses are all empty. "
f"Are variables {tr.get_variable_names()} correct for '{optics_params}'?"
)
return response
def __init__(self,
accel_inst,
variable_categories,
varmap_or_path,
at_elements="bpms"):
LOG.debug("Initializing TwissResponse.")
with timeit(lambda t: LOG.debug(
f" Time initializing TwissResponse: {t} s")):
# Get input
self._twiss = self._get_model_twiss(accel_inst)
self._beam_sign = 1 if getattr(accel_inst, "beam", 1) == 1 else -1
self._variables = accel_inst.get_variables(
classes=variable_categories)
self._var_to_el = self._get_variable_mapping(varmap_or_path)
self._elements_in = self._get_input_elements()
self._elements_out = self._get_output_elements(at_elements)
# calculate all phase advances
self._phase_advances = get_phase_advances(self._twiss)
# All responses are calcluated as needed, see getters below!
# slots for response matrices
self._beta = None
self._dispersion = None
self._phase = None
self._phase_adv = None
self._tune = None
self._coupling = None
self._beta_beat = None
self._norm_dispersion = None
# slots for mapped response matrices
self._coupling_mapped = None
self._beta_mapped = None
self._dispersion_mapped = None
self._phase_mapped = None
self._phase_adv_mapped = None
self._tune_mapped = None
self._beta_beat_mapped = None
self._norm_dispersion_mapped = None
def _add_coupling(dict_of_tfs: Dict[str, tfs.TfsDataFrame]) -> Dict[str, tfs.TfsDataFrame]:
"""
For each TfsDataFrame in the input dictionary, computes the coupling RDTs and adds a column for
the real and imaginary parts of the computed coupling RDTs. Returns a copy of the input dictionary with
the aforementioned computed columns added for each TfsDataFrame.
Args:
dict_of_tfs (Dict[str, tfs.TfsDataFrame]): dictionary of Twiss dataframes.
Returns:
An identical dictionary of Twiss dataframes, with the computed columns added.
"""
result_dict_of_tfs = copy.deepcopy(dict_of_tfs)
with timeit(lambda elapsed: LOG.debug(f" Time adding coupling: {elapsed} s")):
for var, tfs_dframe in result_dict_of_tfs.items(): # already copies, so it's safe to act on them
coupling_rdts_df = coupling_via_cmatrix(tfs_dframe)
tfs_dframe[f"{F1001}R"] = np.real(coupling_rdts_df[f"{F1001}"]).astype(np.float64)
tfs_dframe[f"{F1001}I"] = np.imag(coupling_rdts_df[f"{F1001}"]).astype(np.float64)
tfs_dframe[f"{F1010}R"] = np.real(coupling_rdts_df[f"{F1010}"]).astype(np.float64)
tfs_dframe[f"{F1010}I"] = np.imag(coupling_rdts_df[f"{F1010}"]).astype(np.float64)
return result_dict_of_tfs
def evaluate_for_variables(
accel_inst: Accelerator,
variable_categories,
order: int = 4,
num_proc: int = multiprocessing.cpu_count(),
temp_dir: Path = None
) -> dict:
""" Generate a dictionary containing response matrices for
beta, phase, dispersion, tune and coupling and saves it to a file.
Args:
accel_inst (Accelerator): Accelerator Instance.
variable_categories (list): Categories of the variables/knobs to use. (from .json)
order (int or tuple): Max or [min, max] of K-value order to use.
num_proc (int): Number of processes to use in parallel.
temp_dir (Path): temporary directory. If ``None``, uses model_dir.
"""
LOG.debug("Generating Fullresponse via Mad-X.")
with timeit(lambda elapsed: LOG.debug(f" Total time generating fullresponse: {elapsed}s")):
if not temp_dir:
temp_dir = Path(accel_inst.model_dir)
temp_dir.mkdir(parents=True, exist_ok=True)
variables = accel_inst.get_variables(classes=variable_categories)
if len(variables) == 0:
raise ValueError("No variables found! Make sure your categories are valid!")
num_proc = num_proc if len(variables) > num_proc else len(variables)
process_pool = multiprocessing.Pool(processes=num_proc)
k_values = _get_orders(order)
try:
_generate_madx_jobs(accel_inst, variables, k_values, num_proc, temp_dir)
_call_madx(process_pool, temp_dir, num_proc)
mapping = _load_madx_results(variables, k_values, process_pool, temp_dir)
finally:
_clean_up(variables, temp_dir, num_proc)
return mapping
def get_phase_advances(twiss_df: pd.DataFrame) -> Dict[str, pd.DataFrame]:
"""
Calculate phase advances between all elements
Returns:
Matrices similar to DPhi(i,j) = Phi(j) - Phi(i)
"""
LOG.debug("Calculating Phase Advances:")
phase_advance_dict = dict.fromkeys(["X", "Y"])
with timeit(lambda elapsed: LOG.debug(
f" Phase Advances calculated in {elapsed}s")):
for plane in PLANES:
colmn_phase = f"{PHASE_ADV}{plane}"
phases_mdl = twiss_df.loc[twiss_df.index, colmn_phase]
# Same convention as in [1]: DAdv(i,j) = Phi(j) - Phi(i)
phase_advances = pd.DataFrame((phases_mdl.to_numpy()[None, :] -
phases_mdl.to_numpy()[:, None]),
index=twiss_df.index,
columns=twiss_df.index)
# Do not calculate dphi and tau here.
# only slices of phase_advances as otherwise super slow
phase_advance_dict[plane] = phase_advances
return phase_advance_dict
def create_fullresponse(
accel_inst: Accelerator,
variable_categories: Sequence[str],
delta_k: float = 2e-5,
num_proc: int = multiprocessing.cpu_count(),
temp_dir: Path = None
) -> Dict[str, pd.DataFrame]:
""" Generate a dictionary containing response matrices for
beta, phase, dispersion, tune and coupling and saves it to a file.
Args:
accel_inst : Accelerator Instance.
variable_categories (list): Categories of the variables/knobs to use. (from .json)
delta_k (float): delta K1L to be applied to quads for sensitivity matrix
num_proc (int): Number of processes to use in parallel.
temp_dir (str): temporary directory. If ``None``, uses folder of original_jobfile.
"""
LOG.debug("Generating Fullresponse via Mad-X.")
with timeit(lambda t: LOG.debug(f" Total time generating fullresponse: {t} s")):
if not temp_dir:
temp_dir = Path(accel_inst.model_dir)
temp_dir.mkdir(parents=True, exist_ok=True)
variables = accel_inst.get_variables(classes=variable_categories)
if len(variables) == 0:
raise ValueError("No variables found! Make sure your categories are valid!")
num_proc = num_proc if len(variables) > num_proc else len(variables)
process_pool = multiprocessing.Pool(processes=num_proc)
incr_dict = _generate_madx_jobs(accel_inst, variables, delta_k, num_proc, temp_dir)
_call_madx(process_pool, temp_dir, num_proc)
_clean_up(temp_dir, num_proc)
var_to_twiss = _load_madx_results(variables, process_pool, incr_dict, temp_dir)
fullresponse = _create_fullresponse_from_dict(var_to_twiss)
return fullresponse
def _calc_coupling_response(self):
"""Response Matrix for coupling."""
LOG.debug("Calculate Coupling Matrix")
with timeit(lambda t: LOG.debug(f" Time needed: {t} s")):
tw = self._twiss
adv = self._phase_advances
el_out = self._elements_out
k1s_el = self._elements_in["K1SL"]
dcoupl = dict.fromkeys(["1001", "1010"])
i2pi = 2j * np.pi
phx = dphi(adv["X"].loc[k1s_el, el_out], tw.Q1).to_numpy()
phy = dphi(adv["Y"].loc[k1s_el, el_out], tw.Q2).to_numpy()
bet_term = np.sqrt(tw.loc[k1s_el, f"{BETA}X"].to_numpy() *
tw.loc[k1s_el, f"{BETA}Y"].to_numpy())
for coupling_rdt in ["1001", "1010"]:
phs_sign = -1 if coupling_rdt == "1001" else 1
dcoupl[coupling_rdt] = pd.DataFrame(
bet_term[:, None] * np.exp(i2pi * (phx + phs_sign * phy)) /
(4 * (1 - np.exp(i2pi * (tw.Q1 + phs_sign * tw.Q2)))),
index=k1s_el,
columns=el_out).transpose()
return dcoupl
def _run_evaluate_and_clean_up(inputs, optics_opt):
with timeit(lambda spanned: print(f"\nTotal time for optics measurements: {spanned}")):
measure_optics.measure_optics(inputs, optics_opt)
evaluate_accuracy(optics_opt.outputdir)
_clean_up(optics_opt.outputdir)
def get_response_for(self, observables=None) -> dict: # Dict[str, ???]
""" Calculates and returns only desired response matrices """
# calling functions for the getters to call functions only if needed
def caller(func, plane):
return func()[plane]
def disp_caller(func, plane):
disp = func()
return response_add(
*[disp[k] for k in disp.keys() if k.startswith(plane)])
def tune_caller(func, _unused):
tune = func()
res = tune["X"].append(tune["Y"])
res.index = [f"{TUNE}1", f"{TUNE}2"]
return res
def couple_caller(func, plane):
# apply() converts empty DataFrames to Series! Cast them back.
# Also: take care of minus-sign convention!
sign = -1 if plane[-1] == "R" else 1
part_func = np.real if plane[-1] == "R" else np.imag
return sign * pd.DataFrame(
func()[plane[:-1]].apply(part_func).astype(np.float64))
# to avoid if-elif-elif-...
obs_map = {
f"{TUNE}": (tune_caller, self.get_tune, None),
f"{BETA}X": (caller, self.get_beta_beat, "X"),
f"{BETA}Y": (caller, self.get_beta_beat, "Y"),
f"{PHASE_ADV}X": (caller, self.get_phase, "X"),
f"{PHASE_ADV}Y": (caller, self.get_phase, "Y"),
f"{DISP}X": (disp_caller, self.get_dispersion, "X"),
f"{DISP}Y": (disp_caller, self.get_dispersion, "Y"),
f"{NORM_DISP}X": (disp_caller, self.get_norm_dispersion, "X"),
f"{NORM_DISP}Y": (disp_caller, self.get_norm_dispersion, "Y"),
f"{F1001}R": (couple_caller, self.get_coupling, "1001R"),
f"{F1001}I": (couple_caller, self.get_coupling, "1001I"),
f"{F1010}R": (couple_caller, self.get_coupling, "1010R"),
f"{F1010}I": (couple_caller, self.get_coupling, "1010I"),
}
if observables is None:
observables = list(obs_map.keys())
else:
# rename to PHASE to MU
observables = [
f"{PHASE_ADV}{key[-1]}" if key.startswith(PHASE) else key
for key in observables
]
LOG.debug(f"Calculating responses for {observables}.")
with timeit(
lambda t: LOG.debug(f"Total time getting responses: {t} s")):
response = dict.fromkeys(observables)
for key in observables:
# response[key] = obs_map[key][0](*obs_map[key][1:3])
# I still don't know why there is a sign-change between B1 and B2 (jdilly, 2021)
response[key] = self._beam_sign * obs_map[key][0](
*obs_map[key][1:3])
return response
def _calc_phase_advance_response(self):
"""Response Matrix for delta DPhi.
Reduced to only phase advances between consecutive elements,
as the 3D-Matrix of all elements exceeds memory space
(~11000^3 = 1331 Giga Elements)
--> w = j-1: DPhi(z,j) = DPhi(x, (j-1)->j)
"""
LOG.debug("Calculate Phase Advance Response Matrix")
with timeit(lambda elapsed: LOG.debug(f" Time needed: {elapsed} s")):
tw = self._twiss
adv = self._phase_advances
k1_el = self._elements_in["K1L"]
el_out_all = [
DUMMY_ID
] + self._elements_out # Add MU[XY] = 0.0 to the start
el_out = el_out_all[1:] # in these we are actually interested
el_out_mm = el_out_all[0:-1] # elements--
if len(k1_el) > 0:
dmu = dict.fromkeys(PLANES)
pi = pd.DataFrame(
tw[f"{S}"].to_numpy()[:, None] <
tw[f"{S}"].to_numpy()[None, :], # pi(i,j) = s(i) < s(j)
index=tw.index,
columns=tw.index,
dtype=int)
pi_term = (
pi.loc[k1_el, el_out].to_numpy() -
pi.loc[k1_el, el_out_mm].to_numpy() +
np.diag(pi.loc[el_out, el_out_mm].to_numpy())[None, :])
for plane in PLANES:
col_beta = f"{BETA}{plane}"
q = tw.Q1 if plane == "X" else tw.Q2
coeff_sign = 1 if plane == "X" else -1
pi2tau = 2 * np.pi * tau(adv[plane].loc[k1_el, el_out_all],
q)
brackets = 2 * pi_term + (
(np.sin(2 * pi2tau.loc[:, el_out].to_numpy()) -
np.sin(2 * pi2tau.loc[:, el_out_mm].to_numpy())) /
np.sin(2 * np.pi * q))
dmu[plane] = pd.DataFrame(
tw.loc[k1_el, col_beta].to_numpy()[:, None] *
brackets * (coeff_sign / (8 * np.pi)),
index=k1_el,
columns=el_out,
).transpose()
else:
LOG.debug(
" No 'K1L' variables found. Phase Response will be empty."
)
dmu = {
"X": pd.DataFrame(None, index=el_out),
"Y": pd.DataFrame(None, index=el_out)
}
return dmu
def _calc_norm_dispersion_response(self):
"""Response Matrix for delta normalized dispersion."""
LOG.debug("Calculate Normalized Dispersion Response Matrix")
with timeit(lambda t: LOG.debug(f" Time needed: {t} s")):
tw = self._twiss
adv = self._phase_advances
el_out = self._elements_out
els_in = self._elements_in
sign_map = {
"X": {
"K0L": 1,
"K1L": -1,
"K1SL": 1
},
"Y": {
"K0SL": -1,
"K1L": 1,
"K1SL": 1
},
}
col_disp_map = {
"X": {
"K1L": f"{DISP}X",
"K1SL": f"{DISP}Y"
},
"Y": {
"K1L": f"{DISP}Y",
"K1SL": f"{DISP}X"
},
}
sign_correct_term = {
"X": {
"K1L": 1
},
"Y": {
"K1L": -1
},
}
q_map = {"X": tw.Q1, "Y": tw.Q2}
disp_resp = dict.fromkeys([
"{p:s}_{t:s}".format(p=p, t=t) for p in sign_map
for t in sign_map[p]
])
for plane in sign_map:
q = q_map[plane]
col_beta = f"{BETA}{plane}"
el_types = sign_map[plane].keys()
els_per_type = [els_in[el_type] for el_type in el_types]
coeff = 1 / (2 * np.sin(np.pi * q))
coeff_corr = 1 / (4 * np.sin(2 * np.pi * q))
for el_in, el_type in zip(els_per_type, el_types):
coeff_sign = sign_map[plane][el_type]
out_str = "{p:s}_{t:s}".format(p=plane, t=el_type)
if len(el_in):
pi2tau = 2 * np.pi * tau(adv[plane].loc[el_in, el_out],
q)
beta_in = tw.loc[el_in, col_beta]
bet_term = np.sqrt(beta_in)
try:
col_disp = col_disp_map[plane][el_type]
except KeyError:
pass
else:
bet_term *= tw.loc[el_in, col_disp]
result = (coeff_sign * coeff * bet_term
).to_numpy()[:, None] * np.cos(pi2tau)
# correction term
try:
sign_corr = sign_correct_term[plane][el_type]
except KeyError:
pass
else:
norm_disp_corr = (
tw.loc[el_out, col_disp] /
np.sqrt(tw.loc[el_out, col_beta]))
result += (sign_corr * coeff_corr *
norm_disp_corr.to_numpy()[None, :] *
beta_in.to_numpy()[:, None] *
np.cos(2 * pi2tau))
disp_resp[out_str] = result.transpose()
else:
LOG.debug(
f" No '{el_type:s}' variables found. "
f"Normalized Dispersion Response '{out_str:s}' will be empty."
)
disp_resp[out_str] = pd.DataFrame(None, index=el_out)
return disp_resp
def _calc_dispersion_response(self):
"""Response Matrix for delta normalized dispersion."""
LOG.debug("Calculate Dispersion Response Matrix")
with timeit(lambda t: LOG.debug(f" Time needed: {t} s")):
tw = self._twiss
adv = self._phase_advances
el_out = self._elements_out
els_in = self._elements_in
sign_map = {
"X": {
"K0L": 1,
"K1L": -1,
"K1SL": 1
},
"Y": {
"K0SL": -1,
"K1L": 1,
"K1SL": 1
},
}
col_disp_map = {
"X": {
"K1L": f"{DISP}X",
"K1SL": f"{DISP}Y"
},
"Y": {
"K1L": f"{DISP}Y",
"K1SL": f"{DISP}X"
},
}
q_map = {"X": tw.Q1, "Y": tw.Q2}
disp_resp = dict.fromkeys(
[f"{p}_{t}" for p in sign_map for t in sign_map[p]])
for plane in sign_map:
q = q_map[plane]
col_beta = f"{BETA}{plane}"
el_types = sign_map[plane].keys()
els_per_type = [els_in[el_type] for el_type in el_types]
coeff = np.sqrt(tw.loc[el_out, col_beta].to_numpy()) / (
2 * np.sin(np.pi * q))
for el_in, el_type in zip(els_per_type, el_types):
coeff_sign = sign_map[plane][el_type]
out_str = f"{plane}_{el_type}"
if len(el_in):
pi2tau = 2 * np.pi * tau(adv[plane].loc[el_in, el_out],
q)
bet_term = np.sqrt(tw.loc[el_in, col_beta])
try:
col_disp = col_disp_map[plane][el_type]
except KeyError:
pass
else:
bet_term *= tw.loc[el_in, col_disp]
disp_resp[out_str] = (coeff_sign * coeff[None, :] *
bet_term.to_numpy()[:, None] *
np.cos(pi2tau)).transpose()
else:
LOG.debug(
f" No '{el_type}' variables found. "
f"Dispersion Response '{out_str}' will be empty.")
disp_resp[out_str] = pd.DataFrame(None, index=el_out)
return disp_resp
