change model handling

Author: JoschD

ir_amplitude_detuning/lhc_detuning_corrections.py

@@ -39,7 +39,13 @@
 )
 from ir_amplitude_detuning.detuning.measurements import MeasureValue
 from ir_amplitude_detuning.detuning.targets import Target
-from ir_amplitude_detuning.simulation.lhc_simulation import FakeLHCBeam, LHCBeam, LHCCorrectors
+from ir_amplitude_detuning.simulation.lhc_simulation import (
+    ACC_MODELS,
+    FakeLHCBeam,
+    LHCBeam,
+    LHCCorrectors,
+    pathstr,
+)
 from ir_amplitude_detuning.utilities.constants import (
     AMPDET_CALC_ERR_ID,
     AMPDET_CALC_ID,
@@ -68,12 +74,20 @@
     from ir_amplitude_detuning.detuning.targets import Target
 
 
-LOG = logging.getLogger(__name__)
-
 LHCBeams: TypeAlias = dict[int, LHCBeam]
 LHCBeamsPerXing: TypeAlias = dict[str, LHCBeams]
 
 
+LOG = logging.getLogger(__name__)
+
+
+def get_optics(year: int) -> str:
+    return {
+        2018: pathstr("optics2018", "PROTON", "opticsfile.22_ctpps2"),
+        2022: pathstr(ACC_MODELS, "strengths", "ATS_Nominal", "2022", "squeeze", "ats_30cm.madx")
+    }[year]
+
+
 @dataclass(slots=True)
 class CorrectionResults:
     """Class to store the results of a correction calculation.
@@ -97,7 +111,7 @@ def create_optics(
     outputdir: Path,
     output_id: str = '',
     xing: dict[str, dict] | None = None,  # default set below
-    optics: str = "round3030",  # 30cm round optics
+    optics: str | Path | None = None,  # defaults to 30cm round optics
     year: int = 2018,  # lhc year
     tune_x: float = 62.28,  # horizontal tune
     tune_y: float = 60.31,  # vertical tune
@@ -120,19 +134,15 @@ def create_optics(
     Returns:
         LHCBeams: The LHC beams, i.e. a dictionary of LHCBeam objects.
     """
-    # set mutable defaults ----
-    if xing is None:
-        xing = {'scheme': 'top'}  # use top-energy crossing scheme
-
     # Setup LHC for both beams -------------------------------------------------
     lhc_beams = {}
     for beam in beams:
         output_subdir = get_label_outputdir(outputdir, output_id, beam)
         lhc_beam = LHCBeam(
             beam=beam,
             outputdir=output_subdir,
-            xing=xing,
-            optics=optics,
+            xing=xing or {'scheme': 'top'},  # use top-energy crossing scheme
+            optics=optics or get_optics(year),
             year=year,
             tune_x=tune_x,
             tune_y=tune_y,
@@ -339,7 +349,7 @@ def check_corrections_ptc(
     # Below only needed if lhc_beams is None ---
     beams: Sequence[int] | None = None,
     xing: dict[str, dict] | None = None,
-    optics: str = "round3030",  # 30cm round optics
+    optics: Path | None = None,  # defaults to 30cm round optics
     year: int = 2018,  # lhc year
     tune_x: float = 62.28,  # horizontal tune
     tune_y: float = 60.31,  # vertical tune
@@ -361,7 +371,7 @@ def check_corrections_ptc(
         lhc_beams (dict[int, LHCBeam]): Pre-run LHC beams.
         beams (Sequence[int]): Beams (if ``lhc_beams`` is None).
         xing (dict[str, dict]): Crossing scheme (if ``lhc_beams`` is `None`).
-        optics (str): Optics (if ``lhc_beams`` is `None`).
+        optics (Path): Path to the optics file (if ``lhc_beams`` is `None`).
         year (int): Year (if ``lhc_beams`` is `None`).
         tune_x (float): Horizontal tune (if ``lhc_beams`` is `None`).
         tune_y (float): Vertical tune (if ``lhc_beams`` is `None`).
@@ -371,15 +381,13 @@ def check_corrections_ptc(
         lhc_beams = {}
         if beams is None:
             raise ValueError("Either lhc_beams or beams must be given.")
-        if xing is None:
-            xing = {'scheme': 'top'}
 
         for beam in beams:
             lhc_beam = LHCBeam(
                 beam=beam,
                 outputdir=get_label_outputdir(outputdir, 'tmp_ptc', beam),
-                xing=xing,
-                optics=optics,
+                xing=xing or {'scheme': 'top'},
+                optics=optics or get_optics(year),
                 year=year,
                 tune_x=tune_x,
                 tune_y=tune_y,

ir_amplitude_detuning/simulation/lhc_simulation.py

@@ -37,11 +37,9 @@
 LOG = logging.getLogger(__name__)  # setup in main()
 LOG_LEVEL = logging.DEBUG
 
-ACC_MODELS = "acc-models-lhc"
+ACC_MODELS: str = "acc-models-lhc"
 
-PATHS = {
-    "db5": Path("/afs/cern.ch/eng/lhc/optics/V6.503"),
-    "optics2016": Path("/afs/cern.ch/eng/lhc/optics/runII/2016"),
+PATHS: dict[str, Path] = {
     "optics2018": Path("/afs/cern.ch/eng/lhc/optics/runII/2018"),
     "optics_repo": Path("/afs/cern.ch/eng/acc-models/lhc"),
     ACC_MODELS: Path(ACC_MODELS),
@@ -62,45 +60,6 @@ def pathstr(key: str, *args: str) -> str:
     return str(PATHS[key].joinpath(*args))
 
 
-def get_optics_path(year: int, name: str | Path):
-    """Get optics by name, i.e. a collection of optics path-strings to the optics files.
-
-     Args:
-         year (int): Year of the optics
-         name (str, Path): Name for the optics or a path to the optics file.
-
-    Returns:
-        str: Path to the optics file.
-     """
-    if isinstance(name, Path):
-        return str(name)
-
-    # Predefined optics paths ---
-    optics_map = {
-        2018: {
-            'inj': pathstr("optics2018", "PROTON", "opticsfile.1"),
-            'flat6015': pathstr("optics2018", 'MDflatoptics2018', 'opticsfile_flattele60cm.21'),
-            'round3030': pathstr("optics2018", "PROTON", "opticsfile.22_ctpps2"),
-        },
-        2022: {
-            'round3030': pathstr(ACC_MODELS, "strengths", "ATS_Nominal", "2022", "squeeze", "ats_30cm.madx")
-        }
-    }
-    return optics_map[year][name]
-
-
-def get_wise_path(seed: int):
-    """Get the wise errordefinition file by seed-number.
-
-    Args:
-        seed (int): Seed for the error realization.
-
-    Returns:
-        str: Path to the wise errortable file.
-    """
-    return pathstr('wise', f"WISE.errordef.{seed:04d}.tfs")
-
-
 def drop_allzero_columns(df: TfsDataFrame, keep: Sequence = ()) -> TfsDataFrame:
     """Drop columns that contain only zeros, to save harddrive space.
 
@@ -145,9 +104,8 @@ class LHCBeam:
     beam: int
     outputdir: Path
     xing: dict
-    optics: str
+    optics: str | Path | None
     year: int = 2018
-    thin: bool = False
     tune_x: float = 62.28
     tune_y: float = 60.31
     chroma: float = 3
@@ -176,6 +134,15 @@ def __post_init__(self):
         # Define Sequence to use
         self.seq_name, self.seq_file, self.bv_flag = get_lhc_sequence_filename_and_bv(self.beam, accel="lhc" if self.year < 2020 else "hllhc")  # `hllhc` just for naming of the sequence file, i.e. without _as_built
 
+        self.path_to_use = "optics2018"
+        if self.year > 2019:  # after 2019, use acc-models
+            self.path_to_use = ACC_MODELS
+
+            acc_models_path = PATHS[ACC_MODELS]
+            if acc_models_path.exists():
+                acc_models_path.unlink()
+            acc_models_path.symlink_to(pathstr("optics_repo", str(self.year)))
+
     # Output Helper ---
     def output_path(self, type_: str, output_id: str, dir_: Path | None = None, suffix: str = ".tfs") -> Path:
         """Returns the output path for standardized tfs names in the default output directory.
@@ -303,30 +270,13 @@ def setup_machine(self):
         Initialized the beam and applies optics, crossing."""
         self.reinstate_loggers()
         madx = self.madx  # shorthand
-        mvars = madx.globals  # shorthand
 
         # Load Macros
-        madx.call(pathstr("optics2018", "toolkit", "macro.madx"))
+        madx.call(pathstr(self.path_to_use, "toolkit", "macro.madx"))
 
         # Lattice Setup ---------------------------------------
         # Load Sequence
-        if self.year > 2019:  # after 2019, use acc-models
-            acc_models_path = PATHS[ACC_MODELS]
-            if acc_models_path.exists():
-                acc_models_path.unlink()
-            acc_models_path.symlink_to(pathstr("optics_repo", str(self.year)))
-            madx.call(pathstr(ACC_MODELS, self.seq_file))
-        else:
-            madx.call(pathstr("optics2018", self.seq_file))
-
-        # Slice Sequence
-        if self.thin:
-            mvars.slicefactor = 4
-            madx.beam()
-            madx.call(pathstr("optics2018", "toolkit", "myslice.madx"))
-            madx.beam()
-            madx.use(sequence=self.seq_name)
-            madx.makethin(sequence=self.seq_name, style="teapot", makedipedge=True)
+        madx.call(pathstr(self.path_to_use, self.seq_file))
 
         # Cycling w.r.t. to IP3 (mandatory to find closed orbit in collision in the presence of errors)
         madx.seqedit(sequence=self.seq_name)
@@ -335,10 +285,7 @@ def setup_machine(self):
         madx.endedit()
 
         # Define Optics and make beam
-        madx.call(get_optics_path(self.year, self.optics))
-        if self.optics == 'inj':
-            mvars.NRJ = 450.000  # not defined in injection optics.1 but in the others
-
+        madx.call(str(self.optics))
         madx.beam(sequence=self.seq_name, bv=self.bv_flag,
                   energy="NRJ", particle="proton", npart=self.n_particles,
                   kbunch=1, ex=self.emittance, ey=self.emittance)

pyproject.toml

@@ -62,6 +62,7 @@ test = [
   "pytest-cov >= 2.9",
   "pytest-timeout >= 1.4",
   "pytest-dependency >= 0.6.0",
+  "gitpython >= 3.1",  # imported as 'git', used for acc-models fixture
 ]
 doc = [
   "sphinx >= 7.0",

tests/conftest.py

@@ -9,11 +9,20 @@
 import sys
 from pathlib import Path
 
+import git
 import pytest
 
+GITLAB_REPO_ACC_MODELS: str = "https://gitlab.cern.ch/acc-models/acc-models-{}.git"
 
 
 def assert_exists_and_not_empty(file_path: Path):
     """Assert that a file exists and is not empty."""
     assert file_path.exists(), f"File {file_path} does not exist."
-    assert file_path.stat().st_size > 0, f"File {file_path} is empty."
+    assert file_path.stat().st_size > 0, f"File {file_path} is empty."
+
+
+def clone_acc_models(tmp_path_factory: pytest.TempPathFactory, accel: str, year: int) -> Path:
+    """ Clone the acc-models directory for the specified accelerator from github into a temporary directory. """
+    tmp_path = tmp_path_factory.mktemp(f"acc-models-{accel}-{year}")
+    git.Repo.clone_from(GITLAB_REPO_ACC_MODELS.format(accel), tmp_path, branch=str(year))
+    return tmp_path