diff --git a/bin/sirius-script-si-loco-run-fitting.py b/bin/sirius-script-si-loco-run-fitting.py index 7858430..6263e2c 100755 --- a/bin/sirius-script-si-loco-run-fitting.py +++ b/bin/sirius-script-si-loco-run-fitting.py @@ -1,5 +1,9 @@ #!/usr/bin/env python-sirius """Script for running LOCO algorithm.""" +# TODO: when selecting particular quads to fit, the saved jacobian should +# be named accordingly and signal or encode the fams selection +# an user may attempt to re-use a previously calculated jacobian and run +# into issues. import os import sys @@ -70,10 +74,13 @@ def move_tunes(model, tunex_goal, tuney_goal): model.radiation_on = rad -def _initialize_config_and_model(): +def _initialize_config_and_model(init_fit_data=None): """Initializes LOCOConfigSI and the accelerator model.""" config = LOCOConfigSI() - config.model = si.create_accelerator() + if init_fit_data is None: + config.model = si.create_accelerator() + else: + config.model = init_fit_data['fit_model'] config.use6dtrack = True return config @@ -114,20 +121,29 @@ def _configure_inversion_and_minimization(config, lambda_lm): config.fixed_lambda = False -def _configure_elements_to_fit(config): +def _configure_elements_to_fit(config, quadfams2fit=None, id_name=None): """Configures which elements to include in the fit.""" config.fit_quadrupoles = True config.fit_sextupoles = False config.fit_dipoles = False - config.quadrupoles_to_fit = None - config.sextupoles_to_fit = None - config.skew_quadrupoles_to_fit = config.famname_skewquadset.copy() - fc2_idx = config.skew_quadrupoles_to_fit.index('FC2') - config.skew_quadrupoles_to_fit.pop(fc2_idx) config.dipoles_to_fit = None + config.sextupoles_to_fit = None + + if id_name is None: + config.quadrupoles_to_fit = quadfams2fit + config.skew_quadrupoles_to_fit = config.famname_skewquadset.copy() + fc2_idx = config.skew_quadrupoles_to_fit.index('FC2') + config.skew_quadrupoles_to_fit.pop(fc2_idx) + return + + config.quadrupoles_to_fit = [] + config.skew_quadrupoles_to_fit = [] + config.quad_indices_kl, config.skew_quad_indices_ksl = _select_id_quads( + id_name, config + ) -def _configure_coupling_and_gains(config): +def _configure_coupling_and_gains(config, init_fit_data=None): """Configures coupling, BPM, and corrector gains fitting.""" config.fit_dipoles_kick = False if config.use_offdiagonal: @@ -143,10 +159,20 @@ def _configure_coupling_and_gains(config): config.fit_roll_bpm = False config.fit_skew_quadrupoles = False config.fit_energy_shift = False + if init_fit_data is not None: + config.gain_bpm = init_fit_data['gain_bpm'] + config.roll_bpm = init_fit_data['roll_bpm'] + config.gain_corr = init_fit_data['gain_corr'] + config.fit_gain_bpm = False # frozen params if + config.fit_gain_corr = False # starting from previous fit + config.fit_roll_bpm = False + return + config.fit_gain_bpm = True config.fit_gain_corr = True + def _configure_measurement_kicks(config): """Configures kicks used in the measurements.""" config.delta_kickx_meas = DELTA_KICK_X_MEAS @@ -176,32 +202,69 @@ def _configure_svd(config): config.parallel = True +def _select_id_quads(id_name, config): + """.""" + from pymodels import si + + model = config.model + famdata = si.get_family_data(model) + + if id_name not in famdata.keys(): + raise ValueError(f'ID family {id_name} not found in the model!') + + id_idc = pyac.lattice.find_indices(model, 'fam_name', id_name)[0] + id_subsec_type = famdata[id_name]['subsection'][0][-1] # A, B or P + quad_fams = [ + fam for fam in config.famname_quadset if id_subsec_type in fam + ] + quad_indices_kl = [] + for fam in quad_fams: + idcs = get_idcs(fam, model, id_idc) + for i in idcs: + quad_indices_kl.append([i]) + quad_indices_ksl = get_idcs('IDQS', model, id_idc) + quad_indices_ksl = [[idx] for idx in quad_indices_ksl] + return quad_indices_kl, quad_indices_ksl + + +def get_idcs(fam_name, mod, idc_delta): + """.""" + idc = np.array(pyac.lattice.find_indices(mod, 'fam_name', fam_name)) + idx = np.argsort(np.abs(idc_delta - idc))[:2] + return idc[idx] + + def create_loco_config( goal_tunes, deltakl_weight=DEFAULT_DELTAKL_WEIGHT, lambda_lm=DEFAULT_LAMBDA_LM, + quadfams2fit=None, + init_fit_data=None, + id_name=None, + no_tune_corr=False ): """Creates and configures the LOCO object.""" - config = _initialize_config_and_model() + config = _initialize_config_and_model(init_fit_data) _configure_cavity_and_radiation(config) _configure_jacobian_elements(config) _configure_fitting_families(config) _configure_constraints(config, deltakl_weight) _configure_inversion_and_minimization(config, lambda_lm) - _configure_elements_to_fit(config) - _configure_coupling_and_gains(config) + _configure_elements_to_fit(config, quadfams2fit, id_name) + _configure_coupling_and_gains(config, init_fit_data) _configure_measurement_kicks(config) _configure_girder_shifts(config) _configure_weights(config) _configure_svd(config) - print('--- changing tunes for nominal model...') - move_tunes( - config.model, - TUNE_X_INTEGER + goal_tunes[0], - TUNE_Y_INTEGER + goal_tunes[1], - ) + if not no_tune_corr: + print('--- changing tunes for nominal model...') + move_tunes( + config.model, + TUNE_X_INTEGER + goal_tunes[0], + TUNE_Y_INTEGER + goal_tunes[1], + ) config.update() return config @@ -214,12 +277,20 @@ def create_loco( goal_tunes=None, deltakl_weight=DEFAULT_DELTAKL_WEIGHT, lambda_lm=DEFAULT_LAMBDA_LM, + quadfams2fit=None, + init_fit_data=None, + id_name=None, + no_tune_corr=False ): """Creates a LOCO object with the given setup.""" config = create_loco_config( goal_tunes=goal_tunes, deltakl_weight=deltakl_weight, lambda_lm=lambda_lm, + quadfams2fit=quadfams2fit, + init_fit_data=init_fit_data, + id_name=id_name, + no_tune_corr=no_tune_corr ) if 'orbmat_name' in loco_setup: @@ -273,6 +344,10 @@ def run_and_save( lambda_lm=DEFAULT_LAMBDA_LM, goal_tunes=None, force_tunes=False, + quadfams2fit=None, + init_fit_data=None, + id_name=None, + no_tune_corr=False ): """Runs the LOCO fitting and saves the results.""" setup = load_data(setup_name) @@ -293,6 +368,10 @@ def run_and_save( goal_tunes=[tunex_goal, tuney_goal], deltakl_weight=deltakl_weight, lambda_lm=lambda_lm, + quadfams2fit=quadfams2fit, + init_fit_data=init_fit_data, + id_name=id_name, + no_tune_corr=no_tune_corr ) loco.run_fit(niter=nriters) if force_tunes: @@ -353,6 +432,11 @@ def main(): type=str, help='Name of the LOCO setup file (.pickle)', ) + parser.add_argument( + '--filename-init-fit', + type=str, + help='File name of a previous LOCO fit.', + ) parser.add_argument( '-f', '--folder', @@ -407,6 +491,22 @@ def main(): default=DEFAULT_LAMBDA_LM, help='Lambda parameter for LevenbergMarquardt. Default is 0.001', ) + parser.add_argument( + '--quadfams2fit', + nargs='+', + help='Names of quadrupoles families to include in the fitting. ' + 'Should be a space separeted list of families names, e.g.: QFP1 Q1...', + ) + parser.add_argument( + '--id-name', + type=str, + help=""" + ID name to perform fitting with local normal & skew quads. Should be + the name of the ID family in the model, e.g. 'DELTA52', 'UE44', ... + If provided, `--quadfams2fit` will be overridden to include only + the quadrupoles adjacent to the specified ID. + """, + ) parser.add_argument( '-r', '--report', @@ -420,6 +520,11 @@ def main(): help='Cleanup .png files. ' 'Default: False, set to True if flag is given.', ) + parser.add_argument( + '--no-tune-corr', + action='store_true', + help='Do not change the model tunes to match the measured tunes.' + ) args = parser.parse_args() @@ -427,6 +532,16 @@ def main(): if not os.path.isfile(fname_setup): raise ValueError(f'LOCO setup {fname_setup} not in current directory!') + if args.filename_init_fit: + fname_init_fit = args.filename_init_fit + if not os.path.isfile(fname_init_fit): + raise ValueError( + f'Initial fit file {fname_init_fit} not in current directory!' + ) + init_fit_data = load_data(fname_init_fit) + else: + init_fit_data = None + folder = args.folder folder_jac = args.jacobianfolder load_jac = True if folder_jac else False @@ -442,6 +557,8 @@ def main(): fname_fit = 'fitting_' + os.path.splitext(fname_setup)[0] fname_fit_path = os.path.join(folder, fname_fit) + id_name = args.id_name + run_and_save( setup_name=fname_setup, file_name=fname_fit_path, @@ -453,6 +570,10 @@ def main(): deltakl_weight=args.deltakl_weight, lambda_lm=args.lambda_lm, force_tunes=args.forcetunes, + quadfams2fit=args.quadfams2fit, + init_fit_data=init_fit_data, + id_name=id_name, + no_tune_corr=args.no_tune_corr ) if args.report: