Add ellipse

pmd_beamphysics/particles.py

@@ -926,6 +926,7 @@ def plot(self, key1='x', key2=None,
              ylim=None,
              return_figure=False, 
              tex=True, nice=True,
+             ellipse=False,
              **kwargs):
         """
         1d or 2d density plot. 
@@ -993,6 +994,7 @@ def plot(self, key1='x', key2=None,
                                 ylim=ylim,
                                 tex=tex,
                                 nice=nice,
+                                ellipse=ellipse,
                                 **kwargs)
 
         if return_figure:

pmd_beamphysics/plot.py

@@ -16,7 +16,7 @@
 from pmd_beamphysics.units import (nice_array, nice_scale_prefix,
                                    plottable_array)
 
-from .statistics import slice_statistics
+from .statistics import slice_statistics, twiss_ellipse_points
 
 CMAP0 = copy(plt.get_cmap('viridis'))
 CMAP0.set_under('white')
@@ -234,6 +234,7 @@ def marginal_plot(particle_group, key1='t', key2='p',
                   ylim=None,
                   tex=True,
                   nice=True,
+                  ellipse=False,
                   **kwargs):
     """
     Density plot and projections
@@ -267,7 +268,10 @@ def marginal_plot(particle_group, key1='t', key2='p',
         Use TEX for labels
         
     nice: bool, default = True
-        
+
+    ellipse: bool, default = True
+        If True, plot an ellipse representing the 
+        2x2 sigma matrix
     
     Returns
     -------
@@ -307,9 +311,19 @@ def marginal_plot(particle_group, key1='t', key2='p',
     ax_marg_y = fig.add_subplot(gs[1:4,3])
     #ax_info = fig.add_subplot(gs[0, 3:4])
     #ax_info.table(cellText=['a'])
-
+
+    # Main plot
     # Proper weighting
     ax_joint.hexbin(x, y, C=w, reduce_C_function=np.sum, gridsize=bins, cmap=CMAP0, vmin=1e-20)
+
+
+    if ellipse:
+        sigma_mat2 = particle_group.cov(key1, key2)
+        x_ellipse, y_ellipse = twiss_ellipse_points(sigma_mat2)
+        x_ellipse += particle_group.avg(key1)
+        y_ellipse += particle_group.avg(key2)
+        ax_joint.plot(x_ellipse/f1, y_ellipse/f2, color='red')
+
 
     # Manual histogramming version
     #H, xedges, yedges = np.histogram2d(x, y, weights=w, bins=bins)

pmd_beamphysics/statistics.py

@@ -56,6 +56,27 @@ def twiss_calc(sigma_mat2):
     return twiss
 
 
+def twiss_ellipse_points(sigma_mat2, n_points=36):
+    """
+    Returns points that will trace a the rms ellipse 
+    from a 2x2 covariance matrix `sigma_mat2`.
+
+    Returns
+    -------
+    vec: np.ndarray with shape (2, n_points)
+        x, p representing the ellipse points.
+    
+    """
+    twiss = twiss_calc(sigma_mat2)
+    A = A_mat_calc(twiss['beta'], twiss['alpha'])
+
+    theta = np.linspace(0, np.pi*2, n_points)
+    zvec0 = np.array([np.cos(theta), np.sin(theta)]) * np.sqrt(2*twiss['emit'])
+
+    zvec1 = np.matmul(A, zvec0)
+    return zvec1
+
+
 
 def twiss_match(x, p, beta0=1, alpha0=0, beta1=1, alpha1=0):
     """