Repository: pylhc/optics_functions

Issue:
Implement calculation from Coupling to R matrix
This is the inverse of the coupling from r matrix calculation

Comments:

Patch:
diff --git a/optics_functions/coupling.py b/optics_functions/coupling.py
index efa91e1..d4e6397 100644
--- a/optics_functions/coupling.py
+++ b/optics_functions/coupling.py
@@ -12,11 +12,11 @@ from contextlib import suppress
 from typing import Sequence
 
 import numpy as np
+from pandas import DataFrame, Series
 from tfs import TfsDataFrame
 
-from optics_functions.constants import (ALPHA, BETA, GAMMA,
-                                        X, Y, TUNE, DELTA,
-                                        MINIMUM, PI2, PHASE_ADV, S, LENGTH)
+from optics_functions.constants import (ALPHA, BETA, GAMMA, X, Y, TUNE, DELTA,
+                                        MINIMUM, PI2, PHASE_ADV, S, LENGTH, IMAG, REAL)
 from optics_functions.rdt import calculate_rdts
 from optics_functions.utils import split_complex_columns, timeit
 
@@ -26,8 +26,11 @@ COUPLING_RDTS = ["F1001", "F1010"]
 LOG = logging.getLogger(__name__)
 
 
-def coupling_via_rdts(df: TfsDataFrame, complex_columns: bool = True, **kwargs):
-    """ Returns the coupling term.
+# Coupling ---------------------------------------------------------------------
+
+
+def coupling_via_rdts(df: TfsDataFrame, complex_columns: bool = True, **kwargs) -> TfsDataFrame:
+    """Returns the coupling terms.
 
     .. warning::
         This function changes sign of the real part of the RDTs compared to
@@ -54,17 +57,16 @@ def coupling_via_rdts(df: TfsDataFrame, complex_columns: bool = True, **kwargs):
 
     if not complex_columns:
         df_res = split_complex_columns(df_res, COUPLING_RDTS)
-
     return df_res
 
 
-def coupling_via_cmatrix(df: TfsDataFrame, complex_columns: bool = True,
-                         output: Sequence[str] = ("rdts", "gamma", "cmatrix")):
-    """ Calculates C matrix then Coupling and Gamma from it.
+def coupling_via_cmatrix(df: DataFrame, complex_columns: bool = True,
+                         output: Sequence[str] = ("rdts", "gamma", "cmatrix")) -> DataFrame:
+    """Calculates C matrix then Coupling and Gamma from it.
     See [CalagaBetatronCoupling2005]_ .
 
     Args:
-        df (TfsDataFrame): Twiss Dataframe
+        df (DataFrame): Twiss Dataframe
         complex_columns (bool): Output complex values in single column of type complex.
                                 If ``False``, split complex columns into two
                                 real-valued columns.
@@ -75,13 +77,14 @@ def coupling_via_cmatrix(df: TfsDataFrame, complex_columns: bool = True,
         New TfsDataFrame with columns as specified in 'output'.
     """
     LOG.info("Calculating coupling from c-matrix.")
-    df_res = TfsDataFrame(index=df.index)
+    df_res = DataFrame(index=df.index)
 
     with timeit("CMatrix calculation", print_fun=LOG.debug):
         n = len(df)
         gx, r, inv_gy = np.zeros((n, 2, 2)), np.zeros((n, 2, 2)), np.zeros((n, 2, 2))
 
-        # rs form after -J R^T J == inv(R)*det|R| == C
+        # Eq. (16)  C = 1 / (1 + |R|) * -J R J
+        # rs form after -J R^T J
         r[:, 0, 0] = df["R22"]
         r[:, 0, 1] = -df["R12"]
         r[:, 1, 0] = -df["R21"]
@@ -89,27 +92,26 @@ def coupling_via_cmatrix(df: TfsDataFrame, complex_columns: bool = True,
 
         r *= 1 / np.sqrt(1 + np.linalg.det(r)[:, None, None])
 
-        # Cbar = Gx * C * Gy^-1  (Eq. (5) in reference)
-        sqrtbetax = np.sqrt(df[f"{BETA}{X}"])
-        sqrtbetay = np.sqrt(df[f"{BETA}{Y}"])
+        # Cbar = Gx * C * Gy^-1,   Eq. (5)
+        sqrt_betax = np.sqrt(df[f"{BETA}{X}"])
+        sqrt_betay = np.sqrt(df[f"{BETA}{Y}"])
 
-        gx[:, 0, 0] = 1 / sqrtbetax
+        gx[:, 0, 0] = 1 / sqrt_betax
         gx[:, 1, 0] = df[f"{ALPHA}{X}"] * gx[:, 0, 0]
-        gx[:, 1, 1] = sqrtbetax
+        gx[:, 1, 1] = sqrt_betax
 
-        inv_gy[:, 1, 1] = 1 / sqrtbetay
+        inv_gy[:, 1, 1] = 1 / sqrt_betay
         inv_gy[:, 1, 0] = -df[f"{ALPHA}{Y}"] * inv_gy[:, 1, 1]
-        inv_gy[:, 0, 0] = sqrtbetay
+        inv_gy[:, 0, 0] = sqrt_betay
 
         c = np.matmul(gx, np.matmul(r, inv_gy))
         gamma = np.sqrt(1 - np.linalg.det(c))
 
     if "rdts" in output:
+        # Eq. (9) and Eq. (10)
         denom = 1 / (4 * gamma)
-        df_res.loc[:, "F1001"] = ((c[:, 0, 0] + c[:, 1, 1]) * 1j +
-                                  (c[:, 0, 1] - c[:, 1, 0])) * denom
-        df_res.loc[:, "F1010"] = ((c[:, 0, 0] - c[:, 1, 1]) * 1j +
-                                  (-c[:, 0, 1]) - c[:, 1, 0]) * denom
+        df_res.loc[:, "F1001"] = denom * (+c[:, 0, 1] - c[:, 1, 0] + (c[:, 0, 0] + c[:, 1, 1]) * 1j)
+        df_res.loc[:, "F1010"] = denom * (-c[:, 0, 1] - c[:, 1, 0] + (c[:, 0, 0] - c[:, 1, 1]) * 1j)
         LOG.info(f"Average coupling amplitude |F1001|: {df_res['F1001'].abs().mean():g}")
         LOG.info(f"Average coupling amplitude |F1010|: {df_res['F1010'].abs().mean():g}")
 
@@ -117,10 +119,10 @@ def coupling_via_cmatrix(df: TfsDataFrame, complex_columns: bool = True,
             df_res = split_complex_columns(df_res, COUPLING_RDTS)
 
     if "cmatrix" in output:
-        df_res.loc[:, "C11"] = c[:, 0, 0]
-        df_res.loc[:, "C12"] = c[:, 0, 1]
-        df_res.loc[:, "C21"] = c[:, 1, 0]
-        df_res.loc[:, "C22"] = c[:, 1, 1]
+        df_res["C11"] = c[:, 0, 0]
+        df_res["C12"] = c[:, 0, 1]
+        df_res["C21"] = c[:, 1, 0]
+        df_res["C22"] = c[:, 1, 1]
 
     if "gamma" in output:
         df_res.loc[:, GAMMA] = gamma
@@ -129,8 +131,88 @@ def coupling_via_cmatrix(df: TfsDataFrame, complex_columns: bool = True,
     return df_res
 
 
-def closest_tune_approach(df: TfsDataFrame, qx: float = None, qy: float = None, method: str = "calaga"):
-    """ Calculates the closest tune approach from coupling resonances.
+# R-Matrix ---------------------------------------------------------------------
+
+
+def rmatrix_from_coupling(df: DataFrame, complex_columns: bool = True) -> DataFrame:
+    """Calculates the R-matrix from a DataFrame containing the coupling columns
+     as well as alpha and beta columns. This is the inverse of
+    :func:`optics_functions.coupling.coupling_via_cmatrix`.
+    See [CalagaBetatronCoupling2005]_ .
+
+    Args:
+        df (DataFrame): Twiss Dataframe
+        complex_columns (bool): Tells the function if the coupling input columns
+                                are complex-valued or split into real and
+                                imaginary parts.
+
+    Returns:
+        New DataFrame containing the R-columns.
+    """
+    LOG.info("Calculating r-matrix from coupling rdts.")
+    df_res = DataFrame(index=df.index)
+
+    with timeit("R-Matrix calculation", print_fun=LOG.debug):
+        if complex_columns:
+            df = split_complex_columns(df, COUPLING_RDTS, drop=False)
+
+        n = len(df)
+
+        # From Eq. (5) in reference:
+        inv_gx, jcj, gy = np.zeros((n, 2, 2)), np.zeros((n, 2, 2)), np.zeros((n, 2, 2))
+
+        sqrt_betax = np.sqrt(df[f"{BETA}{X}"])
+        sqrt_betay = np.sqrt(df[f"{BETA}{Y}"])
+
+        inv_gx[:, 1, 1] = 1 / sqrt_betax
+        inv_gx[:, 1, 0] = -df[f"{ALPHA}{X}"] * inv_gx[:, 1, 1]
+        inv_gx[:, 0, 0] = sqrt_betax
+
+        gy[:, 0, 0] = 1 / sqrt_betay
+        gy[:, 1, 0] = df[f"{ALPHA}{Y}"] * gy[:, 0, 0]
+        gy[:, 1, 1] = sqrt_betay
+
+        # Eq. (15)
+        if complex_columns:
+            abs_squared_diff = df["F1001"].abs()**2 - df["F1010"].abs()**2
+        else:
+            abs_squared_diff = (df[f"F1001{REAL}"]**2 + df[f"F1001{IMAG}"]**2 -
+                                df[f"F1010{REAL}"]**2 - df[f"F1010{IMAG}"]**2)
+
+        gamma = np.sqrt(1.0 / (1.0 + 4.0 * abs_squared_diff))
+
+        # Eq. (11) and Eq. (12)
+        cbar = np.zeros((n, 2, 2))
+        cbar[:, 0, 0] = (df[f"F1001{IMAG}"] + df[f"F1010{IMAG}"]).to_numpy()
+        cbar[:, 0, 1] = -(df[f"F1010{REAL}"] - df[f"F1001{REAL}"]).to_numpy()
+        cbar[:, 1, 0] = -(df[f"F1010{REAL}"] + df[f"F1001{REAL}"]).to_numpy()
+        cbar[:, 1, 1] = (df[f"F1001{IMAG}"] - df[f"F1010{IMAG}"]).to_numpy()
+        cbar = 2 * gamma.to_numpy()[:, None, None] * cbar
+
+        # Gx^-1 * Cbar * Gy = C  (Eq. (5) inverted)
+        c = np.matmul(inv_gx, np.matmul(cbar, gy))
+
+        # from above: -J R^T J == inv(R)*det|R| == C
+        # therefore -J C^T J = R
+        jcj[:, 0, 0] = c[:, 1, 1]
+        jcj[:, 0, 1] = -c[:, 0, 1]
+        jcj[:, 1, 0] = -c[:, 1, 0]
+        jcj[:, 1, 1] = c[:, 0, 0]
+
+        rmat = jcj * np.sqrt(1 / (1 - np.linalg.det(jcj))[:, None, None])
+        df_res["R11"] = rmat[:, 0, 0]
+        df_res["R12"] = rmat[:, 0, 1]
+        df_res["R21"] = rmat[:, 1, 0]
+        df_res["R22"] = rmat[:, 1, 1]
+
+    return df_res
+
+
+# Closest Tune Approach --------------------------------------------------------
+
+def closest_tune_approach(df: TfsDataFrame, qx: float = None, qy: float = None,
+                          method: str = "calaga") -> TfsDataFrame:
+    """Calculates the closest tune approach from coupling resonances.
 
     A complex F1001 column is assumed to be present in the DataFrame.
     This can be calculated by :func:`~optics_functions.rdt.rdts`
@@ -157,7 +239,7 @@ def closest_tune_approach(df: TfsDataFrame, qx: float = None, qy: float = None,
                       Choices: 'calaga', 'franchi', 'persson' and 'persson_alt'.
 
     Returns:
-        New DataFrame with closest tune approach (DELTAQMIN) column.
+        New TfsDataFrame with closest tune approach (DELTAQMIN) column.
         The value is real for 'calaga' and 'franchi' methods,
     """
     method_map = {
@@ -181,35 +263,36 @@ def closest_tune_approach(df: TfsDataFrame, qx: float = None, qy: float = None,
     return df_res
 
 
-def _cta_franchi(df: TfsDataFrame, qx_frac: float, qy_frac: float):
+def _cta_franchi(df: TfsDataFrame, qx_frac: float, qy_frac: float) -> Series:
     """ Closest tune approach calculated by Eq. (1) in [PerssonImprovedControlCoupling2014]_ . """
-    return 4 * (qx_frac - qy_frac) * df['F1001'].abs()
+    return 4 * (qx_frac - qy_frac) * df["F1001"].abs()
 
 
-def _cta_persson_alt(df: TfsDataFrame, qx_frac: float, qy_frac: float):
-    """ Closest tune approach calculated by Eq. (2) in [PerssonImprovedControlCoupling2014]_ .
-
-    The exp(i(Qx-Qy)s/R) term is omitted. """
+def _cta_persson_alt(df: TfsDataFrame, qx_frac: float, qy_frac: float) -> Series:
+    """Closest tune approach calculated by Eq. (2) in [PerssonImprovedControlCoupling2014]_ .
+    The exp(i(Qx-Qy)s/R) term is omitted.
+    """
     deltaq = qx_frac - qy_frac  # fractional tune split
-    return 4 * deltaq * df['F1001'] * np.exp(-1j * (df[f"{PHASE_ADV}{X}"] - df[f"{PHASE_ADV}{Y}"]))
+    return 4 * deltaq * df["F1001"] * np.exp(-1j * (df[f"{PHASE_ADV}{X}"] - df[f"{PHASE_ADV}{Y}"]))
 
 
-def _cta_persson(df: TfsDataFrame, qx_frac: float, qy_frac: float):
+def _cta_persson(df: TfsDataFrame, qx_frac: float, qy_frac: float) -> Series:
     """ Closest tune approach calculated by Eq. (2) in [PerssonImprovedControlCoupling2014]_ . """
     deltaq = qx_frac - qy_frac  # fractional tune split
-    return 4 * deltaq * df['F1001'] * np.exp(1j *
-           ((deltaq * df[S] / (df.headers[LENGTH] / PI2)) - (df[f"{PHASE_ADV}{X}"] - df[f"{PHASE_ADV}{Y}"])))
-
+    exponential_term = ((deltaq * df[S] / (df.headers[LENGTH] / PI2)) - (df[f"{PHASE_ADV}{X}"] - df[f"{PHASE_ADV}{Y}"]))
+    return 4 * deltaq * df['F1001'] * np.exp(1j * exponential_term)
 
-def _cta_calaga(df: TfsDataFrame, qx_frac: float, qy_frac: float):
-    """ Closest tune approach calculated by Eq. (27) in [CalagaBetatronCoupling2005]_ .
 
+def _cta_calaga(df: TfsDataFrame, qx_frac: float, qy_frac: float) -> Series:
+    """Closest tune approach calculated by Eq. (27) in [CalagaBetatronCoupling2005]_ .
     If F1010 is not given, it is assumed to be zero.
     """
     f_diff = df["F1001"].abs() ** 2
     with suppress(KeyError):
         f_diff -= df["1010"].abs() ** 2
 
-    return ((np.cos(PI2 * qx_frac) - np.cos(PI2 * qy_frac))
-            / (np.pi * (np.sin(PI2 * qx_frac) + np.sin(PI2 * qy_frac)))
-            * (4 * np.sqrt(f_diff) / (1 + 4*f_diff)))
+    return (
+        (np.cos(PI2 * qx_frac) - np.cos(PI2 * qy_frac))
+        / (np.pi * (np.sin(PI2 * qx_frac) + np.sin(PI2 * qy_frac)))
+        * (4 * np.sqrt(f_diff) / (1 + 4 * f_diff))
+    )
diff --git a/optics_functions/utils.py b/optics_functions/utils.py
index 49dfad8..b06e601 100644
--- a/optics_functions/utils.py
+++ b/optics_functions/utils.py
@@ -11,14 +11,13 @@ import logging
 import string
 from contextlib import contextmanager
 from time import time
-from typing import Sequence, Tuple, Iterable
+from typing import Iterable, Sequence, Tuple
 
 import numpy as np
 import pandas as pd
 from tfs import TfsDataFrame
 
-from optics_functions.constants import (REAL, IMAG, PLANES, PHASE_ADV,
-                                        X, Y, S, NAME, DELTA_ORBIT)
+from optics_functions.constants import DELTA_ORBIT, IMAG, NAME, PHASE_ADV, PLANES, REAL, S, X, Y
 
 LOG = logging.getLogger(__name__)
 D = DELTA_ORBIT
@@ -26,12 +25,14 @@ D = DELTA_ORBIT
 # DataFrames -------------------------------------------------------------------
 
 
-def prepare_twiss_dataframe(df_twiss: TfsDataFrame,
-                            df_errors: pd.DataFrame = None,
-                            invert_signs_madx: bool = False,
-                            max_order: int = 16,
-                            join: str = "inner") -> TfsDataFrame:
-    """ Prepare dataframe to use with the optics functions.
+def prepare_twiss_dataframe(
+    df_twiss: TfsDataFrame,
+    df_errors: pd.DataFrame = None,
+    invert_signs_madx: bool = False,
+    max_order: int = 16,
+    join: str = "inner",
+) -> TfsDataFrame:
+    """Prepare dataframe to use with the optics functions.
 
     - Adapt Beam 4 signs.
     - Add missing K(S)L and orbit columns.
@@ -87,14 +88,14 @@ def prepare_twiss_dataframe(df_twiss: TfsDataFrame,
     for name, df_old in (("twiss", df_twiss), ("errors", df_errors)):
         dropped_columns = set(df_old.columns) - set(df.columns)
         if dropped_columns:
-            LOG.warning(f"The following {name}-columns were"
-                        f" dropped on merge: {seq2str(dropped_columns)}")
+            LOG.warning(f"The following {name}-columns were dropped on merge: {seq2str(dropped_columns)}")
     return df
 
 
-def split_complex_columns(df: pd.DataFrame, columns: Sequence[str],
-                          drop: bool = True) -> TfsDataFrame:
-    """ Splits the given complex columns into two real-values columns containing the
+def split_complex_columns(
+    df: pd.DataFrame, columns: Sequence[str], drop: bool = True
+) -> TfsDataFrame:
+    """Splits the given complex columns into two real-values columns containing the
     real and imaginary parts of the original columns.
 
     Args:
@@ -115,28 +116,46 @@ def split_complex_columns(df: pd.DataFrame, columns: Sequence[str],
     return df
 
 
-def switch_signs_for_beam4(df_twiss: pd.DataFrame,
-                           df_errors: pd.DataFrame = None) -> Tuple[TfsDataFrame, TfsDataFrame]:
-    """ Switch the signs for Beam 4 optics.
+def merge_complex_columns(
+    df: pd.DataFrame, columns: Sequence[str], drop: bool = True
+) -> TfsDataFrame:
+    """Merges the given real and imag columns into complex columns.
+
+    Args:
+        df (pd.DataFrame): DataFrame containing the original columns.
+        columns (Sequence[str]): List of complex columns names to be created.
+        drop (bool): Original columns are not present in resulting DataFrame.
+
+    Returns:
+        Original TfsDataFrame with added columns.
+    """
+    df = df.copy()
+    for column in columns:
+        df[column] = df[f"{column}{REAL}"] + 1j * df[f"{column}{IMAG}"]
+
+    if drop:
+        df = df.drop(columns=[f"{column}{part}" for column in columns for part in (REAL, IMAG)])
+    return df
+
+
+def switch_signs_for_beam4(
+    df_twiss: pd.DataFrame, df_errors: pd.DataFrame = None
+) -> Tuple[TfsDataFrame, TfsDataFrame]:
+    """Switch the signs for Beam 4 optics.
     This is due to the switch in direction for this beam and
     (anti-) symmetry after a rotation of 180deg around the y-axis of magnets,
     combined with the fact that the KL values in MAD-X twiss do not change sign,
     but in the errors they do (otherwise it would compensate).
     Magnet orders that show anti-symmetry are: a1 (K0SL), b2 (K1L), a3 (K2SL), b4 (K3L) etc.
     Also the sign for (delta) X is switched back to have the same orientation as beam2."""
-    LOG.debug(
-        f"Switching signs for X and K(S)L values when needed, to match Beam 4 to Beam 2."
-    )
+    LOG.debug(f"Switching signs for X and K(S)L values when needed, to match Beam 4 to Beam 2.")
     df_twiss, df_errors = df_twiss.copy(), df_errors.copy()
     df_twiss[X] = -df_twiss[X]
 
     if df_errors is not None:
         df_errors[f"{D}{X}"] = -df_errors[f"{D}{X}"]
         max_order = (
-            df_errors.columns.str.extract(r"^K(\d+)S?L$", expand=False)
-            .dropna()
-            .astype(int)
-            .max()
+            df_errors.columns.str.extract(r"^K(\d+)S?L$", expand=False).dropna().astype(int).max()
         )
         for order in range(max_order + 1):
             name = f"K{order:d}{'' if order % 2 else 'S'}L"  # odd -> '', even -> S
@@ -149,7 +168,7 @@ def switch_signs_for_beam4(df_twiss: pd.DataFrame,
 
 
 def get_all_phase_advances(df: pd.DataFrame) -> dict:
-    """ Calculate phase advances between all elements.
+    """Calculate phase advances between all elements.
     Will result in a elements x elements matrix, that might be very large!
 
     Args:
@@ -175,7 +194,7 @@ def get_all_phase_advances(df: pd.DataFrame) -> dict:
 
 
 def dphi(data: np.ndarray, q: float) -> np.ndarray:
-    """ Return dphi from phase advances in data, see Eq. (8) in [FranchiAnalyticFormulas2017]_ .
+    """Return dphi from phase advances in data, see Eq. (8) in [FranchiAnalyticFormulas2017]_ .
 
     Args:
         data (DataFrame, Series): Phase-Advance data.
@@ -201,7 +220,7 @@ def tau(data: np.ndarray, q: float) -> np.ndarray:
 
 
 def dphi_at_element(df: pd.DataFrame, element: str, qx: float, qy: float) -> dict:
-    """ Return dphis for both planes at the given element.
+    """Return dphis for both planes at the given element.
     See Eq. (8) in [FranchiAnalyticFormulas2017]_ .
 
     Args:
@@ -239,7 +258,7 @@ def add_missing_columns(df: pd.DataFrame, columns: Iterable) -> pd.DataFrame:
 
 @contextmanager
 def timeit(text: str = "Time used {:.3f}s", print_fun=LOG.debug):
-    """ Timing context with logging/printing output.
+    """Timing context with logging/printing output.
 
     Args:
         text (str): Text to print. If it contains an unnamed format key, this
@@ -282,8 +301,6 @@ def set_name_index(df: pd.DataFrame, df_name="") -> pd.DataFrame:
         df = df.set_index(NAME)
 
     if not all(isinstance(indx, str) for indx in df.index):
-        raise TypeError(
-            f"Index of the {df_name} Dataframe should be string (i.e. from {NAME})"
-        )
+        raise TypeError(f"Index of the {df_name} Dataframe should be string (i.e. from {NAME})")
 
     return df


Test patch:
diff --git a/tests/unit/test_coupling.py b/tests/unit/test_coupling.py
index c3731e0..aa4ce44 100644
--- a/tests/unit/test_coupling.py
+++ b/tests/unit/test_coupling.py
@@ -6,7 +6,7 @@ import tfs
 
 from optics_functions.constants import NAME, S, ALPHA, Y, BETA, X, GAMMA, REAL, IMAG, TUNE, PHASE_ADV
 from optics_functions.coupling import (closest_tune_approach, coupling_via_rdts,
-                                       coupling_via_cmatrix, COUPLING_RDTS)
+                                       coupling_via_cmatrix, COUPLING_RDTS, rmatrix_from_coupling)
 from optics_functions.utils import prepare_twiss_dataframe
 from tests.unit.test_rdt import arrays_are_close_almost_everywhere
 
@@ -23,9 +23,10 @@ def test_cmatrix():
     assert all(c in df_res.columns for c in ("F1001", "F1010", "C11", "C12", "C21", "C22", GAMMA))
     assert not df_res.isna().any().any()
 
-    detC = (df_res["C11"]*df_res["C22"] - df_res["C12"]*df_res["C21"])
+    # Checks based on CalagaBetatronCoupling2005
+    detC = (df_res["C11"] * df_res["C22"] - df_res["C12"] * df_res["C21"])
     fsq_diff = np.abs(df_res["F1001"])**2 - np.abs(df_res["F1010"])**2
-    f_term = 1/(1+4*fsq_diff)
+    f_term = 1/(1 + 4 * fsq_diff)
     g_sq = df_res[GAMMA]**2
     assert all(np.abs(detC + g_sq - 1) < 1e-15)
     assert all(np.abs(detC / (4 * g_sq) - fsq_diff) < 1e-15)  # Eq. (13)
@@ -33,6 +34,29 @@ def test_cmatrix():
     assert all(np.abs(g_sq - f_term) < 1e-15)  # Eq. (14)
 
 
+@pytest.mark.basic
+@pytest.mark.parametrize('source', ['real', 'fake'])
+def test_rmatrix_to_coupling_to_rmatrix(source):
+    if source == "fake":
+        n = 5
+        np.random.seed(487423872)
+        df = get_df(n)
+    else:
+        df = tfs.read(INPUT / "coupling_bump" / f"twiss.lhc.b1.coupling_bump.tfs", index=NAME)
+
+    df_coupling = coupling_via_cmatrix(df)
+    for col in ("ALFX", "BETX", "ALFY", "BETY"):
+        df_coupling[col] = df[col]
+
+    df_res = rmatrix_from_coupling(df_coupling)
+
+    for col in ("R11", "R12", "R21", "R22"):
+        # For debugging:
+        # print(col)
+        # print(max(np.abs(df[col] - df_res[col])))
+        assert all(np.abs(df[col] - df_res[col]) < 5e-15)
+
+
 @pytest.mark.basic
 def test_real_output():
     n = 7
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index 68162a5..96efe9a 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -10,7 +10,7 @@ from optics_functions.constants import PHASE_ADV, Y, X, REAL, IMAG, NAME, DELTA_
 from optics_functions.utils import (add_missing_columns, dphi,
                                     get_all_phase_advances, tau, seq2str, i_pow,
                                     prepare_twiss_dataframe, switch_signs_for_beam4,
-                                    get_format_keys, dphi_at_element, split_complex_columns)
+                                    get_format_keys, dphi_at_element, split_complex_columns, merge_complex_columns)
 
 INPUT = Path(__file__).parent.parent / "inputs"
 
@@ -93,6 +93,21 @@ def test_split_complex_columns():
             assert (fun(df[col]) == df_split[f"{col}{part}"]).all()
 
 
+@pytest.mark.basic
+def test_merge_complex_columns():
+    df = pd.DataFrame([1+2j, 3j + 4], columns=["Col"], index=["A", "B"])
+    df_split = split_complex_columns(df, df.columns, drop=True)
+
+    df_merged = merge_complex_columns(df_split, df.columns, drop=False)
+    assert len(df_merged.columns) == 3
+
+    df_merged = merge_complex_columns(df_split, df.columns, drop=True)
+    assert len(df_merged.columns) == 1
+
+    for col in df.columns:
+        assert (df[col] == df_merged[col]).all()
+
+
 @pytest.mark.basic
 def test_dphi():
     n, qx, qy = 5, 4, 5