Feature constant detuning term

pywit/landau_damping.py

@@ -1,6 +1,6 @@
 import numpy as np
 from scipy.special import exp1
-from scipy.optimize import newton
+from scipy.optimize import bisect, newton
 
 from typing import Sequence, Tuple
 
@@ -59,26 +59,32 @@ def dispersion_integral_2d(tune_shift: np.ndarray, b_direct: float, b_cross: flo
     return -i
 
 
-def find_detuning_coeffs_threshold(tune_shift: complex, q_s: float, b_direct_ref: float, b_cross_ref: float,
+def find_detuning_coeffs_threshold(tune_shift: complex, q_s: float, reference_b_direct: float, reference_b_cross: float,
+                                   added_b_direct: float = 0, added_b_cross: float = 0,
                                    fraction_of_qs_allowed_on_positive_side: float = 0.05,
-                                   distribution: str = 'gaussian', tolerance=1e-10):
+                                   distribution: str = 'gaussian', algorithm: str = 'newton', tolerance=1e-10):
     """
     Compute the detuning coefficients (multiplied by sigma) corresponding to stability diagram threshold for a complex
     tune shift.
-    It keeps fixed the ratio between b_direct_ref and b_cross_ref.
+    It keeps fixed the ratio between reference_b_direct and reference_b_cross.
     :param tune_shift: the tune shift for which the octupole threshold is computed
     :param q_s: the synchrotron tune
-    :param b_direct_ref: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in
+    :param reference_b_direct: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in
     the x plane or $\alpha_y \sigma_y$ if working in the y plane)
-    :param b_cross_ref: the cross detuning coefficient multiplied by sigma (i.e. $\alpha_{xy} \sigma_y$ if working in
+    :param reference_b_cross: the cross detuning coefficient multiplied by sigma (i.e. $\alpha_{xy} \sigma_y$ if working in
     the x plane or $\alpha_{yx} \sigma_x$ if working in the y plane)
+    :param added_b_direct: an additional contribution to the direct detuning coefficient which is kept fixed in the
+    procedure to find the stability threshold
+    :param added_b_cross: an additional contribution to the cross detuning coefficient which is kept fixed in the
+    procedure to find the stability threshold
     :param distribution: the transverse distribution of the beam. It can be 'gaussian' or 'parabolic'
     :param fraction_of_qs_allowed_on_positive_side: to determine azimuthal mode number l_mode (around which is drawn the
     stability diagram), one can consider positive tune shift up to this fraction of q_s (default=5%)
-    :param tolerance: tolerance on difference w.r.t stability diagram, for Newton's root finding
+    :param algorithm: root-solving algorithm ('bisect' or 'newton')
+    :param tolerance: tolerance on difference w.r.t stability diagram, for root finding
     and for the final check that the roots are actually proper roots.
     :return: the detuning coefficients corresponding to the stability diagram threshold if the corresponding mode is
-    unstable, 0 if the corresponding mode is stable or np.nan if the threshold cannot be found (failure of Newton's
+    unstable, 0 if the corresponding mode is stable or np.nan if the threshold cannot be found (failure of root-solving
     algorithm).
     """
     # evaluate azimuthal mode number
@@ -88,23 +94,35 @@ def find_detuning_coeffs_threshold(tune_shift: complex, q_s: float, b_direct_ref
     # take away the shift from azimuthal mode number
     tune_shift -= q_s * l_mode
 
-    b_ratio = b_cross_ref/b_direct_ref
+    b_ratio = reference_b_cross / reference_b_direct
     if tune_shift.imag < 0.:
 
         # function to solve (distance in imag. part w.r.t stab. diagram, as a function of oct. current)
         def f(b_direct):
-            b_direct_i = b_direct
-            b_cross_i = b_ratio * b_direct
+            b_direct_i = b_direct + added_b_direct
+            b_cross_i = b_ratio * b_direct + added_b_cross
             stab = [-1. / dispersion_integral_2d(t_s, b_direct_i, b_cross_i, distribution=distribution) for e in (-1, 1)
                     for t_s in b_direct_i * e * 10. ** np.arange(-3, 2, 0.01)[::e]]
             # note: one has to reverse the table to get the interpolation right, for negative polarity (np.interp always
             # wants monotonically increasing abscissae)
-            return tune_shift.imag - np.interp(tune_shift.real, np.real(stab)[::int(np.sign(b_direct_ref))],
-                                               np.imag(stab)[::int(np.sign(b_direct_ref))])
+            return tune_shift.imag - np.interp(tune_shift.real, np.real(stab)[::int(np.sign(b_direct_i))],
+                                               np.imag(stab)[::int(np.sign(b_direct_i))])
 
-        # Newton root finding
+        # root finding
         try:
-            b_direct_new = newton(f, b_direct_ref, tol=tolerance)
+            if algorithm=='bisect':
+                # we use 1e-15 as a bound instead of 0 because 0 would cause a divsion by 0 in dispersion_integral_2d
+                if reference_b_direct > 0:
+                    b_direct_new = bisect(f, 1e-15, 100 * reference_b_direct)
+                else:
+                    b_direct_new = bisect(f, 100 * reference_b_direct, -1e-15)
+
+            elif algorithm=='newton':
+                b_direct_new = newton(f, reference_b_direct, tol=tolerance)
+
+            else:
+                raise ValueError("algorithm must be either 'newton' or 'bisect'")
+
         except RuntimeError:
             b_direct_new = np.nan
         else:
@@ -123,33 +141,40 @@ def abs_first_item_or_nan(tup: Tuple):
         return np.nan
 
 
-def find_detuning_coeffs_threshold_many_tune_shifts(tune_shifts: Sequence[complex], q_s: float, b_direct_ref: float,
-                                                    b_cross_ref: float, distribution: str = 'gaussian',
+def find_detuning_coeffs_threshold_many_tune_shifts(tune_shifts: Sequence[complex], q_s: float,
+                                                    reference_b_direct: float, reference_b_cross: float,
+                                                    added_b_direct: float = 0, added_b_cross: float = 0,
+                                                    distribution: str = 'gaussian',
                                                     fraction_of_qs_allowed_on_positive_side: float = 0.05,
-                                                    tolerance=1e-10):
+                                                    algorithm: str = 'newton', tolerance = 1e-10):
     """
     Compute the detuning coefficients corresponding to the most stringent stability diagram threshold for a sequence of
-    complex tune shifts. It keeps fixed the ratio between b_direct_ref and b_cross_ref.
+    complex tune shifts. It keeps fixed the ratio between reference_b_direct and reference_b_cross.
     :param tune_shifts: the sequence of complex tune shifts
     :param q_s: the synchrotron tune
-    :param b_direct_ref: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in
+    :param reference_b_direct: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in
     the x plane or $\alpha_y \sigma_y$ if working in the y plane)
-    :param b_cross_ref: the cross detuning coefficient multiplied by sigma (i.e. $\alpha_{xy} \sigma_y$ if working in
+    :param reference_b_cross: the cross detuning coefficient multiplied by sigma (i.e. $\alpha_{xy} \sigma_y$ if working in
     the x plane or $\alpha_{yx} \sigma_x$ if working in the y plane)
+    :param added_b_direct: an additional contribution to the direct detuning coefficient which is kept fixed in the
+    procedure to find the stability threshold
+    :param added_b_cross: an additional contribution to the cross detuning coefficient which is kept fixed in the
+    procedure to find the stability threshold
     :param distribution: the transverse distribution of the beam. It can be 'gaussian' or 'parabolic'
     :param fraction_of_qs_allowed_on_positive_side: to determine azimuthal mode number l_mode (around which is drawn the
     stability diagram), one can consider positive tuneshift up to this fraction of q_s (default=5%)
-    :param tolerance: tolerance on difference w.r.t stability diagram, for Newton's root finding
+    :param algorithm: root-solving algorithm ('bisect' or 'newton')
+    :param tolerance: tolerance on difference w.r.t stability diagram, for root finding
     and for the final check that the roots are actually proper roots.
     :return: the detuning coefficients corresponding to the most stringent stability diagram threshold for all the
     given tune shifts if the corresponding mode is unstable, 0 if all modes are stable or np.nan if the
-    no threshold can be found (failure of Newton's algorithm).
+    no threshold can be found (failure of root-solving algorithm).
     """
     # find max octupole current required from a list of modes, given their tuneshifts
     b_coefficients = np.array([find_detuning_coeffs_threshold(
-        tune_shift=tune_shift, q_s=q_s, b_direct_ref=b_direct_ref,
-        b_cross_ref=b_cross_ref, distribution=distribution,
+        tune_shift=tune_shift, q_s=q_s, reference_b_direct=reference_b_direct,
+        reference_b_cross=reference_b_cross, added_b_direct=added_b_direct, added_b_cross=added_b_cross, distribution=distribution,
         fraction_of_qs_allowed_on_positive_side=fraction_of_qs_allowed_on_positive_side,
-        tolerance=tolerance) for tune_shift in tune_shifts if tune_shift is not np.nan])
+        algorithm=algorithm, tolerance=tolerance) for tune_shift in tune_shifts if tune_shift is not np.nan])
 
     return max(b_coefficients, key=abs_first_item_or_nan)

