From 5b0aa9382e7c404e58ad011152c9f8e95f443242 Mon Sep 17 00:00:00 2001
From: Watanyu Foosang <watanyu.f@gmail.com>
Date: Fri, 26 Feb 2021 10:05:34 +0100
Subject: [PATCH] Add Courant-Snyder invariant calculation method to Bunch
class
---
tracking/element.py | 23 ++++++++++++++++++++++-
tracking/monitors/__init__.py | 3 ++-
tracking/monitors/monitors.py | 2 +-
tracking/optics.py | 3 ++-
tracking/particles.py | 10 ++++++++++
tracking/synchrotron.py | 15 +++++++++++++++
6 files changed, 52 insertions(+), 4 deletions(-)
diff --git a/tracking/element.py b/tracking/element.py
index 0ef071c..00f75ed 100644
--- a/tracking/element.py
+++ b/tracking/element.py
@@ -125,7 +125,7 @@ class SynchrotronRadiation(Element):
if (self.switch[0] == True):
rand = np.random.normal(size=len(bunch))
bunch["delta"] = ((1 - 2*self.ring.T0/self.ring.tau[2])*bunch["delta"] +
- 2*self.ring.sigma_delta*(self.ring.T0/self.ring.tau[2])**0.5*rand)
+ 2*self.ring.sigma_delta*(self.ring.T0/self.ring.tau[2])**0.5*rand)
if (self.switch[1] == True):
rand = np.random.normal(size=(len(bunch),2))
@@ -142,6 +142,27 @@ class SynchrotronRadiation(Element):
# Reset energy change to 0 for next turn
bunch.energy_change = 0
+ #--------------------------------------
+ # if (self.switch[0] == True):
+ # rand = np.random.normal(size=len(bunch))
+ # bunch["delta"] = ((1 - 2*self.ring.T0/self.ring.tau[2])*bunch["delta"] +
+ # 2*self.ring.sigma_delta*(self.ring.T0/self.ring.tau[2])**0.5*rand)
+
+ # if (self.switch[1] == True):
+ # rand = np.random.normal(size=(len(bunch),2))
+ # bunch["x"] += self.ring.sigma[0]*(2*self.ring.T0/self.ring.tau[0])**0.5*rand[:,0]
+ # bunch["xp"] = (1 + bunch["delta"])/(1 + bunch["delta"] + bunch.energy_change)*bunch["xp"]
+ # bunch["xp"] += self.ring.sigma[1]*(2*self.ring.T0/self.ring.tau[0])**0.5*rand[:,1]
+
+ # if (self.switch[2] == True):
+ # rand = np.random.normal(size=(len(bunch),2))
+ # bunch["y"] += self.ring.sigma[2]*(2*self.ring.T0/self.ring.tau[1])**0.5*rand[:,0]
+ # bunch["yp"] = (1 + bunch["delta"])/(1 + bunch["delta"] + bunch.energy_change)*bunch["yp"]
+ # bunch["yp"] += self.ring.sigma[3]*(2*self.ring.T0/self.ring.tau[1])**0.5*rand[:,1]
+
+ # # Reset energy change to 0 for next turn
+ # bunch.energy_change = 0
+ #----------------------------------------
class TransverseMap(Element):
"""
Transverse map for a single turn in the synchrotron.
diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index 1d186da..bd39dbd 100644
--- a/tracking/monitors/__init__.py
+++ b/tracking/monitors/__init__.py
@@ -15,4 +15,5 @@ from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
plot_phasespacedata,
plot_profiledata,
plot_beamdata,
- plot_wakedata)
\ No newline at end of file
+ plot_wakedata,
+ plot_tunedata)
\ No newline at end of file
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 6a46e38..7a2b209 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -16,7 +16,7 @@ from mbtrack2.tracking.element import Element
from mbtrack2.tracking.particles import Bunch, Beam
from scipy.interpolate import interp1d
from abc import ABCMeta
-from mpi4py import MPI
+#from mpi4py import MPI
class Monitor(Element, metaclass=ABCMeta):
"""
diff --git a/tracking/optics.py b/tracking/optics.py
index 58049c5..865a136 100644
--- a/tracking/optics.py
+++ b/tracking/optics.py
@@ -160,7 +160,7 @@ class Optics:
kind='linear')
self.disppY = interp1d(self.position, self.dispersion_array[3,:],
kind='linear')
-
+
@property
def local_beta(self):
"""
@@ -215,6 +215,7 @@ class Optics:
"""
return self._local_gamma
+
def beta(self, position):
"""
Return beta functions at specific locations given by position. If no
diff --git a/tracking/particles.py b/tracking/particles.py
index 630e232..87988e8 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -231,6 +231,16 @@ class Bunch:
np.mean(self['tau']*self['delta'])**2)**(0.5)
return np.array([emitX, emitY, emitS])
+ @property
+ def cs_invariant(self):
+ Jx = (self.ring.optics.local_gamma * np.mean(self['x']**2)) + \
+ (2*self.ring.optics.local_alpha * np.mean(self['x']*self['xp'])) + \
+ (self.ring.optics.local_beta * np.mean(self['xp']**2))
+ Jy = (self.ring.optics.local_gamma * np.mean(self['y']**2)) + \
+ (2*self.ring.optics.local_alpha * np.mean(self['y']*self['yp'])) + \
+ (self.ring.optics.local_beta * np.mean(self['yp']**2))
+ return np.array((Jx,Jy))
+
@property
def energy_change(self):
"""Store the particle relative energy change in the last turn. Used by
diff --git a/tracking/synchrotron.py b/tracking/synchrotron.py
index f5e2d69..1c45d8c 100644
--- a/tracking/synchrotron.py
+++ b/tracking/synchrotron.py
@@ -265,6 +265,21 @@ class Synchrotron:
sigma[3,:] = (self.emit[1]*self.optics.alpha(position)[1] +
self.optics.dispersion(position)[3]**2*self.sigma_delta)**0.5
return sigma
+ #------------------------------------------
+ # @property
+ # def sigma(self):
+ # """RMS beam size at equilibrium"""
+ # sigma = np.zeros((4,))
+ # sigma[0] = (self.emit[0]*self.optics.beta(10)[0] +
+ # self.optics.dispersion(10)[0]**2*self.sigma_delta)**0.5
+ # sigma[1] = (self.emit[0]*self.optics.alpha(10)[0] +
+ # self.optics.dispersion(10)[1]**2*self.sigma_delta)**0.5
+ # sigma[2] = (self.emit[1]*self.optics.beta(10)[1] +
+ # self.optics.dispersion(10)[2]**2*self.sigma_delta)**0.5
+ # sigma[3] = (self.emit[1]*self.optics.alpha(10)[1] +
+ # self.optics.dispersion(10)[3]**2*self.sigma_delta)**0.5
+ # return sigma
+ #------------------------------------------
def synchrotron_tune(self, Vrf):
"""
--
GitLab