Skip to content
GitLab
Explore
Sign in
Register
Primary navigation
Search or go to…
Project
M
mbtrack2
Manage
Activity
Members
Labels
Plan
Issues
Issue boards
Milestones
Iterations
Wiki
Requirements
Code
Merge requests
Repository
Branches
Commits
Tags
Repository graph
Compare revisions
Snippets
Locked files
Deploy
Releases
Package registry
Model registry
Operate
Terraform modules
Analyze
Model experiments
Help
Help
Support
GitLab documentation
Compare GitLab plans
GitLab community forum
Contribute to GitLab
Provide feedback
Keyboard shortcuts
?
Snippets
Groups
Projects
Show more breadcrumbs
PA
Collective Effects
mbtrack2
Commits
60325b95
Commit
60325b95
authored
Sep 28, 2020
by
Gamelin Alexis
Browse files
Options
Downloads
Patches
Plain Diff
Move ImpedanceModel to a separate file to avoid circular import
parent
622b5a26
No related branches found
No related tags found
No related merge requests found
Changes
3
Show whitespace changes
Inline
Side-by-side
Showing
3 changed files
collective_effects/__init__.py
+3
-2
3 additions, 2 deletions
collective_effects/__init__.py
collective_effects/impedance_model.py
+424
-0
424 additions, 0 deletions
collective_effects/impedance_model.py
collective_effects/wakefield.py
+1
-414
1 addition, 414 deletions
collective_effects/wakefield.py
with
428 additions
and
416 deletions
collective_effects/__init__.py
+
3
−
2
View file @
60325b95
...
@@ -9,8 +9,9 @@ from mbtrack2.collective_effects.instabilities import mbi_threshold, cbi_thresho
...
@@ -9,8 +9,9 @@ from mbtrack2.collective_effects.instabilities import mbi_threshold, cbi_thresho
from
mbtrack2.collective_effects.resistive_wall
import
skin_depth
,
CircularResistiveWall
from
mbtrack2.collective_effects.resistive_wall
import
skin_depth
,
CircularResistiveWall
from
mbtrack2.collective_effects.resonator
import
Resonator
from
mbtrack2.collective_effects.resonator
import
Resonator
from
mbtrack2.collective_effects.tapers
import
StupakovRectangularTaper
,
StupakovCircularTaper
from
mbtrack2.collective_effects.tapers
import
StupakovRectangularTaper
,
StupakovCircularTaper
from
mbtrack2.collective_effects.wakefield
import
ComplexData
,
Impedance
,
WakeFunction
,
ImpedanceModel
,
WakeField
from
mbtrack2.collective_effects.wakefield
import
ComplexData
,
Impedance
,
WakeFunction
,
WakeField
from
mbtrack2.collective_effects.utilities
import
read_CST
,
read_IW2D
,
spectral_density
from
mbtrack2.collective_effects.utilities
import
read_CST
,
read_IW2D
,
spectral_density
from
mbtrack2.collective_effects.utilities
import
gaussian_bunch_spectrum
,
beam_spectrum
,
effective_impedance
from
mbtrack2.collective_effects.utilities
import
gaussian_bunch_spectrum
,
beam_spectrum
,
effective_impedance
from
mbtrack2.collective_effects.utilities
import
yokoya_elliptic
,
beam_loss_factor
from
mbtrack2.collective_effects.utilities
import
yokoya_elliptic
,
beam_loss_factor
from
mbtrack2.collective_effects.utilities
import
double_sided_impedance
from
mbtrack2.collective_effects.utilities
import
double_sided_impedance
from
mbtrack2.collective_effects.impedance_model
import
ImpedanceModel
\ No newline at end of file
This diff is collapsed.
Click to expand it.
collective_effects/impedance_model.py
0 → 100644
+
424
−
0
View file @
60325b95
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 28 15:18:23 2020
@author: gamelina
"""
import
pandas
as
pd
import
numpy
as
np
import
matplotlib.pyplot
as
plt
from
scipy.integrate
import
trapz
from
mpl_toolkits.axes_grid1.inset_locator
import
inset_axes
from
mbtrack2.collective_effects.utilities
import
(
beam_spectrum
,
gaussian_bunch_spectrum
,
beam_loss_factor
,
spectral_density
,
effective_impedance
)
from
mbtrack2.collective_effects.wakefield
import
WakeField
from
mbtrack2.tracking.element
import
Element
class
ImpedanceModel
(
Element
):
"""
Define the impedance model of the machine.
Parameters
----------
ring : Synchrotron object
wakefield_list : list of WakeField objects
WakeFields to add to the model.
wakefiled_postions : list
Longitudinal positions corresponding to the added Wakfields.
Attributes
----------
wakefields : list of WakeField objects
WakeFields in the model.
positions : array
WakeFields positions.
names : array
Names of (unique) WakeField objects.
sum : WakeField
Sum of every WakeField in the model.
sum_
"
name
"
: WakeField
Sum of every Wakefield with the same
"
name
"
.
sum_names : array
Names of attributes where the WakeFields are summed by name.
Methods
-------
add(wakefield_list, wakefiled_postions)
Add elements to the model.
add_multiple_elements(wakefield, wakefiled_postions)
Add the same element at different locations to the model.
sum_elements()
Sum all WakeFields into self.sum.
update_name_list()
Update self.names with uniques names of self.wakefields.
find_wakefield(name)
Return indexes of WakeFields with the same name in self.wakefields.
sum_by_name(name)
Sum the elements with the same name in the model into sum_name.
sum_by_name_all(name)
Sum all the elements with the same name in the model into sum_name.
plot_area(Z_type=
"
Zlong
"
, component=
"
real
"
, sigma=None, attr_list=None)
Plot the contributions of different kind of WakeFields.
"""
def
__init__
(
self
,
ring
,
wakefield_list
=
None
,
wakefiled_postions
=
None
):
self
.
ring
=
ring
self
.
optics
=
self
.
ring
.
optics
self
.
wakefields
=
[]
self
.
positions
=
np
.
array
([])
self
.
names
=
np
.
array
([])
self
.
sum_names
=
np
.
array
([])
self
.
add
(
wakefield_list
,
wakefiled_postions
)
def
track
(
self
,
beam
):
"""
Track a beam object through this Element.
Parameters
----------
beam : Beam object
"""
raise
NotImplementedError
def
sum_elements
(
self
):
"""
Sum all WakeFields into self.sum
"""
beta
=
self
.
optics
.
beta
(
self
.
positions
)
self
.
sum
=
WakeField
.
add_several_wakefields
(
self
.
wakefields
,
beta
)
if
"
sum
"
not
in
self
.
sum_names
:
self
.
sum_names
=
np
.
append
(
self
.
sum_names
,
"
sum
"
)
def
update_name_list
(
self
):
"""
Update self.names with uniques names of self.wakefields.
"""
for
wakefield
in
self
.
wakefields
:
if
wakefield
.
name
is
None
:
continue
if
wakefield
.
name
not
in
self
.
names
:
self
.
names
=
np
.
append
(
self
.
names
,
wakefield
.
name
)
def
count_elements
(
self
):
"""
Count number of each type of WakeField in the model.
"""
self
.
count
=
np
.
zeros
(
len
(
self
.
names
))
for
wakefield
in
self
.
wakefields
:
if
wakefield
.
name
is
not
None
:
self
.
count
+=
(
wakefield
.
name
==
self
.
names
).
astype
(
int
)
def
find_wakefield
(
self
,
name
):
"""
Return indexes of WakeFields with the same name in self.wakefields.
Parameters
----------
name : str
WakeField name.
Returns
-------
index : list
Index of positions in self.wakefields.
"""
index
=
[]
for
i
,
wakefield
in
enumerate
(
self
.
wakefields
):
if
wakefield
.
name
==
name
:
index
.
append
(
i
)
return
index
def
sum_by_name
(
self
,
name
):
"""
Sum the elements with the same name in the model into sum_name.
Parameters
----------
name : str
Name of the WakeField to sum.
"""
attribute_name
=
"
sum_
"
+
name
index
=
self
.
find_wakefield
(
name
)
beta
=
self
.
optics
.
beta
(
self
.
positions
[
index
])
wakes
=
[]
for
i
in
index
:
wakes
.
append
(
self
.
wakefields
[
i
])
wake_sum
=
WakeField
.
add_several_wakefields
(
wakes
,
beta
)
setattr
(
self
,
attribute_name
,
wake_sum
)
if
attribute_name
not
in
self
.
sum_names
:
self
.
sum_names
=
np
.
append
(
self
.
sum_names
,
attribute_name
)
def
sum_by_name_all
(
self
):
"""
Sum all the elements with the same name in the model into sum_name.
"""
for
name
in
self
.
names
:
self
.
sum_by_name
(
name
)
def
add
(
self
,
wakefield_list
,
wakefiled_postions
):
"""
Add elements to the model.
Parameters
----------
wakefield_list : list of WakeField objects
WakeFields to add to the model.
wakefiled_postions : list
Longitudinal positions corresponding to the added Wakfields.
"""
if
(
wakefield_list
is
not
None
)
and
(
wakefiled_postions
is
not
None
):
for
wakefield
in
wakefield_list
:
self
.
wakefields
.
append
(
wakefield
)
for
position
in
wakefiled_postions
:
self
.
positions
=
np
.
append
(
self
.
positions
,
position
)
self
.
update_name_list
()
def
add_multiple_elements
(
self
,
wakefield
,
wakefiled_postions
):
"""
Add the same element at different locations to the model.
Parameters
----------
WakeField : WakeField object
WakeField to add to the model.
wakefiled_postions : list
Longitudinal positions corresponding to the added Wakfield.
"""
for
position
in
wakefiled_postions
:
self
.
positions
=
np
.
append
(
self
.
positions
,
position
)
self
.
wakefields
.
append
(
wakefield
)
self
.
update_name_list
()
def
plot_area
(
self
,
Z_type
=
"
Zlong
"
,
component
=
"
real
"
,
sigma
=
None
,
attr_list
=
None
,
zoom
=
False
):
"""
Plot the contributions of different kind of WakeFields.
Parameters
----------
Z_type : str, optional
Type of impedance to plot.
component : str, optional
Component to plot, can be
"
real
"
or
"
imag
"
.
sigma : float, optional
RMS bunch length in [s] to use for the spectral density. If equal
to None, the spectral density is not plotted.
attr_list : list or array of str, optional
Attributes to plot.
zoom : bool
If True, add a zoomed plot on top right corner.
"""
if
attr_list
is
None
:
attr_list
=
self
.
sum_names
[
self
.
sum_names
!=
"
sum
"
]
# manage legend
Ztype_dict
=
{
"
Zlong
"
:
0
,
"
Zxdip
"
:
1
,
"
Zydip
"
:
2
,
"
Zxquad
"
:
3
,
"
Zyquad
"
:
4
}
scale
=
[
1e-3
,
1e-6
,
1e-6
,
1e-6
,
1e-6
]
label_list
=
[
r
"
$Z_{long} \; [k\Omega]$
"
,
r
"
$\sum_{j} \beta_{x,j} Z_{x,j}^{Dip} \; [M\Omega]$
"
,
r
"
$\sum_{j} \beta_{y,j} Z_{y,j}^{Dip} \; [M\Omega]$
"
,
r
"
$\sum_{j} \beta_{x,j} Z_{x,j}^{Quad} \; [M\Omega]$
"
,
r
"
$\sum_{j} \beta_{y,j} Z_{y,j}^{Quad} \; [M\Omega]$
"
]
leg
=
Ztype_dict
[
Z_type
]
# sort plot by decresing area size
area
=
np
.
zeros
((
len
(
attr_list
),))
for
index
,
attr
in
enumerate
(
attr_list
):
try
:
sum_imp
=
getattr
(
getattr
(
self
,
attr
),
Z_type
)
area
[
index
]
=
trapz
(
sum_imp
.
data
[
component
],
sum_imp
.
data
.
index
)
except
AttributeError
:
pass
sorted_index
=
area
.
argsort
()[::
-
1
]
# Init fig
fig
=
plt
.
figure
()
ax
=
fig
.
gca
()
zero_impedance
=
getattr
(
self
.
sum
,
Z_type
)
*
0
total_imp
=
0
legend
=
[]
if
sigma
is
not
None
:
legend
.
append
(
"
Spectral density for sigma =
"
+
str
(
sigma
)
+
"
s
"
)
# Main plot
for
index
in
sorted_index
:
attr
=
attr_list
[
index
]
# Set all impedances with common indexes using + zero_impedance
try
:
sum_imp
=
getattr
(
getattr
(
self
,
attr
),
Z_type
)
+
zero_impedance
ax
.
fill_between
(
sum_imp
.
data
.
index
*
1e-9
,
total_imp
,
total_imp
+
sum_imp
.
data
[
component
]
*
scale
[
leg
])
total_imp
+=
sum_imp
.
data
[
component
]
*
scale
[
leg
]
if
attr
[:
4
]
==
"
sum_
"
:
legend
.
append
(
attr
[
4
:])
else
:
legend
.
append
(
attr
)
except
AttributeError
:
pass
if
sigma
is
not
None
:
spect
=
spectral_density
(
zero_impedance
.
data
.
index
,
sigma
)
spect
=
spect
/
spect
.
max
()
*
total_imp
.
max
()
ax
.
plot
(
sum_imp
.
data
.
index
*
1e-9
,
spect
,
'
r
'
,
linewidth
=
2.5
)
ax
.
legend
(
legend
,
loc
=
"
upper left
"
)
ax
.
set_xlabel
(
"
Frequency [GHz]
"
)
ax
.
set_ylabel
(
label_list
[
leg
]
+
"
-
"
+
component
+
"
part
"
)
ax
.
set_title
(
label_list
[
leg
]
+
"
-
"
+
component
+
"
part
"
)
if
zoom
is
True
:
in_ax
=
inset_axes
(
ax
,
width
=
"
30%
"
,
# width = 30% of parent_bbox
height
=
1.5
,
# height : 1 inch
loc
=
1
)
total_imp
=
0
for
index
in
sorted_index
:
attr
=
attr_list
[
index
]
# Set all impedances with common indexes using + zero_impedance
try
:
sum_imp
=
getattr
(
getattr
(
self
,
attr
),
Z_type
)
+
zero_impedance
in_ax
.
fill_between
(
sum_imp
.
data
.
index
*
1e-3
,
total_imp
,
total_imp
+
sum_imp
.
data
[
component
]
*
1e-9
)
total_imp
+=
sum_imp
.
data
[
component
]
*
1e-9
except
AttributeError
:
pass
in_ax
.
set_xlim
([
0
,
200
])
in_ax
.
set_xlabel
(
"
Frequency [kHz]
"
)
in_ax
.
set_ylabel
(
r
"
$[G\Omega]$
"
)
return
fig
def
effective_impedance
(
self
,
m
,
mu
,
sigma
,
M
,
tuneS
,
xi
=
None
,
mode
=
"
Hermite
"
):
"""
Compute the longitudinal and transverse effective impedance.
Parameters
----------
mu : int
coupled bunch mode number, goes from 0 to (M-1) where M is the
number of bunches
m : int
head-tail (or azimutal/synchrotron) mode number
sigma : float
RMS bunch length in [s]
M : int
Number of bunches.
tuneS : float
Synchrotron tune.
xi : float, optional
(non-normalized) chromaticity
mode: str, optional
type of the mode taken into account for the computation:
-
"
Hermite
"
modes for Gaussian bunches
Returns
-------
summary : DataFrame
Longitudinal and transverse effective impedance.
"""
attr_list
=
self
.
sum_names
eff_array
=
np
.
zeros
((
len
(
attr_list
),
3
),
dtype
=
complex
)
for
i
,
attr
in
enumerate
(
attr_list
):
try
:
impedance
=
getattr
(
getattr
(
self
,
attr
),
"
Zlong
"
)
eff_array
[
i
,
0
]
=
effective_impedance
(
self
.
ring
,
impedance
,
m
,
mu
,
sigma
,
M
,
tuneS
,
xi
,
mode
)
except
AttributeError
:
pass
try
:
impedance
=
getattr
(
getattr
(
self
,
attr
),
"
Zxdip
"
)
eff_array
[
i
,
1
]
=
effective_impedance
(
self
.
ring
,
impedance
,
m
,
mu
,
sigma
,
M
,
tuneS
,
xi
,
mode
)
except
AttributeError
:
pass
try
:
impedance
=
getattr
(
getattr
(
self
,
attr
),
"
Zydip
"
)
eff_array
[
i
,
2
]
=
effective_impedance
(
self
.
ring
,
impedance
,
m
,
mu
,
sigma
,
M
,
tuneS
,
xi
,
mode
)
except
AttributeError
:
pass
eff_array
[:,
0
]
=
eff_array
[:,
0
]
*
self
.
ring
.
omega0
*
1e3
eff_array
[:,
1
]
=
eff_array
[:,
1
]
*
1e-3
eff_array
[:,
2
]
=
eff_array
[:,
2
]
*
1e-3
summary
=
pd
.
DataFrame
(
eff_array
,
index
=
attr_list
,
columns
=
[
"
Z/n [mOhm]
"
,
"
sum betax x Zxeff [kOhm]
"
,
"
sum betay x Zyeff [kOhm]
"
])
return
summary
def
energy_loss
(
self
,
sigma
,
M
,
bunch_spacing
,
I
,
n_points
=
10e6
):
"""
Compute the beam and bunch loss factor and energy losses for each type
of element in the model assuming Gaussian bunches and constant spacing
between bunches.
Parameters
----------
sigma : float
RMS bunch length in [s].
M : int
Number of bunches in the beam.
bunch_spacing : float
Time between two bunches in [s].
I : float
Total beam current in [A].
n_points : float, optional
Number of points used in the frequency spectrums.
Returns
-------
summary : Dataframe
Contains the beam and bunch loss factor and energy loss for the
full model and for each type of different component.
"""
fmax
=
self
.
sum
.
Zlong
.
data
.
index
.
max
()
fmin
=
self
.
sum
.
Zlong
.
data
.
index
.
min
()
Q
=
I
*
self
.
ring
.
T0
/
M
if
fmin
>=
0
:
fmin
=
-
1
*
fmax
f
=
np
.
linspace
(
fmin
,
fmax
,
int
(
n_points
))
beam_spect
=
beam_spectrum
(
f
,
M
,
bunch_spacing
,
sigma
=
sigma
)
bunch_spect
=
gaussian_bunch_spectrum
(
f
,
sigma
)
attr_list
=
self
.
sum_names
loss_array
=
np
.
zeros
((
len
(
attr_list
),
2
))
for
i
,
attr
in
enumerate
(
attr_list
):
try
:
impedance
=
getattr
(
getattr
(
self
,
attr
),
"
Zlong
"
)
loss_array
[
i
,
0
]
=
beam_loss_factor
(
impedance
,
f
,
beam_spect
,
self
.
ring
)
loss_array
[
i
,
1
]
=
beam_loss_factor
(
impedance
,
f
,
bunch_spect
,
self
.
ring
)
except
AttributeError
:
pass
loss_array
=
loss_array
*
1e-12
summary
=
pd
.
DataFrame
(
loss_array
,
index
=
attr_list
,
columns
=
[
"
loss factor (beam) [V/pC]
"
,
"
loss factor (bunch) [V/pC]
"
])
summary
[
"
P (beam) [W]
"
]
=
summary
[
"
loss factor (beam) [V/pC]
"
]
*
1e12
*
Q
**
2
/
(
self
.
ring
.
T0
)
summary
[
"
P (bunch) [W]
"
]
=
summary
[
"
loss factor (bunch) [V/pC]
"
]
*
1e12
*
Q
**
2
/
(
self
.
ring
.
T0
)
*
M
return
summary
This diff is collapsed.
Click to expand it.
collective_effects/wakefield.py
+
1
−
414
View file @
60325b95
...
@@ -11,17 +11,11 @@ import warnings
...
@@ -11,17 +11,11 @@ import warnings
import
re
import
re
import
pandas
as
pd
import
pandas
as
pd
import
numpy
as
np
import
numpy
as
np
import
matplotlib.pyplot
as
plt
import
scipy
as
sc
import
scipy
as
sc
from
copy
import
deepcopy
from
copy
import
deepcopy
from
scipy.interpolate
import
interp1d
from
scipy.interpolate
import
interp1d
from
scipy.integrate
import
trapz
from
scipy.integrate
import
trapz
from
scipy.constants
import
c
from
scipy.constants
import
c
from
mpl_toolkits.axes_grid1.inset_locator
import
inset_axes
from
mbtrack2.tracking.element
import
Element
from
mbtrack2.collective_effects.utilities
import
(
beam_spectrum
,
gaussian_bunch_spectrum
,
beam_loss_factor
,
spectral_density
,
effective_impedance
)
class
ComplexData
:
class
ComplexData
:
"""
"""
...
@@ -771,410 +765,3 @@ class WakeField:
...
@@ -771,410 +765,3 @@ class WakeField:
else
:
else
:
raise
ValueError
(
"
Error in input.
"
)
raise
ValueError
(
"
Error in input.
"
)
\ No newline at end of file
class
ImpedanceModel
(
Element
):
"""
Define the impedance model of the machine.
Parameters
----------
ring : Synchrotron object
wakefield_list : list of WakeField objects
WakeFields to add to the model.
wakefiled_postions : list
Longitudinal positions corresponding to the added Wakfields.
Attributes
----------
wakefields : list of WakeField objects
WakeFields in the model.
positions : array
WakeFields positions.
names : array
Names of (unique) WakeField objects.
sum : WakeField
Sum of every WakeField in the model.
sum_
"
name
"
: WakeField
Sum of every Wakefield with the same
"
name
"
.
sum_names : array
Names of attributes where the WakeFields are summed by name.
Methods
-------
add(wakefield_list, wakefiled_postions)
Add elements to the model.
add_multiple_elements(wakefield, wakefiled_postions)
Add the same element at different locations to the model.
sum_elements()
Sum all WakeFields into self.sum.
update_name_list()
Update self.names with uniques names of self.wakefields.
find_wakefield(name)
Return indexes of WakeFields with the same name in self.wakefields.
sum_by_name(name)
Sum the elements with the same name in the model into sum_name.
sum_by_name_all(name)
Sum all the elements with the same name in the model into sum_name.
plot_area(Z_type=
"
Zlong
"
, component=
"
real
"
, sigma=None, attr_list=None)
Plot the contributions of different kind of WakeFields.
"""
def
__init__
(
self
,
ring
,
wakefield_list
=
None
,
wakefiled_postions
=
None
):
self
.
ring
=
ring
self
.
optics
=
self
.
ring
.
optics
self
.
wakefields
=
[]
self
.
positions
=
np
.
array
([])
self
.
names
=
np
.
array
([])
self
.
sum_names
=
np
.
array
([])
self
.
add
(
wakefield_list
,
wakefiled_postions
)
def
track
(
self
,
beam
):
"""
Track a beam object through this Element.
Parameters
----------
beam : Beam object
"""
raise
NotImplementedError
def
sum_elements
(
self
):
"""
Sum all WakeFields into self.sum
"""
beta
=
self
.
optics
.
beta
(
self
.
positions
)
self
.
sum
=
WakeField
.
add_several_wakefields
(
self
.
wakefields
,
beta
)
if
"
sum
"
not
in
self
.
sum_names
:
self
.
sum_names
=
np
.
append
(
self
.
sum_names
,
"
sum
"
)
def
update_name_list
(
self
):
"""
Update self.names with uniques names of self.wakefields.
"""
for
wakefield
in
self
.
wakefields
:
if
wakefield
.
name
is
None
:
continue
if
wakefield
.
name
not
in
self
.
names
:
self
.
names
=
np
.
append
(
self
.
names
,
wakefield
.
name
)
def
count_elements
(
self
):
"""
Count number of each type of WakeField in the model.
"""
self
.
count
=
np
.
zeros
(
len
(
self
.
names
))
for
wakefield
in
self
.
wakefields
:
if
wakefield
.
name
is
not
None
:
self
.
count
+=
(
wakefield
.
name
==
self
.
names
).
astype
(
int
)
def
find_wakefield
(
self
,
name
):
"""
Return indexes of WakeFields with the same name in self.wakefields.
Parameters
----------
name : str
WakeField name.
Returns
-------
index : list
Index of positions in self.wakefields.
"""
index
=
[]
for
i
,
wakefield
in
enumerate
(
self
.
wakefields
):
if
wakefield
.
name
==
name
:
index
.
append
(
i
)
return
index
def
sum_by_name
(
self
,
name
):
"""
Sum the elements with the same name in the model into sum_name.
Parameters
----------
name : str
Name of the WakeField to sum.
"""
attribute_name
=
"
sum_
"
+
name
index
=
self
.
find_wakefield
(
name
)
beta
=
self
.
optics
.
beta
(
self
.
positions
[
index
])
wakes
=
[]
for
i
in
index
:
wakes
.
append
(
self
.
wakefields
[
i
])
wake_sum
=
WakeField
.
add_several_wakefields
(
wakes
,
beta
)
setattr
(
self
,
attribute_name
,
wake_sum
)
if
attribute_name
not
in
self
.
sum_names
:
self
.
sum_names
=
np
.
append
(
self
.
sum_names
,
attribute_name
)
def
sum_by_name_all
(
self
):
"""
Sum all the elements with the same name in the model into sum_name.
"""
for
name
in
self
.
names
:
self
.
sum_by_name
(
name
)
def
add
(
self
,
wakefield_list
,
wakefiled_postions
):
"""
Add elements to the model.
Parameters
----------
wakefield_list : list of WakeField objects
WakeFields to add to the model.
wakefiled_postions : list
Longitudinal positions corresponding to the added Wakfields.
"""
if
(
wakefield_list
is
not
None
)
and
(
wakefiled_postions
is
not
None
):
for
wakefield
in
wakefield_list
:
self
.
wakefields
.
append
(
wakefield
)
for
position
in
wakefiled_postions
:
self
.
positions
=
np
.
append
(
self
.
positions
,
position
)
self
.
update_name_list
()
def
add_multiple_elements
(
self
,
wakefield
,
wakefiled_postions
):
"""
Add the same element at different locations to the model.
Parameters
----------
WakeField : WakeField object
WakeField to add to the model.
wakefiled_postions : list
Longitudinal positions corresponding to the added Wakfield.
"""
for
position
in
wakefiled_postions
:
self
.
positions
=
np
.
append
(
self
.
positions
,
position
)
self
.
wakefields
.
append
(
wakefield
)
self
.
update_name_list
()
def
plot_area
(
self
,
Z_type
=
"
Zlong
"
,
component
=
"
real
"
,
sigma
=
None
,
attr_list
=
None
,
zoom
=
False
):
"""
Plot the contributions of different kind of WakeFields.
Parameters
----------
Z_type : str, optional
Type of impedance to plot.
component : str, optional
Component to plot, can be
"
real
"
or
"
imag
"
.
sigma : float, optional
RMS bunch length in [s] to use for the spectral density. If equal
to None, the spectral density is not plotted.
attr_list : list or array of str, optional
Attributes to plot.
zoom : bool
If True, add a zoomed plot on top right corner.
"""
if
attr_list
is
None
:
attr_list
=
self
.
sum_names
[
self
.
sum_names
!=
"
sum
"
]
# manage legend
Ztype_dict
=
{
"
Zlong
"
:
0
,
"
Zxdip
"
:
1
,
"
Zydip
"
:
2
,
"
Zxquad
"
:
3
,
"
Zyquad
"
:
4
}
scale
=
[
1e-3
,
1e-6
,
1e-6
,
1e-6
,
1e-6
]
label_list
=
[
r
"
$Z_{long} \; [k\Omega]$
"
,
r
"
$\sum_{j} \beta_{x,j} Z_{x,j}^{Dip} \; [M\Omega]$
"
,
r
"
$\sum_{j} \beta_{y,j} Z_{y,j}^{Dip} \; [M\Omega]$
"
,
r
"
$\sum_{j} \beta_{x,j} Z_{x,j}^{Quad} \; [M\Omega]$
"
,
r
"
$\sum_{j} \beta_{y,j} Z_{y,j}^{Quad} \; [M\Omega]$
"
]
leg
=
Ztype_dict
[
Z_type
]
# sort plot by decresing area size
area
=
np
.
zeros
((
len
(
attr_list
),))
for
index
,
attr
in
enumerate
(
attr_list
):
try
:
sum_imp
=
getattr
(
getattr
(
self
,
attr
),
Z_type
)
area
[
index
]
=
trapz
(
sum_imp
.
data
[
component
],
sum_imp
.
data
.
index
)
except
AttributeError
:
pass
sorted_index
=
area
.
argsort
()[::
-
1
]
# Init fig
fig
=
plt
.
figure
()
ax
=
fig
.
gca
()
zero_impedance
=
getattr
(
self
.
sum
,
Z_type
)
*
0
total_imp
=
0
legend
=
[]
if
sigma
is
not
None
:
legend
.
append
(
"
Spectral density for sigma =
"
+
str
(
sigma
)
+
"
s
"
)
# Main plot
for
index
in
sorted_index
:
attr
=
attr_list
[
index
]
# Set all impedances with common indexes using + zero_impedance
try
:
sum_imp
=
getattr
(
getattr
(
self
,
attr
),
Z_type
)
+
zero_impedance
ax
.
fill_between
(
sum_imp
.
data
.
index
*
1e-9
,
total_imp
,
total_imp
+
sum_imp
.
data
[
component
]
*
scale
[
leg
])
total_imp
+=
sum_imp
.
data
[
component
]
*
scale
[
leg
]
if
attr
[:
4
]
==
"
sum_
"
:
legend
.
append
(
attr
[
4
:])
else
:
legend
.
append
(
attr
)
except
AttributeError
:
pass
if
sigma
is
not
None
:
spect
=
spectral_density
(
zero_impedance
.
data
.
index
,
sigma
)
spect
=
spect
/
spect
.
max
()
*
total_imp
.
max
()
ax
.
plot
(
sum_imp
.
data
.
index
*
1e-9
,
spect
,
'
r
'
,
linewidth
=
2.5
)
ax
.
legend
(
legend
,
loc
=
"
upper left
"
)
ax
.
set_xlabel
(
"
Frequency [GHz]
"
)
ax
.
set_ylabel
(
label_list
[
leg
]
+
"
-
"
+
component
+
"
part
"
)
ax
.
set_title
(
label_list
[
leg
]
+
"
-
"
+
component
+
"
part
"
)
if
zoom
is
True
:
in_ax
=
inset_axes
(
ax
,
width
=
"
30%
"
,
# width = 30% of parent_bbox
height
=
1.5
,
# height : 1 inch
loc
=
1
)
total_imp
=
0
for
index
in
sorted_index
:
attr
=
attr_list
[
index
]
# Set all impedances with common indexes using + zero_impedance
try
:
sum_imp
=
getattr
(
getattr
(
self
,
attr
),
Z_type
)
+
zero_impedance
in_ax
.
fill_between
(
sum_imp
.
data
.
index
*
1e-3
,
total_imp
,
total_imp
+
sum_imp
.
data
[
component
]
*
1e-9
)
total_imp
+=
sum_imp
.
data
[
component
]
*
1e-9
except
AttributeError
:
pass
in_ax
.
set_xlim
([
0
,
200
])
in_ax
.
set_xlabel
(
"
Frequency [kHz]
"
)
in_ax
.
set_ylabel
(
r
"
$[G\Omega]$
"
)
return
fig
def
effective_impedance
(
self
,
m
,
mu
,
sigma
,
M
,
tuneS
,
xi
=
None
,
mode
=
"
Hermite
"
):
"""
Compute the longitudinal and transverse effective impedance.
Parameters
----------
mu : int
coupled bunch mode number, goes from 0 to (M-1) where M is the
number of bunches
m : int
head-tail (or azimutal/synchrotron) mode number
sigma : float
RMS bunch length in [s]
M : int
Number of bunches.
tuneS : float
Synchrotron tune.
xi : float, optional
(non-normalized) chromaticity
mode: str, optional
type of the mode taken into account for the computation:
-
"
Hermite
"
modes for Gaussian bunches
Returns
-------
summary : DataFrame
Longitudinal and transverse effective impedance.
"""
attr_list
=
self
.
sum_names
eff_array
=
np
.
zeros
((
len
(
attr_list
),
3
),
dtype
=
complex
)
for
i
,
attr
in
enumerate
(
attr_list
):
try
:
impedance
=
getattr
(
getattr
(
self
,
attr
),
"
Zlong
"
)
eff_array
[
i
,
0
]
=
effective_impedance
(
self
.
ring
,
impedance
,
m
,
mu
,
sigma
,
M
,
tuneS
,
xi
,
mode
)
except
AttributeError
:
pass
try
:
impedance
=
getattr
(
getattr
(
self
,
attr
),
"
Zxdip
"
)
eff_array
[
i
,
1
]
=
effective_impedance
(
self
.
ring
,
impedance
,
m
,
mu
,
sigma
,
M
,
tuneS
,
xi
,
mode
)
except
AttributeError
:
pass
try
:
impedance
=
getattr
(
getattr
(
self
,
attr
),
"
Zydip
"
)
eff_array
[
i
,
2
]
=
effective_impedance
(
self
.
ring
,
impedance
,
m
,
mu
,
sigma
,
M
,
tuneS
,
xi
,
mode
)
except
AttributeError
:
pass
eff_array
[:,
0
]
=
eff_array
[:,
0
]
*
self
.
ring
.
omega0
*
1e3
eff_array
[:,
1
]
=
eff_array
[:,
1
]
*
1e-3
eff_array
[:,
2
]
=
eff_array
[:,
2
]
*
1e-3
summary
=
pd
.
DataFrame
(
eff_array
,
index
=
attr_list
,
columns
=
[
"
Z/n [mOhm]
"
,
"
sum betax x Zxeff [kOhm]
"
,
"
sum betay x Zyeff [kOhm]
"
])
return
summary
def
energy_loss
(
self
,
sigma
,
M
,
bunch_spacing
,
I
,
n_points
=
10e6
):
"""
Compute the beam and bunch loss factor and energy losses for each type
of element in the model assuming Gaussian bunches and constant spacing
between bunches.
Parameters
----------
sigma : float
RMS bunch length in [s].
M : int
Number of bunches in the beam.
bunch_spacing : float
Time between two bunches in [s].
I : float
Total beam current in [A].
n_points : float, optional
Number of points used in the frequency spectrums.
Returns
-------
summary : Dataframe
Contains the beam and bunch loss factor and energy loss for the
full model and for each type of different component.
"""
fmax
=
self
.
sum
.
Zlong
.
data
.
index
.
max
()
fmin
=
self
.
sum
.
Zlong
.
data
.
index
.
min
()
Q
=
I
*
self
.
ring
.
T0
/
M
if
fmin
>=
0
:
fmin
=
-
1
*
fmax
f
=
np
.
linspace
(
fmin
,
fmax
,
int
(
n_points
))
beam_spect
=
beam_spectrum
(
f
,
M
,
bunch_spacing
,
sigma
=
sigma
)
bunch_spect
=
gaussian_bunch_spectrum
(
f
,
sigma
)
attr_list
=
self
.
sum_names
loss_array
=
np
.
zeros
((
len
(
attr_list
),
2
))
for
i
,
attr
in
enumerate
(
attr_list
):
try
:
impedance
=
getattr
(
getattr
(
self
,
attr
),
"
Zlong
"
)
loss_array
[
i
,
0
]
=
beam_loss_factor
(
impedance
,
f
,
beam_spect
,
self
.
ring
)
loss_array
[
i
,
1
]
=
beam_loss_factor
(
impedance
,
f
,
bunch_spect
,
self
.
ring
)
except
AttributeError
:
pass
loss_array
=
loss_array
*
1e-12
summary
=
pd
.
DataFrame
(
loss_array
,
index
=
attr_list
,
columns
=
[
"
loss factor (beam) [V/pC]
"
,
"
loss factor (bunch) [V/pC]
"
])
summary
[
"
P (beam) [W]
"
]
=
summary
[
"
loss factor (beam) [V/pC]
"
]
*
1e12
*
Q
**
2
/
(
self
.
ring
.
T0
)
summary
[
"
P (bunch) [W]
"
]
=
summary
[
"
loss factor (bunch) [V/pC]
"
]
*
1e12
*
Q
**
2
/
(
self
.
ring
.
T0
)
*
M
return
summary
This diff is collapsed.
Click to expand it.
Preview
0%
Loading
Try again
or
attach a new file
.
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Save comment
Cancel
Please
register
or
sign in
to comment