From ac8a21bf86384ac2c98ee295d05b07617e32a660 Mon Sep 17 00:00:00 2001
From: Alexis Gamelin <alexis.gamelin@synchrotron-soleil.fr>
Date: Mon, 25 Mar 2024 17:41:17 +0100
Subject: [PATCH] Bump version to 0.6.0 and apply formatters
---
LICENSE | 2 +-
mbtrack2/__init__.py | 2 +-
mbtrack2/impedance/csr.py | 2 +
mbtrack2/impedance/impedance_model.py | 1 +
mbtrack2/impedance/resistive_wall.py | 2 +
mbtrack2/impedance/resonator.py | 13 +++--
mbtrack2/impedance/tapers.py | 5 ++
mbtrack2/impedance/wakefield.py | 4 ++
mbtrack2/instability/instabilities.py | 60 +++++++++++-----------
mbtrack2/tracking/aperture.py | 4 ++
mbtrack2/tracking/element.py | 7 +++
mbtrack2/tracking/feedback.py | 2 +
mbtrack2/tracking/monitors/monitors.py | 18 +++++--
mbtrack2/tracking/parallel.py | 1 +
mbtrack2/tracking/particles.py | 5 ++
mbtrack2/tracking/rf.py | 70 +++++++++++++++-----------
mbtrack2/tracking/synchrotron.py | 32 ++++++------
mbtrack2/tracking/wakepotential.py | 12 +++--
mbtrack2/utilities/beamloading.py | 5 +-
mbtrack2/utilities/optics.py | 2 +
pyproject.toml | 4 +-
tests/conftest.py | 6 ++-
tests/test_bunch.py | 4 +-
tests/test_optics.py | 4 +-
tests/test_synchrotron.py | 6 ++-
25 files changed, 173 insertions(+), 100 deletions(-)
diff --git a/LICENSE b/LICENSE
index 8c09e5d..b184fd1 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,6 +1,6 @@
BSD 3-Clause License
-Copyright (c) 2021-2023, SOLEIL Synchrotron
+Copyright (c) 2021-2024, SOLEIL Synchrotron
All rights reserved.
Redistribution and use in source and binary forms, with or without
diff --git a/mbtrack2/__init__.py b/mbtrack2/__init__.py
index 190f0ac..d3383ef 100644
--- a/mbtrack2/__init__.py
+++ b/mbtrack2/__init__.py
@@ -1,5 +1,5 @@
# -*- coding: utf-8 -*-
-__version__ = "0.5.0"
+__version__ = "0.6.0"
from mbtrack2.impedance import *
from mbtrack2.instability import *
from mbtrack2.tracking import *
diff --git a/mbtrack2/impedance/csr.py b/mbtrack2/impedance/csr.py
index 592a4da..b196110 100644
--- a/mbtrack2/impedance/csr.py
+++ b/mbtrack2/impedance/csr.py
@@ -38,6 +38,7 @@ class FreeSpaceCSR(WakeField):
Beams 7.5 (2004): 054403.
"""
+
def __init__(self, time, frequency, length, radius):
super().__init__()
@@ -119,6 +120,7 @@ class ParallelPlatesCSR(WakeField):
Beams 7.5 (2004): 054403.
"""
+
def __init__(self, time, frequency, length, radius, distance):
super().__init__()
diff --git a/mbtrack2/impedance/impedance_model.py b/mbtrack2/impedance/impedance_model.py
index 96efc7f..713fe46 100644
--- a/mbtrack2/impedance/impedance_model.py
+++ b/mbtrack2/impedance/impedance_model.py
@@ -79,6 +79,7 @@ class ImpedanceModel():
load(file)
Load impedance model from file.
"""
+
def __init__(self, ring):
self.ring = ring
self.optics = self.ring.optics
diff --git a/mbtrack2/impedance/resistive_wall.py b/mbtrack2/impedance/resistive_wall.py
index 1676ba0..dd99bf1 100644
--- a/mbtrack2/impedance/resistive_wall.py
+++ b/mbtrack2/impedance/resistive_wall.py
@@ -77,6 +77,7 @@ class CircularResistiveWall(WakeField):
Detectors and Associated Equipment 806 (2016): 221-230.
"""
+
def __init__(self,
time,
frequency,
@@ -252,6 +253,7 @@ class CircularResistiveWall(WakeField):
class Coating(WakeField):
+
def __init__(self,
frequency,
length,
diff --git a/mbtrack2/impedance/resonator.py b/mbtrack2/impedance/resonator.py
index d8dc346..b89a818 100644
--- a/mbtrack2/impedance/resonator.py
+++ b/mbtrack2/impedance/resonator.py
@@ -10,6 +10,7 @@ from mbtrack2.impedance.wakefield import Impedance, WakeField, WakeFunction
class Resonator(WakeField):
+
def __init__(self, time, frequency, Rs, fr, Q, plane, atol=1e-20):
"""
Resonator model WakeField element which computes the impedance and the
@@ -111,14 +112,14 @@ class Resonator(WakeField):
def transverse_wake_function(self, t):
if self.Q >= 0.5:
wt = (self.wr * self.Rs / self.Q_p *
- np.exp(-1 * t * self.wr / 2 / self.Q_p) *
- np.sin(self.wr_p * t))
+ np.exp(-1 * t * self.wr / 2 / self.Q_p) *
+ np.sin(self.wr_p * t))
else:
wt = (self.wr * self.Rs / self.Q_p *
- np.exp(-1 * t * self.wr / 2 / self.Q_p) *
- np.sinh(self.wr_p * t))
+ np.exp(-1 * t * self.wr / 2 / self.Q_p) *
+ np.sinh(self.wr_p * t))
if np.any(t < 0):
- wt[t<0] = 0
+ wt[t < 0] = 0
return wt
@@ -139,6 +140,7 @@ class PureInductive(WakeField):
Maximum frequency used in the impedance.
nout, trim : see Impedance.to_wakefunction
"""
+
def __init__(self,
L,
n_wake=1e6,
@@ -181,6 +183,7 @@ class PureResistive(WakeField):
Maximum frequency used in the impedance.
nout, trim : see Impedance.to_wakefunction
"""
+
def __init__(self,
R,
n_wake=1e6,
diff --git a/mbtrack2/impedance/tapers.py b/mbtrack2/impedance/tapers.py
index 088fa4b..411248d 100644
--- a/mbtrack2/impedance/tapers.py
+++ b/mbtrack2/impedance/tapers.py
@@ -25,6 +25,7 @@ class StupakovRectangularTaper(WakeField):
length: taper length in [m]
width : full horizontal width of the taper in [m]
"""
+
def __init__(self,
frequency,
gap_entrance,
@@ -101,6 +102,7 @@ class StupakovRectangularTaper(WakeField):
return np.imag(self.long(1)) * f0
def ydip(self):
+
def G1(x, m_max):
m = np.arange(0, m_max)
phi = np.outer(pi * x / 2, 2*m + 1)
@@ -118,6 +120,7 @@ class StupakovRectangularTaper(WakeField):
return -1j * pi * self.width * self.Z0 / 4 * integral
def xdip(self):
+
def G3(x, m_max):
m = np.arange(0, m_max)
phi = np.outer(pi * x, m)
@@ -135,6 +138,7 @@ class StupakovRectangularTaper(WakeField):
return -1j * pi * self.Z0 / 4 * integral
def quad(self):
+
def G2(x, m_max):
m = np.arange(0, m_max)
phi = np.outer(pi * x / 2, 2*m + 1)
@@ -166,6 +170,7 @@ class StupakovCircularTaper(WakeField):
radius_exit : radius at taper exit in [m]
length : taper length in [m]
"""
+
def __init__(self,
frequency,
radius_entrance,
diff --git a/mbtrack2/impedance/wakefield.py b/mbtrack2/impedance/wakefield.py
index 6df7cfc..6f221f2 100644
--- a/mbtrack2/impedance/wakefield.py
+++ b/mbtrack2/impedance/wakefield.py
@@ -29,6 +29,7 @@ class ComplexData:
function : list or numpy array of comp lex numbers
contains the values taken by the complex function
"""
+
def __init__(self,
variable=np.array([-1e15, 1e15]),
function=np.array([0, 0])):
@@ -363,6 +364,7 @@ class WakeFunction(ComplexData):
loss_factor(sigma)
Compute the loss factor or the kick factor for a Gaussian bunch.
"""
+
def __init__(self,
variable=np.array([-1e15, 1e15]),
function=np.array([0, 0]),
@@ -599,6 +601,7 @@ class Impedance(ComplexData):
plot()
Plot the impedance data.
"""
+
def __init__(self,
variable=np.array([-1e15, 1e15]),
function=np.array([0, 0]),
@@ -817,6 +820,7 @@ class WakeField:
load(file)
Load WakeField element from file.
"""
+
def __init__(self, structure_list=None, name=None):
self.list_to_attr(structure_list)
self.name = name
diff --git a/mbtrack2/instability/instabilities.py b/mbtrack2/instability/instabilities.py
index 75988eb..692c465 100644
--- a/mbtrack2/instability/instabilities.py
+++ b/mbtrack2/instability/instabilities.py
@@ -194,14 +194,14 @@ def lcbi_growth_rate_mode(ring,
return factor * sum_val
-def lcbi_growth_rate(ring,
- I,
+def lcbi_growth_rate(ring,
+ I,
M,
synchrotron_tune=None,
Vrf=None,
- fr=None,
- RL=None,
- QL=None,
+ fr=None,
+ RL=None,
+ QL=None,
Z=None):
"""
Compute the maximum growth rate for longitudinal coupled bunch instability
@@ -249,16 +249,17 @@ def lcbi_growth_rate(ring,
"""
growth_rates = np.zeros(M)
for i in range(M):
- growth_rates[i] = lcbi_growth_rate_mode(ring,
- I,
- M,
- i,
- synchrotron_tune=synchrotron_tune,
- Vrf=Vrf,
- fr=fr,
- RL=RL,
- QL=QL,
- Z=Z)
+ growth_rates[i] = lcbi_growth_rate_mode(
+ ring,
+ I,
+ M,
+ i,
+ synchrotron_tune=synchrotron_tune,
+ Vrf=Vrf,
+ fr=fr,
+ RL=RL,
+ QL=QL,
+ Z=Z)
growth_rate = np.max(growth_rates)
mu = np.argmax(growth_rates)
@@ -266,11 +267,11 @@ def lcbi_growth_rate(ring,
return growth_rate, mu, growth_rates
-def lcbi_stability_diagram(ring,
- I,
- M,
- modes,
- cavity_list,
+def lcbi_stability_diagram(ring,
+ I,
+ M,
+ modes,
+ cavity_list,
detune_range,
synchrotron_tune=None,
Vrf=None):
@@ -317,15 +318,16 @@ def lcbi_stability_diagram(ring,
for mu in modes:
fixed_gr = 0
for cav in cavity_list[:-1]:
- fixed_gr += lcbi_growth_rate_mode(ring,
- I=I,
- mu=mu,
- fr=cav.fr,
- RL=cav.RL,
- QL=cav.QL,
- M=M,
- synchrotron_tune=synchrotron_tune,
- Vrf=Vrf)
+ fixed_gr += lcbi_growth_rate_mode(
+ ring,
+ I=I,
+ mu=mu,
+ fr=cav.fr,
+ RL=cav.RL,
+ QL=cav.QL,
+ M=M,
+ synchrotron_tune=synchrotron_tune,
+ Vrf=Vrf)
cav = cavity_list[-1]
for i, det in enumerate(detune_range):
diff --git a/mbtrack2/tracking/aperture.py b/mbtrack2/tracking/aperture.py
index 7f3645c..c9316d1 100644
--- a/mbtrack2/tracking/aperture.py
+++ b/mbtrack2/tracking/aperture.py
@@ -18,6 +18,7 @@ class CircularAperture(Element):
radius : float
radius of the circle in [m]
"""
+
def __init__(self, radius):
self.radius = radius
self.radius_squared = radius**2
@@ -50,6 +51,7 @@ class ElipticalAperture(Element):
Y_radius : float
vertical radius of the elipse in [m]
"""
+
def __init__(self, X_radius, Y_radius):
self.X_radius = X_radius
self.X_radius_squared = X_radius**2
@@ -88,6 +90,7 @@ class RectangularAperture(Element):
Y_bottom : float, optional
bottom vertical aperture of the rectangle in [m]
"""
+
def __init__(self, X_right, Y_top, X_left=None, Y_bottom=None):
self.X_right = X_right
self.X_left = X_left
@@ -135,6 +138,7 @@ class LongitudinalAperture(Element):
tau_low : float, optional
Lower longitudinal bound in [s].
"""
+
def __init__(self, tau_up, tau_low=None):
self.tau_up = tau_up
if tau_low is None:
diff --git a/mbtrack2/tracking/element.py b/mbtrack2/tracking/element.py
index 7fab92b..47cc401 100644
--- a/mbtrack2/tracking/element.py
+++ b/mbtrack2/tracking/element.py
@@ -19,6 +19,7 @@ class Element(metaclass=ABCMeta):
Abstract Element class used for subclass inheritance to define all kinds
of objects which intervene in the tracking.
"""
+
@abstractmethod
def track(self, beam):
"""
@@ -53,6 +54,7 @@ class Element(metaclass=ABCMeta):
track method of an Element subclass which takes a Beam object or a
Bunch object as input
"""
+
@wraps(track)
def track_wrapper(*args, **kwargs):
if isinstance(args[1], Beam):
@@ -79,6 +81,7 @@ class LongitudinalMap(Element):
----------
ring : Synchrotron object
"""
+
def __init__(self, ring):
self.ring = ring
@@ -109,6 +112,7 @@ class SynchrotronRadiation(Element):
switch : bool array of shape (3,), optional
allow to choose on which plane the synchrotron radiation is active
"""
+
def __init__(self, ring, switch=np.ones((3, ), dtype=bool)):
self.ring = ring
self.switch = switch
@@ -154,6 +158,7 @@ class TransverseMap(Element):
----------
ring : Synchrotron object
"""
+
def __init__(self, ring):
self.ring = ring
self.alpha = self.ring.optics.local_alpha
@@ -253,6 +258,7 @@ class SkewQuadrupole:
Integrated strength of the skew quadrupole [m].
"""
+
def __init__(self, strength):
self.strength = strength
@@ -304,6 +310,7 @@ class TransverseMapSector(Element):
for details. The default is None.
"""
+
def __init__(self,
ring,
alpha0,
diff --git a/mbtrack2/tracking/feedback.py b/mbtrack2/tracking/feedback.py
index 77bef8a..351402c 100644
--- a/mbtrack2/tracking/feedback.py
+++ b/mbtrack2/tracking/feedback.py
@@ -26,6 +26,7 @@ class ExponentialDamper(Element):
Phase difference between the damper and the position monitor in [rad].
"""
+
def __init__(self, ring, plane, damping_time, phase_diff):
self.ring = ring
self.damping_time = damping_time
@@ -117,6 +118,7 @@ class FIRDamper(Element):
2004. Transverse bunch by bunch feedback system for the Spring-8
storage ring.
"""
+
def __init__(self,
ring,
plane,
diff --git a/mbtrack2/tracking/monitors/monitors.py b/mbtrack2/tracking/monitors/monitors.py
index 01794eb..151fe2a 100644
--- a/mbtrack2/tracking/monitors/monitors.py
+++ b/mbtrack2/tracking/monitors/monitors.py
@@ -302,6 +302,7 @@ class BunchMonitor(Monitor):
track(object_to_save)
Save data
"""
+
def __init__(self,
bunch_number,
save_every,
@@ -378,6 +379,7 @@ class PhaseSpaceMonitor(Monitor):
track(object_to_save)
Save data
"""
+
def __init__(self,
bunch_number,
mp_number,
@@ -476,6 +478,7 @@ class BeamMonitor(Monitor):
track(beam)
Save data
"""
+
def __init__(self,
h,
save_every,
@@ -682,6 +685,7 @@ class ProfileMonitor(Monitor):
track(object_to_save)
Save data.
"""
+
def __init__(self,
bunch_number,
save_every,
@@ -813,6 +817,7 @@ class WakePotentialMonitor(Monitor):
track(object_to_save, wake_potential_to_save)
Save data.
"""
+
def __init__(self,
bunch_number,
wake_types,
@@ -1023,6 +1028,7 @@ class BunchSpectrumMonitor(Monitor):
Save spectrum data.
"""
+
def __init__(self,
ring,
bunch_number,
@@ -1332,6 +1338,7 @@ class BeamSpectrumMonitor(Monitor):
Save spectrum data.
"""
+
def __init__(self,
ring,
save_every,
@@ -1537,6 +1544,7 @@ class CavityMonitor(Monitor):
track(beam, cavity)
Save data
"""
+
def __init__(self,
cavity_name,
ring,
@@ -1567,8 +1575,8 @@ class CavityMonitor(Monitor):
ring.h,
buffer_size,
),
- "DFB_ig_phasor":(
- ring.h,
+ "DFB_ig_phasor": (
+ ring.h,
buffer_size,
),
"detune": (buffer_size, ),
@@ -1599,8 +1607,8 @@ class CavityMonitor(Monitor):
ring.h,
total_size,
),
- "DFB_ig_phasor":(
- ring.h,
+ "DFB_ig_phasor": (
+ ring.h,
total_size,
),
"detune": (total_size, ),
@@ -1619,7 +1627,7 @@ class CavityMonitor(Monitor):
"beam_phasor_record": complex,
"generator_phasor_record": complex,
"ig_phasor_record": complex,
- "DFB_ig_phasor":complex,
+ "DFB_ig_phasor": complex,
"detune": float,
"psi": float,
"Vg": float,
diff --git a/mbtrack2/tracking/parallel.py b/mbtrack2/tracking/parallel.py
index 58c45b1..083bfdb 100644
--- a/mbtrack2/tracking/parallel.py
+++ b/mbtrack2/tracking/parallel.py
@@ -58,6 +58,7 @@ class Mpi:
Computing using Python, Advances in Water Resources, 34(9):1124-1139, 2011.
"""
+
def __init__(self, filling_pattern):
from mpi4py import MPI
self.MPI = MPI
diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py
index f334be8..7ba97ff 100644
--- a/mbtrack2/tracking/particles.py
+++ b/mbtrack2/tracking/particles.py
@@ -24,6 +24,7 @@ class Particle:
E_rest : float
particle rest energy in [eV]
"""
+
def __init__(self, mass, charge):
self.mass = mass
self.charge = charge
@@ -35,12 +36,14 @@ class Particle:
class Electron(Particle):
""" Define an electron"""
+
def __init__(self):
super().__init__(m_e, -1 * e)
class Proton(Particle):
""" Define a proton"""
+
def __init__(self):
super().__init__(m_p, e)
@@ -113,6 +116,7 @@ class Bunch:
[1] Wiedemann, H. (2015). Particle accelerator physics. 4th edition.
Springer, Eq.(8.39) of p224.
"""
+
def __init__(self,
ring,
mp_number=1e3,
@@ -586,6 +590,7 @@ class Beam:
load(file_name, mpi)
Load data from a HDF5 file recorded by Beam save method.
"""
+
def __init__(self, ring, bunch_list=None):
self.ring = ring
self.mpi_switch = False
diff --git a/mbtrack2/tracking/rf.py b/mbtrack2/tracking/rf.py
index c01e63c..a92c4b3 100644
--- a/mbtrack2/tracking/rf.py
+++ b/mbtrack2/tracking/rf.py
@@ -7,8 +7,9 @@ import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.legend_handler import HandlerPatch
-from mbtrack2.tracking.element import Element
+
from mbtrack2.instability import lcbi_growth_rate
+from mbtrack2.tracking.element import Element
class RFCavity(Element):
@@ -26,6 +27,7 @@ class RFCavity(Element):
theta : float
Phase of Cavity voltage
"""
+
def __init__(self, ring, m, Vc, theta):
self.ring = ring
self.m = m
@@ -200,6 +202,7 @@ class CavityResonator():
of Longitudinal Beam Dynamics With Harmonic Cavities by Using the Code
Mbtrack." IPAC’19, Melbourne, Australia, 2019.
"""
+
def __init__(self,
ring,
m,
@@ -579,7 +582,7 @@ class CavityResonator():
if isinstance(FB, (ProportionalIntegralLoop, DirectFeedback)):
return FB.ig_phasor_record
return np.zeros(self.ring.h)
-
+
@property
def DFB_ig_phasor(self):
"""Last direct feedback current generator phasor of each bunch in [A]"""
@@ -887,6 +890,7 @@ class CavityResonator():
Figure.
"""
+
def make_legend_arrow(legend, orig_handle, xdescent, ydescent, width,
height, fontsize):
p = mpatches.FancyArrow(0,
@@ -938,7 +942,7 @@ class CavityResonator():
return fig
- def is_CBI_stable(self,
+ def is_CBI_stable(self,
I0,
synchrotron_tune=None,
modes=None,
@@ -981,27 +985,28 @@ class CavityResonator():
instability growth rate.
"""
- growth_rate, mu, growth_rates = lcbi_growth_rate(self.ring,
- I0,
- M=self.ring.h,
- synchrotron_tune=synchrotron_tune,
- Vrf=self.Vc,
- fr=self.fr,
- RL=self.RL,
- QL=self.QL)
-
+ growth_rate, mu, growth_rates = lcbi_growth_rate(
+ self.ring,
+ I0,
+ M=self.ring.h,
+ synchrotron_tune=synchrotron_tune,
+ Vrf=self.Vc,
+ fr=self.fr,
+ RL=self.RL,
+ QL=self.QL)
+
if modes is not None:
- growth_rates = growth_rates[modes]
+ growth_rates = growth_rates[modes]
return growth_rates
-
+
if bool_return:
- if growth_rate > 1/self.ring.tau[2]:
+ if growth_rate > 1 / self.ring.tau[2]:
return True
else:
return False
else:
return growth_rate, mu
-
+
def is_DC_Robinson_stable(self, I0):
"""
Check DC Robinson stability - Eq. (6.1.1) [1]
@@ -1161,6 +1166,7 @@ class ProportionalLoop():
Must be supperior or equal to 1.
"""
+
def __init__(self, ring, cav_res, gain_A, gain_P, delay):
self.ring = ring
self.cav_res = cav_res
@@ -1222,6 +1228,7 @@ class TunerLoop():
Tuning offset in [rad].
"""
+
def __init__(self, ring, cav_res, gain=0.01, avering_period=0, offset=0):
self.ring = ring
self.cav_res = cav_res
@@ -1366,14 +1373,15 @@ class ProportionalIntegralLoop():
of synchrotron light sources. PRAB, 21(1), 012001.
"""
- def __init__(self,
- ring,
- cav_res,
- gain,
- sample_num,
- every,
- delay,
- IIR_cutoff=0,
+
+ def __init__(self,
+ ring,
+ cav_res,
+ gain,
+ sample_num,
+ every,
+ delay,
+ IIR_cutoff=0,
FF=True):
self.ring = ring
self.cav_res = cav_res
@@ -1521,8 +1529,8 @@ class ProportionalIntegralLoop():
else:
omega = 2.0 * np.pi * cutoff
T = self.ring.T1 * self.every
- alpha = np.cos(omega*T)-1
- tmp = alpha * alpha - 2 * alpha
+ alpha = np.cos(omega * T) - 1
+ tmp = alpha*alpha - 2*alpha
if tmp > 0:
self.IIRcoef = alpha + np.sqrt(tmp)
else:
@@ -1531,14 +1539,17 @@ class ProportionalIntegralLoop():
def IIR(self, input):
"""Return IIR filter output."""
- self.IIRout = (1 - self.IIRcoef)*self.IIRout + self.IIRcoef*input
+ self.IIRout = (1 - self.IIRcoef) * self.IIRout + self.IIRcoef * input
return self.IIRout
-
+
@property
def IIRcutoff(self):
"""Return IIR cutoff frequency in [Hz]."""
T = self.ring.T1 * self.every
- return 1.0/2.0/np.pi/T * np.arccos((2-2*self.IIRcoef-self.IIRcoef*self.IIRcoef)/2/(1-self.IIRcoef))
+ return 1.0 / 2.0 / np.pi / T * np.arccos(
+ (2 - 2 * self.IIRcoef - self.IIRcoef * self.IIRcoef) / 2 /
+ (1 - self.IIRcoef))
+
class DirectFeedback(ProportionalIntegralLoop):
"""
@@ -1612,6 +1623,7 @@ class DirectFeedback(ProportionalIntegralLoop):
ring. In Proc. IPAC'23. doi:10.18429/JACoW-IPAC2023-WEPL161
"""
+
def __init__(self,
DFB_gain,
DFB_phase_shift,
diff --git a/mbtrack2/tracking/synchrotron.py b/mbtrack2/tracking/synchrotron.py
index c519d50..b6e36fa 100644
--- a/mbtrack2/tracking/synchrotron.py
+++ b/mbtrack2/tracking/synchrotron.py
@@ -117,6 +117,7 @@ class Synchrotron:
to_pyat(Vrf)
Return a pyAT simple_ring element from the Synchrotron element data.
"""
+
def __init__(self, h, optics, particle, **kwargs):
self._h = h
self.particle = particle
@@ -508,7 +509,7 @@ class Synchrotron:
self.long_gamma = long_gamma
else:
return tuneS, long_alpha, long_beta, long_gamma
-
+
def to_pyat(self, Vrf, harmonic_number=None, TimeLag=False):
"""
Return a pyAT simple_ring element from the Synchrotron element data.
@@ -534,29 +535,30 @@ class Synchrotron:
"""
from at import simple_ring
optics = self.optics
-
+
if (harmonic_number is None) and isinstance(Vrf, (float, int)):
- harmonic_number=self.h
-
+ harmonic_number = self.h
+
if isinstance(Vrf, (list, np.ndarray)):
if (harmonic_number is None) or (TimeLag is None):
- raise ValueError("If sevral cavities are provided, "
- "harmonic_number and TimeLag should be provided.")
-
+ raise ValueError(
+ "If sevral cavities are provided, "
+ "harmonic_number and TimeLag should be provided.")
+
if self.adts is not None:
try:
- A1 = self.adts[0][-2]*2
- A2 = self.adts[1][-2]*2
- A3 = self.adts[3][-2]*2
+ A1 = self.adts[0][-2] * 2
+ A2 = self.adts[1][-2] * 2
+ A3 = self.adts[3][-2] * 2
except IndexError:
A1 = None
A2 = None
A3 = None
- else:
+ else:
A1 = None
A2 = None
A3 = None
-
+
at_simple_ring = simple_ring(energy=self.E0,
circumference=self.L,
harmonic_number=harmonic_number,
@@ -576,9 +578,9 @@ class Synchrotron:
emitx=self.emit[0],
emity=self.emit[1],
espread=self.sigma_delta,
- taux=self.tau[0]/self.T0,
- tauy=self.tau[1]/self.T0,
- tauz=self.tau[2]/self.T0,
+ taux=self.tau[0] / self.T0,
+ tauy=self.tau[1] / self.T0,
+ tauz=self.tau[2] / self.T0,
U0=self.U0,
TimeLag=TimeLag)
return at_simple_ring
diff --git a/mbtrack2/tracking/wakepotential.py b/mbtrack2/tracking/wakepotential.py
index 46abfbd..3003b1f 100644
--- a/mbtrack2/tracking/wakepotential.py
+++ b/mbtrack2/tracking/wakepotential.py
@@ -76,6 +76,7 @@ class WakePotential(Element):
Reduce wake function samping by an integer factor.
"""
+
def __init__(self, ring, wakefield, n_bin=80, interp_on_postion=True):
self.wakefield = wakefield
self.types = self.wakefield.wake_components
@@ -303,12 +304,12 @@ class WakePotential(Element):
for wake_type in self.types:
tau0, Wp = self.get_wakepotential(bunch, wake_type)
if self.interp_on_postion:
- Wp_interp = np.interp(bunch["tau"], tau0 + self.tau_mean, Wp,
- 0, 0)
+ Wp_interp = np.interp(bunch["tau"], tau0 + self.tau_mean,
+ Wp, 0, 0)
else:
- Wp_interp = np.interp(self.center, tau0 + self.tau_mean, Wp,
- 0, 0)
- Wp_interp= Wp_interp[self.sorted_index]
+ Wp_interp = np.interp(self.center, tau0 + self.tau_mean,
+ Wp, 0, 0)
+ Wp_interp = Wp_interp[self.sorted_index]
if wake_type == "Wlong":
bunch["delta"] += Wp_interp * bunch.charge / self.ring.E0
elif wake_type == "Wxdip":
@@ -645,6 +646,7 @@ class LongRangeResistiveWall(Element):
[1] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and
interbunch collective beam motions in storage rings." NIM.A (2016).
"""
+
def __init__(self,
ring,
beam,
diff --git a/mbtrack2/utilities/beamloading.py b/mbtrack2/utilities/beamloading.py
index caecc15..19e51e9 100644
--- a/mbtrack2/utilities/beamloading.py
+++ b/mbtrack2/utilities/beamloading.py
@@ -31,6 +31,7 @@ class BeamLoadingEquilibrium():
B1 : lower intergration boundary
B2 : upper intergration boundary
"""
+
def __init__(self,
ring,
cavity_list,
@@ -326,8 +327,8 @@ class BeamLoadingEquilibrium():
print("The final center of mass offset is " +
str(self.center_of_mass() * 1e12) + " ps")
else:
- print("The final energy balance is " + str(self.energy_balance()) +
- " eV")
+ print("The final energy balance is " +
+ str(self.energy_balance()) + " eV")
if not sol.success:
print("The algorithm has converged: " + str(sol.success))
diff --git a/mbtrack2/utilities/optics.py b/mbtrack2/utilities/optics.py
index 806d52b..3c7ff92 100644
--- a/mbtrack2/utilities/optics.py
+++ b/mbtrack2/utilities/optics.py
@@ -54,6 +54,7 @@ class Optics:
plot(self, var, option, n_points=1000)
Plot optical variables.
"""
+
def __init__(self,
lattice_file=None,
local_beta=None,
@@ -483,6 +484,7 @@ class PhysicalModel:
Return the effective radius of the chamber for resistive wall
calculations.
"""
+
def __init__(self,
ring,
x_right,
diff --git a/pyproject.toml b/pyproject.toml
index ae962b4..9d1b011 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "mbtrack2"
-version = "0.5.0"
+version = "0.6.0"
description = "A coherent object-oriented framework to work on collective effects in synchrotrons."
authors = ["Alexis Gamelin <alexis.gamelin@synchrotron-soleil.fr>"]
license = "BSD-3-Clause"
@@ -16,7 +16,7 @@ h5py = "^3.6"
mpi4py = "^3.1"
matplotlib = "^3.5"
mpmath = "^1.2.1"
-accelerator-toolbox = ">= 0.3.0"
+accelerator-toolbox = ">= 0.5.0"
seaborn = "^0.12"
[tool.poetry.group.dev.dependencies]
diff --git a/tests/conftest.py b/tests/conftest.py
index cee5c6b..356dff4 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,6 +1,8 @@
-import pytest
import numpy as np
-from mbtrack2 import Optics, Synchrotron, Electron, Bunch
+import pytest
+
+from mbtrack2 import Bunch, Electron, Optics, Synchrotron
+
@pytest.fixture
def local_optics():
diff --git a/tests/test_bunch.py b/tests/test_bunch.py
index 19aceaf..7c1499a 100644
--- a/tests/test_bunch.py
+++ b/tests/test_bunch.py
@@ -1,8 +1,10 @@
-import pytest
import numpy as np
+import pytest
from scipy.constants import e
+
from mbtrack2 import Bunch
+
def test_bunch_values(demo_ring):
mp_number = 10
current = 20e-3
diff --git a/tests/test_optics.py b/tests/test_optics.py
index 2889154..ceabcba 100644
--- a/tests/test_optics.py
+++ b/tests/test_optics.py
@@ -1,7 +1,9 @@
-import pytest
import numpy as np
+import pytest
+
from mbtrack2 import Optics
+
def test_local_optics():
beta = np.array([1, 1])
alpha = np.array([0, 0])
diff --git a/tests/test_synchrotron.py b/tests/test_synchrotron.py
index 32572ab..ad152a3 100644
--- a/tests/test_synchrotron.py
+++ b/tests/test_synchrotron.py
@@ -1,7 +1,9 @@
-import pytest
import numpy as np
+import pytest
from scipy.constants import c, e
-from mbtrack2 import Synchrotron, Electron
+
+from mbtrack2 import Electron, Synchrotron
+
def test_synchrotron_values(local_optics):
h = 20
--
GitLab