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

@@ -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]])

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):
@@ -161,6 +161,60 @@ class Optics:
         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