diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 2fd220554163a01855cb1fa3cc1dff3cb7c6323b..c0dd36ec019dff2c0073084c893efa8a1b06608c 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -12,7 +12,7 @@ stages:
include:
- template: Jobs/Dependency-Scanning.gitlab-ci.yml
-testing:
+unit_test:
extends: .python_job
stage: test
rules:
@@ -22,11 +22,29 @@ testing:
- if: '$CI_MERGE_REQUEST_TARGET_BRANCH_NAME == "develop"'
script:
- cd tests
- - poetry run pytest --junitxml=report.xml
+ - poetry run pytest -m 'unit' --junitxml=unit_test.xml
artifacts:
when: always
reports:
- junit: /builds/PA/collective-effects/mbtrack2/tests/report.xml
+ junit: /builds/PA/collective-effects/mbtrack2/tests/unit_test.xml
+
+include:
+ - template: Jobs/Dependency-Scanning.gitlab-ci.yml
+phys_test:
+ extends: .python_job
+ stage: test
+ rules:
+ - if: '$CI_COMMIT_BRANCH == "stable"'
+ - if: '$CI_COMMIT_BRANCH == "develop"'
+ - if: '$CI_MERGE_REQUEST_TARGET_BRANCH_NAME == "stable"'
+ - if: '$CI_MERGE_REQUEST_TARGET_BRANCH_NAME == "develop"'
+ script:
+ - cd tests
+ - poetry run pytest -m 'phys' --junitxml=phys_test.xml
+ artifacts:
+ when: always
+ reports:
+ junit: /builds/PA/collective-effects/mbtrack2/tests/phys_test.xml
formatters:
extends: .python_job
diff --git a/mbtrack2/impedance/impedance_model.py b/mbtrack2/impedance/impedance_model.py
index 70729d4a8b1fb91381377942e5a1ce7a55d7d305..a15d6f31b0819b60ed550a816d7e0212a50d8520 100644
--- a/mbtrack2/impedance/impedance_model.py
+++ b/mbtrack2/impedance/impedance_model.py
@@ -39,7 +39,15 @@ class ImpedanceModel():
Parameters
----------
ring : Synchrotron object
-
+ average_beta : array-like of shape (2,), optional
+ Average beta function used for global wakes normalization in [m].
+ The global wakes are normalized by average_beta / local_beta.
+ If None and an AT lattice is loaded, average_beta is computed from
+ the lattice.
+ If None and an AT lattice is not loaded, average_beta is taken to be
+ equal to local_beta, i.e. no normalization.
+ The default is None.
+
Attributes
----------
wakefields : list of WakeField objects
@@ -65,7 +73,7 @@ class ImpedanceModel():
Add the same WakeField object at different locations to the model.
add_global(wakefield, name)
Add a "global" WakeField object which will added to the sum WakeField
- but not weighted by beta functions.
+ but weighted only by the average beta functions.
sum_beta(wake, beta)
Weight a WakeField object by an array of beta functions.
compute_sum_names()
@@ -80,7 +88,7 @@ class ImpedanceModel():
Load impedance model from file.
"""
- def __init__(self, ring):
+ def __init__(self, ring, average_beta=None):
self.ring = ring
self.optics = self.ring.optics
self.wakefields = []
@@ -90,6 +98,12 @@ class ImpedanceModel():
self.globals_names = []
self.sum_names = []
+ # global wake beta normalization
+ if average_beta is None:
+ average_beta = self.ring.optics.average_beta
+
+ self.average_beta = np.expand_dims(np.array(average_beta), 1)
+
def add(self, wakefield, positions, name=None):
"""
Add the same WakeField object at different locations to the model.
@@ -123,7 +137,7 @@ class ImpedanceModel():
def add_global(self, wakefield, name=None):
"""
Add a "global" WakeField object which will added to the sum WakeField
- but not weighted by beta functions.
+ but weighted only by the average beta functions.
To use with "distributed" elements, for example a resistive wall
wakefield computed from an effective radius (so which has already been
@@ -216,6 +230,7 @@ class ImpedanceModel():
except AttributeError:
setattr(self.sum, component_name2, comp2)
for i, wake2 in enumerate(self.globals):
+ wake2 = self.sum_beta(wake2, self.average_beta)
name = self.globals_names[i]
setattr(self, name, wake2)
self.sum_names.append(name)
@@ -783,7 +798,8 @@ class ImpedanceModel():
"positions": self.positions,
"names": self.names,
"globals": self.globals,
- "globals_names": self.globals_names
+ "globals_names": self.globals_names,
+ "average_beta": self.average_beta
}
with open(file, "wb") as f:
pickle.dump(to_save, f)
@@ -814,4 +830,5 @@ class ImpedanceModel():
self.names = to_load["names"]
self.globals = to_load["globals"]
self.globals_names = to_load["globals_names"]
+ self.average_beta = to_load["average_beta"]
self.compute_sum()
diff --git a/mbtrack2/tracking/wakepotential.py b/mbtrack2/tracking/wakepotential.py
index 47581c48bcd7a1cac3f4cb5347fbbb17b9ccb2a1..a95c5e2b5b821ecf1074745d27cc809dcd5164d7 100644
--- a/mbtrack2/tracking/wakepotential.py
+++ b/mbtrack2/tracking/wakepotential.py
@@ -643,6 +643,14 @@ class LongRangeResistiveWall(Element):
y3 : float, optional
Vertical effective radius of the 3rd power in [m], as Eq.27 in [1].
The default is radius.
+ average_beta : array-like of shape (2,), optional
+ Average beta function used for kick normalization in [m].
+ The transverse kick is normalized by average_beta / local_beta.
+ If None and an AT lattice is loaded, average_beta is computed from
+ the lattice.
+ If None and an AT lattice is not loaded, average_beta is taken to be
+ equal to local_beta, i.e. no normalization.
+ The default is None.
References
----------
@@ -659,7 +667,8 @@ class LongRangeResistiveWall(Element):
types=["Wlong", "Wxdip", "Wydip"],
nt=50,
x3=None,
- y3=None):
+ y3=None,
+ average_beta=None):
# parameters
self.ring = ring
self.length = length
@@ -682,6 +691,13 @@ class LongRangeResistiveWall(Element):
else:
self.y3 = radius
+ # kick beta normalization
+ if average_beta is None:
+ average_beta = self.ring.optics.average_beta
+
+ self.norm_x = average_beta[0] / self.ring.optics.local_beta[0]
+ self.norm_y = average_beta[1] / self.ring.optics.local_beta[1]
+
# constants
self.Z0 = mu_0 * c
self.t0 = (2 * self.rho * self.radius**2 / self.Z0)**(1 / 3) / c
@@ -852,13 +868,13 @@ class LongRangeResistiveWall(Element):
if wake_type == "Wlong":
bunch["delta"] -= kick / self.ring.E0
elif wake_type == "Wxdip":
- bunch["xp"] += kick / self.ring.E0
+ bunch["xp"] += kick / self.ring.E0 * self.norm_x
elif wake_type == "Wydip":
- bunch["yp"] += kick / self.ring.E0
+ bunch["yp"] += kick / self.ring.E0 * self.norm_y
elif wake_type == "Wxquad":
- bunch["xp"] += (bunch["x"] * kick / self.ring.E0)
+ bunch["xp"] += (bunch["x"] * kick / self.ring.E0) * self.norm_x
elif wake_type == "Wyquad":
- bunch["yp"] += (bunch["y"] * kick / self.ring.E0)
+ bunch["yp"] += (bunch["y"] * kick / self.ring.E0) * self.norm_y
def track(self, beam):
"""
diff --git a/mbtrack2/utilities/optics.py b/mbtrack2/utilities/optics.py
index 2aba39543025ca81d7131a3784015feebf190b3a..d37aafd6030bd57631bafc20c19650e0fbbcc741 100644
--- a/mbtrack2/utilities/optics.py
+++ b/mbtrack2/utilities/optics.py
@@ -18,6 +18,11 @@ class Optics:
lattice_file : str, optional if local_beta, local_alpha and
local_dispersion are specified.
AT lattice file path.
+ tracking_loc : float, optional
+ Longitudinal position where the tracking is in the AT lattice in [m].
+ Only used if an AT lattice is loaded to compute the optic functions at
+ this location.
+ Default is 0.
local_beta : array of shape (2,), optional if lattice_file is specified.
Beta function at the location of the tracking. Default is mean beta if
lattice has been loaded.
@@ -36,6 +41,8 @@ class Optics:
local_gamma : array of shape (2,)
Gamma function at the location of the tracking.
lattice : AT lattice
+ average_beta : array of shape (2,)
+ H and V average beta functions.
Methods
-------
@@ -57,6 +64,7 @@ class Optics:
def __init__(self,
lattice_file=None,
+ tracking_loc=0,
local_beta=None,
local_alpha=None,
local_dispersion=None,
@@ -64,23 +72,25 @@ class Optics:
if lattice_file is not None:
self.use_local_values = False
+ self.tracking_loc = tracking_loc
self.load_from_AT(lattice_file, **kwargs)
if local_beta is None:
- self._local_beta = np.mean(self.beta_array, axis=1)
+ self._local_beta = self.beta(self.tracking_loc)
else:
self._local_beta = local_beta
if local_alpha is None:
- self._local_alpha = np.mean(self.alpha_array, axis=1)
+ self._local_alpha = self.alpha(self.tracking_loc)
else:
self._local_alpha = local_alpha
if local_dispersion is None:
- self.local_dispersion = np.zeros((4, ))
+ self.local_dispersion = self.dispersion(self.tracking_loc)
else:
self.local_dispersion = local_dispersion
self._local_gamma = (1 + self._local_alpha**2) / self._local_beta
else:
self.use_local_values = True
+ self.tracking_loc = None
self._local_beta = local_beta
self._local_alpha = local_alpha
self._local_gamma = (1 + self._local_alpha**2) / self._local_beta
@@ -426,6 +436,30 @@ class Optics:
return fig
+ @property
+ def average_beta(self):
+ """
+ Return average beta functions.
+ If self.use_local_values, self.local_beta is returned.
+
+ Returns
+ -------
+ average_beta : array of shape (2,)
+ H and V average beta functions.
+
+ """
+ if self.use_local_values:
+ return self.local_beta
+ else:
+ L = self.position[-1]
+ position = np.linspace(0, L, int(self.n_points))
+ length = position[1:] - position[:-1]
+ center = (position[1:] + position[:-1]) / 2
+ beta = self.beta(center)
+ beta_H_star = 1 / L * (length * beta[0, :]).sum()
+ beta_V_star = 1 / L * (length * beta[1, :]).sum()
+ return np.array([beta_H_star, beta_V_star])
+
class PhysicalModel:
"""
@@ -541,6 +575,10 @@ class PhysicalModel:
Total length considered in [m].
rho_star : float
Effective resistivity of the chamber material in [ohm.m].
+ beta_H_star : float
+ Effective horizontal beta function in [m].
+ beta_V_star : float
+ Effective vertical beta function in [m].
a1_L : float
Effective longitudinal radius of the chamber of the first power in
[m].
diff --git a/pyproject.toml b/pyproject.toml
index f4f25eb39fe66332a95628ba57ae8c55b4f43ae0..463b6d2ff1e4a4595242f796cdeff211988cece4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,3 +37,9 @@ blank_line_before_nested_class_or_def = true
[tool.isort]
multi_line_output = 3
include_trailing_comma = true
+
+[tool.pytest.ini_options]
+markers = [
+ "unit: mark a test as a unit-test.",
+ "phys: mark a test as a physical-test.",
+]
\ No newline at end of file
diff --git a/tests/physics/test_ibs_phys.py b/tests/physics/test_ibs_phys.py
index a8a668989e2c836764af372eb0b1713786fe6abe..026369073ae1e3ae080108437e64ca8e01735cd0 100644
--- a/tests/physics/test_ibs_phys.py
+++ b/tests/physics/test_ibs_phys.py
@@ -1,6 +1,6 @@
import numpy as np
import pytest
-
+pytestmark = pytest.mark.phys
from mbtrack2 import Electron, Synchrotron
from mbtrack2.tracking import (
Bunch,
diff --git a/tests/physics/test_lrrw_phys.py b/tests/physics/test_lrrw_phys.py
new file mode 100644
index 0000000000000000000000000000000000000000..d4c92898a9b8efc189e9d5fd0405dd7d703ca1ca
--- /dev/null
+++ b/tests/physics/test_lrrw_phys.py
@@ -0,0 +1,114 @@
+import numpy as np
+from mbtrack2 import BeamMonitor, rwmbi_growth_rate
+from mbtrack2 import Beam, LongRangeResistiveWall, RFCavity, LongitudinalMap, TransverseMap
+from utility_test_growth_rate import get_growth_beam
+import pytest
+pytestmark = pytest.mark.phys
+
+class TestLongRangeResistiveWallPhys:
+
+ def test_llrw_growth_rate_chroma0(self, demo_ring, tmp_path):
+
+ ring = demo_ring
+ mp_num = int(1)
+ n_turns = int(5e3)
+ nt = 50
+ rho = 1.7e-8
+ radius = 4e-3
+ Itot = 1.0
+ Vc = 1.0e6
+ name = str(tmp_path / 'save')
+ ring.chro = [0,0]
+
+ growth_rate = rwmbi_growth_rate(ring, Itot, radius, rho, "y")
+
+ RF = RFCavity(ring, 1, Vc, np.arccos(ring.U0/Vc))
+ long = LongitudinalMap(ring)
+ trans = TransverseMap(ring)
+
+ bb = Beam(ring)
+ filling_pattern = np.ones((ring.h,))*Itot/ring.h
+
+ bb.init_beam(filling_pattern, mp_per_bunch=mp_num, track_alive=False, mpi=False)
+
+ rw = LongRangeResistiveWall(ring, beam=bb, length=ring.L, rho=rho,
+ radius=radius, types=["Wydip"], nt=nt)
+
+ save_every = 10
+ bbmon = BeamMonitor(h=ring.h, save_every=save_every, buffer_size=10,
+ total_size=int(n_turns/save_every),
+ file_name=name, mpi_mode=False)
+
+ for i in range(n_turns):
+ if (i % 1000 == 0):
+ print(i)
+
+ RF.track(bb)
+ long.track(bb)
+ trans.track(bb)
+
+ rw.track(bb)
+ bbmon.track(bb)
+
+ bbmon.close()
+
+ path = name + '.hdf5'
+ gr, gr_std = get_growth_beam(ring, path, 1e3, 5e3)
+ rel_error = np.abs(growth_rate-gr)*100/growth_rate
+ rel_error_std = np.abs(gr_std)*100/growth_rate
+ print(f"Relative error: {rel_error} % (std: {rel_error_std} %)")
+ assert rel_error < 2
+
+ # def test_llrw_growth_rate_chroma0_average_beta(self,
+ # generate_ring_with_at_lattice,
+ # tmp_path):
+ # tracking_loc = np.random.rand()*350
+ # ring = generate_ring_with_at_lattice(tracking_loc=tracking_loc)
+ # mp_num = int(1)
+ # n_turns = int(1e3)
+ # nt = 50
+ # rho=1.7e-8
+ # radius=4e-3
+ # Itot = 2
+ # Vc = 2.8e6
+ # name = str(tmp_path / 'save')
+ # ring.chro = [0,0]
+
+ # growth_rate = rwmbi_growth_rate(ring, Itot, radius, rho, "y")
+
+ # RF = RFCavity(ring, 1, Vc, np.arccos(ring.U0/Vc))
+ # long = LongitudinalMap(ring)
+ # trans = TransverseMap(ring)
+
+ # bb = Beam(ring)
+ # filling_pattern = np.ones((ring.h,))*Itot/ring.h
+
+ # bb.init_beam(filling_pattern, mp_per_bunch=mp_num, track_alive=False, mpi=False)
+
+ # rw = LongRangeResistiveWall(ring, beam=bb, length=ring.L, rho=rho,
+ # radius=radius, types=["Wydip"], nt=nt)
+
+ # save_every = 10
+ # bbmon = BeamMonitor(h=ring.h, save_every=save_every, buffer_size=10,
+ # total_size=int(n_turns/save_every),
+ # file_name=name, mpi_mode=False)
+
+ # for i in range(n_turns):
+ # if (i % 100 == 0):
+ # print(i)
+
+ # RF.track(bb)
+ # long.track(bb)
+ # trans.track(bb)
+
+ # rw.track(bb)
+ # bbmon.track(bb)
+
+ # bbmon.close()
+
+ # path = name + '.hdf5'
+ # gr, gr_std = get_growth_beam(ring, path, 500, 1e3)
+ # rel_error = np.abs(growth_rate-gr)*100/growth_rate
+ # rel_error_std = np.abs(gr_std)*100/growth_rate
+ # print(f"Relative error: {rel_error} % (std: {rel_error_std} %)")
+ # assert rel_error < 2
\ No newline at end of file
diff --git a/tests/test_imports.py b/tests/test_imports.py
index 11a9edfb665ae6c357eb2635d5e9602ce17bb7ac..b02158fcc70017ae00b4684412b5e6660369a5f1 100644
--- a/tests/test_imports.py
+++ b/tests/test_imports.py
@@ -1,5 +1,6 @@
+import pytest
+pytestmark = pytest.mark.unit
from mbtrack2 import *
-
def test_imports():
assert True
\ No newline at end of file
diff --git a/tests/unit/impedance/test_impedance_model.py b/tests/unit/impedance/test_impedance_model.py
index 02e07cf5e3401842cf666253e85086945f9d4cdc..8c593f2982b34551a22de73fd57f1a1b05e588d9 100644
--- a/tests/unit/impedance/test_impedance_model.py
+++ b/tests/unit/impedance/test_impedance_model.py
@@ -7,16 +7,17 @@ component_list = ComplexData.name_and_coefficients_table().columns.to_list()
import pytest
+pytestmark = pytest.mark.unit
class TestImpedanceModel:
# Adding WakeField objects with unique names should succeed without errors
@pytest.mark.parametrize('pos1,pos2', [([0],[1]), (0,1), (0.5,1.5)])
def test_add_wakefield(self,
- ring_with_at_lattice,
- generate_wakefield,
- pos1,
- pos2):
+ ring_with_at_lattice,
+ generate_wakefield,
+ pos1,
+ pos2):
model = ImpedanceModel(ring_with_at_lattice)
wakefield1 = generate_wakefield(name="wake1")
wakefield2 = generate_wakefield(name="wake2")
@@ -36,6 +37,32 @@ class TestImpedanceModel:
model.add_global(global_wakefield)
assert "global_wake" in model.globals_names
assert global_wakefield in model.globals
+
+ # Test weighted global sum with average_beta
+ def test_global_wakefield_weighted(self, ring_with_at_lattice, generate_wakefield):
+ ring_with_at_lattice.optics.local_beta = np.array([0.5,1])
+ model = ImpedanceModel(ring_with_at_lattice, average_beta=np.array([4, 2]))
+
+ global_wakefield = generate_wakefield(name = "global_wake",
+ wake_function_params=[{'component_type': "xdip",
+ 'function': np.array([1, 2])},
+ {'component_type': "ydip",
+ 'function': np.array([1, 2])}],
+ impedance_params=[{'component_type': "xdip",
+ 'function': np.array([1, 2])},
+ {'component_type': "ydip",
+ 'function': np.array([1, 2])}])
+
+ model.add_global(global_wakefield)
+ model.compute_sum()
+
+ excepted_x = np.array([4/0.5*1, 4/0.5*2])
+ excepted_y = np.array([2/1*1, 2/1*2])
+
+ np.testing.assert_allclose(model.sum.Wxdip.data["real"], excepted_x)
+ np.testing.assert_allclose(model.sum.Zxdip.data["real"], excepted_x)
+ np.testing.assert_allclose(model.sum.Wydip.data["real"], excepted_y)
+ np.testing.assert_allclose(model.sum.Zydip.data["real"], excepted_y)
# Test beta weighting summation
@pytest.mark.parametrize('comp,expected',
@@ -115,6 +142,7 @@ class TestImpedanceModel:
assert [0] in new_model.positions
assert "global_wake" in new_model.globals_names
assert isinstance(new_model.globals[0], WakeField)
+ assert new_model.average_beta.shape == (2,1)
# Plotting the contributions of WakeFields should generate a valid plot
@pytest.mark.parametrize('comp',[('long'),('xdip'),('ydip'),('xquad'),('yquad')])
diff --git a/tests/unit/impedance/test_wakefield.py b/tests/unit/impedance/test_wakefield.py
index 3e22aa05bae320bf71604b1b6d4392e70b676414..fbc7dd08aa98dd114594fea27e3fa8f26b0ec33c 100644
--- a/tests/unit/impedance/test_wakefield.py
+++ b/tests/unit/impedance/test_wakefield.py
@@ -1,5 +1,6 @@
import numpy as np
import pytest
+pytestmark = pytest.mark.unit
from mbtrack2 import ComplexData, WakeFunction, Impedance, WakeField
component_list = ComplexData.name_and_coefficients_table().columns.to_list()
diff --git a/tests/unit/tracking/test_aperture.py b/tests/unit/tracking/test_aperture.py
index df2a54ca9ce81d62348783ffb304f7aaabcdf23b..fd8a1a2589efcaeabaca2987d70a5726a5b757a2 100644
--- a/tests/unit/tracking/test_aperture.py
+++ b/tests/unit/tracking/test_aperture.py
@@ -1,4 +1,5 @@
import pytest
+pytestmark = pytest.mark.unit
from mbtrack2 import (CircularAperture,
ElipticalAperture,
RectangularAperture,
diff --git a/tests/unit/tracking/test_beam_ion_effects.py b/tests/unit/tracking/test_beam_ion_effects.py
index 3a23e0465702001cb605f96ec1ed029444223ca5..368d063b0ac15aeb93aba8a52e24ed574f59a5fa 100644
--- a/tests/unit/tracking/test_beam_ion_effects.py
+++ b/tests/unit/tracking/test_beam_ion_effects.py
@@ -4,6 +4,7 @@ import os
import h5py as hp
import numpy as np
import pytest
+pytestmark = pytest.mark.unit
@pytest.fixture
def generate_ion_particles(demo_ring):
diff --git a/tests/unit/tracking/test_element.py b/tests/unit/tracking/test_element.py
index 209c5d2e05495311f910a215a3beb774aeea22ed..6fcbe8d15a38406e45adebbacfed1e19725892d1 100644
--- a/tests/unit/tracking/test_element.py
+++ b/tests/unit/tracking/test_element.py
@@ -1,5 +1,6 @@
import numpy as np
import pytest
+pytestmark = pytest.mark.unit
from scipy.special import factorial
from utility_test_functions import assert_attr_changed
from mbtrack2 import (Element,
@@ -247,7 +248,7 @@ class TestTransverseMap:
----------
bunch : Bunch or Beam object
"""
-
+ print("x 0", bunch['x'][0])
# Compute phase advance which depends on energy via chromaticity and ADTS
if self.ring.adts is None:
phase_advance_x = (
@@ -256,6 +257,7 @@ class TestTransverseMap:
phase_advance_y = (
2 * np.pi *
(self.ring.tune[1] + self.ring.chro[1] * bunch["delta"]))
+ print("old phase advance ",phase_advance_x[0]/(2*np.pi), phase_advance_y[0]/(2*np.pi))
else:
Jx = ((self.ring.optics.local_gamma[0] * bunch["x"]**2) +
(2 * self.ring.optics.local_alpha[0] * bunch["x"] *
@@ -288,6 +290,8 @@ class TestTransverseMap:
matrix[1, 1, :] = c_x - self.alpha[0] * s_x
matrix[1, 2, :] = self.dispersion[1]
matrix[2, 2, :] = 1
+
+ print(matrix[0, 0, :], matrix[0, 1, :], matrix[0, 2, :])
# Vertical
c_y = np.cos(phase_advance_y)
@@ -338,9 +342,22 @@ class TestTransverseMap:
def test_trans_map_adts(self, ring_with_at_lattice, small_bunch):
ring_with_at_lattice.get_adts()
+
+ # Needed as old map was not handling dispersion properly!
+ ring_with_at_lattice.optics.local_dispersion = np.array([0,0,0,0])
+
old_map = self.Old_TransverseMap(ring_with_at_lattice)
current_map = TransverseMap(ring_with_at_lattice)
-
+
+ assert np.array_equal(old_map.gamma, current_map.gamma0)
+ assert np.array_equal(old_map.gamma, current_map.gamma1)
+ assert np.array_equal(old_map.beta, current_map.beta0)
+ assert np.array_equal(old_map.beta, current_map.beta1)
+ assert np.array_equal(old_map.alpha, current_map.alpha0)
+ assert np.array_equal(old_map.alpha, current_map.alpha1)
+ assert np.array_equal(old_map.dispersion, current_map.dispersion0)
+ assert np.array_equal(old_map.dispersion, current_map.dispersion1)
+
attrs = ["x", "xp", "y","yp"]
initial_values = {attr: small_bunch[attr].copy() for attr in attrs}
@@ -356,5 +373,4 @@ class TestTransverseMap:
for attr in attrs:
assert not np.array_equal(initial_values[attr], current[attr])
assert not np.array_equal(initial_values[attr], old[attr])
- assert np.allclose(current[attr], old[attr])
-
+ np.testing.assert_allclose(current[attr], old[attr])
diff --git a/tests/unit/tracking/test_emfields.py b/tests/unit/tracking/test_emfields.py
index e7eea3b21a62c7cecbe549d6a53ef5ce882607e3..40cbea56b12e8a09e47af0238d7d61c2926974f1 100644
--- a/tests/unit/tracking/test_emfields.py
+++ b/tests/unit/tracking/test_emfields.py
@@ -11,6 +11,7 @@ from mbtrack2.tracking.emfields import (
)
import pytest
+pytestmark = pytest.mark.unit
class Test_particles_electromagnetic_fields:
diff --git a/tests/unit/tracking/test_excite.py b/tests/unit/tracking/test_excite.py
index d4a8042d3b3c2684f53fc515e9293501e4dc50d9..f698a18ccbf1c41a59d19d6f1c26ad09b730cde8 100644
--- a/tests/unit/tracking/test_excite.py
+++ b/tests/unit/tracking/test_excite.py
@@ -1,5 +1,6 @@
import numpy as np
import pytest
+pytestmark = pytest.mark.unit
from mbtrack2 import Sweep
from utility_test_functions import assert_attr_changed
diff --git a/tests/unit/tracking/test_ibs.py b/tests/unit/tracking/test_ibs.py
index 974f1e868e2a77f2666f6a4a0ecd6b8ad8ad49a8..c119720053daf3d3344ea70d3a6e34a4cc80fd6f 100644
--- a/tests/unit/tracking/test_ibs.py
+++ b/tests/unit/tracking/test_ibs.py
@@ -1,6 +1,7 @@
from mbtrack2 import IntrabeamScattering
import numpy as np
import pytest
+pytestmark = pytest.mark.unit
from utility_test_functions import assert_attr_changed
@pytest.fixture
diff --git a/tests/unit/tracking/test_parallel.py b/tests/unit/tracking/test_parallel.py
index 3404b869b41392ed1a6d15bf08b96bacc45bdb81..dd75e633103bec871170206a77c91fafd7785f6e 100644
--- a/tests/unit/tracking/test_parallel.py
+++ b/tests/unit/tracking/test_parallel.py
@@ -1,4 +1,5 @@
import pytest
+pytestmark = pytest.mark.unit
import numpy as np
from mbtrack2.tracking.parallel import Mpi
from mpi4py import MPI
diff --git a/tests/unit/tracking/test_particle.py b/tests/unit/tracking/test_particle.py
index 7a1a51ff5807f97dcdeef264f9761f8c35c338a3..0d527f752eace2f74ff9a584f789a745e7615ff4 100644
--- a/tests/unit/tracking/test_particle.py
+++ b/tests/unit/tracking/test_particle.py
@@ -1,6 +1,7 @@
import numpy as np
import matplotlib.pyplot as plt
import pytest
+pytestmark = pytest.mark.unit
from scipy.constants import e, m_p, c
from mbtrack2 import Particle, Bunch, Beam
diff --git a/tests/unit/tracking/test_rf.py b/tests/unit/tracking/test_rf.py
index d41e51e429bb21d73a0d33e7f4dd3544f1436390..1f89e38c3be4c556a5b8ea3177c7ea5cabfbdbd2 100644
--- a/tests/unit/tracking/test_rf.py
+++ b/tests/unit/tracking/test_rf.py
@@ -1,4 +1,5 @@
import pytest
+pytestmark = pytest.mark.unit
import numpy as np
from utility_test_functions import assert_attr_changed
from mbtrack2 import (RFCavity,
diff --git a/tests/unit/tracking/test_spacecharge.py b/tests/unit/tracking/test_spacecharge.py
index b60a2156cbfb50b2a77e45b542fb8b4519b752ee..20d3073fad389887d1cab0e962287fe4511f01c9 100644
--- a/tests/unit/tracking/test_spacecharge.py
+++ b/tests/unit/tracking/test_spacecharge.py
@@ -1,4 +1,5 @@
import pytest
+pytestmark = pytest.mark.unit
from mbtrack2 import TransverseSpaceCharge
from utility_test_functions import assert_attr_changed
import numpy as np
diff --git a/tests/unit/tracking/test_synchrotron.py b/tests/unit/tracking/test_synchrotron.py
index dc59e4a9da53491b6c96664240f2c3c8e3f16bfe..fa0d42e7d1bb0071bb22ad752d0d510e8f839f82 100644
--- a/tests/unit/tracking/test_synchrotron.py
+++ b/tests/unit/tracking/test_synchrotron.py
@@ -1,5 +1,6 @@
import numpy as np
import pytest
+pytestmark = pytest.mark.unit
import at
from scipy.constants import c, e
from mbtrack2 import Electron, Synchrotron
@@ -31,15 +32,23 @@ def demo_ring_h1(demo_ring):
return demo_ring
@pytest.fixture
-def ring_with_at_lattice(at_optics):
- h = 416
- tau = np.array([6.56e-3, 6.56e-3, 3.27e-3])
- emit = np.array([3.9e-9, 3.9e-9*0.01])
- sigma_0 = 15e-12
- sigma_delta = 1.025e-3
- particle = Electron()
- ring = Synchrotron(h, at_optics, particle, tau=tau, emit=emit,
- sigma_0=sigma_0, sigma_delta=sigma_delta)
+def generate_ring_with_at_lattice(generate_at_optics):
+ def generate(**kwargs):
+ optics = generate_at_optics(**kwargs)
+ h = 416
+ tau = np.array([6.56e-3, 6.56e-3, 3.27e-3])
+ emit = np.array([3.9e-9, 3.9e-9*0.01])
+ sigma_0 = 15e-12
+ sigma_delta = 1.025e-3
+ particle = Electron()
+ ring = Synchrotron(h, optics, particle, tau=tau, emit=emit,
+ sigma_0=sigma_0, sigma_delta=sigma_delta, **kwargs)
+ return ring
+ return generate
+
+@pytest.fixture
+def ring_with_at_lattice(generate_ring_with_at_lattice):
+ ring = generate_ring_with_at_lattice()
return ring
class TestSynchrotron:
diff --git a/tests/unit/tracking/test_wakepotential.py b/tests/unit/tracking/test_wakepotential.py
index 8581839cc4aeae2e9b9e96286e9dc9133a62dfa5..403f711d3331af1234000bae1b6ed85d1787c102 100644
--- a/tests/unit/tracking/test_wakepotential.py
+++ b/tests/unit/tracking/test_wakepotential.py
@@ -1,4 +1,5 @@
import pytest
+pytestmark = pytest.mark.unit
import numpy as np
import pandas as pd
from mbtrack2 import WakePotential, Resonator, LongRangeResistiveWall
@@ -208,7 +209,8 @@ def generate_lrrw(demo_ring, beam_uniform):
types=["Wlong", "Wxdip", "Wydip"],
nt=50,
x3=None,
- y3=None):
+ y3=None,
+ average_beta=None):
lrrw = LongRangeResistiveWall(ring=ring,
beam=beam,
length=length,
@@ -217,7 +219,8 @@ def generate_lrrw(demo_ring, beam_uniform):
types=types,
nt=nt,
x3=x3,
- y3=y3)
+ y3=y3,
+ average_beta=average_beta)
return lrrw
return generate
@@ -313,4 +316,15 @@ class TestLongRangeResistiveWall:
if "Wxdip" in lrrw.types:
assert not np.array_equal(bunch["xp"], initial_xp), "Horizontal dipole kick not applied correctly"
if "Wydip" in lrrw.types:
- assert not np.array_equal(bunch["yp"], initial_yp), "Vertical dipole kick not applied correctly"
\ No newline at end of file
+ assert not np.array_equal(bunch["yp"], initial_yp), "Vertical dipole kick not applied correctly"
+
+ def test_average_beta_local(self, generate_lrrw):
+ lrrw = generate_lrrw()
+ assert lrrw.norm_x == 1
+ assert lrrw.norm_y == 1
+
+ def test_average_beta_mean(self, generate_lrrw, generate_ring_with_at_lattice):
+ ring = generate_ring_with_at_lattice(tracking_loc = 5)
+ lrrw = generate_lrrw(ring=ring)
+ assert lrrw.norm_x != 1
+ assert lrrw.norm_y != 1
\ No newline at end of file
diff --git a/tests/unit/utilities/test_optics.py b/tests/unit/utilities/test_optics.py
index c3d7fe0213258483822d91ea14e32a1877abbdce..c5a0eea554e1a048420628c27cf89fd6f2fa5564 100644
--- a/tests/unit/utilities/test_optics.py
+++ b/tests/unit/utilities/test_optics.py
@@ -1,5 +1,6 @@
import numpy as np
import pytest
+pytestmark = pytest.mark.unit
import at
import matplotlib.pyplot as plt
from pathlib import Path
@@ -15,11 +16,28 @@ def local_optics():
return local_optics
@pytest.fixture
-def at_optics():
- path_to_file = Path(__file__).parent
- lattice_file = path_to_file / ".." / ".." / ".." / "examples" / "SOLEIL_OLD.mat"
- at_optics = Optics(lattice_file=lattice_file)
- return at_optics
+def generate_at_optics():
+ def generate(lattice_file=None,
+ tracking_loc=0,
+ local_beta=None,
+ local_alpha=None,
+ local_dispersion=None,
+ **kwargs):
+ if lattice_file is None:
+ path_to_file = Path(__file__).parent
+ lattice_file = path_to_file / ".." / ".." / ".." / "examples" / "SOLEIL_OLD.mat"
+ optics = Optics(lattice_file=lattice_file,
+ tracking_loc=tracking_loc,
+ local_beta=local_beta,
+ local_alpha=local_alpha,
+ local_dispersion=local_dispersion,
+ **kwargs)
+ return optics
+ return generate
+
+@pytest.fixture
+def at_optics(generate_at_optics):
+ return generate_at_optics()
class TestOptics:
@@ -70,6 +88,15 @@ class TestOptics:
fig = local_optics.plot('beta', 'x')
assert fig is not None
+ @pytest.mark.parametrize('tracking_loc', [(0), (5), (100)])
+ def test_optics_tracking_loc(self, generate_at_optics, tracking_loc):
+ at_optics = generate_at_optics(tracking_loc=tracking_loc, n_points=1e4)
+
+ np.testing.assert_allclose(at_optics.beta(tracking_loc), at_optics.local_beta)
+ np.testing.assert_allclose(at_optics.alpha(tracking_loc), at_optics.local_alpha)
+ np.testing.assert_allclose(at_optics.gamma(tracking_loc), at_optics.local_gamma, rtol=1e-3) # high interpol. error
+ np.testing.assert_allclose(at_optics.dispersion(tracking_loc), at_optics.local_dispersion)
+
class TestPhysicalModel:
@pytest.fixture
diff --git a/tests/utility_test_growth_rate.py b/tests/utility_test_growth_rate.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd3fee16b676072b6417e6e8f986d1c799e04065
--- /dev/null
+++ b/tests/utility_test_growth_rate.py
@@ -0,0 +1,60 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import h5py as hp
+
+def r2(y, y_fit):
+ # residual sum of squares
+ ss_res = np.sum((y - y_fit) ** 2)
+
+ # total sum of squares
+ ss_tot = np.sum((y - np.mean(y)) ** 2)
+
+ # r-squared
+ r2 = 1 - (ss_res / ss_tot)
+ return r2
+
+def get_growth(ring, start, end, bunch_num, g, plane="y", plot=False):
+ plane_dict = {"x":0, "y":1, "long":2}
+ time = np.array(g["Beam"]["time"])*ring.T0
+ data = np.array(g["Beam"]["cs_invariant"][plane_dict[plane],bunch_num,:])
+ idx = (time > ring.T0*start) & (time < ring.T0*end)
+ time_fit = time[idx] - start*ring.T0
+ data_fit = data[idx]
+
+ ploy = np.polynomial.polynomial.Polynomial.fit(time_fit, np.log(data_fit), 1, w=np.sqrt(data_fit))
+
+ if plot:
+ plt.plot(time/ring.T0, np.log(data), '-')
+
+ arr = ploy.convert().coef
+ fit_vals = [np.exp(arr[0]) * np.exp(arr[1] *x) for x in time_fit]
+
+ return (arr[0], arr[1]/2, r2(data_fit, fit_vals))
+
+def get_growth_beam(ring, path, start, end, plane="y", plot=False):
+ g = hp.File(path,"r")
+ gr_cs = np.zeros(ring.h)
+ cste = np.zeros(ring.h)
+ r2_cs = np.zeros(ring.h)
+
+ for i in range(ring.h):
+ cste[i], gr_cs[i], r2_cs[i] = get_growth(ring, start, end, i, g, plane, plot)
+
+ time = np.array(g["Beam"]["time"])
+ idx = (time > start) & (time < end)
+ time_fit = time[idx]
+ fit = np.mean(cste) + np.mean(gr_cs)*2*(time_fit-start)*ring.T0
+
+ print(f"growth rate = {np.mean(gr_cs)} +- {np.std(gr_cs)} s-1")
+ print(f"r2 = {np.mean(r2_cs)} +- {np.std(r2_cs)}")
+ if plot:
+ plt.plot(time_fit, fit, 'k--')
+ ymax = plt.gca().get_ylim()[1]
+ ymin = plt.gca().get_ylim()[0]
+ plt.vlines(start, ymin, ymax, "r")
+ plt.vlines(end, ymin, ymax, "r")
+ plt.grid()
+ plt.xlabel("Turn number")
+ plt.ylabel("log(CS_invariant)")
+ g.close()
+ return (np.mean(gr_cs), np.std(gr_cs))
\ No newline at end of file