Fix LHC BestKnowledge Model fetcher (#487)

* b2 error check
* modified for abs/rel paths, added source

Author: JoschD

CHANGELOG.md

@@ -1,5 +1,10 @@
 # OMC3 Changelog
 
+#### 2025-03-06 - v0.22.1 - _jdilly_
+
+- Fixed:
+  - Error in list b2 error choices for lhc best knowledge models.
+
 #### 2025-03-05 - v0.22.0 - _jdilly_
 
 - Added:

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.22.0"
+__version__ = "0.22.1"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"

omc3/model/accelerators/lhc.py

@@ -144,11 +144,6 @@ def get_parameters():
             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):
@@ -158,7 +153,6 @@ def __init__(self, *args, **kwargs):
         self.year = opt.year
         self.ats = opt.ats
         self.b2_errors = opt.b2_errors
-        self.list_b2_errors = opt.list_b2_errors
         self.beam = opt.beam
         beam_to_beam_direction = {1: 1, 2: -1}
         self.beam_direction = beam_to_beam_direction[self.beam]

omc3/model/model_creators/abstract_model_creator.py

@@ -289,55 +289,6 @@ def _find_modifier(self, modifier: Path | str) -> Path:
         raise FileNotFoundError(msg)
 
 
-def check_folder_choices(
-    parent: Path,
-    msg: str,
-    selection: str,
-    list_choices: bool = False,
-    predicate=iotools.always_true,
-    ) -> Path:
-    """
-    A helper function that scans a selected folder for children, which will then be displayed as possible choices.
-    This funciton allows the model-creator to get only the file/folder names, check
-    in the desired folder if the choice is present and return the full path to the selected folder.
-    
-    Args:
-        parent (Path): The folder to scan.
-        msg (str): The message to display, on failure.
-        selection (str): The current selection.
-        list_choices (bool): Whether to just list the choices. 
-                             In that case `None` is returned, instead of an error
-        predicate (callable): A function that takes a path and returns True.
-                              if the path results in a valid choice.
-    
-    Returns:
-       Path: Full path of the selected choice in `parent`.
-
-    Examples:
-        Let's say we expect a choice for a sequence file in the folder `model_root`.
-
-        ```
-        check_folder_choices(model_root, "Expected sequence file", predicate=lambda p: p.suffix == ".seq")
-        ```
-
-        Or we want all subfolder of `scenarios`
-
-        ```
-        check_folder_choices(scenarios, "Expected scenario folder", predicate=lambda p: p.is_dir())
-        ```
-    """
-    choices = [d.name for d in parent.iterdir() if predicate(d)]
-
-    if selection in choices:
-        return parent / selection
-
-    if list_choices:
-        for choice in choices:
-            print(choice)
-    raise AcceleratorDefinitionError(f"{msg}.\nSelected: '{selection}'.\nChoices: [{', '.join(choices)}]")
-
-
-
 class SegmentCreator(ModelCreator, ABC):
     """ Model creator for Segments, to be used in the Segment-by-Segment algorithm. 
     These segments propagate the measured values from the beginning of the segment to the end.
@@ -435,3 +386,55 @@ def __init__(self, accel: Accelerator, twiss_out: Path | str, corr_files: Sequen
         self.logfile= self.twiss_out.parent.absolute() / f"job.create_{self.twiss_out.stem}.log"
         self.corr_files = corr_files
         self.update_dpp = update_dpp
+
+
+# Helper functions -------------------------------------------------------------
+
+def check_folder_choices(
+    parent: Path,
+    msg: str,
+    selection: str,
+    list_choices: bool = False,
+    predicate=iotools.always_true,
+    stem_only: bool = False,
+    ) -> Path | str:
+    """
+    A helper function that scans a selected folder for children, which will then be displayed as possible choices.
+    This funciton allows the model-creator to get only the file/folder names, check
+    in the desired folder if the choice is present and return the full path to the selected folder.
+    
+    Args:
+        parent (Path): The folder to scan.
+        msg (str): The message to display, on failure.
+        selection (str): The current selection.
+        list_choices (bool): Whether to just list the choices. 
+                             In that case `None` is returned, instead of an error
+        predicate (callable): A function that takes a path and returns True.
+                              if the path results in a valid choice.
+        stem_only (bool): If True, only the stem of the path is checked.
+    
+    Returns:
+       Path: Full path of the selected choice in `parent`. 
+
+    Examples:
+        Let's say we expect a choice for a sequence file in the folder `model_root`.
+
+        ```
+        check_folder_choices(model_root, "Expected sequence file", predicate=lambda p: p.suffix == ".seq")
+        ```
+
+        Or we want all subfolder of `scenarios`
+
+        ```
+        check_folder_choices(scenarios, "Expected scenario folder", predicate=lambda p: p.is_dir())
+        ```
+    """
+    choices = [d.stem if stem_only else d.name for d in parent.iterdir() if predicate(d)]
+
+    if selection in choices:
+        return parent / selection
+
+    if list_choices:
+        for choice in choices:
+            print(choice)
+    raise AcceleratorDefinitionError(f"{msg}.\nSelected: '{selection}'.\nChoices: [{', '.join(choices)}]")

omc3/model/model_creators/lhc_model_creator.py

@@ -51,7 +51,8 @@
     SegmentCreator,
     check_folder_choices,
 )
-from omc3.utils.iotools import create_dirs, get_check_suffix_func
+from omc3.optics_measurements.constants import NAME
+from omc3.utils.iotools import create_dirs, get_check_by_regex_func, get_check_suffix_func
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -104,7 +105,7 @@ def prepare_options(self, opt):
             return
 
         # list optics choices ---
-        if opt.list_choices:  # assumes we want to list optics. Invoked even if modifiers are given!
+        if opt.list_choices:  # assumes we only want to list optics. Invoked even if modifiers are given!
             check_folder_choices(  
                 acc_model_path / OPTICS_SUBDIR,
                 msg="No modifier given",
@@ -399,16 +400,32 @@ class LhcBestKnowledgeCreator(LhcModelCreator):  # -----------------------------
     jobfile: str = JOB_MODEL_MADX_BEST_KNOWLEDGE
     files_to_check: tuple[str] = (TWISS_BEST_KNOWLEDGE_DAT, TWISS_ELEMENTS_BEST_KNOWLEDGE_DAT)
 
-    def check_options(self, opt) -> bool:
+    def prepare_options(self, opt) -> bool:
         accel: Lhc = self.accel
-        if accel.list_b2_errors:
-            errors_dir = AFS_B2_ERRORS_ROOT / f"Beam{accel.beam}"
-            for d in errors_dir.iterdir():
-                if d.suffix==".errors" and d.name.startswith("MB2022"):
-                    print(d.stem)
-            return False
 
-        return super().check_options(opt)
+        # Check given B2 errors ---
+        exists = False
+        if accel.b2_errors is not None: 
+            # check if user gave absolute path (! with_suffix does not work, due to '.' in name)
+            exists = Path(f"{accel.b2_errors!s}.errors").is_file()
+
+        if not exists:
+            afs_dir = AFS_B2_ERRORS_ROOT / f"Beam{accel.beam}"
+            if not AFS_B2_ERRORS_ROOT.is_dir():
+                raise AcceleratorDefinitionError(
+                    f"Given b2 errors do not exist and neither does the AFS path: '{afs_dir!s}'."
+                )
+
+            accel.b2_errors = check_folder_choices(  # returns full path
+                afs_dir,
+                msg="No valid b2 errors given.",
+                selection=accel.b2_errors, 
+                list_choices=opt.list_choices,
+                predicate=get_check_by_regex_func(r"MB2022.+\.errors$"),
+                stem_only=True,
+            )  # raises AcceleratorDefinitionError
+
+        return super().prepare_options(opt)
 
     def check_accelerator_instance(self):
         accel: Lhc = self.accel
@@ -430,24 +447,29 @@ def prepare_run(self):
 
         LOGGER.debug("Copying B2 error tables")
 
-        b2_error_path = AFS_B2_ERRORS_ROOT / f"Beam{accel.beam}" / f"{accel.b2_errors}.errors"
-        b2_madx_path = AFS_B2_ERRORS_ROOT / f"Beam{accel.beam}" / f"{accel.b2_errors}.madx"
-        shutil.copy(
-            b2_madx_path,
-            accel.model_dir / B2_SETTINGS_MADX,
-        )
-        b2_table = tfs.read(b2_error_path, index="NAME")
+        b2_stem_path = Path(f"{accel.b2_errors!s}.fakesuffix")  # makes `with_suffix`` work
+        # hint: if b2_stem_path is absolute, the following AFS parts are ignored
+        b2_error_path = AFS_B2_ERRORS_ROOT / f"Beam{accel.beam}" / b2_stem_path.with_suffix(".errors")
+        b2_madx_path = AFS_B2_ERRORS_ROOT / f"Beam{accel.beam}" / b2_stem_path.with_suffix(".madx")
+
+        # copy madx ---
+        output_madx = accel.model_dir / B2_SETTINGS_MADX
+        content = f"! Source: {b2_madx_path}\n{b2_madx_path.read_text()}"
+        output_madx.write_text(content)
+
+        # copy errors table ---
+        b2_table = tfs.read(b2_error_path, index=NAME)
         gen_df = pd.DataFrame(
-            data=np.zeros((b2_table.index.size, len(_b2_columns()))),
+            data=0.,
             index=b2_table.index,
             columns=_b2_columns(),
         )
         gen_df["K1L"] = b2_table.loc[:, "K1L"].to_numpy()
         tfs.write(
             accel.model_dir / B2_ERRORS_TFS,
             gen_df,
-            headers_dict={"NAME": "EFIELD", "TYPE": "EFIELD"},
-            save_index="NAME",
+            headers_dict={"NAME": "EFIELD", "TYPE": "EFIELD", "SOURCE": b2_error_path},
+            save_index=NAME,
         )
 
     def get_madx_script(self) -> str:

omc3/utils/iotools.py

@@ -264,8 +264,13 @@ def always_true(*args, **kwargs) -> bool:
 
 
 def get_check_suffix_func(suffix: str) -> Callable[[Path],bool]:
-    """ Returns a function that checks the suffix of a given path agains 
-    the suffix. """
+    """ Returns a function that checks the suffix of a given path against the suffix. """
     def check_suffix(path: Path) -> bool:
         return path.suffix == suffix
     return check_suffix
+
+def get_check_by_regex_func(pattern: str) -> Callable[[Path],bool]:
+    """ Returns a function that checks the name of a given path against the pattern. """
+    def check(path: Path) -> bool:
+        return re.match(pattern, path.name) is not None
+    return check