Arc corr and glob check fix

CHANGELOG.md

@@ -1,5 +1,19 @@
 # OMC3 Changelog
 
+#### IN PROGRESS - v0.18.0 - _jdilly_, _fscarlier_, _fesoubel_
+
+- Fixed:
+  - Plot Optics Measurements: Added extra mpl style for clearer plots
+  - LHC exciter BPM not found: Tells you which BPMs were searched for
+  - Plot Spectrum: Correct error handling for Single-Plane BPMs
+
+- Added:
+  - Global Correction: Total phase correction arc-by-arc
+  - Global Correction: New MQM-knob categories `MQM_INJ_2024` and `MQM_TOP_2024` without Q4 and Q4-6 
+  - Script to copy KMod results into optics directory
+  - Cleaning: Filter BPMs with NaNs
+  - Cleaning: Log bad BPMs with reasons before raising errors.
+
 #### 2024-11-14 - v0.20.0 - _jdilly_, _awegsche_
 
 - Added:

omc3/check_corrections.py

@@ -246,7 +246,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
@@ -481,8 +481,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("Differences to nominal model calculated.")
 
@@ -520,12 +522,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,
             )
@@ -577,7 +584,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,98 @@
+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, plane):
+    bpms = meas_dict[f'PHASE{plane}'].index
+    beam = bpms[0][-1]
+    bpm_pairs = {}
+    bpm_pairs_with_ips = {}
+
+    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 opt.include_ips_in_arc_by_arc == '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 opt.include_ips_in_arc_by_arc == '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, "X")
+    bpm_pairs_y = get_arc_by_arc_bpm_pairs(meas_dict, opt, "Y")
+    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,7 +18,7 @@
 from sklearn.linear_model import OrthogonalMatchingPursuit
 
 import omc3.madx_wrapper as madx_wrapper
-from omc3.correction import filters, model_appenders, response_twiss, response_madx
+from omc3.correction import filters, model_appenders, response_twiss, response_madx, arc_by_arc as abba
 from omc3.correction.constants import DIFF, ERROR, VALUE, WEIGHT, ORBIT_DPP
 from omc3.correction.model_appenders import add_coupling_to_model
 from omc3.correction.response_io import read_fullresponse
@@ -69,6 +69,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

@@ -175,7 +175,7 @@
                                                 NORM_DISPERSION, PHASE, TUNE)
 from omc3.model import manager
 from omc3.utils import logging_tools
-from omc3.utils.iotools import PathOrStr, save_config
+from omc3.utils.iotools import PathOrStr, OptionalStr, save_config
 
 if TYPE_CHECKING:
     from generic_parser import DotDict
@@ -190,14 +190,14 @@
                          f"{F1001}R", f"{F1001}I", f"{F1010}R", f"{F1010}I")
 
 CORRECTION_DEFAULTS = {
-    "optics_file": None,
     "output_filename": "changeparameters_iter",
     "svd_cut": 0.01,
     "optics_params": OPTICS_PARAMS_CHOICES[:6],
     "variable_categories": ["MQM", "MQT", "MQTL", "MQY"],
     "beta_filename": "beta_phase_",
     "method": "pinv",
     "iterations": 4,
+    "include_ips_in_arc_by_arc": None,
 }
 
 
@@ -276,6 +276,14 @@ 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",
+                         action="store_true",
+                         help="Set to perform arc-by-arc total phase correction.", )
+    params.add_parameter(name="include_ips_in_arc_by_arc",
+                         type=str,
+                         choices=("left", "right"),
+                         default=CORRECTION_DEFAULTS["include_ips_in_arc_by_arc"],
+                         help="If not specified only takes pure arcs. Otherwise it includes IPs left or right of arcs.", )
     return params
 
 

omc3/harpy/clean.py

@@ -54,14 +54,20 @@ def _cut_cleaning(harpy_input, bpm_data, model):
     bpm_flatness = _detect_flat_bpms(bpm_data, harpy_input.peak_to_peak)
     bpm_spikes = _detect_bpms_with_spikes(bpm_data, harpy_input.max_peak)
     exact_zeros = _detect_bpms_with_exact_zeros(bpm_data, harpy_input.keep_exact_zeros)
-    all_bad_bpms = _index_union(known_bad_bpms, bpm_flatness, bpm_spikes, exact_zeros)
+    bpm_nans = _detect_bpms_with_nans(bpm_data)
+    all_bad_bpms = _index_union(known_bad_bpms, bpm_flatness, bpm_spikes, exact_zeros, bpm_nans)
     original_bpms = bpm_data.index
 
     bpm_data = bpm_data.loc[bpm_data.index.difference(all_bad_bpms, sort=False)]
     bad_bpms_with_reasons = _get_bad_bpms_summary(
-        harpy_input, known_bad_bpms, bpm_flatness, bpm_spikes, exact_zeros
+        harpy_input, 
+        known=known_bad_bpms, 
+        flat=bpm_flatness, 
+        spike=bpm_spikes, 
+        zero=exact_zeros, 
+        nan=bpm_nans,
     )
-    _report_clean_stats(original_bpms.size, bpm_data.index.size)
+    _report_clean_stats(original_bpms.size, bpm_data.index.size, bad_bpms_with_reasons)
     bpm_data = _fix_polarity(harpy_input.wrong_polarity_bpms, bpm_data)
     if model is not None and harpy_input.first_bpm is not None:
         bpm_data = _resync_bpms(harpy_input, bpm_data, model)
@@ -110,42 +116,81 @@ def _detect_bpms_with_spikes(bpm_data, max_peak_cut):
     too_low = bpm_data[bpm_data.min(axis=1) < -max_peak_cut].index
     bpm_spikes = too_high.union(too_low)
     if bpm_spikes.size:
-        LOGGER.debug(f"Spikes > {max_peak_cut} detected. BPMs removed: {bpm_spikes.size}")
+        LOGGER.debug(
+            f"Spikes > {max_peak_cut} detected. "
+            f"{bpm_spikes.size} BPMs removed: {', '.join(bpm_spikes)}"
+        )
     return bpm_spikes
 
 
+def _detect_bpms_with_nans(bpm_data) -> pd.Index:
+    """Detects BPMs with NaN values."""
+    nan_bpms = bpm_data[bpm_data.isna().any(axis=1)].index
+    if nan_bpms.size:
+        LOGGER.warning(
+            f"NaN BPMs detected. "
+            f"{nan_bpms.size} BPMs removed: {', '.join(nan_bpms)}"
+        )
+    return nan_bpms
+
+
 def _detect_bpms_with_exact_zeros(bpm_data, keep_exact_zeros):
     """Detects BPMs with exact zeros due to OP workaround."""
     if keep_exact_zeros:
         LOGGER.debug("Skipped exact zero check")
         return pd.Index([])
     exact_zeros = bpm_data[~np.all(bpm_data, axis=1)].index
     if exact_zeros.size:
-        LOGGER.debug(f"Exact zeros detected. BPMs removed: {exact_zeros.size}")
+        LOGGER.debug(
+            f"Exact zeros detected. "
+            f"{exact_zeros.size} BPMs removed: {', '.join(exact_zeros)} "
+        )
     return exact_zeros
 
 
-def _get_bad_bpms_summary(harpy_input, known_bad_bpms, bpm_flatness, bpm_spikes, exact_zeros):
-    return ([f"{bpm_name} Known bad BPM" for bpm_name in known_bad_bpms] +
-            [f"{bpm_name} Flat BPM, the difference between min/max is smaller than "
-            f"{harpy_input.peak_to_peak}" for bpm_name in bpm_flatness] +
-            [f"{bpm_name} Spiky BPM, found spike higher than "
-            f"{harpy_input.max_peak}" for bpm_name in bpm_spikes] +
-            [f"{bpm_name} Found an exact zero" for bpm_name in exact_zeros])
+def _get_bad_bpms_summary(harpy_input, **kwargs):
+    human_readable = {
+        "known": "Known bad BPM",
+        "flat": f"Flat BPM, the difference between min/max is smaller than {harpy_input.peak_to_peak}",
+        "spike": f"Spiky BPM, found spike higher than {harpy_input.max_peak}",
+        "zero": "Found an exact zero",
+        "nan": "Found NaN",
+    }
+
+    # Quick check that catches coding errors/typos:
+    unknown_kwargs = [kwarg for kwarg in kwargs if kwarg not in human_readable.keys()]
+    if len(unknown_kwargs):
+        raise NameError(f"Unknown reason(s) for Bad-BPMs: {unknown_kwargs}")
+
+    return [
+        f"{bpm_name} {msg}" for reason, msg in human_readable.items() for bpm_name in kwargs[reason]
+    ]
 
 
-def _report_clean_stats(n_total_bpms, n_good_bpms):
+def _report_clean_stats(n_total_bpms, n_good_bpms, bad_bpms_with_reasons):
     LOGGER.debug("Filtering done:")
+
+    # As it not written out yet, provide more info in case of raising errors:
+    bad_bpms_message = "Bad BPMs found:\n"
+    bad_bpms_message += "\n".join(bad_bpms_with_reasons)
+
+    # If all BPMs have been bad ---
     if n_good_bpms == 0:
+        LOGGER.info(bad_bpms_message)
         raise ValueError("Total Number of BPMs after filtering is zero.")
+
+    # The good, the bad and the ugly BPMs ---
     n_bad_bpms = n_total_bpms - n_good_bpms
     LOGGER.debug(f"(Statistics for file reading) Total BPMs: {n_total_bpms}, "
                  f"Good BPMs: {n_good_bpms} ({(100 * n_good_bpms / n_total_bpms):2.2f}%), "
                  f"Bad BPMs: {n_bad_bpms} ({(100 * n_bad_bpms / n_total_bpms):2.2f}%)")
     if (n_good_bpms / n_total_bpms) < 0.5:
+        LOGGER.info(bad_bpms_message)
         raise ValueError("More than half of BPMs are bad. "
                          "This could be because a bunch not present in the machine has been "
                          "selected or because of a problem with the phasing of the BPMs.")
+
+    LOGGER.debug(bad_bpms_message)
 
 
 def _index_union(*indices):

omc3/model/accelerators/generic.py

@@ -0,0 +1,17 @@
+
+from generic_parser import EntryPoint
+
+from omc3.model.accelerators.accelerator import AccElementTypes, Accelerator
+
+
+class Generic(Accelerator):
+    NAME = "generic"
+    RE_DICT = {AccElementTypes.BPMS: r"B",
+               AccElementTypes.MAGNETS: r".*",
+               AccElementTypes.ARC_BPMS: r"B",
+               }
+
+    def __init__(self, *args, **kwargs):
+        parser = EntryPoint(self.get_parameters(), strict=True)
+        opt = parser.parse(*args, **kwargs)
+        super().__init__(opt)