Global Correction Fixes (#458)

* test, fix and annotations
* new variables logic
* testing more correctors
* 2024 LHC Correctors
* allow correctos json files in model dir

CHANGELOG.md

@@ -1,5 +1,21 @@
 # OMC3 Changelog
 
+#### 2024-09-18 - v0.16.0 - _jdilly_
+
+- Added:
+  - Global Correction for LHC:
+    - The correction variables in the LHC accelerator class are now handled differently internally,
+      allowing new variable classes to be added to each lhc-year and user-given files in the model-directory.
+    - Variable categories `MQM_ALL` added to all LHC years.
+    - Variable categories `MQM_INJ_2024` and `MQM_TOP_2024` added to LHC 2024.
+    - Adding a "-" in front of a given correction variable name removes this variable from the correction. Does not work for whole variable categories.
+  - Tests for running `global_correction` with `omp` and `pinv` correction methods.
+
+- Fixed:
+  - Orthogonal Matching Pursuit (`omp`) in global correction runs again ([#448](https://github.com/pylhc/omc3/issues/448))
+  - Corrections Check window stop-iteration issue: fixed. Single-plane files, i.e. normalized dispersion, now accepted ([#447](https://github.com/pylhc/omc3/issues/447))
+  - LHC Correction variables now logged, if not found in the variable categories ([#446](https://github.com/pylhc/omc3/issues/446))
+
 #### 2024-09-16 - v0.15.4 - _jdilly_
 
 - Fixed:

omc3/__init__.py

@@ -11,7 +11,7 @@
 __title__ = "omc3"
 __description__ = "An accelerator physics tools package for the OMC team at CERN."
 __url__ = "https://github.com/pylhc/omc3"
-__version__ = "0.15.4"
+__version__ = "0.16.0"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"

omc3/check_corrections.py

@@ -228,14 +228,15 @@
     Limits on the y axis (Tupel)
 
 """
+from __future__ import annotations
+
 import copy
 from pathlib import Path
-from typing import Dict, Sequence, Any, List
+from typing import TYPE_CHECKING, Any
 
 import pandas as pd
 
 import tfs
-from generic_parser import DotDict
 from generic_parser.entrypoint_parser import EntryPointParameters, entrypoint
 from omc3.correction import filters
 from omc3.correction import handler as global_correction
@@ -258,6 +259,11 @@
 from omc3.utils.stats import rms, circular_rms
 from tfs import TfsDataFrame
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from generic_parser import DotDict
+
+
 LOG = logging_tools.get_logger(__name__)
 
 
@@ -390,7 +396,7 @@ def _check_opt_add_dicts(opt: DotDict) -> DotDict:
     return opt
 
 
-def _get_corrections(corrections: Sequence[Path], file_pattern: str = None) -> Dict[str, Sequence[Path]]:
+def _get_corrections(corrections: Sequence[Path], file_pattern: str = None) -> dict[str, Sequence[Path]]:
     """ Sort the given correction files:
     If given by individual files, they all go into one bucket,
     if given by folders (i.e. scenarios) they are sorted by its name.
@@ -414,14 +420,14 @@ def _get_corrections(corrections: Sequence[Path], file_pattern: str = None) -> D
     return corr_dict
 
 
-def _glob_regex_paths(path: Path, pattern: str) -> List[Path]:
+def _glob_regex_paths(path: Path, pattern: str) -> list[Path]:
     """ Filter the files in path by pattern and return a list of paths. """
     return [path / f for f in glob_regex(path, pattern)]
 
 
 # Main and Output --------------------------------------------------------------
 
-def _get_measurement_filter(nominal_model: TfsDataFrame, opt: DotDict) -> Dict[str, pd.Index]:
+def _get_measurement_filter(nominal_model: TfsDataFrame, opt: DotDict) -> dict[str, pd.Index]:
     """ Get the filtered measurement based on the cuts as done in the correction calculation.
     As we need this only for RMS calculations later on, we only care about the
     BPM-names. So the returned dict contains the index to be used for this
@@ -456,7 +462,7 @@ def _get_measurement_filter(nominal_model: TfsDataFrame, opt: DotDict) -> Dict[s
 
 def _create_model_and_write_diff_to_measurements(
         output_dir: Path, measurement: OpticsMeasurement, correction_name: str, correction_files: Sequence[Path],
-        accel_inst: Accelerator, rms_masks: Dict) -> OpticsMeasurement:
+        accel_inst: Accelerator, rms_masks: dict) -> OpticsMeasurement:
     """ Create a new model with the corrections (well, the "matchings") applied and calculate
     the difference to the nominal model, i.e. the expected improvement of the measurements
     (for detail see main docstring in this file).
@@ -528,7 +534,7 @@ def _create_model_and_write_diff_to_measurements(
 
 def _create_check_columns(measurement: OpticsMeasurement, output_measurement: OpticsMeasurement, diff_models: TfsDataFrame,
                           colmap_meas: ColumnsAndLabels, colmap_model: ColumnsAndLabels, attribute: str,
-                          rms_mask: Dict = None) -> None:
+                          rms_mask: dict = None) -> None:
     """Creates the columns in the measurements, that allow for checking the corrections.
     These are:
         diff_correction_column: Difference between the corrected and uncorrected model,
@@ -643,7 +649,7 @@ def _maybe_add_coupling_to_model(model: tfs.TfsDataFrame, measurement: OpticsMea
 
 # Plotting ---------------------------------------------------------------------
 
-def _do_plots(corrections: Dict[str, Any], opt: DotDict):
+def _do_plots(corrections: dict[str, Any], opt: DotDict):
     """ Plot the differences of the matched models to the measurement. """
     opt_plot = {k: v for k, v in opt.items() if k in get_plotting_style_parameters().keys()}
     opt_plot["input_dir"] = opt.output_dir

omc3/correction/filters.py

@@ -13,28 +13,44 @@
 i.e. beta, phase etc. In this implementation most of it is handled by
 the `_get_filtered_generic` function.
 """
+from __future__ import annotations
+
 from collections import defaultdict
-from typing import Callable, Dict, Sequence
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
 import tfs
 from generic_parser import DotDict
 
 from omc3.correction.constants import ERROR, VALUE, WEIGHT
-from omc3.optics_measurements.constants import DELTA, ERR, NAME2, PHASE, PHASE_ADV, TUNE
 from omc3.definitions.constants import PLANES
-from omc3.optics_measurements.constants import AMPLITUDE, F1001, F1010, IMAG, REAL
+from omc3.optics_measurements.constants import (
+    AMPLITUDE,
+    DELTA,
+    ERR,
+    F1001,
+    F1010,
+    IMAG,
+    NAME2,
+    PHASE,
+    PHASE_ADV,
+    REAL,
+    TUNE,
+)
 from omc3.utils import logging_tools, stats
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
 LOG = logging_tools.get_logger(__name__)
 
 
 # Measurement Filter -----------------------------------------------------------
 
 
 def filter_measurement(
-    keys: Sequence[str], meas: Dict[str, pd.DataFrame], model: pd.DataFrame, opt: DotDict
+    keys: Sequence[str], meas: dict[str, pd.DataFrame], model: pd.DataFrame, opt: DotDict
 ) -> dict:
     """Filters measurements in `keys` based on the dict-entries (keys as in `keys`)
     in `opt.errorcut`, `opt.modelcut` and `opt.weights` and unifies the
@@ -103,9 +119,9 @@ def _get_filtered_generic(col: str, meas: pd.DataFrame, model: pd.DataFrame, opt
 
     # if opt.automatic_model_cut:  # TODO automated model cut
     #     model_filter = _get_smallest_data_mask(np.abs(meas.loc[:, f"{DELTA}{col}"].to_numpy()), portion=0.95)
-    if f"{PHASE}" in col:
+    if PHASE in col:
         new[NAME2] = meas.loc[:, NAME2].to_numpy()
-        second_bpm_exists = np.in1d(new.loc[:, NAME2].to_numpy(), new.index.to_numpy())
+        second_bpm_exists = np.isin(new.loc[:, NAME2].to_numpy(), new.index.to_numpy())
         good_bpms = error_mask & model_mask & second_bpm_exists
         good_bpms[-1] = False  # TODO not sure why, ask Lukas? (jdilly)
     else:
@@ -162,7 +178,7 @@ def _get_errorbased_weights(key: str, weights, errors):
 # Response Matrix Filter -------------------------------------------------------
 
 
-def filter_response_index(response: Dict, measurement: Dict, keys: Sequence[str]):
+def filter_response_index(response: dict, measurement: dict, keys: Sequence[str]):
     """Filters the index of the response matrices `response` by the respective entries in `measurement`."""
     # rename MU to PHASE as we create a PHASE-Response afterwards
     # easier to do here, than to check eveywhere below. (jdilly)
@@ -182,7 +198,7 @@ def filter_response_index(response: Dict, measurement: Dict, keys: Sequence[str]
     return new_response
 
 
-def _get_response_filters() -> Dict[str, Callable]:
+def _get_response_filters() -> dict[str, Callable]:
     """
     Returns a dict with the respective `_get_*_response` functions that defaults
     to `_get_generic_response`.

omc3/correction/handler.py

@@ -4,30 +4,35 @@
 
 This module contains high-level functions to manage most functionality of the corrections calculations.
 """
+from __future__ import annotations
+
 import datetime
 import time
 from pathlib import Path
-from typing import Callable, Dict, List, Sequence, Tuple, Union
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
 import tfs
-from generic_parser import DotDict
 from sklearn.linear_model import OrthogonalMatchingPursuit
 
 import omc3.madx_wrapper as madx_wrapper
 from omc3.correction import filters, model_appenders, response_twiss
-from omc3.optics_measurements.constants import (BETA, DELTA, DISPERSION, F1001,
-                                                F1010, NORM_DISPERSION, PHASE, TUNE,
-                                                DISPERSION_NAME, EXT, NORM_DISP_NAME, PHASE_NAME, NAME
-                                                )
-from omc3.correction.constants import ERROR, VALUE, WEIGHT, DIFF
+from omc3.correction.constants import DIFF, ERROR, VALUE, WEIGHT
 from omc3.correction.model_appenders import add_coupling_to_model
 from omc3.correction.response_io import read_fullresponse
 from omc3.model.accelerators.accelerator import Accelerator
+from omc3.optics_measurements.constants import (BETA, DELTA, DISPERSION, DISPERSION_NAME, EXT,
+                                                F1001, F1010, NAME, NORM_DISP_NAME, NORM_DISPERSION,
+                                                PHASE, PHASE_NAME, TUNE)
 from omc3.utils import logging_tools
 from omc3.utils.stats import rms
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+    from generic_parser import DotDict
+
+
 LOG = logging_tools.get_logger(__name__)
 
 
@@ -137,8 +142,8 @@ def get_measurement_data(
         keys: Sequence[str],
         meas_dir: Path,
         beta_filename: str,
-        w_dict: Dict[str, float] = None,
-) -> Tuple[List[str], Dict[str, tfs.TfsDataFrame]]:
+        w_dict: dict[str, float] = None,
+) -> tuple[list[str], dict[str, tfs.TfsDataFrame]]:
     """ Loads all measurements defined by `keys` into a dictionary. """
     measurement = {}
     filtered_keys = keys
@@ -211,7 +216,7 @@ def _maybe_add_coupling_to_model(model: tfs.TfsDataFrame, keys: Sequence[str]) -
     return model
 
 
-def _create_corrected_model(twiss_out: Union[Path, str], change_params: Sequence[Path], accel_inst: Accelerator) -> tfs.TfsDataFrame:
+def _create_corrected_model(twiss_out: Path | str, change_params: Sequence[Path], accel_inst: Accelerator) -> tfs.TfsDataFrame:
     """ Use the calculated deltas in changeparameters.madx to create a corrected model """
     madx_script: str = accel_inst.get_update_correction_script(twiss_out, change_params)
     twiss_out_path = Path(twiss_out)
@@ -268,7 +273,7 @@ def _orthogonal_matching_pursuit(response_mat: pd.DataFrame, diff_vec, opt: DotD
     if opt.n_correctors is None:
         raise ValueError("n_correctors setting needed for orthogonal matching pursuit.")
 
-    res = OrthogonalMatchingPursuit(opt.n_correctors).fit(response_mat, diff_vec)
+    res = OrthogonalMatchingPursuit(n_nonzero_coefs=opt.n_correctors).fit(response_mat, diff_vec)
     coef = res.coef_
     LOG.debug(f"Orthogonal Matching Pursuit Results: \n"
               f"  Chosen variables: {response_mat.columns.to_numpy()[coef.nonzero()]}\n"

omc3/global_correction.py

@@ -164,10 +164,11 @@
  * Missing a part that treats the output from LSA
 
 """
+from __future__ import annotations
+
 from pathlib import Path
-from typing import Dict
+from typing import TYPE_CHECKING
 
-from generic_parser import DotDict
 from generic_parser.entrypoint_parser import EntryPointParameters, entrypoint
 from omc3.correction import handler
 from omc3.optics_measurements.constants import (BETA, DISPERSION, F1001, F1010,
@@ -176,6 +177,9 @@
 from omc3.utils import logging_tools
 from omc3.utils.iotools import PathOrStr, save_config
 
+if TYPE_CHECKING:
+    from generic_parser import DotDict
+
 LOG = logging_tools.get_logger(__name__)
 
 OPTICS_PARAMS_CHOICES = (f"{PHASE}X", f"{PHASE}Y",
@@ -316,7 +320,7 @@ def _add_hardcoded_paths(opt: DotDict) -> DotDict:  # acts inplace...
 
 
 # Define functions here, to new optics params
-def _get_default_values() -> Dict[str, Dict[str, float]]:
+def _get_default_values() -> dict[str, dict[str, float]]:
     return {
         "modelcut": {
             f"{PHASE}X": 0.05,