diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py index 6f76fa053d0f624483d8b84fad143d67c5ade966..9a6d0263c55a20d8ec002fe9d58d8c09ea898d26 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 0f16258355fe4abed9d4e1b08c50eec39189215c..58049c53d30e36c3307eae3baf3d35c5a1a9d456 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