Tests/source and beam

abel/__init__.py

@@ -44,6 +44,7 @@
 from .classes.source.impl.source_combiner import *
 from .classes.source.impl.source_from_file import *
 from .classes.source.impl.source_flattop import *
+from .classes.source.impl.source_capsule import *
 from .classes.stage.impl.stage_basic import *
 from .classes.stage.impl.stage_nonlinear_1d import *
 from .classes.stage.impl.stage_hipace import *

abel/classes/source/impl/source_from_file.py

@@ -1,12 +1,27 @@
 from abel import Source, Beam
+import os
 
 class SourceFromFile(Source):
 
     def __init__(self, length=0, charge=None, energy=None, accel_gradient=None, wallplug_efficiency=1, file=None, x_offset=0, y_offset=0, x_angle=0, y_angle=0, waist_shift_x=0, waist_shift_y=0):
+
+        if file is not None and not os.path.exists(file):
+            raise FileNotFoundError(f"Error: The file '{file}' was not found.")
 
-        self.file = file
+        self._file = file
 
         super().__init__(length, charge, energy, accel_gradient, wallplug_efficiency, x_offset, y_offset, x_angle, y_angle, waist_shift_x, waist_shift_y)
+
+
+    @property
+    def file(self) -> str | None:
+        return self._file
+    @file.setter
+    def file(self, file_path : str):
+        if file_path is not None and not os.path.exists(file_path):
+            raise FileNotFoundError(f"Error: The file '{file_path}' was not found.")
+        else:
+            self._file = file_path
 
 
     def track(self, _=None, savedepth=0, runnable=None, verbose=False):

abel/classes/stage/impl/stage_basic.py

@@ -9,11 +9,12 @@
 
 class StageBasic(Stage):
 
-    def __init__(self, nom_accel_gradient=None, nom_energy_gain=None, plasma_density=None, driver_source=None, ramp_beta_mag=None, transformer_ratio=1):
+    def __init__(self, nom_accel_gradient=None, nom_energy_gain=None, plasma_density=None, driver_source=None, ramp_beta_mag=None, transformer_ratio=1, calc_evolution=False):
 
         super().__init__(nom_accel_gradient=nom_accel_gradient, nom_energy_gain=nom_energy_gain, plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag)
 
         self.transformer_ratio = transformer_ratio
+        self.calc_evolution = calc_evolution
 
 
     def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
@@ -85,7 +86,12 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
 
 
         # ========== Betatron oscillations ==========
-        beam.apply_betatron_motion(self.length_flattop, self.plasma_density, self.nom_energy_gain_flattop, x0_driver=driver0.x_offset(), y0_driver=driver0.y_offset())
+        deltaEs = np.full(len(beam.Es()), self.nom_energy_gain_flattop)  # Homogeneous energy gain for all macroparticles.
+        if self.calc_evolution:
+            _, evol = beam.apply_betatron_motion(self.length_flattop, self.plasma_density, deltaEs, x0_driver=driver0.x_offset(), y0_driver=driver0.y_offset(), calc_evolution=self.calc_evolution)
+            self.evolution.beam = evol
+        else:
+            beam.apply_betatron_motion(self.length_flattop, self.plasma_density, deltaEs, x0_driver=driver0.x_offset(), y0_driver=driver0.y_offset())
 
 
         # ========== Accelerate beam with homogeneous energy gain ==========

abel/utilities/relativity.py

@@ -54,8 +54,10 @@ def proper_velocity2gamma(u):
     return np.sign(u) * np.sqrt(1+(u/SI.c)**2)
 
 # convert proper velocity to momentum [kg m/s]
-def proper_velocity2momentum(u):
-    return gamma2momentum(proper_velocity2gamma(u))
+# def proper_velocity2momentum(u):
+#     return gamma2momentum(proper_velocity2gamma(u))
+def proper_velocity2momentum(u, m=SI.m_e):
+    return m*u
 
 # convert proper velocity to energy
 def proper_velocity2energy(u, unit=defaultUnitE, m=SI.m_e):

abel/utilities/statistics.py

@@ -20,14 +20,20 @@ def prct_clean2d(xs, ys, enable=True, cut_percentile=5):
     return x[mask], y[mask]
 
 def weighted_mean(values, weights, clean=False, cut_percentile=5):
-    mask = clean_mask(values, enable=clean, cut_percentile=5)
+    if np.sum(weights) == 0.0:
+        raise ZeroDivisionError("Weights sum to zero, can't be normalized.")
+    mask = clean_mask(values, enable=clean, cut_percentile=cut_percentile)
     return np.average(values[mask], weights=weights[mask])
 
 def weighted_std(values, weights, clean=False, cut_percentile=5):
-    mask = clean_mask(values, enable=clean, cut_percentile=5)
+    if np.sum(weights) == 0.0:
+        raise ZeroDivisionError("Weights sum to zero, can't be normalized.")
+    mask = clean_mask(values, enable=clean, cut_percentile=cut_percentile)
     return np.sqrt(np.cov(values[mask], aweights=weights[mask]))
 
 def weighted_cov(xs, ys, weights, clean=False, cut_percentile=5):
+    if np.sum(weights) == 0.0:
+        raise ZeroDivisionError("Weights sum to zero, can't be normalized.")
     mask_x = clean_mask(xs, enable=clean, cut_percentile=cut_percentile/2)
     mask_y = clean_mask(ys, enable=clean, cut_percentile=cut_percentile/2)
     mask = np.logical_and(mask_x, mask_y)

pyproject.toml

@@ -51,4 +51,6 @@ addopts = [
 markers = [
   "core: Test core functionality",
   "stageGeometry: Test stage geometry calculation functionality",
+  "sources: Test the various Source subclasses generate beams with desired properties",
+  "beam: Tests for the Beam class",
 ]