test/test_landau_damping.py

@@ -1,86 +1,79 @@
 import pytest
+from pytest import mark, raises
 import numpy as np
 
 from pywit.landau_damping import dispersion_integral_2d, find_detuning_coeffs_threshold
 from pywit.landau_damping import find_detuning_coeffs_threshold_many_tune_shifts
 
 
-def test_dispersion_integral_2d():
-
-    distribution = 'gaussian'
-    tune_shift = -7.3622423693e-05-3.0188372754e-06j
+@mark.parametrize('distribution, expected_value',
+                  [['gaussian', 7153.859519171599 - 2722.966574677259j],
+                   ['parabolic', 6649.001778623455 - 3168.4257879737443j],
+                   ])
+def test_dispersion_integral_2d(distribution: str, expected_value: complex):
+    # reference values obtained with DELPHI
+    tune_shift = -7.3622423693e-05 - 3.0188372754e-06j
     b_direct = 7.888357197059519e-05
     b_cross = -5.632222163778799e-05
 
-    # reference value obtained with DELPHI
-    expected_value = 7153.859519171599-2722.966574677259j
-
-    assert np.isclose(np.real(expected_value), np.real(dispersion_integral_2d(tune_shift=tune_shift,
-                                                                              b_direct=b_direct,
-                                                                              b_cross=b_cross,
-                                                                              distribution=distribution)))
-    assert np.isclose(np.imag(expected_value), np.imag(dispersion_integral_2d(tune_shift=tune_shift,
-                                                                              b_direct=b_direct,
-                                                                              b_cross=b_cross,
-                                                                              distribution=distribution)))
+    test_value = dispersion_integral_2d(tune_shift=tune_shift, b_direct=b_direct, b_cross=b_cross,
+                                        distribution=distribution)
 
