From c1f993b248973cf1df65f01ee85a6a33894b2641 Mon Sep 17 00:00:00 2001 From: Matheus Date: Fri, 17 Apr 2026 14:58:22 -0300 Subject: [PATCH 1/3] WIP (LOCO): allow fitting of specific quad families & starting from a previous LOCO fit - allows selection of quadrupoles families for performing the fit: currently the fitting does not seem to be working, under investigation (find orbit error during new parameters evaluation within the LOCO loop) - allows starting point of the fitting to be performed from an initial LOCO fit --- bin/sirius-script-si-loco-run-fitting.py | 57 ++++++++++++++++++++---- 1 file changed, 49 insertions(+), 8 deletions(-) diff --git a/bin/sirius-script-si-loco-run-fitting.py b/bin/sirius-script-si-loco-run-fitting.py index 7858430..23759ba 100755 --- a/bin/sirius-script-si-loco-run-fitting.py +++ b/bin/sirius-script-si-loco-run-fitting.py @@ -70,10 +70,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,12 +117,12 @@ 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): """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.quadrupoles_to_fit = quadfams2fit config.sextupoles_to_fit = None config.skew_quadrupoles_to_fit = config.famname_skewquadset.copy() fc2_idx = config.skew_quadrupoles_to_fit.index('FC2') @@ -127,7 +130,7 @@ def _configure_elements_to_fit(config): config.dipoles_to_fit = None -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: @@ -146,6 +149,11 @@ def _configure_coupling_and_gains(config): config.fit_gain_bpm = True config.fit_gain_corr = True + 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'] + def _configure_measurement_kicks(config): """Configures kicks used in the measurements.""" @@ -180,17 +188,19 @@ def create_loco_config( goal_tunes, deltakl_weight=DEFAULT_DELTAKL_WEIGHT, lambda_lm=DEFAULT_LAMBDA_LM, + quadfams2fit=None, + init_fit_data=None, ): """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) + _configure_coupling_and_gains(config, init_fit_data) _configure_measurement_kicks(config) _configure_girder_shifts(config) _configure_weights(config) @@ -214,12 +224,16 @@ def create_loco( goal_tunes=None, deltakl_weight=DEFAULT_DELTAKL_WEIGHT, lambda_lm=DEFAULT_LAMBDA_LM, + quadfams2fit=None, + init_fit_data=None, ): """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, ) if 'orbmat_name' in loco_setup: @@ -273,6 +287,8 @@ def run_and_save( lambda_lm=DEFAULT_LAMBDA_LM, goal_tunes=None, force_tunes=False, + quadfams2fit=None, + init_fit_data=None, ): """Runs the LOCO fitting and saves the results.""" setup = load_data(setup_name) @@ -293,6 +309,8 @@ 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, ) loco.run_fit(niter=nriters) if force_tunes: @@ -353,6 +371,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 +430,12 @@ 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( '-r', '--report', @@ -427,6 +456,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 @@ -453,6 +492,8 @@ 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, ) if args.report: From 3db5a120769b457b5389233529cca129537a30af Mon Sep 17 00:00:00 2001 From: Matheus Date: Mon, 27 Apr 2026 17:50:41 -0300 Subject: [PATCH 2/3] ENH (LOCO): support for ID local magnets fitting - new arg `id-name` forces fit to be carried out with quad magnets adjacent to the corresponding ID, as well as the ID's skew quads - if using a previous fitted model, will not fit gains and couplings --- bin/sirius-script-si-loco-run-fitting.py | 88 +++++++++++++++++++++--- 1 file changed, 78 insertions(+), 10 deletions(-) diff --git a/bin/sirius-script-si-loco-run-fitting.py b/bin/sirius-script-si-loco-run-fitting.py index 23759ba..95fecb4 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 @@ -117,17 +121,26 @@ def _configure_inversion_and_minimization(config, lambda_lm): config.fixed_lambda = False -def _configure_elements_to_fit(config, quadfams2fit=None): +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 = quadfams2fit - 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, init_fit_data=None): @@ -146,13 +159,18 @@ def _configure_coupling_and_gains(config, init_fit_data=None): config.fit_roll_bpm = False config.fit_skew_quadrupoles = False config.fit_energy_shift = False - config.fit_gain_bpm = True - config.fit_gain_corr = True - 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): @@ -184,12 +202,45 @@ 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, ): """Creates and configures the LOCO object.""" config = _initialize_config_and_model(init_fit_data) @@ -199,7 +250,7 @@ def create_loco_config( _configure_fitting_families(config) _configure_constraints(config, deltakl_weight) _configure_inversion_and_minimization(config, lambda_lm) - _configure_elements_to_fit(config, quadfams2fit) + _configure_elements_to_fit(config, quadfams2fit, id_name) _configure_coupling_and_gains(config, init_fit_data) _configure_measurement_kicks(config) _configure_girder_shifts(config) @@ -226,6 +277,7 @@ def create_loco( lambda_lm=DEFAULT_LAMBDA_LM, quadfams2fit=None, init_fit_data=None, + id_name=None, ): """Creates a LOCO object with the given setup.""" config = create_loco_config( @@ -234,6 +286,7 @@ def create_loco( lambda_lm=lambda_lm, quadfams2fit=quadfams2fit, init_fit_data=init_fit_data, + id_name=id_name, ) if 'orbmat_name' in loco_setup: @@ -289,6 +342,7 @@ def run_and_save( force_tunes=False, quadfams2fit=None, init_fit_data=None, + id_name=None, ): """Runs the LOCO fitting and saves the results.""" setup = load_data(setup_name) @@ -311,6 +365,7 @@ def run_and_save( lambda_lm=lambda_lm, quadfams2fit=quadfams2fit, init_fit_data=init_fit_data, + id_name=id_name, ) loco.run_fit(niter=nriters) if force_tunes: @@ -436,6 +491,16 @@ def main(): 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', @@ -481,6 +546,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, @@ -494,6 +561,7 @@ def main(): force_tunes=args.forcetunes, quadfams2fit=args.quadfams2fit, init_fit_data=init_fit_data, + id_name=id_name, ) if args.report: From caa9292e07617db8bdd6d9b3956ef4f36c6b0ba7 Mon Sep 17 00:00:00 2001 From: Matheus Date: Fri, 8 May 2026 13:09:16 -0300 Subject: [PATCH 3/3] EN (LOCO): allow skipping tune corrections - when starting from a previous loco fitted model, the tunes will not match measured ones (same old tune problem) - if tune is corrected, quadrupole changes associated with this correction may be difficult to isolate from the actual changes needed to fit the data starting fom the previous fit model - the new arg allows to bypass the tune correction --- bin/sirius-script-si-loco-run-fitting.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/bin/sirius-script-si-loco-run-fitting.py b/bin/sirius-script-si-loco-run-fitting.py index 95fecb4..6263e2c 100755 --- a/bin/sirius-script-si-loco-run-fitting.py +++ b/bin/sirius-script-si-loco-run-fitting.py @@ -241,6 +241,7 @@ def create_loco_config( quadfams2fit=None, init_fit_data=None, id_name=None, + no_tune_corr=False ): """Creates and configures the LOCO object.""" config = _initialize_config_and_model(init_fit_data) @@ -257,12 +258,13 @@ def create_loco_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 @@ -278,6 +280,7 @@ def create_loco( 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( @@ -287,6 +290,7 @@ def create_loco( quadfams2fit=quadfams2fit, init_fit_data=init_fit_data, id_name=id_name, + no_tune_corr=no_tune_corr ) if 'orbmat_name' in loco_setup: @@ -343,6 +347,7 @@ def run_and_save( 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) @@ -366,6 +371,7 @@ def run_and_save( 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: @@ -514,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() @@ -562,6 +573,7 @@ def main(): quadfams2fit=args.quadfams2fit, init_fit_data=init_fit_data, id_name=id_name, + no_tune_corr=args.no_tune_corr ) if args.report: