diff --git a/.gitignore b/.gitignore index ee79bd46aae84b6042273ef60da06371d79fe0fa..0f302ec1bbec39d73f7a8fd36ff34b8200fe59c9 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ *__pycache__* *.ipynb_checkpoints* test_*.py +*.ipynb +*.hdf5 diff --git a/SOLEIL_U_74BA_HOA_SYM02_V0313_cleaned.mat b/SOLEIL_U_74BA_HOA_SYM02_V0313_cleaned.mat new file mode 100644 index 0000000000000000000000000000000000000000..f021674c7c458deeff2a876be27f7853ba1a1b3b Binary files /dev/null and b/SOLEIL_U_74BA_HOA_SYM02_V0313_cleaned.mat differ diff --git a/collective_effects/data/Yokoya_elliptic_from_Elias_USPAS.csv b/collective_effects/data/Yokoya_elliptic_from_Elias_USPAS.csv new file mode 100644 index 0000000000000000000000000000000000000000..88feaf61e520a3fa8d7f528a56669a4589e7963f --- /dev/null +++ b/collective_effects/data/Yokoya_elliptic_from_Elias_USPAS.csv @@ -0,0 +1,119 @@ +x,dipy,dipx,quady,quadx +-1.38778e-17,1,1.00307,-0.00270479,0.00429341 +7.94873e-06,0.999988,1.00304,-0.00268229,0.0042735 +0.00131154,0.99802,0.998283,0.00100816,0.00100816 +0.0184543,0.972133,0.935779,0.0495391,-0.0419323 +0.0195168,0.970529,0.931905,0.0521597,-0.0445937 +0.0221637,0.966532,0.922254,0.0586882,-0.0511943 +0.0233454,0.96535,0.917945,0.0616028,-0.0541411 +0.036898,0.9518,0.86594,0.0950297,-0.0879369 +0.0376001,0.951098,0.863246,0.0965459,-0.0896878 +0.0390229,0.949675,0.857042,0.0996185,-0.0932359 +0.0404405,0.948258,0.850861,0.10268,-0.0967708 +0.046659,0.942615,0.823745,0.116109,-0.112278 +0.0579595,0.932362,0.789359,0.140512,-0.140458 +0.0585291,0.931845,0.787626,0.141742,-0.141878 +0.0596446,0.930833,0.784232,0.144151,-0.143994 +0.0639448,0.926931,0.774182,0.153438,-0.152153 +0.0765568,0.919567,0.744707,0.180674,-0.176081 +0.0790209,0.918129,0.737585,0.185995,-0.180757 +0.079368,0.917926,0.736581,0.18656,-0.181415 +0.0796277,0.917774,0.735831,0.186982,-0.181819 +0.098641,0.90472,0.680875,0.217927,-0.211402 +0.101367,0.902849,0.676101,0.222364,-0.215644 +0.101839,0.902525,0.675275,0.222994,-0.216378 +0.102833,0.902214,0.673535,0.224321,-0.217924 +0.121457,0.896392,0.640927,0.24919,-0.2469 +0.126297,0.893069,0.632452,0.255654,-0.254432 +0.126324,0.893051,0.632405,0.255689,-0.254465 +0.135123,0.887009,0.616999,0.267106,-0.265333 +0.143668,0.881143,0.605496,0.278192,-0.275888 +0.147338,0.880067,0.600556,0.282953,-0.28042 +0.148477,0.879733,0.599023,0.28408,-0.281827 +0.164594,0.875006,0.577327,0.300014,-0.295503 +0.169035,0.872782,0.571349,0.304404,-0.299271 +0.174896,0.869847,0.564512,0.310198,-0.304243 +0.177209,0.868689,0.561814,0.311798,-0.306206 +0.182895,0.865841,0.555182,0.31573,-0.31085 +0.201121,0.860497,0.53392,0.328337,-0.325736 +0.203331,0.859849,0.531343,0.329443,-0.32754 +0.203821,0.859705,0.530771,0.329688,-0.327805 +0.205573,0.859461,0.528728,0.330565,-0.328749 +0.226064,0.856601,0.513301,0.340824,-0.339797 +0.231266,0.855831,0.509385,0.343427,-0.342601 +0.237314,0.854935,0.504831,0.346309,-0.345862 +0.246069,0.853639,0.498241,0.350479,-0.349956 +0.246999,0.853501,0.497793,0.350922,-0.35039 +0.266624,0.850405,0.488348,0.36027,-0.359567 +0.269271,0.849578,0.487075,0.361531,-0.360804 +0.276534,0.847307,0.483579,0.364202,-0.3642 +0.286242,0.844273,0.478907,0.367772,-0.36641 +0.290515,0.843676,0.476851,0.369344,-0.367383 +0.302025,0.84207,0.47313,0.373577,-0.370004 +0.308485,0.841169,0.471041,0.375354,-0.371474 +0.317085,0.841142,0.468261,0.377721,-0.373432 +0.328118,0.841109,0.464694,0.380756,-0.375867 +0.328441,0.841064,0.46459,0.380845,-0.375938 +0.334771,0.840181,0.462842,0.382587,-0.377335 +0.350361,0.838005,0.458536,0.384111,-0.380775 +0.358946,0.837979,0.456166,0.384951,-0.382669 +0.364884,0.837961,0.454526,0.385532,-0.383162 +0.372612,0.837937,0.452392,0.386087,-0.383804 +0.372912,0.837844,0.452309,0.386108,-0.383829 +0.39223,0.831805,0.447911,0.387495,-0.385433 +0.396896,0.831791,0.446849,0.38783,-0.385821 +0.405467,0.831765,0.444897,0.388445,-0.387712 +0.413464,0.83174,0.443077,0.389451,-0.389476 +0.432805,0.831682,0.440122,0.391884,-0.393744 +0.438756,0.831664,0.439213,0.392633,-0.395057 +0.452603,0.831622,0.437097,0.394375,-0.395967 +0.454023,0.831617,0.43688,0.39457,-0.396061 +0.455057,0.831614,0.436727,0.394712,-0.396129 +0.487177,0.828947,0.431971,0.399129,-0.39824 +0.493006,0.828463,0.431108,0.399931,-0.398668 +0.495812,0.828454,0.430692,0.400316,-0.398874 +0.495892,0.828454,0.43068,0.400322,-0.39888 +0.530363,0.828349,0.428076,0.402796,-0.401408 +0.533582,0.828339,0.427833,0.403028,-0.401631 +0.536396,0.828331,0.427621,0.40323,-0.401826 +0.537768,0.828327,0.427517,0.403225,-0.401921 +0.573045,0.828219,0.424852,0.403118,-0.404367 +0.576166,0.82821,0.424616,0.40335,-0.404584 +0.579393,0.8282,0.424372,0.40359,-0.404593 +0.579645,0.828181,0.424353,0.403609,-0.404594 +0.612319,0.825713,0.424254,0.406035,-0.404693 +0.619359,0.825181,0.424233,0.406014,-0.404715 +0.62022,0.825116,0.42423,0.406011,-0.404717 +0.62127,0.825037,0.424075,0.406008,-0.40472 +0.652895,0.822647,0.419392,0.405912,-0.404816 +0.662089,0.821953,0.41803,0.405884,-0.404844 +0.663146,0.821873,0.417939,0.405881,-0.404848 +0.670406,0.821851,0.417316,0.405859,-0.40487 +0.683,0.821813,0.416234,0.405821,-0.405677 +0.69873,0.821765,0.414882,0.407431,-0.406685 +0.699795,0.821762,0.414879,0.40754,-0.406753 +0.711803,0.821725,0.414843,0.40877,-0.407523 +0.720136,0.8217,0.414817,0.408745,-0.408057 +0.733826,0.821658,0.414776,0.408703,-0.408099 +0.740615,0.821008,0.414755,0.408682,-0.408119 +0.756305,0.819504,0.413494,0.408635,-0.408167 +0.763329,0.81883,0.412929,0.408613,-0.408188 +0.766541,0.818523,0.412671,0.408604,-0.408389 +0.779873,0.818482,0.4116,0.408563,-0.409223 +0.79819,0.818427,0.411544,0.408508,-0.410368 +0.808425,0.818395,0.411513,0.409101,-0.411008 +0.809978,0.818511,0.411508,0.409191,-0.411105 +0.814368,0.818837,0.410816,0.409446,-0.41138 +0.829603,0.819969,0.408412,0.41033,-0.411426 +0.8477,0.821313,0.410106,0.411379,-0.411481 +0.847936,0.821296,0.410128,0.411393,-0.411482 +0.861024,0.820389,0.411353,0.411353,-0.411521 +0.861488,0.820357,0.411352,0.411352,-0.411523 +0.88982,0.818392,0.411266,0.411266,-0.409555 +0.893503,0.818137,0.411255,0.411255,-0.409299 +0.903381,0.818107,0.411225,0.411225,-0.408614 +0.912071,0.818081,0.411198,0.411198,-0.410191 +0.931705,0.818021,0.411139,0.411139,-0.413756 +0.937396,0.818004,0.411121,0.411121,-0.41479 +0.94455,0.817982,0.4111,0.4111,-0.416089 +0.946102,0.817977,0.411095,0.411095,-0.416371 diff --git a/collective_effects/resistive_wall.py b/collective_effects/resistive_wall.py index 764a59e61ae17515fcab23fe19a6460322f7fbbc..95ae59f68698a0d0bc8f51374bd0b33d66493714 100644 --- a/collective_effects/resistive_wall.py +++ b/collective_effects/resistive_wall.py @@ -8,7 +8,7 @@ Define resistive wall elements based on the Wakefield class. import numpy as np from scipy.constants import mu_0, epsilon_0, c -from CollectiveEffect.wakefield import Wakefield, Impedance +from collective_effects.wakefield import Wakefield, Impedance def skin_depth(omega, rho, mu_r = 1, epsilon_r = 1): """ General formula for the skin depth diff --git a/collective_effects/tapers.py b/collective_effects/tapers.py new file mode 100644 index 0000000000000000000000000000000000000000..0e09b36bd476ef0d97915b6d92b861bc800375d2 --- /dev/null +++ b/collective_effects/tapers.py @@ -0,0 +1,212 @@ +# -*- coding: utf-8 -*- +""" +Created on Tue Mar 31 13:15:36 2020 + +@author: gamelina +""" + +from scipy.constants import mu_0, epsilon_0, c, pi +import numpy as np +from scipy.integrate import trapz +from collective_effects.wakefield import Wakefield, Impedance + +class StupakovRectangularTaper(Wakefield): + """ + Rectangular vertical taper Wakefield element, using the low frequency + approxmiation. Assume constant taper angle. Formulas from [1]. + + ! Valid for low w + + Parameters + ---------- + frequency: frequency points where the impedance will be evaluated in [Hz] + gap_entrance : full vertical gap at taper entrance in [m] + gap_exit: full vertical gap at taper exit in [m] + length: taper length in [m] + width : full horizontal width of the taper in [m] + """ + + def __init__(self, frequency, gap_entrance, gap_exit, length, width, + m_max=100, n_points=int(1e4)): + super().__init__() + + self.frequency = frequency + self.gap_entrance = gap_entrance + self.gap_exit = gap_exit + self.length = length + self.width = width + self.m_max = m_max + self.n_points = n_points + + Zlong = Impedance(variable = frequency, function = self.long(), impedance_type='long') + Zxdip = Impedance(variable = frequency, function = self.xdip(), impedance_type='xdip') + Zydip = Impedance(variable = frequency, function = self.ydip(), impedance_type='ydip') + Zxquad = Impedance(variable = frequency, function = -1*self.quad(), impedance_type='xquad') + Zyquad = Impedance(variable = frequency, function = self.quad(), impedance_type='yquad') + + super().append_to_model(Zlong) + super().append_to_model(Zxdip) + super().append_to_model(Zydip) + super().append_to_model(Zxquad) + super().append_to_model(Zyquad) + + @property + def gap_prime(self): + return (self.gap_entrance-self.gap_exit)/self.length + + @property + def angle(self): + return np.arctan((self.gap_entrance/2 - self.gap_exit/2)/self.length) + + @property + def Z0(self): + return mu_0*c + + def long(self, frequency=None): + + if frequency is None: + frequency = self.frequency + + def F(x, m_max): + m = np.arange(0, m_max) + phi = np.outer(pi*x/2, 2*m+1) + val = 1/(2*m+1)/(np.cosh(phi)**2)*np.tanh(phi) + return val.sum(1) + + z = np.linspace(0, self.length, self.n_points) + g = np.linspace(self.gap_entrance, self.gap_exit, self.n_points) + + to_integrate = self.gap_prime**2*F(g/self.width, self.m_max) + integral = trapz(to_integrate,x=z) + + return -1j*frequency*self.Z0/(2*c)*integral + + def Z_over_n(self, f0): + 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) + val = (2*m+1)/(np.sinh(phi)**2)/np.tanh(phi) + val = x[:,None]**3*val + return val.sum(1) + + z = np.linspace(0, self.length, self.n_points) + g = np.linspace(self.gap_entrance, self.gap_exit, self.n_points) + + to_integrate = self.gap_prime**2/(g**3)*G1(g/self.width, self.m_max) + integral = trapz(to_integrate, x=z) + + 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) + val = 2*m/(np.cosh(phi)**2)*np.tanh(phi) + val = x[:,None]**2*val + return val.sum(1) + + z = np.linspace(0, self.length, self.n_points) + g = np.linspace(self.gap_entrance, self.gap_exit, self.n_points) + + to_integrate = self.gap_prime**2/(g**2)*G3(g/self.width, self.m_max) + integral = trapz(to_integrate, x=z) + + 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) + val = (2*m+1)/(np.cosh(phi)**2)*np.tanh(phi) + val = x[:,None]**2*val + return val.sum(1) + + z = np.linspace(0, self.length, self.n_points) + g = np.linspace(self.gap_entrance, self.gap_exit, self.n_points) + + to_integrate = self.gap_prime**2/(g**2)*G2(g/self.width, self.m_max) + integral = trapz(to_integrate, x=z) + + return -1j*pi*self.Z0/4*integral + +class StupakovCircularTaper(Wakefield): + """ + Circular taper Wakefield element, using the low frequency + approxmiation. Assume constant taper angle. Formulas from [1]. + + ! Valid for low w + + Parameters + ---------- + frequency: frequency points where the impedance will be evaluated in [Hz] + radius_entrance : radius at taper entrance in [m] + radius_exit : radius at taper exit in [m] + length : taper length in [m] + """ + + def __init__(self, frequency, radius_entrance, radius_exit, length, + m_max=100, n_points=int(1e4)): + super().__init__() + + self.frequency = frequency + self.radius_entrance = radius_entrance + self.radius_exit = radius_exit + self.length = length + self.m_max = m_max + self.n_points = n_points + + Zlong = Impedance(variable = frequency, function = self.long(), impedance_type='long') + Zxdip = Impedance(variable = frequency, function = self.dip(), impedance_type='xdip') + Zydip = Impedance(variable = frequency, function = self.dip(), impedance_type='ydip') + + super().append_to_model(Zlong) + super().append_to_model(Zxdip) + super().append_to_model(Zydip) + + @property + def angle(self): + return np.arctan((self.radius_entrance-self.radius_exit)/self.length) + + @property + def radius_prime(self): + return (self.radius_entrance-self.radius_exit)/self.length + + @property + def Z0(self): + return mu_0*c + + def long(self, frequency=None): + + if frequency is None: + frequency = self.frequency + + return (self.Z0/(2*pi)*np.log(self.radius_entrance/self.radius_exit) + - 1j*self.Z0*frequency/(2*c)*self.radius_prime**2*self.length) + + def Z_over_n(self, f0): + return np.imag(self.long(1))*f0 + + def dip(self, frequency=None): + + if frequency is None: + frequency = self.frequency + + z = np.linspace(0, self.length, self.n_points) + r = np.linspace(self.radius_entrance, self.radius_exit, self.n_points) + + to_integrate = self.radius_prime**2/(r**2) + integral = trapz(to_integrate, x=z) + + return (self.Z0*c/(4*pi**2*frequency)*(1/(self.radius_exit**2) - + 1/(self.radius_entrance**2)) - 1j*self.Z0/(2*pi)*integral) + + + + diff --git a/collective_effects/tools.py b/collective_effects/tools.py index 3f3f63362666aeb42311c565131a0d0fcc4266f9..d758e62323b2d7df5b40c498b827fc1be788cdca 100644 --- a/collective_effects/tools.py +++ b/collective_effects/tools.py @@ -1,22 +1,41 @@ # -*- coding: utf-8 -*- """ -Tools and utilities functions for the Wakefield and impedances classes. +Tools and utilities functions for the Wakefield and Impedances classes. @author: Alexis Gamelin -@date: 14/01/2020 +@date: 24/04/2020 """ import pandas as pd import numpy as np from scipy.integrate import trapz -from CollectiveEffect.wakefield import Wakefield, Impedance, WakeFunction -#from ..Tracking.synchrotron import Synchrotron +from collective_effects.wakefield import Wakefield, Impedance, WakeFunction -def read_CST(file, impedance_type='long'): - """Read CST file format into an Impedance object.""" +def read_CST(file, impedance_type='long', divide_by=None): + """ + Read CST file format into an Impedance object. + + Parameters + ---------- + file : str + Path to the file to read. + impedance_type : str, optional + Type of the Impedance object. + divide_by : float, optional + Divide the impedance by a value. Mainly used to normalize transverse + impedance by displacement. + + Returns + ------- + result : Impedance object + Data from file. + """ df = pd.read_csv(file, comment="#", header = None, sep = "\t", names = ["Frequency","Real","Imaginary"]) df["Frequency"] = df["Frequency"]*1e9 + if divide_by is not None: + df["Real"] = df["Real"]/divide_by + df["Imaginary"] = df["Imaginary"]/divide_by df.set_index("Frequency", inplace = True) result = Impedance(variable = df.index, function = df["Real"] + 1j*df["Imaginary"], @@ -24,7 +43,21 @@ def read_CST(file, impedance_type='long'): return result def read_IW2D(file, impedance_type='long'): - """Read IW2D file format into an Impedance object.""" + """ + Read IW2D file format into an Impedance object. + + Parameters + ---------- + file : str + Path to the file to read. + impedance_type : str, optional + Type of the Impedance object. + + Returns + ------- + result : Impedance object + Data from file. + """ df = pd.read_csv(file, sep = ' ', header = None, names = ["Frequency","Real","Imaginary"], skiprows=1) df.set_index("Frequency", inplace = True) @@ -61,9 +94,15 @@ def loss_factor(wake, sigma): if(wake.impedance_type == "long"): kloss = trapz(x = omega, y = 1/np.pi*wake.data["real"][positive_index]*gaussian_spectrum) - if(wake.impedance_type == "xdip" or wake.impedance_type == "ydip"): + elif(wake.impedance_type == "xdip" or wake.impedance_type == "ydip"): + kloss = trapz(x = omega, + y = 1/np.pi*wake.data["imag"][positive_index]*gaussian_spectrum) + elif(wake.impedance_type == "xquad" or wake.impedance_type == "yquad"): kloss = trapz(x = omega, y = 1/np.pi*wake.data["imag"][positive_index]*gaussian_spectrum) + else: + raise TypeError("Impedance type not recognized.") + if isinstance(wake, WakeFunction): # To be implemented. pass @@ -183,3 +222,60 @@ def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"): " and ydip impedance type.") return Zeff + + +def Yokoya_elliptic(x_radius , y_radius): + """ + Compute Yokoya factors for an elliptic beam pipe. + Function adapted from N. Mounet IW2D. + + Parameters + ---------- + x_radius : float + Horizontal semi-axis of the ellipse in [m]. + y_radius : float + Vertical semi-axis of the ellipse in [m]. + + Returns + ------- + yoklong : float + Yokoya factor for the longitudinal impedance. + yokxdip : float + Yokoya factor for the dipolar horizontal impedance. + yokydip : float + Yokoya factor for the dipolar vertical impedance. + yokxquad : float + Yokoya factor for the quadrupolar horizontal impedance. + yokyquad : float + Yokoya factor for the quadrupolar vertical impedance. + """ + if y_radius < x_radius: + small_semiaxis = y_radius + large_semiaxis = x_radius + else: + small_semiaxis = x_radius + large_semiaxis = y_radius + + # read Yokoya factors interpolation file + # BEWARE: columns are ratio, dipy, dipx, quady, quadx + yokoya_file = pd.read_csv("collective_effects/data/" + + "Yokoya_elliptic_from_Elias_USPAS.csv") + ratio_col = yokoya_file["x"] + # compute semi-axes ratio (first column of this file) + ratio = (large_semiaxis - small_semiaxis)/(large_semiaxis + small_semiaxis) + + # interpolate Yokoya file at the correct ratio + yoklong = 1 + + if y_radius < x_radius: + yokydip = np.interp(ratio, ratio_col, yokoya_file["dipy"]) + yokxdip = np.interp(ratio, ratio_col, yokoya_file["dipx"]) + yokyquad = np.interp(ratio, ratio_col, yokoya_file["quady"]) + yokxquad = np.interp(ratio, ratio_col, yokoya_file["quadx"]) + else: + yokxdip = np.interp(ratio, ratio_col, yokoya_file["dipy"]) + yokydip = np.interp(ratio, ratio_col, yokoya_file["dipx"]) + yokxquad = np.interp(ratio, ratio_col, yokoya_file["quady"]) + yokyquad = np.interp(ratio, ratio_col, yokoya_file["quadx"]) + + return (yoklong, yokxdip, yokydip, yokxquad, yokyquad) diff --git a/collective_effects/wakefield.py b/collective_effects/wakefield.py index dea656d1b3bb4ac92c4c312f2bdb877a981a2d06..e604d9082a8c8ac8b6de3b9ad5b4f2ec80e49c24 100644 --- a/collective_effects/wakefield.py +++ b/collective_effects/wakefield.py @@ -11,7 +11,12 @@ import warnings import re import pandas as pd import numpy as np +import matplotlib.pyplot as plt +import collective_effects.tools as tools from scipy.interpolate import interp1d +from tracking.element import Element +from copy import deepcopy +from scipy.integrate import trapz class ComplexData: """ @@ -34,7 +39,8 @@ class ComplexData: index=variable) self.data.index.name = 'variable' - def add(self, structure_to_add, method='zero', interp_kind='cubic'): + def add(self, structure_to_add, method='zero', interp_kind='cubic', + index_name="variable"): """ Method to add two structures. If the data don't have the same length, different cases are handled. @@ -55,6 +61,8 @@ class ComplexData: interp_kind : str, optional interpolation method which is passed to pandas and to scipy.interpolate.interp1d. + index_name : str, optional + name of the dataframe index passed to the method Returns ------- @@ -66,21 +74,22 @@ class ComplexData: # from the two input data if isinstance(structure_to_add, (int, float, complex)): - structure_to_add = ComplexData( - variable=self.data.index, - function=structure_to_add * np.ones(len(self.data.index))) - - initial_data = pd.concat( - [self.data, - structure_to_add.data.index.to_frame().set_index('variable')], - sort=True) + structure_to_add = ComplexData(variable=self.data.index, + function=(structure_to_add * + np.ones(len(self.data.index)))) + + data_to_concat = structure_to_add.data.index.to_frame().set_index(index_name) + + initial_data = pd.concat([self.data, data_to_concat], sort=True) initial_data = initial_data[~initial_data.index.duplicated(keep='first')] + initial_data = initial_data.sort_index() data_to_add = pd.concat( [structure_to_add.data, - self.data.index.to_frame().set_index('variable')], + self.data.index.to_frame().set_index(index_name)], sort=True) data_to_add = data_to_add[~data_to_add.index.duplicated(keep='first')] + data_to_add = data_to_add.sort_index() if method == 'common': max_variable = min(structure_to_add.data.index.max(), @@ -249,12 +258,15 @@ class Impedance(ComplexData): compared so that the addition is self-consistent. Beta functions can be precised as well. """ + if isinstance(structure_to_add, (int, float, complex)): - result = super().add(structure_to_add, method=method) + result = super().add(structure_to_add, method=method, + index_name="frequency [Hz]") elif isinstance(structure_to_add, Impedance): if (self.impedance_type == structure_to_add.impedance_type): weight = (beta_x ** self.power_x) * (beta_y ** self.power_y) - result = super().add(structure_to_add * weight, method=method) + result = super().add(structure_to_add * weight, method=method, + index_name="frequency [Hz]") else: warnings.warn(('The two Impedance objects do not have the ' 'same coordinates or plane or type. ' @@ -317,17 +329,45 @@ class Impedance(ComplexData): class Wakefield: """ - Define a general Wakefield machine element. - It contains Imepdance or WakeFunction objects which defines the wakefield - and other informations: beta functions. + Defines a Wakefield which corresponds to a single physical element which + produces different types of wakes, represented by their Impedance or + WakeFunction objects. + + Parameters + ---------- + structure_list : list, optional + list of Impedance/WakeFunction objects to add to the Wakefield. + name : str, optional + + Attributes + ---------- + impedance_components : array of str + Impedance component names present for this element. + + Methods + ------- + append_to_model(structure_to_add) + Add Impedance/WakeFunction to Wakefield. + list_to_attr(structure_list) + Add list of Impedance/WakeFunction to Wakefield. + add_wakefields(wake1, beta1, wake2, beta2) + Add two Wakefields taking into account beta functions. + add_several_wakefields(wakefields, beta) + Add a list of Wakefields taking into account beta functions. """ - def __init__(self, betax=1, betay=1, structure_list=[]): - self.betax = betax - self.betay = betay + def __init__(self, structure_list=None, name=None): self.list_to_attr(structure_list) + self.name = name def append_to_model(self, structure_to_add): + """ + Add Impedance/WakeFunction component to Wakefield. + + Parameters + ---------- + structure_to_add : Impedance or WakeFunction object + """ list_of_attributes = dir(self) if isinstance(structure_to_add, Impedance): attribute_name = "Z" + structure_to_add.impedance_type @@ -347,79 +387,353 @@ class Wakefield: raise ValueError("{} is not an Impedance nor a WakeFunction.".format(structure_to_add)) def list_to_attr(self, structure_list): - for component in structure_list: - self.append_to_model(component) + """ + Add list of Impedance/WakeFunction components to Wakefield. + + Parameters + ---------- + structure_list : list of Impedance or WakeFunction objects. + """ + if structure_list is not None: + for component in structure_list: + self.append_to_model(component) @property - def list_impedance_components(self): + def impedance_components(self): """ - Return the list of impedance components for the element, - based on the attributes names. + Return an array of the impedance component names for the element. """ return np.array([comp for comp in dir(self) if re.match(r'[Z]', comp)]) - def add(self, element_to_add): + @staticmethod + def add_wakefields(wake1, beta1, wake2, beta2): + """ + Add two Wakefields taking into account beta functions. + + Parameters + ---------- + wake1 : Wakefield + Wakefield to add. + beta1 : array of shape (2,) + Beta functions at wake1 postion. + wake2 : Wakefield + Wakefield to add. + beta2 : array of shape (2,) + Beta functions at wake2 postion. + + Returns + ------- + wake_sum : Wakefield + Wakefield with summed components. + """ - This function adds two Wakefield objects. - The different impedance components are weighted and added. - The result is a Wakefield object with all the components but the - beta function equal to 1. + wake_sum = deepcopy(wake1) + for component_name1 in wake1.impedance_components: + impedance1 = getattr(wake_sum, component_name1) + weight1 = ((beta1[0] ** impedance1.power_x) * + (beta1[1] ** impedance1.power_y)) + setattr(wake_sum, component_name1, weight1*impedance1) + + for component_name2 in wake2.impedance_components: + impedance2 = getattr(wake2, component_name2) + weight2 = ((beta2[0] ** impedance2.power_x) * + (beta2[1] ** impedance2.power_y)) + try: + impedance1 = getattr(wake_sum, component_name2) + setattr(wake_sum, component_name2, impedance1 + + weight2*impedance2) + except AttributeError: + setattr(wake_sum, component_name2, weight2*impedance2) + + return wake_sum + + @staticmethod + def add_several_wakefields(wakefields, beta): """ + Add a list of Wakefields taking into account beta functions. - if not isinstance(element_to_add, Wakefield): - raise TypeError("You can only add Wakefield objects to other" - "Wakefields objects.") + Parameters + ---------- + wakefields : list of Wakefield + Wakefields to add. + beta : array of shape (len(wakefields), 2) + Beta functions at the Wakefield postions. + + Returns + ------- + wake_sum : Wakefield + Wakefield with summed components.. + + """ + if len(wakefields) == 1: + return wakefields[0] + elif len(wakefields) > 1: + wake_sum = Wakefield.add_wakefields(wakefields[0], beta[:,0], + wakefields[1], beta[:,1]) + for i in range(len(wakefields) - 2): + wake_sum = Wakefield.add_wakefields(wake_sum, [1 ,1], + wakefields[i+2], beta[:,i+2]) + return wake_sum + else: + raise ValueError("Error in input.") - list_of_summed_impedances = [] - list_of_components_added = [] + +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) - betax1 = self.betax - betay1 = self.betay - betax2 = element_to_add.betax - betay2 = element_to_add.betay + def track(self, beam): + """ + Track a beam object through this Element. - # We first go through element1 impedance components and add them - # to the element2 components. If the component is missing in element2, - # we then simply weight the element1 component and add it to the model. - for component_name in self.list_impedance_components: + Parameters + ---------- + beam : Beam object + """ + raise NotImplementedError - element1_impedance = self.__getattribute__(component_name) + 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. - element1_weight = ((betax1 ** element1_impedance.power_x) - * (betay1 ** element1_impedance.power_y)) + 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) - try: - element2_impedance = element_to_add.__getattribute__(component_name) - list_of_components_added.append(component_name) - element12_impedance_sum = element1_weight * element1_impedance - element12_impedance_sum = element12_impedance_sum.add(element2_impedance, betax2, betay2) - except: - element12_impedance_sum = element1_weight * element1_impedance + 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): + """ + 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. + + """ + if attr_list is None: + attr_list = self.sum_names[self.sum_names != "sum"] - list_of_summed_impedances.append(element12_impedance_sum) + # sort plot by decresing area size + area = np.zeros((len(attr_list),)) + for index, attr in enumerate(attr_list): + sum_imp = getattr(getattr(self, attr), Z_type) + area[index] = trapz(sum_imp.data[component], sum_imp.data.index) + sorted_index = area.argsort()[::-1] - # We now go through the components which are unique to element2 and - # add them to the impedance model, with the beta function weighting - missing_components_list = list(set(element_to_add.list_impedance_components) - - set(list_of_components_added)) - - for component_name in missing_components_list: - - element2_impedance = element_to_add.__getattribute__(component_name) - element2_weight = ((betax2 ** element2_impedance.power_x) - * (betay2 ** element2_impedance.power_y)) - element12_impedance_sum = element2_weight * element2_impedance + # 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 + sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance + ax.fill_between(sum_imp.data.index, total_imp, + total_imp + sum_imp.data[component]) + total_imp += sum_imp.data[component] + if attr_list is None: + legend.append(attr[4:]) + else: + legend.append(attr) - list_of_summed_impedances.append(element12_impedance_sum) + if sigma is not None: + spect = tools.spectral_density(zero_impedance.data.index, sigma) + spect = spect/spect.max()*total_imp.max() + ax.plot(sum_imp.data.index, spect, 'r', linewidth=2.5) + + ax.legend(legend) + ax.set_xlabel("Frequency [Hz]") - # Gather everything in an Wakefield object - sum_wakefield = Wakefield(betax=1., betay=1., - structure_list=list_of_summed_impedances) + if Z_type == "Zlong": + ax.set_ylabel(Z_type + " [Ohm] - " + component + " part") + else: + ax.set_ylabel(Z_type + " [Ohm/m] - " + component + " part") - return sum_wakefield + def summary(self, sigma, attr_list=None, list_components=None): - def __radd__(self, structure_to_add): - return self.add(structure_to_add) - - def __add__(self, structure_to_add): - return self.add(structure_to_add) + if attr_list is None: + attr_list = self.sum_names + if list_components is None: + list_components = self.sum.impedance_components + + shape_array = (len(list_components), len(attr_list)) + loss_array = np.zeros(shape_array) + + for i, attr in enumerate(attr_list): + for j, component_name in enumerate(list_components): + try: + impedance = getattr(getattr(self, attr), component_name) + loss_array[j,i] = tools.loss_factor(impedance, sigma) + except AttributeError: + pass + + summary = pd.DataFrame(loss_array.T, index=attr_list, + columns=list_components) + return summary + \ No newline at end of file diff --git a/machines/soleil.py b/machines/soleil.py index d2fc712fab174388805c9711e5d4c11e1a4994f2..c8ac723c058f7c2d1860dc7859939cd7b0adbf55 100644 --- a/machines/soleil.py +++ b/machines/soleil.py @@ -32,30 +32,32 @@ def soleil(mode = 'Uniform'): # mean values beta = np.array([3, 1.3]) alpha = np.array([0, 0]) - disp = np.array([0, 0]) - dispp = np.array([0, 0]) - mean_val = Optics(beta, alpha, disp, dispp) - - ring = Synchrotron(h, L, E0, particle, ac=ac, U0=U0, tau=tau, - mean_optics=mean_val, emit=emit, tune=tune, - sigma_delta=sigma_delta, sigma_0=sigma_0, - chro=chro) - - if mode == 'Uniform': - beam = Beam(ring, ) - elif mode == 'Hybrid': - pass - elif mode == 'Low alpha/25': - pass - elif mode == 'Low alpha/25': - pass - elif mode == '8 bunches': - pass - elif mode == 'Single': - pass - else: - raise ValueError("{} is not a correct operation mode.".format(mode)) + dispersion = np.array([0, 0, 0, 0]) + optics = Optics(local_beta=beta, local_alpha=alpha, + local_dispersion=dispersion) + ring = Synchrotron(h, optics, particle, L=L, E0=E0, ac=ac, U0=U0, tau=tau, + emit=emit, tune=tune, sigma_delta=sigma_delta, + sigma_0=sigma_0, chro=chro) return ring +def v0313(): + """ + + """ + + h = 416 + particle = Electron() + tau = np.array([7.3e-3, 13.1e-3, 11.7e-3]) + emit = np.array([53.47e-12, 53.47e-12]) + sigma_0 = 9.2e-12 + sigma_delta = 9e-4 + + lattice_file = "SOLEIL_U_74BA_HOA_SYM02_V0313_cleaned.mat" + optics = Optics(lattice_file=lattice_file) + + ring = Synchrotron(h, optics, particle, tau=tau, emit=emit, + sigma_0=sigma_0, sigma_delta=sigma_delta) + + return ring \ No newline at end of file diff --git a/tracking/element.py b/tracking/element.py index c7811a7c9a39e894479e76f9069595cd4a8914e5..bee031ce02323679e53832f4745196ef33dc1666 100644 --- a/tracking/element.py +++ b/tracking/element.py @@ -153,11 +153,10 @@ class TransverseMap(Element): def __init__(self, ring): self.ring = ring - self.alpha = self.ring.mean_optics.alpha - self.beta = self.ring.mean_optics.beta - self.gamma = self.ring.mean_optics.gamma - self.disp = self.ring.mean_optics.disp - self.dispp = self.ring.mean_optics.dispp + self.alpha = self.ring.optics.local_alpha + self.beta = self.ring.optics.local_beta + self.gamma = self.ring.optics.local_gamma + self.dispersion = self.ring.optics.local_dispersion self.phase_advance = self.ring.tune[0:2]*2*np.pi @Element.parallel @@ -182,19 +181,19 @@ class TransverseMap(Element): # Horizontal matrix[0,0,:] = np.cos(phase_advance_x) + self.alpha[0]*np.sin(phase_advance_x) matrix[0,1,:] = self.beta[0]*np.sin(phase_advance_x) - matrix[0,2,:] = self.disp[0] + matrix[0,2,:] = self.dispersion[0] matrix[1,0,:] = -1*self.gamma[0]*np.sin(phase_advance_x) matrix[1,1,:] = np.cos(phase_advance_x) - self.alpha[0]*np.sin(phase_advance_x) - matrix[1,2,:] = self.dispp[0] + matrix[1,2,:] = self.dispersion[1] matrix[2,2,:] = 1 # Vertical matrix[3,3,:] = np.cos(phase_advance_y) + self.alpha[1]*np.sin(phase_advance_y) matrix[3,4,:] = self.beta[1]*np.sin(phase_advance_y) - matrix[3,5,:] = self.disp[1] + matrix[3,5,:] = self.dispersion[2] matrix[4,3,:] = -1*self.gamma[1]*np.sin(phase_advance_y) matrix[4,4,:] = np.cos(phase_advance_y) - self.alpha[1]*np.sin(phase_advance_y) - matrix[4,5,:] = self.dispp[1] + matrix[4,5,:] = self.dispersion[3] matrix[5,5,:] = 1 x = matrix[0,0,:]*bunch["x"] + matrix[0,1,:]*bunch["xp"] + matrix[0,2,:]*bunch["delta"] diff --git a/tracking/monitors.py b/tracking/monitors.py new file mode 100644 index 0000000000000000000000000000000000000000..c33db1f1f6fdb0d395da5b1a7bebf02c4950ed6f --- /dev/null +++ b/tracking/monitors.py @@ -0,0 +1,467 @@ +# -*- coding: utf-8 -*- +""" +This module defines the different monitor class which are used to save data +during tracking. + +@author: Alexis Gamelin +@date: 24/03/2020 + +""" + +import numpy as np +import h5py as hp +from tracking.element import Element +from tracking.particles import Bunch, Beam +from abc import ABCMeta +# from mpi4py import MPI + +class Monitor(Element, metaclass=ABCMeta): + """ + Abstract Monitor class used for subclass inheritance to define all the + different kind of monitors objects. + + The Monitor class is based on h5py module to be able to write data on + structured binary files. The class provides a common file where the + different Monitor subclass can write. + + Attributes + ---------- + file : HDF5 file + Common file where all monitors, Monitor subclass elements, write the + saved data. Based on class attribute _file_storage. + file_name : string + Name of the HDF5 file where the data is stored. Based on class + attribute _file_name_storage. + + Methods + ------- + monitor_init(group_name, save_every, buffer_size, total_size, + dict_buffer, dict_file, file_name=None, mpi_mode=True) + Method called to initialize Monitor subclass. + write() + Write data from buffer to the HDF5 file. + to_buffer(object_to_save) + Save data to buffer. + close() + Close the HDF5 file shared by all Monitor subclass, must be called + by at least an instance of a Montior subclass at the end of the + tracking. + """ + + _file_name_storage = [] + _file_storage = [] + + @property + def file_name(self): + """Common file where all monitors, Monitor subclass elements, write the + saved data.""" + try: + return self._file_name_storage[0] + except IndexError: + print("The HDF5 file name for monitors is not set.") + raise ValueError + + @property + def file(self): + """Name of the HDF5 file where the data is stored.""" + try: + return self._file_storage[0] + except IndexError: + print("The HDF5 file to store data is not set.") + raise ValueError + + def monitor_init(self, group_name, save_every, buffer_size, total_size, + dict_buffer, dict_file, file_name=None, mpi_mode=True): + """ + Method called to initialize Monitor subclass. + + Parameters + ---------- + group_name : string + Name of the HDF5 group in which the data for the current monitor + will be saved. + save_every : int or float + Set the frequency of the save. The data is saved every save_every + call of the montior. + buffer_size : int or float + Size of the save buffer. + total_size : int or float + Total size of the save. The following relationships between the + parameters must exist: + total_size % buffer_size == 0 + number of call to track / save_every == total_size + dict_buffer : dict + Dictionary with keys as the attribute name to save and values as + the shape of the buffer to create to hold the attribute, like + (key.shape, buffer_size) + dict_file : dict + Dictionary with keys as the attribute name to save and values as + the shape of the dataset to create to hold the attribute, like + (key.shape, total_size) + file_name : string, optional + Name of the HDF5 where the data will be stored. Must be specified + the first time a subclass of Monitor is instancied and must be None + the following times. + mpi_mode : bool, optional + If True, open the HDF5 file in parallel mode, which is needed to + allow several cores to write in the same file at the same time. + If False, open the HDF5 file in standard mode. + """ + + # setup and open common file for all monitors + if file_name is not None: + if len(self._file_name_storage) == 0: + self._file_name_storage.append(file_name + ".hdf5") + if len(self._file_storage) == 0: + if mpi_mode == True: + f = hp.File(self.file_name, "w", libver='latest', + driver='mpio', comm=MPI.COMM_WORLD) + else: + f = hp.File(self.file_name, "w", libver='latest') + self._file_storage.append(f) + else: + raise ValueError("File is already open.") + else: + raise ValueError("File name for monitors is already attributed.") + + self.group_name = group_name + self.save_every = int(save_every) + self.total_size = int(total_size) + self.buffer_size = int(buffer_size) + if total_size % buffer_size != 0: + raise ValueError("total_size must be divisible by buffer_size.") + self.buffer_count = 0 + self.write_count = 0 + self.track_count = 0 + + # setup attribute buffers from values given in dict_buffer + for key, value in dict_buffer.items(): + self.__setattr__(key,np.zeros(value)) + self.time = np.zeros((self.buffer_size,), dtype=int) + + # create HDF5 groups and datasets to save data from group_name and + # dict_file + self.g = self.file.require_group(self.group_name) + self.g.require_dataset("time", (self.total_size,), dtype=int) + for key, value in dict_file.items(): + self.g.require_dataset(key, value, dtype=float) + + # create a dictionary which + slice_dict = {} + for key, value in dict_file.items(): + slice_dict[key] = [] + for i in range(len(value)-1): + slice_dict[key].append(slice(None)) + self.slice_dict = slice_dict + + def write(self): + """Write data from buffer to the HDF5 file.""" + + self.file[self.group_name]["time"][self.write_count*self.buffer_size:( + self.write_count+1)*self.buffer_size] = self.time + for key, value in self.dict_buffer.items(): + slice_list = list(self.slice_dict[key]) + slice_list.append(slice(self.write_count*self.buffer_size, + (self.write_count+1)*self.buffer_size)) + slice_tuple = tuple(slice_list) + self.file[self.group_name][key][slice_tuple] = self.__getattribute__(key) + + self.file.flush() + self.write_count += 1 + + def to_buffer(self, object_to_save): + """ + Save data to buffer. + + Parameters + ---------- + object_to_save : python object + Depends on the Monitor subclass, typically a Beam or Bunch object. + """ + self.time[self.buffer_count] = self.track_count + for key, value in self.dict_buffer.items(): + slice_list = list(self.slice_dict[key]) + slice_list.append(self.buffer_count) + slice_tuple = tuple(slice_list) + self.__getattribute__(key)[slice_tuple] = object_to_save.__getattribute__(key) + self.buffer_count += 1 + + if self.buffer_count == self.buffer_size: + self.write() + self.buffer_count = 0 + + def close(self): + """ + Close the HDF5 file shared by all Monitor subclass, must be called + by at least an instance of a Montior subclass at the end of the + tracking. + """ + try: + self.file.close() + except ValueError: + pass + +class BunchMonitor(Monitor): + """ + Monitor a single bunch and save attributes (mean, std, emit and current). + + Parameters + ---------- + bunch_number : int + Bunch to monitor + file_name : string, optional + Name of the HDF5 where the data will be stored. Must be specified + the first time a subclass of Monitor is instancied and must be None + the following times. + save_every : int or float, optional + Set the frequency of the save. The data is saved every save_every + call of the montior. + buffer_size : int or float, optional + Size of the save buffer. + total_size : int or float, optional + Total size of the save. The following relationships between the + parameters must exist: + total_size % buffer_size == 0 + number of call to track / save_every == total_size + mpi_mode : bool, optional + If True, open the HDF5 file in parallel mode, which is needed to + allow several cores to write in the same file at the same time. + If False, open the HDF5 file in standard mode. + + Methods + ------- + track(object_to_save) + Save data + """ + + def __init__(self, bunch_number, file_name=None, save_every=5, + buffer_size=500, total_size=2e4, mpi_mode=True): + + self.bunch_number = bunch_number + group_name = "BunchData_" + str(self.bunch_number) + dict_buffer = {"mean":(6, buffer_size), "std":(6, buffer_size), + "emit":(3, buffer_size), "current":(buffer_size,)} + dict_file = {"mean":(6, total_size), "std":(6, total_size), + "emit":(3, total_size), "current":(total_size,)} + self.monitor_init(group_name, save_every, buffer_size, total_size, + dict_buffer, dict_file, file_name, mpi_mode) + + self.dict_buffer = dict_buffer + self.dict_file = dict_file + + def track(self, object_to_save): + """ + Save data + + Parameters + ---------- + object_to_save : Bunch or Beam object + """ + if self.track_count % self.save_every == 0: + if isinstance(object_to_save, Beam): + if (object_to_save.mpi_switch == True): + if object_to_save.mpi.bunch_num == self.bunch_number: + self.to_buffer(object_to_save[object_to_save.mpi.bunch_num]) + else: + self.to_buffer(object_to_save[self.bunch_number]) + elif isinstance(object_to_save, Bunch): + self.to_buffer(object_to_save) + else: + raise TypeError("object_to_save should be a Beam or Bunch object.") + self.track_count += 1 + + +class PhaseSpaceMonitor(Monitor): + """ + Monitor a single bunch and save the full phase space. + + Parameters + ---------- + bunch_number : int + Bunch to monitor + mp_number : int or float + Number of macroparticle in the phase space to save. + file_name : string, optional + Name of the HDF5 where the data will be stored. Must be specified + the first time a subclass of Monitor is instancied and must be None + the following times. + save_every : int or float, optional + Set the frequency of the save. The data is saved every save_every + call of the montior. + buffer_size : int or float, optional + Size of the save buffer. + total_size : int or float, optional + Total size of the save. The following relationships between the + parameters must exist: + total_size % buffer_size == 0 + number of call to track / save_every == total_size + mpi_mode : bool, optional + If True, open the HDF5 file in parallel mode, which is needed to + allow several cores to write in the same file at the same time. + If False, open the HDF5 file in standard mode. + + Methods + ------- + track(object_to_save) + Save data + """ + + def __init__(self, bunch_number, mp_number, file_name=None, save_every=1e3, + buffer_size=10, total_size=100, mpi_mode=True): + + self.bunch_number = bunch_number + self.mp_number = int(mp_number) + group_name = "PhaseSpaceData_" + str(self.bunch_number) + dict_buffer = {"particles":(self.mp_number, 6, buffer_size), + "alive":(self.mp_number, buffer_size)} + dict_file = {"particles":(self.mp_number, 6, total_size), + "alive":(self.mp_number, total_size)} + self.monitor_init(group_name, save_every, buffer_size, total_size, + dict_buffer, dict_file, file_name, mpi_mode) + + self.dict_buffer = dict_buffer + self.dict_file = dict_file + + def track(self, object_to_save): + """ + Save data + + Parameters + ---------- + object_to_save : Bunch or Beam object + """ + + if self.track_count % self.save_every == 0: + if isinstance(object_to_save, Beam): + if (object_to_save.mpi_switch == True): + if object_to_save.mpi.bunch_num == self.bunch_number: + self.to_buffer(object_to_save[object_to_save.mpi.bunch_num]) + else: + self.to_buffer(object_to_save[self.bunch_number]) + elif isinstance(object_to_save, Bunch): + self.to_buffer(object_to_save) + else: + raise TypeError("object_to_save should be a Beam or Bunch object.") + self.track_count += 1 + + +class BeamMonitor(Monitor): + """ + Monitor the full beam and save each bunch attributes (mean, std, emit and + current). + + Parameters + ---------- + file_name : string, optional + Name of the HDF5 where the data will be stored. Must be specified + the first time a subclass of Monitor is instancied and must be None + the following times. + save_every : int or float, optional + Set the frequency of the save. The data is saved every save_every + call of the montior. + buffer_size : int or float, optional + Size of the save buffer. + total_size : int or float, optional + Total size of the save. The following relationships between the + parameters must exist: + total_size % buffer_size == 0 + number of call to track / save_every == total_size + mpi_mode : bool, optional + If True, open the HDF5 file in parallel mode, which is needed to + allow several cores to write in the same file at the same time. + If False, open the HDF5 file in standard mode. + + Methods + ------- + track(beam) + Save data + """ + + def __init__(self, file_name=None, save_every=5, buffer_size=500, + total_size=2e4, mpi_mode=True): + + group_name = "Beam" + dict_buffer = {} + dict_file = {} + self.monitor_init(group_name, save_every, buffer_size, total_size, + dict_buffer, dict_file, file_name, mpi_mode) + + self.mean = np.zeros((6, self.buffer_size), dtype=float) + self.std = np.zeros((6, self.buffer_size), dtype=float) + self.emit = np.zeros((3, self.buffer_size), dtype=float) + self.current = np.zeros((self.buffer_size,), dtype=float) + + self.g.require_dataset("mean", (6, 416, self.total_size,), dtype=float) + self.g.require_dataset("std", (6, 416, self.total_size,), dtype=float) + self.g.require_dataset("emit", (3, 416, self.total_size,), dtype=float) + self.g.require_dataset("current", (416, self.total_size,), dtype=float) + + def track(self, beam): + """ + Save data + + Parameters + ---------- + beam : Beam object + """ + if self.track_count % self.save_every == 0: + if (beam.mpi_switch == True): + self.to_buffer(beam[beam.mpi.bunch_num], beam.mpi.bunch_num) + else: + raise NotImplementedError + + self.track_count += 1 + + def to_buffer(self, bunch, bunch_num): + """ + Save data to buffer. + + Parameters + ---------- + bunch : Bunch object + bunch_num : int + """ + + self.time[self.buffer_count] = self.track_count + self.mean[:, self.buffer_count] = bunch.mean + self.std[:, self.buffer_count] = bunch.std + self.emit[:, self.buffer_count] = bunch.emit + self.current[self.buffer_count] = bunch.current + + self.buffer_count += 1 + + if self.buffer_count == self.buffer_size: + self.write(bunch_num) + self.buffer_count = 0 + + def write(self, bunch_num): + """ + Write data from buffer to the HDF5 file. + + Parameters + ---------- + bunch_num : int + """ + self.file[self.group_name]["time"][self.write_count*self.buffer_size:( + self.write_count+1)*self.buffer_size] = self.time + + self.file[self.group_name]["mean"][:, bunch_num, + self.write_count*self.buffer_size:(self.write_count+1) * + self.buffer_size] = self.mean + + self.file[self.group_name]["std"][:, bunch_num, + self.write_count*self.buffer_size:(self.write_count+1) * + self.buffer_size] = self.std + + self.file[self.group_name]["emit"][:, bunch_num, + self.write_count*self.buffer_size:(self.write_count+1) * + self.buffer_size] = self.emit + + self.file[self.group_name]["current"][bunch_num, + self.write_count*self.buffer_size:(self.write_count+1) * + self.buffer_size] = self.current + + self.file.flush() + self.write_count += 1 + + \ No newline at end of file diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c67145e06e038f4a48ceda33f888dc60ab733eea --- /dev/null +++ b/tracking/monitors/__init__.py @@ -0,0 +1,7 @@ +# -*- coding: utf-8 -*- +""" +Created on Tue Jan 14 18:11:33 2020 + +@author: gamelina +""" + diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py new file mode 100644 index 0000000000000000000000000000000000000000..af79efa10ad7cdb73e4fd78c23641c79e88845ea --- /dev/null +++ b/tracking/monitors/monitors.py @@ -0,0 +1,467 @@ +# -*- coding: utf-8 -*- +""" +This module defines the different monitor class which are used to save data +during tracking. + +@author: Alexis Gamelin +@date: 24/03/2020 + +""" + +import numpy as np +import h5py as hp +from tracking.element import Element +from tracking.particles import Bunch, Beam +from abc import ABCMeta +from mpi4py import MPI + +class Monitor(Element, metaclass=ABCMeta): + """ + Abstract Monitor class used for subclass inheritance to define all the + different kind of monitors objects. + + The Monitor class is based on h5py module to be able to write data on + structured binary files. The class provides a common file where the + different Monitor subclass can write. + + Attributes + ---------- + file : HDF5 file + Common file where all monitors, Monitor subclass elements, write the + saved data. Based on class attribute _file_storage. + file_name : string + Name of the HDF5 file where the data is stored. Based on class + attribute _file_name_storage. + + Methods + ------- + monitor_init(group_name, save_every, buffer_size, total_size, + dict_buffer, dict_file, file_name=None, mpi_mode=True) + Method called to initialize Monitor subclass. + write() + Write data from buffer to the HDF5 file. + to_buffer(object_to_save) + Save data to buffer. + close() + Close the HDF5 file shared by all Monitor subclass, must be called + by at least an instance of a Montior subclass at the end of the + tracking. + """ + + _file_name_storage = [] + _file_storage = [] + + @property + def file_name(self): + """Common file where all monitors, Monitor subclass elements, write the + saved data.""" + try: + return self._file_name_storage[0] + except IndexError: + print("The HDF5 file name for monitors is not set.") + raise ValueError + + @property + def file(self): + """Name of the HDF5 file where the data is stored.""" + try: + return self._file_storage[0] + except IndexError: + print("The HDF5 file to store data is not set.") + raise ValueError + + def monitor_init(self, group_name, save_every, buffer_size, total_size, + dict_buffer, dict_file, file_name=None, mpi_mode=True): + """ + Method called to initialize Monitor subclass. + + Parameters + ---------- + group_name : string + Name of the HDF5 group in which the data for the current monitor + will be saved. + save_every : int or float + Set the frequency of the save. The data is saved every save_every + call of the montior. + buffer_size : int or float + Size of the save buffer. + total_size : int or float + Total size of the save. The following relationships between the + parameters must exist: + total_size % buffer_size == 0 + number of call to track / save_every == total_size + dict_buffer : dict + Dictionary with keys as the attribute name to save and values as + the shape of the buffer to create to hold the attribute, like + (key.shape, buffer_size) + dict_file : dict + Dictionary with keys as the attribute name to save and values as + the shape of the dataset to create to hold the attribute, like + (key.shape, total_size) + file_name : string, optional + Name of the HDF5 where the data will be stored. Must be specified + the first time a subclass of Monitor is instancied and must be None + the following times. + mpi_mode : bool, optional + If True, open the HDF5 file in parallel mode, which is needed to + allow several cores to write in the same file at the same time. + If False, open the HDF5 file in standard mode. + """ + + # setup and open common file for all monitors + if file_name is not None: + if len(self._file_name_storage) == 0: + self._file_name_storage.append(file_name + ".hdf5") + if len(self._file_storage) == 0: + if mpi_mode == True: + f = hp.File(self.file_name, "w", libver='latest', + driver='mpio', comm=MPI.COMM_WORLD) + else: + f = hp.File(self.file_name, "w", libver='latest') + self._file_storage.append(f) + else: + raise ValueError("File is already open.") + else: + raise ValueError("File name for monitors is already attributed.") + + self.group_name = group_name + self.save_every = int(save_every) + self.total_size = int(total_size) + self.buffer_size = int(buffer_size) + if total_size % buffer_size != 0: + raise ValueError("total_size must be divisible by buffer_size.") + self.buffer_count = 0 + self.write_count = 0 + self.track_count = 0 + + # setup attribute buffers from values given in dict_buffer + for key, value in dict_buffer.items(): + self.__setattr__(key,np.zeros(value)) + self.time = np.zeros((self.buffer_size,), dtype=int) + + # create HDF5 groups and datasets to save data from group_name and + # dict_file + self.g = self.file.require_group(self.group_name) + self.g.require_dataset("time", (self.total_size,), dtype=int) + for key, value in dict_file.items(): + self.g.require_dataset(key, value, dtype=float) + + # create a dictionary which + slice_dict = {} + for key, value in dict_file.items(): + slice_dict[key] = [] + for i in range(len(value)-1): + slice_dict[key].append(slice(None)) + self.slice_dict = slice_dict + + def write(self): + """Write data from buffer to the HDF5 file.""" + + self.file[self.group_name]["time"][self.write_count*self.buffer_size:( + self.write_count+1)*self.buffer_size] = self.time + for key, value in self.dict_buffer.items(): + slice_list = list(self.slice_dict[key]) + slice_list.append(slice(self.write_count*self.buffer_size, + (self.write_count+1)*self.buffer_size)) + slice_tuple = tuple(slice_list) + self.file[self.group_name][key][slice_tuple] = self.__getattribute__(key) + + self.file.flush() + self.write_count += 1 + + def to_buffer(self, object_to_save): + """ + Save data to buffer. + + Parameters + ---------- + object_to_save : python object + Depends on the Monitor subclass, typically a Beam or Bunch object. + """ + self.time[self.buffer_count] = self.track_count + for key, value in self.dict_buffer.items(): + slice_list = list(self.slice_dict[key]) + slice_list.append(self.buffer_count) + slice_tuple = tuple(slice_list) + self.__getattribute__(key)[slice_tuple] = object_to_save.__getattribute__(key) + self.buffer_count += 1 + + if self.buffer_count == self.buffer_size: + self.write() + self.buffer_count = 0 + + def close(self): + """ + Close the HDF5 file shared by all Monitor subclass, must be called + by at least an instance of a Montior subclass at the end of the + tracking. + """ + try: + self.file.close() + except ValueError: + pass + +class BunchMonitor(Monitor): + """ + Monitor a single bunch and save attributes (mean, std, emit and current). + + Parameters + ---------- + bunch_number : int + Bunch to monitor + file_name : string, optional + Name of the HDF5 where the data will be stored. Must be specified + the first time a subclass of Monitor is instancied and must be None + the following times. + save_every : int or float, optional + Set the frequency of the save. The data is saved every save_every + call of the montior. + buffer_size : int or float, optional + Size of the save buffer. + total_size : int or float, optional + Total size of the save. The following relationships between the + parameters must exist: + total_size % buffer_size == 0 + number of call to track / save_every == total_size + mpi_mode : bool, optional + If True, open the HDF5 file in parallel mode, which is needed to + allow several cores to write in the same file at the same time. + If False, open the HDF5 file in standard mode. + + Methods + ------- + track(object_to_save) + Save data + """ + + def __init__(self, bunch_number, file_name=None, save_every=5, + buffer_size=500, total_size=2e4, mpi_mode=True): + + self.bunch_number = bunch_number + group_name = "BunchData_" + str(self.bunch_number) + dict_buffer = {"mean":(6, buffer_size), "std":(6, buffer_size), + "emit":(3, buffer_size), "current":(buffer_size,)} + dict_file = {"mean":(6, total_size), "std":(6, total_size), + "emit":(3, total_size), "current":(total_size,)} + self.monitor_init(group_name, save_every, buffer_size, total_size, + dict_buffer, dict_file, file_name, mpi_mode) + + self.dict_buffer = dict_buffer + self.dict_file = dict_file + + def track(self, object_to_save): + """ + Save data + + Parameters + ---------- + object_to_save : Bunch or Beam object + """ + if self.track_count % self.save_every == 0: + if isinstance(object_to_save, Beam): + if (object_to_save.mpi_switch == True): + if object_to_save.mpi.bunch_num == self.bunch_number: + self.to_buffer(object_to_save[object_to_save.mpi.bunch_num]) + else: + self.to_buffer(object_to_save[self.bunch_number]) + elif isinstance(object_to_save, Bunch): + self.to_buffer(object_to_save) + else: + raise TypeError("object_to_save should be a Beam or Bunch object.") + self.track_count += 1 + + +class PhaseSpaceMonitor(Monitor): + """ + Monitor a single bunch and save the full phase space. + + Parameters + ---------- + bunch_number : int + Bunch to monitor + mp_number : int or float + Number of macroparticle in the phase space to save. + file_name : string, optional + Name of the HDF5 where the data will be stored. Must be specified + the first time a subclass of Monitor is instancied and must be None + the following times. + save_every : int or float, optional + Set the frequency of the save. The data is saved every save_every + call of the montior. + buffer_size : int or float, optional + Size of the save buffer. + total_size : int or float, optional + Total size of the save. The following relationships between the + parameters must exist: + total_size % buffer_size == 0 + number of call to track / save_every == total_size + mpi_mode : bool, optional + If True, open the HDF5 file in parallel mode, which is needed to + allow several cores to write in the same file at the same time. + If False, open the HDF5 file in standard mode. + + Methods + ------- + track(object_to_save) + Save data + """ + + def __init__(self, bunch_number, mp_number, file_name=None, save_every=1e3, + buffer_size=10, total_size=100, mpi_mode=True): + + self.bunch_number = bunch_number + self.mp_number = int(mp_number) + group_name = "PhaseSpaceData_" + str(self.bunch_number) + dict_buffer = {"particles":(self.mp_number, 6, buffer_size), + "alive":(self.mp_number, buffer_size)} + dict_file = {"particles":(self.mp_number, 6, total_size), + "alive":(self.mp_number, total_size)} + self.monitor_init(group_name, save_every, buffer_size, total_size, + dict_buffer, dict_file, file_name, mpi_mode) + + self.dict_buffer = dict_buffer + self.dict_file = dict_file + + def track(self, object_to_save): + """ + Save data + + Parameters + ---------- + object_to_save : Bunch or Beam object + """ + + if self.track_count % self.save_every == 0: + if isinstance(object_to_save, Beam): + if (object_to_save.mpi_switch == True): + if object_to_save.mpi.bunch_num == self.bunch_number: + self.to_buffer(object_to_save[object_to_save.mpi.bunch_num]) + else: + self.to_buffer(object_to_save[self.bunch_number]) + elif isinstance(object_to_save, Bunch): + self.to_buffer(object_to_save) + else: + raise TypeError("object_to_save should be a Beam or Bunch object.") + self.track_count += 1 + + +class BeamMonitor(Monitor): + """ + Monitor the full beam and save each bunch attributes (mean, std, emit and + current). + + Parameters + ---------- + file_name : string, optional + Name of the HDF5 where the data will be stored. Must be specified + the first time a subclass of Monitor is instancied and must be None + the following times. + save_every : int or float, optional + Set the frequency of the save. The data is saved every save_every + call of the montior. + buffer_size : int or float, optional + Size of the save buffer. + total_size : int or float, optional + Total size of the save. The following relationships between the + parameters must exist: + total_size % buffer_size == 0 + number of call to track / save_every == total_size + mpi_mode : bool, optional + If True, open the HDF5 file in parallel mode, which is needed to + allow several cores to write in the same file at the same time. + If False, open the HDF5 file in standard mode. + + Methods + ------- + track(beam) + Save data + """ + + def __init__(self, file_name=None, save_every=5, buffer_size=500, + total_size=2e4, mpi_mode=True): + + group_name = "Beam" + dict_buffer = {} + dict_file = {} + self.monitor_init(group_name, save_every, buffer_size, total_size, + dict_buffer, dict_file, file_name, mpi_mode) + + self.mean = np.zeros((6, self.buffer_size), dtype=float) + self.std = np.zeros((6, self.buffer_size), dtype=float) + self.emit = np.zeros((3, self.buffer_size), dtype=float) + self.current = np.zeros((self.buffer_size,), dtype=float) + + self.g.require_dataset("mean", (6, 416, self.total_size,), dtype=float) + self.g.require_dataset("std", (6, 416, self.total_size,), dtype=float) + self.g.require_dataset("emit", (3, 416, self.total_size,), dtype=float) + self.g.require_dataset("current", (416, self.total_size,), dtype=float) + + def track(self, beam): + """ + Save data + + Parameters + ---------- + beam : Beam object + """ + if self.track_count % self.save_every == 0: + if (beam.mpi_switch == True): + self.to_buffer(beam[beam.mpi.bunch_num], beam.mpi.bunch_num) + else: + raise NotImplementedError + + self.track_count += 1 + + def to_buffer(self, bunch, bunch_num): + """ + Save data to buffer. + + Parameters + ---------- + bunch : Bunch object + bunch_num : int + """ + + self.time[self.buffer_count] = self.track_count + self.mean[:, self.buffer_count] = bunch.mean + self.std[:, self.buffer_count] = bunch.std + self.emit[:, self.buffer_count] = bunch.emit + self.current[self.buffer_count] = bunch.current + + self.buffer_count += 1 + + if self.buffer_count == self.buffer_size: + self.write(bunch_num) + self.buffer_count = 0 + + def write(self, bunch_num): + """ + Write data from buffer to the HDF5 file. + + Parameters + ---------- + bunch_num : int + """ + self.file[self.group_name]["time"][self.write_count*self.buffer_size:( + self.write_count+1)*self.buffer_size] = self.time + + self.file[self.group_name]["mean"][:, bunch_num, + self.write_count*self.buffer_size:(self.write_count+1) * + self.buffer_size] = self.mean + + self.file[self.group_name]["std"][:, bunch_num, + self.write_count*self.buffer_size:(self.write_count+1) * + self.buffer_size] = self.std + + self.file[self.group_name]["emit"][:, bunch_num, + self.write_count*self.buffer_size:(self.write_count+1) * + self.buffer_size] = self.emit + + self.file[self.group_name]["current"][bunch_num, + self.write_count*self.buffer_size:(self.write_count+1) * + self.buffer_size] = self.current + + self.file.flush() + self.write_count += 1 + + \ No newline at end of file diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py new file mode 100644 index 0000000000000000000000000000000000000000..9ae1449936bf8c12e56b951c543bb48391a3e677 --- /dev/null +++ b/tracking/monitors/plotting.py @@ -0,0 +1,157 @@ +# -*- coding: utf-8 -*- +""" +Module for plotting the data recorded by the monitor module during the +tracking. + +@author: Watanyu Foosang +@Date: 10/04/2020 +""" + +import numpy as np +import matplotlib.pyplot as plt +import seaborn as sns +import h5py as hp + +def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"): + """ + Plot data recorded from a BunchMonitor. + + Parameters + ---------- + filename : str + Name of the HDF5 file that contains the data. + bunch_number : int + The bunch number whose data has been saved in the HDF5 file. + This has to be identical to 'bunch_number' parameter in 'BunchMonitor' object. + detaset : str {"current","emit","mean","std"} + HDF5 file's dataset to be plotted. + option : str, optional + If dataset is "emit", "mean", or "std", the variable name to be plotted + needs to be specified : + for "emit", option = {"x","y","s"} + for "mean" and "std", option = {"x","xp","y","yp","tau","delta"} + x_var : str, optional + Variable to be plotted on the horizontal axis. The default is "time". + + """ + + file = hp.File(filename, "r") + + group = "BunchData_{0}".format(bunch_number) # Data group of the HDF5 file + + if dataset == "current": + y_var = file[group][dataset][:]*1e3 + label = "current (mA)" + + elif dataset == "emit": + # Specifying the axis + # horizontal axis (x) -> 0 + # vertical axis (y) -> 1 + # longitudinal axis (s) -> 2 + + emit_axis = {"x":0, "y":1, "s":2} + + y_var = file[group][dataset][emit_axis[option]]*1e9 + + if option == "x": label = "hor. emittance (nm.rad)" + elif option == "y": label = "ver. emittance (nm.rad)" + elif option == "s": label = "long. emittance (nm.rad)" + + + elif dataset == "mean" or "std": + # Specify the variable + # x -> 0 + # xp -> 1 + # y -> 2 + # yp -> 3 + # tau -> 4 + # delta -> 5 + + var_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5} + scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1] + label_list = ["x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)", + "$\\tau$ (ps)", "$\\delta$"] + + y_var = file[group][dataset][var_dict[option]]*scale[var_dict[option]] + label = label_list[var_dict[option]] + + + plt.plot(file[group]["time"][:],y_var) + plt.xlabel("number of turns") + plt.ylabel(label) + + file.close() + +def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn, + only_alive=True): + """ + Plot data recorded from a PhaseSpaceMonitor. + + Parameters + ---------- + filename : str + Name of the HDF5 file that contains the data. + bunch_number : int + Number of the bunch whose data has been saved in the HDF5 file. + This has to be identical to 'bunch_number' parameter in + 'PhaseSpaceMonitor' object. + x_var, y_var : str {"x", "xp", "y", "yp", "tau", "delta"} + If dataset is "particles", the variables to be plotted on the + horizontal and the vertical axes need to be specified. + turn : int + Turn at which the data will be extracted. + only_alive : bool, optional + When only_alive is True, only alive particles are plotted and dead + particles will be discarded. + """ + + file = hp.File(filename, "r") + + group = "PhaseSpaceData_{0}".format(bunch_number) + dataset = "particles" + + # Specify the parameter + # x -> 0 + # xp -> 1 + # y -> 2 + # yp -> 3 + # tau -> 4 + # delta -> 5 + + var_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5} + scale = [1e3,1e3,1e3,1e3,1e12,1] + label = ["x (mm)","x' (mrad)","y (mm)","y' (mrad)","$\\tau$ (ps)", + "$\\delta$"] + + # find the index of "turn" in time array + turn_index = np.where(file["PhaseSpaceData_0"]["time"][:]==turn) + + if len(turn_index[0]) == 0: + raise ValueError("Turn {0} is not found. Enter turn from {1}.". + format(turn, file[group]["time"][:])) + + path = file[group][dataset] + + if only_alive is False: + # format : sns.jointplot(x_axis, yaxis, kind) + x_axis = path[:,var_dict[x_var],turn_index[0][0]] + y_axis = path[:,var_dict[y_var],turn_index[0][0]] + + elif only_alive is True: + alive_index = np.where(file[group]["alive"][:,turn_index]) + + x_axis = path[alive_index[0],var_dict[x_var],turn_index[0][0]] + y_axis = path[alive_index[0],var_dict[y_var],turn_index[0][0]] + + sns.jointplot(x_axis*scale[var_dict[x_var]], + y_axis*scale[var_dict[y_var]], kind="kde") + + plt.xlabel(label[var_dict[x_var]]) + plt.ylabel(label[var_dict[y_var]]) + + file.close() + + + + + diff --git a/tracking/optics.py b/tracking/optics.py index 39b8fba50ac50932b9ec092c782c974151c4f854..2e3483657181ebab0073f5d9eb135ac7808d07dd 100644 --- a/tracking/optics.py +++ b/tracking/optics.py @@ -1,18 +1,516 @@ # -*- coding: utf-8 -*- """ -Created on Wed Mar 11 18:00:28 2020 +Module where the class to store the optic functions and the lattice physical +parameters are defined. -@author: gamelina +@author: Alexis Gamelin +@date: 10/04/2020 """ +import at +import numpy as np +import matplotlib.pyplot as plt +from scipy.interpolate import interp1d + + class Optics: - """Handle optic functions""" - def __init__(self, beta, alpha, disp, dispp): - self.beta = beta - self.alpha = alpha - self.disp = disp - self.dispp = dispp - - @property - def gamma(self): - return (1 + self.alpha**2)/self.beta \ No newline at end of file + """ + Class used to handle optic functions. + + Parameters + ---------- + lattice_file : str, optional if local_beta, local_alpha and + local_dispersion are specified. + AT lattice file path. + 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. + local_alpha : array of shape (2,), optional if lattice_file is specified. + Alpha function at the location of the tracking. Default is mean alpha + if lattice has been loaded. + local_dispersion : array of shape (4,), optional if lattice_file is + specified. + Dispersion function and derivative at the location of the tracking. + Default is zero if lattice has been loaded. + + Attributes + ---------- + use_local_values : bool + True if no lattice has been loaded. + local_gamma : array of shape (2,) + Gamma function at the location of the tracking. + lattice : AT lattice + + Methods + ------- + load_from_AT(lattice_file, **kwargs) + Load a lattice from accelerator toolbox (AT). + setup_interpolation() + Setup interpolation of the optic functions. + beta(position) + Return beta functions at specific locations given by position. + alpha(position) + Return alpha functions at specific locations given by position. + gamma(position) + Return gamma functions at specific locations given by position. + dispersion(position) + Return dispersion functions at specific locations given by position. + """ + + def __init__(self, lattice_file=None, local_beta=None, local_alpha=None, + local_dispersion=None, **kwargs): + + if lattice_file is not None: + self.use_local_values = False + self.load_from_AT(lattice_file, **kwargs) + if local_beta is None: + self.local_beta = np.mean(self.beta_array) + else: + self.local_beta = local_beta + if local_alpha is None: + self.local_alpha = np.mean(self.alpha_array) + else: + self.local_alpha = local_alpha + if local_dispersion is None: + self.local_dispersion = np.zeros((4,)) + else: + self.local_dispersion = local_dispersion + self.local_gamma = (1 + self.local_alpha**2)/self.local_beta + + else: + self.use_local_values = True + self.local_beta = local_beta + self.local_alpha = local_alpha + self.local_gamma = (1 + self.local_alpha**2)/self.local_beta + self.local_dispersion = local_dispersion + + def load_from_AT(self, lattice_file, **kwargs): + """ + Load a lattice from accelerator toolbox (AT). + + Parameters + ---------- + lattice_file : str + AT lattice file path. + n_points : int or float, optional + Minimum number of points to use for the optic function arrays. + periodicity : int, optional + Lattice periodicity, if not specified the AT lattice periodicity is + used. + """ + self.n_points = int(kwargs.get("n_points", 1e3)) + periodicity = kwargs.get("periodicity") + + self.lattice = at.load_lattice(lattice_file) + lattice = self.lattice.slice(slices=self.n_points) + refpts = np.arange(0, len(lattice)) + twiss0, tune, chrom, twiss = at.get_twiss(lattice, refpts=refpts, + get_chrom=True) + + if periodicity is None: + self.periodicity = lattice.periodicity + else: + self.periodicity = periodicity + + for i in range(self.periodicity-1): + pos = np.append(twiss.s_pos, twiss.s_pos + twiss.s_pos[-1]*(i+1)) + + self.position = pos + self.beta_array = np.tile(twiss.beta.T, self.periodicity) + self.alpha_array = np.tile(twiss.alpha.T, self.periodicity) + self.dispersion_array = np.tile(twiss.dispersion.T, self.periodicity) + + self.position = np.append(self.position, self.lattice.circumference) + self.beta_array = np.append(self.beta_array, self.beta_array[:,0:1], + axis=1) + self.alpha_array = np.append(self.alpha_array, self.alpha_array[:,0:1], + axis=1) + self.dispersion_array = np.append(self.dispersion_array, + self.dispersion_array[:,0:1], axis=1) + + self.gamma_array = (1 + self.alpha_array**2)/self.beta_array + self.tune = tune * self.periodicity + self.chro = chrom * self.periodicity + self.ac = at.get_mcf(self.lattice) + + self.setup_interpolation() + + + def setup_interpolation(self): + """Setup interpolation of the optic functions.""" + self.betaX = interp1d(self.position, self.beta_array[0,:], + kind='linear') + self.betaY = interp1d(self.position, self.beta_array[1,:], + kind='linear') + self.alphaX = interp1d(self.position, self.alpha_array[0,:], + kind='linear') + self.alphaY = interp1d(self.position, self.alpha_array[1,:], + kind='linear') + self.gammaX = interp1d(self.position, self.gamma_array[0,:], + kind='linear') + self.gammaY = interp1d(self.position, self.gamma_array[1,:], + kind='linear') + self.dispX = interp1d(self.position, self.dispersion_array[0,:], + kind='linear') + self.disppX = interp1d(self.position, self.dispersion_array[1,:], + kind='linear') + self.dispY = interp1d(self.position, self.dispersion_array[2,:], + kind='linear') + self.disppY = interp1d(self.position, self.dispersion_array[3,:], + kind='linear') + + def beta(self, position): + """ + Return beta functions at specific locations given by position. If no + lattice has been loaded, local values are returned. + + Parameters + ---------- + position : array or float + Longitudinal position at which the beta functions are returned. + + Returns + ------- + beta : array + Beta functions. + """ + if self.use_local_values: + return np.outer(self.local_beta, np.ones_like(position)) + else: + beta = [self.betaX(position), self.betaY(position)] + return np.array(beta) + + def alpha(self, position): + """ + Return alpha functions at specific locations given by position. If no + lattice has been loaded, local values are returned. + + Parameters + ---------- + position : array or float + Longitudinal position at which the alpha functions are returned. + + Returns + ------- + alpha : array + Alpha functions. + """ + if self.use_local_values: + return np.outer(self.local_alpha, np.ones_like(position)) + else: + alpha = [self.alphaX(position), self.alphaY(position)] + return np.array(alpha) + + def gamma(self, position): + """ + Return gamma functions at specific locations given by position. If no + lattice has been loaded, local values are returned. + + Parameters + ---------- + position : array or float + Longitudinal position at which the gamma functions are returned. + + Returns + ------- + gamma : array + Gamma functions. + """ + if self.use_local_values: + return np.outer(self.local_gamma, np.ones_like(position)) + else: + gamma = [self.gammaX(position), self.gammaY(position)] + return np.array(gamma) + + def dispersion(self, position): + """ + Return dispersion functions at specific locations given by position. + If no lattice has been loaded, local values are returned. + + Parameters + ---------- + position : array or float + Longitudinal position at which the dispersion functions are + returned. + + Returns + ------- + dispersion : array + Dispersion functions. + """ + if self.use_local_values: + return np.outer(self.local_dispersion, np.ones_like(position)) + else: + dispersion = [self.dispX(position), self.disppX(position), + self.dispY(position), self.disppY(position)] + return np.array(dispersion) + +class PhyisicalModel: + """ + Store the lattice physical parameters such as apperture and resistivity. + + Parameters + ---------- + ring : Synchrotron object + x_right : float + Default value for the right horizontal aperture in [m]. + y_top : float + Default value for the top vertical aperture in [m]. + shape : str + Default value for the shape of the chamber cross section + (elli/circ/rect). + rho : float + Default value for the resistivity of the chamber material [ohm.m]. + x_left : float, optional + Default value for the left horizontal aperture in [m]. + y_bottom : float, optional + Default value for the bottom vertical aperture in [m]. + n_points : int or float, optional + Number of points to use in class arrays + + Attributes + ---------- + position : array of shape (n_points,) + Longitudinal position in [m]. + x_right : array of shape (n_points,) + Right horizontal aperture in [m]. + y_top : array of shape (n_points,) + Top vertical aperture in [m]. + shape : array of shape (n_points - 1,) + Shape of the chamber cross section (elli/circ/rect). + rho : array of shape (n_points - 1,) + Resistivity of the chamber material in [ohm.m]. + x_left : array of shape (n_points,) + Left horizontal aperture in [m]. + y_bottom : array of shape (n_points,) + Bottom vertical aperture in [m]. + length : array of shape (n_points - 1,) + Length of each segments between two longitudinal positions in [m]. + center : array of shape (n_points - 1,) + Center of each segments between two longitudinal positions in [m]. + + Methods + ------- + change_values(start_position, end_position, x_right, y_top, shape, rho) + Change the physical parameters between start_position and end_position. + taper(start_position, end_position, x_right_start, x_right_end, + y_top_start, y_top_end, shape, rho) + Change the physical parameters to have a tapered transition between + start_position and end_position. + plot_aperture() + Plot horizontal and vertical apertures. + resistive_wall_effective_radius(optics) + Return the effective radius of the chamber for resistive wall + calculations. + """ + def __init__(self, ring, x_right, y_top, shape, rho, x_left=None, + y_bottom=None, n_points=1e4): + + self.n_points = int(n_points) + self.position = np.linspace(0, ring.L, self.n_points) + self.x_right = np.ones_like(self.position)*x_right + self.y_top = np.ones_like(self.position)*y_top + + if x_left is None: + self.x_left = np.ones_like(self.position)*-1*x_right + else: + self.x_left = np.ones_like(self.position)*x_left + + if y_bottom is None: + self.y_bottom = np.ones_like(self.position)*-1*y_top + else: + self.y_bottom = np.ones_like(self.position)*y_bottom + + self.length = self.position[1:] - self.position[:-1] + self.center = (self.position[1:] + self.position[:-1])/2 + self.rho = np.ones_like(self.center)*rho + + self.shape = np.repeat(np.array([shape]), self.n_points-1) + + def resistive_wall_effective_radius(self, optics, x_right=True, + y_top=True): + """ + Return the effective radius of the chamber for resistive wall + calculations as defined in Eq. 27 of [1]. + + Parameters + ---------- + optics : Optics object + x_right : bool, optional + If True, x_right is used, if Fasle, x_left is used. + y_top : TYPE, optional + If True, y_top is used, if Fasle, y_bottom is used. + + Returns + ------- + rho_star : float + Effective resistivity of the chamber material in [ohm.m]. + a1_L : float + Effective longitudinal radius of the chamber of the first power in + [m]. + a2_L : float + Effective longitudinal radius of the chamber of the second power + in [m]. + a3_H : float + Effective horizontal radius of the chamber of the third power in + [m]. + a4_H : float + Effective horizontal radius of the chamber of the fourth power in + [m]. + a3_V : float + Effective vertical radius of the chamber of the third power in [m]. + a4_V : float + Effective vertical radius of the chamber of the fourth power in + [m]. + + References + ---------- + [1] Skripka, Galina, et al. "Simultaneous computation of intrabunch + and interbunch collective beam motions in storage rings." Nuclear + Instruments and Methods in Physics Research Section A: Accelerators, + Spectrometers, Detectors and Associated Equipment 806 (2016): 221-230. + """ + + if x_right is True: + a0 = (self.x_right[1:] + self.x_right[:-1])/2 + else: + a0 = np.abs((self.x_left[1:] + self.x_left[:-1])/2) + + if y_top is True: + b0 = (self.y_top[1:] + self.y_top[:-1])/2 + else: + b0 = np.abs((self.y_bottom[1:] + self.y_bottom[:-1])/2) + + beta = optics.beta(self.center) + L = self.position[-1] + sigma = 1/self.rho + beta_H_star = 1/L*(self.length*beta[0,:]).sum() + beta_V_star = 1/L*(self.length*beta[1,:]).sum() + sigma_star = 1/L*(self.length*sigma).sum() + + a1_H = (((self.length*beta[0,:]/(np.sqrt(sigma)*(a0)**1)).sum())**(-1)*L*beta_H_star/np.sqrt(sigma_star))**(1/1) + a2_H = (((self.length*beta[0,:]/(np.sqrt(sigma)*(a0)**2)).sum())**(-1)*L*beta_H_star/np.sqrt(sigma_star))**(1/2) + + a1_V = (((self.length*beta[1,:]/(np.sqrt(sigma)*(b0)**1)).sum())**(-1)*L*beta_V_star/np.sqrt(sigma_star))**(1/1) + a2_V = (((self.length*beta[1,:]/(np.sqrt(sigma)*(b0)**2)).sum())**(-1)*L*beta_V_star/np.sqrt(sigma_star))**(1/2) + + a3_H = (((self.length*beta[0,:]/(np.sqrt(sigma)*(a0)**3)).sum())**(-1)*L*beta_H_star/np.sqrt(sigma_star))**(1/3) + a4_H = (((self.length*beta[0,:]/(np.sqrt(sigma)*(a0)**4)).sum())**(-1)*L*beta_H_star/np.sqrt(sigma_star))**(1/4) + + a3_V = (((self.length*beta[1,:]/(np.sqrt(sigma)*(b0)**3)).sum())**(-1)*L*beta_V_star/np.sqrt(sigma_star))**(1/3) + a4_V = (((self.length*beta[1,:]/(np.sqrt(sigma)*(b0)**4)).sum())**(-1)*L*beta_V_star/np.sqrt(sigma_star))**(1/4) + + a1_L = min((a1_H,a1_V)) + a2_L = min((a2_H,a2_V)) + + return (1/sigma_star, a1_L, a2_L, a3_H, a4_H, a3_V, a4_V) + + def change_values(self, start_position, end_position, x_right, y_top, + shape, rho, x_left=None, y_bottom=None): + """ + Change the physical parameters between start_position and end_position. + + Parameters + ---------- + start_position : float + end_position : float + x_right : float + Right horizontal aperture in [m]. + y_top : float + Top vertical aperture in [m]. + shape : str + Shape of the chamber cross section (elli/circ/rect). + rho : float + Resistivity of the chamber material [ohm.m]. + x_left : float, optional + Left horizontal aperture in [m]. + y_bottom : float, optional + Bottom vertical aperture in [m]. + """ + ind = (self.position > start_position) & (self.position < end_position) + self.x_right[ind] = x_right + self.y_top[ind] = y_top + + if x_left is None: + self.x_left[ind] = -1*x_right + else: + self.x_left[ind] = x_left + + if y_bottom is None: + self.y_bottom[ind] = -1*y_top + else: + self.y_bottom[ind] = y_bottom + + ind2 = ((self.position[:-1] > start_position) & + (self.position[1:] < end_position)) + self.rho[ind2] = rho + self.shape[ind2] = shape + + def taper(self, start_position, end_position, x_right_start, x_right_end, + y_top_start, y_top_end, shape, rho, x_left_start=None, + x_left_end=None, y_bottom_start=None, y_bottom_end=None): + """ + Change the physical parameters to have a tapered transition between + start_position and end_position. + + Parameters + ---------- + start_position : float + end_position : float + x_right_start : float + Right horizontal aperture at the start of the taper in [m]. + x_right_end : float + Right horizontal aperture at the end of the taper in [m]. + y_top_start : float + Top vertical aperture at the start of the taper in [m]. + y_top_end : float + Top vertical aperture at the end of the taper in [m]. + shape : str + Shape of the chamber cross section (elli/circ/rect). + rho : float + Resistivity of the chamber material [ohm.m]. + x_left_start : float, optional + Left horizontal aperture at the start of the taper in [m]. + x_left_end : float, optional + Left horizontal aperture at the end of the taper in [m]. + y_bottom_start : float, optional + Bottom vertical aperture at the start of the taper in [m]. + y_bottom_end : float, optional + Bottom vertical aperture at the end of the taper in [m]. + """ + ind = (self.position > start_position) & (self.position < end_position) + self.x_right[ind] = np.linspace(x_right_start, x_right_end, sum(ind)) + self.y_top[ind] = np.linspace(y_top_start, y_top_end, sum(ind)) + + if (x_left_start is not None) and (x_left_end is not None): + self.x_left[ind] = np.linspace(x_left_start, x_left_end, sum(ind)) + else: + self.x_left[ind] = -1*np.linspace(x_right_start, x_right_end, + sum(ind)) + + if (y_bottom_start is not None) and (y_bottom_end is not None): + self.y_bottom[ind] = np.linspace(y_bottom_start, y_bottom_end, + sum(ind)) + else: + self.y_bottom[ind] = -1*np.linspace(y_top_start, y_top_end, + sum(ind)) + + ind2 = ((self.position[:-1] > start_position) + & (self.position[1:] < end_position)) + self.rho[ind2] = rho + self.shape[ind2] = shape + + def plot_aperture(self): + """Plot horizontal and vertical apertures.""" + fig, axs = plt.subplots(2) + axs[0].plot(self.position,self.x_right*1e3) + axs[0].plot(self.position,self.x_left*1e3) + axs[0].set(xlabel="Longitudinal position [m]", + ylabel="Horizontal aperture [mm]") + axs[0].legend(["Right","Left"]) + + axs[1].plot(self.position,self.y_top*1e3) + axs[1].plot(self.position,self.y_bottom*1e3) + axs[1].set(xlabel="Longitudinal position [m]", + ylabel="Vertical aperture [mm]") + axs[1].legend(["Top","Bottom"]) diff --git a/tracking/particles.py b/tracking/particles.py index d28669a5c615f642b00fadb0bd0a02f4134c2219..cb4c0468306f63b81fb9d45fc56f0127713b6f83 100644 --- a/tracking/particles.py +++ b/tracking/particles.py @@ -11,12 +11,12 @@ import numpy as np import pandas as pd import matplotlib.pyplot as plt from tracking.parallel import Mpi -from scipy.constants import c, m_e, m_p, e import seaborn as sns -import h5py as hp +from scipy.constants import c, m_e, m_p, e class Particle: - """ Define a particle object + """ + Define a particle object Attributes ---------- @@ -88,6 +88,8 @@ class Bunch: ------- init_gaussian(cov=None, mean=None, **kwargs) Initialize bunch particles with 6D gaussian phase space. + plot_phasespace(x_var="tau", y_var="delta", plot_type="j") + Plot phase space. References ---------- @@ -252,17 +254,17 @@ class Bunch: if cov is None: sigma_0 = kwargs.get("sigma_0", self.ring.sigma_0) sigma_delta = kwargs.get("sigma_delta", self.ring.sigma_delta) - optics = kwargs.get("optics", self.ring.mean_optics) + optics = kwargs.get("optics", self.ring.optics) cov = np.zeros((6,6)) - cov[0,0] = self.ring.emit[0]*optics.beta[0] - cov[1,1] = self.ring.emit[0]*optics.gamma[0] - cov[0,1] = -1*self.ring.emit[0]*optics.alpha[0] - cov[1,0] = -1*self.ring.emit[0]*optics.alpha[0] - cov[2,2] = self.ring.emit[1]*optics.beta[1] - cov[3,3] = self.ring.emit[1]*optics.gamma[1] - cov[2,3] = -1*self.ring.emit[1]*optics.alpha[1] - cov[3,2] = -1*self.ring.emit[1]*optics.alpha[1] + cov[0,0] = self.ring.emit[0]*optics.local_beta[0] + cov[1,1] = self.ring.emit[0]*optics.local_gamma[0] + cov[0,1] = -1*self.ring.emit[0]*optics.local_alpha[0] + cov[1,0] = -1*self.ring.emit[0]*optics.local_alpha[0] + cov[2,2] = self.ring.emit[1]*optics.local_beta[1] + cov[3,3] = self.ring.emit[1]*optics.local_gamma[1] + cov[2,3] = -1*self.ring.emit[1]*optics.local_alpha[1] + cov[3,2] = -1*self.ring.emit[1]*optics.local_alpha[1] cov[4,4] = sigma_0**2 cov[5,5] = sigma_delta**2 @@ -274,18 +276,65 @@ class Bunch: self.particles["tau"] = values[:,4] self.particles["delta"] = values[:,5] - def plot_phasespace(self, x_var, y_var, plot_type="j"): + def binning(self, dimension="tau", n_bin=75): + """ + Bin macro-particles. + + Parameters + ---------- + dimension : str, optional + Dimension in which the binning is done. The default is "tau". + n_bin : int, optional + Number of bins. The default is 75. + + Returns + ------- + bins : IntervalIndex of shape (n_bin,) + Bins where the particles are sorted. + sorted_index : Series of shape (self.mp_number,) + Bin number of each macro-particles. + profile : Series of shape (n_bin,) + Number of marco-particles in each bin. + + """ + bins = pd.interval_range(start=min(self[dimension]), + end=max(self[dimension]), + periods=n_bin, + name=dimension) + sorted_index = pd.cut(self[dimension], bins) + sorted_index = sorted_index.reindex(self.alive.index) + profile = sorted_index.value_counts().sort_index() + return (bins, sorted_index, profile) + + def plot_profile(self, dimension="tau", n_bin=75): + """ + Plot bunch profile. + + Parameters + ---------- + dimension : str, optional + Dimension to plot. The default is "tau". + n_bin : int, optional + Number of bins. The default is 75. + + """ + bins, sorted_index, profile = self.binning(dimension, n_bin) + fig = plt.figure() + ax = fig.gca() + ax.plot(bins.mid, profile) + + def plot_phasespace(self, x_var="tau", y_var="delta", plot_type="j"): """ - Plot longitudinal phase space. + Plot phase space. Parameters ---------- x_var : str - name from Bunch object to plot on horizontal axis. + Dimension to plot on horizontal axis. y_var : str - name from Bunch object to plot on ver. axis. + Dimension to plot on vertical axis. plot_type : str {"j" , "sc"} - type of the plot. The defualt value is "j" for a joint plot. + Type of the plot. The defualt value is "j" for a joint plot. Can be modified to "sc" for a scatter plot. """ @@ -305,8 +354,8 @@ class Bunch: plt.xlabel(label_dict[x_var]) plt.ylabel(label_dict[y_var]) - else: raise ValueError("Plot type not recognised") - + else: + raise ValueError("Plot type not recognised.") class Beam: """ @@ -344,7 +393,6 @@ class Beam: Status of MPI parallelisation, should not be changed directly but with mpi_init() and mpi_close() - Methods ------ init_beam(filling_pattern, current_per_bunch=1e-3, mp_per_bunch=1e3) @@ -358,6 +406,8 @@ class Beam: all processors. Rather slow mpi_close() Call mpi_gather and switch off MPI parallelisation + plot(var, option=None) + Plot variables with respect to bunch number. """ def __init__(self, ring, bunch_list=None): @@ -392,12 +442,29 @@ class Beam: @property def not_empty(self): - """Return a generator to iterate over not empty bunches""" - for bunch in self: - if bunch.current == 0: - pass + """Return a generator to iterate over not empty bunches.""" + for index, value in enumerate(self.filling_pattern): + if value == True: + yield self[index] else: - yield bunch + pass + + @property + def distance_between_bunches(self): + filling_pattern = self.filling_pattern + distance = np.zeros(filling_pattern.shape) + for index, value in enumerate(filling_pattern): + if value == False: + pass + elif value == True: + count = 1 + for value2 in filling_pattern[index+1:]: + if value2 == False: + count += 1 + elif value2 == True: + break + distance[index] = count + return distance def init_beam(self, filling_pattern, current_per_bunch=1e-3, mp_per_bunch=1e3): @@ -444,20 +511,25 @@ class Beam: raise TypeError("{} should be bool or float64".format(filling_pattern.dtype)) self.bunch_list = bunch_list + self.update_filling_pattern() for bunch in self.not_empty: bunch.init_gaussian() - @property - def filling_pattern(self): - """Return an array with the filling pattern of the beam as bool""" + def update_filling_pattern(self): + """Update the beam filling pattern.""" filling_pattern = [] for bunch in self: if bunch.current != 0: filling_pattern.append(True) else: filling_pattern.append(False) - return np.array(filling_pattern) + self._filling_pattern = np.array(filling_pattern) + + @property + def filling_pattern(self): + """Return an array with the filling pattern of the beam as bool""" + return self._filling_pattern @property def bunch_current(self): @@ -542,13 +614,13 @@ class Beam: self.mpi_switch = False self.mpi = None - def plot_bunchnumber(self, var, option=None): + def plot(self, var, option=None): """ - Plot varviables with respect to bunch number. + Plot variables with respect to bunch number. Parameters ---------- - var : str {"bunch_currebt", "bunch_charge", "bunch_particle", + var : str {"bunch_current", "bunch_charge", "bunch_particle", "bunch_mean", "bunch_std", "bunch_emit"} Variable to be plotted. option : str, optional @@ -558,7 +630,6 @@ class Beam: option = {"x","xp","y","yp","tau","delta"}. For "bunch_emit", option = {"x","y","s"}. The default is None. - """ var_dict = {"bunch_current":self.bunch_current, @@ -630,4 +701,4 @@ class Beam: - \ No newline at end of file + diff --git a/tracking/synchrotron.py b/tracking/synchrotron.py index bc3d7ada308940c0af3d446c80c619b99233d2c3..a622c0458f5ba5efc77ac2492c3f0dc4f245223b 100644 --- a/tracking/synchrotron.py +++ b/tracking/synchrotron.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- """ -General elements +Module where the synchrotron class is defined. @author: Alexis Gamelin @date: 15/01/2020 @@ -8,25 +8,92 @@ General elements import numpy as np from scipy.constants import c, e - class Synchrotron: - """ Synchrotron class to store main properties. """ - def __init__(self, h, L, E0, particle, **kwargs): - self.particle = particle + """ + Synchrotron class to store main properties. + + Optional parameters are optional only if the Optics object passed to the + class uses a loaded lattice. + + Parameters + ---------- + h : int + Harmonic number of the accelerator. + optics : Optics object + Object where the optic functions are stored. + particle : Particle object + Particle which is accelerated. + tau : array of shape (3,) + Horizontal, vertical and longitudinal damping times in [s]. + sigma_delta : float + Equilibrium energy spread. + sigma_0 : float + Natural bunch length in [s]. + emit : array of shape (2,) + Horizontal and vertical equilibrium emittance in [m.rad]. + L : float, optional + Ring circumference in [m]. + E0 : float, optional + Nominal (total) energy of the ring in [eV]. + ac : float, optional + Momentum compaction factor. + tune : array of shape (2,), optional + Horizontal and vertical tunes. + chro : array of shape (2,), optional + Horizontal and vertical (non-normalized) chromaticities. + U0 : float, optional + Energy loss per turn in [eV]. + + Attributes + ---------- + T0 : float + Revolution time in [s]. + f0 : float + Revolution frequency in [Hz]. + omega0 : float + Angular revolution frequency in [Hz.rad] + T1 : flaot + Fundamental RF period in [s]. + f1 : float + Fundamental RF frequency in [Hz]. + omega1 : float + Fundamental RF angular frequency in [Hz.rad]. + k1 : float + Fundamental RF wave number in [m**-1]. + gamma : float + Relativistic Lorentz gamma. + beta : float + Relativistic Lorentz beta. + eta : float + Momentum compaction. + sigma : array of shape (4,) + RMS beam size at equilibrium in [m]. + """ + def __init__(self, h, optics, particle, **kwargs): self._h = h - self.L = L - self.E0 = E0 # Nominal (total) energy of the ring [eV] + self.particle = particle + self.optics = optics - self.ac = kwargs.get('ac') # Momentum compaction factor - self.U0 = kwargs.get('U0') # Energy loss per turn [eV] + if self.optics.use_local_values == False: + self.L = kwargs.get('L', self.optics.lattice.circumference) + self.E0 = kwargs.get('E0', self.optics.lattice.energy) + self.ac = kwargs.get('ac', self.optics.ac) + self.tune = kwargs.get('tune', self.optics.tune) + self.chro = kwargs.get('chro', self.optics.chro) + self.U0 = kwargs.get('U0', self.optics.lattice.energy_loss) + else: + self.L = kwargs.get('L') # Ring circumference [m] + self.E0 = kwargs.get('E0') # Nominal (total) energy of the ring [eV] + self.ac = kwargs.get('ac') # Momentum compaction factor + self.tune = kwargs.get('tune') # X/Y/S tunes + self.chro = kwargs.get('chro') # X/Y (non-normalized) chromaticities + self.U0 = kwargs.get('U0') # Energy loss per turn [eV] + self.tau = kwargs.get('tau') # X/Y/S damping times [s] self.sigma_delta = kwargs.get('sigma_delta') # Equilibrium energy spread self.sigma_0 = kwargs.get('sigma_0') # Natural bunch length [s] - self.tune = kwargs.get('tune') # X/Y/S tunes self.emit = kwargs.get('emit') # X/Y emittances in [m.rad] - self.chro = kwargs.get('chro') # X/Y (non-normalized) chromaticities - self.mean_optics = kwargs.get('mean_optics') # Optics object with mean values @property def h(self): @@ -34,7 +101,7 @@ class Synchrotron: return self._h @h.setter - def h(self,value): + def h(self, value): self._h = value self.L = self.L # call setter @@ -47,6 +114,7 @@ class Synchrotron: def L(self,value): self._L = value self._T0 = self.L/c + self._T1 = self.T0/self.h self._f0 = 1/self.T0 self._omega0 = 2*np.pi*self.f0 self._f1 = self.h*self.f0 @@ -62,6 +130,15 @@ class Synchrotron: def T0(self, value): self.L = c*value + @property + def T1(self): + """"Fundamental RF period [s]""" + return self._T1 + + @T1.setter + def T1(self, value): + self.L = c*value*self.h + @property def f0(self): """Revolution frequency [Hz]""" @@ -145,8 +222,12 @@ class Synchrotron: def sigma(self): """RMS beam size at equilibrium""" sigma = np.zeros((4,)) - sigma[0] = (self.emit[0]*self.mean_optics.beta[0] + self.mean_optics.disp[0]**2*self.sigma_delta)**0.5 - sigma[1] = (self.emit[0]*self.mean_optics.alpha[0] + self.mean_optics.dispp[0]**2*self.sigma_delta)**0.5 - sigma[2] = (self.emit[1]*self.mean_optics.beta[1] + self.mean_optics.disp[1]**2*self.sigma_delta)**0.5 - sigma[3] = (self.emit[1]*self.mean_optics.alpha[1] + self.mean_optics.dispp[1]**2*self.sigma_delta)**0.5 + sigma[0] = (self.emit[0]*self.optics.local_beta[0] + + self.optics.local_dispersion[0]**2*self.sigma_delta)**0.5 + sigma[1] = (self.emit[0]*self.optics.local_alpha[0] + + self.optics.local_dispersion[1]**2*self.sigma_delta)**0.5 + sigma[2] = (self.emit[1]*self.optics.local_beta[1] + + self.optics.local_dispersion[2]**2*self.sigma_delta)**0.5 + sigma[3] = (self.emit[1]*self.optics.local_alpha[1] + + self.optics.local_dispersion[3]**2*self.sigma_delta)**0.5 return sigma \ No newline at end of file