SbS Dispersion Fix (local)

.zenodo.json

@@ -4,6 +4,11 @@
             "name": "OMC-Team",
             "affiliation": "CERN"
         },
+        {
+            "name": "Yannis Angelis",
+            "affiliation": "CERN",
+            "orcid": "0009-0003-9057-5165"
+        },
         {
             "name": "Felix Simon Carlier",
             "affiliation": "CERN",
@@ -65,6 +70,10 @@
         {
             "name": "Andreas Wegscheider",
             "affiliation": "CERN"
+        },
+        {
+            "name": "Wietse Van Goethem",
+            "affiliation": "CERN"
         }
     ],
     "title": "OMC3",

CHANGELOG.md

@@ -1,5 +1,10 @@
 # OMC3 Changelog
 
+#### 2025-08-26 - v0.24.5 - _yangelis_
+
+- Fixed/Added:
+  - Segement-by-Segment dispersion propagation by adding momentum dispersion propagable.
+
 #### 2025-04-28 - v0.24.4 - _wvangoet_
 
 - Fixed:
@@ -346,7 +351,7 @@
     (over moving average window)
   - Action error calculated from error on the spectral line
     (which in turn is the same as NOISE)
-  
+
 #### 2022-11-01 - v0.6.6
 
 - Bugfixes:
@@ -477,7 +482,7 @@
   - more usages of constants for strings (e.g. column names)
   - introducing pathlib.Path in some places
   - output-paths in model job-files are relative
-  
+
 - Fixed:
   - Matplotlib warnings for `set_window_title`
   - excluded Windows and MacOS py3.9 from normal testing, due to installation issues of pyTables

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

omc3/optics_measurements/constants.py

@@ -20,6 +20,7 @@
 DRIVEN_PHASE_NAME: str = f"{PHASE_NAME}driven_"
 DRIVEN_TOTAL_PHASE_NAME: str = f"{TOTAL_PHASE_NAME}driven_"
 DISPERSION_NAME: str = "dispersion_"
+MOMENTUM_DISPERSION_NAME: str = 'momentum_dispersion_' # Segment-by-Segment only
 NORM_DISP_NAME: str = "normalised_dispersion_"
 ORBIT_NAME: str = "orbit_"
 KICK_NAME: str = "kick_"
@@ -66,6 +67,7 @@
 ORBIT: str = ""  # Column is plane (X or Y) in files
 CLOSED_ORBIT: str = "CO"
 DISPERSION: str = "D"
+MOMENTUM_DISPERSION: str = "DP"
 NORM_DISPERSION: str = "ND"
 
 MEASUREMENT: str = "MEAS"

omc3/scripts/fake_measurement_from_model.py

@@ -93,6 +93,7 @@
     IMAG,
     MDL,
     MODEL_DIRECTORY,
+    MOMENTUM_DISPERSION,
     NAME,
     NAME2,
     NORM_DISP_NAME,