-    distribution = 'parabolic'
+    assert np.isclose(np.real(expected_value), np.real(test_value))
+    assert np.isclose(np.imag(expected_value), np.imag(test_value))
 
-    # reference value obtained with DELPHI
-    expected_value = 6649.001778623455-3168.4257879737443j
 
-    assert np.isclose(np.real(expected_value), np.real(dispersion_integral_2d(tune_shift=tune_shift,
-                                                                              b_direct=b_direct,
-                                                                              b_cross=b_cross,
-                                                                              distribution=distribution)))
-    assert np.isclose(np.imag(expected_value), np.imag(dispersion_integral_2d(tune_shift=tune_shift,
-                                                                              b_direct=b_direct,
-                                                                              b_cross=b_cross,
-                                                                              distribution=distribution)))
+def test_wrong_algo_in_find_detuning_coeffs_threshold():
+    with raises(ValueError, match="algorithm must be either 'newton' or 'bisect'"):
+        _ = find_detuning_coeffs_threshold(tune_shift=1e-5-1j*1e-5, q_s=2e-3,
+                                           reference_b_direct=1e-15,
+                                           reference_b_cross=1e-15,
+                                           algorithm='newto')
 
 
-def test_find_detuning_coeffs_threshold():
+@mark.parametrize('algorithm',['newton','bisect'])
+@mark.parametrize('b_direct, b_cross, added_b_direct, added_b_cross',
+                  [[1.980315192200037e-05, -1.4139287608495406e-05, 0, 0],
+                   [-1.2718161244917965e-05, 9.08066253298482e-06, 0, 0],
+                   [1.980315192200037e-05, -1.4139287608495406e-05, 1.980315192200037e-05 / 3,
+                    -1.4139287608495406e-05 / 3],
+                   [-1.2718161244917965e-05, 9.08066253298482e-06, -1.2718161244917965e-05 / 3,
+                    9.08066253298482e-06 / 3],
+                   ])
+def test_find_detuning_coeffs_threshold(algorithm: str, b_direct: float,
+            b_cross: float, added_b_direct: float, added_b_cross: float):
     # reference values obtained with the old impedance wake model https://gitlab.cern.ch/IRIS/HLLHC_IW_model
     # test positive octupole polarity
     tune_shift = -7.3622423693e-05 - 3.0188372754e-06j
     q_s = 2e-3
