Arc corr and glob check fix

omc3/check_corrections.py

@@ -245,7 +245,7 @@
 from omc3.correction.response_twiss import PLANES
 from omc3.definitions.optics import (
     OpticsMeasurement, ColumnsAndLabels,
-    FILE_COLUMN_MAPPING, RDT_COLUMN_MAPPING, TUNE_COLUMN
+    FILE_COLUMN_MAPPING, RDT_COLUMN_MAPPING, TUNE_COLUMN, TOTAL_PHASE_NAME, MU_COLUMN
 )
 from omc3.global_correction import _get_default_values, CORRECTION_DEFAULTS, OPTICS_PARAMS_CHOICES
 from omc3.model import manager
@@ -475,8 +475,10 @@ def _create_model_and_write_diff_to_measurements(
     diff_columns = (
             list(OPTICS_PARAMS_CHOICES[:-4]) +
             [col for col in corr_model_elements.columns if col.startswith("F1")] +
-            list(PLANES)
+            list(PLANES) +
+            ['MUX', 'MUY']
     )
+
     diff_models = diff_twiss_parameters(corr_model_elements, accel_inst.model, parameters=diff_columns)
     LOG.debug(f"Differences to nominal model calculated.")
 
@@ -514,12 +516,17 @@ def _create_model_and_write_diff_to_measurements(
             LOG.debug(f"Checking correction for {attribute}")
             plane = filename[-1].upper()
             cols = colmap.set_plane(plane)
+            if filename[:-1] == TOTAL_PHASE_NAME:
+                model_cols = MU_COLUMN.set_plane(plane)
+            else:
+                model_cols = cols
+
             _create_check_columns(
                 measurement=measurement,
                 output_measurement=output_measurement,
                 diff_models=diff_models,
                 colmap_meas=cols,
-                colmap_model=cols,
+                colmap_model=model_cols,
                 attribute=attribute,
                 rms_mask=rms_mask,
             )
@@ -571,7 +578,7 @@ def _create_check_columns(measurement: OpticsMeasurement, output_measurement: Op
     diff = diff_models.loc[df.index, colmap_model.delta_column]
 
     df[colmap_meas.diff_correction_column] = diff
-    if colmap_meas.column == PHASE:
+    if colmap_meas.column[:-1] == PHASE:
         df[colmap_meas.expected_column] = pd.to_numeric(ang_diff(df[colmap_meas.delta_column], diff))  # assumes period 1
         df.headers[colmap_meas.delta_rms_header] =  circular_rms(df[colmap_meas.delta_column], period=1)
         df.headers[colmap_meas.expected_rms_header] = circular_rms(df[colmap_meas.expected_column], period=1)

omc3/correction/arc_by_arc.py

@@ -0,0 +1,100 @@
+import pickle
+import numpy as np
+import pandas as pd
+
+
+def identify_closest_arc_bpm_to_ip(ip, side, beam, bpms):
+    indices = range(1,15)
+    for ii in indices:
+        bpm = f'BPM.{ii}{side}{ip}.B{beam}'
+        if bpm in bpms:
+            return bpm
+
+
+def get_left_right_pair(arc, beam, bpms):
+    left_of_arc = identify_closest_arc_bpm_to_ip(int(arc[0]), 'R', beam, bpms)
+    right_of_arc = identify_closest_arc_bpm_to_ip(int(arc[1]), 'L', beam, bpms)
+    return [left_of_arc, right_of_arc]
+
+
+def get_arc_by_arc_bpm_pairs(meas_dict, opt):
+    bpms = meas_dict['PHASEX'].index
+    beam = bpms[0][-1]
+    bpm_pairs = {}
+    bpm_pairs_with_ips = {}
+    include_ips = False
+
+    arcs_to_cycle = ['81', '12', '23', '34', '45', '56', '67', '78']
+
+    for lhc_arc in arcs_to_cycle:
+        bpm_pairs[lhc_arc] = get_left_right_pair(lhc_arc, beam, bpms)
+
+    if include_ips == 'left': 
+        bpm_pairs_with_ips['81'] =  [bpm_pairs['78'][1], bpm_pairs['81'][1]]
+        bpm_pairs_with_ips['12'] =  [bpm_pairs['81'][1], bpm_pairs['12'][1]]
+        bpm_pairs_with_ips['23'] =  [bpm_pairs['12'][1], bpm_pairs['23'][1]]
+        bpm_pairs_with_ips['34'] =  [bpm_pairs['23'][1], bpm_pairs['34'][1]]
+        bpm_pairs_with_ips['45'] =  [bpm_pairs['34'][1], bpm_pairs['45'][1]]
+        bpm_pairs_with_ips['56'] =  [bpm_pairs['45'][1], bpm_pairs['56'][1]]
+        bpm_pairs_with_ips['67'] =  [bpm_pairs['56'][1], bpm_pairs['67'][1]]
+        bpm_pairs_with_ips['78'] =  [bpm_pairs['67'][1], bpm_pairs['78'][1]]
+        bpm_pairs = bpm_pairs_with_ips
+    elif include_ips == 'right': 
+        bpm_pairs_with_ips['81'] =  [bpm_pairs['78'][0], bpm_pairs['81'][0]]
+        bpm_pairs_with_ips['12'] =  [bpm_pairs['81'][0], bpm_pairs['12'][0]]
+        bpm_pairs_with_ips['23'] =  [bpm_pairs['12'][0], bpm_pairs['23'][0]]
+        bpm_pairs_with_ips['34'] =  [bpm_pairs['23'][0], bpm_pairs['34'][0]]
+        bpm_pairs_with_ips['45'] =  [bpm_pairs['34'][0], bpm_pairs['45'][0]]
+        bpm_pairs_with_ips['56'] =  [bpm_pairs['45'][0], bpm_pairs['56'][0]]
+        bpm_pairs_with_ips['67'] =  [bpm_pairs['56'][0], bpm_pairs['67'][0]]
+        bpm_pairs_with_ips['78'] =  [bpm_pairs['67'][0], bpm_pairs['78'][0]]
+        bpm_pairs = bpm_pairs_with_ips
+
+    return bpm_pairs
+
+
+def circular_sum_phase(phase_df, tune, bpm_pair, key):
+    idx_0 = phase_df[key].index.get_loc(bpm_pair[0])
+    idx_1 = phase_df[key].index.get_loc(bpm_pair[1])
+    if idx_0 > idx_1:
+        inverted_result = sum(phase_df[key][bpm_pair[1]:bpm_pair[0]])
+        result = tune - inverted_result
+    else:
+        result = sum(phase_df[key][bpm_pair[0]:bpm_pair[1]])
+    return result
+
+
+def circular_sum_phase_error(phase_df, bpm_pair):
+    idx_0 = phase_df['ERROR'].index.get_loc(bpm_pair[0])
+    idx_1 = phase_df['ERROR'].index.get_loc(bpm_pair[1])
+    if idx_0 > idx_1:
+        selection = pd.concat([phase_df['ERROR'].loc[:bpm_pair[1]], phase_df['ERROR'].loc[bpm_pair[0]:]])
+        result = np.sqrt(np.sum(selection**2))
+    else:
+        result = np.sqrt(np.sum(phase_df['ERROR'][bpm_pair[0]:bpm_pair[1]]**2))
+    return result
+
+def get_arc_phases(bpm_pairs, meas_dict, tune, plane):
+    arc_meas = []
+    for arc, bpm_pair in bpm_pairs.items():
+        results = {}
+        results['NAME'] = bpm_pair[0]
+        results['NAME2'] = bpm_pair[1]
+        results['WEIGHT'] = meas_dict[f'PHASE{plane}'].loc[bpm_pair[0], 'WEIGHT']
+        results['VALUE'] = circular_sum_phase(meas_dict[f'PHASE{plane}'], tune, bpm_pair, 'VALUE')
+        results['MODEL'] = circular_sum_phase(meas_dict[f'PHASE{plane}'], tune, bpm_pair, 'MODEL')
+        results['ERROR'] = circular_sum_phase_error(meas_dict[f'PHASE{plane}'], bpm_pair)
+        results['DIFF'] = results['VALUE'] - results['MODEL']
+        arc_meas.append(results)
+
+    meas_dict[f'PHASE{plane}'] = pd.DataFrame(arc_meas).set_index('NAME')
+
+    return meas_dict
+
+
+def reduce_to_arc_extremities(meas_dict, nominal_model, opt):
+    bpm_pairs_x = get_arc_by_arc_bpm_pairs(meas_dict, opt)
+    bpm_pairs_y = get_arc_by_arc_bpm_pairs(meas_dict, opt)
+    meas_dict = get_arc_phases(bpm_pairs_x, meas_dict, nominal_model.headers['Q1'], 'X')
+    meas_dict = get_arc_phases(bpm_pairs_y, meas_dict, nominal_model.headers['Q2'], 'Y')
+    return meas_dict

omc3/correction/handler.py

@@ -18,6 +18,7 @@
 
 import omc3.madx_wrapper as madx_wrapper
 from omc3.correction import filters, model_appenders, response_twiss
+import omc3.correction.arc_by_arc as abba
 from omc3.optics_measurements.constants import (BETA, DELTA, DISPERSION, F1001,
                                                 F1010, NORM_DISPERSION, PHASE, TUNE,
                                                 DISPERSION_NAME, EXT, REAL, IMAG,
@@ -63,6 +64,9 @@ def correct(accel_inst: Accelerator, opt: DotDict) -> None:
     meas_dict = filters.filter_measurement(optics_params, meas_dict, nominal_model, opt)
     meas_dict = model_appenders.add_differences_to_model_to_measurements(nominal_model, meas_dict)
 
+    if opt.arc_by_arc_phase and accel_inst.NAME == 'lhc':
+        meas_dict = abba.reduce_to_arc_extremities(meas_dict, nominal_model, opt)
+
     resp_dict = filters.filter_response_index(resp_dict, meas_dict, optics_params)
     resp_matrix = _join_responses(resp_dict, optics_params, vars_list)
     delta = tfs.TfsDataFrame(0., index=vars_list, columns=[DELTA])

omc3/correction/model_appenders.py

@@ -67,10 +67,16 @@ def _get_model_generic(model: pd.DataFrame, meas: pd.DataFrame, key: str) -> pd.
 
 
 def _get_model_phases(model: pd.DataFrame, meas: pd.DataFrame, key: str) -> pd.DataFrame:
-    model_column = f"{PHASE_ADV}{key[-1]}"
+    plane = key[-1]
+    tunes = {'X':model.headers['Q1'],
+             'Y':model.headers['Q2'],
+             }
+    model_column = f"{PHASE_ADV}{plane}"
     with logging_tools.log_pandas_settings_with_copy(LOG.debug):
-        meas[MODEL] = (model.loc[meas["NAME2"].to_numpy(), model_column].to_numpy() -
-                       model.loc[meas.index.to_numpy(), model_column].to_numpy())
+        model_phases_advances = (model.loc[meas["NAME2"].to_numpy(), model_column].to_numpy() -
+                                 model.loc[meas.index.to_numpy(), model_column].to_numpy())
+        model_phases_advances[model_phases_advances < 0] += tunes[plane]
+        meas[MODEL] = model_phases_advances
         meas[DIFF] = df_diff(meas, VALUE, MODEL)
     return meas
 

omc3/correction/model_diff.py

@@ -54,7 +54,6 @@ def _get_mapping():
     return defaultdict(
         lambda: _diff,
         {
-            PHASE_ADV: _phase_advance_diff,
             PHASE: _phase_advance_diff,
             BETA: _beta_beating,
             NORM_DISPERSION: _normalized_dispersion_diff,

omc3/definitions/optics.py

@@ -195,6 +195,7 @@ def delta_label(self):
 NORM_DISPERSION_COLUMN = ColumnsAndLabels(NORM_DISPERSION, _label=ylabels['norm_dispersion'], _text_label='normalized dispersion')
 PHASE_COLUMN =           ColumnsAndLabels(PHASE, _label=ylabels['phase'], _text_label='phase')
 TOTAL_PHASE_COLUMN =     ColumnsAndLabels(PHASE, _label=ylabels['phase'], _text_label='total phase')
+MU_COLUMN =              ColumnsAndLabels(PHASE_ADV, _label=ylabels['phasetot'], _text_label='total phase')
 PHASE_ADVANCE_COLUMN =   ColumnsAndLabels(f'{PHASE_ADV}{{0}}{MDL}', _label=r'Phase Advance [$2 \pi$]', _text_label='phase advance')
 S_COLUMN =               ColumnsAndLabels(S, _label='Location [m]', _text_label='longitudinal location', needs_plane=False)
 

omc3/global_correction.py

@@ -194,6 +194,8 @@
     "beta_filename": "beta_phase_",
     "method": "pinv",
     "iterations": 4,
+    "arc_by_arc_phase": False,
+    "include_ips": None,
 }
 
 
@@ -272,6 +274,15 @@ def correction_params():
     params.add_parameter(name="update_response",
                          action="store_true",
                          help="Update the (analytical) response per iteration.", )
+    params.add_parameter(name="arc_by_arc_phase",
+                         type=bool,
+                         default=CORRECTION_DEFAULTS["arc_by_arc_phase"],
+                         help="Set to True for arc-by-arc total pahse correction.", )
+    params.add_parameter(name="include_ips",
+                         type=str,
+                         choices=("left", "right"),
+                         default=CORRECTION_DEFAULTS["include_ips"],
+                         help="Include ips left or right of arcs in arcs.", )
     return params