more comments

pylhc · Nov 12, 2024 · aee9255 · aee9255
1 parent 0f15388
commit aee9255
Show file tree

Hide file tree

Showing 13 changed files with 127 additions and 83 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,6 @@
 # OMC3 Changelog
 
-#### 2024-11-12 - v0.19.0 - _fscarlier_, _jdilly_
+#### 2024-11-13 - v0.19.0 - _fscarlier_, _jdilly_
 
 - Added K-Modulation tools:
   - `kmod_importer`: Main function to call the following scripts in succession.

diff --git a/doc/entrypoints/other.rst b/doc/entrypoints/other.rst
@@ -16,21 +16,6 @@ Other
     :noindex:
 
 
-.. automodule:: omc3.kmod_averages
-    :members:
-    :noindex:
-
-
-.. automodule:: omc3.kmod_import
-    :members:
-    :noindex:
-
-
-.. automodule:: omc3.kmod_lumi_imbalance
-    :members:
-    :noindex:
-
-
 .. automodule:: omc3.madx_wrapper
     :members:
     :noindex:

diff --git a/doc/entrypoints/scripts.rst b/doc/entrypoints/scripts.rst
@@ -6,7 +6,17 @@ Scripts
     :noindex:
 
 
-.. automodule:: omc3.scripts.merge_kmod_results
+.. automodule:: omc3.kmod_averages
+    :members:
+    :noindex:
+
+
+.. automodule:: omc3.kmod_import
+    :members:
+    :noindex:
+
+
+.. automodule:: omc3.kmod_lumi_imbalance
     :members:
     :noindex:
 

diff --git a/omc3/kmod_importer.py b/omc3/kmod_importer.py
@@ -67,7 +67,7 @@
     BETA,
     ERR,
     EXT,
-    IMBALACE,
+    IMBALANCE,
     LSA_FILE_NAME,
     LUMINOSITY,
     NAME,