-    b_direct = 1.980315192200037e-05
-    b_cross = -1.4139287608495406e-05
-    assert np.isclose(b_direct, find_detuning_coeffs_threshold(tune_shift=tune_shift, q_s=q_s, b_direct_ref=b_direct,
-                                                               b_cross_ref=b_cross)[0])
-    assert np.isclose(b_cross, find_detuning_coeffs_threshold(tune_shift=tune_shift, q_s=q_s, b_direct_ref=b_direct,
-                                                              b_cross_ref=b_cross)[1])
-    # test negative octupole polarity
-    b_direct = -1.2718161244917965e-05
-    b_cross = 9.08066253298482e-06
-
-    assert np.isclose(b_direct, find_detuning_coeffs_threshold(tune_shift=tune_shift, q_s=q_s, b_direct_ref=b_direct,
-                                                               b_cross_ref=b_cross)[0])
-    assert np.isclose(b_cross, find_detuning_coeffs_threshold(tune_shift=tune_shift, q_s=q_s, b_direct_ref=b_direct,
-                                                              b_cross_ref=b_cross)[1])
-
-
-def test_find_detuning_coeffs_threshold_many_tune_shifts():
+
+    b_direct_test, b_cross_test = find_detuning_coeffs_threshold(tune_shift=tune_shift, q_s=q_s,
+                                                                 reference_b_direct=b_direct,
+                                                                 reference_b_cross=b_cross,
+                                                                 added_b_direct=added_b_direct,
+                                                                 added_b_cross=added_b_cross,
+                                                                 algorithm=algorithm)
+
+    assert np.isclose(b_direct - added_b_direct, b_direct_test)
+    assert np.isclose(b_cross - added_b_cross, b_cross_test)
+
+
+@mark.parametrize('algorithm',['newton','bisect'])
+@mark.parametrize('b_direct, b_cross',
+                  [[3.960630384598084e-05, -2.8278575218404585e-05],
+                   [-2.5436322491034757e-05, 1.816132506682559e-05],
+                   ])
+def test_find_detuning_coeffs_threshold_many_tune_shifts(algorithm: str, b_direct: float, b_cross: float):
     # reference values obtained with the old impedance wake model https://gitlab.cern.ch/IRIS/HLLHC_IW_model
     tune_shift = -7.3622423693e-05 - 3.0188372754e-06j
     q_s = 2e-3
-    tune_shifts = [np.nan, tune_shift, 2*tune_shift]
+    tune_shifts = [np.nan, tune_shift, 2 * tune_shift]
+
+    b_direct_test, b_cross_test = find_detuning_coeffs_threshold_many_tune_shifts(tune_shifts=tune_shifts, q_s=q_s,
+                                                                                  reference_b_direct=b_direct,
+                                                                                  reference_b_cross=b_cross,
+                                                                                  algorithm=algorithm)
+    assert np.isclose(b_direct, b_direct_test)
+    assert np.isclose(b_cross, b_cross_test)
+
 
-    # test positive octupole polarity
-    b_direct = 3.960630384598084e-05
-    b_cross = -2.8278575218404585e-05
-    assert np.isclose(b_direct, find_detuning_coeffs_threshold_many_tune_shifts(tune_shifts=tune_shifts, q_s=q_s,
-                                                                                b_direct_ref=b_direct,
-                                                                                b_cross_ref=b_cross)[0])
-    assert np.isclose(b_cross, find_detuning_coeffs_threshold_many_tune_shifts(tune_shifts=tune_shifts, q_s=q_s,
-                                                                               b_direct_ref=b_direct,
-                                                                               b_cross_ref=b_cross)[1])
-    # test negative octupole polarity
-    b_direct = -2.5436322491034757e-05
-    b_cross = 1.816132506682559e-05
-    assert np.isclose(b_direct, find_detuning_coeffs_threshold_many_tune_shifts(tune_shifts=tune_shifts, q_s=q_s,
-                                                                                b_direct_ref=b_direct,
-                                                                                b_cross_ref=b_cross)[0])
-    assert np.isclose(b_cross, find_detuning_coeffs_threshold_many_tune_shifts(tune_shifts=tune_shifts, q_s=q_s,
-                                                                               b_direct_ref=b_direct,
-                                                                               b_cross_ref=b_cross)[1])