Acc-Models Creation (#432)

- Added:

    - complete overhaul of model creation, uses now acc-models for LHC, PS and PSB and prints useful information about available model parameters. Can load from either a user defined path (--path <PATH>) or from the afs copy of acc-models (--afs)


---------

Co-authored-by: JoschD <[email protected]>

Author: awegsche

CHANGELOG.md

@@ -1,5 +1,12 @@
 # OMC3 Changelog
 
+#### 2023-12-07 - v0.13.0 - _awegsche_
+
+- Added:
+ - complete overhaul of model creation, uses now `acc-models` for LHC, PS and PSB and prints
+   useful information about available model parameters. Can load from either a user defined path
+   (`--path <PATH>`) or from the afs copy of acc-models (`--afs`)
+
 #### 2023-11-29 - v0.12.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.12.1"
+__version__ = "0.13.0"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"

omc3/model/accelerators/accelerator.py

@@ -6,6 +6,7 @@
 It contains entrypoint the parent `Accelerator` class as well as other support classes.
 """
 import re
+import os
 from pathlib import Path
 from typing import List, Union, Sequence
 
@@ -16,7 +17,7 @@
 
 from omc3.model.constants import (
     ERROR_DEFFS_TXT,
-    JOB_MODEL_MADX,
+    JOB_MODEL_MADX_NOMINAL,
     MODIFIER_TAG,
     MODIFIERS_MADX,
     TWISS_AC_DAT,
@@ -59,6 +60,8 @@ class Accelerator:
         AccElementTypes.ARC_BPMS: r".*",
     }
     BPM_INITIAL = "B"
+    NAME=None
+    REPOSITORY=None
 
     @staticmethod
     def get_parameters():
@@ -95,11 +98,11 @@ def get_parameters():
         params.add_parameter(
             name="energy",
             type=float,
-            help="Energy in Tev.",
+            help="Energy in GeV.",
         )
         params.add_parameter(
             name="modifiers",
-            type=Path,
+            type=PathOrStr,
             nargs="*",
             help="Path to the optics file to use (modifiers file).",
         )
@@ -118,6 +121,7 @@ def __init__(self, opt) -> None:
         self.model_best_knowledge = None
         self.elements = None
         self.error_defs_file = None
+        self.acc_model_path = None
         self.modifiers = None
         self._beam_direction = 1
         self._beam = None
@@ -137,11 +141,7 @@ def __init__(self, opt) -> None:
         else:
             self.init_from_options(opt)
 
-    def init_from_options(self, opt) -> None:
-        if opt.nat_tunes is None:
-            raise AcceleratorDefinitionError("Argument 'nat_tunes' is required.")
-        if (opt.drv_tunes is None) and (opt.driven_excitation is not None):
-            raise AcceleratorDefinitionError("Argument 'drv_tunes' is required.")
+    def init_from_options(self, opt):
         self.nat_tunes = opt.nat_tunes
 
         if opt.driven_excitation is not None:
@@ -193,6 +193,17 @@ def init_from_model_dir(self, model_dir: Path) -> None:
         if best_knowledge_path.is_file():
             self.model_best_knowledge = tfs.read(best_knowledge_path, index="NAME")
 
+        # Base Model ########################################
+        if self.REPOSITORY is not None:
+            acc_models = model_dir / self.REPOSITORY
+
+            if acc_models.is_dir():
+                if acc_models.is_symlink():
+                    self.acc_model_path = Path(os.readlink(acc_models)).absolute()
+                else:
+                    self.acc_model_path = acc_models
+            # else this wasn't an acc-models based model
+
         # Modifiers #########################################
         self.modifiers = _get_modifiers_from_modeldir(model_dir)
 
@@ -256,7 +267,12 @@ def verify_object(self):
         """
         Verifies that this instance of an `Accelerator` is properly instantiated.
         """
-        raise NotImplementedError("A function should have been overwritten, check stack trace.")
+        # since we removed `required` args, we check here if everything has been passed
+        if self.model_dir is None:
+            if self.nat_tunes is None:
+                raise AttributeError("Natural tunes not set (missing `--nat_tunes` flag?)")
+            if self.excitation != AccExcitationMode.FREE and self.drv_tunes is None:
+                raise AttributeError("Driven excitation selected but no driven tunes given (missing `--drv_tunes` flag?)")
 
     def get_exciter_bpm(self, plane: str, commonbpms: List[str]):
         """
@@ -355,13 +371,16 @@ class AcceleratorDefinitionError(Exception):
 
 def _get_modifiers_from_modeldir(model_dir: Path) -> List[Path]:
     """Parse modifiers from job.create_model.madx or use modifiers.madx file."""
-    job_file = model_dir / JOB_MODEL_MADX
+    job_file = model_dir / JOB_MODEL_MADX_NOMINAL
     if job_file.exists():
         job_madx = job_file.read_text()
 
         # find modifier tag in lines and return called file in these lines
+        # the modifier tag is used by the model creator to mark which line defines modifiers
+        # see e.g. `get_base_madx_script()` in `lhc.py`
+        # example for a match to the regex: `call, file = 'modifiers.madx'; MODIFIER_TAG`
         modifiers = re.findall(
-            fr"\s+call,\s*file\s*=\s*[\"\']?([^;\'\"]+)[\"\']?\s*;\s*{MODIFIER_TAG}",
+            fr"\s*call,\s*file\s*=\s*[\"\']?([^;\'\"]+)[\"\']?\s*;\s*{MODIFIER_TAG}",
             job_madx,
             flags=re.IGNORECASE,
         )

omc3/model/accelerators/lhc.py

@@ -113,6 +113,7 @@ class Lhc(Accelerator):
     """
 
     NAME = "lhc"
+    REPOSITORY = "acc-models-lhc"
     RE_DICT: Dict[str, str] = {
         AccElementTypes.BPMS: r"BPM",
         AccElementTypes.MAGNETS: r"M",
@@ -132,15 +133,23 @@ def get_parameters():
         params.add_parameter(
             name="year",
             type=str,
-            required=True,
-            choices=("2012", "2015", "2016", "2017", "2018", "2022", "hllhc1.3"),
             help="Year of the optics (or hllhc1.x version).",
         )
         params.add_parameter(
             name="ats",
             action="store_true",
             help="Force use of ATS macros and knobs for years which are not ATS by default.",
-            )
+        )
+        params.add_parameter(
+            name="b2_errors",
+            type=str,
+            help="The B2 error table to load for the best knowledge model.",
+        )
+        params.add_parameter(
+            name="list_b2_errors",
+            action="store_true",
+            help="Lists all available b2 error tables",
+        )
         return params
 
     def __init__(self, *args, **kwargs):
@@ -150,36 +159,38 @@ def __init__(self, *args, **kwargs):
         self.correctors_dir = "2012"
         self.year = opt.year
         self.ats = opt.ats
+        self.b2_errors = opt.b2_errors
+        self.list_b2_errors = opt.list_b2_errors
         if self.year == "hllhc1.3":
             self.correctors_dir = "hllhc1.3"
         self.beam = opt.beam
         beam_to_beam_direction = {1: 1, 2: -1}
         self.beam_direction = beam_to_beam_direction[self.beam]
-        self.verify_object()
 
     def verify_object(self) -> None:  # TODO: Maybe more checks?
         """
         Verifies if everything is defined which should be defined.
         Will Raise an ``AcceleratorDefinitionError`` if one of the checks is invalid.
         """
         LOGGER.debug("Accelerator class verification")
+
+        Accelerator.verify_object(self)
         _ = self.beam
 
         if self.model_dir is None and self.xing is None:
             raise AcceleratorDefinitionError("Crossing on or off not set.")
 
-        if self.excitation is None:
-            raise AcceleratorDefinitionError("Excitation mode not set.")
-        if (self.excitation != AccExcitationMode.FREE) and (self.drv_tunes is None):
-            raise AcceleratorDefinitionError("An excitation mode was given but driven tunes are not set.")
-
         # TODO: write more output prints
-        LOGGER.debug("... verification passed. \nSome information about the accelerator:")
+        LOGGER.debug(
+            "... verification passed. \nSome information about the accelerator:"
+        )
         LOGGER.debug(f"Class name       {self.__class__.__name__}")
         LOGGER.debug(f"Beam             {self.beam}")
         LOGGER.debug(f"Beam direction   {self.beam_direction}")
         if self.modifiers:
-            LOGGER.debug(f"Modifiers        {', '.join([str(m) for m in self.modifiers])}")
+            LOGGER.debug(
+                f"Modifiers        {', '.join([str(m) for m in self.modifiers])}"
+            )
 
     @property
     def beam(self) -> int:
@@ -216,7 +227,11 @@ def get_variables(self, frm: float = None, to: float = None, classes: Iterable[s
             LOGGER.debug("The following classes are not found as corrector/variable classes and "
                          f"are assumed to be the variable names directly instead:\n{str(unknown_classes)}")
 
-        vars = list(set(_flatten_list(all_vars_by_class[corr_cls] for corr_cls in known_classes)))
+        vars = list( set(
+            _flatten_list(
+                all_vars_by_class[corr_cls]
+                    for corr_cls in known_classes))
+            )
         vars = vars + unknown_classes
 
         # Sort variables by S (nice for comparing different files)
@@ -297,6 +312,11 @@ def log_status(self) -> None:
         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:
@@ -323,19 +343,28 @@ def get_exciter_bpm(self, plane: str, commonbpms: List[str]):
         if self.excitation == AccExcitationMode.ACD:
             try:
                 return (
-                    _is_one_of_in([f"BPMY{a_b}.6L4.B{beam}", f"BPM.7L4.B{beam}"], commonbpms),
+                    _is_one_of_in(
+                        [f"BPMY{a_b}.6L4.B{beam}", f"BPM.7L4.B{beam}"], commonbpms
+                    ),
                     f"MKQA.6L4.B{beam}",
                 )
             except KeyError as e:
-                raise KeyError("AC-Dipole BPM not found in the common BPMs. Maybe cleaned?") from e
+                raise KeyError(
+                    "AC-Dipole BPM not found in the common BPMs. Maybe cleaned?"
+                ) from e
         if self.excitation == AccExcitationMode.ADT:
             try:
                 return (
-                    _is_one_of_in([f"BPMWA.B5{l_r}4.B{beam}", f"BPMWA.A5{l_r}4.B{beam}"], commonbpms),
+                    _is_one_of_in(
+                        [f"BPMWA.B5{l_r}4.B{beam}", f"BPMWA.A5{l_r}4.B{beam}"],
+                        commonbpms,
+                    ),
                     f"ADTK{adt}5{l_r}4.B{beam}",
                 )
             except KeyError as e:
-                raise KeyError("ADT BPM not found in the common BPMs. Maybe cleaned?") from e
+                raise KeyError(
+                    "ADT BPM not found in the common BPMs. Maybe cleaned?"
+                ) from e
         return None
 
     def important_phase_advances(self) -> List[List[str]]:
@@ -379,34 +408,46 @@ def get_base_madx_script(self, best_knowledge: bool = False) -> str:
         high_beta = False
         madx_script = (
             f"{self._get_madx_script_info_comments()}"
-            f"! ----- Calling Sequence and Optics -----\n"
             f"call, file = '{self.model_dir / MACROS_DIR / GENERAL_MACROS}';\n"
             f"call, file = '{self.model_dir / MACROS_DIR / LHC_MACROS}';\n"
-            )
+        )
+        madx_script += f"omc3_beam_energy = {self.energy};\n"
+        madx_script += "exec, define_nominal_beams();\n\n"
         if self._uses_run3_macros():
-            LOGGER.debug("According to the optics year, Run 3 versions of the macros will be used")
+            LOGGER.debug(
+                "According to the optics year, Run 3 versions of the macros will be used"
+            )
             madx_script += (
                 f"call, file = '{self.model_dir / MACROS_DIR / LHC_MACROS_RUN3}';\n"
             )
 
-        madx_script += (
-            f"{self.load_main_seq_madx()}\n"
-            f"exec, define_nominal_beams();\n"
-        )
+        madx_script += "! ----- Calling Sequence and Optics -----\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"
+
         if self.modifiers is not None:
             madx_script += "".join(
                 f"call, file = '{self.model_dir / modifier}'; {MODIFIER_TAG}\n"
                 for modifier in self.modifiers
             )
+
+        if self.year in ['2012', '2015', '2016', '2017', '2018', '2021', 'hllhc1.3']:
+            # backwards compatibility with pre acc-models optics
+            madx_script += (
+                f"\n! ----- Defining Configuration Specifics -----\n"
+                f"xing_angles = {'1' if self.xing else '0'};\n"
+                f"if(xing_angles==1){{\n"
+                f"    exec, set_crossing_scheme_ON();\n"
+                f"}}else{{\n"
+                f"    exec, set_default_crossing_scheme();\n"
+                f"}}\n"
+            )
+        else:
+            madx_script += 'call, file="knobs.madx";\n\n'
+
         madx_script += (
-            f"\n! ----- Defining Configuration Specifics -----\n"
-            f"exec, cycle_sequences();\n"
-            f"xing_angles = {'1' if self.xing else '0'};\n"
-            f"if(xing_angles==1){{\n"
-            f"    exec, set_crossing_scheme_ON();\n"
-            f"}}else{{\n"
-            f"    exec, set_default_crossing_scheme();\n"
-            f"}}\n"
+            "exec, cycle_sequences();\n"
             f"use, sequence = LHCB{self.beam};\n"
             f"option, echo;\n"
         )
@@ -422,14 +463,19 @@ def get_base_madx_script(self, best_knowledge: bool = False) -> str:
             madx_script += "exec, high_beta_matcher();\n"
 
         madx_script += f"\n! ----- Matching Knobs and Output Files -----\n"
-        if self._uses_ats_knobs():
-            LOGGER.debug("According to the optics year or the --ats flag being provided, ATS macros and knobs will be used")
-            madx_script += f"exec, match_tunes_ats({self.nat_tunes[0]}, {self.nat_tunes[1]}, {self.beam});\n"
-            madx_script += f"exec, coupling_knob_ats({self.beam});\n"
-        else:
-            madx_script += f"exec, match_tunes({self.nat_tunes[0]}, {self.nat_tunes[1]}, {self.beam});\n"
-            madx_script += f"exec, coupling_knob({self.beam});\n"
-        
+
+        # in the best knowledge case, all knobs are loaded from actual knowledge
+        if not best_knowledge:
+            if self._uses_ats_knobs():
+                LOGGER.debug(
+                    "According to the optics year or the --ats flag being provided, ATS macros and knobs will be used"
+                )
+                madx_script += f"exec, match_tunes_ats({self.nat_tunes[0]}, {self.nat_tunes[1]}, {self.beam});\n"
+                madx_script += f"exec, coupling_knob_ats({self.beam});\n"
+            else:
+                madx_script += f"exec, match_tunes({self.nat_tunes[0]}, {self.nat_tunes[1]}, {self.beam});\n"
+                madx_script += f"exec, coupling_knob({self.beam});\n"
+
         if ats_md:
             madx_script += "exec, full_response_ats();\n"
 
@@ -461,13 +507,16 @@ def _uses_run3_macros(self) -> bool:
         except ValueError:  # if a "hllhc1.x" year is given
             return False
 
+
 # General functions ##########################################################
 
 
 def _get_call_main_for_year(year: str) -> str:
     call_main = f"call, file = '{_get_file_for_year(year, 'main.seq')}';\n"
     if year == "2012":
-        call_main += f"call, file = '{LHC_DIR / '2012' / 'install_additional_elements.madx'}';\n"
+        call_main += (
+            f"call, file = '{LHC_DIR / '2012' / 'install_additional_elements.madx'}';\n"
+        )
     if year == "hllhc1.3":
         call_main += f"call, file = '{LHC_DIR / 'hllhc1.3' / 'main_update.seq'}';\n"
     return call_main