Add support to LOCO fitting with ID's local quads and skew quads

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,13 +520,28 @@ 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()
 
     fname_setup = args.filename_setup
     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: