Fixes main import and deprecation warnings (#14)

* fix absolute import
* update gitignore
* Silence rcond warning in lstsq in test
* use math instead of np.math
* added commandline tests
* bumpversion

.gitignore

@@ -260,4 +260,7 @@ Temporary Items
 *.out
 *.synctex.*
 *.toc
-*.fdb_latexmk
+*.fdb_latexmk
+
+# Personal testing scripts
+tst_*

irnl_rdt_correction/__init__.py

@@ -13,7 +13,7 @@
 __title__ = "irnl-rdt-correction"
 __description__ = "Correction script to power the nonlinear correctors in the (HL-)LHC insertion regions based on RDTs."
 __url__ = "https://github.com/pylhc/irnl_rdt_correction"
-__version__ = "1.1.1"
+__version__ = "1.1.2"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"

irnl_rdt_correction/__main__.py

@@ -1,5 +1,5 @@
 from irnl_rdt_correction.main import irnl_rdt_correction
-from utilities import log_setup
+from irnl_rdt_correction.utilities import log_setup
 
 if __name__ == '__main__':
     log_setup()

irnl_rdt_correction/equation_system.py

@@ -7,6 +7,7 @@
 
 """
 import logging
+import math
 from typing import Dict, List, Sequence, Set, Tuple
 
 import numpy as np
@@ -354,7 +355,7 @@ def get_elements_integral(rdt: RDT, ip: int, optics: Optics, feeddown: int) -> f
             iksl_err = 1j*errors_df.loc[elements, f"K{n_mad:d}SL"]
 
             k_sum += ((kl_opt + kl_err + iksl_opt + iksl_err) *
-                      (dx_idy**q) / np.math.factorial(q))
+                      (dx_idy**q) / math.factorial(q))
 
         # note the minus sign before the sum!
         integral += -sum(np.real(i_pow(lm) * k_sum.to_numpy()) * (side_sign * betax * betay).to_numpy())
@@ -406,7 +407,7 @@ def get_corrector_coefficient(rdt: RDT, corrector: IRCorrector, optics: Optics)
         dx = twiss_df.loc[corrector.name, X] + errors_df.loc[corrector.name, f"{DELTA}{X}"]
         dy = twiss_df.loc[corrector.name, Y] + errors_df.loc[corrector.name, f"{DELTA}{Y}"]
         dx_idy = dx + 1j*dy
-        z_cmplx = (dx_idy**p) / np.math.factorial(p)  # Eq. (32)
+        z_cmplx = (dx_idy**p) / math.factorial(p)  # Eq. (32)
 
         # Get the correct part of z_cmplx, see Eq. (36) in [DillyNonlinearIRCorrections2023]_
         if (corrector.skew and is_odd(lm)) or (not corrector.skew and is_even(lm)):

irnl_rdt_correction/io_handling.py

@@ -220,8 +220,9 @@ def _check_dfs(beams: Sequence[int], twiss_dfs: Sequence[DataFrame], errors_dfs:
         - All needed field strengths are present in twiss
           -> either Error or Warning depending on ``ignore_missing_columns``.
     """
-    twiss_dfs, errors_dfs = tuple(tuple(df.copy() for df in dfs) 
-                                  for dfs in (twiss_dfs, errors_dfs))
+    twiss_dfs, errors_dfs = tuple(tuple(
+        convert_numeric_columns_to_float(df) for df in dfs) for dfs in (twiss_dfs, errors_dfs)
+    )
 
     needed_twiss = list(PLANES)
     needed_errors = [f"{DELTA}{p}" for p in PLANES]
@@ -258,7 +259,7 @@ def _check_dfs(beams: Sequence[int], twiss_dfs: Sequence[DataFrame], errors_dfs:
                       f"the errors: {list2str(not_found_optics.to_list())}."
                       f"They are assumed to be zero.")
             for indx in not_found_optics:
-                errors.loc[indx, :] = 0
+                errors.loc[indx, :] = 0.0
 
         for df, needed_cols, file_type in ((
             twiss, needed_twiss, "twiss"), (errors, needed_errors, "error")
@@ -273,10 +274,21 @@ def _check_dfs(beams: Sequence[int], twiss_dfs: Sequence[DataFrame], errors_dfs:
                     raise IOError(text)
                 LOG.warning(text + " They are assumed to be zero.")
                 for kl in missing_cols:
-                    df[kl] = 0
+                    df[kl] = 0.0
     return beams, twiss_dfs, errors_dfs
 
 
+def convert_numeric_columns_to_float(df: TfsDataFrame) -> TfsDataFrame:
+    """ Convert numeric columns in df to float. 
+    This avoids that the user accidentally gives columns with dtype int (e.g. all 0), 
+    in which case assigning float values will fail in some pandas version after 2.1.0 
+    (which had a deprecation warning). """
+    df = df.copy()
+    numeric_columns = df.select_dtypes(include=['number']).columns
+    df[numeric_columns] = df[numeric_columns].astype('float64')
+    return df
+
+
 def maybe_switch_signs(optics: Optics):
     """ Switch the signs for Beam optics.
      This is due to the switch in direction between beam and

tests/helpers.py

@@ -60,7 +60,7 @@ def read_lhc_model(beam: int) -> TfsDataFrame:
 
 
 def generate_pseudo_model(n_ips: int, n_magnets: int, accel: str,
-                          betax: float = 1, betay: float = 1, x: float = 0, y: float = 0) -> pd.DataFrame:
+                          betax: float = 1.0, betay: float = 1.0, x: float = 0.0, y: float = 0.0) -> pd.DataFrame:
     """Generate a Twiss-Like DataFrame with magnets as index and Beta and Orbit columns."""
     df = pd.DataFrame(
         index=(
@@ -77,7 +77,7 @@ def generate_pseudo_model(n_ips: int, n_magnets: int, accel: str,
     return df
 
 
-def generate_errortable(index: Sequence, value: float = 0, max_order: int = MAX_N) -> pd.DataFrame:
+def generate_errortable(index: Sequence, value: float = 0.0, max_order: int = MAX_N) -> pd.DataFrame:
     """Return DataFrame from index and KN(S)L + D[XY] columns."""
     return pd.DataFrame(value,
                         index=index,

tests/test_commandline.py

@@ -0,0 +1,141 @@
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import tfs
+
+from tests.helpers import (CIRCUIT, EPS, IP, MAX_N, STRENGTH, VALUE, generate_errortable,
+                           generate_pseudo_model, get_some_magnet_names)
+
+
+def test_package_execution_no_arguments():
+    """ Tests if the package can be run as a module. 
+    Without arguments it should abort with an error message stating required arguments. """
+    # Run the package as a module using subprocess
+    result = subprocess.run([sys.executable, "-m", "irnl_rdt_correction"], capture_output=True, text=True)
+
+    # Check if the command executed with sysexit (exit code 2)
+    assert result.returncode == 2
+
+    # Check for expected output
+    expected_output = "error: the following arguments are required: --beams, --twiss"
+    assert expected_output in result.stderr
+
+
+def test_package_execution_fix_v1_1_2():
+    """ Tests if the package can be run as a module. This failed in <1.1.2 with an import error. """
+    # Run the package as a module using subprocess
+    result = subprocess.run([sys.executable, "-m", "irnl_rdt_correction"], capture_output=True, text=True)
+
+    # Check if the command executed not with an import error (exit code 1)
+    assert result.returncode != 1
+
+    # Check for a module not found error
+    not_expected_output = "ModuleNotFoundError: No module named 'utilities'"  
+    assert not_expected_output not in result.stderr
+
+
+@pytest.mark.parametrize('accel', ('lhc', 'hllhc'))
+def test_basic_correction(tmp_path: Path, accel: str):
+    """Tests the basic correction functionality and performs some sanity checks.
+    Same as in tets_standard_correction, but less testcases and called from commandline.
+
+    Operates on a pseudo-model so that the corrector values are easily known.
+    Sanity Checks:
+    - all correctors found
+    - correctors have the correct value (as set by errors or zero)
+    - all corrector circuits present in madx-script
+    """
+    # Parameters -----------------------------------------------------------
+    order = 4
+    orientation = ""
+
+    correct_ips = (1, 3)
+    error_value = 2
+    n_magnets = 4
+    n_ips = 4
+
+    n_correct_ips = len(correct_ips)
+    n_sides = len("LR")
+    n_orientation = len(["S", ""])
+
+    # Setup ----------------------------------------------------------------
+    twiss = generate_pseudo_model(accel=accel, n_ips=n_ips, n_magnets=n_magnets)
+    errors = generate_errortable(index=get_some_magnet_names(n_ips=n_ips, n_magnets=n_magnets))
+    error_component = f"K{order-1}{orientation}L"
+    errors[error_component] = error_value
+
+    if order % 2:  # order is odd -> sides have different sign in rdt
+        left_hand_magnets = errors.index.str.match(r".*L\d$")
+        errors.loc[left_hand_magnets, error_component] = errors.loc[left_hand_magnets, error_component] / 2  # so they don't fully compensate
+
+    twiss_path = tmp_path / "twiss_input.tfs"
+    errors_path = tmp_path / "errors_input.tfs"
+    result_path = tmp_path / "correct"
+
+    tfs.write(twiss_path, twiss, save_index="NAME")
+    tfs.write(errors_path, errors, save_index="NAME")
+
+    # Correction -----------------------------------------------------------
+    subprocess.run([sys.executable, "-m", 
+                    "irnl_rdt_correction",
+                    "--accel", accel,
+                    "--twiss", twiss_path,
+                    "--errors", errors_path,
+                    "--beams", "1",
+                    "--output", result_path,
+                    "--feeddown", "0",
+                    "--ips", *[str(ip) for ip in correct_ips],
+                    "--ignore_missing_columns",
+                    "--iterations", "1",
+                    ],
+                   capture_output=True, text=True
+    )
+
+    madx_corrections = result_path.with_suffix(".madx").read_text()
+    df_corrections = tfs.read(result_path.with_suffix(".tfs"))
+
+    # Testing --------------------------------------------------------------
+    # Same as in test_standard_correction - TODO: make function?
+    # Check output data ---
+    assert len(list(tmp_path.glob("correct.*"))) == 2
+
+    # Check all found correctors ---
+    if accel == 'lhc':
+        assert len(df_corrections.index) == (
+                n_orientation * n_sides * n_correct_ips * len("SO") +
+                n_sides * n_correct_ips * len("T")
+        )
+
+    if accel == 'hllhc':
+        assert len(df_corrections.index) == n_orientation * n_sides * n_correct_ips * len("SODT")
+
+    # All circuits in madx script ---
+    for circuit in df_corrections[CIRCUIT]:
+        assert circuit in madx_corrections
+
+    # Check corrector values ---
+    for test_order in range(3, MAX_N+1):
+        for test_orientation in ("S", ""):
+            for ip in correct_ips:
+                mask = (
+                        (df_corrections[STRENGTH] == f"K{test_order-1}{test_orientation}L") &
+                        (df_corrections[IP] == ip)
+                )
+                if (test_order == order) and (test_orientation == orientation):
+                    if order % 2:
+                        corrector_strengths = sum(df_corrections.loc[mask, VALUE])
+                        assert abs(corrector_strengths) < EPS  # correctors should be equally distributed
+
+                        corrector_strengths = -sum(df_corrections.loc[mask, VALUE].abs())
+                        # as beta cancels out (and is 1 anyway)
+                        error_strengths = n_magnets * error_value / 2  # account for partial compensation (from above)
+                    else:
+                        corrector_strengths = sum(df_corrections.loc[mask, VALUE])
+                        assert all(abs(df_corrections.loc[mask, VALUE] - corrector_strengths / n_sides) < EPS)
+                        # as beta cancels out (and is 1 anyway)
+                        error_strengths = (n_sides * n_magnets * error_value)
+                    assert abs(corrector_strengths + error_strengths) < EPS  # compensation of RDT
+                else:
+                    assert all(df_corrections.loc[mask, VALUE] == 0.)

tests/test_dual_optics_correction.py

@@ -85,7 +85,8 @@ def test_dual_optics(tmp_path: Path):
     b1 = beta**(exp_x+exp_y)
     b2 = beta2**(exp_x+exp_y)
     dual_correction = np.linalg.lstsq(np.array([[b1, b1], [b2, b2]]),
-                                      np.array([-b1*error_strengths1, -b2*error_strengths2]))[0]
+                                      np.array([-b1*error_strengths1, -b2*error_strengths2]), 
+                                      rcond=None)[0]
 
     assert all(np.abs(dual_correction) > 0)  # just for safety, that there is a solution
 

tests/test_standard_correction.py

@@ -228,8 +228,8 @@ class BeamResult:
     # here: 2 == sextupole
     kl_columns = [f"K{order}{orientation}L" for order in orders for orientation in orientations]
 
-    twiss_0.loc[:, kl_columns]  = 0
-    errors_0.loc[:, kl_columns]  = 0
+    twiss_0.loc[:, kl_columns]  = 0.0
+    errors_0.loc[:, kl_columns]  = 0.0
 
     # orbit at corrector should not matter, as we don't use feed-down to correct here)
     # twiss_0.loc[:, [X, Y]] = 0.5
@@ -243,23 +243,23 @@ class BeamResult:
             if magnet_order in ("D", "DS") and is_lhc:
                 continue
             corrector_magnets = twiss_0.index.str.match(rf"MC{magnet_order}X") 
-            twiss_0.loc[corrector_magnets, kl] = 1
+            twiss_0.loc[corrector_magnets, kl] = 1.0
 
     # Power other magnets and assign errors
     non_corrector_magnets = twiss_0.index.str.match(r"M.\.")  # M[A-Z]. is created above
-    twiss_0.loc[non_corrector_magnets, kl_columns] = 1
+    twiss_0.loc[non_corrector_magnets, kl_columns] = 1.0
     twiss_0.loc[non_corrector_magnets, [X, Y]] = 0.5
 
     non_corrector_magnets = errors_0.index.str.match(r"M.\.")  # M[A-Z]. is created above
-    errors_0.loc[non_corrector_magnets, kl_columns] = 1
+    errors_0.loc[non_corrector_magnets, kl_columns] = 1.0
     errors_0.loc[non_corrector_magnets, [f"{DELTA}{X}", f"{DELTA}{Y}"]] = 0.5
 
     # modify the left side, so they don't fully compensate for even (madx)-orders
     left_hand_magnets = twiss_0.index.str.match(r".*L\d$")
-    twiss_0.loc[left_hand_magnets, f"{BETA}{Y}"] = 2 * twiss_0.loc[left_hand_magnets, f"{BETA}{Y}"]
+    twiss_0.loc[left_hand_magnets, f"{BETA}{Y}"] = 2.0 * twiss_0.loc[left_hand_magnets, f"{BETA}{Y}"]
 
     left_hand_magnets = errors_0.index.str.match(r".*L\d$")
-    errors_0.loc[left_hand_magnets, kl_columns] = errors_0.loc[left_hand_magnets, kl_columns] / 2  
+    errors_0.loc[left_hand_magnets, kl_columns] = errors_0.loc[left_hand_magnets, kl_columns] / 2.0  
 
     # Pre-calculate the integral based on this setup. 
     integral_left = 1.5 * n_magnets