@@ -281,6 +282,11 @@ def create_dispersion(df_twiss: pd.DataFrame, df_model: pd.DataFrame, parameter:
     plane = parameter[-1]
     df = create_measurement(df_twiss, parameter, relative_error, randomize)
     df = append_model_param(df, df_model, parameter)
+
+    df_dp = create_measurement(df_twiss, f'{MOMENTUM_DISPERSION}{plane}', relative_error, randomize)
+    df_dp = append_model_param(df_dp, df_model, f'{MOMENTUM_DISPERSION}{plane}')
+    df = tfs.concat([df, df_dp], axis=1, join='inner')
+
     df = append_model_s_and_phaseadv(df, df_model, planes=plane)
     df.headers = headers.copy()
     return {f'{DISPERSION_NAME}{plane.lower()}': df}

omc3/segment_by_segment/propagables/__init__.py

@@ -16,8 +16,8 @@
 from omc3.segment_by_segment.propagables.alpha import AlphaPhase
 from omc3.segment_by_segment.propagables.beta import BetaPhase
 from omc3.segment_by_segment.propagables.coupling import F1001, F1010
-from omc3.segment_by_segment.propagables.dispersion import Dispersion
+from omc3.segment_by_segment.propagables.dispersion import Dispersion, MomentumDispersion
 from omc3.segment_by_segment.propagables.phase import Phase
 from omc3.segment_by_segment.propagables.utils import PropagableColumns  # noqa: F401 -> expose
 
-ALL_PROPAGABLES: tuple[type[Propagable], ...] = (Phase, BetaPhase, AlphaPhase, Dispersion, F1001, F1010)
+ALL_PROPAGABLES: tuple[type[Propagable], ...] = (Phase, BetaPhase, AlphaPhase, Dispersion, MomentumDispersion, F1001, F1010)

omc3/segment_by_segment/propagables/dispersion.py

@@ -9,7 +9,10 @@
 
 from typing import TYPE_CHECKING
 
-from omc3.optics_measurements.constants import DISPERSION
+import numpy as np
+from pandas import Series
+
+from omc3.optics_measurements.constants import DISPERSION, MOMENTUM_DISPERSION
 from omc3.segment_by_segment import math as sbs_math
 from omc3.segment_by_segment.propagables.abstract import Propagable
 from omc3.segment_by_segment.propagables.phase import Phase
@@ -119,3 +122,97 @@ def _compute_elements(self, plane: str, seg_model: pd.DataFrame, forward: bool):
         model_phase = Phase.get_segment_phase(seg_model, plane, forward)
         propagated_err = sbs_math.propagate_error_dispersion(model_disp, model_phase, init_condition)
         return model_disp, propagated_err
+
+
+class MomentumDispersion(Propagable):
+
+    _init_pattern = "dp{}_{}"  # format(plane, ini/end)
+    columns: PropagableColumns = PropagableColumns(MOMENTUM_DISPERSION)
+
+    @classmethod
+    def get_at(cls, names: IndexType, meas: OpticsMeasurement, plane: str) -> ValueErrorType:
+        c = cls.columns.planed(plane)
+        momentum_dispersion = meas.dispersion[plane].loc[names, c.column]
+        model_momentum_dispersion_err = np.nan # No error in the measured dpx,dpy
+        return momentum_dispersion, model_momentum_dispersion_err
+
+    @classmethod
+    def in_measurement(cls, meas: OpticsMeasurement) -> bool:
+        """ Check if the dispersion is in the measurement data. """
+        try:
+            meas.dispersion_x
+            meas.dispersion_y
+        except FileNotFoundError:
+            return False
+        return True
+
+    def init_conditions_dict(self):
+        # needs to be inverted for backward propagation, i.e. the end-init
+        init_cond = super().init_conditions_dict()
+        for key, value in init_cond.items():
+            if "end" in key:
+                init_cond[key] = -value
+        return init_cond
+
+    def add_differences(self, segment_diffs: SegmentDiffs):
+        dfs = self.get_difference_dataframes()
+        for plane, df in dfs.items():
+            # save to diffs/write to file (if allow_write is set)
+            segment_diffs.momentum_dispersion[plane] = df
+
+    def _compute_measured(self,
+            plane: str,
+            seg_model: TfsDataFrame,
+            forward: bool
+        ) -> tuple[pd.Series, pd.Series]:
+        """ Compute the momentum dispersion difference between the given segment model and the measured values."""
+
+        # get the measured values
+        names = self.get_segment_observation_points(plane)
+        momentum_dispersion, err_pdisp = self.get_at(names, self._meas, plane)
+
+        # get the propagated values
+        model_momentum_dispersion = seg_model.loc[names, f"{MOMENTUM_DISPERSION}{plane}"]
+
+        if not forward:
+            model_momentum_dispersion = -model_momentum_dispersion
+
+        # calculate difference
+        momentum_dispersion_diff = momentum_dispersion - model_momentum_dispersion
+
+        propagated_err = Series(np.nan, index=model_momentum_dispersion.index) # No error propagation for now
+        return momentum_dispersion_diff, propagated_err
+
+    def _compute_correction(
+            self,
+            plane: str,
+            seg_model: pd.DataFrame,
+            seg_model_corr: pd.DataFrame,
+            forward: bool,
+        ) -> tuple[pd.Series, pd.Series]:
+        """Compute the momentum dispersion differennce between the nominal and the corrected model."""
+
+        model_momentum_dispersion = seg_model.loc[:, f"{MOMENTUM_DISPERSION}{plane}"]
+        corrected_momentum_dispersion = seg_model_corr.loc[:, f"{MOMENTUM_DISPERSION}{plane}"]
+
+        momentum_dispersion_diff = corrected_momentum_dispersion - model_momentum_dispersion
+        if not forward:
+            momentum_dispersion_diff = -momentum_dispersion_diff
+
+        propagated_err = Series(np.nan, index=model_momentum_dispersion.index) # No error propagation for now
+        return momentum_dispersion_diff, propagated_err
+
+    def _compute_elements(self, plane: str, seg_model: pd.DataFrame, forward: bool):
+        """ Compute get the propagated momentum dispersion values from the segment model and calculate the propagated error.  """
+
+        model_momentum_dispersion = seg_model.loc[:, f"{MOMENTUM_DISPERSION}{plane}"]
+
+        pderror = Series(np.nan, index=model_momentum_dispersion.index) # No error propagation for now
+        return model_momentum_dispersion, pderror
+
+    def get_segment_observation_points(self, plane: str):
+        """ Return the measurement points for the given plane, that are in the segment. """
+        return common_indices(
+            self.segment_models.forward.index,
+            self._meas.dispersion[plane].index
+        )

omc3/segment_by_segment/segments.py

@@ -20,6 +20,7 @@
     F1001_NAME,
     F1010_NAME,
     KMOD_BETA_NAME,
+    MOMENTUM_DISPERSION_NAME,
     NORM_DISP_NAME,
     PHASE_NAME,
 )
