Rework Optics element

Optic functions as arrays
Uses PyAT lattice for import
Adapt other files for the change

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

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"]

tracking/optics.py

@@ -5,14 +5,121 @@ Created on Wed Mar 11 18:00:28 2020
 @author: gamelina
 """
 
+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
+    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):
+        
+        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.gamma_array = (1 + self.alpha_array**2)/self.beta_array
+        self.dispersion_array = np.tile(twiss.dispersion.T, self.periodicity)
+        self.tune = tune * self.periodicity
+        self.chro = chrom * self.periodicity
+        self.ac = at.get_mcf(self.lattice)
+        
+        self.setup_interpolation()
+        
+        
+    def setup_interpolation(self):        
+        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):
+        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):
+        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):
+        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):
+        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)
+    
 
-    @property
-    def gamma(self):
-        return (1 + self.alpha**2)/self.beta
\ No newline at end of file

tracking/particles.py

@@ -248,17 +248,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
             

tracking/synchrotron.py

@@ -12,21 +12,33 @@ 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
+    
+    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 +46,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
         
@@ -145,8 +157,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