Merge remote-tracking branch 'upstream/master' into cgull-corrector-m…

…odeling

.github/workflows/ci.yml

@@ -14,10 +14,11 @@ jobs:
         python-version: [3.9]
     steps:
       - uses: actions/checkout@v3
-      - uses: conda-incubator/setup-miniconda@v2
+      - uses: conda-incubator/setup-miniconda@v3
         with:
           python-version: ${{ matrix.python-version }}
-          mamba-version: "*"
+          miniforge-version: latest
+          use-mamba: true
           channels: conda-forge
           activate-environment: pmd_beamphysics-dev
           environment-file: environment.yml

pmd_beamphysics/fields/analysis.py

@@ -7,12 +7,19 @@
 import numpy as np
 
 
+# Numpy migration per https://numpy.org/doc/stable/numpy_2_0_migration_guide.html
+if np.lib.NumpyVersion(np.__version__) >= '2.0.0':
+    from numpy import trapezoid
+else:
+    # Support 'trapz' from numpy 1.0
+    from numpy import trapz as trapezoid
+
 
 #----------------------
 # Analysis
 
 def accelerating_voltage_and_phase(z, Ez, frequency):
-    """
+    r"""
     Computes the accelerating voltage and phase for a v=c positively charged particle in an accelerating cavity field.
     
         Z = \int Ez * e^{-i k z} dz 
@@ -36,7 +43,7 @@ def accelerating_voltage_and_phase(z, Ez, frequency):
     fz =Ez*np.exp(-1j*k*z)
 
     # Integrate
-    Z = np.trapz(fz, z)
+    Z = trapezoid(fz, z)
 
     # Max voltage at phase
     voltage = np.abs(Z)
@@ -57,7 +64,7 @@ def track_field_1d(z,
                    debug=False,
                    max_step=None,
                   ):
-    """
+    r"""
     Tracks a particle in a 1d complex electric field Ez, oscillating as Ez * exp(-i omega t)
     
     Uses scipy.integrate.solve_ivp to track the particle. 
@@ -181,7 +188,7 @@ def track_field_1df(Ez_f,
                    max_step=None,
                     method='RK23'
                   ):
-    """
+    r"""
     Similar to track_field_1d, execpt uses a function Ez_f
     
     Tracks a particle in a 1d electric field Ez(z, t)

pmd_beamphysics/labels.py

@@ -122,7 +122,7 @@ def texlabel(key: str):
     if key.startswith('bunching'):
         wavelength = parse_bunching_str(key)
         x, _, prefix = nice_array(wavelength)
-        return f'\mathrm{{bunching~at}}~{x:.1f}~\mathrm{{ {prefix}m }}'
+        return fr'\mathrm{{bunching~at}}~{x:.1f}~\mathrm{{ {prefix}m }}'
 
     return None
 

pmd_beamphysics/particles.py

@@ -678,7 +678,7 @@ def average_current(self):
         return self.charge / dt
 
     def bunching(self, wavelength):
-        """
+        r"""
         Calculate the normalized bunching parameter, which is the magnitude of the 
         complex sum of weighted exponentials at a given point.
     
@@ -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. 
@@ -964,6 +965,10 @@ def plot(self, key1='x', key2=None,
             
         nice: bool, default = True
             Scale to nice units
+
+        ellipse: bool, default = True
+            If True, plot an ellipse representing the 
+            2x2 sigma matrix            
             
         return_figure: bool, default = False
             If true, return a matplotlib.figure.Figure object
@@ -993,6 +998,7 @@ def plot(self, key1='x', key2=None,
                                 ylim=ylim,
                                 tex=tex,
                                 nice=nice,
+                                ellipse=ellipse,
                                 **kwargs)
 
         if return_figure:

pmd_beamphysics/plot.py

@@ -18,7 +18,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')
@@ -236,6 +236,7 @@ def marginal_plot(particle_group, key1='t', key2='p',
                   ylim=None,
                   tex=True,
                   nice=True,
+                  ellipse=False,
                   **kwargs):
     """
     Density plot and projections
@@ -269,7 +270,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
     -------
@@ -309,9 +313,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):
     """