@@ -284,7 +284,7 @@ def calculate_all_lumi_imbalances(
             continue
 
         # Print luminosity imbalance
-        imb, err_imb = df.headers[f"{LUMINOSITY}{IMBALACE}"], df.headers[f"{ERR}{LUMINOSITY}{IMBALACE}"]
+        imb, err_imb = df.headers[f"{LUMINOSITY}{IMBALANCE}"], df.headers[f"{ERR}{LUMINOSITY}{IMBALANCE}"]
         LOG.info(f"Luminosity imbalance between {ipA} and {ipB}: {imb:.2e} +/- {err_imb:.2e}")
 
 

diff --git a/omc3/optics_measurements/constants.py b/omc3/optics_measurements/constants.py
@@ -91,7 +91,7 @@
 # Columns and Column-Prefixes
 EFFECTIVE: str = "EFF"
 LUMINOSITY: str = "LUMI"
-IMBALACE: str = "IMB"
+IMBALANCE: str = "IMBALANCE"
 S_LOCATION: str = "_S_LOCATION"
 
 BEAM: str = "BEAM"
@@ -104,14 +104,14 @@
 KMOD_PHASE_ADV: str = "PHASEADV"
 
 # file names from kmod-application
-BEAM_DIR: str = 'B'
-LSA_FILE_NAME: str = 'lsa_results'  # contains beta-per-BPM (and IP) results
-RESULTS_FILE_NAME: str = 'results'  # contains betastar results
+BEAM_DIR: str = "B"
+LSA_FILE_NAME: str = "lsa_results"  # contains beta-per-BPM (and IP) results
+RESULTS_FILE_NAME: str = "results"  # contains betastar results
 
 # file names defined by omc3
-AVERAGED_BETASTAR_FILENAME: str = 'averaged_ip{ip}_beta{betastar_x:.2f}m{betastar_y:.2f}m'
-AVERAGED_BPM_FILENAME: str = 'averaged_bpm_beam{beam}_ip{ip}_beta{betastar_x:.2f}m{betastar_y:.2f}m'
-EFFECTIVE_BETAS_FILENAME: str = 'effective_betas_beta{betastar_x:.2f}m{betastar_y:.2f}m'
+AVERAGED_BETASTAR_FILENAME: str = "averaged_ip{ip}_beta{betastar_x:.2f}m{betastar_y:.2f}m"
+AVERAGED_BPM_FILENAME: str = "averaged_bpm_beam{beam}_ip{ip}_beta{betastar_x:.2f}m{betastar_y:.2f}m"
+EFFECTIVE_BETAS_FILENAME: str = "effective_betas_beta{betastar_x:.2f}m{betastar_y:.2f}m"
 
-BETA_KMOD_FILENAME: str = 'beta_kmod_'
-BETA_STAR_FILENAME: str = 'betastar_'
+BETA_KMOD_FILENAME: str = "beta_kmod_"
+BETA_STAR_FILENAME: str = "betastar_"
diff --git a/omc3/scripts/kmod_average.py b/omc3/scripts/kmod_average.py
@@ -222,7 +222,7 @@ def get_average_betastar_results(meas_paths: Sequence[Path]) -> tfs.TfsDataFrame
 
         mean_df =  _get_averaged_df(all_dfs)
         if LABEL in mean_df.columns:
-            mean_df[NAME] = f"IP{mean_df[LABEL].iloc[0][-1]}"
+            mean_df[NAME] = mean_df[LABEL].apply(lambda s: f"IP{s[-1]}")
             mean_df = mean_df.drop(columns=[LABEL])
 
         mean_df[BEAM] = beam

diff --git a/omc3/scripts/kmod_import.py b/omc3/scripts/kmod_import.py
@@ -91,7 +91,7 @@ def _get_params() -> EntryPointParameters:
     params.add_parameter(
         name="measurements",
         required=True,
-        nargs='+',
+        nargs="+",
         type=PathOrStrOrDataFrame,
         help="Paths to the K-modulation results files to import. "
              "Can be either the TFS-files directly or a path to a folder containing them."
@@ -131,7 +131,7 @@ def import_kmod_data(opt: DotDict) -> dict[str, tfs.TfsDataFrame]:
             path to the folder containing multiple of such files.
         
         model (Path|str):
-            Path to the model twiss file, or a folder containing 'twiss_elemtents.dat'.
+            Path to the model twiss file, or a folder containing 'twiss_elements.dat'.
             Determines which elements to keep, i.e. `twiss.dat` keeps only the BPMs,
             `twiss_elements.dat` keeps BPMs and IPs.
         
@@ -187,7 +187,7 @@ def convert_bpm_results(
 
     # merge files
     bpm_results_list = [df.set_index(NAME, drop=True) if NAME in df.columns else df for df in bpm_results_list]
-    kmod_results = tfs.concat(bpm_results_list, join='inner',)
+    kmod_results = tfs.concat(bpm_results_list, join="inner",)
     df_model = _sync_model_index(kmod_results, df_model) 
 
     dfs = {}
@@ -196,24 +196,24 @@ def convert_bpm_results(
 
         # copy s and beta
         beta_kmod.loc[:, S] = df_model.loc[:, S]
-        beta_kmod.loc[:, f'{BETA}{plane}{MDL}'] = df_model.loc[:, f'{BETA}{plane}']
-        beta_kmod.loc[:, f'{BETA}{plane}'] = kmod_results.loc[:, f'{BETA}{plane}']
-        beta_kmod.loc[:, f'{ERR}{BETA}{plane}'] = kmod_results.loc[:, f'{ERR}{BETA}{plane}']
+        beta_kmod.loc[:, f"{BETA}{plane}{MDL}"] = df_model.loc[:, f"{BETA}{plane}"]
+        beta_kmod.loc[:, f"{BETA}{plane}"] = kmod_results.loc[:, f"{BETA}{plane}"]
+        beta_kmod.loc[:, f"{ERR}{BETA}{plane}"] = kmod_results.loc[:, f"{ERR}{BETA}{plane}"]
 
         # model-delta and beta-beating
-        beta_kmod.loc[:, f'{DELTA}{BETA}{plane}{MDL}'] = (
-            beta_kmod[f'{BETA}{plane}'] - beta_kmod[f'{BETA}{plane}{MDL}']
+        beta_kmod.loc[:, f"{DELTA}{BETA}{plane}{MDL}"] = (
+            beta_kmod[f"{BETA}{plane}"] - beta_kmod[f"{BETA}{plane}{MDL}"]
         )
-        beta_kmod.loc[:, f'{DELTA}{BETA}{plane}'] = (
-            beta_kmod[f'{DELTA}{BETA}{plane}{MDL}'] / beta_kmod[f'{BETA}{plane}{MDL}']
+        beta_kmod.loc[:, f"{DELTA}{BETA}{plane}"] = (
+            beta_kmod[f"{DELTA}{BETA}{plane}{MDL}"] / beta_kmod[f"{BETA}{plane}{MDL}"]
         )
-        beta_kmod.loc[:, f'{ERR}{DELTA}{BETA}{plane}'] = (
-            beta_kmod[f'{ERR}{BETA}{plane}'] / beta_kmod[f'{BETA}{plane}{MDL}']
+        beta_kmod.loc[:, f"{ERR}{DELTA}{BETA}{plane}"] = (
+            beta_kmod[f"{ERR}{BETA}{plane}"] / beta_kmod[f"{BETA}{plane}{MDL}"]
         )
 
         # tune
-        beta_kmod.headers[f'{TUNE}1'] = df_model.headers[f'{TUNE}1'] % 1
-        beta_kmod.headers[f'{TUNE}2'] = df_model.headers[f'{TUNE}2'] % 1
+        beta_kmod.headers[f"{TUNE}1"] = df_model.headers[f"{TUNE}1"] % 1
+        beta_kmod.headers[f"{TUNE}2"] = df_model.headers[f"{TUNE}2"] % 1
 
         beta_kmod = beta_kmod.sort_values(by=S)
 
@@ -243,7 +243,7 @@ def convert_betastar_results(
     LOG.debug("Converting K-modulation BetaStar results")
 
     # merge files and set index
-    kmod_results = tfs.concat(betastar_results_list, join='inner',)
+    kmod_results = tfs.concat(betastar_results_list, join="inner",)
     if BEAM in kmod_results.columns or kmod_results.index.name == BEAM:  # averaged file
         if beam is None:
             raise ValueError("Need to give beam when importing averaged betastar files.")
@@ -270,31 +270,31 @@ def convert_betastar_results(
 
         # copy s and beta from model
         beta_kmod.loc[:, S] = df_model.loc[:, S]
-        beta_kmod.loc[:, f'{BETASTAR}{plane}{MDL}'] = df_model.loc[:, f'{BETA}{plane}']
+        beta_kmod.loc[:, f"{BETASTAR}{plane}{MDL}"] = df_model.loc[:, f"{BETA}{plane}"]
 
         # copy columns with errors from results
         columns = [
             f"{column}{plane}" for column in (BETASTAR, BETAWAIST, WAIST, PHASEADV)] + [f"{WAIST}{plane}{S_LOCATION}"
         ]
         for column in columns:
-            beta_kmod.loc[:, f'{column}'] = kmod_results[f'{column}']
-            beta_kmod.loc[:, f'{ERR}{column}'] = kmod_results[f'{ERR}{column}']
+            beta_kmod.loc[:, f"{column}"] = kmod_results[f"{column}"]
+            beta_kmod.loc[:, f"{ERR}{column}"] = kmod_results[f"{ERR}{column}"]
 
 
         # model-delta and beta-beating
-        beta_kmod.loc[:, f'{DELTA}{BETASTAR}{plane}{MDL}'] = (
-            beta_kmod[f'{BETASTAR}{plane}'] - beta_kmod[f'{BETASTAR}{plane}{MDL}']
+        beta_kmod.loc[:, f"{DELTA}{BETASTAR}{plane}{MDL}"] = (
+            beta_kmod[f"{BETASTAR}{plane}"] - beta_kmod[f"{BETASTAR}{plane}{MDL}"]
         )
-        beta_kmod.loc[:, f'{DELTA}{BETASTAR}{plane}'] = (
-            beta_kmod[f'{DELTA}{BETASTAR}{plane}{MDL}'] / beta_kmod[f'{BETASTAR}{plane}{MDL}']
+        beta_kmod.loc[:, f"{DELTA}{BETASTAR}{plane}"] = (
+            beta_kmod[f"{DELTA}{BETASTAR}{plane}{MDL}"] / beta_kmod[f"{BETASTAR}{plane}{MDL}"]
         )
-        beta_kmod.loc[:, f'{ERR}{DELTA}{BETASTAR}{plane}'] = (
-            beta_kmod[f'{ERR}{BETASTAR}{plane}'] / beta_kmod[f'{BETASTAR}{plane}{MDL}']
+        beta_kmod.loc[:, f"{ERR}{DELTA}{BETASTAR}{plane}"] = (
+            beta_kmod[f"{ERR}{BETASTAR}{plane}"] / beta_kmod[f"{BETASTAR}{plane}{MDL}"]
         )
 
         # tune
-        beta_kmod.headers[f'{TUNE}1'] = df_model.headers[f'{TUNE}1'] % 1
-        beta_kmod.headers[f'{TUNE}2'] = df_model.headers[f'{TUNE}2'] % 1
+        beta_kmod.headers[f"{TUNE}1"] = df_model.headers[f"{TUNE}1"] % 1
+        beta_kmod.headers[f"{TUNE}2"] = df_model.headers[f"{TUNE}2"] % 1
 
         beta_kmod = beta_kmod.sort_values(by=S)
 
@@ -304,10 +304,10 @@ def convert_betastar_results(
 
 
 def _sync_model_index(kmod_results: tfs.TfsDataFrame, df_model: tfs.TfsDataFrame):
-    missing_elmnts = kmod_results.index.difference(df_model.index)
-    if len(missing_elmnts):
+    missing_elements = kmod_results.index.difference(df_model.index)
+    if len(missing_elements):
         msg =  (
-            f"Elements {missing_elmnts} not found in the model. "
+            f"Elements {missing_elements} not found in the model. "
             "Make sure to use the `elements` file as input!"
         )
         raise NameError(msg)
@@ -387,7 +387,7 @@ def _write_output(dfs: dict[str, tfs.TfsDataFrame], output_dir: Path):
             if df is None:
                 continue
 
-            tfs.write(output_dir / f'{filename}{plane.lower()}{EXT}', df, save_index=NAME)
+            tfs.write(output_dir / f"{filename}{plane.lower()}{EXT}", df, save_index=NAME)
 
 
 # Script Mode ------------------------------------------------------------------

diff --git a/omc3/scripts/kmod_lumi_imbalance.py b/omc3/scripts/kmod_lumi_imbalance.py
@@ -63,7 +63,7 @@
     EFFECTIVE_BETAS_FILENAME,
     ERR,
     EXT,
-    IMBALACE,
+    IMBALANCE,
     LUMINOSITY,
     MDL,
     NAME,
@@ -200,8 +200,8 @@ def get_lumi_imbalance_df(**kwargs) -> tfs.TfsDataFrame:
         *tuple(df_effective_betas.loc[ip_b, :])
     )
 
-    df_effective_betas.headers[f'{LUMINOSITY}{IMBALACE}'] = lumi_imb
-    df_effective_betas.headers[f'{ERR}{LUMINOSITY}{IMBALACE}'] = lumi_imb_err
+    df_effective_betas.headers[f'{LUMINOSITY}{IMBALANCE}'] = lumi_imb
+    df_effective_betas.headers[f'{ERR}{LUMINOSITY}{IMBALANCE}'] = lumi_imb_err
     return df_effective_betas
 
 

diff --git a/tests/inputs/kmod/references/effective_betas_beta0.22m0.22m.tfs b/tests/inputs/kmod/references/effective_betas_beta0.22m0.22m.tfs
@@ -1,5 +1,5 @@
-@ LUMIIMB              %le   0.9998570043832672
-@ ERRLUMIIMB           %le    0.026940272475189
+@ LUMIIMBALANCE        %le   0.9998570043832672
+@ ERRLUMIIMBALANCE     %le    0.026940272475189
 * NAME                              BETSTAR           ERRBETSTAR
 $ %s                                    %le                  %le
   "IP1"                      0.226162919812     0.00415582295069

diff --git a/tests/unit/test_kmod_averaging.py b/tests/unit/test_kmod_averaging.py
@@ -1,22 +1,28 @@
+from collections.abc import Sequence
+import logging
 from pathlib import Path
+import shutil
 
 import pandas.testing as pdt
 import pytest
 import tfs
 
-from omc3.optics_measurements.constants import BETA, NAME
+from omc3.optics_measurements.constants import BEAM, BEAM_DIR, BETA, NAME
 from omc3.scripts.kmod_average import (
     AVERAGED_BETASTAR_FILENAME,
     AVERAGED_BPM_FILENAME,
     EXT,
     average_kmod_results,
 )
-from omc3.plotting.plot_kmod_results import PARAM_BETA, PARAM_WAIST
+from omc3.plotting.plot_kmod_results import PARAM_BETA, PARAM_BETABEAT, PARAM_WAIST
 from tests.conftest import INPUTS, ids_str
 
 KMOD_INPUT_DIR = INPUTS / "kmod"
 REFERENCE_DIR = KMOD_INPUT_DIR / "references"
 
+
+# Tests ----
+
 @pytest.mark.basic
 @pytest.mark.parametrize("ip", [1, 5], ids=ids_str("ip{}"))
 @pytest.mark.parametrize("n_files", [1, 2], ids=ids_str("{}files"))
@@ -40,9 +46,51 @@ def test_kmod_averaging(tmp_path, ip, n_files):
         pdt.assert_frame_equal(out_file, ref_file, check_like=True)
 
 
-def _assert_correct_files_are_present(outputdir: Path, ip: int, beta: float) -> None:
+@pytest.mark.extended
+@pytest.mark.parametrize("beam", [1, 2], ids=ids_str("beam{}"))
+def test_kmod_averaging_single_beam(tmp_path, beam, caplog):
+    ip = 1
+    n_files = 2
+
+    beta = get_betastar_model(beam=beam, ip=ip)
+    ref_output_dir = get_reference_dir(ip, n_files)
+
+    meas_paths = [get_measurement_dir(ip, i+1)  for i in range(n_files)]
+    new_meas_paths = [tmp_path / f"single_beam_meas_{i+1}" for i in range(n_files)]
+
+    for tmp_meas, old_meas in zip(new_meas_paths, meas_paths):
+        beam_dir = f"{BEAM_DIR}{beam}"
+        (tmp_meas / beam_dir).mkdir(parents=True, exist_ok=True)
+        for tfs_file in (old_meas / beam_dir).glob("*"):
+            shutil.copy(tfs_file, tmp_meas / beam_dir)
+
+    with caplog.at_level(logging.WARNING):
+        average_kmod_results(
+            meas_paths=new_meas_paths, 
+            output_dir=tmp_path,
+            ip=ip, 
+            betastar=beta,
+            plot=True
+        )
+    assert f"Could not find all results for beam {1 if beam == 2 else 2}" in caplog.text
+    _assert_correct_files_are_present(tmp_path, ip, beta[0], beams=[beam])
+
+    for out_name in get_all_tfs_filenames(ip, beta[0], beams=[beam]):
+        out_file = tfs.read(tmp_path / out_name)
+        ref_file = tfs.read(ref_output_dir / out_name)
+        if BEAM in ref_file.columns:
+            ref_file = ref_file.loc[ref_file[BEAM] == beam, :].reset_index(drop=True)
+        pdt.assert_frame_equal(out_file, ref_file, check_like=True)
+
+
+# Helper ---
+
+def _assert_correct_files_are_present(outputdir: Path, ip: int, beta: float, beams: Sequence[int] = (1, 2)) -> None:
     """Simply checks the expected converted files are present in the outputdir"""
-    all_files = get_all_tfs_filenames(ip, beta) + [f"ip{ip}_{PARAM_WAIST}.pdf", f"ip{ip}_{PARAM_BETA}.pdf"]
+    all_files = (
+        get_all_tfs_filenames(ip, beta, beams) + 
+        [f"ip{ip}_{PARAM_WAIST}.pdf", f"ip{ip}_{PARAM_BETA}.pdf", f"ip{ip}_{PARAM_BETABEAT}.pdf"]
+    )
     for file_name in all_files:
         assert (outputdir / file_name).is_file()
 
@@ -64,10 +112,11 @@ def get_betastar_model(beam: int, ip: int) -> Path:
     return model.loc[f"IP{ip}", [f"{BETA}Y", f"{BETA}Y"]].tolist()
 
 
-def get_all_tfs_filenames(ip: int, beta:float) -> list[str]:
+def get_all_tfs_filenames(ip: int, beta: float, beams: Sequence[int] = (1, 2)) -> list[str]:
     return [
-        f"{AVERAGED_BPM_FILENAME.format(betastar_x=beta, betastar_y=beta, ip=ip, beam=1)}{EXT}",
-        f"{AVERAGED_BPM_FILENAME.format(betastar_x=beta, betastar_y=beta, ip=ip, beam=2)}{EXT}",
+        f"{AVERAGED_BPM_FILENAME.format(betastar_x=beta, betastar_y=beta, ip=ip, beam=beam)}{EXT}"
+        for beam in beams
+    ] + [
         f"{AVERAGED_BETASTAR_FILENAME.format(betastar_x=beta, betastar_y=beta, ip=ip)}{EXT}",
     ]