@@ -117,6 +118,7 @@ class SegmentDiffs(TfsCollection):
     beta_kmod = Tfs(f"{PREFIX}{KMOD_BETA_NAME}{{plane}}_{{name}}{EXT}")
     beta_amp = Tfs(f"{PREFIX}{AMP_BETA_NAME}{{plane}}_{{name}}{EXT}")
     dispersion = Tfs(f"{PREFIX}{DISPERSION_NAME}{{plane}}_{{name}}{EXT}")
+    momentum_dispersion = Tfs(f"{PREFIX}{MOMENTUM_DISPERSION_NAME}{{plane}}_{{name}}{EXT}")
     norm_dispersion = Tfs(f"{PREFIX}{NORM_DISP_NAME}{{plane}}_{{name}}{EXT}")
     f1001 = Tfs(f"{PREFIX}{F1001_NAME}_{{name}}{EXT}", two_planes=False)
     f1010 = Tfs(f"{PREFIX}{F1010_NAME}_{{name}}{EXT}", two_planes=False)

tests/accuracy/test_sbs.py

@@ -25,6 +25,7 @@
     AlphaPhase,
     BetaPhase,
     Dispersion,
+    MomentumDispersion,
     Phase,
     PropagableColumns,
 )
@@ -59,7 +60,8 @@ class TestSbSLHC:
         Phase: TestCfg(1e-2, 1e-8, 5e-4, "phase"),
         BetaPhase: TestCfg(9e-2, 1e-10, 5e-4, "beta_phase"),
         AlphaPhase: TestCfg(1e-1, 1e-8, 8e-2, "alpha_phase"),
-        Dispersion: TestCfg(1e-2, None, None, "dispersion"),  # not really working at the moment, see https://github.com/pylhc/omc3/issues/498
+        Dispersion: TestCfg(1e-6, 1e-4, 1e-4, "dispersion"),
+        MomentumDispersion: TestCfg(1e-7, 1e-4, 1e-4, "momentum_dispersion"),
         F1001: TestCfg(5e-4, 5e-7, None, "f1001"),
         F1010: TestCfg(5e-4, 5e-6, None, "f1010"),
     }
@@ -87,7 +89,7 @@ def test_lhc_segment_creation(self,
             "beam": beam,
             "nat_tunes": [0.31, 0.32],
             "dpp": 0.0,
-            "energy": 6500,
+            "energy": 6800,
             "modifiers": [acc_models_lhc_2025 / OPTICS_SUBDIR / OPTICS_30CM_FLAT],
         }
 
@@ -203,8 +205,6 @@ def test_lhc_propagation_sbs(self,
             ]
 
         for propagable in ALL_PROPAGABLES:
-            if propagable is Dispersion:
-                continue  # TODO: not working, see https://github.com/pylhc/omc3/issues/498
 
             cfg: TestCfg = self.config_map[propagable]
             planes = [REAL, IMAG, AMPLITUDE, PHASE] if propagable.is_rdt() else "xy"
@@ -230,6 +230,13 @@ def test_lhc_propagation_sbs(self,
                         INPUT_SBS / f"measurement_b{beam}" / f"{beta_file}_{plane}.tfs",
                         index=NAME
                     )
+                elif propagable is MomentumDispersion:
+                    # momentum dispersion is in the dispersion-measurement file
+                    dispersion_file = self.config_map[Dispersion].file_name
+                    meas_df = tfs.read(
+                        INPUT_SBS / f"measurement_b{beam}" / f"{dispersion_file}_{plane}.tfs",
+                        index=NAME
+                    )
                 else:
                     meas_df = tfs.read(
                         INPUT_SBS / f"measurement_b{beam}" / f"{full_file_name}.tfs",
@@ -264,7 +271,6 @@ def test_lhc_propagation_sbs(self,
                             assert sbs_df[col].abs().min() < self.eps_rdt_min  # for RDTs the init value is calculated,
                         else:
                             assert sbs_df[col].abs().min() == 0  # at least the first entry should show no difference
-
                         assert sbs_df[col].abs().max() > cfg.diff_max
                         assert sbs_df[err_col].all()
 
@@ -277,7 +283,7 @@ def test_lhc_propagation_sbs(self,
                                 # TODO: check backward coupling, see https://github.com/pylhc/omc3/issues/498
                                 continue
 
-                            if propagable in [Phase, F1001, F1010]:  # check absolute difference
+                            if propagable in [Phase, F1001, F1010, Dispersion, MomentumDispersion]:  # check absolute difference
                                 assert_all_close(sbs_df, correction, col, atol=eps)
                                 assert_all_close(sbs_df, expected, 0, atol=eps)
                             else:  # check relative difference