diff --git a/tests/test_landau_damping.py b/tests/test_landau_damping.py index b4472307..1a3f1ee6 100644 --- a/tests/test_landau_damping.py +++ b/tests/test_landau_damping.py @@ -4,88 +4,81 @@ # ######################################### # 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]) diff --git a/xwakes/wit/landau_damping.py b/xwakes/wit/landau_damping.py index 104aedc1..49baa83b 100644 --- a/xwakes/wit/landau_damping.py +++ b/xwakes/wit/landau_damping.py @@ -5,7 +5,7 @@ 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 @@ -64,26 +64,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 @@ -93,23 +99,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: @@ -128,33 +146,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)