fixed fake meas alfx bug

pylhc · Mar 3, 2025 · 5d39359 · 5d39359
1 parent 0f894b5
commit 5d39359
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/omc3/scripts/fake_measurement_from_model.py b/omc3/scripts/fake_measurement_from_model.py
@@ -157,7 +157,7 @@ def get_params():
     )
     params.add_parameter(
         name="parameters",
-        help="Optics parameters to use",
+        help="Optics parameters to use.",
         choices=list(OUTPUTNAMES_MAP.keys()),
         default=list(OUTPUTNAMES_MAP.keys()),
         type=str,
@@ -256,14 +256,16 @@ def create_beta(df_twiss: pd.DataFrame, df_model: pd.DataFrame, parameter: str,
     LOG.info(f"Creating fake beta for {parameter}.")
     plane = parameter[-1]
 
+    # create beta
     df = create_measurement(df_twiss, parameter, relative_error, randomize)
     df[parameter] = np.abs(df[parameter])
     df = append_model_param(df, df_model, parameter, beat=True)
 
+    # create alpha
     df_alpha = create_measurement(df_twiss, f'{ALPHA}{plane}', relative_error, randomize)
-    df_alpha = append_model_param(df_alpha, df_model, f'{ALPHA}{plane}', plane)
-
+    df_alpha = append_model_param(df_alpha, df_model, f'{ALPHA}{plane}')
     df = tfs.concat([df, df_alpha], axis=1, join='inner')
+
     df = append_model_s_and_phaseadv(df, df_model, planes=plane)
 
     df.headers = headers.copy()
@@ -483,7 +485,7 @@ def append_model_s_and_phaseadv(df: pd.DataFrame, df_model: pd.DataFrame, planes
 
 # Other Functions --------------------------------------------------------------
 
-def _get_data(twiss: tfs.TfsDataFrame, model: tfs.TfsDataFrame = None,
+def _get_data(twiss: tfs.TfsDataFrame, model: tfs.TfsDataFrame | None = None,
               add_coupling: bool = False) -> tuple[tfs.TfsDataFrame, tfs.TfsDataFrame]:
     """ Gets the input data as TfsDataFrames. """
     # Helper ---
@@ -513,7 +515,7 @@ def try_reading(df_or_path):
     return twiss, model
 
 
-def _get_loop_parameters(parameters: Sequence[str], errors: Sequence[float]) -> list[str]:
+def _get_loop_parameters(parameters: Sequence[str], errors: Sequence[float] | None) -> list[str]:
     """ Special care for normalized dispersion"""
     parameters = list(parameters)
     if errors is None:

diff --git a/tests/unit/test_fake_measurement.py b/tests/unit/test_fake_measurement.py
@@ -167,14 +167,14 @@ def _test_error_columns(name, df, randomized, error_val):
             if name.startswith(TOTAL_PHASE_NAME):
                 idx = df.index[1:]
 
-            if name[:-1] in (PHASE_NAME, TOTAL_PHASE_NAME):
+            if param[:-1] in (PHASE, PHASE_ADV):
                 # phase errors are equal to the relative error
                 assert not any(df.loc[idx, col] - error_val)
-            elif name[:-1] in (BETA_NAME, AMP_BETA_NAME) and DELTA in col:
+            elif param[:-1] in (BETA, ) and DELTA in col:
                 # errdeltabet (beating) errors are also equal to the relative error,
                 # but with less precision
                 assert all(np.abs(df.loc[idx, col] - error_val) < EPS)
-            elif name[:-1] in (NORM_DISP_NAME):
+            elif param[:-1] in (NORM_DISPERSION, DISPERSION):
                 # not sure how to test this, but should already be tested with disp and beta
                 assert all(df.loc[idx, col])
             else: