Use modifiers for knobs madx (#441)

* remove hardcoded knobs.madx
* zip response output
* save config-files for model creation, response creation
* acc-models-path replacement and code cleanup

Author: JoschD

CHANGELOG.md

@@ -1,5 +1,17 @@
 # OMC3 Changelog
 
+#### 2024-07-08 - v0.15.0 - _jdilly_
+
+- PINNING NUMPY TO < 2.0.0
+
+- Changed:
+  - Model creation:
+    - removed hard-coded `knobs.madx` from `lhc`
+    - removed `corrections.madx` from `lhc` best-knowledge model
+    - zip up log-output files in `response_madx.py`
+    - keep 0th output file in `response_madx.py` for reference of the model setup
+    - Sequence and modifiers use the acc-models symlink in madx-jobs where applicable.
+
 #### 2024-06-05 - v0.14.1 - _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.14.1"
+__version__ = "0.15.0"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"

omc3/correction/response_madx.py

@@ -18,6 +18,7 @@
 from typing import Dict, Sequence, Tuple, List
 
 import numpy as np
+import zipfile
 import pandas as pd
 import tfs
 from optics_functions.coupling import coupling_via_cmatrix
@@ -139,9 +140,16 @@ def _clean_up(temp_dir: Path, num_proc: int) -> None:
         log_path = job_path.with_name(f"{job_path.name}.log")
         full_log += log_path.read_text()
         log_path.unlink()
-        job_path.unlink()
-    full_log_path = temp_dir / "response_madx_full.log"
-    full_log_path.write_text(full_log)
+        if index:  # keep 0th for reference
+            job_path.unlink()
+    
+    # write compressed full log file
+    full_log_name = "response_madx_full.log"
+    zipfile.ZipFile(
+        temp_dir / f"{full_log_name}.zip", 
+        mode="w", 
+        compression=zipfile.ZIP_DEFLATED
+    ).writestr(full_log_name, full_log)
 
 
 def _load_madx_results(

omc3/knob_extractor.py

@@ -33,8 +33,6 @@
     Specify user-defined output path. This should probably be
     `model_dir/knobs.madx`
 
-    default: ``knobs.madx``
-
 
 - **state**:
 
@@ -461,16 +459,16 @@ def _write_knobsfile(output: Union[Path, str], collected_knobs: tfs.TfsDataFrame
 def _get_knobs_def_file(user_defined: Optional[Union[Path, str]] = None) -> Path:
     """ Check which knobs-definition file is appropriate to take. """
     if user_defined is not None:
-        LOGGER.info(f"Using user defined knobs.txt: '{user_defined}")
+        LOGGER.info(f"Using user knobs-definition file: '{user_defined}")
         return Path(user_defined)
 
     if KNOBS_FILE_ACC_MODELS.is_file():
-        LOGGER.info(f"Using model folder's knobs.txt: '{KNOBS_FILE_ACC_MODELS}")
+        LOGGER.info(f"Using given acc-models folder's knobs.txt as knobsdefinition file: '{KNOBS_FILE_ACC_MODELS}")
         return KNOBS_FILE_ACC_MODELS
 
     if KNOBS_FILE_AFS.is_file():
         # if all fails, fall back to lhc acc-models
-        LOGGER.info(f"Using fallback knobs.txt: '{KNOBS_FILE_AFS}'")
+        LOGGER.info(f"Using afs-fallback acc-models folder's knobs.txt as knobs-definition file: '{KNOBS_FILE_AFS}'")
         return KNOBS_FILE_AFS
 
     raise FileNotFoundError("None of the knobs-definition files are available.")

omc3/model/accelerators/lhc.py

@@ -311,29 +311,6 @@ def log_status(self) -> None:
         LOGGER.info(f"> Driven Tune X     [{self.drv_tunes[0]:10.3f}]")
         LOGGER.info(f"> Driven Tune Y     [{self.drv_tunes[1]:10.3f}]")
 
-    def load_main_seq_madx(self) -> str:
-        if self.acc_model_path is not None:
-            main_call = f'call, file = \'{self.acc_model_path / "lhc.seq"}\';'
-            if self.year.startswith('hl'):
-                main_call += f'\ncall, file = \'{self.acc_model_path / "hllhc_sequence.madx"}\';'
-            return main_call
-        try:
-            return _get_call_main_for_year(self.year)
-        except AttributeError:
-            raise AcceleratorDefinitionError(
-                "The accelerator definition is incomplete, mode "
-                "has to be specified (--lhcmode option missing?)."
-            )
-
-    # Private Methods ##########################################################
-
-    def _get_triplet_correctors_file(self) -> Path:
-        correctors_dir = LHC_DIR / self.correctors_dir / "correctors"
-        return correctors_dir / "triplet_correctors.json"
-
-    def _get_corrector_elems(self) -> Path:
-        correctors_dir = LHC_DIR / self.correctors_dir / "correctors"
-        return correctors_dir / f"corrector_elems_b{self.beam}.tfs"
 
     def get_exciter_bpm(self, plane: str, commonbpms: List[str]):
         beam = self.beam
@@ -384,30 +361,11 @@ def get_synch_BPMs(self, index):
         elif self.beam == 2:
             return [i in index for i in self.model.loc["BPMSW.33R8.B2":].index]
 
-    def _get_madx_script_info_comments(self) -> str:
-        info_comments = (
-            f'title, "LHC Model created by omc3";\n'
-            f"! Model directory: {Path(self.model_dir).absolute()}\n"
-            f"! Natural Tune X         [{self.nat_tunes[0]:10.3f}]\n"
-            f"! Natural Tune Y         [{self.nat_tunes[1]:10.3f}]\n"
-            f"! Best Knowledge:        [{'NO' if self.model_best_knowledge is None else 'OK':>10s}]\n"
-        )
-        if self.excitation == AccExcitationMode.FREE:
-            info_comments += f"! Excitation             [{'NO':>10s}]\n\n"
-            return info_comments
-        else:
-            info_comments += (
-                f"! Excitation             [{'ACD' if self.excitation == AccExcitationMode.ACD else 'ADT':>10s}]\n"
-                f"! > Driven Tune X        [{self.drv_tunes[0]:10.3f}]\n"
-                f"! > Driven Tune Y        [{self.drv_tunes[1]:10.3f}]\n\n"
-            )
-        return info_comments
+    # MAD-X Methods ############################################################
 
-    def get_base_madx_script(self, best_knowledge: bool = False) -> str:
-        ats_md = False
-        high_beta = False
+    def get_base_madx_script(self, best_knowledge: bool = False, ats_md: bool = False, high_beta: bool = False) -> str:
         madx_script = (
-            f"{self._get_madx_script_info_comments()}"
+            f"{self._get_madx_script_info_comments()}\n\n"
             f"call, file = '{self.model_dir / MACROS_DIR / GENERAL_MACROS}';\n"
             f"call, file = '{self.model_dir / MACROS_DIR / LHC_MACROS}';\n"
         )
@@ -421,10 +379,12 @@ def get_base_madx_script(self, best_knowledge: bool = False) -> str:
                 f"call, file = '{self.model_dir / MACROS_DIR / LHC_MACROS_RUN3}';\n"
             )
 
-        madx_script += "! ----- Calling Sequence and Optics -----\n"
+        madx_script += "\n! ----- Calling Sequence -----\n"
         madx_script += "option, -echo;  ! suppress output from base sequence loading to keep the log small\n"
-        madx_script += self.load_main_seq_madx()
-        madx_script += "\n\n"
+        madx_script += self._get_madx_main_sequence_loading()
+        madx_script += "\noption, echo;  ! re-enable output to see the optics settings\n"
+        
+        madx_script += "\n! ---- Call optics and other modifiers ----\n"
 
         if self.modifiers is not None:
             madx_script += "".join(
@@ -434,6 +394,7 @@ def get_base_madx_script(self, best_knowledge: bool = False) -> str:
 
         if self.year in ['2012', '2015', '2016', '2017', '2018', '2021', 'hllhc1.3']:
             # backwards compatibility with pre acc-models optics
+            # WARNING: This might override values extracted via the knob-extractor.
             madx_script += (
                 f"\n! ----- Defining Configuration Specifics -----\n"
                 f"xing_angles = {'1' if self.xing else '0'};\n"
@@ -443,14 +404,13 @@ def get_base_madx_script(self, best_knowledge: bool = False) -> str:
                 f"    exec, set_default_crossing_scheme();\n"
                 f"}}\n"
             )
-        else:
-            madx_script += 'call, file="knobs.madx";\n\n'
 
         madx_script += (
+            "\n! ----- Finalize Sequence -----\n"
             "exec, cycle_sequences();\n"
             f"use, sequence = LHCB{self.beam};\n"
-            f"option, echo;\n"
         )
+
         if best_knowledge:
             # madx_script += f"exec, load_average_error_table({self.energy}, {self.beam});\n"
             madx_script += (
@@ -459,6 +419,7 @@ def get_base_madx_script(self, best_knowledge: bool = False) -> str:
                 f"seterr, table=errtab;\n"
                 f"call, file = '{self.model_dir / B2_SETTINGS_MADX}';\n"
             )
+            
         if high_beta:
             madx_script += "exec, high_beta_matcher();\n"
 
@@ -480,6 +441,39 @@ def get_base_madx_script(self, best_knowledge: bool = False) -> str:
             madx_script += "exec, full_response_ats();\n"
 
         return madx_script
+        
+    def _get_madx_script_info_comments(self) -> str:
+        info_comments = (
+            f'title, "LHC Model created by omc3";\n'
+            f"! Model directory: {Path(self.model_dir).absolute()}\n"
+            f"! Natural Tune X         [{self.nat_tunes[0]:10.3f}]\n"
+            f"! Natural Tune Y         [{self.nat_tunes[1]:10.3f}]\n"
+            f"! Best Knowledge:        [{'NO' if self.model_best_knowledge is None else 'YES':>10s}]\n"
+        )
+        if self.excitation == AccExcitationMode.FREE:
+            info_comments += f"! Excitation             [{'NO':>10s}]\n"
+        else:
+            info_comments += (
+                f"! Excitation             [{'ACD' if self.excitation == AccExcitationMode.ACD else 'ADT':>10s}]\n"
+                f"! > Driven Tune X        [{self.drv_tunes[0]:10.3f}]\n"
+                f"! > Driven Tune Y        [{self.drv_tunes[1]:10.3f}]\n"
+
+            )
+        return info_comments
+
+    def _get_madx_main_sequence_loading(self) -> str:
+        if self.acc_model_path is not None:
+            main_call = f'call, file = \'{self.acc_model_path / "lhc.seq"}\';'
+            if self.year.startswith('hl'):
+                main_call += f'\ncall, file = \'{self.acc_model_path / "hllhc_sequence.madx"}\';'
+            return main_call
+        try:
+            return _get_call_main_for_year(self.year)
+        except AttributeError:
+            raise AcceleratorDefinitionError(
+                "The accelerator definition is incomplete, mode "
+                "has to be specified (--lhcmode option missing?)."
+            )
 
     def get_update_correction_script(self, outpath: Union[Path, str], corr_files: Sequence[Union[Path, str]]) -> str:
         madx_script = self.get_base_madx_script()
@@ -488,6 +482,8 @@ def get_update_correction_script(self, outpath: Union[Path, str], corr_files: Se
         madx_script += f"exec, do_twiss_elements(LHCB{self.beam}, '{str(outpath)}', {self.dpp});\n"
         return madx_script
 
+    # Private Methods ##########################################################
+
     def _uses_ats_knobs(self) -> bool:
         """
         Returns wether the ATS knobs and macros should be used, based on the instance's year.
@@ -507,6 +503,14 @@ def _uses_run3_macros(self) -> bool:
         except ValueError:  # if a "hllhc1.x" year is given
             return False
 
+    def _get_triplet_correctors_file(self) -> Path:
+        correctors_dir = LHC_DIR / self.correctors_dir / "correctors"
+        return correctors_dir / "triplet_correctors.json"
+
+    def _get_corrector_elems(self) -> Path:
+        correctors_dir = LHC_DIR / self.correctors_dir / "correctors"
+        return correctors_dir / f"corrector_elems_b{self.beam}.tfs"
+
 
 # General functions ##########################################################
 

omc3/model/constants.py

@@ -38,6 +38,6 @@
 PLANE_TO_HV = dict(X="H", Y="V")
 
 AFS_ACCELERATOR_MODEL_REPOSITORY = Path("/afs/cern.ch/eng/acc-models")
-MODIFIER_SUBDIR = "operation/optics"
+OPTICS_SUBDIR = "operation/optics"
 
 AFS_B2_ERRORS_ROOT = Path("/afs/cern.ch/eng/sl/lintrack/error_tables/")