From 6bbd6fef581d3099580e4b8274229e50ab60a091 Mon Sep 17 00:00:00 2001
From: Watanyu Foosang <watanyu.f@gmail.com>
Date: Thu, 25 Feb 2021 16:34:23 +0100
Subject: [PATCH] Make optical functions as property objects
Main change
- Variable local_beta, local_alpha, and local_gamma are made property objects.
Minor change
- Add an option to plot phase space in different style.
---
tracking/monitors/plotting.py | 6 ++-
tracking/optics.py | 70 +++++++++++++++++++++++++++++++----
2 files changed, 66 insertions(+), 10 deletions(-)
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index 6f76fa0..9a6d026 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -190,7 +190,7 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
return fig
def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn,
- only_alive=True, plot_size=1):
+ only_alive=True, plot_size=1, plot_kind='kde'):
"""
Plot data recorded by PhaseSpaceMonitor.
@@ -213,6 +213,8 @@ def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn,
Number of macro-particles to plot relative to the total number
of macro-particles recorded. This option helps reduce processing time
when the data is big.
+ plot_kind : {'scatter', 'kde', 'hex', 'reg', 'resid'}, optional
+ The plot style. The default is 'kde'.
Return
------
@@ -258,7 +260,7 @@ def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn,
y_axis = path[samples,option_dict[y_var],turn_index[0][0]]
fig = sns.jointplot(x_axis*scale[option_dict[x_var]],
- y_axis*scale[option_dict[y_var]], kind="kde")
+ y_axis*scale[option_dict[y_var]], kind=plot_kind)
plt.xlabel(label[option_dict[x_var]])
plt.ylabel(label[option_dict[y_var]])
diff --git a/tracking/optics.py b/tracking/optics.py
index 0f16258..58049c5 100644
--- a/tracking/optics.py
+++ b/tracking/optics.py
@@ -66,24 +66,24 @@ class Optics:
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, axis=1)
+ self._local_beta = np.mean(self.beta_array, axis=1)
else:
- self.local_beta = local_beta
+ self._local_beta = local_beta
if local_alpha is None:
- self.local_alpha = np.mean(self.alpha_array, axis=1)
+ self._local_alpha = np.mean(self.alpha_array, axis=1)
else:
- self.local_alpha = local_alpha
+ 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
+ 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_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):
@@ -160,7 +160,61 @@ class Optics:
kind='linear')
self.disppY = interp1d(self.position, self.dispersion_array[3,:],
kind='linear')
+
+ @property
+ def local_beta(self):
+ """
+ Return beta function at the location defined by the lattice file.
+
+ """
+ return self._local_beta
+
+ @local_beta.setter
+ def local_beta(self, beta_array):
+ """
+ Set the values of beta function. Gamma function is automatically
+ recalculated after the new value of beta function is set.
+
+ Parameters
+ ----------
+ beta_array : array of shape (2,)
+ Beta function in the horizontal and vertical plane.
+
+ """
+ self._local_beta = beta_array
+ self._local_gamma = (1+self._local_alpha**2) / self._local_beta
+ @property
+ def local_alpha(self):
+ """
+ Return alpha function at the location defined by the lattice file.
+
+ """
+ return self._local_alpha
+
+ @local_alpha.setter
+ def local_alpha(self, alpha_array):
+ """
+ Set the value of beta functions. Gamma function is automatically
+ recalculated after the new value of alpha function is set.
+
+ Parameters
+ ----------
+ alpha_array : array of shape (2,)
+ Alpha function in the horizontal and vertical plane.
+
+ """
+ self._local_alpha = alpha_array
+ self._local_gamma = (1+self._local_alpha**2) / self._local_beta
+
+ @property
+ def local_gamma(self):
+ """
+ Return beta function at the location defined by the lattice file.
+
+ """
+ return self._local_gamma
+
def beta(self, position):
"""
Return beta functions at specific locations given by position. If no
--
GitLab