Num beta osc stage length

pyproject.toml

@@ -66,6 +66,7 @@ markers = [
   "StageBasic_linac: Integration tests for StageBasic linacs",
   "StageQuasistatic2d: Tests for the StageQuasistatic2d class",
   "StageWakeT: Tests for the StageWakeT class",
+  "StageHipace: Tets for the StageHipace class",
   "presets: Tests for collider presets",
   "impactx: Tests for ImpactX",
 ]

src/abel/classes/stage/impl/stage_reduced_models.py

@@ -682,9 +682,9 @@ def track_upramp(self, beam0, driver0, shot_path=None):
                 raise TypeError('upramp is not a StageReducedModels.')
 
             # Set a new time step for the ramp
-            n_steps = 25  # Do n_steps time steps in the tracking of the ramp
+            n_steps = 15  # Do n_steps time steps in the tracking of the ramp
             lambda_beta = upramp.matched_beta_function_flattop(beam0_energy) * 2*np.pi  # [m], betatron wavelength
-            upramp.time_step_mod = min(upramp.length_flattop / (lambda_beta*n_steps), 0.02)  # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c.
+            upramp.time_step_mod = min(upramp.length_flattop / (lambda_beta*n_steps), 0.04)  # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c.
 
         elif type(self.upramp) is Stage:
             upramp = self.upramp  # Allow for other types of ramps
@@ -782,9 +782,9 @@ def track_downramp(self, beam0, driver0, shot_path=None):
                 raise TypeError('downramp is not a StageReducedModels.')
 
             # Set a new time step for the ramp
-            n_steps = 25  # Do n_steps time steps in the tracking of the ramp
+            n_steps = 15  # Do n_steps time steps in the tracking of the ramp
             lambda_beta = downramp.matched_beta_function_flattop(beam0_energy) * 2*np.pi  # [m], betatron wavelength
-            downramp.time_step_mod = min(downramp.length_flattop / (lambda_beta*n_steps), 0.02)  # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c.
+            downramp.time_step_mod = min(downramp.length_flattop / (lambda_beta*n_steps), 0.04)  # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c.
 
         elif type(self.downramp) is Stage:
             downramp = self.downramp  # Allow for other types of ramps

src/abel/utilities/beam_physics.py

@@ -5,6 +5,7 @@
 # License: GPL-3.0-or-later
 
 import numpy as np
+import scipy.constants as SI
 
 
 # =============================================
@@ -1152,3 +1153,113 @@ def evolve_chromatic_amplitude(ls, inv_rhos, ks, ms, taus, beta0, alpha0=0, Dx0=
 
     return W, evolution
 
+
+# =============================================
+def phase_advance_integrate_betas(ss, betas):
+    """
+    Calculate the phase advance in one dimesion by using the composite Simpson’s 
+    rule (:func:`scipy.integrate.simpson() <scipy.integrate.simpson>`) to 
+    integrate two arrays containing the location and the beta function.
+    """
+
+    from scipy import integrate
+    inv_betas = 1/betas
+    return integrate.simpson(y=inv_betas, x=ss)
+
+
+# =============================================
+def phase_advance_traj_data(pos, angles, orbit_pos=None, orbit_angles=None):
+    """
+    Compute the unwrapped betatron phase advance from transverse position and 
+    angle data by mapping the phase-space trajectory onto a normalised unit 
+    circle and evaluating the continuous phase angle evolution.
+
+    Parameters
+    ----------
+    pos : [m] 1D float ndarray
+        Transverse position along the beamline.
+
+    angles : 1D float ndarray
+        Transverse angle dpos/ds.
+
+    orbit_pos : [m] 1D float ndarray, optional
+        Reference orbit transverse position (e.g., driver trajectory). If 
+        provided, these are subtracted before phase calculation. Defaults to 
+        `None`.
+
+    orbit_angles : 1D float ndarray, optional
+        Reference orbit transverse angles (e.g., driver trajectory angles). If 
+        provided, these are subtracted before phase calculation. Defaults to 
+        `None`.
+
+    Returns
+    -------
+    phase_advance : float
+        Unwrapped phase angle in radians, representing the continuous betatron
+        phase advance at the end of the beamline
+    """
+
+    if orbit_pos is not None:
+        pos = pos - orbit_pos
+    if orbit_angles is not None:
+        angles = angles - orbit_angles
+
+    # Compute the phase space angle by normalising coordinates to a unit circle 
+    # and taking the arctangent before unwrapping
+    phase = np.unwrap(np.arctan2(angles / np.max(np.abs(angles)), 
+                                 pos / np.max(np.abs(pos))))
+
+    return np.abs(phase[-1] - phase[0]) 
+
+
+# =============================================
+def length2num_beta_osc(length, initial_energy, accel_gradient, foc_gradient, q=SI.e):
+        """
+        Calculate the number of betatron oscillations a particle can undergo in 
+        the stage (excluding ramps).
+
+        Will take into account the contribution from an external linear magnetic 
+        field B=[gy,-gx,0] if :attr:`self.external_focusing_gradient <abel.Stage.external_focusing>`
+        is not ``None``.
+
+        Parameters
+        ----------
+        length : [m] float
+            Length of a plasma stage excluding ramps that the particle can 
+            perform betatron oscillations in. 
+
+        initial_energy : [eV] float
+            The initial energy of the particle at the start of the plasma stage. 
+
+        accel_gradient : [V/m] float
+            Nominal accelerating gradient of the plasma stage exclusing ramps. 
+
+        foc_gradient : [T/m] float
+            The gradient of the total transverse focusing field.
+
+        q : [C] float, optional
+            Particle charge. q * accel_gradient must be positive. 
+            Defaults to elementary charge.
+
+
+        Returns
+        -------
+        num_beta_osc : float
+            Total number of betatron oscillations that the particle will perform
+            across the plasma stage.
+        """
+
+        if q * accel_gradient < 0:
+            raise ValueError('q * accel_gradient must be positive.')
+
+        if accel_gradient < 1e-15: # Need to treat very small gradients separately. Often the case for plasma density ramps.
+            return np.sqrt(np.abs(q)*SI.c*foc_gradient/(initial_energy*SI.e)) * length/(2*np.pi)
+        else:
+
+            prefactor = 2*np.sqrt(np.abs(q)*foc_gradient*SI.c) / (q*accel_gradient)
+            energy_scaling = np.sqrt(initial_energy*SI.e + q*accel_gradient*length) - np.sqrt(initial_energy*SI.e)
+
+            num_beta_osc = prefactor * energy_scaling / (2*np.pi)
+
+            return num_beta_osc
+

src/abel/wrappers/hipace/hipace_wrapper.py

@@ -156,7 +156,9 @@ def hipace_write_inputs(filename_input, filename_beam, filename_driver, plasma_d
         print('>> HiPACE++: Changing from', num_cell_xy, 'to', new_num_cell_xy, ' (i.e., 2^n-1) for better performance.')
         num_cell_xy = new_num_cell_xy
 
-    # plasma-density profile from file
+    # gradient for external magnetic field
+    if external_focusing_gradient is None:
+        external_focusing_gradient = '#'
     if abs(external_focusing_gradient) > 0:
         external_focusing_comment = ''
     else: