Add Calculation of DPP in full response and global correction (#466)

* Add `ORBIT_DPP` as a parameter seen by `response_madx.create_fullresponse` and `global_correction` in `variable_categories`. This allows a calculation of the effective deltap/p change when the orbit changes. * Add `response_madx.create_fullresponse` as an option when updating the full response in the global correction * Fix the calculation of`response_twiss.create_response` when updating the full response in the global correction * Test all of above * some maintenance changes * Update changelog and bump version to 0.18.0 --------- Co-authored-by: jgray-19 <[email protected]> Co-authored-by: JoschD <[email protected]>
pylhc · Nov 1, 2024 · 351b2a7 · 351b2a7
1 parent 4bb17b1
commit 351b2a7
Show file tree

Hide file tree

Showing 20 changed files with 2,923 additions and 80 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,17 @@
 # OMC3 Changelog
 
+#### 2024-10-29 - v0.18.0 - _jgray_
+
+- Added:
+  - Add the ability to calculate a deltap/p offset caused by a change in orbit.
+  - Add the ability to use `response_madx` to calculate the updated response matrix for the global correction.
+  - Tests for the calculation of the deltap/p and corresponding correction.
+  - Tests for the calculation of the updated response matrix for the global correction.
+
+- Fixed:
+  - Fixed the `response_twiss` when updating the response matrix when calculating the global correction.
+
+
 #### 2024-10-29 - v0.17.0 - _jdilly_ 
 
 - Added:

diff --git a/omc3/__init__.py b/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.17.0"
+__version__ = "0.18.0"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"

diff --git a/omc3/correction/constants.py b/omc3/correction/constants.py
@@ -28,6 +28,7 @@
 
 # For FullResponse
 INCR: str = "incr"
+ORBIT_DPP: str = "orbit_dpp"
 
 # Correction Test Constants ----------------------------------------------------
 MODEL_NOMINAL_FILENAME: str = "twiss_nominal.tfs"  # using twiss from model for now

diff --git a/omc3/correction/filters.py b/omc3/correction/filters.py
@@ -167,7 +167,7 @@ def _get_errorbased_weights(key: str, weights, errors):
     # TODO case without errors used may corrupt the correction (typical error != 1)
     w2 = stats.weights_from_errors(errors)
     if w2 is None:
-        LOG.warn(
+        LOG.warning(
             f"Weights will not be based on errors for '{key}'"
             f", zeros of NaNs were found. Maybe don't use --errorbars."
         )

diff --git a/omc3/correction/handler.py b/omc3/correction/handler.py
@@ -10,15 +10,16 @@
 import time
 from pathlib import Path
 from typing import TYPE_CHECKING
+import copy
 
 import numpy as np
 import pandas as pd
 import tfs
 from sklearn.linear_model import OrthogonalMatchingPursuit
 
 import omc3.madx_wrapper as madx_wrapper
-from omc3.correction import filters, model_appenders, response_twiss
-from omc3.correction.constants import DIFF, ERROR, VALUE, WEIGHT
+from omc3.correction import filters, model_appenders, response_twiss, response_madx
+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
 from omc3.model.accelerators.accelerator import Accelerator
@@ -48,6 +49,8 @@ def correct(accel_inst: Accelerator, opt: DotDict) -> None:
     method_options = opt.get_subdict(["svd_cut", "n_correctors"])
     # read data from files
     vars_list = _get_varlist(accel_inst, opt.variable_categories)
+    update_deltap = ORBIT_DPP in vars_list
+
     optics_params, meas_dict = get_measurement_data(
         opt.optics_params,
         opt.meas_dir,
@@ -76,10 +79,10 @@ def correct(accel_inst: Accelerator, opt: DotDict) -> None:
 
         # ######### Update Model and Response ######### #
         if iteration > 0:
-            LOG.debug("Updating model via MADX.")
+            LOG.debug("Updating model via MAD-X.")
             corr_model_path = opt.output_dir / f"twiss_{iteration}{EXT}"
 
-            corr_model_elements = _create_corrected_model(corr_model_path, [opt.change_params_path], accel_inst)
+            corr_model_elements = _create_corrected_model(corr_model_path, [opt.change_params_path], accel_inst, update_deltap)
             corr_model_elements = _maybe_add_coupling_to_model(corr_model_elements, optics_params)
 
             bpms_index_mask = accel_inst.get_element_types_mask(corr_model_elements.index, types=["bpm"])
@@ -88,11 +91,15 @@ def correct(accel_inst: Accelerator, opt: DotDict) -> None:
             meas_dict = model_appenders.add_differences_to_model_to_measurements(corr_model, meas_dict)
 
             if opt.update_response:
-                LOG.debug("Updating response.")
-                # please look away for the next two lines.
-                accel_inst._model = corr_model
-                accel_inst._elements = corr_model_elements
-                resp_dict = response_twiss.create_response(accel_inst, opt.variable_categories, optics_params)
+                resp_dict = _update_response(
+                    accel_inst=accel_inst,
+                    corrected_elements=corr_model_elements,
+                    optics_params=optics_params,
+                    corr_files=[opt.change_params_path],
+                    variable_categories=opt.variable_categories,
+                    update_dpp=update_deltap,
+                    update_response=opt.update_response,
+                )
                 resp_dict = filters.filter_response_index(resp_dict, meas_dict, optics_params)
                 resp_matrix = _join_responses(resp_dict, optics_params, vars_list)
 
@@ -109,6 +116,47 @@ def correct(accel_inst: Accelerator, opt: DotDict) -> None:
     LOG.info("Finished Iterative Global Correction.")
 
 
+
+def _update_response(
+    accel_inst: Accelerator, 
+    corrected_elements: pd.DataFrame,
+    optics_params: Sequence[str],
+    corr_files: Sequence[Path],
+    variable_categories: Sequence[str], 
+    update_dpp: bool, 
+    update_response: bool | str,
+    ) -> dict[str, pd.DataFrame]:
+    """ Create an updated response matrix.
+    
+    If we are to compute the response including the DPP, then we have to do so from MAD-X, 
+    as we do not have the analytical formulae. This therefore requires correction files to be
+    provided.
+    Otherwise we go through the way of computing the response the user requested.
+
+    All other parameters are taken care of in the model/elements for the response_twiss only.
+    """
+    # update model by creating a copy of the accelerator instance
+    accel_inst_cp = copy.copy(accel_inst)
+
+    # Modifiers is None or list, if none, we need to make a list before extending it with the correction files
+    accel_inst_cp.modifiers = list(accel_inst_cp.modifiers or []) + corr_files
+
+    if update_dpp:
+        LOG.info("Updating response via MAD-X, due to delta dpp requested.")
+        resp_dict = response_madx.create_fullresponse(accel_inst_cp, variable_categories)
+    else:
+        if update_response == "madx":
+            LOG.info("Updating response via MAD-X.")
+            resp_dict = response_madx.create_fullresponse(accel_inst_cp, variable_categories)
+        else:
+            LOG.info("Updating response via analytical formulae.")
+            accel_inst_cp.elements = corrected_elements
+            # accel_inst_cp.model = corrected_model # - Not needed, don't think it's used by response_twiss (jgray 2024)
+            resp_dict = response_twiss.create_response(accel_inst_cp, variable_categories, optics_params)
+
+    return resp_dict
+
+
 # Input ------------------------------------------------------------------------
 
 
@@ -216,9 +264,9 @@ def _maybe_add_coupling_to_model(model: tfs.TfsDataFrame, keys: Sequence[str]) -
     return model
 
 
-def _create_corrected_model(twiss_out: Path | str, change_params: Sequence[Path], accel_inst: Accelerator) -> tfs.TfsDataFrame:
+def _create_corrected_model(twiss_out: Path | str, corr_files: Sequence[Path], accel_inst: Accelerator, update_dpp: bool = False) -> 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)
+    madx_script: str = accel_inst.get_update_correction_script(twiss_out, corr_files, update_dpp)
     twiss_out_path = Path(twiss_out)
     madx_script = f"! Based on model '{accel_inst.model_dir}'\n" + madx_script
     madx_wrapper.run_string(
@@ -314,11 +362,11 @@ def _calculate_delta(
 def _print_rms(meas: dict, diff_w, r_delta_w) -> None:
     """ Prints current RMS status """
     f_str = "{:>20s} : {:.5e}"
-    LOG.debug("RMS Measure - Model (before correction, w/o weigths):")
+    LOG.debug("RMS Measure - Model (before correction, w/o weights):")
     for key in meas:
         LOG.debug(f_str.format(key, rms(meas[key].loc[:, DIFF].to_numpy())))
 
-    LOG.info("RMS Measure - Model (before correction, w/ weigths):")
+    LOG.info("RMS Measure - Model (before correction, w/ weights):")
     for key in meas:
         LOG.info(f_str.format(key, rms(meas[key].loc[:, DIFF].to_numpy() * meas[key].loc[:, WEIGHT].to_numpy())))
 

diff --git a/omc3/correction/response_madx.py b/omc3/correction/response_madx.py
@@ -12,28 +12,41 @@
 
 :author: Lukas Malina, Joschua Dilly, Jaime (...) Coello de Portugal
 """
+from __future__ import annotations
+
 import copy
 import multiprocessing
+import zipfile
 from pathlib import Path
-from typing import Dict, Sequence, Tuple, List
+from typing import TYPE_CHECKING
 
 import numpy as np
-import zipfile
 import pandas as pd
 import tfs
 from numpy.exceptions import ComplexWarning
 from optics_functions.coupling import coupling_via_cmatrix
 
 import omc3.madx_wrapper as madx_wrapper
-from omc3.optics_measurements.constants import (BETA, DISPERSION, F1001, F1010,
-                                                NORM_DISPERSION, PHASE_ADV, TUNE)
-from omc3.correction.constants import INCR
-from omc3.model.accelerators.accelerator import Accelerator, AccElementTypes
+from omc3.correction.constants import INCR, ORBIT_DPP
+from omc3.model.accelerators.accelerator import AccElementTypes, Accelerator
+from omc3.optics_measurements.constants import (
+    BETA,
+    DISPERSION,
+    F1001,
+    F1010,
+    NAME,
+    NORM_DISPERSION,
+    PHASE_ADV,
+    TUNE,
+)
 from omc3.utils import logging_tools
 from omc3.utils.contexts import suppress_warnings, timeit
 
 LOG = logging_tools.get_logger(__name__)
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 
 # Full Response Mad-X ##########################################################
 
@@ -44,7 +57,7 @@ def create_fullresponse(
     delta_k: float = 2e-5,
     num_proc: int = multiprocessing.cpu_count(),
     temp_dir: Path = None
-) -> Dict[str, pd.DataFrame]:
+) -> dict[str, pd.DataFrame]:
     """ Generate a dictionary containing response matrices for
         beta, phase, dispersion, tune and coupling and saves it to a file.
 
@@ -83,31 +96,47 @@ def _generate_madx_jobs(
     delta_k: float,
     num_proc: int,
     temp_dir: Path
-) -> Dict[str, float]:
+) -> dict[str, float]:
     """ Generates madx job-files """
-    LOG.debug("Generating MADX jobfiles.")
+    LOG.debug("Generating MAD-X jobfiles.")
     incr_dict = {'0': 0.0}
-    vars_per_proc = int(np.ceil(len(variables) / num_proc))
+    compute_deltap: bool = ORBIT_DPP in variables
+    no_dpp_vars = [var for var in variables if var != ORBIT_DPP]
+    vars_per_proc = int(np.ceil(len(no_dpp_vars) / num_proc))
 
     madx_job = _get_madx_job(accel_inst)
+    deltap_twiss = ""
+    if compute_deltap:
+        # This is here only for multiple iteration of the global correction
+        # By including dpp here, it means that if deltap is in variables and dpp is not 0, the orbit and tune magnets change
+        # We have to be very careful that DELTAP_NAME is not used ANYWHERE else in MAD-X
+        madx_job += f"{ORBIT_DPP} = {accel_inst.dpp};\n" # Set deltap to 0
+        madx_job += accel_inst.get_update_deltap_script(deltap=ORBIT_DPP)
+        deltap_twiss = f", deltap={ORBIT_DPP}"
 
     for proc_idx in range(num_proc):
         jobfile_path = _get_jobfiles(temp_dir, proc_idx)
 
         current_job = madx_job
         for i in range(vars_per_proc):
             var_idx = proc_idx * vars_per_proc + i
-            if var_idx >= len(variables):
+            if var_idx >= len(no_dpp_vars):
                 break
-            var = variables[var_idx]
+            var = no_dpp_vars[var_idx]
             incr_dict[var] = delta_k
             current_job += f"{var} = {var}{delta_k:+.15e};\n"
-            current_job += f"twiss, file='{str(temp_dir / f'twiss.{var}')}';\n"
+            current_job += f"twiss, file='{str(temp_dir / f'twiss.{var}')}'{deltap_twiss};\n"
             current_job += f"{var} = {var}{-delta_k:+.15e};\n\n"
 
         if proc_idx == num_proc - 1:
-            current_job += f"twiss, file='{str(temp_dir / 'twiss.0')}';\n"
-
+            current_job += f"twiss, file='{str(temp_dir / 'twiss.0')}'{deltap_twiss};\n"
+
+            if compute_deltap: # If ORBIT_DPP is in variables, we run this in the last iteration
+                # Due to the match and correction of the orbit, this needs to be run at the end of the process
+                incr_dict[ORBIT_DPP] = delta_k
+                current_job += f"{ORBIT_DPP} = {ORBIT_DPP}{delta_k:+.15e};\n"
+                current_job += accel_inst.get_update_deltap_script(deltap=ORBIT_DPP) # Do twiss, correct, match
+                current_job += f"twiss, deltap={ORBIT_DPP}, file='{str(temp_dir/f'twiss.{ORBIT_DPP}')}';\n"
         jobfile_path.write_text(current_job)
     return incr_dict
 
@@ -154,11 +183,11 @@ def _clean_up(temp_dir: Path, num_proc: int) -> None:
 
 
 def _load_madx_results(
-    variables: List[str],
+    variables: list[str],
     process_pool,
     incr_dict: dict,
     temp_dir: Path
-) -> Dict[str, tfs.TfsDataFrame]:
+) -> dict[str, tfs.TfsDataFrame]:
     """ Load the madx results in parallel and return var-tfs dictionary """
     LOG.debug("Loading Madx Results.")
     vars_and_paths = []
@@ -171,7 +200,7 @@ def _load_madx_results(
     return var_to_twiss
 
 
-def _create_fullresponse_from_dict(var_to_twiss: Dict[str, tfs.TfsDataFrame]) -> Dict[str, pd.DataFrame]:
+def _create_fullresponse_from_dict(var_to_twiss: dict[str, tfs.TfsDataFrame]) -> dict[str, pd.DataFrame]:
     """ Convert var-tfs dictionary to fullresponse dictionary. """
     var_to_twiss = _add_coupling(var_to_twiss)
     keys = list(var_to_twiss.keys())
@@ -243,25 +272,25 @@ def _launch_single_job(inputfile_path: Path) -> None:
     madx_wrapper.run_file(inputfile_path, log_file=log_file, cwd=inputfile_path.parent)
 
 
-def _load_and_remove_twiss(var_and_path: Tuple[str, Path]) -> Tuple[str, tfs.TfsDataFrame]:
+def _load_and_remove_twiss(var_and_path: tuple[str, Path]) -> tuple[str, tfs.TfsDataFrame]:
     """ Function for pool to retrieve results """
     (var, path) = var_and_path
     twissfile = path / f"twiss.{var}"
-    tfs_data = tfs.read(twissfile, index="NAME")
+    tfs_data = tfs.read(twissfile, index=NAME)
     tfs_data[f"{TUNE}1"] = tfs_data.Q1
     tfs_data[f"{TUNE}2"] = tfs_data.Q2
     twissfile.unlink()
     return var, tfs_data
 
 
-def _add_coupling(dict_of_tfs: Dict[str, tfs.TfsDataFrame]) -> Dict[str, tfs.TfsDataFrame]:
+def _add_coupling(dict_of_tfs: dict[str, tfs.TfsDataFrame]) -> dict[str, tfs.TfsDataFrame]:
     """
     For each TfsDataFrame in the input dictionary, computes the coupling RDTs and adds a column for
     the real and imaginary parts of the computed coupling RDTs. Returns a copy of the input dictionary with
     the aforementioned computed columns added for each TfsDataFrame.
 
     Args:
-        dict_of_tfs (Dict[str, tfs.TfsDataFrame]): dictionary of Twiss dataframes.
+        dict_of_tfs (dict[str, tfs.TfsDataFrame]): dictionary of Twiss dataframes.
 
     Returns:
         An identical dictionary of Twiss dataframes, with the computed columns added.

diff --git a/omc3/correction/response_twiss.py b/omc3/correction/response_twiss.py
@@ -117,7 +117,8 @@
     https://cds.cern.ch/record/2632945/
 
 """
-from typing import Dict, List, Sequence
+from __future__ import annotations
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -131,6 +132,9 @@
 from omc3.utils import logging_tools
 from omc3.utils.contexts import timeit
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 LOG = logging_tools.get_logger(__name__)
 
 DUMMY_ID = "DUMMY_PLACEHOLDER"
@@ -801,7 +805,7 @@ def upper(list_of_strings: Sequence[str]) -> Sequence[str]:
     return [item.upper() for item in list_of_strings]
 
 
-def get_phase_advances(twiss_df: pd.DataFrame) -> Dict[str, pd.DataFrame]:
+def get_phase_advances(twiss_df: pd.DataFrame) -> dict[str, pd.DataFrame]:
     """
     Calculate phase advances between all elements
 
@@ -840,7 +844,7 @@ def tau(data, q):
 def create_response(
     accel_inst: Accelerator,
     vars_categories: Sequence[str],
-    optics_params: List[str],
+    optics_params: list[str],
 ) -> dict:
     """ Wrapper to create response via TwissResponse """
     LOG.debug("Creating response via TwissResponse.")