From 413fddeda019edd5341369f32c12747399495203 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Wed, 19 Mar 2025 23:30:09 +0100
Subject: [PATCH 01/51] Added methods for comparing beams and for printing out
 summary in the Beam class.

---
 abel/classes/beam.py | 129 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 129 insertions(+)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index d1d72215..5ee240b0 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -231,6 +231,101 @@ def deltas(self, pz0=None):
     def ts(self):
         return self.zs()/SI.c
     
+
+    def comp_beams(beam1, beam2, comp_location=False, rtol=1e-05, atol=1e-08):
+        """
+        Compare the phase spaces of two beams. Chekcks if all arrays are element-wise equal within given tolerances.
+
+        Parameters
+        ----------
+        beam1, beam2 : ABEL ``Beam`` object
+            Beams to be compared
+
+        comp_location: bool, optional
+            Flag for comparing the location of the beams. Default set to ``False``.
+            
+        rtol : float, optional
+            The relative tolerance parameter (see [1]_). Default set to 1e-5.
+            
+        rb_fit_obj : float, optional
+            The absolute tolerance parameter (see [1]_). Default set to 1e-8.
+        
+            
+        Returns
+        ----------
+        ``None``
+
+        
+        References
+        ----------
+        .. [1] ``numpy.allclose()`` https://numpy.org/doc/stable/reference/generated/numpy.allclose.html
+        """
+
+        if comp_location:
+            assert np.allclose(beam1.location, beam2.location, rtol, atol)
+        assert np.allclose(beam1.qs(), beam2.qs(), rtol, atol)
+        assert np.allclose(beam1.weightings(), beam2.weightings(), rtol, atol)
+        assert np.allclose(beam1.xs(), beam2.xs(), rtol, atol)
+        assert np.allclose(beam1.ys(), beam2.ys(), rtol, atol)
+        assert np.allclose(beam1.zs(), beam2.zs(), rtol, atol)
+        assert np.allclose(beam1.uxs(), beam2.uxs(), rtol, atol)
+        assert np.allclose(beam1.uys(), beam2.uys(), rtol, atol)
+        assert np.allclose(beam1.uzs(), beam2.uzs(), rtol, atol)
+        assert np.allclose(beam1.particle_mass, beam2.particle_mass, rtol, atol)
+
+
+    def comp_beam_params(beam1, beam2, comp_location=False):
+        """
+        Compare the parameters of two beams to see if they are equal within given tolerances.
+
+        Parameters
+        ----------
+        beam1, beam2 : ABEL ``Beam`` object
+            Beams to be compared
+
+        comp_location: bool, optional
+            Flag for comparing the location of the beams. Default set to ``False``.
+        
+            
+        Returns
+        ----------
+        ``None``
+
+        
+        References
+        ----------
+        .. [1] ``numpy.allclose()`` https://numpy.org/doc/stable/reference/generated/numpy.allclose.html
+        """
+
+        if comp_location:
+            assert np.allclose(beam1.location, beam2.location, rtol=0.0, atol=0.3)
+        assert np.allclose(beam1.particle_mass, beam2.particle_mass, rtol=1e-05, atol=1e-08)
+        assert np.allclose(beam1.total_particles(), beam2.total_particles(), rtol=1e-05, atol=1e-08)
+        assert np.allclose(beam1.charge(), beam2.charge(), rtol=1e-05, atol=1e-08)
+        assert np.allclose(beam1.energy(), beam2.energy(), rtol=0.0, atol=0.6e9)  # Usually rather large discrepancy in energy.
+        assert np.allclose(beam1.energy_spread(), beam2.energy_spread(), rtol=0.0, atol=0.6e9)
+        assert np.allclose(beam1.rel_energy_spread(), beam2.rel_energy_spread(), rtol=0.0, atol=5e-4)
+        assert np.allclose(beam1.bunch_length(), beam2.bunch_length(), rtol=0.0, atol=1e-6)
+        assert np.allclose(beam1.beam_size_x(), beam2.beam_size_x(), rtol=0.0, atol=0.5e-6)
+        assert np.allclose(beam1.beam_size_y(), beam2.beam_size_y(), rtol=0.0, atol=0.05e-6)
+        assert np.allclose(beam1.geom_emittance_x(), beam2.geom_emittance_x(), rtol=1e-05, atol=1e-08)
+        assert np.allclose(beam1.geom_emittance_y(), beam2.geom_emittance_y(), rtol=1e-05, atol=1e-08)
+        assert np.allclose(beam1.norm_emittance_x(), beam2.norm_emittance_x(), rtol=0.0, atol=1e-6)
+        assert np.allclose(beam1.norm_emittance_y(), beam2.norm_emittance_y(), rtol=0.0, atol=0.5e-6)
+        assert np.allclose(beam1.peak_current(), beam2.peak_current(), rtol=0.0, atol=0.7e3)
+        assert np.allclose(beam1.beta_x(), beam2.beta_x(), rtol=0.0, atol=50e-3)
+        assert np.allclose(beam1.beta_y(), beam2.beta_y(), rtol=0.0, atol=50e-3)
+        assert np.allclose(beam1.gamma_x(), beam2.gamma_x(), rtol=0.0, atol=0.1)
+        assert np.allclose(beam1.gamma_y(), beam2.gamma_y(), rtol=0.0, atol=0.1)        
+        assert np.allclose(beam1.z_offset(), beam2.z_offset(), rtol=0.0, atol=3e-6)
+        assert np.allclose(beam1.x_offset(), beam2.x_offset(), rtol=0.0, atol=3e-6)
+        assert np.allclose(beam1.y_offset(), beam2.y_offset(), rtol=0.0, atol=3e-6)
+        assert np.allclose(beam1.x_angle(), beam2.x_angle(), rtol=0.0, atol=3e-6)
+        assert np.allclose(beam1.y_angle(), beam2.y_angle(), rtol=0.0, atol=3e-6)
+        assert np.allclose(beam1.divergence_x(), beam2.divergence_x(), rtol=0.0, atol=0.2e-5)
+        assert np.allclose(beam1.divergence_y(), beam2.divergence_y(), rtol=0.0, atol=0.5e-6)
+
+
     # vector of transverse positions and angles: (x, x', y, y')
     def transverse_vector(self):
         vector = np.zeros((4,len(self)))
@@ -1669,3 +1764,37 @@ def density_map_diags(self):
         extent_energ[2] = extent_energ[2]/1e9  # [GeV]
         extent_energ[3] = extent_energ[3]/1e9  # [GeV]
         self.distribution_plot_2D(arr1=zs, arr2=Es, weights=weights, hist_bins=hist_bins, hist_range=hist_range_energ, axes=axs[2][2], extent=extent_energ, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=energ_lab, clab=r'$\partial^2 N/\partial \xi \partial\mathcal{E}$ [$\mathrm{m}^{-1}$ $\mathrm{eV}^{-1}$]', origin='lower', interpolation='nearest')
+
+
+        # ==================================================
+    def print_summary(self):
+
+        print('---------------------------------------------------')
+        print('Quantity \t\t\t\t Value')
+        print('---------------------------------------------------')
+        if hasattr(self, 'beam_name'):
+            print(f"Beam name:\t\t\t\t {self.beam_name :s}")
+        print(f"Number of macroparticles:\t\t {len(self) :d}")
+        print(f"Mean gamma:\t\t\t\t {self.gamma() :.3f}")
+        print(f"Mean energy [GeV]:\t\t\t {self.energy()/1e9 :.3f}")
+        print(f"rms energy spread [%]:\t\t\t {self.rel_energy_spread()*1e2 :.3f}\n")
+
+        print(f"Location [m]:\t\t\t\t {self.location :.3f}")
+        print(f"x offset [um]:\t\t\t\t {self.x_offset()*1e6 :.3f}")
+        print(f"y offset [um]:\t\t\t\t {self.y_offset()*1e6 :.3f}")
+        print(f"z offset [um]:\t\t\t\t {self.z_offset()*1e6 :.3f}\n")
+
+        print(f"x angular offset [urad]:\t\t {self.x_angle()*1e6 :.3f}")
+        print(f"y angular offset [urad]:\t\t {self.y_angle()*1e6 :.3f}\n")
+
+        print(f"Normalised x emittance [mm mrad]:\t {self.norm_emittance_x()*1e6 :.3f}")
+        print(f"Normalised y emittance [mm mrad]:\t {self.norm_emittance_y()*1e6 :.3f}")
+        print(f"Angular momentum [mm mrad]:\t\t {self.angular_momentum()*1e6 :.3f}\n")
+        print(f"x beta function [mm]:\t\t\t {self.beta_x()*1e3 :.3f} \t\t")
+        print(f"y beta function [mm]:\t\t\t {self.beta_y()*1e3 :.3f} \t\t\n")
+
+        print(f"x beam size [um]:\t\t\t {self.beam_size_x()*1e6 :.3f} \t\t\t")
+        print(f"y beam size [um]:\t\t\t {self.beam_size_y()*1e6 :.3f} \t\t\t")
+        print(f"rms beam length [um]:\t\t\t {self.bunch_length()*1e6 :.3f} \t\t")
+        print(f"Peak current [kA]:\t\t\t {self.peak_current()/1e3 :.3f} \t\t")
+        print('---------------------------------------------------')

From fd58b691aa616ae2f39e74a1e70f884ad32a066c Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Wed, 19 Mar 2025 23:31:14 +0100
Subject: [PATCH 02/51] Added test for testing SourceBasic

---
 tests/test_source.py | 79 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)
 create mode 100644 tests/test_source.py

diff --git a/tests/test_source.py b/tests/test_source.py
new file mode 100644
index 00000000..427910cf
--- /dev/null
+++ b/tests/test_source.py
@@ -0,0 +1,79 @@
+"""
+ABEL : Source class tests
+=======================================
+
+This file is a part of ABEL.
+Copyright 2022– C.A.Lindstrøm, B.Chen, O.G. Finnerud,
+D. Kallvik, E. Hørlyk, K.N. Sjobak, E.Adli, University of Oslo
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful, but
+WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program. If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import pytest
+from abel import *
+
+def test_SourceBasic2Beam():
+    "Check that the generated ``Beam`` from a ``SourceBasic`` has the desired properties."
+
+    np.random.seed(42)
+
+    plasma_density = 6.0e+20                                                      # [m^-3]
+    ramp_beta_mag = 5.0
+
+    main = SourceBasic()
+    main.bunch_length = 40.0e-06                                                  # [m], rms.
+    main.num_particles = 10000                                               
+    main.charge = -e * 1.0e10                                                     # [C]
+
+    # Energy parameters
+    main.energy = 3e9                                                             # [eV]
+    main.rel_energy_spread = 0.02                                                 # Relative rms energy spread
+
+    # Emittances
+    main.emit_nx, main.emit_ny = 90.0e-6, 0.32e-6                                 # [m rad], budget value
+
+    # Beta functions
+    main.beta_x = beta_matched(plasma_density, main.energy) * ramp_beta_mag       # [m]
+    main.beta_y = main.beta_x                                                     # [m]
+
+    # Offsets
+    main.z_offset = 0.00e-6                                                       # [m]
+    main.x_offset = 1.50e-6                                                       # [m]
+    main.y_offset = -1.0e-6                                                       # [m]
+    main.x_angle = -0.1e-6                                                         # [rad]
+    main.y_angle = 0.7e-6                                                         # [rad]
+
+    # Other
+    main.symmetrize_6d = True
+
+    # Track and check
+    beam = main.track()
+
+    beam.print_summary()
+
+    assert np.allclose(beam.particle_mass, SI.m_e, rtol=1e-05, atol=1e-08)
+    assert np.allclose(len(beam), main.num_particles, rtol=1e-05, atol=1e-08)
+    assert np.allclose(beam.charge(), main.charge, rtol=1e-05, atol=1e-08)
+    assert np.allclose(beam.energy(), main.energy, rtol=0.0, atol=0.1e9)
+    assert np.allclose(beam.rel_energy_spread(), main.rel_energy_spread, rtol=0.0, atol=0.005)
+    assert np.allclose(beam.bunch_length(), main.bunch_length, rtol=0.0, atol=2e-6)
+    assert np.allclose(beam.norm_emittance_x(), main.emit_nx, rtol=0.0, atol=1e-6)
+    assert np.allclose(beam.norm_emittance_y(), main.emit_ny, rtol=0.0, atol=0.05e-6)
+    assert np.allclose(beam.beta_x(), main.beta_x, rtol=0.0, atol=5e-3)
+    assert np.allclose(beam.beta_y(), main.beta_y, rtol=0.0, atol=5e-3)     
+    assert np.allclose(beam.z_offset(), main.z_offset, rtol=0.0, atol=1e-8)
+    assert np.allclose(beam.x_offset(), main.x_offset, rtol=0.0, atol=1e-8)
+    assert np.allclose(beam.y_offset(), main.y_offset, rtol=0.0, atol=1e-8)
+    assert np.allclose(beam.x_angle(), main.x_angle, rtol=0.0, atol=1e-8)
+    assert np.allclose(beam.y_angle(), main.y_angle, rtol=0.0, atol=1e-8)
\ No newline at end of file

From e3ea280be6caff0ccbf059e15ef150f3c3f046d7 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Thu, 20 Mar 2025 15:40:11 +0100
Subject: [PATCH 03/51] Updates to test_sources.py

---
 pyproject.toml        |   1 +
 tests/test_source.py  |  79 -----------------
 tests/test_sources.py | 194 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 195 insertions(+), 79 deletions(-)
 delete mode 100644 tests/test_source.py
 create mode 100644 tests/test_sources.py

diff --git a/pyproject.toml b/pyproject.toml
index debde464..8fd3fb88 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -51,4 +51,5 @@ addopts = [
 markers = [
   "core: Test core functionality",
   "stageGeometry: Test stage geometry calculation functionality",
+  "sources: Test the various Source subclasses generate beams with desired properties",
 ]
diff --git a/tests/test_source.py b/tests/test_source.py
deleted file mode 100644
index 427910cf..00000000
--- a/tests/test_source.py
+++ /dev/null
@@ -1,79 +0,0 @@
-"""
-ABEL : Source class tests
-=======================================
-
-This file is a part of ABEL.
-Copyright 2022– C.A.Lindstrøm, B.Chen, O.G. Finnerud,
-D. Kallvik, E. Hørlyk, K.N. Sjobak, E.Adli, University of Oslo
-
-This program is free software: you can redistribute it and/or modify
-it under the terms of the GNU General Public License as published by
-the Free Software Foundation, either version 3 of the License, or
-(at your option) any later version.
-
-This program is distributed in the hope that it will be useful, but
-WITHOUT ANY WARRANTY; without even the implied warranty of
-MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
-General Public License for more details.
-
-You should have received a copy of the GNU General Public License
-along with this program. If not, see <https://www.gnu.org/licenses/>.
-"""
-
-import pytest
-from abel import *
-
-def test_SourceBasic2Beam():
-    "Check that the generated ``Beam`` from a ``SourceBasic`` has the desired properties."
-
-    np.random.seed(42)
-
-    plasma_density = 6.0e+20                                                      # [m^-3]
-    ramp_beta_mag = 5.0
-
-    main = SourceBasic()
-    main.bunch_length = 40.0e-06                                                  # [m], rms.
-    main.num_particles = 10000                                               
-    main.charge = -e * 1.0e10                                                     # [C]
-
-    # Energy parameters
-    main.energy = 3e9                                                             # [eV]
-    main.rel_energy_spread = 0.02                                                 # Relative rms energy spread
-
-    # Emittances
-    main.emit_nx, main.emit_ny = 90.0e-6, 0.32e-6                                 # [m rad], budget value
-
-    # Beta functions
-    main.beta_x = beta_matched(plasma_density, main.energy) * ramp_beta_mag       # [m]
-    main.beta_y = main.beta_x                                                     # [m]
-
-    # Offsets
-    main.z_offset = 0.00e-6                                                       # [m]
-    main.x_offset = 1.50e-6                                                       # [m]
-    main.y_offset = -1.0e-6                                                       # [m]
-    main.x_angle = -0.1e-6                                                         # [rad]
-    main.y_angle = 0.7e-6                                                         # [rad]
-
-    # Other
-    main.symmetrize_6d = True
-
-    # Track and check
-    beam = main.track()
-
-    beam.print_summary()
-
-    assert np.allclose(beam.particle_mass, SI.m_e, rtol=1e-05, atol=1e-08)
-    assert np.allclose(len(beam), main.num_particles, rtol=1e-05, atol=1e-08)
-    assert np.allclose(beam.charge(), main.charge, rtol=1e-05, atol=1e-08)
-    assert np.allclose(beam.energy(), main.energy, rtol=0.0, atol=0.1e9)
-    assert np.allclose(beam.rel_energy_spread(), main.rel_energy_spread, rtol=0.0, atol=0.005)
-    assert np.allclose(beam.bunch_length(), main.bunch_length, rtol=0.0, atol=2e-6)
-    assert np.allclose(beam.norm_emittance_x(), main.emit_nx, rtol=0.0, atol=1e-6)
-    assert np.allclose(beam.norm_emittance_y(), main.emit_ny, rtol=0.0, atol=0.05e-6)
-    assert np.allclose(beam.beta_x(), main.beta_x, rtol=0.0, atol=5e-3)
-    assert np.allclose(beam.beta_y(), main.beta_y, rtol=0.0, atol=5e-3)     
-    assert np.allclose(beam.z_offset(), main.z_offset, rtol=0.0, atol=1e-8)
-    assert np.allclose(beam.x_offset(), main.x_offset, rtol=0.0, atol=1e-8)
-    assert np.allclose(beam.y_offset(), main.y_offset, rtol=0.0, atol=1e-8)
-    assert np.allclose(beam.x_angle(), main.x_angle, rtol=0.0, atol=1e-8)
-    assert np.allclose(beam.y_angle(), main.y_angle, rtol=0.0, atol=1e-8)
\ No newline at end of file
diff --git a/tests/test_sources.py b/tests/test_sources.py
new file mode 100644
index 00000000..6100e71d
--- /dev/null
+++ b/tests/test_sources.py
@@ -0,0 +1,194 @@
+"""
+ABEL : Source class tests
+=======================================
+
+This file is a part of ABEL.
+Copyright 2022– C.A.Lindstrøm, B.Chen, O.G. Finnerud,
+D. Kallvik, E. Hørlyk, K.N. Sjobak, E.Adli, University of Oslo
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful, but
+WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program. If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import pytest
+from abel import *
+
+def check_beam_source_parameters(beam, source):
+    assert np.allclose(beam.particle_mass, SI.m_e, rtol=1e-05, atol=1e-08)
+    assert np.allclose(len(beam), source.num_particles, rtol=1e-05, atol=1e-08)
+    assert np.allclose(beam.charge(), source.charge, rtol=1e-05, atol=1e-08)
+    assert np.allclose(beam.energy(), source.energy, rtol=0.0, atol=0.1e9)
+    assert np.allclose(beam.rel_energy_spread(), source.rel_energy_spread, rtol=0.0, atol=0.005)
+    assert np.allclose(beam.norm_emittance_x(), source.emit_nx, rtol=0.01, atol=1e-6)
+    assert np.allclose(beam.norm_emittance_y(), source.emit_ny, rtol=0.01, atol=0.05e-6)
+    assert np.allclose(beam.beta_x(), source.beta_x, rtol=0.01, atol=5e-3)
+    assert np.allclose(beam.beta_y(), source.beta_y, rtol=0.01, atol=5e-3)
+
+    if isinstance(source, SourceTrapezoid) or isinstance(source, SourceFlatTop):
+        assert np.allclose(beam.zs().max(), source.z_offset, rtol=0.01, atol=2e-6)
+    else:
+        assert np.allclose(beam.z_offset(), source.z_offset, rtol=0.0, atol=1e-8)
+    assert np.allclose(beam.x_offset(), source.x_offset, rtol=0.0, atol=1e-8)
+    assert np.allclose(beam.y_offset(), source.y_offset, rtol=0.0, atol=1e-8)
+    assert np.allclose(beam.x_angle(), source.x_angle, rtol=0.0, atol=1e-8)
+    assert np.allclose(beam.y_angle(), source.y_angle, rtol=0.0, atol=1e-8)
+
+    if isinstance(source, SourceTrapezoid) or isinstance(source, SourceFlatTop):
+        assert np.allclose(beam.zs().max()-beam.zs().min(), source.bunch_length, rtol=0.01, atol=0.0)
+    else:
+        assert np.allclose(beam.bunch_length(), source.bunch_length, rtol=0.01, atol=2e-6)
+
+
+@pytest.mark.sources
+def test_SourceBasic2Beam():
+    "Check that the generated ``Beam`` from a ``SourceBasic`` has the desired properties."
+
+    np.random.seed(42)
+
+    plasma_density = 6.0e+20                                                      # [m^-3]
+    ramp_beta_mag = 5.0
+
+    source = SourceBasic()
+    source.bunch_length = 40.0e-06                                                # [m], rms.
+    source.num_particles = 10000                                               
+    source.charge = -e * 1.0e10                                                   # [C]
+
+    # Energy parameters
+    source.energy = 3e9                                                           # [eV]
+    source.rel_energy_spread = 0.02                                               # Relative rms energy spread
+
+    # Emittances
+    source.emit_nx, source.emit_ny = 90.0e-6, 0.32e-6                             # [m rad], budget value
+
+    # Beta functions
+    source.beta_x = beta_matched(plasma_density, source.energy) * ramp_beta_mag   # [m]
+    source.beta_y = source.beta_x                                                 # [m]
+
+    # Offsets
+    source.z_offset = 0.00e-6                                                     # [m]
+    source.x_offset = 1.50e-6                                                     # [m]
+    source.y_offset = -1.0e-6                                                     # [m]
+    source.x_angle = -0.1e-6                                                      # [rad]
+    source.y_angle = 0.7e-6                                                       # [rad]
+
+    # Other
+    source.symmetrize_6d = True
+
+    # Track and check
+    beam = source.track()
+    beam.print_summary()
+    check_beam_source_parameters(beam, source)
+
+
+@pytest.mark.sources
+def test_SourceTrapezoid2Beam():
+    "Check that the generated ``Beam`` from a ``SourceTrapezoid`` has the desired properties."
+
+    np.random.seed(42)
+
+    source = SourceTrapezoid()
+    source.bunch_length = 1050e-6                                                 # [m], rms.
+    source.num_particles = 30000                                               
+    source.charge = -e * 5.0e10                                                   # [C]
+    source.current_head = 0.1e3                                                   # [A]
+
+    # Energy parameters
+    source.energy = 4.0e9                                                         # [eV]
+    source.rel_energy_spread = 0.01                                               # Relative rms energy spread
+
+    # Emittances
+    source.emit_nx, source.emit_ny = 50e-6, 100e-6                                # [m rad], budget value
+
+    # Beta functions
+    source.beta_x, source.beta_y = 0.5, 0.5                                       # [m]
+
+    # Offsets
+    source.z_offset = 1615e-6                                                     # [m]
+    source.x_offset = 1.50e-6                                                     # [m]
+    source.y_offset = -1.0e-6                                                     # [m]
+    source.x_angle = -0.1e-6                                                      # [rad]
+    source.y_angle = 0.7e-6                                                       # [rad]
+
+    # Other
+    source.symmetrize = True
+    #source.gaussian_blur = 50e-6                                                  # [m], excluding blur makes it easier to control z_offset.
+
+    # Track and check
+    beam = source.track()
+    beam.print_summary()
+    check_beam_source_parameters(beam, source)
+
+
+@pytest.mark.sources
+def test_SourceFlatopBeam():
+    "Check that the generated ``Beam`` from a ``SourceFlatTop`` has the desired properties."
+
+    np.random.seed(42)
+
+    source = SourceFlatTop()
+    source.bunch_length = 1050e-6                                                 # [m], rms.
+    source.num_particles = 30000                                               
+    source.charge = -e * 5.0e10                                                   # [C]
+
+    # Energy parameters
+    source.energy = 4.0e9                                                         # [eV]
+    source.rel_energy_spread = 0.01                                               # Relative rms energy spread
+
+    # Emittances
+    source.emit_nx, source.emit_ny = 50e-6, 100e-6                                # [m rad], budget value
+
+    # Beta functions
+    source.beta_x, source.beta_y = 0.5, 0.5                                       # [m]
+
+    # Offsets
+    source.z_offset = 1615e-6                                                     # [m]
+    source.x_offset = 1.50e-6                                                     # [m]
+    source.y_offset = -1.0e-6                                                     # [m]
+    source.x_angle = -0.1e-6                                                      # [rad]
+    source.y_angle = 0.7e-6                                                       # [rad]
+
+    # Other
+    source.symmetrize = True
+
+    # Track and check
+    beam = source.track()
+    beam.print_summary()
+    check_beam_source_parameters(beam, source)
+
+
+@pytest.mark.sources
+def test_SourceCapsule2Beam():
+    "Check that ``SourceCapsule`` retuns returns the same ``Beam`` as the input ``Beam``."
+
+    beam_file = 'tests/data/test_StagePrtclTransWakeInstability_beamline/test_baseline_linac/shot_000/beam_003_00048.558626.h5'
+
+    source = SourceCapsule()
+    ref_beam = Beam.load(beam_file)
+    source.beam = ref_beam
+    beam = source.track()
+    Beam.comp_beams(beam, ref_beam)
+
+
+@pytest.mark.sources
+def test_SourceFromFile2Beam():
+    "Check that ``SourceFromFile`` retuns returns the same ``Beam`` as the input ``Beam``."
+
+    beam_file = 'tests/data/test_StagePrtclTransWakeInstability_beamline/test_baseline_linac/shot_000/beam_003_00048.558626.h5'
+
+    source = SourceFromFile()
+    source.file = beam_file
+    beam = source.track()
+    ref_beam = Beam.load(beam_file)
+    Beam.comp_beams(beam, ref_beam)
+
+

From ea6b68921c407c72ce77fc770d18d04308ef9326 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Thu, 20 Mar 2025 23:53:51 +0100
Subject: [PATCH 04/51] Added test_beam.py for testing the Beam class.

---
 tests/test_beam.py | 246 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 246 insertions(+)
 create mode 100644 tests/test_beam.py

diff --git a/tests/test_beam.py b/tests/test_beam.py
new file mode 100644
index 00000000..26af401b
--- /dev/null
+++ b/tests/test_beam.py
@@ -0,0 +1,246 @@
+"""
+ABEL : Beam class tests
+=======================================
+
+This file is a part of ABEL.
+Copyright 2022– C.A.Lindstrøm, B.Chen, O.G. Finnerud,
+D. Kallvik, E. Hørlyk, K.N. Sjobak, E.Adli, University of Oslo
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful, but
+WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program. If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import pytest
+from abel import *
+import random
+
+
+############# Basic tests #############
+
+@pytest.mark.beam
+def test_initialization():
+    "Test if the ``Beam`` object initialises correctly."
+    beam = Beam(num_particles=1)
+    assert beam.num_bunches_in_train == 1
+    assert beam.bunch_separation == 0.0
+    assert beam.trackable_number == -1
+    assert beam.stage_number == 0
+    assert beam.location == 0
+
+
+@pytest.mark.beam
+def test_set_phase_space():
+    "Verify that the phase space is set correctly."
+    beam = Beam()
+    xs = np.random.rand(1000)
+    ys = np.random.rand(1000)
+    zs = np.random.rand(1000)
+    uxs = np.random.rand(1000)
+    uys = np.random.rand(1000)
+    uzs = np.random.rand(1000)
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    assert len(beam) == 1000
+    assert np.allclose(beam.xs(), xs)
+    assert np.allclose(beam.ys(), ys)
+    assert np.allclose(beam.zs(), zs)
+    assert np.allclose(beam.uxs(), uxs)
+    assert np.allclose(beam.uys(), uys)
+    assert np.allclose(beam.uzs(), uzs)
+    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/1000)
+    assert np.allclose(beam.charge(), Q)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/1000/(-SI.e))
+
+
+@pytest.mark.beam
+def test_reset_phase_space():
+    "Test reset_phase_space to ensure it initializes an 8xN zero matrix for the specified number of particles."
+    beam = Beam()
+    beam.reset_phase_space(10)
+    assert beam._Beam__phasespace.shape == (8, 10)
+    assert (beam._Beam__phasespace == 0).all()
+
+
+@pytest.mark.beam
+def test_delitem():
+    "Test deleting macroparticles by indices or masks."
+    beam = Beam()
+    beam.reset_phase_space(5)
+    beam.__delitem__([0, 2])
+    assert beam._Beam__phasespace.shape[1] == 3  # Two particles removed
+
+
+@pytest.mark.beam
+def test_remove_nans():
+    "Verify that remove_nans removes particles with any NaN value."
+    beam = Beam()
+    beam.reset_phase_space(10)
+    beam._Beam__phasespace[0, 1] = float('nan')  # Introduce NaN in one particle
+    beam._Beam__phasespace[0, 4] = float('nan')  # Introduce NaN in one particle
+    beam._Beam__phasespace[1, 4] = float('nan')  # Introduce NaN in one particle
+    beam._Beam__phasespace[1, 5] = float('nan')  # Introduce NaN in one particle
+    beam._Beam__phasespace[5, 7] = float('nan')  # Introduce NaN in one particle
+    beam.remove_nans()
+    assert len(beam) == 6
+
+
+@pytest.mark.beam
+def test_add_beams():
+    "Test the __add__ operator."
+    beam1 = Beam()
+    num_particles1 = 1021
+    beam1.reset_phase_space(num_particles1)
+    beam2 = Beam()
+    num_particles2 = 42
+    beam2.reset_phase_space(num_particles2)
+    beam3 = beam1 + beam2
+    assert len(beam3) == num_particles1 + num_particles2
+
+
+@pytest.mark.beam
+def test_iadd_beams():
+    "Test the __iadd__ operator."
+    beam1 = Beam()
+    num_particles1 = 3044
+    beam1.reset_phase_space(num_particles1)
+    beam2 = Beam()
+    num_particles2 = 4444
+    beam2.reset_phase_space(num_particles2)
+    beam1 += beam2
+    assert len(beam1) == num_particles1 + num_particles2
+
+
+@pytest.mark.beam
+def test_getitem():
+    "Test retrieving a single particle's data by index."
+
+    beam = Beam()
+    xs = np.random.rand(1000)
+    ys = np.random.rand(1000)
+    zs = np.random.rand(1000)
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs)
+    random_integer = random.randint(0, 999)
+    particle = beam[random_integer]
+    assert particle.shape == (8,)
+    assert np.allclose(particle[0], xs[random_integer])
+    assert np.allclose(particle[1], ys[random_integer])
+    assert np.allclose(particle[2], zs[random_integer])
+
+
+@pytest.mark.beam
+def test_len():
+    "Verify that __len__ returns the correct number of particles."
+    random_integer = random.randint(0, 30000)
+    beam = Beam()
+    beam.reset_phase_space(random_integer)
+    assert len(beam) == random_integer
+
+
+# @pytest.mark.beam
+# def test_str():
+#     "Test the __str__ method for correct formatting of beam properties."
+#     beam = Beam()
+#     beam.reset_phase_space(5)
+#     string_repr = str(beam)
+#     assert "Beam:" in string_repr
+#     assert "5 macroparticles" in string_repr
+
+
+@pytest.mark.beam
+def test_copy_particle_charge():
+
+    num_particles = 8051
+
+    beam1 = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.rand(num_particles)
+    Q = -SI.e * 1.0e10
+    beam1.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    beam2 = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.rand(num_particles)
+    Q = -SI.e * 3.4e10
+    ws = np.random.rand(num_particles)*Q/(-SI.e)
+    beam2.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs, weightings=ws)
+    median_charge = np.median(beam2._Beam__phasespace[6, :])  # Extract expected median charge
+
+    beam1.copy_particle_charge(beam2)
+    assert (beam1._Beam__phasespace[6, :] == median_charge).all()  # Verify charge is copied correctly
+
+
+# def test_scale_charge():
+#     "Test correct scaling of particle charges, scaling to a larger or smaller total charge. Edge case: scaling to zero charge."
+
+#     beam = Beam()
+#     beam.set_phase_space(1e-9, [1.0, 2.0], [0.5, 1.5], [0.2, 0.3])  # Total charge = 1e-9 C
+#     original_charges = beam._Beam__phasespace[6, :].copy()
+
+#     # Scale up charge
+#     beam.scale_charge(2e-9)  # Scale to 2 nC
+#     scaled_charges = beam._Beam__phasespace[6, :]
+#     assert np.isclose(scaled_charges.sum(), 2e-9)  # Check total charge
+#     assert np.allclose(scaled_charges / original_charges, 2.0)  # Charges scaled proportionally
+
+#     # Scale down charge
+#     beam.scale_charge(5e-10)  # Scale to 0.5 nC
+#     scaled_charges = beam._Beam__phasespace[6, :]
+#     assert np.isclose(scaled_charges.sum(), 5e-10)
+
+#     # Edge case: zero charge
+#     beam.scale_charge(0)
+#     scaled_charges = beam._Beam__phasespace[6, :]
+#     assert np.isclose(scaled_charges.sum(), 0)
+
+
+# def test_scale_energy():
+#     "Test correct scaling of particle energies, scaling to a higher or lower total energy. Edge case: scaling to zero energy."
+
+#     beam = Beam()
+#     beam.set_phase_space(1e-9, [1.0, 2.0], [0.5, 1.5], [0.2, 0.3], Es=[1e6, 2e6])  # Initialize with energies in eV
+#     original_energies = beam._Beam__phasespace[5, :].copy()  # Assuming energy stored in index 5
+
+#     # Scale up energy
+#     beam.scale_energy(4e6)  # Scale total energy to 4 MeV
+#     scaled_energies = beam._Beam__phasespace[5, :]
+#     assert np.isclose(scaled_energies.sum(), 4e6)
+#     assert np.allclose(scaled_energies / original_energies, 2.0)  # Energies scaled proportionally
+
+#     # Scale down energy
+#     beam.scale_energy(1e6)  # Scale total energy to 1 MeV
+#     scaled_energies = beam._Beam__phasespace[5, :]
+#     assert np.isclose(scaled_energies.sum(), 1e6)
+
+#     # Edge case: zero energy
+#     try:
+#         beam.scale_energy(0)
+#     except ValueError as e:
+#         assert str(e) == "Energy cannot be scaled to zero"
+
+
+############# Tests on bunch pattern #############
+
+
+
+
+############# ... #############
\ No newline at end of file

From 0e7d7d9383d6edb14c22a09097e63e330ec25fb7 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Fri, 21 Mar 2025 00:51:02 +0100
Subject: [PATCH 05/51] Small change in Beam.copy_particle_charge() to handle
 one of the beams being empty.

---
 abel/classes/beam.py |  2 ++
 tests/test_beam.py   | 45 +++++++++++++++++++++++---------------------
 2 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index 5ee240b0..80d87f9f 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -194,6 +194,8 @@ def weightings(self):
     
     # copy another beam's macroparticle charge
     def copy_particle_charge(self, beam):
+        if beam.__phasespace is None or self.__phasespace is None:
+            raise ValueError('One of the beams is empty.')
         self.set_qs(np.median(beam.qs()))
 
     def scale_charge(self, Q):
diff --git a/tests/test_beam.py b/tests/test_beam.py
index 26af401b..5577ae68 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -189,28 +189,31 @@ def test_copy_particle_charge():
     assert (beam1._Beam__phasespace[6, :] == median_charge).all()  # Verify charge is copied correctly
 
 
-# def test_scale_charge():
-#     "Test correct scaling of particle charges, scaling to a larger or smaller total charge. Edge case: scaling to zero charge."
+def test_scale_charge():
+    "Test correct scaling of particle charges, scaling to a larger or smaller total charge. Edge case: scaling to zero charge."
 
-#     beam = Beam()
-#     beam.set_phase_space(1e-9, [1.0, 2.0], [0.5, 1.5], [0.2, 0.3])  # Total charge = 1e-9 C
-#     original_charges = beam._Beam__phasespace[6, :].copy()
-
-#     # Scale up charge
-#     beam.scale_charge(2e-9)  # Scale to 2 nC
-#     scaled_charges = beam._Beam__phasespace[6, :]
-#     assert np.isclose(scaled_charges.sum(), 2e-9)  # Check total charge
-#     assert np.allclose(scaled_charges / original_charges, 2.0)  # Charges scaled proportionally
-
-#     # Scale down charge
-#     beam.scale_charge(5e-10)  # Scale to 0.5 nC
-#     scaled_charges = beam._Beam__phasespace[6, :]
-#     assert np.isclose(scaled_charges.sum(), 5e-10)
-
-#     # Edge case: zero charge
-#     beam.scale_charge(0)
-#     scaled_charges = beam._Beam__phasespace[6, :]
-#     assert np.isclose(scaled_charges.sum(), 0)
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs)
+    original_charges = copy.deepcopy(beam.qs())
+
+    # Scale up charge
+    beam.scale_charge(-SI.e * 2.0e10)
+    scaled_charges = beam.qs()
+    assert np.isclose(scaled_charges.sum(), -SI.e * 2.0e10)  # Check total charge
+    assert np.allclose(scaled_charges / original_charges, 2.0)  # Charges scaled proportionally
+
+    # Scale down charge
+    beam.scale_charge(-SI.e * 0.5e10)
+    assert np.isclose(scaled_charges.sum(), -SI.e * 0.5e10)
+
+    # Edge case: zero charge
+    beam.scale_charge(0)
+    scaled_charges = beam.qs()
+    assert np.isclose(scaled_charges.sum(), 0)
 
 
 # def test_scale_energy():

From 9543833c1b1d0f1d672cafaad70e3f50ba03a1cd Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Fri, 21 Mar 2025 00:52:46 +0100
Subject: [PATCH 06/51] Added more tests in test_beam.py.

---
 tests/test_beam.py | 54 +++++++++++++++++++++++++++-------------------
 1 file changed, 32 insertions(+), 22 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 5577ae68..dd4992f0 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -193,6 +193,7 @@ def test_scale_charge():
     "Test correct scaling of particle charges, scaling to a larger or smaller total charge. Edge case: scaling to zero charge."
 
     beam = Beam()
+    num_particles = 4352
     xs = np.random.rand(num_particles)
     ys = np.random.rand(num_particles)
     zs = np.random.rand(num_particles)
@@ -216,29 +217,38 @@ def test_scale_charge():
     assert np.isclose(scaled_charges.sum(), 0)
 
 
-# def test_scale_energy():
-#     "Test correct scaling of particle energies, scaling to a higher or lower total energy. Edge case: scaling to zero energy."
+def test_scale_energy():
+    "Test correct scaling of particle energies, scaling to a higher or lower total energy. Edge case: scaling to zero energy."
 
-#     beam = Beam()
-#     beam.set_phase_space(1e-9, [1.0, 2.0], [0.5, 1.5], [0.2, 0.3], Es=[1e6, 2e6])  # Initialize with energies in eV
-#     original_energies = beam._Beam__phasespace[5, :].copy()  # Assuming energy stored in index 5
-
-#     # Scale up energy
-#     beam.scale_energy(4e6)  # Scale total energy to 4 MeV
-#     scaled_energies = beam._Beam__phasespace[5, :]
-#     assert np.isclose(scaled_energies.sum(), 4e6)
-#     assert np.allclose(scaled_energies / original_energies, 2.0)  # Energies scaled proportionally
-
-#     # Scale down energy
-#     beam.scale_energy(1e6)  # Scale total energy to 1 MeV
-#     scaled_energies = beam._Beam__phasespace[5, :]
-#     assert np.isclose(scaled_energies.sum(), 1e6)
-
-#     # Edge case: zero energy
-#     try:
-#         beam.scale_energy(0)
-#     except ValueError as e:
-#         assert str(e) == "Energy cannot be scaled to zero"
+    beam = Beam()
+    num_particles = 4352
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    #Es = np.ones_like(num_particles) * 1e9
+    Es = np.random.normal(1e9, 0.02*1e9, num_particles)
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs, Es=Es)
+    original_energies = copy.deepcopy(beam.Es())
+
+    # Scale up energy
+    beam.scale_energy(4e9)
+    scaled_energies = beam.Es()
+    scaled_mean_energy = weighted_mean(scaled_energies, beam.weightings())
+    assert np.isclose(scaled_mean_energy, 4e9, rtol=0.01, atol=1e8)
+    assert np.allclose(scaled_energies / original_energies, 4.0, rtol=0.01, atol=0.1)  # Energies scaled proportionally
+
+    # Scale down energy
+    beam.scale_energy(1.7e8)
+    scaled_energies = beam.Es()
+    scaled_mean_energy = weighted_mean(scaled_energies, beam.weightings())
+    assert np.isclose(scaled_mean_energy, 1.7e8, rtol=0.01, atol=1e8)
+
+    # # Edge case: zero energy
+    # try:
+    #     beam.scale_energy(0)
+    # except ValueError as e:
+    #     assert str(e) == "Energy cannot be scaled to zero"
 
 
 ############# Tests on bunch pattern #############

From 5eae0e884c093c6882b065b49d9cc7eea8fc07b5 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Mon, 24 Mar 2025 16:01:54 +0100
Subject: [PATCH 07/51] Added Beam class initiation tests to test_beam.py and
 added another tag to pyproject.toml.

---
 pyproject.toml     |   1 +
 tests/test_beam.py | 136 +++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 125 insertions(+), 12 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8fd3fb88..b346be53 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,4 +52,5 @@ 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",
 ]
diff --git a/tests/test_beam.py b/tests/test_beam.py
index dd4992f0..44436355 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -42,25 +42,136 @@ def test_initialization():
 def test_set_phase_space():
     "Verify that the phase space is set correctly."
     beam = Beam()
-    xs = np.random.rand(1000)
-    ys = np.random.rand(1000)
-    zs = np.random.rand(1000)
-    uxs = np.random.rand(1000)
-    uys = np.random.rand(1000)
-    uzs = np.random.rand(1000)
+    num_particles = 10042
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.rand(num_particles)
     Q = -SI.e * 1.0e10
     beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
 
-    assert len(beam) == 1000
+    assert len(beam) == num_particles
+    assert np.isclose(beam.particle_mass, SI.m_e)
     assert np.allclose(beam.xs(), xs)
     assert np.allclose(beam.ys(), ys)
     assert np.allclose(beam.zs(), zs)
     assert np.allclose(beam.uxs(), uxs)
     assert np.allclose(beam.uys(), uys)
     assert np.allclose(beam.uzs(), uzs)
-    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/1000)
-    assert np.allclose(beam.charge(), Q)
-    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/1000/(-SI.e))
+    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
+    assert np.isclose(beam.charge(), Q)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
+
+    assert np.allclose(beam.xps(), uxs/uzs)
+    assert np.allclose(beam.yps(), uys/uzs)
+    assert np.allclose(beam.pxs(), SI.m_e*uxs)
+    assert np.allclose(beam.pys(), SI.m_e*uys)
+    assert np.allclose(beam.pzs(), SI.m_e*uzs)
+    assert np.allclose(beam.Es(), np.sqrt((SI.m_e*uzs*SI.c)**2 + (SI.m_e*SI.c**2)**2)/SI.e )
+
+
+@pytest.mark.beam
+def test_set_phase_space2():
+    "Verify that the phase space is set correctly."
+    beam = Beam()
+    num_particles = 10042
+    xs = np.random.normal(0.0, 2.0e-6, num_particles)
+    ys = np.random.normal(0.0, 1.0e-6, num_particles)
+    zs = np.random.normal(0.0, 50.0e-6, num_particles)
+    xps = np.random.normal(0.0, 1.5e-4, num_particles)
+    yps = np.random.normal(0.0, 1.0e-5, num_particles)
+    Es = np.random.normal(500e9, 0.02*500e9, num_particles)
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs, xps=xps, yps=yps, Es=Es)
+
+    assert len(beam) == num_particles
+    assert np.isclose(beam.particle_mass, SI.m_e)
+    assert np.allclose(beam.xs(), xs)
+    assert np.allclose(beam.ys(), ys)
+    assert np.allclose(beam.zs(), zs)
+    assert np.allclose(beam.xps(), xps, rtol=1e-05, atol=1e-08)
+    assert np.allclose(beam.yps(), yps, rtol=1e-05, atol=1e-08)
+    assert np.allclose(beam.Es(), Es)
+    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
+    assert np.isclose(beam.charge(), Q)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
+
+    assert np.allclose(beam.uzs(), np.sqrt((Es*SI.e/SI.m_e/SI.c)**2 - SI.c**2) )
+    assert np.allclose(beam.uxs(), xps*beam.uzs())
+    assert np.allclose(beam.uys(), yps*beam.uzs())
+    assert np.allclose(beam.pxs(), SI.m_e*xps*beam.uzs(), rtol=1e-05, atol=1e-25)
+    assert np.allclose(beam.pys(), SI.m_e*yps*beam.uzs(), rtol=1e-05, atol=1e-25)
+    assert np.allclose(beam.pzs(), SI.m_e*beam.uzs(), rtol=1e-05, atol=1e-19)
+
+
+@pytest.mark.beam
+def test_set_phase_space3():
+    "Verify that the phase space is set correctly."
+    beam = Beam()
+    num_particles = 10042
+    xs = np.random.normal(1.0, 2.4e-6, num_particles)
+    ys = np.random.normal(0.3, 1.1e-6, num_particles)
+    zs = np.random.normal(10.0, 42.0e-6, num_particles)
+    pxs = np.random.normal(0.0, 5.0e-22, num_particles)
+    pys = np.random.normal(0.0, 5.0e-23, num_particles)
+    pzs = np.random.normal(2.67e-16, 5.4e-18, num_particles)  # Corresponds to 500 GeV.
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs, pxs=pxs, pys=pys, pzs=pzs)
+
+    assert len(beam) == num_particles
+    assert np.isclose(beam.particle_mass, SI.m_e)
+    assert np.allclose(beam.xs(), xs)
+    assert np.allclose(beam.ys(), ys)
+    assert np.allclose(beam.zs(), zs)
+    assert np.allclose(beam.pxs(), pxs, rtol=1e-05, atol=1e-25)
+    assert np.allclose(beam.pys(), pys, rtol=1e-05, atol=1e-25)
+    assert np.allclose(beam.pzs(), pzs, rtol=1e-05, atol=1e-19)
+    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
+    assert np.isclose(beam.charge(), Q)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
+
+    assert np.allclose(beam.xps(), pxs/pzs)
+    assert np.allclose(beam.yps(), pys/pzs)
+    assert np.allclose(beam.uxs(), pxs/SI.m_e)
+    assert np.allclose(beam.uys(), pys/SI.m_e)
+    assert np.allclose(beam.uzs(), pzs/SI.m_e)
+    assert np.allclose(beam.Es(), np.sqrt((pzs*SI.c)**2 + (SI.m_e*SI.c**2)**2)/SI.e )
+
+
+@pytest.mark.beam
+def test_set_phase_space4():
+    "Verify that the phase space is set correctly."
+    beam = Beam()
+    num_particles = 10042
+    xs = np.random.normal(0.0, 2.0e-6, num_particles)
+    ys = np.random.normal(0.0, 1.0e-6, num_particles)
+    zs = np.random.normal(0.0, 50.0e-6, num_particles)
+    uxs = np.random.normal(0.0, 250.0e6, num_particles)
+    yps = np.random.normal(0.0, 1.0e-5, num_particles)
+    pzs = np.random.normal(1.6e-18, 3.0e-20, num_particles)
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs=uxs, yps=yps, pzs=pzs)
+
+    assert len(beam) == num_particles
+    assert np.isclose(beam.particle_mass, SI.m_e)
+    assert np.allclose(beam.xs(), xs)
+    assert np.allclose(beam.ys(), ys)
+    assert np.allclose(beam.zs(), zs)
+    assert np.allclose(beam.uxs(), uxs)
+    assert np.allclose(beam.yps(), yps)
+    assert np.allclose(beam.pzs(), pzs, rtol=1e-05, atol=1e-25)
+    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
+    assert np.isclose(beam.charge(), Q)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
+
+    assert np.allclose(beam.uzs(), pzs/SI.m_e)
+    assert np.allclose(beam.xps(), uxs/beam.uzs())
+    assert np.allclose(beam.uys(), yps*beam.uzs())
+    assert np.allclose(beam.pxs(), SI.m_e*uxs, rtol=1e-05, atol=1e-25)
+    assert np.allclose(beam.pys(), SI.m_e*yps*beam.uzs(), rtol=1e-05, atol=1e-25)    
+    assert np.allclose(beam.Es(), np.sqrt((pzs*SI.c)**2 + (SI.m_e*SI.c**2)**2)/SI.e )
 
 
 @pytest.mark.beam
@@ -189,6 +300,7 @@ def test_copy_particle_charge():
     assert (beam1._Beam__phasespace[6, :] == median_charge).all()  # Verify charge is copied correctly
 
 
+@pytest.mark.beam
 def test_scale_charge():
     "Test correct scaling of particle charges, scaling to a larger or smaller total charge. Edge case: scaling to zero charge."
 
@@ -217,6 +329,7 @@ def test_scale_charge():
     assert np.isclose(scaled_charges.sum(), 0)
 
 
+@pytest.mark.beam
 def test_scale_energy():
     "Test correct scaling of particle energies, scaling to a higher or lower total energy. Edge case: scaling to zero energy."
 
@@ -225,7 +338,6 @@ def test_scale_energy():
     xs = np.random.rand(num_particles)
     ys = np.random.rand(num_particles)
     zs = np.random.rand(num_particles)
-    #Es = np.ones_like(num_particles) * 1e9
     Es = np.random.normal(1e9, 0.02*1e9, num_particles)
     Q = -SI.e * 1.0e10
     beam.set_phase_space(Q, xs, ys, zs, Es=Es)
@@ -256,4 +368,4 @@ def test_scale_energy():
 
 
 
-############# ... #############
\ No newline at end of file
+############# Tests of rotation methods #############
\ No newline at end of file

From bcc9aa7a7a6824f03cbfc6c0c727b519ff89075d Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Mon, 24 Mar 2025 16:02:41 +0100
Subject: [PATCH 08/51] Edited __init__.py to add SourceCapsule.

---
 abel/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/abel/__init__.py b/abel/__init__.py
index f5aaad51..553baebf 100644
--- a/abel/__init__.py
+++ b/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 *

From 18887a3e936cd6d06f7892b1fa64ec3f7d56dbdc Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 25 Mar 2025 00:39:06 +0100
Subject: [PATCH 09/51] Added a function for setting up a SourceBasic in
 test_beam.py. Also added more subtests to the Beam initialisation tests.

---
 tests/test_beam.py | 81 +++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 69 insertions(+), 12 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 44436355..260d0e3c 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -25,6 +25,34 @@
 import random
 
 
+def setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, bunch_length=40.0e-06, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0):
+    main = SourceBasic()
+    main.bunch_length = bunch_length                                              # [m], rms. Standard value
+    main.num_particles = 10000                                               
+    main.charge = -e * 1.0e10                                                     # [C]
+
+    # Energy parameters
+    main.energy = energy                                                          # [eV], HALHF v2 last stage nominal input energy
+    main.rel_energy_spread = 0.02                                                 # Relative rms energy spread
+
+    # Emittances
+    main.emit_nx, main.emit_ny = 15e-6, 0.1e-6                                    # [m rad]
+
+    # Beta functions
+    main.beta_x = beta_matched(plasma_density, main.energy) * ramp_beta_mag       # [m]
+    main.beta_y = main.beta_x                                                     # [m]
+
+    # Offsets
+    main.z_offset = z_offset                                                      # [m]
+    main.x_offset = x_offset                                                      # [m]
+    main.y_offset = y_offset                                                      # [m]
+
+    # Other
+    main.symmetrize_6d = True
+
+    return main
+
+
 ############# Basic tests #############
 
 @pytest.mark.beam
@@ -54,6 +82,8 @@ def test_set_phase_space():
 
     assert len(beam) == num_particles
     assert np.isclose(beam.particle_mass, SI.m_e)
+    assert np.isclose(beam.charge(), Q)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
     assert np.allclose(beam.xs(), xs)
     assert np.allclose(beam.ys(), ys)
     assert np.allclose(beam.zs(), zs)
@@ -61,8 +91,7 @@ def test_set_phase_space():
     assert np.allclose(beam.uys(), uys)
     assert np.allclose(beam.uzs(), uzs)
     assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
-    assert np.isclose(beam.charge(), Q)
-    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
+
 
     assert np.allclose(beam.xps(), uxs/uzs)
     assert np.allclose(beam.yps(), uys/uzs)
@@ -77,17 +106,34 @@ def test_set_phase_space2():
     "Verify that the phase space is set correctly."
     beam = Beam()
     num_particles = 10042
-    xs = np.random.normal(0.0, 2.0e-6, num_particles)
-    ys = np.random.normal(0.0, 1.0e-6, num_particles)
-    zs = np.random.normal(0.0, 50.0e-6, num_particles)
-    xps = np.random.normal(0.0, 1.5e-4, num_particles)
-    yps = np.random.normal(0.0, 1.0e-5, num_particles)
+    xs = np.random.normal(1.1, 2.0e-6, num_particles)
+    ys = np.random.normal(0.3, 1.0e-6, num_particles)
+    zs = np.random.normal(5.6, 50.0e-6, num_particles)
+    xps = np.random.normal(1.1e-5, 1.5e-7, num_particles)
+    yps = np.random.normal(3.2e-6, 1.0e-8, num_particles)
     Es = np.random.normal(500e9, 0.02*500e9, num_particles)
+    #weightings =  # TODO: make an array for varying weightings
     Q = -SI.e * 1.0e10
     beam.set_phase_space(Q, xs, ys, zs, xps=xps, yps=yps, Es=Es)
 
     assert len(beam) == num_particles
     assert np.isclose(beam.particle_mass, SI.m_e)
+    assert np.isclose(beam.charge(), Q)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
+    assert np.isclose(beam.energy(), 500e9, rtol=0.0001*500e9, atol=1e9)
+    assert np.isclose(beam.gamma(), 500e9*SI.e/SI.m_e/SI.c**2, rtol=0.0001*500e9*SI.e/SI.m_e/SI.c**2, atol=200)
+    assert np.isclose(beam.rel_energy_spread(), 0.02, rtol=0.001, atol=0.001)
+    assert np.isclose(beam.z_offset(), 5.6)
+    assert np.isclose(beam.x_offset(), 1.1)
+    assert np.isclose(beam.y_offset(), 0.3)
+    assert np.isclose(beam.bunch_length(), 50.0e-6, rtol=0.005, atol=0.03)
+    assert np.isclose(beam.beam_size_x(), 2.0e-6, rtol=0.005, atol=0.03)
+    assert np.isclose(beam.beam_size_y(), 1.0e-6, rtol=0.005, atol=0.03)
+    assert np.isclose(beam.x_angle(), 1.1e-5)
+    assert np.isclose(beam.y_angle(), 3.2e-6)
+    assert np.isclose(beam.divergence_x(), 1.5e-7)
+    assert np.isclose(beam.divergence_y(), 1.0e-8)
+
     assert np.allclose(beam.xs(), xs)
     assert np.allclose(beam.ys(), ys)
     assert np.allclose(beam.zs(), zs)
@@ -95,8 +141,6 @@ def test_set_phase_space2():
     assert np.allclose(beam.yps(), yps, rtol=1e-05, atol=1e-08)
     assert np.allclose(beam.Es(), Es)
     assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
-    assert np.isclose(beam.charge(), Q)
-    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
 
     assert np.allclose(beam.uzs(), np.sqrt((Es*SI.e/SI.m_e/SI.c)**2 - SI.c**2) )
     assert np.allclose(beam.uxs(), xps*beam.uzs())
@@ -122,6 +166,8 @@ def test_set_phase_space3():
 
     assert len(beam) == num_particles
     assert np.isclose(beam.particle_mass, SI.m_e)
+    assert np.isclose(beam.charge(), Q)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
     assert np.allclose(beam.xs(), xs)
     assert np.allclose(beam.ys(), ys)
     assert np.allclose(beam.zs(), zs)
@@ -129,8 +175,6 @@ def test_set_phase_space3():
     assert np.allclose(beam.pys(), pys, rtol=1e-05, atol=1e-25)
     assert np.allclose(beam.pzs(), pzs, rtol=1e-05, atol=1e-19)
     assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
-    assert np.isclose(beam.charge(), Q)
-    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
 
     assert np.allclose(beam.xps(), pxs/pzs)
     assert np.allclose(beam.yps(), pys/pzs)
@@ -368,4 +412,17 @@ def test_scale_energy():
 
 
 
-############# Tests of rotation methods #############
\ No newline at end of file
+############# Tests of rotation methods #############
+
+#def test_x_tilt_angle():
+
+def test_rotate_coord_sys_3D():
+    "Test correct rotation of beam coordinate system."
+
+    source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
+    beam = source.track()
+    x_axis = np.array([0, 1, 0])  # Axis as an unit vector. Axis permutaton is zxy.
+    y_axis = np.array([0, 0, 1])
+
+    beam.rotate_coord_sys_3D(y_axis, np.pi/2, x_axis, 0.0, invert=False)  # Rotate beam 90 degrees around the y-axis.
+    assert np.isclose(beam.x_tilt_angle(), np.pi/2)
\ No newline at end of file

From 1e4f5c0f793a68eb78d8e09f65625082c3519290 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 25 Mar 2025 01:45:21 +0100
Subject: [PATCH 10/51] Saving some tests for beam coordinate rotation. Work in
 progress...

---
 tests/test_beam.py | 52 ++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 48 insertions(+), 4 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 260d0e3c..a335791d 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -112,7 +112,7 @@ def test_set_phase_space2():
     xps = np.random.normal(1.1e-5, 1.5e-7, num_particles)
     yps = np.random.normal(3.2e-6, 1.0e-8, num_particles)
     Es = np.random.normal(500e9, 0.02*500e9, num_particles)
-    #weightings =  # TODO: make an array for varying weightings
+    #weightings =  # TODO: make an array with varying weightings
     Q = -SI.e * 1.0e10
     beam.set_phase_space(Q, xs, ys, zs, xps=xps, yps=yps, Es=Es)
 
@@ -413,8 +413,51 @@ def test_scale_energy():
 
 
 ############# Tests of rotation methods #############
+# def test_slice_centroids():
+#     source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
+#     beam = source.track()
+#     beam_quant = beam.xs()  # Use x positions as the beam quantity
+#     x_slices, z_centroids = beam.slice_centroids(beam_quant, bin_number=5)
+
+#     # Check that the output shapes are correct
+#     assert x_slices.shape[0] == z_centroids.shape[0]
+#     assert len(x_slices) == 5  # Should match the number of bins
+
+#     # Check that the centroids are within the expected range
+#     assert np.all(z_centroids >= 0)
+#     assert np.all(z_centroids <= 10)
+
+
+def test_x_tilt_angle():
+    "Test the retrieval of the beam tilt angle in the zx-plane."
+    source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
+    beam = source.track()
+    nominal_x_angle = random.uniform(-np.pi/2, np.pi/2)
+    beam.set_xs(beam.zs() * np.tan(nominal_x_angle))  # Add a tilt in x
+    x_angle = beam.x_tilt_angle()
+    assert np.isclose(x_angle, nominal_x_angle, rtol=1e-7, atol=1e-3)
+
+
+def test_y_tilt_angle():
+    "Test the retrieval of the beam tilt angle in the zy-plane."
+    source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
+    beam = source.track()
+    nominal_y_angle = random.uniform(-np.pi/2, np.pi/2)
+    beam.set_ys(beam.zs() * np.tan(nominal_y_angle))  # Add a tilt in x
+    y_angle = beam.y_tilt_angle()
+    assert np.isclose(y_angle, nominal_y_angle, rtol=1e-7, atol=1e-3)
+
+
+def test_xy_rotate_coord_sys1():
+    "Test correct rotation of beam coordinate system around the y-axis."
+    source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
+    beam = source.track()
+    old_beam = copy.deepcopy(beam)
+    nom_x_angle = random.uniform(-np.pi/2, np.pi/2)
+    beam.xy_rotate_coord_sys(x_angle=nom_x_angle, y_angle=None, invert=False)
+    assert np.allclose(beam.xs(), -old_beam.zs()*np.sin(nom_x_angle) + old_beam.xs()*np.cos(nom_x_angle), rtol=1e-7, atol=1e-13)
+    assert np.allclose(beam.zs(), old_beam.zs()*np.cos(nom_x_angle) + old_beam.xs()*np.sin(nom_x_angle), rtol=1e-7, atol=1e-13)
 
-#def test_x_tilt_angle():
 
 def test_rotate_coord_sys_3D():
     "Test correct rotation of beam coordinate system."
@@ -424,5 +467,6 @@ def test_rotate_coord_sys_3D():
     x_axis = np.array([0, 1, 0])  # Axis as an unit vector. Axis permutaton is zxy.
     y_axis = np.array([0, 0, 1])
 
-    beam.rotate_coord_sys_3D(y_axis, np.pi/2, x_axis, 0.0, invert=False)  # Rotate beam 90 degrees around the y-axis.
-    assert np.isclose(beam.x_tilt_angle(), np.pi/2)
\ No newline at end of file
+    nom_x_angle = random.uniform(-np.pi/2, np.pi/2)
+    beam.rotate_coord_sys_3D(y_axis, nom_x_angle, x_axis, 0.0, invert=False)  # Rotate beam 90 degrees around the y-axis.
+    assert np.isclose(beam.x_tilt_angle(), nom_x_angle)
\ No newline at end of file

From 02a45a7046d313aed3337a2c001174a384687e37 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 25 Mar 2025 18:07:19 +0100
Subject: [PATCH 11/51] Added controls to the Beam class to ensure that energy,
 longitudinal proper velocity and momentum are above a given threshold. Also
 added controls for the number of particles, number of bunches in a train and
 bunch separation.

---
 abel/classes/beam.py | 38 +++++++++++++++++++++++++++++++++++---
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index 80d87f9f..d28759fa 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -5,7 +5,7 @@
 from abel.CONFIG import CONFIG
 from types import SimpleNamespace
 import scipy.constants as SI
-from abel.utilities.relativity import energy2proper_velocity, proper_velocity2energy, momentum2proper_velocity, proper_velocity2momentum, proper_velocity2gamma, energy2gamma, gamma2proper_velocity
+from abel.utilities.relativity import energy2proper_velocity, proper_velocity2energy, momentum2proper_velocity, proper_velocity2momentum, proper_velocity2gamma, energy2gamma, gamma2momentum
 from abel.utilities.statistics import weighted_mean, weighted_std, weighted_cov
 from abel.utilities.plasma_physics import k_p, wave_breaking_field, beta_matched
 from abel.physics_models.hills_equation import evolve_hills_equation_analytic
@@ -21,6 +21,14 @@ class Beam():
     
     def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1, bunch_separation=0.0):
 
+        # check the inputs
+        if num_particles < 1:
+            raise ValueError('num_particles cannot be lower than 1.')
+        if num_bunches_in_train < 1:
+            raise ValueError('num_bunches_in_train cannot be lower than 1.')
+        if bunch_separation < 0.0:
+            raise ValueError('bunch_separation cannot be negative.')
+
         # the phase space variable is private
         if phasespace is not None:
             self.__phasespace = phasespace
@@ -38,6 +46,9 @@ def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1,
     
     # reset phase space
     def reset_phase_space(self, num_particles):
+        if num_particles < 1:
+            raise ValueError('num_particles cannot be lower than 1.')
+        
         self.__phasespace = np.zeros((8, num_particles))
     
     # filter out macroparticles based on a mask (true means delete)
@@ -62,13 +73,27 @@ def set_phase_space(self, Q, xs, ys, zs, uxs=None, uys=None, uzs=None, pxs=None,
         self.set_xs(xs)
         self.set_ys(ys)
         self.set_zs(zs)
+
+        # minimum thresholds for energy, uz and pz
+        energy_thres = 100*particle_mass*SI.c**2/SI.e  # [eV], 100 * particle rest energy.
+        uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=particle_mass)
+        pz_thres = gamma2momentum(energy2gamma(energy_thres, unit='eV', m=particle_mass))
         
         # add momenta
         if uzs is None:
+            
             if pzs is not None:
+                if np.any(pzs < pz_thres):
+                    raise ValueError('pzs contains values that are too small.')
                 uzs = momentum2proper_velocity(pzs)
+
             elif Es is not None:
+                if np.any(Es < energy_thres):
+                    raise ValueError('Es contains values that are too small.')
                 uzs = energy2proper_velocity(Es)
+        else:
+            if np.any(uzs < uz_thres):
+                raise ValueError('uzs contains values that are too small.')
         self.__phasespace[5,:] = uzs
         
         if uxs is None:
@@ -175,6 +200,10 @@ def set_uxs(self, uxs):
     def set_uys(self, uys):
         self.__phasespace[4,:] = uys
     def set_uzs(self, uzs):
+        energy_thres = 100*self.particle_mass*SI.c**2/SI.e  # [eV], 100 * particle rest energy.
+        uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=self.particle_mass)
+        if np.any(uzs < uz_thres):
+            raise ValueError('uzs contains values that are too small.')
         self.__phasespace[5,:] = uzs
         
     def set_xps(self, xps):
@@ -182,6 +211,9 @@ def set_xps(self, xps):
     def set_yps(self, yps):
         self.set_uys(yps*self.uzs())
     def set_Es(self, Es):
+        energy_thres = 100*self.particle_mass*SI.c**2/SI.e  # [eV], 100 * particle rest energy.
+        if np.any(Es < energy_thres):
+            raise ValueError('Es contains values that are too small.')
         self.set_uzs(energy2proper_velocity(Es))
         
     def set_qs(self, qs):
@@ -423,7 +455,7 @@ def beam_alignment_angles(self):
         
         Parameters
         ----------
-        ...
+        N/A
         
             
         Returns
@@ -468,7 +500,7 @@ def beam_alignment_angles(self):
     # ==================================================
     def xy_rotate_coord_sys(self, x_angle=None, y_angle=None, invert=False):
         """
-        Rotates the coordinate system of the beam first with x_angle around the y-axis then with y_angle around the x-axis.
+        Rotates the coordinate system of the beam first with ``x_angle`` around the y-axis then with ``y_angle`` around the x-axis.
         
         Parameters
         ----------

From 05610948b421b2b58e601b8790985c00a8cc27fc Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 25 Mar 2025 18:10:10 +0100
Subject: [PATCH 12/51] Edited tests in test_beam.py to meet the requirement of
 energy, longitudinal proper velocity and momentum being above a given
 threshold. Also some more edts on the rotation tests.

---
 tests/test_beam.py | 67 ++++++++++++++++++++++++++++++++++------------
 1 file changed, 50 insertions(+), 17 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index a335791d..c3070f66 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -27,7 +27,7 @@
 
 def setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, bunch_length=40.0e-06, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0):
     main = SourceBasic()
-    main.bunch_length = bunch_length                                              # [m], rms. Standard value
+    main.bunch_length = bunch_length                                              # [m], rms.
     main.num_particles = 10000                                               
     main.charge = -e * 1.0e10                                                     # [C]
 
@@ -71,13 +71,17 @@ def test_set_phase_space():
     "Verify that the phase space is set correctly."
     beam = Beam()
     num_particles = 10042
+    Q = -SI.e * 1.0e10
     xs = np.random.rand(num_particles)
     ys = np.random.rand(num_particles)
     zs = np.random.rand(num_particles)
     uxs = np.random.rand(num_particles)
     uys = np.random.rand(num_particles)
-    uzs = np.random.rand(num_particles)
-    Q = -SI.e * 1.0e10
+
+    energy_thres = 1001*SI.m_e*SI.c**2/SI.e  # [eV], 1000 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
+
     beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
 
     assert len(beam) == num_particles
@@ -92,7 +96,6 @@ def test_set_phase_space():
     assert np.allclose(beam.uzs(), uzs)
     assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
 
-
     assert np.allclose(beam.xps(), uxs/uzs)
     assert np.allclose(beam.yps(), uys/uzs)
     assert np.allclose(beam.pxs(), SI.m_e*uxs)
@@ -297,10 +300,10 @@ def test_getitem():
 @pytest.mark.beam
 def test_len():
     "Verify that __len__ returns the correct number of particles."
-    random_integer = random.randint(0, 30000)
+    num_particles = 30001
     beam = Beam()
-    beam.reset_phase_space(random_integer)
-    assert len(beam) == random_integer
+    beam.reset_phase_space(num_particles)
+    assert len(beam) == num_particles
 
 
 # @pytest.mark.beam
@@ -317,14 +320,16 @@ def test_len():
 def test_copy_particle_charge():
 
     num_particles = 8051
-
+    energy_thres = 1001*SI.m_e*SI.c**2/SI.e  # [eV], 1000 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
     beam1 = Beam()
     xs = np.random.rand(num_particles)
     ys = np.random.rand(num_particles)
     zs = np.random.rand(num_particles)
     uxs = np.random.rand(num_particles)
     uys = np.random.rand(num_particles)
-    uzs = np.random.rand(num_particles)
+    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
     Q = -SI.e * 1.0e10
     beam1.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
 
@@ -334,7 +339,7 @@ def test_copy_particle_charge():
     zs = np.random.rand(num_particles)
     uxs = np.random.rand(num_particles)
     uys = np.random.rand(num_particles)
-    uzs = np.random.rand(num_particles)
+    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
     Q = -SI.e * 3.4e10
     ws = np.random.rand(num_particles)*Q/(-SI.e)
     beam2.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs, weightings=ws)
@@ -412,7 +417,14 @@ def test_scale_energy():
 
 
 
-############# Tests of rotation methods #############
+############# Tests of coordinate rotation methods #############
+def passive_rotate_vector(vec, x_angle, yangle):
+    "Rotate a vector ``vec`` with three componets (along xyz) first with ``x_angle`` around the y-axis then with ``y_angle`` around the x-axis."
+    rotated_z_offset = main_beam.z_offset() * np.cos(test_x_angle)*np.cos(test_y_angle) + main_beam.x_offset() * np.sin(test_x_angle)*np.cos(test_y_angle) - main_beam.y_offset() * np.sin(test_y_angle)
+    rotated_x_offset = -main_beam.z_offset() * np.sin(test_x_angle) + main_beam.x_offset() * np.cos(test_x_angle)
+    rotated_y_offset = main_beam.z_offset() * np.cos(test_x_angle)*np.sin(test_y_angle) + main_beam.x_offset() * np.sin(test_x_angle)*np.sin(test_y_angle) + main_beam.y_offset() * np.cos(test_y_angle)
+
+
 # def test_slice_centroids():
 #     source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
 #     beam = source.track()
@@ -435,7 +447,9 @@ def test_x_tilt_angle():
     nominal_x_angle = random.uniform(-np.pi/2, np.pi/2)
     beam.set_xs(beam.zs() * np.tan(nominal_x_angle))  # Add a tilt in x
     x_angle = beam.x_tilt_angle()
+    y_angle = beam.y_tilt_angle()
     assert np.isclose(x_angle, nominal_x_angle, rtol=1e-7, atol=1e-3)
+    assert np.isclose(y_angle, 0.0, rtol=1e-7, atol=1e-8)
 
 
 def test_y_tilt_angle():
@@ -443,25 +457,42 @@ def test_y_tilt_angle():
     source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
     beam = source.track()
     nominal_y_angle = random.uniform(-np.pi/2, np.pi/2)
-    beam.set_ys(beam.zs() * np.tan(nominal_y_angle))  # Add a tilt in x
+    beam.set_ys(beam.zs() * np.tan(nominal_y_angle))  # Add a tilt in y
+    x_angle = beam.x_tilt_angle()
     y_angle = beam.y_tilt_angle()
     assert np.isclose(y_angle, nominal_y_angle, rtol=1e-7, atol=1e-3)
+    assert np.isclose(x_angle, 0.0, rtol=1e-7, atol=1e-8)
+
+
+def test_beam_alignment_angles():
+    ""
 
+def test_add_pointing_tilts():
+    "Test correct extraction of the beam pointing angles."
 
 def test_xy_rotate_coord_sys1():
     "Test correct rotation of beam coordinate system around the y-axis."
-    source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
+
+    start_x_angle = random.uniform(-np.pi/2, np.pi/2)  # Define an initial tilt in x.
+    source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=start_x_angle, y_angle=0.0)
     beam = source.track()
+    
+    beam.add_pointing_tilts()  # Add an initial tilt in x.
     old_beam = copy.deepcopy(beam)
-    nom_x_angle = random.uniform(-np.pi/2, np.pi/2)
+    nom_x_angle = random.uniform(-np.pi/2, np.pi/2)  # TODO: rather have a loop over different angles. Do not use random angles.
     beam.xy_rotate_coord_sys(x_angle=nom_x_angle, y_angle=None, invert=False)
-    assert np.allclose(beam.xs(), -old_beam.zs()*np.sin(nom_x_angle) + old_beam.xs()*np.cos(nom_x_angle), rtol=1e-7, atol=1e-13)
     assert np.allclose(beam.zs(), old_beam.zs()*np.cos(nom_x_angle) + old_beam.xs()*np.sin(nom_x_angle), rtol=1e-7, atol=1e-13)
+    assert np.allclose(beam.xs(), -old_beam.zs()*np.sin(nom_x_angle) + old_beam.xs()*np.cos(nom_x_angle), rtol=1e-7, atol=1e-13)
+    assert np.allclose(beam.ys(), old_beam.ys(), rtol=1e-7, atol=1e-13)
+    assert np.allclose(beam.uzs(), old_beam.uzs()*np.cos(nom_x_angle) + old_beam.uxs()*np.sin(nom_x_angle), rtol=1e-7, atol=1e-13)
+    assert np.allclose(beam.uxs(), -old_beam.uzs()*np.sin(nom_x_angle) + old_beam.uxs()*np.cos(nom_x_angle), rtol=1e-7, atol=1e-13)
+    assert np.allclose(beam.uys(), old_beam.uys(), rtol=1e-7, atol=1e-13)
+    assert np.isclose(beam.x_angle,  np.arcsin(rotated_ux_offset/rotated_uz_offset))
+    assert np.isclose(beam.x_tilt_angle(), )
 
 
 def test_rotate_coord_sys_3D():
     "Test correct rotation of beam coordinate system."
-
     source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
     beam = source.track()
     x_axis = np.array([0, 1, 0])  # Axis as an unit vector. Axis permutaton is zxy.
@@ -469,4 +500,6 @@ def test_rotate_coord_sys_3D():
 
     nom_x_angle = random.uniform(-np.pi/2, np.pi/2)
     beam.rotate_coord_sys_3D(y_axis, nom_x_angle, x_axis, 0.0, invert=False)  # Rotate beam 90 degrees around the y-axis.
-    assert np.isclose(beam.x_tilt_angle(), nom_x_angle)
\ No newline at end of file
+    assert np.isclose(beam.x_tilt_angle(), nom_x_angle)
+
+    # TODO: test un-rotate
\ No newline at end of file

From fcc2ce9ed879a5dcce4291138d9b8e5129080756 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Thu, 27 Mar 2025 00:57:46 +0100
Subject: [PATCH 13/51] Moved a chunk of code (the rotation part) further down
 in beam.py. Also minor edits on some docstrings.

---
 abel/classes/beam.py | 478 ++++++++++++++++++++++---------------------
 1 file changed, 245 insertions(+), 233 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index d28759fa..d5d4f008 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -384,8 +384,238 @@ def norm_transverse_vector(self):
         vector[3,:] = self.uys()/SI.c
         return vector
 
+    
+
+    ## BEAM STATISTICS
+
+    def total_particles(self):
+        return int(np.nansum(self.weightings()))
+    
+    def charge(self):
+        return np.nansum(self.qs())
+    
+    def abs_charge(self):
+        return abs(self.charge())
+    
+    def charge_sign(self):
+        if self.charge() == 0:
+            return 1.0
+        else:
+            return self.charge()/abs(self.charge())
+    
+    def energy(self, clean=False):
+        return weighted_mean(self.Es(), self.weightings(), clean)
+    
+    def gamma(self, clean=False):
+        return weighted_mean(self.gammas(), self.weightings(), clean)
+    
+    def total_energy(self):
+        return SI.e * np.nansum(self.weightings()*self.Es())
+    
+    def energy_spread(self, clean=False):
+        return weighted_std(self.Es(), self.weightings(), clean)
+    
+    def rel_energy_spread(self, clean=False):
+        return self.energy_spread(clean)/self.energy(clean)
+    
+    def z_offset(self, clean=False):
+        return weighted_mean(self.zs(), self.weightings(), clean)
+    
+    def bunch_length(self, clean=False):
+        return weighted_std(self.zs(), self.weightings(), clean)
+    
+    def x_offset(self, clean=False):
+        return weighted_mean(self.xs(), self.weightings(), clean)
+    
+    def beam_size_x(self, clean=False):
+        return weighted_std(self.xs(), self.weightings(), clean)
+
+    def y_offset(self, clean=False):
+        return weighted_mean(self.ys(), self.weightings(), clean)
+
+    def beam_size_y(self, clean=False):
+        return weighted_std(self.ys(), self.weightings(), clean)
+    
+    def x_angle(self, clean=False):
+        return weighted_mean(self.xps(), self.weightings(), clean)
+    
+    def divergence_x(self, clean=False):
+        return weighted_std(self.xps(), self.weightings(), clean)
 
-    ## Rotate the coordinate system of the beam
+    def y_angle(self, clean=False):
+        return weighted_mean(self.yps(), self.weightings(), clean)
+    
+    def divergence_y(self, clean=False):
+        return weighted_std(self.yps(), self.weightings(), clean)
+    
+    def ux_offset(self, clean=False):
+        return weighted_mean(self.uxs(), self.weightings(), clean)
+    
+    def uy_offset(self, clean=False):
+        return weighted_mean(self.uys(), self.weightings(), clean)
+    
+    def uz_offset(self, clean=False):
+        return weighted_mean(self.uzs(), self.weightings(), clean)
+
+    
+    def geom_emittance_x(self, clean=False):
+        return np.sqrt(np.linalg.det(weighted_cov(self.xs(), self.xps(), self.weightings(), clean)))
+    
+    def geom_emittance_y(self, clean=False):
+        return np.sqrt(np.linalg.det(weighted_cov(self.ys(), self.yps(), self.weightings(), clean)))
+    
+    def norm_emittance_x(self, clean=False):
+        return np.sqrt(np.linalg.det(weighted_cov(self.xs(), self.uxs()/SI.c, self.weightings(), clean)))
+    
+    def norm_emittance_y(self, clean=False):
+        return np.sqrt(np.linalg.det(weighted_cov(self.ys(), self.uys()/SI.c, self.weightings(), clean)))
+    
+    def beta_x(self, clean=False):
+        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
+        return covx[0,0]/np.sqrt(np.linalg.det(covx))
+    
+    def beta_y(self, clean=False):
+        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
+        return covy[0,0]/np.sqrt(np.linalg.det(covy))
+    
+    def alpha_x(self, clean=False):
+        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
+        return -covx[1,0]/np.sqrt(np.linalg.det(covx))
+    
+    def alpha_y(self, clean=False):
+        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
+        return -covy[1,0]/np.sqrt(np.linalg.det(covy))
+    
+    def gamma_x(self, clean=False):
+        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
+        return covx[1,1]/np.sqrt(np.linalg.det(covx))
+    
+    def gamma_y(self, clean=False):
+        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
+        return covy[1,1]/np.sqrt(np.linalg.det(covy))
+
+    def intrinsic_emittance(self):
+        covxy = np.cov(self.norm_transverse_vector(), aweights=self.weightings())
+        return np.sqrt(np.sqrt(np.linalg.det(covxy)))
+
+    def angular_momentum(self):
+        covxy = np.cov(self.norm_transverse_vector(), aweights=self.weightings())
+        det_covxy_cross = np.linalg.det(covxy[2:4,0:2])
+        return np.sign(covxy[3,0]-covxy[2,1])*np.sqrt(np.abs(det_covxy_cross))
+
+    def eigen_emittance_max(self):
+        return np.sqrt(self.norm_emittance_x()*self.norm_emittance_y()) + self.angular_momentum()
+
+    def eigen_emittance_min(self):
+        return np.sqrt(self.norm_emittance_x()*self.norm_emittance_y()) - self.angular_momentum()
+
+    def norm_amplitude_x(self, plasma_density=None, clean=False):
+        if plasma_density is not None:
+            beta_x = beta_matched(plasma_density, self.energy())
+            alpha_x = 0
+        else:
+            covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
+            emgx = np.sqrt(np.linalg.det(covx))
+            beta_x = covx[0,0]/emgx
+            alpha_x = -covx[1,0]/emgx
+        return np.sqrt(self.gamma()/beta_x)*np.sqrt(self.x_offset()**2 + (self.x_offset()*alpha_x + self.x_angle()*beta_x)**2)
+        
+    def norm_amplitude_y(self, plasma_density=None, clean=False):
+        if plasma_density is not None:
+            beta_y = beta_matched(plasma_density, self.energy())
+            alpha_y = 0
+        else:
+            covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
+            emgy = np.sqrt(np.linalg.det(covy))
+            beta_y = covy[0,0]/emgy
+            alpha_y = -covy[1,0]/emgy
+        return np.sqrt(self.gamma()/beta_y)*np.sqrt(self.y_offset()**2 + (self.y_offset()*alpha_y + self.y_angle()*beta_y)**2)
+        
+    def peak_density(self):  # TODO: this is only valid for Gaussian beams.
+        return (self.charge()/SI.e)/(np.sqrt(2*SI.pi)**3*self.beam_size_x()*self.beam_size_y()*self.bunch_length())
+    
+    def peak_current(self):
+        Is, _ = self.current_profile()
+        return max(abs(Is))
+    
+
+    ## BEAM HALO CLEANING (EXTREME OUTLIERS)
+    def remove_halo_particles(self, nsigma=20):
+        xfilter = np.abs(self.xs()-self.x_offset(clean=True)) > nsigma*self.beam_size_x(clean=True)
+        xpfilter = np.abs(self.xps()-self.x_angle(clean=True)) > nsigma*self.divergence_x(clean=True)
+        yfilter = np.abs(self.ys()-self.y_offset(clean=True)) > nsigma*self.beam_size_y(clean=True)
+        ypfilter = np.abs(self.yps()-self.y_angle(clean=True)) > nsigma*self.divergence_y(clean=True)
+        filter = np.logical_or(np.logical_or(xfilter, xpfilter), np.logical_or(yfilter, ypfilter))
+        del self[filter]
+
+    
+    ## BEAM PROJECTIONS
+    
+    def projected_density(self, fcn, bins=None):
+        if bins is None:
+            Nbins = int(np.sqrt(len(self)/2))
+            bins = np.linspace(min(fcn()), max(fcn()), Nbins)
+        counts, edges = np.histogram(fcn(), weights=self.qs(), bins=bins)
+        ctrs = (edges[0:-1] + edges[1:])/2
+        proj = counts/np.diff(edges)
+        return proj, ctrs
+        
+    def current_profile(self, bins=None):
+        return self.projected_density(self.ts, bins=bins)
+    
+    def longitudinal_num_density(self, bins=None):
+        dQdz, zs = self.projected_density(self.zs, bins=bins)
+        dNdz = dQdz / SI.e / self.charge_sign()
+        return dNdz, zs
+    
+    def energy_spectrum(self, bins=None):
+        return self.projected_density(self.Es, bins=bins)
+    
+    def rel_energy_spectrum(self, nom_energy=None, bins=None):
+        if nom_energy is None:
+            nom_energy = self.energy()
+        return self.projected_density(lambda: self.Es()/nom_energy-1, bins=bins)
+    
+    def transverse_profile_x(self, bins=None):
+        return self.projected_density(self.xs, bins=bins)
+    
+    def transverse_profile_y(self, bins=None):
+        return self.projected_density(self.ys, bins=bins)
+
+    def transverse_profile_xp(self, bins=None):
+        return self.projected_density(self.xps, bins=bins)
+    
+    def transverse_profile_yp(self, bins=None):
+        return self.projected_density(self.yps, bins=bins)
+    
+    ## phase spaces
+    
+    def phase_space_density(self, hfcn, vfcn, hbins=None, vbins=None):
+        self.remove_nans()
+        if hbins is None:
+            hbins = round(np.sqrt(len(self))/2)
+        if vbins is None:
+            vbins = round(np.sqrt(len(self))/2)
+        counts, hedges, vedges = np.histogram2d(hfcn(), vfcn(), weights=self.qs(), bins=(hbins, vbins))
+        hctrs = (hedges[0:-1] + hedges[1:])/2
+        vctrs = (vedges[0:-1] + vedges[1:])/2
+        density = (counts/np.diff(vedges)).T/np.diff(hedges)
+
+        #dx = np.diff(hedges)
+        #dy = np.diff(vedges)
+        #bin_areas = dx[:, None] * dy[None, :]
+        #density = counts/bin_areas
+        #print(np.sum(density*np.diff(vedges)*np.diff(hedges))/self.charge())
+        return density, hctrs, vctrs
+    
+    def density_lps(self, hbins=None, vbins=None):
+        return self.phase_space_density(self.zs, self.Es, hbins=hbins, vbins=vbins)
+    
+    def density_transverse(self, hbins=None, vbins=None):
+        return self.phase_space_density(self.xs, self.ys, hbins=hbins, vbins=vbins)
+
+
+## Rotate the coordinate system of the beam
     # ==================================================
     def rotate_coord_sys_3D(self, axis1, angle1, axis2=np.array([0, 1, 0]), angle2=0.0, axis3=np.array([1, 0, 0]), angle3=0.0, invert=False):
         """
@@ -538,11 +768,11 @@ def add_pointing_tilts(self, align_x_angle=None, align_y_angle=None):
         
         Parameters
         ----------
-         align_x_angle: [rad] float
-            Beam coordinates with in the zx-plane are rotated with this angle.
+         align_x_angle: [rad] float, optional
+            Beam coordinates in the zx-plane are rotated with this angle. If ``None``, will align the beam using its angular offset.
         
-        align_y_angle: [rad] float
-            Beam coordinates with in the zy-plane are rotated with this angle. Note that due to the right hand rule, a positive rotation angle in the zy-plane corresponds to rotation from z-axis towards negative y. I.e. the opposite sign convention of beam.yps().
+        align_y_angle: [rad] float, optional
+            Beam coordinates in the zy-plane are rotated with this angle. Note that due to the right hand rule, a positive rotation angle in the zy-plane corresponds to rotation from z-axis towards negative y. I.e. the opposite sign convention of ``beam.yps()``. If ``None``, will align the beam using its angular offset.
         
             
         Returns
@@ -654,7 +884,12 @@ def slice_centroids(self, beam_quant, bin_number=None, cut_off=None, make_plot=F
         
     
     # ==================================================
+    # def x_tilt_angle(self, clean=False):
+    #     #return np.arcsin(self.ux_offset(clean=clean)/self.uz_offset(clean=clean))
+    #     return np.arcsin(self.x_offset(clean=clean)/self.z_offset(clean=clean))
+    
     def x_tilt_angle(self, z_cutoff=None):
+        "Retrieve the tilt angle of the beam in the zx-plane. WARNING: becomes unreliable around > 1e-4 rad."
         if z_cutoff is None:
             z_cutoff = 1.5 * self.bunch_length()
             
@@ -666,8 +901,12 @@ def x_tilt_angle(self, z_cutoff=None):
         return np.arctan(slope)
 
     
-    # ==================================================
+    # # ==================================================
+    # def y_tilt_angle(self, clean=False):
+    #     return np.arcsin(self.uy_offset(clean=clean)/self.uz_offset(clean=clean))
+    
     def y_tilt_angle(self, z_cutoff=None):
+        "Retrieve the tilt angle of the beam in the zy-plane. WARNING: becomes unreliable around > 1e-4 rad."
         if z_cutoff is None:
             z_cutoff = 1.5 * self.bunch_length()
         y_centroids, z_centroids = self.slice_centroids(self.ys(), cut_off=z_cutoff, make_plot=False)
@@ -676,234 +915,7 @@ def y_tilt_angle(self, z_cutoff=None):
         slope, _ = np.polyfit(z_centroids, y_centroids, 1)
         
         return np.arctan(slope)
-
-
-    
-    ## BEAM STATISTICS
-
-    def total_particles(self):
-        return int(np.nansum(self.weightings()))
-    
-    def charge(self):
-        return np.nansum(self.qs())
-    
-    def abs_charge(self):
-        return abs(self.charge())
-    
-    def charge_sign(self):
-        if self.charge() == 0:
-            return 1.0
-        else:
-            return self.charge()/abs(self.charge())
-    
-    def energy(self, clean=False):
-        return weighted_mean(self.Es(), self.weightings(), clean)
-    
-    def gamma(self, clean=False):
-        return weighted_mean(self.gammas(), self.weightings(), clean)
-    
-    def total_energy(self):
-        return SI.e * np.nansum(self.weightings()*self.Es())
-    
-    def energy_spread(self, clean=False):
-        return weighted_std(self.Es(), self.weightings(), clean)
     
-    def rel_energy_spread(self, clean=False):
-        return self.energy_spread(clean)/self.energy(clean)
-    
-    def z_offset(self, clean=False):
-        return weighted_mean(self.zs(), self.weightings(), clean)
-    
-    def bunch_length(self, clean=False):
-        return weighted_std(self.zs(), self.weightings(), clean)
-    
-    def x_offset(self, clean=False):
-        return weighted_mean(self.xs(), self.weightings(), clean)
-    
-    def beam_size_x(self, clean=False):
-        return weighted_std(self.xs(), self.weightings(), clean)
-
-    def y_offset(self, clean=False):
-        return weighted_mean(self.ys(), self.weightings(), clean)
-
-    def beam_size_y(self, clean=False):
-        return weighted_std(self.ys(), self.weightings(), clean)
-    
-    def x_angle(self, clean=False):
-        return weighted_mean(self.xps(), self.weightings(), clean)
-    
-    def divergence_x(self, clean=False):
-        return weighted_std(self.xps(), self.weightings(), clean)
-
-    def y_angle(self, clean=False):
-        return weighted_mean(self.yps(), self.weightings(), clean)
-    
-    def divergence_y(self, clean=False):
-        return weighted_std(self.yps(), self.weightings(), clean)
-    
-    def ux_offset(self, clean=False):
-        return weighted_mean(self.uxs(), self.weightings(), clean)
-    
-    def uy_offset(self, clean=False):
-        return weighted_mean(self.uys(), self.weightings(), clean)
-
-    
-    def geom_emittance_x(self, clean=False):
-        return np.sqrt(np.linalg.det(weighted_cov(self.xs(), self.xps(), self.weightings(), clean)))
-    
-    def geom_emittance_y(self, clean=False):
-        return np.sqrt(np.linalg.det(weighted_cov(self.ys(), self.yps(), self.weightings(), clean)))
-    
-    def norm_emittance_x(self, clean=False):
-        return np.sqrt(np.linalg.det(weighted_cov(self.xs(), self.uxs()/SI.c, self.weightings(), clean)))
-    
-    def norm_emittance_y(self, clean=False):
-        return np.sqrt(np.linalg.det(weighted_cov(self.ys(), self.uys()/SI.c, self.weightings(), clean)))
-    
-    def beta_x(self, clean=False):
-        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
-        return covx[0,0]/np.sqrt(np.linalg.det(covx))
-    
-    def beta_y(self, clean=False):
-        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
-        return covy[0,0]/np.sqrt(np.linalg.det(covy))
-    
-    def alpha_x(self, clean=False):
-        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
-        return -covx[1,0]/np.sqrt(np.linalg.det(covx))
-    
-    def alpha_y(self, clean=False):
-        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
-        return -covy[1,0]/np.sqrt(np.linalg.det(covy))
-    
-    def gamma_x(self, clean=False):
-        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
-        return covx[1,1]/np.sqrt(np.linalg.det(covx))
-    
-    def gamma_y(self, clean=False):
-        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
-        return covy[1,1]/np.sqrt(np.linalg.det(covy))
-
-    def intrinsic_emittance(self):
-        covxy = np.cov(self.norm_transverse_vector(), aweights=self.weightings())
-        return np.sqrt(np.sqrt(np.linalg.det(covxy)))
-
-    def angular_momentum(self):
-        covxy = np.cov(self.norm_transverse_vector(), aweights=self.weightings())
-        det_covxy_cross = np.linalg.det(covxy[2:4,0:2])
-        return np.sign(covxy[3,0]-covxy[2,1])*np.sqrt(np.abs(det_covxy_cross))
-
-    def eigen_emittance_max(self):
-        return np.sqrt(self.norm_emittance_x()*self.norm_emittance_y()) + self.angular_momentum()
-
-    def eigen_emittance_min(self):
-        return np.sqrt(self.norm_emittance_x()*self.norm_emittance_y()) - self.angular_momentum()
-
-    def norm_amplitude_x(self, plasma_density=None, clean=False):
-        if plasma_density is not None:
-            beta_x = beta_matched(plasma_density, self.energy())
-            alpha_x = 0
-        else:
-            covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
-            emgx = np.sqrt(np.linalg.det(covx))
-            beta_x = covx[0,0]/emgx
-            alpha_x = -covx[1,0]/emgx
-        return np.sqrt(self.gamma()/beta_x)*np.sqrt(self.x_offset()**2 + (self.x_offset()*alpha_x + self.x_angle()*beta_x)**2)
-        
-    def norm_amplitude_y(self, plasma_density=None, clean=False):
-        if plasma_density is not None:
-            beta_y = beta_matched(plasma_density, self.energy())
-            alpha_y = 0
-        else:
-            covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
-            emgy = np.sqrt(np.linalg.det(covy))
-            beta_y = covy[0,0]/emgy
-            alpha_y = -covy[1,0]/emgy
-        return np.sqrt(self.gamma()/beta_y)*np.sqrt(self.y_offset()**2 + (self.y_offset()*alpha_y + self.y_angle()*beta_y)**2)
-        
-    def peak_density(self):  # TODO: this is only valid for Gaussian beams.
-        return (self.charge()/SI.e)/(np.sqrt(2*SI.pi)**3*self.beam_size_x()*self.beam_size_y()*self.bunch_length())
-    
-    def peak_current(self):
-        Is, _ = self.current_profile()
-        return max(abs(Is))
-    
-
-    ## BEAM HALO CLEANING (EXTREME OUTLIERS)
-    def remove_halo_particles(self, nsigma=20):
-        xfilter = np.abs(self.xs()-self.x_offset(clean=True)) > nsigma*self.beam_size_x(clean=True)
-        xpfilter = np.abs(self.xps()-self.x_angle(clean=True)) > nsigma*self.divergence_x(clean=True)
-        yfilter = np.abs(self.ys()-self.y_offset(clean=True)) > nsigma*self.beam_size_y(clean=True)
-        ypfilter = np.abs(self.yps()-self.y_angle(clean=True)) > nsigma*self.divergence_y(clean=True)
-        filter = np.logical_or(np.logical_or(xfilter, xpfilter), np.logical_or(yfilter, ypfilter))
-        del self[filter]
-
-    
-    ## BEAM PROJECTIONS
-    
-    def projected_density(self, fcn, bins=None):
-        if bins is None:
-            Nbins = int(np.sqrt(len(self)/2))
-            bins = np.linspace(min(fcn()), max(fcn()), Nbins)
-        counts, edges = np.histogram(fcn(), weights=self.qs(), bins=bins)
-        ctrs = (edges[0:-1] + edges[1:])/2
-        proj = counts/np.diff(edges)
-        return proj, ctrs
-        
-    def current_profile(self, bins=None):
-        return self.projected_density(self.ts, bins=bins)
-    
-    def longitudinal_num_density(self, bins=None):
-        dQdz, zs = self.projected_density(self.zs, bins=bins)
-        dNdz = dQdz / SI.e / self.charge_sign()
-        return dNdz, zs
-    
-    def energy_spectrum(self, bins=None):
-        return self.projected_density(self.Es, bins=bins)
-    
-    def rel_energy_spectrum(self, nom_energy=None, bins=None):
-        if nom_energy is None:
-            nom_energy = self.energy()
-        return self.projected_density(lambda: self.Es()/nom_energy-1, bins=bins)
-    
-    def transverse_profile_x(self, bins=None):
-        return self.projected_density(self.xs, bins=bins)
-    
-    def transverse_profile_y(self, bins=None):
-        return self.projected_density(self.ys, bins=bins)
-
-    def transverse_profile_xp(self, bins=None):
-        return self.projected_density(self.xps, bins=bins)
-    
-    def transverse_profile_yp(self, bins=None):
-        return self.projected_density(self.yps, bins=bins)
-    
-    ## phase spaces
-    
-    def phase_space_density(self, hfcn, vfcn, hbins=None, vbins=None):
-        self.remove_nans()
-        if hbins is None:
-            hbins = round(np.sqrt(len(self))/2)
-        if vbins is None:
-            vbins = round(np.sqrt(len(self))/2)
-        counts, hedges, vedges = np.histogram2d(hfcn(), vfcn(), weights=self.qs(), bins=(hbins, vbins))
-        hctrs = (hedges[0:-1] + hedges[1:])/2
-        vctrs = (vedges[0:-1] + vedges[1:])/2
-        density = (counts/np.diff(vedges)).T/np.diff(hedges)
-
-        #dx = np.diff(hedges)
-        #dy = np.diff(vedges)
-        #bin_areas = dx[:, None] * dy[None, :]
-        #density = counts/bin_areas
-        #print(np.sum(density*np.diff(vedges)*np.diff(hedges))/self.charge())
-        return density, hctrs, vctrs
-    
-    def density_lps(self, hbins=None, vbins=None):
-        return self.phase_space_density(self.zs, self.Es, hbins=hbins, vbins=vbins)
-    
-    def density_transverse(self, hbins=None, vbins=None):
-        return self.phase_space_density(self.xs, self.ys, hbins=hbins, vbins=vbins)
-
     
     # ==================================================
     # TODO: Currently does not reproduce the correct peak density for Gaussian beams unless the bin numbers are adjusted manually.

From df96f88c25fc00f8cf413235b2c6bfd3553fefe3 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Thu, 27 Mar 2025 00:59:11 +0100
Subject: [PATCH 14/51] Added several major tests for the beam coordinate
 system rotations.

---
 tests/test_beam.py | 318 +++++++++++++++++++++++++++++++++++++--------
 1 file changed, 265 insertions(+), 53 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index c3070f66..10320657 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -25,32 +25,68 @@
 import random
 
 
-def setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, bunch_length=40.0e-06, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0):
-    main = SourceBasic()
-    main.bunch_length = bunch_length                                              # [m], rms.
-    main.num_particles = 10000                                               
-    main.charge = -e * 1.0e10                                                     # [C]
+def setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, bunch_length=40.0e-06, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0):
+    source = SourceBasic()
+    source.bunch_length = bunch_length                                              # [m], rms.
+    source.num_particles = 10000                                               
+    source.charge = -e * 1.0e10                                                     # [C]
 
     # Energy parameters
-    main.energy = energy                                                          # [eV], HALHF v2 last stage nominal input energy
-    main.rel_energy_spread = 0.02                                                 # Relative rms energy spread
+    source.energy = energy                                                          # [eV]
+    source.rel_energy_spread = 0.02                                                 # Relative rms energy spread
 
     # Emittances
-    main.emit_nx, main.emit_ny = 15e-6, 0.1e-6                                    # [m rad]
+    source.emit_nx, source.emit_ny = 15e-6, 0.1e-6                                  # [m rad]
 
     # Beta functions
-    main.beta_x = beta_matched(plasma_density, main.energy) * ramp_beta_mag       # [m]
-    main.beta_y = main.beta_x                                                     # [m]
+    source.beta_x = beta_matched(plasma_density, source.energy) * ramp_beta_mag     # [m]
+    source.beta_y = source.beta_x                                                   # [m]
 
     # Offsets
-    main.z_offset = z_offset                                                      # [m]
-    main.x_offset = x_offset                                                      # [m]
-    main.y_offset = y_offset                                                      # [m]
+    source.z_offset = z_offset                                                      # [m]
+    source.x_offset = x_offset                                                      # [m]
+    source.y_offset = y_offset                                                      # [m]
+    source.x_angle = x_angle                                                        # [rad]
+    source.y_angle = y_angle                                                        # [rad]
 
     # Other
-    main.symmetrize_6d = True
+    source.symmetrize_6d = True
 
-    return main
+    return source
+
+
+def setup_trapezoid_source(current_head=0.1e3, bunch_length=1050e-6, z_offset=1615e-6, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0, enable_xy_jitter=False, enable_xpyp_jitter=False):
+    source = SourceTrapezoid()
+    source.current_head = current_head                                              # [A]
+    source.bunch_length = bunch_length                                              # [m]
+
+    source.num_particles = 30000                                                 
+    source.charge = 5.0e10 * -SI.e                                                  # [C]
+    source.energy = 4.9e9                                                           # [eV] 
+    source.gaussian_blur = 50e-6                                                    # [m]
+    source.rel_energy_spread = 0.01                                              
+
+    source.emit_nx, source.emit_ny = 50e-6, 100e-6                                  # [m rad]
+    source.beta_x, source.beta_y = 0.5, 0.5                                         # [m]
+
+     # Offsets
+    source.z_offset = z_offset                                                      # [m]
+    source.x_offset = x_offset                                                      # [m]
+    source.y_offset = y_offset                                                      # [m]
+    source.x_angle = x_angle                                                        # [rad]
+    source.y_angle = y_angle                                                        # [rad]
+
+    if enable_xy_jitter:
+        source.jitter.x = 100e-9                                                    # [m], std
+        source.jitter.y = 100e-9                                                    # [m], std
+
+    if enable_xpyp_jitter:
+        source.jitter.xp = 1.0e-6                                                   # [rad], std
+        source.jitter.yp = 1.0e-6                                                   # [rad], std
+
+    source.symmetrize = True
+
+    return source
 
 
 ############# Basic tests #############
@@ -415,18 +451,36 @@ def test_scale_energy():
 ############# Tests on bunch pattern #############
 
 
+############# Tests for beam statistics #############
+# Need tests for beam parameter calculations such as offsets, emittance, beta functions, peak current
+
+
+############# Tests for beam projections #############
+# ...
+
+
 
 
 ############# Tests of coordinate rotation methods #############
-def passive_rotate_vector(vec, x_angle, yangle):
-    "Rotate a vector ``vec`` with three componets (along xyz) first with ``x_angle`` around the y-axis then with ``y_angle`` around the x-axis."
-    rotated_z_offset = main_beam.z_offset() * np.cos(test_x_angle)*np.cos(test_y_angle) + main_beam.x_offset() * np.sin(test_x_angle)*np.cos(test_y_angle) - main_beam.y_offset() * np.sin(test_y_angle)
-    rotated_x_offset = -main_beam.z_offset() * np.sin(test_x_angle) + main_beam.x_offset() * np.cos(test_x_angle)
-    rotated_y_offset = main_beam.z_offset() * np.cos(test_x_angle)*np.sin(test_y_angle) + main_beam.x_offset() * np.sin(test_x_angle)*np.sin(test_y_angle) + main_beam.y_offset() * np.cos(test_y_angle)
+def active_rotate_arrs(x_comps, y_comps, z_comps, x_angle, y_angle):
+    "Active rotation of ndarrays ``x_comps``, ``y_comps`` and ``z_comps`` containing the xyz-components of vectors such as position and proper velocity first with ``x_angle`` around the y-axis then with ``y_angle`` around the x-axis."
+    rotated_z_comps = z_comps * np.cos(x_angle)*np.cos(y_angle) - x_comps * np.sin(x_angle)*np.cos(y_angle) + y_comps * np.sin(y_angle)
+    rotated_x_comps = z_comps * np.sin(x_angle) + x_comps * np.cos(x_angle)
+    rotated_y_comps = -z_comps * np.cos(x_angle)*np.sin(y_angle) + x_comps * np.sin(x_angle)*np.sin(y_angle) + y_comps * np.cos(y_angle)
+
+    return rotated_x_comps, rotated_y_comps, rotated_z_comps
+
+def passive_rotate_arrs(x_comps, y_comps, z_comps, x_angle, y_angle):
+    "Passive rotation of ndarrays ``x_comps``, ``y_comps`` and ``z_comps`` containing the xyz-components of vectors such as position and proper velocity first with ``x_angle`` around the y-axis then with ``y_angle`` around the x-axis."
+    rotated_z_comps = z_comps * np.cos(x_angle)*np.cos(y_angle) + x_comps * np.sin(x_angle)*np.cos(y_angle) - y_comps * np.sin(y_angle)
+    rotated_x_comps = -z_comps * np.sin(x_angle) + x_comps * np.cos(x_angle)
+    rotated_y_comps = z_comps * np.cos(x_angle)*np.sin(y_angle) + x_comps * np.sin(x_angle)*np.sin(y_angle) + y_comps * np.cos(y_angle)
+
+    return rotated_x_comps, rotated_y_comps, rotated_z_comps
 
 
 # def test_slice_centroids():
-#     source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
+#     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
 #     beam = source.track()
 #     beam_quant = beam.xs()  # Use x positions as the beam quantity
 #     x_slices, z_centroids = beam.slice_centroids(beam_quant, bin_number=5)
@@ -440,66 +494,224 @@ def passive_rotate_vector(vec, x_angle, yangle):
 #     assert np.all(z_centroids <= 10)
 
 
+@pytest.mark.beam
 def test_x_tilt_angle():
     "Test the retrieval of the beam tilt angle in the zx-plane."
-    source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
+    source = setup_trapezoid_source(current_head=0.1e3, bunch_length=1050e-6, z_offset=1615e-6, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0, enable_xy_jitter=False, enable_xpyp_jitter=False)
     beam = source.track()
-    nominal_x_angle = random.uniform(-np.pi/2, np.pi/2)
-    beam.set_xs(beam.zs() * np.tan(nominal_x_angle))  # Add a tilt in x
+    nom_x_angle = -2.82e-8  # [rad]
+
+    # Add tilt in x
+    rotated_xs, rotated_ys, rotated_zs = passive_rotate_arrs(beam.xs(), beam.ys(), beam.zs(), nom_x_angle, y_angle=0.0)
+    rotated_uxs, rotated_uys, rotated_uzs = passive_rotate_arrs(beam.uxs(), beam.uys(), beam.uzs(), nom_x_angle, y_angle=0.0)
+    beam.set_phase_space(beam.charge(), rotated_xs, rotated_ys, rotated_zs, rotated_uxs, rotated_uys, rotated_uzs)
+
+    # Compare the angles
     x_angle = beam.x_tilt_angle()
     y_angle = beam.y_tilt_angle()
-    assert np.isclose(x_angle, nominal_x_angle, rtol=1e-7, atol=1e-3)
+    assert np.isclose(x_angle, -nom_x_angle, rtol=1e-3, atol=1e-8)
     assert np.isclose(y_angle, 0.0, rtol=1e-7, atol=1e-8)
 
 
+@pytest.mark.beam
 def test_y_tilt_angle():
     "Test the retrieval of the beam tilt angle in the zy-plane."
-    source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
+    source = setup_trapezoid_source(current_head=0.1e3, bunch_length=1050e-6, z_offset=1615e-6, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0, enable_xy_jitter=False, enable_xpyp_jitter=False)
     beam = source.track()
-    nominal_y_angle = random.uniform(-np.pi/2, np.pi/2)
-    beam.set_ys(beam.zs() * np.tan(nominal_y_angle))  # Add a tilt in y
+    nom_y_angle = 7.04e-9  # [rad]
+
+    # Add tilt in y
+    rotated_xs, rotated_ys, rotated_zs = passive_rotate_arrs(beam.xs(), beam.ys(), beam.zs(), 0.0, y_angle=nom_y_angle)
+    rotated_uxs, rotated_uys, rotated_uzs = passive_rotate_arrs(beam.uxs(), beam.uys(), beam.uzs(), 0.0, y_angle=nom_y_angle)
+    beam.set_phase_space(beam.charge(), rotated_xs, rotated_ys, rotated_zs, rotated_uxs, rotated_uys, rotated_uzs)
+
+    # Compare the angles
     x_angle = beam.x_tilt_angle()
     y_angle = beam.y_tilt_angle()
-    assert np.isclose(y_angle, nominal_y_angle, rtol=1e-7, atol=1e-3)
+    assert np.isclose(y_angle, nom_y_angle, rtol=1e-3, atol=1e-11)
     assert np.isclose(x_angle, 0.0, rtol=1e-7, atol=1e-8)
 
 
+@pytest.mark.beam
 def test_beam_alignment_angles():
-    ""
+    "Test ``Beam.beam_alignment_angles()``, which calculates the angles for rotation around the y- and x-axis to align the z-axis to the beam proper velocity."
+    
+    start_x_angle = 1.2e-7  # [rad], define an angle in the zx-plane.
+    start_y_angle = -2.3e-7  # [rad], define an angle in the zy-plane.
 
+    # Create a source with angular offsets, but not tilted in space
+    source = setup_trapezoid_source(current_head=0.1e3, bunch_length=1050e-6, z_offset=1615e-6, x_offset=0.0, y_offset=0.0, x_angle=start_x_angle, y_angle=start_y_angle, enable_xy_jitter=False, enable_xpyp_jitter=False)
+
+    # Generate a beam
+    beam = source.track()
+
+    # Get the angles
+    align_x_angle, align_y_angle = beam.beam_alignment_angles()
+
+    # Compare to the chosen angular offsets
+    assert np.isclose(align_x_angle, start_x_angle, rtol=1e-5, atol=1e-13)
+    assert np.isclose(align_y_angle, start_y_angle, rtol=1e-5, atol=1e-13)
+    
+
+
+@pytest.mark.beam
 def test_add_pointing_tilts():
-    "Test correct extraction of the beam pointing angles."
+    "Test ``Beam.add_pointing_tilts()``, which uses active transformation to tilt the beam in the zx- and zy-planes."
+
+    start_x_angle = -2.82328e-08 # [rad], define an angle in the zx-plane.
+    start_y_angle = 7.040999e-09   # [rad], define an angle in the zy-plane.
+
+    # Create a source with angular offsets, but not tilted in space
+    source = setup_trapezoid_source(current_head=0.1e3, bunch_length=1050e-6, z_offset=1615e-6, x_offset=0.0, y_offset=0.0, x_angle=start_x_angle, y_angle=start_y_angle, enable_xy_jitter=False, enable_xpyp_jitter=False)
 
+    # Generate a beam
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+
+    # Tilt the beam using its angular offset
+    beam.add_pointing_tilts()
+
+    # Rotate the positions and proper velocities of all particles with written out formulas for comparison. Note the sign convention for y_angle.
+    rotated_xs, rotated_ys, rotated_zs = active_rotate_arrs(initial_beam.xs(), initial_beam.ys(), initial_beam.zs(), start_x_angle, -start_y_angle)
+
+    # Examine the resulting beam
+    assert np.isclose(beam.x_angle(), start_x_angle, rtol=1e-3, atol=1e-11)
+    assert np.isclose(beam.y_angle(), start_y_angle, rtol=1e-3, atol=1e-11)
+    assert np.allclose(beam.uxs(), initial_beam.uxs(), rtol=1e-15, atol=0.0)  # The proper velocities should not change.
+    assert np.allclose(beam.uys(), initial_beam.uys(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uzs(), initial_beam.uzs(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.xs(), rotated_xs, rtol=1e-8, atol=0.0)
+    assert np.allclose(beam.ys(), rotated_ys, rtol=1e-8, atol=0.0)
+    assert np.allclose(beam.zs(), rotated_zs, rtol=1e-8, atol=0.0)
+
+
+@pytest.mark.beam
 def test_xy_rotate_coord_sys1():
     "Test correct rotation of beam coordinate system around the y-axis."
 
-    start_x_angle = random.uniform(-np.pi/2, np.pi/2)  # Define an initial tilt in x.
-    source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=start_x_angle, y_angle=0.0)
-    beam = source.track()
+    np.random.seed(42)
+
+    start_x_angle = 1e-7  # [rad], define an angle in the zx-plane.
+
+    # Source with angular offset, but no tilt
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=start_x_angle, y_angle=0.0)
+    
+    nom_x_angles = np.concatenate((-np.logspace(-9, np.log10(0.11), 10), np.logspace(-9, -1, 10)))  # [rad]
+
+    for i in range(len(nom_x_angles)):
+        
+        beam = source.track()
+        beam.add_pointing_tilts()  # Add an initial tilt in the zx-plane.
+        old_beam = copy.deepcopy(beam)
+
+        nom_x_angle = nom_x_angles[i]
+        beam.xy_rotate_coord_sys(x_angle=nom_x_angle, y_angle=None, invert=False)
+
+        # Rotate the positions and proper velocities of all particles with written out formulas for comparison
+        rotated_xs, rotated_ys, rotated_zs = passive_rotate_arrs(old_beam.xs(), old_beam.ys(), old_beam.zs(), nom_x_angle, y_angle=0.0)
+        rotated_uxs, rotated_uys, rotated_uzs = passive_rotate_arrs(old_beam.uxs(), old_beam.uys(), old_beam.uzs(), nom_x_angle, y_angle=0.0)
+        rotated_ux_offset, rotated_uy_offset, rotated_uz_offset = passive_rotate_arrs(old_beam.ux_offset(), old_beam.uy_offset(), old_beam.uz_offset(), nom_x_angle, y_angle=0.0)
+
+        # Compare the rotated arrays
+        assert np.allclose(beam.zs(), rotated_zs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.xs(), rotated_xs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.ys(), rotated_ys, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.uzs(), rotated_uzs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.uxs(), rotated_uxs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.uys(), rotated_uys, rtol=1e-7, atol=1e-13)
+        assert np.isclose(beam.x_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=np.abs(beam.x_angle())*1e-3, atol=np.abs(beam.x_angle())*2e-3)
+        assert np.isclose(beam.y_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=1e-7, atol=1e-13)
+
+        # print(nom_x_angle, beam.x_tilt_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset))
+        # assert np.isclose(beam.x_tilt_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=np.abs(beam.x_angle())*1e-3, atol=np.abs(beam.x_angle())*2e-3)
+        # assert np.isclose(beam.y_tilt_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=1e-7, atol=1e-13)
+
+
+@pytest.mark.beam
+def test_xy_rotate_coord_sys2():
+    "Test correct rotation of beam coordinate system around the y-axis."
+
+    np.random.seed(42)
+
+    start_y_angle = 1e-6  # [rad], define an angle in the zy-plane.
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=start_y_angle)
     
-    beam.add_pointing_tilts()  # Add an initial tilt in x.
-    old_beam = copy.deepcopy(beam)
-    nom_x_angle = random.uniform(-np.pi/2, np.pi/2)  # TODO: rather have a loop over different angles. Do not use random angles.
-    beam.xy_rotate_coord_sys(x_angle=nom_x_angle, y_angle=None, invert=False)
-    assert np.allclose(beam.zs(), old_beam.zs()*np.cos(nom_x_angle) + old_beam.xs()*np.sin(nom_x_angle), rtol=1e-7, atol=1e-13)
-    assert np.allclose(beam.xs(), -old_beam.zs()*np.sin(nom_x_angle) + old_beam.xs()*np.cos(nom_x_angle), rtol=1e-7, atol=1e-13)
-    assert np.allclose(beam.ys(), old_beam.ys(), rtol=1e-7, atol=1e-13)
-    assert np.allclose(beam.uzs(), old_beam.uzs()*np.cos(nom_x_angle) + old_beam.uxs()*np.sin(nom_x_angle), rtol=1e-7, atol=1e-13)
-    assert np.allclose(beam.uxs(), -old_beam.uzs()*np.sin(nom_x_angle) + old_beam.uxs()*np.cos(nom_x_angle), rtol=1e-7, atol=1e-13)
-    assert np.allclose(beam.uys(), old_beam.uys(), rtol=1e-7, atol=1e-13)
-    assert np.isclose(beam.x_angle,  np.arcsin(rotated_ux_offset/rotated_uz_offset))
-    assert np.isclose(beam.x_tilt_angle(), )
+    nom_y_angles = np.concatenate((-np.logspace(-9, np.log10(0.11), 10), np.logspace(-9, -1, 10)))  # [rad]
+
+    for i in range(len(nom_y_angles)):
+        
+        beam = source.track()
+        beam.add_pointing_tilts()  # Add an initial tilt in the zy-plane.
+        old_beam = copy.deepcopy(beam)
+
+        nom_y_angle = nom_y_angles[i]
+        beam.xy_rotate_coord_sys(x_angle=None, y_angle=nom_y_angle, invert=False)
+
+        # Roatate the positions and proper velocities of all particles with written out formulas for comparison
+        rotated_xs, rotated_ys, rotated_zs = passive_rotate_arrs(old_beam.xs(), old_beam.ys(), old_beam.zs(), 0.0, y_angle=nom_y_angle)
+        rotated_uxs, rotated_uys, rotated_uzs = passive_rotate_arrs(old_beam.uxs(), old_beam.uys(), old_beam.uzs(), 0.0, y_angle=nom_y_angle)
+        rotated_ux_offset, rotated_uy_offset, rotated_uz_offset = passive_rotate_arrs(old_beam.ux_offset(), old_beam.uy_offset(), old_beam.uz_offset(), 0.0, y_angle=nom_y_angle)
+
+        # Compare the rotated arrays
+        assert np.allclose(beam.zs(), rotated_zs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.xs(), rotated_xs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.ys(), rotated_ys, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.uzs(), rotated_uzs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.uxs(), rotated_uxs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.uys(), rotated_uys, rtol=1e-7, atol=1e-13)
+        assert np.isclose(beam.x_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=1e-7, atol=1e-13)
+        assert np.isclose(beam.y_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=np.abs(nom_y_angle)*1e-3, atol=np.abs(nom_y_angle)*2e-3)
 
+        # print(nom_y_angle, beam.y_tilt_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset))
+        # assert np.isclose(beam.y_tilt_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=np.abs(nom_y_angle)*1e-3, atol=np.abs(nom_y_angle)*2e-3)
+        # assert np.isclose(beam.x_tilt_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=1e-7, atol=1e-13)
 
+
+@pytest.mark.beam
 def test_rotate_coord_sys_3D():
-    "Test correct rotation of beam coordinate system."
-    source = setup_basic_main_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
-    beam = source.track()
+    "Test correct passive rotation of beam coordinate system."
+
+    np.random.seed(42)
+
+    start_x_angle = -2.82328e-08 # [rad], define an angle in the zx-plane.
+    start_y_angle = 7.040999e-09   # [rad], define an angle in the zy-plane.
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=start_x_angle, y_angle=start_y_angle)
+
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
+    old_beam = source.track()
+    old_beam.add_pointing_tilts()  # Add an initial tilt in the zy-plane.
+    #old_beam = copy.deepcopy(beam)
+
     x_axis = np.array([0, 1, 0])  # Axis as an unit vector. Axis permutaton is zxy.
     y_axis = np.array([0, 0, 1])
-
-    nom_x_angle = random.uniform(-np.pi/2, np.pi/2)
-    beam.rotate_coord_sys_3D(y_axis, nom_x_angle, x_axis, 0.0, invert=False)  # Rotate beam 90 degrees around the y-axis.
-    assert np.isclose(beam.x_tilt_angle(), nom_x_angle)
+    nom_x_angles = np.concatenate((-np.logspace(-9, np.log10(0.11), 10), np.logspace(-9, -1, 10)))  # [rad]
+    nom_y_angles = np.concatenate((-np.logspace(-9, np.log10(0.11), 10), np.logspace(-9, -1, 10)))  # [rad]
+
+    for i in range(len(nom_x_angles)):
+
+        beam = copy.deepcopy(old_beam)
+
+        nom_x_angle = nom_x_angles[i]
+        nom_y_angle = nom_y_angles[i]
+        beam.rotate_coord_sys_3D(y_axis, nom_x_angle, x_axis, nom_y_angle, invert=False)
+
+        # Roatate the positions and proper velocities of all particles with written out formulas for comparison
+        rotated_xs, rotated_ys, rotated_zs = passive_rotate_arrs(old_beam.xs(), old_beam.ys(), old_beam.zs(), nom_x_angle, nom_y_angle)
+        rotated_uxs, rotated_uys, rotated_uzs = passive_rotate_arrs(old_beam.uxs(), old_beam.uys(), old_beam.uzs(), nom_x_angle, nom_y_angle)
+        rotated_ux_offset, rotated_uy_offset, rotated_uz_offset = passive_rotate_arrs(old_beam.ux_offset(), old_beam.uy_offset(), old_beam.uz_offset(), nom_x_angle, nom_y_angle)
+
+        # Compare the rotated arrays
+        assert np.allclose(beam.zs(), rotated_zs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.xs(), rotated_xs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.ys(), rotated_ys, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.uzs(), rotated_uzs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.uxs(), rotated_uxs, rtol=1e-7, atol=1e-13)
+        assert np.allclose(beam.uys(), rotated_uys, rtol=1e-7, atol=1e-13)
+        assert np.isclose(beam.x_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=np.abs(nom_x_angle)*1e-3, atol=np.abs(nom_x_angle)*2e-3)
+        assert np.isclose(beam.y_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=np.abs(nom_y_angle)*1e-3, atol=np.abs(nom_y_angle)*2e-3)
+
+    # nom_x_angle = random.uniform(-np.pi/2, np.pi/2)
+    # beam.rotate_coord_sys_3D(y_axis, nom_x_angle, x_axis, 0.0, invert=False)  # Rotate beam 90 degrees around the y-axis.
+    # assert np.isclose(beam.x_tilt_angle(), nom_x_angle)
 
     # TODO: test un-rotate
\ No newline at end of file

From 00afa1fafd199a761dc3031c69abe23ed59e3853 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Thu, 27 Mar 2025 14:29:58 +0100
Subject: [PATCH 15/51] Updates to the rotation tets in test_beam.py.

---
 tests/test_beam.py | 132 ++++++++++++++++++++++++++++++++-------------
 1 file changed, 95 insertions(+), 37 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 10320657..136e1f0e 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -571,7 +571,7 @@ def test_add_pointing_tilts():
     # Tilt the beam using its angular offset
     beam.add_pointing_tilts()
 
-    # Rotate the positions and proper velocities of all particles with written out formulas for comparison. Note the sign convention for y_angle.
+    # Rotate the positions of all particles with written out formulas for comparison. Note the sign convention for y_angle.
     rotated_xs, rotated_ys, rotated_zs = active_rotate_arrs(initial_beam.xs(), initial_beam.ys(), initial_beam.zs(), start_x_angle, -start_y_angle)
 
     # Examine the resulting beam
@@ -622,10 +622,6 @@ def test_xy_rotate_coord_sys1():
         assert np.isclose(beam.x_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=np.abs(beam.x_angle())*1e-3, atol=np.abs(beam.x_angle())*2e-3)
         assert np.isclose(beam.y_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=1e-7, atol=1e-13)
 
-        # print(nom_x_angle, beam.x_tilt_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset))
-        # assert np.isclose(beam.x_tilt_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=np.abs(beam.x_angle())*1e-3, atol=np.abs(beam.x_angle())*2e-3)
-        # assert np.isclose(beam.y_tilt_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=1e-7, atol=1e-13)
-
 
 @pytest.mark.beam
 def test_xy_rotate_coord_sys2():
@@ -662,9 +658,63 @@ def test_xy_rotate_coord_sys2():
         assert np.isclose(beam.x_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=1e-7, atol=1e-13)
         assert np.isclose(beam.y_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=np.abs(nom_y_angle)*1e-3, atol=np.abs(nom_y_angle)*2e-3)
 
-        # print(nom_y_angle, beam.y_tilt_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset))
-        # assert np.isclose(beam.y_tilt_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=np.abs(nom_y_angle)*1e-3, atol=np.abs(nom_y_angle)*2e-3)
-        # assert np.isclose(beam.x_tilt_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=1e-7, atol=1e-13)
+
+@pytest.mark.beam
+def test_xy_rotate_coord_sys3():
+    "Test adding creating a tilted beam and then align the z-axis to the drive beam propagation axis using passive transformation to rotate the frame of the beam."
+
+    np.random.seed(42)
+
+    start_x_angle = -2.82328e-08 # [rad], define an angle in the zx-plane.
+    start_y_angle = 7.040999e-09   # [rad], define an angle in the zy-plane.
+
+    # Create a source with angular offsets, but not tilted in space
+    source = setup_trapezoid_source(current_head=0.1e3, bunch_length=1050e-6, z_offset=1615e-6, x_offset=0.0, y_offset=0.0, x_angle=start_x_angle, y_angle=start_y_angle, enable_xy_jitter=False, enable_xpyp_jitter=False)
+
+    # Generate a beam
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)  # Has angular offsets, but not tilted in space
+
+    # Tilt the beam using its angular offset
+    beam.add_pointing_tilts()
+
+    # Rotate the positions of all particles with written out formulas for comparison. Note the sign convention for y_angle.
+    rotated_xs, rotated_ys, rotated_zs = active_rotate_arrs(initial_beam.xs(), initial_beam.ys(), initial_beam.zs(), start_x_angle, -start_y_angle)
+
+    # Examine the tilted beam
+    assert np.isclose(beam.x_angle(), start_x_angle, rtol=1e-3, atol=1e-11)
+    assert np.isclose(beam.y_angle(), start_y_angle, rtol=1e-3, atol=1e-11)
+    assert np.allclose(beam.uxs(), initial_beam.uxs(), rtol=1e-15, atol=0.0)  # The proper velocities should not change.
+    assert np.allclose(beam.uys(), initial_beam.uys(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uzs(), initial_beam.uzs(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.xs(), rotated_xs, rtol=1e-8, atol=0.0)
+    assert np.allclose(beam.ys(), rotated_ys, rtol=1e-8, atol=0.0)
+    assert np.allclose(beam.zs(), rotated_zs, rtol=1e-8, atol=0.0)
+
+    # Calculate the angles that will be used to rotate the beams' frame
+    rotation_angle_x, rotation_angle_y = beam.beam_alignment_angles()
+    rotation_angle_y = -rotation_angle_y  # Minus due to right hand rule.
+
+    assert np.isclose(rotation_angle_x, start_x_angle, rtol=1e-10, atol=0.0)
+    assert np.isclose(rotation_angle_y, -start_y_angle, rtol=1e-10, atol=0.0)
+
+
+    # Use passive transformation to rotate the frame of the beam
+    beam.xy_rotate_coord_sys(rotation_angle_x, rotation_angle_y)  # Align the z-axis to the drive beam propagation.
+
+    # Rotate the proper velocities of all particles with written out formulas for comparison. Note the sign convention for y_angle.
+    rotated_uxs, rotated_uys, rotated_uzs = passive_rotate_arrs(initial_beam.uxs(), initial_beam.uys(), initial_beam.uzs(), rotation_angle_x, rotation_angle_y)
+
+    # Examine the aligned beam
+    assert beam.x_angle() < 1e-15  # The angular offsets should be very small
+    assert beam.y_angle() < 1e-15
+    assert np.allclose(beam.uxs(), rotated_uxs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), rotated_uys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uzs(), rotated_uzs, rtol=1e-15, atol=0.0)
+    # The position arrays should have been rotated back to when the beam only had angular offsets, but not tilted in space
+    assert np.allclose(beam.xs(), initial_beam.xs(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.ys(), initial_beam.ys(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.zs(), initial_beam.zs(), rtol=1e-15, atol=0.0)
 
 
 @pytest.mark.beam
@@ -680,38 +730,46 @@ def test_rotate_coord_sys_3D():
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
     old_beam = source.track()
     old_beam.add_pointing_tilts()  # Add an initial tilt in the zy-plane.
-    #old_beam = copy.deepcopy(beam)
 
     x_axis = np.array([0, 1, 0])  # Axis as an unit vector. Axis permutaton is zxy.
     y_axis = np.array([0, 0, 1])
-    nom_x_angles = np.concatenate((-np.logspace(-9, np.log10(0.11), 10), np.logspace(-9, -1, 10)))  # [rad]
-    nom_y_angles = np.concatenate((-np.logspace(-9, np.log10(0.11), 10), np.logspace(-9, -1, 10)))  # [rad]
+    nom_x_angles = np.concatenate((-np.logspace(np.log10(0.11), -9, 10), np.logspace(-8.9, -1, 10)))  # [rad]
+    nom_y_angles = np.concatenate((-np.logspace(np.log10(0.09), -8.1, 10), np.logspace(-9, -0.97, 10)))  # [rad]
 
     for i in range(len(nom_x_angles)):
+        for j in range(len(nom_y_angles)):
+
+            beam = copy.deepcopy(old_beam)
+
+            nom_x_angle = nom_x_angles[i]
+            nom_y_angle = nom_y_angles[j]
+            beam.rotate_coord_sys_3D(y_axis, nom_x_angle, x_axis, nom_y_angle, invert=False)
+
+            # Roatate the positions and proper velocities of all particles with written out formulas for comparison
+            rotated_xs, rotated_ys, rotated_zs = passive_rotate_arrs(old_beam.xs(), old_beam.ys(), old_beam.zs(), nom_x_angle, nom_y_angle)
+            rotated_uxs, rotated_uys, rotated_uzs = passive_rotate_arrs(old_beam.uxs(), old_beam.uys(), old_beam.uzs(), nom_x_angle, nom_y_angle)
+            rotated_ux_offset, rotated_uy_offset, rotated_uz_offset = passive_rotate_arrs(old_beam.ux_offset(), old_beam.uy_offset(), old_beam.uz_offset(), nom_x_angle, nom_y_angle)
+
+            # Compare the rotated arrays
+            assert np.allclose(beam.zs(), rotated_zs, rtol=1e-7, atol=1e-13)
+            assert np.allclose(beam.xs(), rotated_xs, rtol=1e-7, atol=1e-13)
+            assert np.allclose(beam.ys(), rotated_ys, rtol=1e-7, atol=1e-13)
+            assert np.allclose(beam.uzs(), rotated_uzs, rtol=1e-7, atol=1e-13)
+            assert np.allclose(beam.uxs(), rotated_uxs, rtol=1e-7, atol=1e-13)
+            assert np.allclose(beam.uys(), rotated_uys, rtol=1e-7, atol=1e-13)
+            assert np.isclose(beam.x_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=np.abs(nom_x_angle)*1e-3, atol=np.abs(nom_x_angle)*2e-3)
+            assert np.isclose(beam.y_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=np.abs(nom_y_angle)*1e-3, atol=np.abs(nom_y_angle)*2e-3)
+
+            # Un-rotate
+            beam.rotate_coord_sys_3D(y_axis, nom_x_angle, x_axis, nom_y_angle, invert=True)
+
+            # Compare the rotated arrays
+            assert np.allclose(beam.zs(), old_beam.zs(), rtol=1e-7, atol=1e-13)
+            assert np.allclose(beam.xs(), old_beam.xs(), rtol=1e-7, atol=1e-13)
+            assert np.allclose(beam.ys(), old_beam.ys(), rtol=1e-7, atol=1e-13)
+            assert np.allclose(beam.uzs(), old_beam.uzs(), rtol=1e-7, atol=1e-13)
+            assert np.allclose(beam.uxs(), old_beam.uxs(), rtol=1e-7, atol=1e-13)
+            assert np.allclose(beam.uys(), old_beam.uys(), rtol=1e-7, atol=1e-13)
+            assert np.isclose(beam.x_angle(), old_beam.x_angle(), rtol=1e-7, atol=1e-13)
+            assert np.isclose(beam.y_angle(), old_beam.y_angle(), rtol=1e-7, atol=1e-13)
 
-        beam = copy.deepcopy(old_beam)
-
-        nom_x_angle = nom_x_angles[i]
-        nom_y_angle = nom_y_angles[i]
-        beam.rotate_coord_sys_3D(y_axis, nom_x_angle, x_axis, nom_y_angle, invert=False)
-
-        # Roatate the positions and proper velocities of all particles with written out formulas for comparison
-        rotated_xs, rotated_ys, rotated_zs = passive_rotate_arrs(old_beam.xs(), old_beam.ys(), old_beam.zs(), nom_x_angle, nom_y_angle)
-        rotated_uxs, rotated_uys, rotated_uzs = passive_rotate_arrs(old_beam.uxs(), old_beam.uys(), old_beam.uzs(), nom_x_angle, nom_y_angle)
-        rotated_ux_offset, rotated_uy_offset, rotated_uz_offset = passive_rotate_arrs(old_beam.ux_offset(), old_beam.uy_offset(), old_beam.uz_offset(), nom_x_angle, nom_y_angle)
-
-        # Compare the rotated arrays
-        assert np.allclose(beam.zs(), rotated_zs, rtol=1e-7, atol=1e-13)
-        assert np.allclose(beam.xs(), rotated_xs, rtol=1e-7, atol=1e-13)
-        assert np.allclose(beam.ys(), rotated_ys, rtol=1e-7, atol=1e-13)
-        assert np.allclose(beam.uzs(), rotated_uzs, rtol=1e-7, atol=1e-13)
-        assert np.allclose(beam.uxs(), rotated_uxs, rtol=1e-7, atol=1e-13)
-        assert np.allclose(beam.uys(), rotated_uys, rtol=1e-7, atol=1e-13)
-        assert np.isclose(beam.x_angle(), np.arcsin(rotated_ux_offset/rotated_uz_offset), rtol=np.abs(nom_x_angle)*1e-3, atol=np.abs(nom_x_angle)*2e-3)
-        assert np.isclose(beam.y_angle(), np.arcsin(rotated_uy_offset/rotated_uz_offset), rtol=np.abs(nom_y_angle)*1e-3, atol=np.abs(nom_y_angle)*2e-3)
-
-    # nom_x_angle = random.uniform(-np.pi/2, np.pi/2)
-    # beam.rotate_coord_sys_3D(y_axis, nom_x_angle, x_axis, 0.0, invert=False)  # Rotate beam 90 degrees around the y-axis.
-    # assert np.isclose(beam.x_tilt_angle(), nom_x_angle)
-
-    # TODO: test un-rotate
\ No newline at end of file

From e9634d2d6e5c22b34284e4b00dfe70d63c1d67b7 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Fri, 28 Mar 2025 00:18:37 +0100
Subject: [PATCH 16/51] Changed proper_velocity2momentum() in relativity.py to
 calculate momentum directly as m*u.

---
 abel/utilities/relativity.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/abel/utilities/relativity.py b/abel/utilities/relativity.py
index 057a2007..9ffe37dd 100644
--- a/abel/utilities/relativity.py
+++ b/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):

From eb906047c1f91eb746d137e42e5ea50cca433556 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Fri, 28 Mar 2025 00:20:38 +0100
Subject: [PATCH 17/51] Changed pxs(), pys() and pzs() to reflect the changes
 done to proper_velocity2momentum() in relativity.py where the momentum is now
 calculated as m*u.

---
 abel/classes/beam.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index d5d4f008..8a4baa2e 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -241,11 +241,11 @@ def rs(self):
         return np.sqrt(self.xs()**2 + self.ys()**2)
     
     def pxs(self):
-        return proper_velocity2momentum(self.uxs())
+        return proper_velocity2momentum(self.uxs(), m=self.particle_mass)
     def pys(self):
-        return proper_velocity2momentum(self.uys())
+        return proper_velocity2momentum(self.uys(), m=self.particle_mass)
     def pzs(self):
-        return proper_velocity2momentum(self.uzs())
+        return proper_velocity2momentum(self.uzs(), m=self.particle_mass)
     
     def xps(self):
         return self.uxs()/self.uzs()

From 95c808506180fc1cf476a66329a75adcadd5b30b Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Fri, 28 Mar 2025 00:23:07 +0100
Subject: [PATCH 18/51] Added more tolerances to phase space tests and added a
 large amount of tests for beam statistics.

---
 tests/test_beam.py | 352 +++++++++++++++++++++++++++++++++------------
 1 file changed, 262 insertions(+), 90 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 136e1f0e..21231b99 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -22,6 +22,7 @@
 
 import pytest
 from abel import *
+from abel.utilities.beam_physics import generate_trace_space
 import random
 
 
@@ -121,33 +122,38 @@ def test_set_phase_space():
     beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
 
     assert len(beam) == num_particles
-    assert np.isclose(beam.particle_mass, SI.m_e)
-    assert np.isclose(beam.charge(), Q)
+    assert np.isclose(beam.particle_mass, SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge(), Q, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.abs_charge(), np.abs(Q), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.total_particles(), 1.0e10, rtol=1e-15, atol=0.0)
     assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
-    assert np.allclose(beam.xs(), xs)
-    assert np.allclose(beam.ys(), ys)
-    assert np.allclose(beam.zs(), zs)
-    assert np.allclose(beam.uxs(), uxs)
-    assert np.allclose(beam.uys(), uys)
-    assert np.allclose(beam.uzs(), uzs)
-    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
-
-    assert np.allclose(beam.xps(), uxs/uzs)
-    assert np.allclose(beam.yps(), uys/uzs)
-    assert np.allclose(beam.pxs(), SI.m_e*uxs)
-    assert np.allclose(beam.pys(), SI.m_e*uys)
-    assert np.allclose(beam.pzs(), SI.m_e*uzs)
-    assert np.allclose(beam.Es(), np.sqrt((SI.m_e*uzs*SI.c)**2 + (SI.m_e*SI.c**2)**2)/SI.e )
+    assert np.allclose(beam.xs(), xs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.ys(), ys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.zs(), zs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), uxs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), uys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uzs(), uzs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles, rtol=1e-15, atol=0.0)
+
+    assert np.allclose(beam.xps(), uxs/uzs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.yps(), uys/uzs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pxs(), SI.m_e*uxs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pys(), SI.m_e*uys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pzs(), SI.m_e*uzs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.Es(), np.sqrt((SI.m_e*uzs*SI.c)**2 + (SI.m_e*SI.c**2)**2)/SI.e, rtol=1e-15, atol=0.0)
 
 
 @pytest.mark.beam
 def test_set_phase_space2():
     "Verify that the phase space is set correctly."
+
+    np.random.seed(42)
+
     beam = Beam()
-    num_particles = 10042
-    xs = np.random.normal(1.1, 2.0e-6, num_particles)
-    ys = np.random.normal(0.3, 1.0e-6, num_particles)
-    zs = np.random.normal(5.6, 50.0e-6, num_particles)
+    num_particles = 100042
+    xs = np.random.normal(1.1e-6, 2.0e-6, num_particles)
+    ys = np.random.normal(0.3e-6, 1.0e-6, num_particles)
+    zs = np.random.normal(5.6e-6, 50.0e-6, num_particles)
     xps = np.random.normal(1.1e-5, 1.5e-7, num_particles)
     yps = np.random.normal(3.2e-6, 1.0e-8, num_particles)
     Es = np.random.normal(500e9, 0.02*500e9, num_particles)
@@ -156,37 +162,40 @@ def test_set_phase_space2():
     beam.set_phase_space(Q, xs, ys, zs, xps=xps, yps=yps, Es=Es)
 
     assert len(beam) == num_particles
-    assert np.isclose(beam.particle_mass, SI.m_e)
-    assert np.isclose(beam.charge(), Q)
-    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
-    assert np.isclose(beam.energy(), 500e9, rtol=0.0001*500e9, atol=1e9)
-    assert np.isclose(beam.gamma(), 500e9*SI.e/SI.m_e/SI.c**2, rtol=0.0001*500e9*SI.e/SI.m_e/SI.c**2, atol=200)
+    assert np.isclose(beam.particle_mass, SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge(), Q, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge_sign(), -1, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.total_energy(), SI.e * np.nansum(beam.weightings()*beam.Es()), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e), rtol=1e-15, atol=0.0)
+
+    assert np.isclose(beam.energy(), 500e9, rtol=1e-4, atol=1e9)
+    assert np.isclose(beam.gamma(), 500e9*SI.e/SI.m_e/SI.c**2, rtol=1e-3, atol=200)
     assert np.isclose(beam.rel_energy_spread(), 0.02, rtol=0.001, atol=0.001)
-    assert np.isclose(beam.z_offset(), 5.6)
-    assert np.isclose(beam.x_offset(), 1.1)
-    assert np.isclose(beam.y_offset(), 0.3)
-    assert np.isclose(beam.bunch_length(), 50.0e-6, rtol=0.005, atol=0.03)
+    assert np.isclose(beam.z_offset(), 5.6e-6, rtol=0.3, atol=0.0)
+    assert np.isclose(beam.x_offset(), 1.1e-6, rtol=0.05, atol=0.0)
+    assert np.isclose(beam.y_offset(), 0.3e-6, rtol=0.03, atol=0.0)
+    assert np.isclose(beam.bunch_length(), 50.0e-6, rtol=1e-4, atol=0.00)
     assert np.isclose(beam.beam_size_x(), 2.0e-6, rtol=0.005, atol=0.03)
     assert np.isclose(beam.beam_size_y(), 1.0e-6, rtol=0.005, atol=0.03)
-    assert np.isclose(beam.x_angle(), 1.1e-5)
-    assert np.isclose(beam.y_angle(), 3.2e-6)
-    assert np.isclose(beam.divergence_x(), 1.5e-7)
-    assert np.isclose(beam.divergence_y(), 1.0e-8)
-
-    assert np.allclose(beam.xs(), xs)
-    assert np.allclose(beam.ys(), ys)
-    assert np.allclose(beam.zs(), zs)
-    assert np.allclose(beam.xps(), xps, rtol=1e-05, atol=1e-08)
-    assert np.allclose(beam.yps(), yps, rtol=1e-05, atol=1e-08)
-    assert np.allclose(beam.Es(), Es)
-    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
-
-    assert np.allclose(beam.uzs(), np.sqrt((Es*SI.e/SI.m_e/SI.c)**2 - SI.c**2) )
-    assert np.allclose(beam.uxs(), xps*beam.uzs())
-    assert np.allclose(beam.uys(), yps*beam.uzs())
-    assert np.allclose(beam.pxs(), SI.m_e*xps*beam.uzs(), rtol=1e-05, atol=1e-25)
-    assert np.allclose(beam.pys(), SI.m_e*yps*beam.uzs(), rtol=1e-05, atol=1e-25)
-    assert np.allclose(beam.pzs(), SI.m_e*beam.uzs(), rtol=1e-05, atol=1e-19)
+    assert np.isclose(beam.x_angle(), 1.1e-5, rtol=1e-5, atol=0.0)
+    assert np.isclose(beam.y_angle(), 3.2e-6, rtol=1e-4, atol=0.0)
+    assert np.isclose(beam.divergence_x(), 1.5e-7, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.divergence_y(), 1.0e-8, rtol=1e-2, atol=0.0)
+
+    assert np.allclose(beam.xs(), xs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.ys(), ys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.zs(), zs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.xps(), xps, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.yps(), yps, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.Es(), Es, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles, rtol=1e-15, atol=0.0)
+
+    assert np.allclose(beam.uzs(), np.sqrt((Es*SI.e/SI.m_e/SI.c)**2 - SI.c**2), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), xps*beam.uzs(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), yps*beam.uzs(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pxs(), SI.m_e*xps*beam.uzs(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pys(), SI.m_e*yps*beam.uzs(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pzs(), SI.m_e*beam.uzs(), rtol=1e-15, atol=0.0)
 
 
 @pytest.mark.beam
@@ -204,23 +213,23 @@ def test_set_phase_space3():
     beam.set_phase_space(Q, xs, ys, zs, pxs=pxs, pys=pys, pzs=pzs)
 
     assert len(beam) == num_particles
-    assert np.isclose(beam.particle_mass, SI.m_e)
-    assert np.isclose(beam.charge(), Q)
-    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
-    assert np.allclose(beam.xs(), xs)
-    assert np.allclose(beam.ys(), ys)
-    assert np.allclose(beam.zs(), zs)
-    assert np.allclose(beam.pxs(), pxs, rtol=1e-05, atol=1e-25)
-    assert np.allclose(beam.pys(), pys, rtol=1e-05, atol=1e-25)
-    assert np.allclose(beam.pzs(), pzs, rtol=1e-05, atol=1e-19)
-    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
-
-    assert np.allclose(beam.xps(), pxs/pzs)
-    assert np.allclose(beam.yps(), pys/pzs)
-    assert np.allclose(beam.uxs(), pxs/SI.m_e)
-    assert np.allclose(beam.uys(), pys/SI.m_e)
-    assert np.allclose(beam.uzs(), pzs/SI.m_e)
-    assert np.allclose(beam.Es(), np.sqrt((pzs*SI.c)**2 + (SI.m_e*SI.c**2)**2)/SI.e )
+    assert np.isclose(beam.particle_mass, SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge(), Q, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.xs(), xs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.ys(), ys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.zs(), zs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pxs(), pxs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pys(), pys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pzs(), pzs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles, rtol=1e-15, atol=0.0)
+
+    assert np.allclose(beam.xps(), pxs/pzs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.yps(), pys/pzs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), pxs/SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), pys/SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uzs(), pzs/SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.Es(), np.sqrt((pzs*SI.c)**2 + (SI.m_e*SI.c**2)**2)/SI.e, rtol=1e-15, atol=0.0)
 
 
 @pytest.mark.beam
@@ -239,22 +248,22 @@ def test_set_phase_space4():
 
     assert len(beam) == num_particles
     assert np.isclose(beam.particle_mass, SI.m_e)
-    assert np.allclose(beam.xs(), xs)
-    assert np.allclose(beam.ys(), ys)
-    assert np.allclose(beam.zs(), zs)
-    assert np.allclose(beam.uxs(), uxs)
-    assert np.allclose(beam.yps(), yps)
+    assert np.allclose(beam.xs(), xs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.ys(), ys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.zs(), zs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), uxs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.yps(), yps, rtol=1e-15, atol=0.0)
     assert np.allclose(beam.pzs(), pzs, rtol=1e-05, atol=1e-25)
-    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles)
-    assert np.isclose(beam.charge(), Q)
-    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e))
+    assert np.allclose(beam.qs(), np.ones_like(xs)*Q/num_particles, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge(), Q, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.weightings(), np.ones_like(xs)*Q/num_particles/(-SI.e), rtol=1e-15, atol=0.0)
 
-    assert np.allclose(beam.uzs(), pzs/SI.m_e)
-    assert np.allclose(beam.xps(), uxs/beam.uzs())
-    assert np.allclose(beam.uys(), yps*beam.uzs())
-    assert np.allclose(beam.pxs(), SI.m_e*uxs, rtol=1e-05, atol=1e-25)
-    assert np.allclose(beam.pys(), SI.m_e*yps*beam.uzs(), rtol=1e-05, atol=1e-25)    
-    assert np.allclose(beam.Es(), np.sqrt((pzs*SI.c)**2 + (SI.m_e*SI.c**2)**2)/SI.e )
+    assert np.allclose(beam.uzs(), pzs/SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.xps(), uxs/beam.uzs(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), yps*beam.uzs(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pxs(), SI.m_e*uxs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.pys(), SI.m_e*yps*beam.uzs(), rtol=1e-15, atol=0.0)    
+    assert np.allclose(beam.Es(), np.sqrt((pzs*SI.c)**2 + (SI.m_e*SI.c**2)**2)/SI.e, rtol=1e-15, atol=0.0)
 
 
 @pytest.mark.beam
@@ -342,7 +351,7 @@ def test_len():
     assert len(beam) == num_particles
 
 
-# @pytest.mark.beam
+# @pytest.mark.beam TODO
 # def test_str():
 #     "Test the __str__ method for correct formatting of beam properties."
 #     beam = Beam()
@@ -401,17 +410,17 @@ def test_scale_charge():
     # Scale up charge
     beam.scale_charge(-SI.e * 2.0e10)
     scaled_charges = beam.qs()
-    assert np.isclose(scaled_charges.sum(), -SI.e * 2.0e10)  # Check total charge
-    assert np.allclose(scaled_charges / original_charges, 2.0)  # Charges scaled proportionally
+    assert np.isclose(scaled_charges.sum(), -SI.e * 2.0e10, rtol=1e-15, atol=0.0)  # Check total charge
+    assert np.allclose(scaled_charges / original_charges, 2.0, rtol=1e-15, atol=0.0)  # Charges scaled proportionally
 
     # Scale down charge
     beam.scale_charge(-SI.e * 0.5e10)
-    assert np.isclose(scaled_charges.sum(), -SI.e * 0.5e10)
+    assert np.isclose(scaled_charges.sum(), -SI.e * 0.5e10, rtol=1e-15, atol=0.0)
 
     # Edge case: zero charge
     beam.scale_charge(0)
     scaled_charges = beam.qs()
-    assert np.isclose(scaled_charges.sum(), 0)
+    assert np.isclose(scaled_charges.sum(), 0, rtol=1e-15, atol=0.0)
 
 
 @pytest.mark.beam
@@ -441,11 +450,11 @@ def test_scale_energy():
     scaled_mean_energy = weighted_mean(scaled_energies, beam.weightings())
     assert np.isclose(scaled_mean_energy, 1.7e8, rtol=0.01, atol=1e8)
 
-    # # Edge case: zero energy
-    # try:
-    #     beam.scale_energy(0)
-    # except ValueError as e:
-    #     assert str(e) == "Energy cannot be scaled to zero"
+    # Edge case: zero energy
+    try:
+        beam.scale_energy(0)
+    except ValueError as err:
+        assert str(err) == 'Es contains values that are too small.'
 
 
 ############# Tests on bunch pattern #############
@@ -454,6 +463,170 @@ def test_scale_energy():
 ############# Tests for beam statistics #############
 # Need tests for beam parameter calculations such as offsets, emittance, beta functions, peak current
 
+@pytest.mark.beam
+def test_param_calcs_generate_trace_space():
+    "Test the beam parameter calculations using a trace space generated with generate_trace_space()."
+    
+    np.random.seed(42)
+
+    alpha_x = -0.239
+    alpha_y = -0.171
+    beta_x = 0.120                                                                    # [m]
+    beta_y = 0.120                                                                    # [m]
+    geo_emitt_x = 2.552865e-09                                                        # [m rad]
+    geo_emitt_y = 1.750833e-11                                                        # [m rad]
+
+    # Generate trace spaces
+    xs, xps = generate_trace_space(geo_emitt_x, beta_x, alpha_x, 10011, symmetrize=False)
+    ys, yps = generate_trace_space(geo_emitt_y, beta_y, alpha_y, 10011, symmetrize=False)
+    zs = np.random.normal(loc=0.0, scale=40.0e-6, size=10011)
+    Es = np.random.normal(loc=3e9, scale=0.02*3e9, size=10011)
+
+    beam = Beam()
+    beam.set_phase_space(xs=xs, ys=ys, zs=zs, xps=xps, yps=yps, Es=Es, Q=-e*1.0e10)
+
+    # Examine the beam parameters
+    assert np.isclose(np.std(xs), np.sqrt(geo_emitt_x*beta_x), rtol=1e-2, atol=0.0)  # Beam size
+    assert np.isclose(np.std(ys), np.sqrt(geo_emitt_y*beta_y), rtol=1e-2, atol=0.0)
+    assert np.isclose(np.std(Es), 0.02*3e9, rtol=0.01, atol=0.0)
+    assert np.isclose(np.std(zs), 50.0e-6, rtol=0.005, atol=0.03)
+    assert np.isclose(beam.beam_size_x(), np.sqrt(geo_emitt_x*beta_x), rtol=1e-2, atol=0.0)  # Beam size
+    assert np.isclose(beam.beam_size_y(), np.sqrt(geo_emitt_y*beta_y), rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.divergence_x(), np.sqrt(geo_emitt_x/beta_x), rtol=5e-2, atol=0.0)
+    assert np.isclose(beam.divergence_y(), np.sqrt(geo_emitt_y/beta_y), rtol=5e-2, atol=0.0)
+    assert np.isclose(beam.energy_spread(), 0.02*3e9, rtol=0.01, atol=0.0)
+    assert np.isclose(beam.bunch_length(), 50.0e-6, rtol=0.005, atol=0.03)
+
+    assert np.isclose(beam.alpha_x(), alpha_x, rtol=6e-2, atol=0.0)
+    assert np.isclose(beam.alpha_y(), alpha_y, rtol=6e-2, atol=0.0)
+    assert np.isclose(beam.beta_x(), beta_x, rtol=2e-3, atol=0.0)
+    assert np.isclose(beam.beta_y(), beta_y, rtol=1.5e-2, atol=0.0)
+    assert np.isclose(beam.geom_emittance_x(), geo_emitt_x, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.geom_emittance_y(), geo_emitt_y, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.norm_emittance_x(), geo_emitt_x*beam.gamma(), rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.norm_emittance_y(), geo_emitt_y*beam.gamma(), rtol=1e-2, atol=0.0)
+
+
+@pytest.mark.beam
+def test_param_calcs_generate_trace_space_xy():
+    "Test the beam parameter calculations using a trace space generated with generate_trace_space_xy()."
+    
+    np.random.seed(42)
+
+    alpha_x = -0.239
+    alpha_y = -0.171
+    beta_x = 0.120                                                                    # [m]
+    beta_y = 0.120                                                                    # [m]
+    geo_emitt_x = 2.552865e-09                                                        # [m rad]
+    geo_emitt_y = 1.750833e-11                                                        # [m rad]
+    num_particles = 200011
+
+    # Generate trace space
+    xs, xps, ys, yps = generate_trace_space_xy(geo_emitt_x, beta_x, alpha_x, geo_emitt_y, beta_y, alpha_y, num_particles, L=0, symmetrize=False)
+    zs = np.random.normal(loc=0.0, scale=40.0e-6, size=num_particles)
+    Es = np.random.normal(loc=3e9, scale=0.02*3e9, size=num_particles)
+
+    beam = Beam()
+    beam.set_phase_space(xs=xs, ys=ys, zs=zs, xps=xps, yps=yps, Es=Es, Q=-e*1.0e10)
+
+    # Examine the beam parameters
+    assert np.isclose(np.std(xs), np.sqrt(geo_emitt_x*beta_x), rtol=1e-2, atol=0.0)  # Beam size
+    assert np.isclose(np.std(ys), np.sqrt(geo_emitt_y*beta_y), rtol=1e-2, atol=0.0)
+    assert np.isclose(np.std(Es), 0.02*3e9, rtol=0.005, atol=0.0)
+    assert np.isclose(np.std(zs), 50.0e-6, rtol=0.005, atol=0.03)
+    assert np.isclose(beam.beam_size_x(), np.sqrt(geo_emitt_x*beta_x), rtol=1e-3, atol=0.0)  # Beam size
+    assert np.isclose(beam.beam_size_y(), np.sqrt(geo_emitt_y*beta_y), rtol=5e-3, atol=0.0)
+    assert np.isclose(beam.divergence_x(), np.sqrt(geo_emitt_x/beta_x), rtol=5e-2, atol=0.0)
+    assert np.isclose(beam.divergence_y(), np.sqrt(geo_emitt_y/beta_y), rtol=5e-2, atol=0.0)
+    assert np.isclose(beam.energy_spread(), 0.02*3e9, rtol=0.005, atol=0.0)
+    assert np.isclose(beam.bunch_length(), 50.0e-6, rtol=0.005, atol=0.03)
+
+    assert np.isclose(beam.alpha_x(), alpha_x, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.alpha_y(), alpha_y, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.beta_x(), beta_x, rtol=1e-3, atol=0.0)
+    assert np.isclose(beam.beta_y(), beta_y, rtol=3e-3, atol=0.0)
+    assert np.isclose(beam.geom_emittance_x(), geo_emitt_x, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.geom_emittance_y(), geo_emitt_y, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.norm_emittance_x(), geo_emitt_x*beam.gamma(), rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.norm_emittance_y(), geo_emitt_y*beam.gamma(), rtol=1e-2, atol=0.0)
+
+    # Perform the same control with a symmetrised phase space
+    xs, xps, ys, yps = generate_trace_space_xy(geo_emitt_x, beta_x, alpha_x, geo_emitt_y, beta_y, alpha_y, num_particles, L=0, symmetrize=True)
+    num_tiling = 4
+    num_particles_actual = round(num_particles/num_tiling)
+    zs = np.random.normal(loc=0.0, scale=40.0e-6, size=num_particles_actual)
+    Es = np.random.normal(loc=3e9, scale=0.02*3e9, size=num_particles_actual)
+    zs = np.tile(zs, num_tiling)
+    Es = np.tile(Es, num_tiling)
+
+    beam = Beam()
+    beam.set_phase_space(xs=xs, ys=ys, zs=zs, xps=xps, yps=yps, Es=Es, Q=-e*1.0e10)
+
+    assert np.isclose(np.std(xs), np.sqrt(geo_emitt_x*beta_x), rtol=1e-2, atol=0.0)  # Beam size
+    assert np.isclose(np.std(ys), np.sqrt(geo_emitt_y*beta_y), rtol=5e-2, atol=0.0)
+    assert np.isclose(np.std(Es), 0.02*3e9, rtol=0.005, atol=0.0)
+    assert np.isclose(np.std(zs), 50.0e-6, rtol=0.005, atol=0.03)
+    assert np.isclose(beam.beam_size_x(), np.sqrt(geo_emitt_x*beta_x), rtol=1e-3, atol=0.0)  # Beam size
+    assert np.isclose(beam.beam_size_y(), np.sqrt(geo_emitt_y*beta_y), rtol=5e-3, atol=0.0)
+    assert np.isclose(beam.divergence_x(), np.sqrt(geo_emitt_x/beta_x), rtol=5e-2, atol=0.0)
+    assert np.isclose(beam.divergence_y(), np.sqrt(geo_emitt_y/beta_y), rtol=5e-2, atol=0.0)
+    assert np.isclose(beam.energy_spread(), 0.02*3e9, rtol=0.01, atol=0.0)
+    assert np.isclose(beam.bunch_length(), 50.0e-6, rtol=0.005, atol=0.03)
+
+    assert np.isclose(beam.alpha_x(), alpha_x, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.alpha_y(), alpha_y, rtol=4e-2, atol=0.0)
+    assert np.isclose(beam.beta_x(), beta_x, rtol=3e-3, atol=0.0)
+    assert np.isclose(beam.beta_y(), beta_y, rtol=5e-3, atol=0.0)
+    assert np.isclose(beam.geom_emittance_x(), geo_emitt_x, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.geom_emittance_y(), geo_emitt_y, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.norm_emittance_x(), geo_emitt_x*beam.gamma(), rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.norm_emittance_y(), geo_emitt_y*beam.gamma(), rtol=1e-2, atol=0.0)
+
+
+@pytest.mark.beam
+def test_param_calcs_generate_symm_trace_space_xyz():
+    "Test the beam parameter calculations using a trace space generated with generate_symm_trace_space_xyz()."
+    
+    np.random.seed(42)
+
+    alpha_x = -0.239
+    alpha_y = -0.171
+    beta_x = 0.120                                                                    # [m]
+    beta_y = 0.120                                                                    # [m]
+    geo_emitt_x = 2.552865e-09                                                        # [m rad]
+    geo_emitt_y = 1.750833e-11                                                        # [m rad]
+    num_particles = 200011
+
+    # Generate trace space
+    xs, xps, ys, yps, zs, Es = generate_symm_trace_space_xyz(geo_emitt_x, beta_x, alpha_x, geo_emitt_y, beta_y, alpha_y, num_particles, bunch_length=50.0e-6, energy_spread=0.02*3e9, L=0)
+    Es += 3e9  # Add offset.
+
+    beam = Beam()
+    beam.set_phase_space(xs=xs, ys=ys, zs=zs, xps=xps, yps=yps, Es=Es, Q=-e*1.0e10)
+
+    # Examine the beam parameters
+    assert np.isclose(np.std(xs), np.sqrt(geo_emitt_x*beta_x), rtol=1e-2, atol=0.0)  # Beam size
+    assert np.isclose(np.std(ys), np.sqrt(geo_emitt_y*beta_y), rtol=1e-2, atol=0.0)
+    assert np.isclose(np.std(Es), 0.02*3e9, rtol=0.005, atol=0.0)
+    assert np.isclose(np.std(zs), 50.0e-6, rtol=0.005, atol=0.03)
+    assert np.isclose(beam.beam_size_x(), np.sqrt(geo_emitt_x*beta_x), rtol=1e-3, atol=0.0)  # Beam size
+    assert np.isclose(beam.beam_size_y(), np.sqrt(geo_emitt_y*beta_y), rtol=5e-3, atol=0.0)
+    assert np.isclose(beam.divergence_x(), np.sqrt(geo_emitt_x/beta_x), rtol=5e-2, atol=0.0)
+    assert np.isclose(beam.divergence_y(), np.sqrt(geo_emitt_y/beta_y), rtol=5e-2, atol=0.0)
+    assert np.isclose(beam.energy_spread(), 0.02*3e9, rtol=0.005, atol=0.0)
+    assert np.isclose(beam.bunch_length(), 50.0e-6, rtol=0.005, atol=0.03)
+    assert np.isclose(beam.peak_current(), -beam.charge()/(np.sqrt(2*np.pi)*beam.bunch_length())*SI.c, rtol=0.0, atol=3e2)
+    assert np.isclose(beam.peak_density(), beam.charge()/(SI.e*np.sqrt(2*np.pi)**3*beam.beam_size_x()*beam.beam_size_y()*beam.bunch_length()), rtol=1e-8, atol=0.0)
+
+    assert np.isclose(beam.alpha_x(), alpha_x, rtol=3e-2, atol=0.0)
+    assert np.isclose(beam.alpha_y(), alpha_y, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.beta_x(), beta_x, rtol=3e-3, atol=0.0)
+    assert np.isclose(beam.beta_y(), beta_y, rtol=3e-3, atol=0.0)
+    assert np.isclose(beam.geom_emittance_x(), geo_emitt_x, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.geom_emittance_y(), geo_emitt_y, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.norm_emittance_x(), geo_emitt_x*beam.gamma(), rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.norm_emittance_y(), geo_emitt_y*beam.gamma(), rtol=1e-2, atol=0.0)
+
 
 ############# Tests for beam projections #############
 # ...
@@ -553,7 +726,6 @@ def test_beam_alignment_angles():
     assert np.isclose(align_y_angle, start_y_angle, rtol=1e-5, atol=1e-13)
     
 
-
 @pytest.mark.beam
 def test_add_pointing_tilts():
     "Test ``Beam.add_pointing_tilts()``, which uses active transformation to tilt the beam in the zx- and zy-planes."

From 94ab1fed9e56d152a1182febc6fcf7179e725701 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Fri, 28 Mar 2025 15:42:57 +0100
Subject: [PATCH 19/51] Added tests for beam projection methods in
 test_beam.py.

---
 tests/test_beam.py | 161 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 158 insertions(+), 3 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 21231b99..e2932547 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -457,12 +457,13 @@ def test_scale_energy():
         assert str(err) == 'Es contains values that are too small.'
 
 
+
 ############# Tests on bunch pattern #############
+# TODO
 
 
-############# Tests for beam statistics #############
-# Need tests for beam parameter calculations such as offsets, emittance, beta functions, peak current
 
+############# Tests for beam statistics #############
 @pytest.mark.beam
 def test_param_calcs_generate_trace_space():
     "Test the beam parameter calculations using a trace space generated with generate_trace_space()."
@@ -628,10 +629,139 @@ def test_param_calcs_generate_symm_trace_space_xyz():
     assert np.isclose(beam.norm_emittance_y(), geo_emitt_y*beam.gamma(), rtol=1e-2, atol=0.0)
 
 
+
 ############# Tests for beam projections #############
-# ...
+@pytest.mark.beam
+def test_current_profile():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=1.0e-6, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+
+    dIdt, ts = beam.current_profile()
+    assert np.isclose(np.max(np.abs(dIdt)), beam.peak_current(), rtol=1e-15, atol=0.0)
+    assert np.isclose(SI.c*ts.mean(), beam.z_offset(), rtol=1e-10, atol=0.0)
+
+    Nbins = int(np.sqrt(len(beam)/2))
+    bins = np.linspace(min(beam.ts()), max(beam.ts()), Nbins)
+    dIdt, ts = beam.current_profile(bins=bins)
+    assert np.isclose(np.max(np.abs(dIdt)), beam.peak_current(), rtol=1e-15, atol=0.0)
+    assert len(dIdt) == Nbins - 1
+    assert len(ts) == Nbins - 1
+    assert np.isclose(SI.c*ts.mean(), beam.z_offset(), rtol=1e-10, atol=0.0)
 
 
+@pytest.mark.beam
+def test_longitudinal_num_density():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=1.0e-6, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    dNdz, zs = beam.longitudinal_num_density()
+    assert np.isclose( np.sum( dNdz*np.diff(zs)[0] ), beam.total_particles(), rtol=1e-10, atol=0.0 )
+    assert np.isclose(zs.mean(), beam.z_offset(), rtol=1e-10, atol=0.0)
+
+
+@pytest.mark.beam
+def test_energy_spectrum():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    dQdE, Es = beam.energy_spectrum()
+    assert np.isclose( np.sum( dQdE*np.diff(Es)[0] ), beam.charge(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( Es.mean(), beam.energy(), rtol=1e-10, atol=0.0 )
+
+
+@pytest.mark.beam
+def test_rel_energy_spectrum():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    dQdE, rel_Es = beam.rel_energy_spectrum()
+    assert np.isclose( np.sum( dQdE*np.diff(rel_Es)[0] ), beam.charge(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( rel_Es.mean(), 0.0, rtol=0.0, atol=1e-10 )
+
+
+@pytest.mark.beam
+def test_transverse_profile_x():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=1.0e-6, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    dQdx, xs = beam.transverse_profile_x()
+    assert np.isclose( np.sum( dQdx*np.diff(xs)[0] ), beam.charge(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( xs.mean(), beam.x_offset(), rtol=1e-10, atol=0.0 )
+
+
+@pytest.mark.beam
+def test_transverse_profile_y():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=1.0e-6, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    dQdy, ys = beam.transverse_profile_y()
+    assert np.isclose( np.sum( dQdy*np.diff(ys)[0] ), beam.charge(), rtol=1e-8, atol=0.0 )
+    assert np.isclose( ys.mean(), beam.y_offset(), rtol=1e-10, atol=0.0 )
+
+
+@pytest.mark.beam
+def test_transverse_profile_xp():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=1.0e-6, y_angle=0.0)
+    beam = source.track()
+    dQdxp, xps = beam.transverse_profile_xp()
+    assert np.isclose( np.sum( dQdxp*np.diff(xps)[0] ), beam.charge(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( xps.mean(), beam.x_angle(), rtol=1e-10, atol=0.0 )
+
+
+@pytest.mark.beam
+def test_transverse_profile_yp():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=1.0e-6)
+    beam = source.track()
+    dQdyp, yps = beam.transverse_profile_yp()
+    assert np.isclose( np.sum( dQdyp*np.diff(yps)[0] ), beam.charge(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( yps.mean(), beam.y_angle(), rtol=1e-10, atol=0.0 )
+
+
+@pytest.mark.beam
+def test_phase_space_density():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=1.2e-6, y_angle=0.3e-6)
+    beam = source.track()
+    density, xps, yps = beam.phase_space_density(beam.xps, beam.yps, hbins=None, vbins=None)
+    density_1d = np.sum(density*np.diff(xps)[0], axis=1)
+    assert np.isclose( np.sum(density_1d*np.diff(yps)[0]), beam.charge(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( xps.mean(), beam.x_angle(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( yps.mean(), beam.y_angle(), rtol=1e-10, atol=0.0 )
+
+    Nbins = int(np.sqrt(len(beam)/2))
+    hbins = np.linspace(min(beam.xps()), max(beam.xps()), Nbins+1)
+    vbins = np.linspace(min(beam.yps()), max(beam.yps()), Nbins+2)
+    density, xps, yps = beam.phase_space_density(beam.xps, beam.yps, hbins, vbins)
+    density_1d = np.sum(density*np.diff(xps)[0], axis=1)
+    assert np.isclose( np.sum(density_1d*np.diff(yps)[0]), beam.charge(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( xps.mean(), beam.x_angle(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( yps.mean(), beam.y_angle(), rtol=1e-10, atol=0.0 )
+
+
+@pytest.mark.beam
+def test_density_lps():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=1.0e-6, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    density, zs, Es = beam.density_lps(hbins=None, vbins=None)
+    density_1d = np.sum(density*np.diff(zs)[0], axis=1)
+    assert np.isclose( np.sum(density_1d*np.diff(Es)[0]), beam.charge(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( Es.mean(), beam.energy(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( zs.mean(), beam.z_offset(), rtol=1e-10, atol=0.0 )
+
+    Nbins = int(np.sqrt(len(beam)/2))
+    hbins = np.linspace(min(beam.zs()), max(beam.zs()), Nbins+1)
+    vbins = np.linspace(min(beam.Es()), max(beam.Es()), Nbins+2)
+    density, zs, Es = beam.density_lps(hbins, vbins)
+    density_1d = np.sum(density*np.diff(zs)[0], axis=1)
+    assert np.isclose( np.sum(density_1d*np.diff(Es)[0]), beam.charge(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( Es.mean(), beam.energy(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( zs.mean(), beam.z_offset(), rtol=1e-10, atol=0.0 )
+
+
+@pytest.mark.beam
+def test_density_transverse():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.2e-6, y_offset=3.3e-6, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    density, xs, ys = beam.density_transverse()
+    density_1d = np.sum(density*np.diff(xs)[0], axis=1)
+    assert np.isclose( np.sum(density_1d*np.diff(ys)[0]), beam.charge(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( xs.mean(), beam.x_offset(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( ys.mean(), beam.y_offset(), rtol=1e-10, atol=0.0 )
+
 
 
 ############# Tests of coordinate rotation methods #############
@@ -643,6 +773,7 @@ def active_rotate_arrs(x_comps, y_comps, z_comps, x_angle, y_angle):
 
     return rotated_x_comps, rotated_y_comps, rotated_z_comps
 
+
 def passive_rotate_arrs(x_comps, y_comps, z_comps, x_angle, y_angle):
     "Passive rotation of ndarrays ``x_comps``, ``y_comps`` and ``z_comps`` containing the xyz-components of vectors such as position and proper velocity first with ``x_angle`` around the y-axis then with ``y_angle`` around the x-axis."
     rotated_z_comps = z_comps * np.cos(x_angle)*np.cos(y_angle) + x_comps * np.sin(x_angle)*np.cos(y_angle) - y_comps * np.sin(y_angle)
@@ -945,3 +1076,27 @@ def test_rotate_coord_sys_3D():
             assert np.isclose(beam.x_angle(), old_beam.x_angle(), rtol=1e-7, atol=1e-13)
             assert np.isclose(beam.y_angle(), old_beam.y_angle(), rtol=1e-7, atol=1e-13)
 
+
+
+############# Tests of in-house beam field calculation methods #############
+@pytest.mark.beam
+def test_charge_density_3D():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=123.4e-6, x_offset=0.1e-6, y_offset=4.5e-6, x_angle=1.2e-6, y_angle=0.3e-6)
+    beam = source.track()
+    dQ_dzdxdy, zctrs, xctrs, yctrs, edges_z, edges_x, edges_y = beam.charge_density_3D(zbins=None, xbins=None, ybins=None)
+
+    # Calculate volume of each bin
+    dz = np.diff(edges_z)
+    dx = np.diff(edges_x)
+    dy = np.diff(edges_y)
+    bin_volumes = dz[:, None, None] * dx[None, :, None] * dy[None, None, :]  # The None indexing is used to add new axes to the differences arrays, allowing them to be broadcasted properly for division with counts. This ensures that each element of counts is divided by the corresponding bin volume (element-wise division).
+
+    assert np.isclose(np.sum(dQ_dzdxdy*bin_volumes), beam.charge(), rtol=1e-15, atol=0.0)
+    assert np.isclose( xctrs.mean(), beam.x_offset(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( yctrs.mean(), beam.y_offset(), rtol=1e-10, atol=0.0 )
+    assert np.isclose( zctrs.mean(), beam.z_offset(), rtol=1e-10, atol=0.0 )
+
+# TODO: test Dirichlet_BC_system_matrix(), Ex_Ey_2D(), Ex_Ey()
+
+
+############# Tests of plotting methods #############
\ No newline at end of file

From 74c833294e7492389d81b9888fbeaa044aa628d6 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Fri, 28 Mar 2025 16:17:47 +0100
Subject: [PATCH 20/51] Minor change to Beam plotting methods and moved these
 further down in beam.py.

---
 abel/classes/beam.py | 158 ++++++++++++++++++++++---------------------
 1 file changed, 80 insertions(+), 78 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index 8a4baa2e..5bf7aa10 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -1198,84 +1198,6 @@ def Ex_Ey(self, x_box_min, x_box_max, y_box_min, y_box_max, dx, dy, num_z_cells=
           
         
     
-    ## PLOTTING
-    def plot_current_profile(self):
-        dQdt, ts = self.current_profile()
-
-        fig, ax = plt.subplots()
-        fig.set_figwidth(6)
-        fig.set_figheight(4)        
-        ax.plot(ts*SI.c*1e6, -dQdt/1e3)
-        ax.set_xlabel('z (um)')
-        ax.set_ylabel('Beam current (kA)')
-    
-    def plot_lps(self):
-        dQdzdE, zs, Es = self.density_lps()
-
-        fig, ax = plt.subplots()
-        fig.set_figwidth(8)
-        fig.set_figheight(5)  
-            
-        p = ax.pcolor(zs*1e6, Es/1e9, -dQdzdE*1e15, cmap=CONFIG.default_cmap, shading='auto')
-        ax.set_xlabel('z (um)')
-        ax.set_ylabel('E (GeV)')
-        ax.set_title('Longitudinal phase space')
-        cb = fig.colorbar(p)
-        cb.ax.set_ylabel('Charge density (pC/um/GeV)')
-        
-    def plot_trace_space_x(self):
-        dQdxdxp, xs, xps = self.phase_space_density(self.xs, self.xps)
-
-        fig, ax = plt.subplots()
-        fig.set_figwidth(8)
-        fig.set_figheight(5)  
-        p = ax.pcolor(xs*1e6, xps*1e3, -dQdxdxp*1e3, cmap=CONFIG.default_cmap, shading='auto')
-        ax.set_xlabel('x (um)')
-        ax.set_ylabel('x'' (mrad)')
-        ax.set_title('Horizontal trace space')
-        cb = fig.colorbar(p)
-        cb.ax.set_ylabel('Charge density (pC/um/mrad)')
-        
-    def plot_trace_space_y(self):
-        dQdydyp, ys, yps = self.phase_space_density(self.ys, self.yps)
-
-        fig, ax = plt.subplots()
-        fig.set_figwidth(8)
-        fig.set_figheight(5)  
-        p = ax.pcolor(ys*1e6, yps*1e3, -dQdydyp*1e3, cmap=CONFIG.default_cmap, shading='auto')
-        ax.set_xlabel('y (um)')
-        ax.set_ylabel('y'' (mrad)')
-        ax.set_title('Vertical trace space')
-        cb = fig.colorbar(p)
-        cb.ax.set_ylabel('Charge density (pC/um/mrad)')
-
-    def plot_transverse_profile(self):
-        dQdxdy, xs, ys = self.phase_space_density(self.xs, self.ys)
-
-        fig, ax = plt.subplots()
-        fig.set_figwidth(8)
-        fig.set_figheight(5)
-        p = ax.pcolor(xs*1e6, ys*1e6, -dQdxdy, cmap=CONFIG.default_cmap, shading='auto')
-        #p = ax.imshow(-dQdxdy, extent=[xs.min()*1e6, xs.max()*1e6, ys.min()*1e6, ys.max()*1e6], 
-        #   origin='lower', cmap=CONFIG.default_cmap, aspect='auto')
-        ax.set_xlabel('x (um)')
-        ax.set_ylabel('y (um)')
-        ax.set_title('Transverse profile')
-        cb = fig.colorbar(p)
-        cb.ax.set_ylabel('Charge density (pC/um^2)')
-
-    
-    def plot_bunch_pattern(self):
-        
-        fig, ax = plt.subplots()
-        fig.set_figwidth(6)
-        fig.set_figheight(4)        
-        ax.plot(ts*SI.c*1e6, -dQdt/1e3)
-        ax.set_xlabel('z (um)')
-        ax.set_ylabel('Beam current (kA)')
-        
-        
-    
     ## CHANGE BEAM
     
     def accelerate(self, energy_gain=0, chirp=0, z_offset=0):
@@ -1432,6 +1354,86 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat
         else:
             return Es_final
         
+
+    
+     ## PLOTTING
+    def plot_current_profile(self):
+        dQdt, ts = self.current_profile()
+
+        fig, ax = plt.subplots()
+        fig.set_figwidth(6)
+        fig.set_figheight(4)        
+        ax.plot(ts*SI.c*1e6, np.abs(dQdt)/1e3)
+        ax.set_xlabel('z (um)')
+        ax.set_ylabel('Beam current (kA)')
+    
+    def plot_lps(self):
+        dQdzdE, zs, Es = self.density_lps()
+
+        fig, ax = plt.subplots()
+        fig.set_figwidth(8)
+        fig.set_figheight(5)  
+            
+        p = ax.pcolor(zs*1e6, Es/1e9, -dQdzdE*1e15, cmap=CONFIG.default_cmap, shading='auto')
+        ax.set_xlabel('z (um)')
+        ax.set_ylabel('E (GeV)')
+        ax.set_title('Longitudinal phase space')
+        cb = fig.colorbar(p)
+        cb.ax.set_ylabel('Charge density (pC/um/GeV)')
+        
+    def plot_trace_space_x(self):
+        dQdxdxp, xs, xps = self.phase_space_density(self.xs, self.xps)
+
+        fig, ax = plt.subplots()
+        fig.set_figwidth(8)
+        fig.set_figheight(5)  
+        p = ax.pcolor(xs*1e6, xps*1e3, -dQdxdxp*1e3, cmap=CONFIG.default_cmap, shading='auto')
+        ax.set_xlabel('x (um)')
+        ax.set_ylabel('x'' (mrad)')
+        ax.set_title('Horizontal trace space')
+        cb = fig.colorbar(p)
+        cb.ax.set_ylabel('Charge density (pC/um/mrad)')
+        
+    def plot_trace_space_y(self):
+        dQdydyp, ys, yps = self.phase_space_density(self.ys, self.yps)
+
+        fig, ax = plt.subplots()
+        fig.set_figwidth(8)
+        fig.set_figheight(5)  
+        p = ax.pcolor(ys*1e6, yps*1e3, -dQdydyp*1e3, cmap=CONFIG.default_cmap, shading='auto')
+        ax.set_xlabel('y (um)')
+        ax.set_ylabel('y'' (mrad)')
+        ax.set_title('Vertical trace space')
+        cb = fig.colorbar(p)
+        cb.ax.set_ylabel('Charge density (pC/um/mrad)')
+
+    def plot_transverse_profile(self):
+        dQdxdy, xs, ys = self.phase_space_density(self.xs, self.ys)
+
+        fig, ax = plt.subplots()
+        fig.set_figwidth(8)
+        fig.set_figheight(5)
+        p = ax.pcolor(xs*1e6, ys*1e6, -dQdxdy, cmap=CONFIG.default_cmap, shading='auto')
+        #p = ax.imshow(-dQdxdy, extent=[xs.min()*1e6, xs.max()*1e6, ys.min()*1e6, ys.max()*1e6], 
+        #   origin='lower', cmap=CONFIG.default_cmap, aspect='auto')
+        ax.set_xlabel('x (um)')
+        ax.set_ylabel('y (um)')
+        ax.set_title('Transverse profile')
+        cb = fig.colorbar(p)
+        cb.ax.set_ylabel('Charge density (pC/um^2)')
+
+    
+    # TODO: unfinished!
+    # def plot_bunch_pattern(self):
+        
+    #     fig, ax = plt.subplots()
+    #     fig.set_figwidth(6)
+    #     fig.set_figheight(4)
+    #     ax.plot(ts*SI.c*1e6, np.abs(dQdt)/1e3)
+    #     ax.set_xlabel('z (um)')
+    #     ax.set_ylabel('Beam current (kA)')
+        
+        
   
     ## SAVE AND LOAD BEAM
     

From 81cffabf99e5df38f2856bbe7fd336201d951be2 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Fri, 28 Mar 2025 16:20:39 +0100
Subject: [PATCH 21/51] Added tests for Beam plotting methods.

---
 tests/test_beam.py | 54 ++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 48 insertions(+), 6 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index e2932547..9dbf71ae 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -636,16 +636,18 @@ def test_current_profile():
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=1.0e-6, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
     beam = source.track()
 
-    dIdt, ts = beam.current_profile()
-    assert np.isclose(np.max(np.abs(dIdt)), beam.peak_current(), rtol=1e-15, atol=0.0)
+    dQdt, ts = beam.current_profile()
+    assert np.isclose(np.max(np.abs(dQdt)), beam.peak_current(), rtol=1e-15, atol=0.0)
+    assert np.isclose( np.sum( dQdt*np.diff(ts)[0] ), beam.charge(), rtol=1e-10, atol=0.0 )
     assert np.isclose(SI.c*ts.mean(), beam.z_offset(), rtol=1e-10, atol=0.0)
 
     Nbins = int(np.sqrt(len(beam)/2))
     bins = np.linspace(min(beam.ts()), max(beam.ts()), Nbins)
-    dIdt, ts = beam.current_profile(bins=bins)
-    assert np.isclose(np.max(np.abs(dIdt)), beam.peak_current(), rtol=1e-15, atol=0.0)
-    assert len(dIdt) == Nbins - 1
+    dQdt, ts = beam.current_profile(bins=bins)
+    assert np.isclose(np.max(np.abs(dQdt)), beam.peak_current(), rtol=1e-15, atol=0.0)
+    assert len(dQdt) == Nbins - 1
     assert len(ts) == Nbins - 1
+    assert np.isclose( np.sum( dQdt*np.diff(ts)[0] ), beam.charge(), rtol=1e-10, atol=0.0 )
     assert np.isclose(SI.c*ts.mean(), beam.z_offset(), rtol=1e-10, atol=0.0)
 
 
@@ -1099,4 +1101,44 @@ def test_charge_density_3D():
 # TODO: test Dirichlet_BC_system_matrix(), Ex_Ey_2D(), Ex_Ey()
 
 
-############# Tests of plotting methods #############
\ No newline at end of file
+############# Tests of plotting methods #############
+@pytest.mark.beam
+def test_plot_current_profile():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    beam.plot_current_profile()
+
+
+@pytest.mark.beam
+def test_plot_lps():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    beam.plot_lps()
+
+
+@pytest.mark.beam
+def test_plot_trace_space_x():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    beam.plot_trace_space_x()
+
+
+@pytest.mark.beam
+def test_plot_trace_space_y():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    beam.plot_trace_space_y()
+
+
+@pytest.mark.beam
+def test_plot_transverse_profile():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    beam.plot_transverse_profile()
+
+
+# @pytest.mark.beam
+# def test_plot_bunch_pattern():
+#     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+#     beam = source.track()
+#     beam.plot_bunch_pattern()
\ No newline at end of file

From 602c6d0646fa219c71c3e49fca065437d8640710 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Fri, 28 Mar 2025 16:59:13 +0100
Subject: [PATCH 22/51] Added tests for Beam.accelerate().

---
 tests/test_beam.py | 59 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 9dbf71ae..8e1905d6 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -1101,12 +1101,67 @@ def test_charge_density_3D():
 # TODO: test Dirichlet_BC_system_matrix(), Ex_Ey_2D(), Ex_Ey()
 
 
+
+############# Tests of methods changing the beam #############
+@pytest.mark.beam
+def test_accelerate_nominal():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    energy_gain = 1e9  # 1 GeV energy gain
+    chirp = 0  # No chirp
+    initial_Es = beam.Es()
+        
+    beam.accelerate(energy_gain=energy_gain, chirp=chirp)
+    expected_Es = initial_Es + energy_gain
+
+    assert np.allclose(beam.Es(), expected_Es, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_accelerate_negative():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    energy_gain = -beam.energy()  # Energy decrease
+    chirp = 0  # No chirp
+
+    # Suppress warnings
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=RuntimeWarning)
+
+        try:
+            beam.accelerate(energy_gain=energy_gain, chirp=chirp)
+        except ValueError as err:
+            assert str(err) == 'uzs contains values that are too small.'
+
+
+@pytest.mark.beam
+def test_accelerate_chirp():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    energy_gain = 1e9  # 1 GeV energy gain
+    chirp = 1e8  # Energy chirp dE/dz
+    zs = beam.zs()
+    initial_Es = beam.Es()
+
+    expected_Es = energy_gain + initial_Es + np.sign(beam.qs()) * zs * chirp
+    beam.accelerate(energy_gain=energy_gain, chirp=chirp)
+
+    assert np.allclose(beam.Es(), expected_Es, rtol=1e-15, atol=0.0)
+
+
+
+
+
+
+
+
 ############# Tests of plotting methods #############
 @pytest.mark.beam
 def test_plot_current_profile():
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
     beam = source.track()
     beam.plot_current_profile()
+    plt.close()
 
 
 @pytest.mark.beam
@@ -1114,6 +1169,7 @@ def test_plot_lps():
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
     beam = source.track()
     beam.plot_lps()
+    plt.close()
 
 
 @pytest.mark.beam
@@ -1121,6 +1177,7 @@ def test_plot_trace_space_x():
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
     beam = source.track()
     beam.plot_trace_space_x()
+    plt.close()
 
 
 @pytest.mark.beam
@@ -1128,6 +1185,7 @@ def test_plot_trace_space_y():
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
     beam = source.track()
     beam.plot_trace_space_y()
+    plt.close()
 
 
 @pytest.mark.beam
@@ -1135,6 +1193,7 @@ def test_plot_transverse_profile():
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
     beam = source.track()
     beam.plot_transverse_profile()
+    plt.close()
 
 
 # @pytest.mark.beam

From ef6c2b59557cf559d6f6b0ccd8340f26b8f54511 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Mon, 31 Mar 2025 16:20:32 +0200
Subject: [PATCH 23/51] Added controls to Beam.scale_norm_emittance_x() and
 Beam.scale_norm_emittance_y() to prevent negative emittances.

---
 abel/classes/beam.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index 5bf7aa10..cc969db4 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -1220,11 +1220,15 @@ def scale_to_length(self, bunch_length):
         self.set_zs(zs_scaled)
 
     def scale_norm_emittance_x(self, norm_emit_nx):
+        if norm_emit_nx < 0:
+            raise ValueError('Normalised emittance cannot be negative.')
         scale_factor = norm_emit_nx/self.norm_emittance_x()
         self.set_xs(self.xs() * np.sqrt(scale_factor))
         self.set_uxs(self.uxs() * np.sqrt(scale_factor))
 
     def scale_norm_emittance_y(self, norm_emit_ny):
+        if norm_emit_ny < 0:
+            raise ValueError('Normalised emittance cannot be negative.')
         scale_factor = norm_emit_ny/self.norm_emittance_y()
         self.set_ys(self.ys() * np.sqrt(scale_factor))
         self.set_uys(self.uys() * np.sqrt(scale_factor))
@@ -1433,7 +1437,7 @@ def plot_transverse_profile(self):
     #     ax.set_xlabel('z (um)')
     #     ax.set_ylabel('Beam current (kA)')
         
-        
+
   
     ## SAVE AND LOAD BEAM
     

From 1904545218adb951a712563fd9e3106eac7251ff Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Mon, 31 Mar 2025 16:22:13 +0200
Subject: [PATCH 24/51] Added a flag calc_evolution to StageBaisc and edited
 StageBasic.track() to enable using parameter evolution tracking when calling
 Beam.apply_betatron_motion().

---
 abel/classes/stage/impl/stage_basic.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/abel/classes/stage/impl/stage_basic.py b/abel/classes/stage/impl/stage_basic.py
index 58dec105..f44b3dec 100644
--- a/abel/classes/stage/impl/stage_basic.py
+++ b/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)
+        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 ==========

From a8c26e04c26f240e78d780d8df72252f5ab8b33f Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Mon, 31 Mar 2025 16:24:36 +0200
Subject: [PATCH 25/51] Added tests for the Beam methods compress(),
 scale_to_length(), scale_norm_emittance_x(), scale_norm_emittance_y() and
 apply_betatron_motion().

---
 tests/test_beam.py | 232 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 230 insertions(+), 2 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 8e1905d6..37f0156f 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -26,7 +26,7 @@
 import random
 
 
-def setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, bunch_length=40.0e-06, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0):
+def setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, bunch_length=40.0e-06, energy=3e9, rel_energy_spread=0.02, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0):
     source = SourceBasic()
     source.bunch_length = bunch_length                                              # [m], rms.
     source.num_particles = 10000                                               
@@ -34,7 +34,7 @@ def setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, bunch_length=40
 
     # Energy parameters
     source.energy = energy                                                          # [eV]
-    source.rel_energy_spread = 0.02                                                 # Relative rms energy spread
+    source.rel_energy_spread = rel_energy_spread                                    # Relative rms energy spread
 
     # Emittances
     source.emit_nx, source.emit_ny = 15e-6, 0.1e-6                                  # [m rad]
@@ -1149,10 +1149,238 @@ def test_accelerate_chirp():
     assert np.allclose(beam.Es(), expected_Es, rtol=1e-15, atol=0.0)
 
 
+@pytest.mark.beam
+def test_compress_nominal():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+
+    R_56 = 0.5**2*SI.c**2*np.sqrt(beam.energy()/10e9)**3/(3*beam.energy()**2)
+    initial_zs = beam.zs()
+    initial_Es = beam.Es()
+    
+    expected_zs = initial_zs + (1 - initial_Es / beam.energy()) * R_56
+    beam.compress(R_56=R_56, nom_energy=beam.energy())
+    
+    assert np.allclose(beam.zs(), expected_zs, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_compress_null():
+
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+
+    R_56 = 0  # No compression
+    initial_zs = beam.zs()
+    
+    beam.compress(R_56=R_56, nom_energy=beam.energy())
+    
+    # zs should remain unchanged
+    assert np.allclose(beam.zs(), initial_zs, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_compress_uniform():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
 
+    R_56 = 0.5**2*SI.c**2*np.sqrt(beam.energy()/10e9)**3/(3*beam.energy()**2)
+    uniform_energy = np.full(len(beam.Es()), beam.energy())
+    beam.set_Es(uniform_energy)
+    initial_zs = beam.zs()
 
+    beam.compress(R_56=R_56, nom_energy=beam.energy())
+    
+    assert np.allclose(beam.zs(), initial_zs, rtol=1e-15, atol=0.0)
 
 
+@pytest.mark.beam
+def test_scale_to_length():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    
+    beam = source.track()
+    new_bunch_length = beam.bunch_length()/2
+    zs_scaled = beam.z_offset() + (beam.zs()-beam.z_offset())*new_bunch_length/beam.bunch_length()
+    beam.scale_to_length(bunch_length=new_bunch_length)
+    assert np.allclose(beam.zs(), zs_scaled, rtol=1e-15, atol=0.0)
+        
+    beam = source.track()
+    new_bunch_length = beam.bunch_length()*5.354
+    zs_scaled = beam.z_offset() + (beam.zs()-beam.z_offset())*new_bunch_length/beam.bunch_length()
+    beam.scale_to_length(bunch_length=new_bunch_length)
+    assert np.allclose(beam.zs(), zs_scaled, rtol=1e-15, atol=0.0)
+
+    beam = source.track()
+    new_bunch_length = beam.bunch_length()*0
+    beam.scale_to_length(bunch_length=new_bunch_length)
+    assert np.allclose(beam.zs(), beam.z_offset(), rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_scale_norm_emittance_x():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    
+    beam = source.track()
+    scale_factor = 15.16
+    expct_emit_nx = scale_factor * beam.norm_emittance_x()
+    expected_xs = beam.xs() * np.sqrt(scale_factor)
+    expected_uxs = beam.uxs() * np.sqrt(scale_factor)
+    beam.scale_norm_emittance_x(scale_factor*beam.norm_emittance_x())
+    assert np.allclose(beam.xs(), expected_xs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), expected_uxs, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_x(), expct_emit_nx, rtol=1e-15, atol=0.0)
+
+    beam = source.track()
+    scale_factor = 0.5354
+    expct_emit_nx = scale_factor * beam.norm_emittance_x()
+    expected_xs = beam.xs() * np.sqrt(scale_factor)
+    expected_uxs = beam.uxs() * np.sqrt(scale_factor)
+    beam.scale_norm_emittance_x(scale_factor*beam.norm_emittance_x())
+    assert np.allclose(beam.xs(), expected_xs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), expected_uxs, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_x(), expct_emit_nx, rtol=1e-15, atol=0.0)
+
+    beam = source.track()
+    scale_factor = 0.0
+    expct_emit_nx = scale_factor * beam.norm_emittance_x()
+    expected_xs = beam.xs() * np.sqrt(scale_factor)
+    expected_uxs = beam.uxs() * np.sqrt(scale_factor)
+    beam.scale_norm_emittance_x(scale_factor*beam.norm_emittance_x())
+    assert np.allclose(beam.xs(), expected_xs, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), expected_uxs, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_x(), expct_emit_nx, rtol=1e-15, atol=0.0)
+
+    try:
+        beam.scale_norm_emittance_x(-8e-6)
+    except ValueError as err:
+        assert str(err) == 'Normalised emittance cannot be negative.'
+
+
+@pytest.mark.beam
+def test_scale_norm_emittance_y():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    
+    beam = source.track()
+    scale_factor = 15.16
+    expct_emit_ny = scale_factor * beam.norm_emittance_y()
+    expected_ys = beam.ys() * np.sqrt(scale_factor)
+    expected_uys = beam.uys() * np.sqrt(scale_factor)
+    beam.scale_norm_emittance_y(scale_factor*beam.norm_emittance_y())
+    assert np.allclose(beam.ys(), expected_ys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), expected_uys, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_y(), expct_emit_ny, rtol=1e-13, atol=0.0)
+
+    beam = source.track()
+    scale_factor = 0.5354
+    expct_emit_ny = scale_factor * beam.norm_emittance_y()
+    expected_ys = beam.ys() * np.sqrt(scale_factor)
+    expected_uys = beam.uys() * np.sqrt(scale_factor)
+    beam.scale_norm_emittance_y(scale_factor*beam.norm_emittance_y())
+    assert np.allclose(beam.ys(), expected_ys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), expected_uys, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_y(), expct_emit_ny, rtol=1e-15, atol=0.0)
+
+    beam = source.track()
+    scale_factor = 0.0
+    expct_emit_ny = scale_factor * beam.norm_emittance_y()
+    expected_ys = beam.ys() * np.sqrt(scale_factor)
+    expected_uys = beam.uys() * np.sqrt(scale_factor)
+    beam.scale_norm_emittance_y(scale_factor*beam.norm_emittance_y())
+    assert np.allclose(beam.ys(), expected_ys, rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), expected_uys, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_y(), expct_emit_ny, rtol=1e-15, atol=0.0)
+
+    try:
+        beam.scale_norm_emittance_y(-8e-6)
+    except ValueError as err:
+        assert str(err) == 'Normalised emittance cannot be negative.'
+
+
+# @pytest.mark.beam
+# def test_apply_betatron_damping():
+    
+    
+
+
+# @pytest.mark.beam
+# def test_magnify_beta_function():
+#     # TODO
+
+
+# @pytest.mark.beam
+# def test_transport():
+#     # TODO
+
+
+# @pytest.mark.beam
+# def test_flip_transverse_phase_spaces():
+#     # TODO
+
+
+@pytest.mark.beam
+def test_apply_betatron_motion_emitt_pres():
+    "Test of ``Beam.apply_betatron_motion()`` with no energy spread and radiation reaction so that the emittances are preserved."
+
+    np.random.seed(42)
+
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+    deltaEs = np.full(len(beam.Es()), 1e9)  # Homogeneous energy gain
+
+    Es_final, evol = beam.apply_betatron_motion(L=1.0, n0=6.0e20, deltaEs=deltaEs, x0_driver=0, y0_driver=0, radiation_reaction=False, calc_evolution=True)
+
+    assert np.allclose(Es_final.mean(), 4e9, rtol=1e-15, atol=0.0)
+    assert np.allclose(np.std(Es_final), 1e-20, rtol=0.0, atol=1e-19)
+    assert np.allclose(np.std(Es_final)/Es_final.mean(), 1e-20, rtol=0.0, atol=1e-19)
+    assert np.allclose(beam.charge(), initial_beam.charge(), rtol=1e-10, atol=0.0)
+    assert np.allclose(beam.z_offset(), initial_beam.z_offset(), rtol=1e-10, atol=0.0)
+    assert np.allclose(beam.bunch_length(), initial_beam.bunch_length(), rtol=1e-10, atol=0.0)
+
+    assert np.allclose(Es_final.mean(), evol.energy[-1], rtol=1e-20, atol=0.0)
+    assert np.allclose(evol.energy_spread[-1], 1e-20, rtol=0.0, atol=1e-19)
+    assert np.allclose(beam.charge(), evol.charge[-1], rtol=1e-10, atol=0.0)
+    assert np.allclose(beam.norm_emittance_x(), evol.emit_nx[-1], rtol=0.0, atol=0.3e-5)
+    assert np.allclose(beam.norm_emittance_y(), evol.emit_ny[-1], rtol=0.0, atol=0.5e-5)
+    
+    # Emittance is preserved when the energy spread is 0 (homogenneoud energy gain) and radiation reaction disabled.
+    #assert np.isclose(beam.rel_energy_spread(), 1e-15, rtol=1e-10, atol=0.0) # Beam energies not updated after apply_betatron_motion(), so cannot check this.
+    assert np.allclose(beam.norm_emittance_x(), initial_beam.norm_emittance_x(), rtol=1e-8, atol=0.0)
+    assert np.allclose(beam.norm_emittance_y(), initial_beam.norm_emittance_y(), rtol=1e-8, atol=0.0)
+    assert np.allclose(evol.emit_nx[0], evol.emit_nx[-1], rtol=1e-8, atol=0.0)
+    assert np.allclose(evol.emit_ny[0], evol.emit_ny[-1], rtol=1e-8, atol=0.0)
+
+
+@pytest.mark.beam
+def test_apply_betatron_motion():
+    "Test of ``Beam.apply_betatron_motion()`` with energy spread and radiation reaction."
+
+    np.random.seed(42)
+
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.01, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+    deltaEs = np.full(len(beam.Es()), 1e9)  # Homogeneous energy gain
+
+    Es_final, evol = beam.apply_betatron_motion(L=1.0, n0=6.0e20, deltaEs=deltaEs, x0_driver=0, y0_driver=0, radiation_reaction=True, calc_evolution=True)
+
+    assert np.allclose(Es_final.mean(), 4e9, rtol=1e-6, atol=0.0)
+    assert np.allclose(np.std(Es_final), 0.01*3e9, rtol=5e-3, atol=0.0)
+    #assert np.allclose(np.std(Es_final)/Es_final.mean(), 1e-20, rtol=0.0, atol=1e-19)
+    assert np.allclose(beam.charge(), initial_beam.charge(), rtol=1e-10, atol=0.0)
+    assert np.allclose(beam.z_offset(), initial_beam.z_offset(), rtol=1e-10, atol=0.0)
+    assert np.allclose(beam.bunch_length(), initial_beam.bunch_length(), rtol=1e-10, atol=0.0)
+
+    assert np.allclose(Es_final.mean(), evol.energy[-1], rtol=1e-3, atol=0.0)
+    assert np.allclose(evol.energy_spread[-1], 0.01*3e9, rtol=5e-2, atol=0.0)
+    assert np.allclose(beam.charge(), evol.charge[-1], rtol=1e-10, atol=0.0)
+    assert np.allclose(beam.norm_emittance_x(), evol.emit_nx[-1], rtol=0.0, atol=0.3e-5)
+    assert np.allclose(beam.norm_emittance_y(), evol.emit_ny[-1], rtol=0.0, atol=0.5e-5)
+
+
+    #mag = np.sqrt(initial_beam.gamma()/energy2gamma(Es_final.mean()))  # Magnification factor
+    #assert np.allclose(beam.beta_x(), initial_beam.beta_x()*mag, rtol=1e-15, atol=0.0)
+    
 
 
 ############# Tests of plotting methods #############

From 5362d2afbfe76e5fb382c24e328aaf9671a93ba6 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Mon, 31 Mar 2025 17:39:52 +0200
Subject: [PATCH 26/51] Added test_magnify_beta_function() to test_beam.py.

---
 tests/test_beam.py | 96 +++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 90 insertions(+), 6 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 37f0156f..4258c4ad 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -520,7 +520,7 @@ def test_param_calcs_generate_trace_space_xy():
     beta_y = 0.120                                                                    # [m]
     geo_emitt_x = 2.552865e-09                                                        # [m rad]
     geo_emitt_y = 1.750833e-11                                                        # [m rad]
-    num_particles = 200011
+    num_particles = 200012
 
     # Generate trace space
     xs, xps, ys, yps = generate_trace_space_xy(geo_emitt_x, beta_x, alpha_x, geo_emitt_y, beta_y, alpha_y, num_particles, L=0, symmetrize=False)
@@ -531,6 +531,7 @@ def test_param_calcs_generate_trace_space_xy():
     beam.set_phase_space(xs=xs, ys=ys, zs=zs, xps=xps, yps=yps, Es=Es, Q=-e*1.0e10)
 
     # Examine the beam parameters
+    assert len(beam) == num_particles
     assert np.isclose(np.std(xs), np.sqrt(geo_emitt_x*beta_x), rtol=1e-2, atol=0.0)  # Beam size
     assert np.isclose(np.std(ys), np.sqrt(geo_emitt_y*beta_y), rtol=1e-2, atol=0.0)
     assert np.isclose(np.std(Es), 0.02*3e9, rtol=0.005, atol=0.0)
@@ -543,7 +544,7 @@ def test_param_calcs_generate_trace_space_xy():
     assert np.isclose(beam.bunch_length(), 50.0e-6, rtol=0.005, atol=0.03)
 
     assert np.isclose(beam.alpha_x(), alpha_x, rtol=1e-2, atol=0.0)
-    assert np.isclose(beam.alpha_y(), alpha_y, rtol=1e-2, atol=0.0)
+    assert np.isclose(beam.alpha_y(), alpha_y, rtol=2e-2, atol=0.0)
     assert np.isclose(beam.beta_x(), beta_x, rtol=1e-3, atol=0.0)
     assert np.isclose(beam.beta_y(), beta_y, rtol=3e-3, atol=0.0)
     assert np.isclose(beam.geom_emittance_x(), geo_emitt_x, rtol=1e-2, atol=0.0)
@@ -563,6 +564,8 @@ def test_param_calcs_generate_trace_space_xy():
     beam = Beam()
     beam.set_phase_space(xs=xs, ys=ys, zs=zs, xps=xps, yps=yps, Es=Es, Q=-e*1.0e10)
 
+    # Examine the beam parameters
+    assert len(beam) == num_particles
     assert np.isclose(np.std(xs), np.sqrt(geo_emitt_x*beta_x), rtol=1e-2, atol=0.0)  # Beam size
     assert np.isclose(np.std(ys), np.sqrt(geo_emitt_y*beta_y), rtol=5e-2, atol=0.0)
     assert np.isclose(np.std(Es), 0.02*3e9, rtol=0.005, atol=0.0)
@@ -596,7 +599,7 @@ def test_param_calcs_generate_symm_trace_space_xyz():
     beta_y = 0.120                                                                    # [m]
     geo_emitt_x = 2.552865e-09                                                        # [m rad]
     geo_emitt_y = 1.750833e-11                                                        # [m rad]
-    num_particles = 200011
+    num_particles = 200016
 
     # Generate trace space
     xs, xps, ys, yps, zs, Es = generate_symm_trace_space_xyz(geo_emitt_x, beta_x, alpha_x, geo_emitt_y, beta_y, alpha_y, num_particles, bunch_length=50.0e-6, energy_spread=0.02*3e9, L=0)
@@ -606,6 +609,7 @@ def test_param_calcs_generate_symm_trace_space_xyz():
     beam.set_phase_space(xs=xs, ys=ys, zs=zs, xps=xps, yps=yps, Es=Es, Q=-e*1.0e10)
 
     # Examine the beam parameters
+    assert len(beam) == num_particles
     assert np.isclose(np.std(xs), np.sqrt(geo_emitt_x*beta_x), rtol=1e-2, atol=0.0)  # Beam size
     assert np.isclose(np.std(ys), np.sqrt(geo_emitt_y*beta_y), rtol=1e-2, atol=0.0)
     assert np.isclose(np.std(Es), 0.02*3e9, rtol=0.005, atol=0.0)
@@ -619,7 +623,7 @@ def test_param_calcs_generate_symm_trace_space_xyz():
     assert np.isclose(beam.peak_current(), -beam.charge()/(np.sqrt(2*np.pi)*beam.bunch_length())*SI.c, rtol=0.0, atol=3e2)
     assert np.isclose(beam.peak_density(), beam.charge()/(SI.e*np.sqrt(2*np.pi)**3*beam.beam_size_x()*beam.beam_size_y()*beam.bunch_length()), rtol=1e-8, atol=0.0)
 
-    assert np.isclose(beam.alpha_x(), alpha_x, rtol=3e-2, atol=0.0)
+    assert np.isclose(beam.alpha_x(), alpha_x, rtol=5e-2, atol=0.0)
     assert np.isclose(beam.alpha_y(), alpha_y, rtol=1e-2, atol=0.0)
     assert np.isclose(beam.beta_x(), beta_x, rtol=3e-3, atol=0.0)
     assert np.isclose(beam.beta_y(), beta_y, rtol=3e-3, atol=0.0)
@@ -1303,8 +1307,88 @@ def test_scale_norm_emittance_y():
 
 
 # @pytest.mark.beam
-# def test_magnify_beta_function():
-#     # TODO
+def test_magnify_beta_function():
+
+    np.random.seed(42)
+
+    x_offset = 5.354e-6
+    y_offset = 0.1516e-6
+    ux_offset = 0.0
+    uy_offset = 1583681.8243787317
+    beta_mag = 5.0  # Magnify
+
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=beta_mag, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=ux_offset, y_angle=uy_offset/energy2proper_velocity(3e9))
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+
+    beam.magnify_beta_function(beta_mag=beta_mag, axis_defining_beam=beam)
+    
+    # calculate beam (not beta) magnification
+    mag = np.sqrt(beta_mag)
+
+    # Examine beam
+    assert len(beam) == len(initial_beam)
+    assert np.isclose(beam.particle_mass, SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge(), initial_beam.charge(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge_sign(), -1, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.total_energy(), initial_beam.total_energy(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.weightings(), initial_beam.weightings(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.energy(), initial_beam.energy(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.rel_energy_spread(), initial_beam.rel_energy_spread(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.z_offset(), initial_beam.z_offset(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.bunch_length(), initial_beam.bunch_length(), rtol=1e-15, atol=0.0)
+
+    assert np.isclose(beam.x_offset(), x_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.y_offset(), y_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beam_size_x(), initial_beam.beam_size_x()*mag, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beam_size_y(), initial_beam.beam_size_y()*mag, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.x_angle(), ux_offset/energy2proper_velocity(3e9), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.y_angle(), uy_offset/energy2proper_velocity(3e9), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.divergence_x(), initial_beam.divergence_x()/mag, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.divergence_y(), initial_beam.divergence_y()/mag, rtol=1e-15, atol=0.0)
+
+    assert np.allclose(beam.xs(), (initial_beam.xs()-initial_beam.x_offset())*mag + initial_beam.x_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.ys(), (initial_beam.ys()-initial_beam.y_offset())*mag + initial_beam.y_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), (initial_beam.uxs()-initial_beam.ux_offset())/mag + initial_beam.ux_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), (initial_beam.uys()-initial_beam.uy_offset())/mag + initial_beam.uy_offset(), rtol=1e-15, atol=0.0)
+
+
+    beta_mag = 1/beta_mag  # De-magnify
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=beta_mag, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=ux_offset, y_angle=uy_offset/energy2proper_velocity(3e9))
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+
+    beam.magnify_beta_function(beta_mag=beta_mag, axis_defining_beam=beam)
+    
+    # calculate beam (not beta) magnification
+    mag = np.sqrt(beta_mag)
+
+    # Examine beam
+    assert len(beam) == len(initial_beam)
+    assert np.isclose(beam.particle_mass, SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge(), initial_beam.charge(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge_sign(), -1, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.total_energy(), initial_beam.total_energy(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.weightings(), initial_beam.weightings(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.energy(), initial_beam.energy(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.rel_energy_spread(), initial_beam.rel_energy_spread(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.z_offset(), initial_beam.z_offset(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.bunch_length(), initial_beam.bunch_length(), rtol=1e-15, atol=0.0)
+
+    assert np.isclose(beam.x_offset(), x_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.y_offset(), y_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beam_size_x(), initial_beam.beam_size_x()*mag, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beam_size_y(), initial_beam.beam_size_y()*mag, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.x_angle(), ux_offset/energy2proper_velocity(3e9), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.y_angle(), uy_offset/energy2proper_velocity(3e9), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.divergence_x(), initial_beam.divergence_x()/mag, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.divergence_y(), initial_beam.divergence_y()/mag, rtol=1e-15, atol=0.0)
+
+    assert np.allclose(beam.xs(), (initial_beam.xs()-initial_beam.x_offset())*mag + initial_beam.x_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.ys(), (initial_beam.ys()-initial_beam.y_offset())*mag + initial_beam.y_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), (initial_beam.uxs()-initial_beam.ux_offset())/mag + initial_beam.ux_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), (initial_beam.uys()-initial_beam.uy_offset())/mag + initial_beam.uy_offset(), rtol=1e-15, atol=0.0)
+
 
 
 # @pytest.mark.beam

From 3b16ce56cf33c94646041b08aaf5355051828467 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 00:39:02 +0200
Subject: [PATCH 27/51] Added test_apply_betatron_damping(),
 test_magnify_beta_function(), test_transport() and
 test_flip_transverse_phase_spaces().

---
 tests/test_beam.py | 221 ++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 210 insertions(+), 11 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 4258c4ad..22cf8255 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -1300,14 +1300,103 @@ def test_scale_norm_emittance_y():
         assert str(err) == 'Normalised emittance cannot be negative.'
 
 
-# @pytest.mark.beam
-# def test_apply_betatron_damping():
-    
-    
+@pytest.mark.beam
+def test_apply_betatron_damping():
 
+    np.random.seed(42)
 
-# @pytest.mark.beam
+    x_offset = 5.354e-6
+    y_offset = 0.1516e-6
+    ux_offset = 739215.7882185677
+    uy_offset = 1583681.8243787317
+    deltaE = 1.516e9
+
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=np.sqrt(3/4.516), energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=ux_offset/energy2proper_velocity(3e9), y_angle=uy_offset/energy2proper_velocity(3e9))
+
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+
+    # calculate beam (not beta) magnification
+    gammasBoosted = energy2gamma(abs(initial_beam.Es() + deltaE))
+    beta_mag = np.sqrt(initial_beam.gammas()/gammasBoosted)
+    assert np.isclose(beta_mag.mean(), np.sqrt(3/4.516), rtol=1e-15, atol=0.0)
+    mag = np.sqrt(beta_mag)
+
+    beam.apply_betatron_damping(deltaE, axis_defining_beam=beam)
+
+    # Examine beam
+    assert len(beam) == len(initial_beam)
+    assert np.isclose(beam.particle_mass, SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge(), initial_beam.charge(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge_sign(), -1, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.total_energy(), initial_beam.total_energy(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.weightings(), initial_beam.weightings(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.energy(), initial_beam.energy(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.rel_energy_spread(), initial_beam.rel_energy_spread(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.z_offset(), initial_beam.z_offset(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.bunch_length(), initial_beam.bunch_length(), rtol=1e-15, atol=0.0)
+
+    assert np.isclose(beam.x_offset(), x_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.y_offset(), y_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beam_size_x(), initial_beam.beam_size_x()*mag.mean(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beam_size_y(), initial_beam.beam_size_y()*mag.mean(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.x_angle(), ux_offset/energy2proper_velocity(3e9), rtol=1e-12, atol=0.0)
+    assert np.isclose(beam.y_angle(), uy_offset/energy2proper_velocity(3e9), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.divergence_x(), initial_beam.divergence_x()/mag.mean(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.divergence_y(), initial_beam.divergence_y()/mag.mean(), rtol=1e-15, atol=0.0)
+
+    assert np.allclose(beam.xs(), (initial_beam.xs()-initial_beam.x_offset())*mag + initial_beam.x_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.ys(), (initial_beam.ys()-initial_beam.y_offset())*mag + initial_beam.y_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), (initial_beam.uxs()-initial_beam.ux_offset())/mag + initial_beam.ux_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), (initial_beam.uys()-initial_beam.uy_offset())/mag + initial_beam.uy_offset(), rtol=1e-15, atol=0.0)
+
+
+    # Test no energy change
+    deltaE = 0.0
+
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=1.0, energy=3e9, rel_energy_spread=0.01, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=ux_offset/energy2proper_velocity(3e9), y_angle=uy_offset/energy2proper_velocity(3e9))
+
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+
+    # calculate beam (not beta) magnification
+    gammasBoosted = energy2gamma(abs(initial_beam.Es() + deltaE))
+    beta_mag = np.sqrt(initial_beam.gammas()/gammasBoosted)
+    assert np.isclose(beta_mag.mean(), 1.0, rtol=1e-15, atol=0.0)
+    mag = np.sqrt(beta_mag)
+
+    beam.apply_betatron_damping(deltaE, axis_defining_beam=beam)
+
+    # Examine beam
+    assert len(beam) == len(initial_beam)
+    assert np.isclose(beam.particle_mass, SI.m_e, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge(), initial_beam.charge(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.charge_sign(), -1, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.total_energy(), initial_beam.total_energy(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.weightings(), initial_beam.weightings(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.energy(), initial_beam.energy(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.rel_energy_spread(), initial_beam.rel_energy_spread(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.z_offset(), initial_beam.z_offset(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.bunch_length(), initial_beam.bunch_length(), rtol=1e-15, atol=0.0)
+
+    assert np.isclose(beam.x_offset(), x_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.y_offset(), y_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beam_size_x(), initial_beam.beam_size_x()*mag.mean(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beam_size_y(), initial_beam.beam_size_y()*mag.mean(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.x_angle(), ux_offset/energy2proper_velocity(3e9), rtol=1e-12, atol=0.0)
+    assert np.isclose(beam.y_angle(), uy_offset/energy2proper_velocity(3e9), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.divergence_x(), initial_beam.divergence_x()/mag.mean(), rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.divergence_y(), initial_beam.divergence_y()/mag.mean(), rtol=1e-15, atol=0.0)
+
+    assert np.allclose(beam.xs(), (initial_beam.xs()-initial_beam.x_offset())*mag + initial_beam.x_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.ys(), (initial_beam.ys()-initial_beam.y_offset())*mag + initial_beam.y_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uxs(), (initial_beam.uxs()-initial_beam.ux_offset())/mag + initial_beam.ux_offset(), rtol=1e-15, atol=0.0)
+    assert np.allclose(beam.uys(), (initial_beam.uys()-initial_beam.uy_offset())/mag + initial_beam.uy_offset(), rtol=1e-15, atol=0.0)
+    
+
+@pytest.mark.beam
 def test_magnify_beta_function():
+    "Tests for ``Beam.magnify_beta_function(beta_mag, axis_defining_beam=None)``, which magnifies beta functions (increases beam size, decreases divergence for beta_mag > 1.0)."
 
     np.random.seed(42)
 
@@ -1390,15 +1479,125 @@ def test_magnify_beta_function():
     assert np.allclose(beam.uys(), (initial_beam.uys()-initial_beam.uy_offset())/mag + initial_beam.uy_offset(), rtol=1e-15, atol=0.0)
 
 
+@pytest.mark.beam
+def test_transport():
+
+    x_offset = 0                                                                    # [m]
+    y_offset = 0                                                                    # [m]
+    x_angle = 0                                                                     # [rad]
+    y_angle = 0                                                                     # [rad]
+    L = 0                                                                           # [m]
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=x_angle, y_angle=y_angle)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
 
-# @pytest.mark.beam
-# def test_transport():
-#     # TODO
+    beam.transport(L)
+    assert np.allclose(beam.xs(), initial_beam.xs()+L*initial_beam.xps())
+    assert np.allclose(beam.ys(), initial_beam.ys()+L*initial_beam.yps())
+
+    x_offset = 5.354e-6                                                             # [m]
+    y_offset = 0.1516e-6                                                            # [m]
+    x_angle = 6.6e-6                                                                # [rad]
+    y_angle = 0.42e-6                                                               # [rad]
+    L = 0.0                                                                         # [m]
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=x_angle, y_angle=y_angle)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+    beam.transport(L)
+    assert np.allclose(beam.xs(), initial_beam.xs()+L*initial_beam.xps())
+    assert np.allclose(beam.ys(), initial_beam.ys()+L*initial_beam.yps())
+
+    x_offset = 5.354e-6                                                             # [m]
+    y_offset = 0.1516e-6                                                            # [m]
+    x_angle = 6.6e-6                                                                # [rad]
+    y_angle = 0.42e-6                                                               # [rad]
+    L = 7.8                                                                         # [m]
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=x_angle, y_angle=y_angle)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+    beam.transport(L)
+    assert np.allclose(beam.xs(), initial_beam.xs()+L*initial_beam.xps())
+    assert np.allclose(beam.ys(), initial_beam.ys()+L*initial_beam.yps())
+
+    x_offset = -5.354e-6                                                            # [m]
+    y_offset = 0.1516e-6                                                            # [m]
+    x_angle = 6.6e-6                                                                # [rad]
+    y_angle = -0.42e-6                                                              # [rad]
+    L = 7.8                                                                         # [m]
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=x_angle, y_angle=y_angle)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+    beam.transport(L)
+    assert np.allclose(beam.xs(), initial_beam.xs()+L*initial_beam.xps())
+    assert np.allclose(beam.ys(), initial_beam.ys()+L*initial_beam.yps())
+
+    x_offset = -5.354e-6                                                            # [m]
+    y_offset = 0.1516e-6                                                            # [m]
+    x_angle = 6.6e-6                                                                # [rad]
+    y_angle = -0.42e-6                                                              # [rad]
+    L = -7.8                                                                        # [m]
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=x_angle, y_angle=y_angle)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+    beam.transport(L)
+    assert np.allclose(beam.xs(), initial_beam.xs()+L*initial_beam.xps())
+    assert np.allclose(beam.ys(), initial_beam.ys()+L*initial_beam.yps())
 
 
-# @pytest.mark.beam
-# def test_flip_transverse_phase_spaces():
-#     # TODO
+@pytest.mark.beam
+def test_flip_transverse_phase_spaces():
+
+    x_offset = 0                                                                    # [m]
+    y_offset = 0                                                                    # [m]
+    x_angle = 0                                                                     # [rad]
+    y_angle = 0                                                                     # [rad]
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=x_angle, y_angle=y_angle)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+    beam.flip_transverse_phase_spaces(flip_momenta=False, flip_positions=False)
+    assert np.allclose(beam.xs(), initial_beam.xs())
+    assert np.allclose(beam.ys(), initial_beam.ys())
+    assert np.allclose(beam.uxs(), initial_beam.uxs())
+    assert np.allclose(beam.uys(), initial_beam.uys())
+
+    x_offset = -5.354e-6                                                            # [m]
+    y_offset = 0.1516e-6                                                            # [m]
+    x_angle = 6.6e-6                                                                # [rad]
+    y_angle = -0.42e-6                                                              # [rad]
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=x_angle, y_angle=y_angle)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+    beam.flip_transverse_phase_spaces(flip_momenta=True, flip_positions=False)
+    assert np.allclose(beam.xs(), initial_beam.xs())
+    assert np.allclose(beam.ys(), initial_beam.ys())
+    assert np.allclose(beam.uxs(), -initial_beam.uxs())
+    assert np.allclose(beam.uys(), -initial_beam.uys())
+
+    x_offset = 5.354e-6                                                             # [m]
+    y_offset = -0.1516e-6                                                           # [m]
+    x_angle = -6.6e-6                                                               # [rad]
+    y_angle = -0.42e-6                                                              # [rad]
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=x_angle, y_angle=y_angle)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+    beam.flip_transverse_phase_spaces(flip_momenta=False, flip_positions=True)
+    assert np.allclose(beam.xs(), -initial_beam.xs())
+    assert np.allclose(beam.ys(), -initial_beam.ys())
+    assert np.allclose(beam.uxs(), initial_beam.uxs())
+    assert np.allclose(beam.uys(), initial_beam.uys())
+
+    x_offset = 5.354e-6                                                             # [m]
+    y_offset = 0.1516e-6                                                            # [m]
+    x_angle = 6.6e-6                                                                # [rad]
+    y_angle = 0.42e-6                                                               # [rad]
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.0, z_offset=0.0, y_offset=y_offset, x_offset=x_offset, x_angle=x_angle, y_angle=y_angle)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+    beam.flip_transverse_phase_spaces(flip_momenta=True, flip_positions=True)
+    assert np.allclose(beam.xs(), -initial_beam.xs())
+    assert np.allclose(beam.ys(), -initial_beam.ys())
+    assert np.allclose(beam.uxs(), -initial_beam.uxs())
+    assert np.allclose(beam.uys(), -initial_beam.uys())
 
 
 @pytest.mark.beam

From ccd580d57e36dac81d1751f7c887f0f3d50fb80e Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 00:45:10 +0200
Subject: [PATCH 28/51] Minor changes in Beam.apply_betatron_damping() and
 Beam.flip_transverse_phase_spaces().

---
 abel/classes/beam.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index cc969db4..84b8c6c7 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -1234,13 +1234,13 @@ def scale_norm_emittance_y(self, norm_emit_ny):
         self.set_uys(self.uys() * np.sqrt(scale_factor))
         
     # betatron damping (must be done before acceleration)
-    def apply_betatron_damping(self, deltaE):
+    def apply_betatron_damping(self, deltaE, axis_defining_beam=None):
         gammasBoosted = energy2gamma(abs(self.Es()+deltaE))
         betamag = np.sqrt(self.gammas()/gammasBoosted)
-        self.magnify_beta_function(betamag)
+        self.magnify_beta_function(betamag, axis_defining_beam)
         
     
-    # magnify beta function (increase beam size, decrease divergence)
+    # magnify beta function (increase beam size, decrease divergence for beta_mag > 1.0)
     def magnify_beta_function(self, beta_mag, axis_defining_beam=None):
         
         # calculate beam (not beta) magnification
@@ -1273,7 +1273,7 @@ def flip_transverse_phase_spaces(self, flip_momenta=True, flip_positions=False):
         if flip_momenta:
             self.set_uxs(-self.uxs())
             self.set_uys(-self.uys())
-        elif flip_positions:
+        if flip_positions:
             self.set_xs(-self.xs())
             self.set_ys(-self.ys())
 

From 7e9b64b886728adeb656d849337a46ee9041a602 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 01:32:05 +0200
Subject: [PATCH 29/51] Changes to the default relative tolerance and absolute
 tolerance of Beam.comp_beams().

---
 abel/classes/beam.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index 84b8c6c7..a8ad786e 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -266,7 +266,7 @@ def ts(self):
         return self.zs()/SI.c
     
 
-    def comp_beams(beam1, beam2, comp_location=False, rtol=1e-05, atol=1e-08):
+    def comp_beams(beam1, beam2, comp_location=False, rtol=1e-15, atol=0.0):
         """
         Compare the phase spaces of two beams. Chekcks if all arrays are element-wise equal within given tolerances.
 
@@ -279,10 +279,10 @@ def comp_beams(beam1, beam2, comp_location=False, rtol=1e-05, atol=1e-08):
             Flag for comparing the location of the beams. Default set to ``False``.
             
         rtol : float, optional
-            The relative tolerance parameter (see [1]_). Default set to 1e-5.
+            The relative tolerance parameter (see [1]_). Default set to 1e-15.
             
         rb_fit_obj : float, optional
-            The absolute tolerance parameter (see [1]_). Default set to 1e-8.
+            The absolute tolerance parameter (see [1]_). Default set to 0.0.
         
             
         Returns

From 2437ed1837022ef797b07217830d318264704188 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 01:32:59 +0200
Subject: [PATCH 30/51] Added test_save_load() to test_beam.py.

---
 tests/test_beam.py | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 22cf8255..56881e1a 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -1711,4 +1711,27 @@ def test_plot_transverse_profile():
 # def test_plot_bunch_pattern():
 #     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
 #     beam = source.track()
-#     beam.plot_bunch_pattern()
\ No newline at end of file
+#     beam.plot_bunch_pattern()
+
+
+
+############# Tests of saving and loading #############
+@pytest.mark.beam
+def test_save_load():
+
+    import os
+    #from abel.classes.beam import comp_beams
+
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.01, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+
+    save_dir = 'tests' + os.sep + 'data' + os.sep + 'test_beam' + os.sep
+    if not os.path.exists(save_dir):
+        os.mkdir(save_dir)
+    filename = save_dir + os.sep + 'beam_test_save.h5'
+
+    beam.save(filename=filename)
+    loaded_beam = Beam.load(filename=filename)
+
+    Beam.comp_beams(beam, loaded_beam, rtol=1e-15, atol=0.0)
+    shutil.rmtree(save_dir)
\ No newline at end of file

From e04994dc8179031da70364fd05cc2d69586595f8 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 12:10:11 +0200
Subject: [PATCH 31/51] Added beam reference file for tests.

---
 .../shot_000/beam_003_00048.558626.h5         | Bin 0 -> 1083528 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 tests/data/test_StagePrtclTransWakeInstability_beamline/test_baseline_linac/shot_000/beam_003_00048.558626.h5

diff --git a/tests/data/test_StagePrtclTransWakeInstability_beamline/test_baseline_linac/shot_000/beam_003_00048.558626.h5 b/tests/data/test_StagePrtclTransWakeInstability_beamline/test_baseline_linac/shot_000/beam_003_00048.558626.h5
new file mode 100644
index 0000000000000000000000000000000000000000..b9b28b154e015a41605e514f2da693737bf600bb
GIT binary patch

From de2a3afff8da0c3dfb3cc33a044311f6af2a88f7 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 15:08:12 +0200
Subject: [PATCH 32/51] Changed the limit for minimum energy in the Beam class
 to gamma=10.

---
 abel/classes/beam.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index a8ad786e..fa5fefeb 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -75,7 +75,7 @@ def set_phase_space(self, Q, xs, ys, zs, uxs=None, uys=None, uzs=None, pxs=None,
         self.set_zs(zs)
 
         # minimum thresholds for energy, uz and pz
-        energy_thres = 100*particle_mass*SI.c**2/SI.e  # [eV], 100 * particle rest energy.
+        energy_thres = 10*particle_mass*SI.c**2/SI.e  # [eV], 10 * particle rest energy. Gives beta=0.995.
         uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=particle_mass)
         pz_thres = gamma2momentum(energy2gamma(energy_thres, unit='eV', m=particle_mass))
         
@@ -200,7 +200,7 @@ def set_uxs(self, uxs):
     def set_uys(self, uys):
         self.__phasespace[4,:] = uys
     def set_uzs(self, uzs):
-        energy_thres = 100*self.particle_mass*SI.c**2/SI.e  # [eV], 100 * particle rest energy.
+        energy_thres = 10*self.particle_mass*SI.c**2/SI.e  # [eV], 10 * particle rest energy. Gives beta=0.995.
         uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=self.particle_mass)
         if np.any(uzs < uz_thres):
             raise ValueError('uzs contains values that are too small.')
@@ -211,7 +211,7 @@ def set_xps(self, xps):
     def set_yps(self, yps):
         self.set_uys(yps*self.uzs())
     def set_Es(self, Es):
-        energy_thres = 100*self.particle_mass*SI.c**2/SI.e  # [eV], 100 * particle rest energy.
+        energy_thres = 10*self.particle_mass*SI.c**2/SI.e  # [eV], 10 * particle rest energy. Gives beta=0.995.
         if np.any(Es < energy_thres):
             raise ValueError('Es contains values that are too small.')
         self.set_uzs(energy2proper_velocity(Es))

From 8f487edbb91df3d04c8cd1bf84b49462aec40e90 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 16:59:07 +0200
Subject: [PATCH 33/51] Added controls to weighted_mean(), weighted_std() and
 weighted_cov() under statistics.py for when the sum of weights is zero.

---
 abel/utilities/statistics.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/abel/utilities/statistics.py b/abel/utilities/statistics.py
index 9dbef183..61543348 100644
--- a/abel/utilities/statistics.py
+++ b/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)

From e21c2393a0f98b8f58c9eb3378ba51079c9553df Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 17:01:06 +0200
Subject: [PATCH 34/51] Edited the str operator and added controls for
 __init__() and reset_phase_space() for when num_particles is an integer
 smaller than 1.

---
 abel/classes/beam.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index fa5fefeb..a7ce4b1e 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -22,9 +22,9 @@ class Beam():
     def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1, bunch_separation=0.0):
 
         # check the inputs
-        if num_particles < 1:
-            raise ValueError('num_particles cannot be lower than 1.')
-        if num_bunches_in_train < 1:
+        if num_particles < 1 or not isinstance(num_particles, int):
+            raise ValueError('num_particles must be an integer larger than 1.')
+        if num_bunches_in_train < 1 or not isinstance(num_bunches_in_train, int):
             raise ValueError('num_bunches_in_train cannot be lower than 1.')
         if bunch_separation < 0.0:
             raise ValueError('bunch_separation cannot be negative.')
@@ -46,8 +46,8 @@ def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1,
     
     # reset phase space
     def reset_phase_space(self, num_particles):
-        if num_particles < 1:
-            raise ValueError('num_particles cannot be lower than 1.')
+        if num_particles < 1 or not isinstance(num_particles, int):
+            raise ValueError('num_particles must be an integer larger than 1.')
         
         self.__phasespace = np.zeros((8, num_particles))
     
@@ -143,7 +143,10 @@ def __len__(self):
     
     # string operator (called when printing)
     def __str__(self):
-        return f"Beam: {len(self)} macroparticles, {self.charge()*1e9:.2f} nC, {self.energy()/1e9:.2f} GeV"
+        if np.sum(self.weightings()) == 0.0:
+            return f"Beam: {len(self)} macroparticles, {self.charge()*1e9:.2f} nC"
+        else:
+            return f"Beam: {len(self)} macroparticles, {self.charge()*1e9:.2f} nC, {self.energy()/1e9:.2f} GeV"
         
         
     ## BUNCH PATTERN

From 7d2b100f1abe5878bdcab312cda8123c40a4a635 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 17:04:25 +0200
Subject: [PATCH 35/51] Edited some tests to purposedly trigger exceptions in
 test_beam.py.

---
 tests/test_beam.py | 75 +++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 64 insertions(+), 11 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 56881e1a..7d58b6f7 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -102,6 +102,18 @@ def test_initialization():
     assert beam.stage_number == 0
     assert beam.location == 0
 
+    # Purposedly trigger exceptions
+    with pytest.raises(ValueError):
+        beam = Beam(num_particles=-1)  # Should raise an error if attempted to initiate with negative number of particles.
+    with pytest.raises(ValueError):
+        beam = Beam(num_particles=1000.2)
+    with pytest.raises(ValueError):
+        beam = Beam(num_bunches_in_train=-1)
+    with pytest.raises(ValueError):
+        beam = Beam(num_bunches_in_train=1300.2)
+    with pytest.raises(ValueError):
+        beam = Beam(bunch_separation=-1)
+
 
 @pytest.mark.beam
 def test_set_phase_space():
@@ -115,7 +127,7 @@ def test_set_phase_space():
     uxs = np.random.rand(num_particles)
     uys = np.random.rand(num_particles)
 
-    energy_thres = 1001*SI.m_e*SI.c**2/SI.e  # [eV], 1000 * particle rest energy.
+    energy_thres = 13*SI.m_e*SI.c**2/SI.e  # [eV], 13 * particle rest energy.
     uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
     uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
 
@@ -274,6 +286,12 @@ def test_reset_phase_space():
     assert beam._Beam__phasespace.shape == (8, 10)
     assert (beam._Beam__phasespace == 0).all()
 
+    # Purposedly trigger exceptions
+    with pytest.raises(ValueError):
+        beam.reset_phase_space(-11)  # Should raise an error if attempted to initiate with negative number of particles.
+    with pytest.raises(ValueError):
+        beam.reset_phase_space(10.001)
+
 
 @pytest.mark.beam
 def test_delitem():
@@ -351,21 +369,38 @@ def test_len():
     assert len(beam) == num_particles
 
 
-# @pytest.mark.beam TODO
-# def test_str():
-#     "Test the __str__ method for correct formatting of beam properties."
-#     beam = Beam()
-#     beam.reset_phase_space(5)
-#     string_repr = str(beam)
-#     assert "Beam:" in string_repr
-#     assert "5 macroparticles" in string_repr
+@pytest.mark.beam
+def test_str():
+    "Test the __str__ method for correct formatting of beam properties."
+
+    beam = Beam()
+    num_particles = 5
+    beam.reset_phase_space(num_particles)
+    beam_str = str(beam)
+    assert "Beam:" in beam_str
+    assert "0.00 nC" in beam_str
+    assert str(num_particles) + " macroparticles" in beam_str
+
+    num_particles = 5005
+    Es = np.random.normal(151.567e9, 0.02*151.567e9, num_particles)
+    qs = np.full(num_particles, -1.5e10*SI.e/num_particles)
+    beam = Beam()
+    beam.reset_phase_space(num_particles)
+    beam.particle_mass = SI.m_e
+    beam.set_qs(qs)
+    beam.set_Es(Es)
+    beam_str = str(beam)
+    assert "Beam:" in beam_str
+    assert "{:.2f}".format(-1.5e10*SI.e*1e9) in beam_str
+    assert str(num_particles) + " macroparticles" in beam_str
+    assert "{:.2f}".format(beam.energy()/1e9) + " GeV" in beam_str
 
 
 @pytest.mark.beam
 def test_copy_particle_charge():
 
     num_particles = 8051
-    energy_thres = 1001*SI.m_e*SI.c**2/SI.e  # [eV], 1000 * particle rest energy.
+    energy_thres = 11*SI.m_e*SI.c**2/SI.e  # [eV], 11 * particle rest energy.
     uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
     
     beam1 = Beam()
@@ -393,6 +428,22 @@ def test_copy_particle_charge():
     beam1.copy_particle_charge(beam2)
     assert (beam1._Beam__phasespace[6, :] == median_charge).all()  # Verify charge is copied correctly
 
+    # Purposedly trigger exceptions
+    with pytest.raises(ValueError):
+        beam1 = Beam()
+        beam2 = Beam()
+        beam1.copy_particle_charge(beam2)
+    with pytest.raises(ValueError):
+        beam1 = Beam()
+        beam2 = Beam()
+        beam2.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs, weightings=ws)
+        beam1.copy_particle_charge(beam2)
+    with pytest.raises(ValueError):
+        beam1 = Beam()
+        beam2 = Beam()
+        beam1.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs, weightings=ws)
+        beam1.copy_particle_charge(beam2)
+
 
 @pytest.mark.beam
 def test_scale_charge():
@@ -458,6 +509,9 @@ def test_scale_energy():
 
 
 
+
+
+
 ############# Tests on bunch pattern #############
 # TODO
 
@@ -1720,7 +1774,6 @@ def test_plot_transverse_profile():
 def test_save_load():
 
     import os
-    #from abel.classes.beam import comp_beams
 
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, rel_energy_spread=0.01, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
     beam = source.track()

From b3556f61d7d80c8950991ad0be09f5fe37d5fd87 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 17:30:05 +0200
Subject: [PATCH 36/51] Added to test_rs(), test_deltas(),
 test_transverse_vector(), test_norm_transverse_vector(), and test_gamma_xy()
 to test_beam.py.

---
 tests/test_beam.py | 111 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 110 insertions(+), 1 deletion(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 7d58b6f7..9b1b1d51 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -508,6 +508,83 @@ def test_scale_energy():
         assert str(err) == 'Es contains values that are too small.'
 
 
+@pytest.mark.beam
+def test_rs():
+    beam = Beam()
+    num_particles = 4352
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs)
+
+    assert np.allclose(beam.rs(), np.sqrt(xs**2 + ys**2), rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_deltas():
+    beam = Beam()
+    num_particles = 4352
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    Es = np.random.normal(1e9, 0.02*1e9, num_particles)
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs, Es=Es)
+
+    assert np.allclose(beam.deltas(), beam.pzs()/np.mean(beam.pzs())-1, rtol=1e-15, atol=0.0)
+    pz0 = np.mean(beam.pzs())*1.1
+    assert np.allclose(beam.deltas(pz0=pz0), beam.pzs()/pz0-1, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_transverse_vector():
+    num_particles = 8989
+    energy_thres = 11*SI.m_e*SI.c**2/SI.e  # [eV], 11 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    vector = np.zeros((4,len(beam))) 
+    vector[0,:] = beam.xs()
+    vector[1,:] = beam.xps()
+    vector[2,:] = beam.ys()
+    vector[3,:] = beam.yps()
+
+    assert np.allclose(beam.transverse_vector(), vector, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_norm_transverse_vector():
+    num_particles = 8989
+    energy_thres = 11*SI.m_e*SI.c**2/SI.e  # [eV], 11 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    vector = np.zeros((4,len(beam))) 
+    vector[0,:] = beam.xs()
+    vector[1,:] = beam.uxs()/SI.c 
+    vector[2,:] = beam.ys()
+    vector[3,:] = beam.uys()/SI.c 
+
+    assert np.allclose(beam.norm_transverse_vector(), vector, rtol=1e-15, atol=0.0)
 
 
 
@@ -687,6 +764,37 @@ def test_param_calcs_generate_symm_trace_space_xyz():
     assert np.isclose(beam.norm_emittance_y(), geo_emitt_y*beam.gamma(), rtol=1e-2, atol=0.0)
 
 
+@pytest.mark.beam
+def test_gamma_xy():
+    num_particles = 8989
+    energy_thres = 11*SI.m_e*SI.c**2/SI.e  # [eV], 11 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
+    Q = -SI.e * 1.0e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    covx = weighted_cov(beam.xs(), beam.xps(), beam.weightings(), clean=False)
+    gamma_x = covx[1,1]/np.sqrt(np.linalg.det(covx))
+    covy = weighted_cov(beam.ys(), beam.yps(), beam.weightings(), clean=False)
+    gamma_y = covy[1,1]/np.sqrt(np.linalg.det(covy))
+    assert np.isclose(beam.gamma_x(), gamma_x, rtol=1e-10, atol=0.0)
+    assert np.isclose(beam.gamma_y(), gamma_y, rtol=1e-10, atol=0.0)
+
+    covx = weighted_cov(beam.xs(), beam.xps(), beam.weightings(), clean=True)
+    gamma_x = covx[1,1]/np.sqrt(np.linalg.det(covx))
+    covy = weighted_cov(beam.ys(), beam.yps(), beam.weightings(), clean=True)
+    gamma_y = covy[1,1]/np.sqrt(np.linalg.det(covy))
+    assert np.isclose(beam.gamma_x(clean=True), gamma_x, rtol=1e-10, atol=0.0)
+    assert np.isclose(beam.gamma_y(clean=True), gamma_y, rtol=1e-10, atol=0.0)
+
+
 
 ############# Tests for beam projections #############
 @pytest.mark.beam
@@ -1786,5 +1894,6 @@ def test_save_load():
     beam.save(filename=filename)
     loaded_beam = Beam.load(filename=filename)
 
-    Beam.comp_beams(beam, loaded_beam, rtol=1e-15, atol=0.0)
+    Beam.comp_beams(beam, loaded_beam, comp_location=True, rtol=1e-15, atol=0.0)
+    Beam.comp_beams(beam, loaded_beam, comp_location=True)
     shutil.rmtree(save_dir)
\ No newline at end of file

From bb597f72e97095b57282813f58caefc036b8d712 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 18:50:48 +0200
Subject: [PATCH 37/51] Added several more tests for beam statistics.

---
 tests/test_beam.py | 438 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 428 insertions(+), 10 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 9b1b1d51..fad34041 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -129,7 +129,7 @@ def test_set_phase_space():
 
     energy_thres = 13*SI.m_e*SI.c**2/SI.e  # [eV], 13 * particle rest energy.
     uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
-    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e), size=num_particles)
 
     beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
 
@@ -409,7 +409,7 @@ def test_copy_particle_charge():
     zs = np.random.rand(num_particles)
     uxs = np.random.rand(num_particles)
     uys = np.random.rand(num_particles)
-    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e), size=num_particles)
     Q = -SI.e * 1.0e10
     beam1.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
 
@@ -419,7 +419,7 @@ def test_copy_particle_charge():
     zs = np.random.rand(num_particles)
     uxs = np.random.rand(num_particles)
     uys = np.random.rand(num_particles)
-    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e), size=num_particles)
     Q = -SI.e * 3.4e10
     ws = np.random.rand(num_particles)*Q/(-SI.e)
     beam2.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs, weightings=ws)
@@ -549,7 +549,7 @@ def test_transverse_vector():
     zs = np.random.rand(num_particles)
     uxs = np.random.rand(num_particles)
     uys = np.random.rand(num_particles)
-    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e), size=num_particles)
     Q = -SI.e * 1.0e10
     beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
 
@@ -574,7 +574,7 @@ def test_norm_transverse_vector():
     zs = np.random.rand(num_particles)
     uxs = np.random.rand(num_particles)
     uys = np.random.rand(num_particles)
-    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e), size=num_particles)
     Q = -SI.e * 1.0e10
     beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
 
@@ -764,6 +764,345 @@ def test_param_calcs_generate_symm_trace_space_xyz():
     assert np.isclose(beam.norm_emittance_y(), geo_emitt_y*beam.gamma(), rtol=1e-2, atol=0.0)
 
 
+@pytest.mark.beam
+def test_total_particles():
+    
+    num_particles = 4344
+    energy_thres = 1200*SI.m_e*SI.c**2/SI.e  # [eV], 1200 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.8e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    total_particles = int(np.nansum(np.full(num_particles, Q/num_particles/-SI.e)))
+    assert beam.total_particles() == total_particles
+
+    # TODO: set some nans in beam
+
+
+@pytest.mark.beam
+def test_charge():
+    
+    num_particles = 4344
+    energy_thres = 1200*SI.m_e*SI.c**2/SI.e  # [eV], 1200 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.8e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    charge = np.nansum(np.full(num_particles, Q/num_particles))
+    assert np.isclose(beam.charge(), charge, rtol=1e-15, atol=0.0)
+
+    # TODO: set some nans in beam
+
+
+@pytest.mark.beam
+def test_energy():
+    
+    num_particles = 4344
+    energy_thres = 1200*SI.m_e*SI.c**2/SI.e  # [eV], 1200 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.8e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    Es = proper_velocity2energy(uzs, unit='eV', m=SI.m_e)
+    energy = weighted_mean(Es, beam.weightings(), clean=False) 
+    assert np.isclose(beam.energy(clean=False), energy, rtol=1e-15, atol=0.0)
+    energy = weighted_mean(Es, beam.weightings(), clean=True) 
+    assert np.isclose(beam.energy(clean=True), energy, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_gamma():
+    
+    num_particles = 4344
+    energy_thres = 1200*SI.m_e*SI.c**2/SI.e  # [eV], 1200 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.8e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    gammas = proper_velocity2gamma(uzs)
+    gamma = weighted_mean(gammas, beam.weightings(), clean=False) 
+    assert np.isclose(beam.gamma(clean=False), gamma, rtol=1e-15, atol=0.0)
+    gamma = weighted_mean(gammas, beam.weightings(), clean=True) 
+    assert np.isclose(beam.gamma(clean=True), gamma, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_total_energy():
+    
+    num_particles = 4344
+    energy_thres = 1200*SI.m_e*SI.c**2/SI.e  # [eV], 1200 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.8e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    total_energy = SI.e * np.nansum(beam.weightings()*beam.Es()) 
+    assert np.isclose(beam.total_energy(), total_energy, rtol=1e-15, atol=0.0)
+
+    # TODO: set some nans in beam
+
+
+@pytest.mark.beam
+def test_energy_spread():
+    
+    num_particles = 4344
+    energy_thres = 1200*SI.m_e*SI.c**2/SI.e  # [eV], 1200 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.8e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    Es = proper_velocity2energy(uzs, unit='eV', m=SI.m_e)
+    energy_spread = weighted_std(Es, beam.weightings(), clean=False) 
+    assert np.isclose(beam.energy_spread(clean=False), energy_spread, rtol=1e-15, atol=0.0)
+    energy_spread = weighted_std(Es, beam.weightings(), clean=True) 
+    assert np.isclose(beam.energy_spread(clean=True), energy_spread, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_rel_energy_spread():
+    
+    num_particles = 4344
+    energy_thres = 1200*SI.m_e*SI.c**2/SI.e  # [eV], 1200 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.8e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    Es = proper_velocity2energy(uzs, unit='eV', m=SI.m_e)
+    rel_energy_spread = weighted_std(Es, beam.weightings(), clean=False)/weighted_mean(Es, beam.weightings(), clean=False) 
+    assert np.isclose(beam.rel_energy_spread(clean=False), rel_energy_spread, rtol=1e-15, atol=0.0)
+    rel_energy_spread = weighted_std(Es, beam.weightings(), clean=True)/weighted_mean(Es, beam.weightings(), clean=True) 
+    assert np.isclose(beam.rel_energy_spread(clean=True), rel_energy_spread, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_xyz_offset():
+    
+    num_particles = 4344
+    energy_thres = 1200*SI.m_e*SI.c**2/SI.e  # [eV], 1200 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.8e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    x_offset = weighted_mean(xs, beam.weightings(), clean=False)
+    y_offset = weighted_mean(ys, beam.weightings(), clean=False)
+    z_offset = weighted_mean(zs, beam.weightings(), clean=False)
+    assert np.isclose(beam.x_offset(clean=False), x_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.y_offset(clean=False), y_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.z_offset(clean=False), z_offset, rtol=1e-15, atol=0.0)
+    x_offset = weighted_mean(xs, beam.weightings(), clean=True)
+    y_offset = weighted_mean(ys, beam.weightings(), clean=True)
+    z_offset = weighted_mean(zs, beam.weightings(), clean=True)
+    assert np.isclose(beam.x_offset(clean=True), x_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.y_offset(clean=True), y_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.z_offset(clean=True), z_offset, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_xpyp_offset():
+    
+    num_particles = 4344
+    energy_thres = 1200*SI.m_e*SI.c**2/SI.e  # [eV], 1200 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    xps = np.random.rand(num_particles)
+    yps = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.8e10
+    beam.set_phase_space(Q, xs, ys, zs, uzs=uzs, xps=xps, yps=yps)
+
+    x_angle = weighted_mean(xps, beam.weightings(), clean=False)
+    y_angle = weighted_mean(yps, beam.weightings(), clean=False)
+    assert np.isclose(beam.x_angle(clean=False), x_angle, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.y_angle(clean=False), y_angle, rtol=1e-15, atol=0.0)
+    x_angle = weighted_mean(xps, beam.weightings(), clean=True)
+    y_angle = weighted_mean(yps, beam.weightings(), clean=True)
+    assert np.isclose(beam.x_angle(clean=True), x_angle, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.y_angle(clean=True), y_angle, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_uxuyuz_offset():
+    
+    num_particles = 4344
+    energy_thres = 1200*SI.m_e*SI.c**2/SI.e  # [eV], 1200 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.8e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    ux_offset = weighted_mean(uxs, beam.weightings(), clean=False)
+    uy_offset = weighted_mean(uys, beam.weightings(), clean=False)
+    uz_offset = weighted_mean(uzs, beam.weightings(), clean=False)
+    assert np.isclose(beam.ux_offset(clean=False), ux_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.uy_offset(clean=False), uy_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.uz_offset(clean=False), uz_offset, rtol=1e-15, atol=0.0)
+    ux_offset = weighted_mean(uxs, beam.weightings(), clean=True)
+    uy_offset = weighted_mean(uys, beam.weightings(), clean=True)
+    uz_offset = weighted_mean(uzs, beam.weightings(), clean=True)
+    assert np.isclose(beam.ux_offset(clean=True), ux_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.uy_offset(clean=True), uy_offset, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.uz_offset(clean=True), uz_offset, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_geom_emittance_xy():
+    num_particles = 4444
+    energy_thres = 1400*SI.m_e*SI.c**2/SI.e  # [eV], 1400 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    xps = np.random.rand(num_particles)
+    yps = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.4e10
+    beam.set_phase_space(Q, xs, ys, zs, uzs=uzs, xps=xps, yps=yps)
+
+    geom_emittance_x = np.sqrt(np.linalg.det(weighted_cov(xs, xps, beam.weightings(), clean=False)))
+    geom_emittance_y = np.sqrt(np.linalg.det(weighted_cov(ys, yps, beam.weightings(), clean=False)))
+    assert np.isclose(beam.geom_emittance_x(clean=False), geom_emittance_x, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.geom_emittance_y(clean=False), geom_emittance_y, rtol=1e-15, atol=0.0)
+
+    geom_emittance_x = np.sqrt(np.linalg.det(weighted_cov(xs, xps, beam.weightings(), clean=True)))
+    geom_emittance_y = np.sqrt(np.linalg.det(weighted_cov(ys, yps, beam.weightings(), clean=True)))
+    assert np.isclose(beam.geom_emittance_x(clean=True), geom_emittance_x, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.geom_emittance_y(clean=True), geom_emittance_y, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_norm_emittance_xy():
+    num_particles = 4444
+    energy_thres = 1400*SI.m_e*SI.c**2/SI.e  # [eV], 1400 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.4e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    norm_emittance_x = np.sqrt(np.linalg.det(weighted_cov(xs, uxs/SI.c, beam.weightings(), clean=False)))
+    norm_emittance_y = np.sqrt(np.linalg.det(weighted_cov(ys, uys/SI.c, beam.weightings(), clean=False)))
+    assert np.isclose(beam.norm_emittance_x(clean=False), norm_emittance_x, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_y(clean=False), norm_emittance_y, rtol=1e-15, atol=0.0)
+
+    norm_emittance_x = np.sqrt(np.linalg.det(weighted_cov(xs, uxs/SI.c, beam.weightings(), clean=True)))
+    norm_emittance_y = np.sqrt(np.linalg.det(weighted_cov(ys, uys/SI.c, beam.weightings(), clean=True)))
+    assert np.isclose(beam.norm_emittance_x(clean=True), norm_emittance_x, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_y(clean=True), norm_emittance_y, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_beta_xy():
+    num_particles = 4444
+    energy_thres = 1400*SI.m_e*SI.c**2/SI.e  # [eV], 1400 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.4e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    covx = weighted_cov(xs, beam.xps(), beam.weightings(), clean=False)
+    beta_x = covx[0,0]/np.sqrt(np.linalg.det(covx))
+    covy = weighted_cov(ys, beam.yps(), beam.weightings(), clean=False)
+    beta_y = covy[0,0]/np.sqrt(np.linalg.det(covy))
+    assert np.isclose(beam.beta_x(clean=False), beta_x, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beta_y(clean=False), beta_y, rtol=1e-15, atol=0.0)
+
+    covx = weighted_cov(xs, beam.xps(), beam.weightings(), clean=True)
+    beta_x = covx[0,0]/np.sqrt(np.linalg.det(covx))
+    covy = weighted_cov(ys, beam.yps(), beam.weightings(), clean=True)
+    beta_y = covy[0,0]/np.sqrt(np.linalg.det(covy))
+    assert np.isclose(beam.beta_x(clean=True), beta_x, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beta_y(clean=True), beta_y, rtol=1e-15, atol=0.0)
+
+
 @pytest.mark.beam
 def test_gamma_xy():
     num_particles = 8989
@@ -776,25 +1115,104 @@ def test_gamma_xy():
     zs = np.random.rand(num_particles)
     uxs = np.random.rand(num_particles)
     uys = np.random.rand(num_particles)
-    uzs = random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e))
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e), size=num_particles)
     Q = -SI.e * 1.0e10
     beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
 
-    covx = weighted_cov(beam.xs(), beam.xps(), beam.weightings(), clean=False)
+    covx = weighted_cov(xs, beam.xps(), beam.weightings(), clean=False)
     gamma_x = covx[1,1]/np.sqrt(np.linalg.det(covx))
-    covy = weighted_cov(beam.ys(), beam.yps(), beam.weightings(), clean=False)
+    covy = weighted_cov(ys, beam.yps(), beam.weightings(), clean=False)
     gamma_y = covy[1,1]/np.sqrt(np.linalg.det(covy))
     assert np.isclose(beam.gamma_x(), gamma_x, rtol=1e-10, atol=0.0)
     assert np.isclose(beam.gamma_y(), gamma_y, rtol=1e-10, atol=0.0)
 
-    covx = weighted_cov(beam.xs(), beam.xps(), beam.weightings(), clean=True)
+    covx = weighted_cov(xs, beam.xps(), beam.weightings(), clean=True)
     gamma_x = covx[1,1]/np.sqrt(np.linalg.det(covx))
-    covy = weighted_cov(beam.ys(), beam.yps(), beam.weightings(), clean=True)
+    covy = weighted_cov(ys, beam.yps(), beam.weightings(), clean=True)
     gamma_y = covy[1,1]/np.sqrt(np.linalg.det(covy))
     assert np.isclose(beam.gamma_x(clean=True), gamma_x, rtol=1e-10, atol=0.0)
     assert np.isclose(beam.gamma_y(clean=True), gamma_y, rtol=1e-10, atol=0.0)
 
 
+@pytest.mark.beam
+def test_beam_size_xy():
+    
+    num_particles = 4444
+    energy_thres = 1400*SI.m_e*SI.c**2/SI.e  # [eV], 1400 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.4e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    beam_size_x = weighted_std(xs, beam.weightings(), clean=False)
+    beam_size_y = weighted_std(ys, beam.weightings(), clean=False)
+    assert np.isclose(beam.beam_size_x(clean=False), beam_size_x, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beam_size_y(clean=False), beam_size_y, rtol=1e-15, atol=0.0)
+
+    beam_size_x = weighted_std(xs, beam.weightings(), clean=True)
+    beam_size_y = weighted_std(ys, beam.weightings(), clean=True) 
+    assert np.isclose(beam.beam_size_x(clean=True), beam_size_x, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.beam_size_y(clean=True), beam_size_y, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_bunch_length():
+    num_particles = 4444
+    energy_thres = 1400*SI.m_e*SI.c**2/SI.e  # [eV], 1400 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e))
+    Q = -SI.e * 1.4e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    bunch_length = weighted_std(zs, beam.weightings(), clean=False)
+    assert np.isclose(beam.bunch_length(clean=False), bunch_length, rtol=1e-15, atol=0.0)
+
+    bunch_length = weighted_std(zs, beam.weightings(), clean=True)
+    assert np.isclose(beam.bunch_length(clean=True), bunch_length, rtol=1e-15, atol=0.0)
+
+
+@pytest.mark.beam
+def test_divergence_xy():
+    
+    num_particles = 4444
+    energy_thres = 1400*SI.m_e*SI.c**2/SI.e  # [eV], 1400 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    
+    beam = Beam()
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), size=num_particles)
+    Q = -SI.e * 1.4e10
+    beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
+
+    divergence_x = weighted_std(beam.xps(), beam.weightings(), clean=False)
+    divergence_y = weighted_std(beam.yps(), beam.weightings(), clean=False)
+    assert np.isclose(beam.divergence_x(clean=False), divergence_x, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.divergence_y(clean=False), divergence_y, rtol=1e-15, atol=0.0)
+
+    divergence_x = weighted_std(beam.xps(), beam.weightings(), clean=True)
+    divergence_y = weighted_std(beam.yps(), beam.weightings(), clean=True)
+    assert np.isclose(beam.divergence_x(clean=True), divergence_x, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.divergence_y(clean=True), divergence_y, rtol=1e-15, atol=0.0)
+
+
 
 ############# Tests for beam projections #############
 @pytest.mark.beam

From 59347277a6ec9c4062ef419ff6e478e18a3c96a3 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 19:04:33 +0200
Subject: [PATCH 38/51] Added exception triggering to
 test_rotate_coord_sys_3D() under test_beam.py.

---
 tests/test_beam.py | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index fad34041..91b9e598 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -1616,7 +1616,6 @@ def test_rotate_coord_sys_3D():
     start_y_angle = 7.040999e-09   # [rad], define an angle in the zy-plane.
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=start_x_angle, y_angle=start_y_angle)
 
-    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0)
     old_beam = source.track()
     old_beam.add_pointing_tilts()  # Add an initial tilt in the zy-plane.
 
@@ -1662,6 +1661,32 @@ def test_rotate_coord_sys_3D():
             assert np.isclose(beam.x_angle(), old_beam.x_angle(), rtol=1e-7, atol=1e-13)
             assert np.isclose(beam.y_angle(), old_beam.y_angle(), rtol=1e-7, atol=1e-13)
 
+    # Purpposedly trigger exceptions
+    with pytest.raises(ValueError):
+        beam = source.track()
+        axis1 = np.array([1, 1.1, 0])  # Axis as an non-unit vector.
+        axis2 = np.array([0, 1, 0])
+        axis3 = np.array([1, 0, 0])
+        beam.rotate_coord_sys_3D(axis1, angle1=nom_x_angle, axis2=axis2, angle2=nom_y_angle, axis3=axis3, angle3=0.0, invert=False)
+    with pytest.raises(ValueError):
+        beam = source.track()
+        axis1 = np.array([1, 0, 0])  # Axis as an non-unit vector.
+        axis2 = np.array([0, 0.45, 0])
+        axis3 = np.array([1, 0, 0])
+        beam.rotate_coord_sys_3D(axis1, angle1=nom_x_angle, axis2=axis2, angle2=nom_y_angle, axis3=axis3, angle3=0.0, invert=False)
+    with pytest.raises(ValueError):
+        beam = source.track()
+        axis1 = np.array([1, 0, 0])  # Axis as an non-unit vector.
+        axis2 = np.array([0, 1, 0])
+        axis3 = np.array([-0.1, 0, 0])
+        beam.rotate_coord_sys_3D(axis1, angle1=nom_x_angle, axis2=axis2, angle2=nom_y_angle, axis3=axis3, angle3=0.0, invert=False)
+    with pytest.raises(ValueError):
+        beam = source.track()
+        axis1 = np.array([1, 0, -3.14])  # Axis as an non-unit vector.
+        axis2 = np.array([0.3, 1, 0])
+        axis3 = np.array([-0.1, 0, 0])
+        beam.rotate_coord_sys_3D(axis1, angle1=nom_x_angle, axis2=axis2, angle2=nom_y_angle, axis3=axis3, angle3=0.0, invert=False)
+
 
 
 ############# Tests of in-house beam field calculation methods #############

From d5b33368f14bd71847e314feeacabac90c47609d Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 19:17:46 +0200
Subject: [PATCH 39/51] Added test_density_map_diags() and test_print_summary()
 to test_beam.py.

---
 tests/test_beam.py | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index 91b9e598..d0c15ce0 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -1827,6 +1827,8 @@ def test_scale_to_length():
 
 @pytest.mark.beam
 def test_scale_norm_emittance_x():
+
+    np.random.seed(42)
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
     
     beam = source.track()
@@ -1837,7 +1839,7 @@ def test_scale_norm_emittance_x():
     beam.scale_norm_emittance_x(scale_factor*beam.norm_emittance_x())
     assert np.allclose(beam.xs(), expected_xs, rtol=1e-15, atol=0.0)
     assert np.allclose(beam.uxs(), expected_uxs, rtol=1e-15, atol=0.0)
-    assert np.isclose(beam.norm_emittance_x(), expct_emit_nx, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_x(), expct_emit_nx, rtol=1e-10, atol=0.0)
 
     beam = source.track()
     scale_factor = 0.5354
@@ -1847,7 +1849,7 @@ def test_scale_norm_emittance_x():
     beam.scale_norm_emittance_x(scale_factor*beam.norm_emittance_x())
     assert np.allclose(beam.xs(), expected_xs, rtol=1e-15, atol=0.0)
     assert np.allclose(beam.uxs(), expected_uxs, rtol=1e-15, atol=0.0)
-    assert np.isclose(beam.norm_emittance_x(), expct_emit_nx, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_x(), expct_emit_nx, rtol=1e-10, atol=0.0)
 
     beam = source.track()
     scale_factor = 0.0
@@ -1867,6 +1869,8 @@ def test_scale_norm_emittance_x():
 
 @pytest.mark.beam
 def test_scale_norm_emittance_y():
+
+    np.random.seed(42)
     source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
     
     beam = source.track()
@@ -1887,7 +1891,7 @@ def test_scale_norm_emittance_y():
     beam.scale_norm_emittance_y(scale_factor*beam.norm_emittance_y())
     assert np.allclose(beam.ys(), expected_ys, rtol=1e-15, atol=0.0)
     assert np.allclose(beam.uys(), expected_uys, rtol=1e-15, atol=0.0)
-    assert np.isclose(beam.norm_emittance_y(), expct_emit_ny, rtol=1e-15, atol=0.0)
+    assert np.isclose(beam.norm_emittance_y(), expct_emit_ny, rtol=1e-13, atol=0.0)
 
     beam = source.track()
     scale_factor = 0.0
@@ -2319,6 +2323,23 @@ def test_plot_transverse_profile():
 #     beam.plot_bunch_pattern()
 
 
+@pytest.mark.beam
+def test_density_map_diags():
+
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    beam.density_map_diags()
+    plt.close()
+
+
+@pytest.mark.beam
+def test_print_summary():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.0, y_offset=0.0, x_angle=0.0, y_angle=0.0)
+    beam = source.track()
+    beam.beam_name = 'beam'
+    beam.print_summary()
+
+
 
 ############# Tests of saving and loading #############
 @pytest.mark.beam

From 97528904932bde4deec5016df48befadd7e4ba08 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 22:46:53 +0200
Subject: [PATCH 40/51] Small correction in test_beam.py

---
 tests/test_beam.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index d0c15ce0..4bf1e603 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -2359,5 +2359,5 @@ def test_save_load():
     loaded_beam = Beam.load(filename=filename)
 
     Beam.comp_beams(beam, loaded_beam, comp_location=True, rtol=1e-15, atol=0.0)
-    Beam.comp_beams(beam, loaded_beam, comp_location=True)
+    Beam.comp_beam_params(beam, loaded_beam, comp_location=True)
     shutil.rmtree(save_dir)
\ No newline at end of file

From 027c5980645c5b04922e3cc7b7f93a1eea560cf6 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Tue, 1 Apr 2025 22:47:40 +0200
Subject: [PATCH 41/51] Rearranged code in beam.py.

---
 abel/classes/beam.py | 323 +++++++++++++++++++++----------------------
 1 file changed, 161 insertions(+), 162 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index a7ce4b1e..15578240 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -1441,167 +1441,6 @@ def plot_transverse_profile(self):
     #     ax.set_ylabel('Beam current (kA)')
         
 
-  
-    ## SAVE AND LOAD BEAM
-    
-    def filename(self, runnable, beam_name):
-        return runnable.shot_path() + "/" + beam_name + "_" + str(self.trackable_number).zfill(3) + "_{:012.6F}".format(self.location) + ".h5"
-    
-    
-    # save beam (to OpenPMD format)
-    def save(self, runnable=None, filename=None, beam_name="beam", series=None):
-        
-        if len(self) == 0:
-            return
-        
-        # make new file if not provided
-        if series is None:
-
-            # open a new file
-            if runnable is not None:
-                filename = self.filename(runnable, beam_name)
-            
-            # open a new file
-            series = io.Series(filename, io.Access.create)
-            
-        
-            # add metadata
-            series.author = "ABEL (the Adaptable Beginning-to-End Linac simulation framework)"
-            series.date = datetime.now(timezone('CET')).strftime('%Y-%m-%d %H:%M:%S %z')
-
-        # make step (only one)
-        index = 0
-        iteration = series.iterations[index]
-        
-        # add attributes
-        iteration.set_attribute("time", self.location/SI.c)
-        for key, value in self.__dict__.items():
-            if not "__phasespace" in key:
-                iteration.set_attribute(key, value)
-                
-        # make beam record
-        particles = iteration.particles[beam_name]
-       
-        # generate datasets
-        dset_z = io.Dataset(self.zs().dtype, extent=self.zs().shape)
-        dset_x = io.Dataset(self.xs().dtype, extent=self.xs().shape)
-        dset_y = io.Dataset(self.ys().dtype, extent=self.ys().shape)
-        dset_zoff = io.Dataset(np.dtype('float64'), extent=[1])
-        dset_xoff = io.Dataset(np.dtype('float64'), extent=[1])
-        dset_yoff = io.Dataset(np.dtype('float64'), extent=[1])
-        dset_uz = io.Dataset(self.uzs().dtype, extent=self.uzs().shape)
-        dset_ux = io.Dataset(self.uxs().dtype, extent=self.uxs().shape)
-        dset_uy = io.Dataset(self.uys().dtype, extent=self.uys().shape)
-        dset_w = io.Dataset(self.weightings().dtype, extent=self.weightings().shape)
-        dset_id = io.Dataset(self.ids().dtype, extent=self.ids().shape)
-        dset_q = io.Dataset(np.dtype('float64'), extent=[1])
-        dset_m = io.Dataset(np.dtype('float64'), extent=[1])
-        
-        dset_n = io.Dataset(self.ids().dtype, extent=[1])
-        dset_f = io.Dataset(np.dtype('float64'), extent=[1])
-        
-        # prepare for writing
-        particles['position']['z'].reset_dataset(dset_z)
-        particles['position']['x'].reset_dataset(dset_x)
-        particles['position']['y'].reset_dataset(dset_y)
-        particles['positionOffset']['z'].reset_dataset(dset_zoff)
-        particles['positionOffset']['x'].reset_dataset(dset_xoff)
-        particles['positionOffset']['y'].reset_dataset(dset_yoff)
-        particles['momentum']['z'].reset_dataset(dset_uz)
-        particles['momentum']['x'].reset_dataset(dset_ux)
-        particles['momentum']['y'].reset_dataset(dset_uy)
-        particles['weighting'][io.Record_Component.SCALAR].reset_dataset(dset_w)
-        particles['id'][io.Record_Component.SCALAR].reset_dataset(dset_id)        
-        particles['charge'][io.Record_Component.SCALAR].reset_dataset(dset_q)
-        particles['mass'][io.Record_Component.SCALAR].reset_dataset(dset_m)
-        
-        # store data
-        particles['position']['z'].store_chunk(self.zs())
-        particles['position']['x'].store_chunk(self.xs())
-        particles['position']['y'].store_chunk(self.ys())
-        particles['positionOffset']['x'].make_constant(0.)
-        particles['positionOffset']['y'].make_constant(0.)
-        particles['positionOffset']['z'].make_constant(0.)
-        particles['momentum']['z'].store_chunk(self.uzs())
-        particles['momentum']['x'].store_chunk(self.uxs())
-        particles['momentum']['y'].store_chunk(self.uys())
-        particles['weighting'][io.Record_Component.SCALAR].store_chunk(self.weightings())
-        particles['id'][io.Record_Component.SCALAR].store_chunk(self.ids())
-        particles['charge'][io.Record_Component.SCALAR].make_constant(self.charge_sign()*SI.e)
-        particles['mass'][io.Record_Component.SCALAR].make_constant(SI.m_e)
-        
-        # set SI units (scaling factor)
-        particles['momentum']['z'].unit_SI = SI.m_e
-        particles['momentum']['x'].unit_SI = SI.m_e
-        particles['momentum']['y'].unit_SI = SI.m_e
-        
-        # set dimensional units
-        particles['position'].unit_dimension = {io.Unit_Dimension.L: 1}
-        particles['positionOffset'].unit_dimension = {io.Unit_Dimension.L: 1}
-        particles['momentum'].unit_dimension = {io.Unit_Dimension.L: 1, io.Unit_Dimension.M: 1, io.Unit_Dimension.T: -1}
-        particles['charge'].unit_dimension = {io.Unit_Dimension.T: 1, io.Unit_Dimension.I: 1}
-        particles['mass'].unit_dimension = {io.Unit_Dimension.M: 1}
-        
-        # save data to file
-        series.flush()
-        
-        return series
-        
-        
-    # load beam (from OpenPMD format)
-    @classmethod
-    def load(_, filename, beam_name='beam'):
-        
-        # load file
-        series = io.Series(filename, io.Access.read_only)
-        
-        # find index (use last one)
-        *_, index = series.iterations
-        
-        # get particle data
-        particles = series.iterations[index].particles[beam_name]
-        
-        # get attributes
-        charge = particles["charge"][io.Record_Component.SCALAR].get_attribute("value")
-        mass = particles["mass"][io.Record_Component.SCALAR].get_attribute("value")
-        
-        # extract phase space
-        ids = particles["id"][io.Record_Component.SCALAR].load_chunk()
-        weightings = particles["weighting"][io.Record_Component.SCALAR].load_chunk()
-        xs = particles['position']['x'].load_chunk()
-        ys = particles['position']['y'].load_chunk()
-        zs = particles['position']['z'].load_chunk()
-        pxs_unscaled = particles['momentum']['x'].load_chunk()
-        pys_unscaled = particles['momentum']['y'].load_chunk()
-        pzs_unscaled = particles['momentum']['z'].load_chunk()
-        series.flush()
-        
-        # apply SI scaling
-        pxs = pxs_unscaled * particles['momentum']['x'].unit_SI
-        pys = pys_unscaled * particles['momentum']['y'].unit_SI
-        pzs = pzs_unscaled * particles['momentum']['z'].unit_SI
-        
-        # make beam
-        beam = Beam()
-        beam.set_phase_space(Q=np.sum(weightings*charge), xs=xs, ys=ys, zs=zs, pxs=pxs, pys=pys, pzs=pzs, weightings=weightings)
-        
-        # add metadata to beam
-        try:
-            beam.trackable_number = series.iterations[index].get_attribute("trackable_number")
-            beam.stage_number = series.iterations[index].get_attribute("stage_number")
-            beam.location = series.iterations[index].get_attribute("location")
-            beam.num_bunches_in_train = series.iterations[index].get_attribute("num_bunches_in_train")
-            beam.bunch_separation = series.iterations[index].get_attribute("bunch_separation")
-        except:
-            beam.trackable_number = None
-            beam.stage_number = None
-            beam.location = None
-            beam.num_bunches_in_train = None
-            beam.bunch_separation = None
-        
-        return beam
-
-
     # ==================================================
     def imshow_plot(self, data, axes=None, extent=None, vmin=None, vmax=None, colmap='seismic', xlab=None, ylab=None, clab='', gridOn=False, origin='lower', interpolation=None, aspect='auto', log_cax=False, reduce_cax_pad=False):
         
@@ -1821,7 +1660,7 @@ def density_map_diags(self):
         self.distribution_plot_2D(arr1=zs, arr2=Es, weights=weights, hist_bins=hist_bins, hist_range=hist_range_energ, axes=axs[2][2], extent=extent_energ, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=energ_lab, clab=r'$\partial^2 N/\partial \xi \partial\mathcal{E}$ [$\mathrm{m}^{-1}$ $\mathrm{eV}^{-1}$]', origin='lower', interpolation='nearest')
 
 
-        # ==================================================
+    # ==================================================
     def print_summary(self):
 
         print('---------------------------------------------------')
@@ -1853,3 +1692,163 @@ def print_summary(self):
         print(f"rms beam length [um]:\t\t\t {self.bunch_length()*1e6 :.3f} \t\t")
         print(f"Peak current [kA]:\t\t\t {self.peak_current()/1e3 :.3f} \t\t")
         print('---------------------------------------------------')
+
+
+         ## SAVE AND LOAD BEAM
+    
+    def filename(self, runnable, beam_name):
+        return runnable.shot_path() + "/" + beam_name + "_" + str(self.trackable_number).zfill(3) + "_{:012.6F}".format(self.location) + ".h5"
+    
+    
+    # save beam (to OpenPMD format)
+    def save(self, runnable=None, filename=None, beam_name="beam", series=None):
+        
+        if len(self) == 0:
+            return
+        
+        # make new file if not provided
+        if series is None:
+
+            # open a new file
+            if runnable is not None:
+                filename = self.filename(runnable, beam_name)
+            
+            # open a new file
+            series = io.Series(filename, io.Access.create)
+            
+        
+            # add metadata
+            series.author = "ABEL (the Adaptable Beginning-to-End Linac simulation framework)"
+            series.date = datetime.now(timezone('CET')).strftime('%Y-%m-%d %H:%M:%S %z')
+
+        # make step (only one)
+        index = 0
+        iteration = series.iterations[index]
+        
+        # add attributes
+        iteration.set_attribute("time", self.location/SI.c)
+        for key, value in self.__dict__.items():
+            if not "__phasespace" in key:
+                iteration.set_attribute(key, value)
+                
+        # make beam record
+        particles = iteration.particles[beam_name]
+       
+        # generate datasets
+        dset_z = io.Dataset(self.zs().dtype, extent=self.zs().shape)
+        dset_x = io.Dataset(self.xs().dtype, extent=self.xs().shape)
+        dset_y = io.Dataset(self.ys().dtype, extent=self.ys().shape)
+        dset_zoff = io.Dataset(np.dtype('float64'), extent=[1])
+        dset_xoff = io.Dataset(np.dtype('float64'), extent=[1])
+        dset_yoff = io.Dataset(np.dtype('float64'), extent=[1])
+        dset_uz = io.Dataset(self.uzs().dtype, extent=self.uzs().shape)
+        dset_ux = io.Dataset(self.uxs().dtype, extent=self.uxs().shape)
+        dset_uy = io.Dataset(self.uys().dtype, extent=self.uys().shape)
+        dset_w = io.Dataset(self.weightings().dtype, extent=self.weightings().shape)
+        dset_id = io.Dataset(self.ids().dtype, extent=self.ids().shape)
+        dset_q = io.Dataset(np.dtype('float64'), extent=[1])
+        dset_m = io.Dataset(np.dtype('float64'), extent=[1])
+        
+        dset_n = io.Dataset(self.ids().dtype, extent=[1])
+        dset_f = io.Dataset(np.dtype('float64'), extent=[1])
+        
+        # prepare for writing
+        particles['position']['z'].reset_dataset(dset_z)
+        particles['position']['x'].reset_dataset(dset_x)
+        particles['position']['y'].reset_dataset(dset_y)
+        particles['positionOffset']['z'].reset_dataset(dset_zoff)
+        particles['positionOffset']['x'].reset_dataset(dset_xoff)
+        particles['positionOffset']['y'].reset_dataset(dset_yoff)
+        particles['momentum']['z'].reset_dataset(dset_uz)
+        particles['momentum']['x'].reset_dataset(dset_ux)
+        particles['momentum']['y'].reset_dataset(dset_uy)
+        particles['weighting'][io.Record_Component.SCALAR].reset_dataset(dset_w)
+        particles['id'][io.Record_Component.SCALAR].reset_dataset(dset_id)        
+        particles['charge'][io.Record_Component.SCALAR].reset_dataset(dset_q)
+        particles['mass'][io.Record_Component.SCALAR].reset_dataset(dset_m)
+        
+        # store data
+        particles['position']['z'].store_chunk(self.zs())
+        particles['position']['x'].store_chunk(self.xs())
+        particles['position']['y'].store_chunk(self.ys())
+        particles['positionOffset']['x'].make_constant(0.)
+        particles['positionOffset']['y'].make_constant(0.)
+        particles['positionOffset']['z'].make_constant(0.)
+        particles['momentum']['z'].store_chunk(self.uzs())
+        particles['momentum']['x'].store_chunk(self.uxs())
+        particles['momentum']['y'].store_chunk(self.uys())
+        particles['weighting'][io.Record_Component.SCALAR].store_chunk(self.weightings())
+        particles['id'][io.Record_Component.SCALAR].store_chunk(self.ids())
+        particles['charge'][io.Record_Component.SCALAR].make_constant(self.charge_sign()*SI.e)
+        particles['mass'][io.Record_Component.SCALAR].make_constant(SI.m_e)
+        
+        # set SI units (scaling factor)
+        particles['momentum']['z'].unit_SI = SI.m_e
+        particles['momentum']['x'].unit_SI = SI.m_e
+        particles['momentum']['y'].unit_SI = SI.m_e
+        
+        # set dimensional units
+        particles['position'].unit_dimension = {io.Unit_Dimension.L: 1}
+        particles['positionOffset'].unit_dimension = {io.Unit_Dimension.L: 1}
+        particles['momentum'].unit_dimension = {io.Unit_Dimension.L: 1, io.Unit_Dimension.M: 1, io.Unit_Dimension.T: -1}
+        particles['charge'].unit_dimension = {io.Unit_Dimension.T: 1, io.Unit_Dimension.I: 1}
+        particles['mass'].unit_dimension = {io.Unit_Dimension.M: 1}
+        
+        # save data to file
+        series.flush()
+        
+        return series
+        
+        
+    # load beam (from OpenPMD format)
+    @classmethod
+    def load(_, filename, beam_name='beam'):
+        
+        # load file
+        series = io.Series(filename, io.Access.read_only)
+        
+        # find index (use last one)
+        *_, index = series.iterations
+        
+        # get particle data
+        particles = series.iterations[index].particles[beam_name]
+        
+        # get attributes
+        charge = particles["charge"][io.Record_Component.SCALAR].get_attribute("value")
+        mass = particles["mass"][io.Record_Component.SCALAR].get_attribute("value")
+        
+        # extract phase space
+        ids = particles["id"][io.Record_Component.SCALAR].load_chunk()
+        weightings = particles["weighting"][io.Record_Component.SCALAR].load_chunk()
+        xs = particles['position']['x'].load_chunk()
+        ys = particles['position']['y'].load_chunk()
+        zs = particles['position']['z'].load_chunk()
+        pxs_unscaled = particles['momentum']['x'].load_chunk()
+        pys_unscaled = particles['momentum']['y'].load_chunk()
+        pzs_unscaled = particles['momentum']['z'].load_chunk()
+        series.flush()
+        
+        # apply SI scaling
+        pxs = pxs_unscaled * particles['momentum']['x'].unit_SI
+        pys = pys_unscaled * particles['momentum']['y'].unit_SI
+        pzs = pzs_unscaled * particles['momentum']['z'].unit_SI
+        
+        # make beam
+        beam = Beam()
+        beam.set_phase_space(Q=np.sum(weightings*charge), xs=xs, ys=ys, zs=zs, pxs=pxs, pys=pys, pzs=pzs, weightings=weightings)
+        
+        # add metadata to beam
+        try:
+            beam.trackable_number = series.iterations[index].get_attribute("trackable_number")
+            beam.stage_number = series.iterations[index].get_attribute("stage_number")
+            beam.location = series.iterations[index].get_attribute("location")
+            beam.num_bunches_in_train = series.iterations[index].get_attribute("num_bunches_in_train")
+            beam.bunch_separation = series.iterations[index].get_attribute("bunch_separation")
+        except:
+            beam.trackable_number = None
+            beam.stage_number = None
+            beam.location = None
+            beam.num_bunches_in_train = None
+            beam.bunch_separation = None
+        
+        return beam
\ No newline at end of file

From 6d69d5aff00ae2d68a6473ac244c96174c51550d Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Wed, 2 Apr 2025 13:14:57 +0200
Subject: [PATCH 42/51] Added checks for input types in Beam.set_phase_space()
 and implemented tests for triggering the exceptions.

---
 abel/classes/beam.py | 120 ++++++++++++++++++++++++++++++++++++-
 tests/test_beam.py   | 137 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 255 insertions(+), 2 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index 15578240..e6818ccd 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -64,9 +64,127 @@ def remove_nans(self):
         
     # set phase space
     def set_phase_space(self, Q, xs, ys, zs, uxs=None, uys=None, uzs=None, pxs=None, pys=None, pzs=None, xps=None, yps=None, Es=None, weightings=None, particle_mass=SI.m_e):
+        """
+        Set the phase space of the beam. All input arrays must have the same lengths.
+
+        Parameters
+        ----------
+        Q : [C], float
+            Total beam charge.
+
+        xs, ys, zs: [m], 1D ndarray
+            Coordinates for the macroparticles.
+            
+        uxs, uys, uzs: [m/s], 1D ndarray, optional
+            Proper velocities for the macroparticles. All uzs values must be above 10*particle rest energy/c/particle mass. Default set to ``None``.
+            
+        pxs, pys, pzs: [kg m/s], 1D ndarray, optional
+            Momenta for the macroparticles. All pzs values must be above 10*particle rest energy/c. Default set to ``None``.
+        
+        xps, yps: [rad], 1D ndarray, optional
+            Angles dx/ds and dy/ds for the macroparticles. Default set to ``None``.
+
+        Es : [eV], 1D ndarray, optional
+            Energies for the macroparticles. All values must be above 10*particle rest energy. Default set to ``None``.
+
+        weightings : 1D ndarray, optional
+            Weights for the macroparticles. Default set to ``None``.
+
+        particle_mass : [kg], float, optional
+            Particle mass for a single real particle. Default set to ``SI.m_e``.
+        
+            
+        Returns
+        ----------
+        ``None``
+        """
+
+        # Check coordinate type and length
+        if not isinstance(xs, np.ndarray)  or not isinstance(ys, np.ndarray) or not isinstance(zs, np.ndarray):
+            raise TypeError('Incompatible input type.')
         
-        # make empty phase space
         num_particles = len(xs)
+        if len(ys) != num_particles or len(zs) != num_particles:
+            raise ValueError('The input arrays must have the same lengths.')
+        
+        # Check proper velocity type and length
+        if uxs is not None:
+            if not isinstance(uxs, np.ndarray):
+                raise TypeError('Incompatible input type.')
+            if len(uxs) != num_particles:
+                raise ValueError('The input arrays must have the same lengths.')
+        if uys is not None:
+            if not isinstance(uys, np.ndarray):
+                raise TypeError('Incompatible input type.')
+            if len(uys) != num_particles:
+                raise ValueError('The input arrays must have the same lengths.')
+        if uzs is not None:
+            if not isinstance(uzs, np.ndarray):
+                raise TypeError('Incompatible input type.')
+            if len(uzs) != num_particles:
+                raise ValueError('The input arrays must have the same lengths.')
+            
+        # Check momentum type and length
+        if pxs is not None:
+            if not isinstance(pxs, np.ndarray):
+                raise TypeError('Incompatible input type.')
+            if len(pxs) != num_particles:
+                raise ValueError('The input arrays must have the same lengths.')
+        if pys is not None:
+            if not isinstance(pys, np.ndarray):
+                raise TypeError('Incompatible input type.')
+            if len(pys) != num_particles:
+                raise ValueError('The input arrays must have the same lengths.')
+        if pzs is not None:
+            if not isinstance(pzs, np.ndarray):
+                raise TypeError('Incompatible input type.')
+            if len(pzs) != num_particles:
+                raise ValueError('The input arrays must have the same lengths.')
+            
+        # Check angle type and length
+        if xps is not None:
+            if not isinstance(xps, np.ndarray):
+                raise TypeError('Incompatible input type.')
+            if len(xps) != num_particles:
+                raise ValueError('The input arrays must have the same lengths.')
+        if yps is not None:
+            if not isinstance(yps, np.ndarray):
+                raise TypeError('Incompatible input type.')
+            if len(yps) != num_particles:
+                raise ValueError('The input arrays must have the same lengths.')
+            
+        # Check energy type and length
+        if Es is not None:
+            if not isinstance(Es, np.ndarray):
+                raise TypeError('Incompatible input type.')
+            if len(Es) != num_particles:
+                raise ValueError('The input arrays must have the same lengths.')
+        
+        # Prevent defining proper velocity, momentum or angle in the same direction
+        if uxs is not None and pxs is not None:
+            raise ValueError('Cannot define proper velocity and momentum in the same direction.')
+        if uxs is not None and xps is not None:
+            raise ValueError('Cannot define proper velocity and angle in the same direction.')
+        if pxs is not None and xps is not None:
+            raise ValueError('Cannot define momentum and angle in the same direction.')
+        if uys is not None and pys is not None:
+            raise ValueError('Cannot define proper velocity and momentum in the same direction.')
+        if uys is not None and yps is not None:
+            raise ValueError('Cannot define proper velocity and angle in the same direction.')
+        if pys is not None and yps is not None:
+            raise ValueError('Cannot define momentum and angle in the same direction.')
+        if uzs is not None and pzs is not None:
+            raise ValueError('Cannot define proper velocity and momentum in the same direction.')
+        if uzs is not None and Es is not None:
+            raise ValueError('Cannot set both uzs and Es.')
+        if pzs is not None and Es is not None:
+            raise ValueError('Cannot set both pzs and Es.')
+        
+        if particle_mass is not None and particle_mass < 0:
+            raise ValueError('Particle mass cannot be negative.')
+
+
+        # make empty phase space
         self.reset_phase_space(num_particles)
         
         # add positions
diff --git a/tests/test_beam.py b/tests/test_beam.py
index 4bf1e603..b650a2e3 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -278,6 +278,141 @@ def test_set_phase_space4():
     assert np.allclose(beam.Es(), np.sqrt((pzs*SI.c)**2 + (SI.m_e*SI.c**2)**2)/SI.e, rtol=1e-15, atol=0.0)
 
 
+@pytest.mark.beam
+def test_set_phase_space5():
+    "Verify that the phase space is set correctly."
+    
+    num_particles = 10042
+    Q = -SI.e * 1.0e10
+    xs = np.random.rand(num_particles)
+    ys = np.random.rand(num_particles)
+    zs = np.random.rand(num_particles)
+    uxs = np.random.rand(num_particles)
+    uys = np.random.rand(num_particles)
+
+    energy_thres = 13*SI.m_e*SI.c**2/SI.e  # [eV], 13 * particle rest energy.
+    uz_thres = energy2proper_velocity(energy_thres, unit='eV', m=SI.m_e)
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(10e12, unit='eV', m=SI.m_e), size=num_particles)
+    pz_thres = gamma2momentum(energy2gamma(energy_thres, unit='eV', m=SI.m_e))
+    pzs = np.random.normal(pz_thres, 0.02*pz_thres, num_particles)
+    Es = np.random.normal(energy_thres, 0.02*energy_thres, num_particles)
+
+
+    ## Purposedly trigger exceptions
+    # Coordinates type and length
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, 3.14, ys, 1.0)
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, 1.0)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, np.random.rand(num_particles+1), ys, zs)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, np.random.rand(num_particles-1), zs)
+
+    # Momenta type and length
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uxs=2.9)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, np.random.rand(num_particles+1), uys, uzs)
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uxs, 15.16, uzs)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uxs, np.random.rand(num_particles-1), uzs)
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uz_thres)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uxs, uys, np.append(uzs, uz_thres))
+
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, pxs=2.9)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, pxs=np.random.rand(num_particles+1), pys=uys, pzs=pzs)
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, pxs=uxs, pys=15.16, pzs=pzs)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, pxs=uxs, pys=np.random.rand(num_particles-1), pzs=pzs)
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, pxs=uxs, pys=uys, pzs=1.6e-18)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, pxs=uxs, pys=uys, pzs=np.append(pzs, 1.6e-18))
+
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, xps=2.9)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, xps=np.random.rand(num_particles+1))
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, xps=uxs, yps=15.16)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, xps=uxs, yps=np.random.rand(num_particles-1), Es=Es)
+    with pytest.raises(TypeError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, xps=uxs, yps=uys, Es=100e9)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, xps=uxs, yps=uys, Es=np.append(Es, 10e9))
+
+    # Define proper velocity, momentum or angle in the same direction
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uxs=uxs, pxs=uxs)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uxs=uxs, xps=uxs)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, pxs=uxs, xps=uxs)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uys=uys, pys=uys)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uys=uys, yps=uys)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, pys=uys, yps=uys)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uzs=uys, pzs=pzs)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uzs=uzs, Es=Es)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, pzs=pzs, Es=Es)
+    
+    # Set energies below the accepted threshold
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs=np.random.rand(num_particles))
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, pzs=pzs*0.01)
+    with pytest.raises(ValueError):
+        beam = Beam()
+        beam.set_phase_space(Q, xs, ys, zs, Es=np.random.rand(num_particles))
+    
+
+
 @pytest.mark.beam
 def test_reset_phase_space():
     "Test reset_phase_space to ensure it initializes an 8xN zero matrix for the specified number of particles."
@@ -1174,7 +1309,7 @@ def test_bunch_length():
     zs = np.random.rand(num_particles)
     uxs = np.random.rand(num_particles)
     uys = np.random.rand(num_particles)
-    uzs = random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e))
+    uzs = np.random.uniform(uz_thres, energy2proper_velocity(energy_thres*1.01, unit='eV', m=SI.m_e), num_particles)
     Q = -SI.e * 1.4e10
     beam.set_phase_space(Q, xs, ys, zs, uxs, uys, uzs)
 

From 4ac59544499a5b48037fd4665a8754aef5e30fb6 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Wed, 2 Apr 2025 15:46:54 +0200
Subject: [PATCH 43/51] Added test_remove_halo_particles() to test_beam.py.

---
 tests/test_beam.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tests/test_beam.py b/tests/test_beam.py
index b650a2e3..d29e0a6c 100644
--- a/tests/test_beam.py
+++ b/tests/test_beam.py
@@ -643,6 +643,30 @@ def test_scale_energy():
         assert str(err) == 'Es contains values that are too small.'
 
 
+@pytest.mark.beam
+def test_remove_halo_particles():
+    source = setup_basic_source(plasma_density=6.0e20, ramp_beta_mag=5.0, energy=3e9, z_offset=0.0, x_offset=0.1516e-6, y_offset=-5.354e-6, x_angle=-1.3e-6, y_angle=0.7e-6)
+    beam = source.track()
+    initial_beam = copy.deepcopy(beam)
+
+    nsigma = 3
+    xfilter = np.abs(beam.xs()-beam.x_offset(clean=True)) > nsigma*beam.beam_size_x(clean=True)
+    xpfilter = np.abs(beam.xps()-beam.x_angle(clean=True)) > nsigma*beam.divergence_x(clean=True)
+    yfilter = np.abs(beam.ys()-beam.y_offset(clean=True)) > nsigma*beam.beam_size_y(clean=True)
+    ypfilter = np.abs(beam.yps()-beam.y_angle(clean=True)) > nsigma*beam.divergence_y(clean=True)
+    filter = np.logical_or(np.logical_or(xfilter, xpfilter), np.logical_or(yfilter, ypfilter))
+    count = np.count_nonzero(filter)  # Number of particles to be filtered away.
+    num_particles_left = len(beam) - count
+
+    beam.remove_halo_particles(nsigma=3)
+
+    assert len(beam) == num_particles_left
+    assert np.all( np.abs(beam.xs()-beam.x_offset(clean=True)) < nsigma*initial_beam.beam_size_x(clean=True) )
+    assert np.all( np.abs(beam.xps()-beam.x_angle(clean=True)) < nsigma*initial_beam.divergence_x(clean=True) )
+    assert np.all( np.abs(beam.ys()-beam.y_offset(clean=True)) < nsigma*initial_beam.beam_size_y(clean=True) )
+    assert np.all( np.abs(beam.yps()-beam.y_angle(clean=True)) < nsigma*initial_beam.divergence_y(clean=True) )
+
+
 @pytest.mark.beam
 def test_rs():
     beam = Beam()

From 3a246fa8405387ef5eda73d83b06ebafe4b28a57 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Thu, 3 Apr 2025 12:23:53 +0200
Subject: [PATCH 44/51] Added doc string to Beam.apply_betatron_motion().

---
 abel/classes/beam.py | 36 ++++++++++++++++++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index e6818ccd..7e996311 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -1399,7 +1399,39 @@ def flip_transverse_phase_spaces(self, flip_momenta=True, flip_positions=False):
             self.set_ys(-self.ys())
 
         
-    def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiation_reaction=False, calc_evolution=False, evolution_samples=None):
+    def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiation_reaction=False, calc_evolution=False):
+        """
+        Evolve the beam by solving Hill's equation.
+
+        Parameters
+        ----------
+        L : [m] float
+            The beam propagation distance.
+
+        n0 : [m^-3] float
+            Plasme number density.
+
+        deltaEs : [m^-3] 1D float ndarray
+            Energy change for the macroparticles.
+
+        x0_driver, y0_driver : [m] float, optional
+            Initial transverse offsets of the drive beam. Defaults set to 0.
+        
+        radiation_reaction : bool
+            Flag for enabling ating radiation reaction.
+        
+        calc_evolution : bool
+            Flag for recording the beam parameter evolution. 
+
+
+        Returns
+        ----------
+        Es_final : [eV] 1D float ndarray
+            The final energies for all macroparticles.
+        
+        evol : SimpleNamespace object
+            only returns when ``calc_evolution=True``. Contains beam parameter evolution data.
+        """
         
         # remove particles with subzero and Nan energy
         neg_indices = self.Es() < 0
@@ -1418,7 +1450,7 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat
         gammas = energy2gamma(Es_final)
         dgamma_ds = (gammas-gamma0s)/L
         
-        if calc_evolution:
+        if calc_evolution:  # TODO: This seems to be very clumsy. Consider re-writing.
                 
             # calculate evolution
             num_evol_steps = max(20, min(400, round(2*L/(beta_matched(n0, self.energy()+min(0,np.mean(deltaEs)))))))

From 40c5fe1703a5e750566f5ff0d714b6d33120cecc Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Thu, 3 Apr 2025 12:24:35 +0200
Subject: [PATCH 45/51] Added minor comment to stage_basic.py.

---
 abel/classes/stage/impl/stage_basic.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/abel/classes/stage/impl/stage_basic.py b/abel/classes/stage/impl/stage_basic.py
index f44b3dec..c5e29b8f 100644
--- a/abel/classes/stage/impl/stage_basic.py
+++ b/abel/classes/stage/impl/stage_basic.py
@@ -86,7 +86,7 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
                 
 
         # ========== Betatron oscillations ==========
-        deltaEs = np.full(len(beam.Es()), self.nom_energy_gain_flattop)
+        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

From 9bb7a736886bdef648118ad3eef36761a34d4a8b Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Thu, 3 Apr 2025 13:06:00 +0200
Subject: [PATCH 46/51] Added controls in source_from_file.py to ensure valid
 file path in both the constructor and file setter.

---
 abel/classes/source/impl/source_from_file.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/abel/classes/source/impl/source_from_file.py b/abel/classes/source/impl/source_from_file.py
index a491825e..1b49d3bc 100644
--- a/abel/classes/source/impl/source_from_file.py
+++ b/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):

From b903b8213a7256faf242489ab8c351ac4d3d4ab5 Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Thu, 3 Apr 2025 13:07:13 +0200
Subject: [PATCH 47/51] Edited test_SourceFromFile2Beam() in test_sources.py to
 trigger exceptions for invalid file path.

---
 tests/test_sources.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tests/test_sources.py b/tests/test_sources.py
index 6100e71d..77de8df7 100644
--- a/tests/test_sources.py
+++ b/tests/test_sources.py
@@ -191,4 +191,13 @@ def test_SourceFromFile2Beam():
     ref_beam = Beam.load(beam_file)
     Beam.comp_beams(beam, ref_beam)
 
+    # Trigger exception for when a file does not exist
+    with pytest.raises(FileNotFoundError):
+        file = 'tests/data/test_StagePrtclTransWakeInstability_beamline/test_baseline_linac/blabla.h5'
+        source = SourceFromFile(file=file)
+    with pytest.raises(FileNotFoundError):
+        file = 'tests/data/test_StagePrtclTransWakeInstability_beamline/test_baseline_linac/shot_000/blabla.h5'
+        source = SourceFromFile()
+        source.file = file
+
 

From 04c546222e16a72fcb25919a1da8ea1b3c9006a5 Mon Sep 17 00:00:00 2001
From: Kyrre Ness Sjobak <k.n.sjobak@fys.uio.no>
Date: Thu, 3 Apr 2025 15:41:42 +0200
Subject: [PATCH 48/51] Revert "Rearranged code in beam.py."

This reverts commit 027c5980645c5b04922e3cc7b7f93a1eea560cf6.
That commit moved functions around in Beam : filename, save, load
---
 abel/classes/beam.py | 323 ++++++++++++++++++++++---------------------
 1 file changed, 162 insertions(+), 161 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index 7e996311..2996ad47 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -1591,6 +1591,167 @@ def plot_transverse_profile(self):
     #     ax.set_ylabel('Beam current (kA)')
         
 
+  
+    ## SAVE AND LOAD BEAM
+    
+    def filename(self, runnable, beam_name):
+        return runnable.shot_path() + "/" + beam_name + "_" + str(self.trackable_number).zfill(3) + "_{:012.6F}".format(self.location) + ".h5"
+    
+    
+    # save beam (to OpenPMD format)
+    def save(self, runnable=None, filename=None, beam_name="beam", series=None):
+        
+        if len(self) == 0:
+            return
+        
+        # make new file if not provided
+        if series is None:
+
+            # open a new file
+            if runnable is not None:
+                filename = self.filename(runnable, beam_name)
+            
+            # open a new file
+            series = io.Series(filename, io.Access.create)
+            
+        
+            # add metadata
+            series.author = "ABEL (the Adaptable Beginning-to-End Linac simulation framework)"
+            series.date = datetime.now(timezone('CET')).strftime('%Y-%m-%d %H:%M:%S %z')
+
+        # make step (only one)
+        index = 0
+        iteration = series.iterations[index]
+        
+        # add attributes
+        iteration.set_attribute("time", self.location/SI.c)
+        for key, value in self.__dict__.items():
+            if not "__phasespace" in key:
+                iteration.set_attribute(key, value)
+                
+        # make beam record
+        particles = iteration.particles[beam_name]
+       
+        # generate datasets
+        dset_z = io.Dataset(self.zs().dtype, extent=self.zs().shape)
+        dset_x = io.Dataset(self.xs().dtype, extent=self.xs().shape)
+        dset_y = io.Dataset(self.ys().dtype, extent=self.ys().shape)
+        dset_zoff = io.Dataset(np.dtype('float64'), extent=[1])
+        dset_xoff = io.Dataset(np.dtype('float64'), extent=[1])
+        dset_yoff = io.Dataset(np.dtype('float64'), extent=[1])
+        dset_uz = io.Dataset(self.uzs().dtype, extent=self.uzs().shape)
+        dset_ux = io.Dataset(self.uxs().dtype, extent=self.uxs().shape)
+        dset_uy = io.Dataset(self.uys().dtype, extent=self.uys().shape)
+        dset_w = io.Dataset(self.weightings().dtype, extent=self.weightings().shape)
+        dset_id = io.Dataset(self.ids().dtype, extent=self.ids().shape)
+        dset_q = io.Dataset(np.dtype('float64'), extent=[1])
+        dset_m = io.Dataset(np.dtype('float64'), extent=[1])
+        
+        dset_n = io.Dataset(self.ids().dtype, extent=[1])
+        dset_f = io.Dataset(np.dtype('float64'), extent=[1])
+        
+        # prepare for writing
+        particles['position']['z'].reset_dataset(dset_z)
+        particles['position']['x'].reset_dataset(dset_x)
+        particles['position']['y'].reset_dataset(dset_y)
+        particles['positionOffset']['z'].reset_dataset(dset_zoff)
+        particles['positionOffset']['x'].reset_dataset(dset_xoff)
+        particles['positionOffset']['y'].reset_dataset(dset_yoff)
+        particles['momentum']['z'].reset_dataset(dset_uz)
+        particles['momentum']['x'].reset_dataset(dset_ux)
+        particles['momentum']['y'].reset_dataset(dset_uy)
+        particles['weighting'][io.Record_Component.SCALAR].reset_dataset(dset_w)
+        particles['id'][io.Record_Component.SCALAR].reset_dataset(dset_id)        
+        particles['charge'][io.Record_Component.SCALAR].reset_dataset(dset_q)
+        particles['mass'][io.Record_Component.SCALAR].reset_dataset(dset_m)
+        
+        # store data
+        particles['position']['z'].store_chunk(self.zs())
+        particles['position']['x'].store_chunk(self.xs())
+        particles['position']['y'].store_chunk(self.ys())
+        particles['positionOffset']['x'].make_constant(0.)
+        particles['positionOffset']['y'].make_constant(0.)
+        particles['positionOffset']['z'].make_constant(0.)
+        particles['momentum']['z'].store_chunk(self.uzs())
+        particles['momentum']['x'].store_chunk(self.uxs())
+        particles['momentum']['y'].store_chunk(self.uys())
+        particles['weighting'][io.Record_Component.SCALAR].store_chunk(self.weightings())
+        particles['id'][io.Record_Component.SCALAR].store_chunk(self.ids())
+        particles['charge'][io.Record_Component.SCALAR].make_constant(self.charge_sign()*SI.e)
+        particles['mass'][io.Record_Component.SCALAR].make_constant(SI.m_e)
+        
+        # set SI units (scaling factor)
+        particles['momentum']['z'].unit_SI = SI.m_e
+        particles['momentum']['x'].unit_SI = SI.m_e
+        particles['momentum']['y'].unit_SI = SI.m_e
+        
+        # set dimensional units
+        particles['position'].unit_dimension = {io.Unit_Dimension.L: 1}
+        particles['positionOffset'].unit_dimension = {io.Unit_Dimension.L: 1}
+        particles['momentum'].unit_dimension = {io.Unit_Dimension.L: 1, io.Unit_Dimension.M: 1, io.Unit_Dimension.T: -1}
+        particles['charge'].unit_dimension = {io.Unit_Dimension.T: 1, io.Unit_Dimension.I: 1}
+        particles['mass'].unit_dimension = {io.Unit_Dimension.M: 1}
+        
+        # save data to file
+        series.flush()
+        
+        return series
+        
+        
+    # load beam (from OpenPMD format)
+    @classmethod
+    def load(_, filename, beam_name='beam'):
+        
+        # load file
+        series = io.Series(filename, io.Access.read_only)
+        
+        # find index (use last one)
+        *_, index = series.iterations
+        
+        # get particle data
+        particles = series.iterations[index].particles[beam_name]
+        
+        # get attributes
+        charge = particles["charge"][io.Record_Component.SCALAR].get_attribute("value")
+        mass = particles["mass"][io.Record_Component.SCALAR].get_attribute("value")
+        
+        # extract phase space
+        ids = particles["id"][io.Record_Component.SCALAR].load_chunk()
+        weightings = particles["weighting"][io.Record_Component.SCALAR].load_chunk()
+        xs = particles['position']['x'].load_chunk()
+        ys = particles['position']['y'].load_chunk()
+        zs = particles['position']['z'].load_chunk()
+        pxs_unscaled = particles['momentum']['x'].load_chunk()
+        pys_unscaled = particles['momentum']['y'].load_chunk()
+        pzs_unscaled = particles['momentum']['z'].load_chunk()
+        series.flush()
+        
+        # apply SI scaling
+        pxs = pxs_unscaled * particles['momentum']['x'].unit_SI
+        pys = pys_unscaled * particles['momentum']['y'].unit_SI
+        pzs = pzs_unscaled * particles['momentum']['z'].unit_SI
+        
+        # make beam
+        beam = Beam()
+        beam.set_phase_space(Q=np.sum(weightings*charge), xs=xs, ys=ys, zs=zs, pxs=pxs, pys=pys, pzs=pzs, weightings=weightings)
+        
+        # add metadata to beam
+        try:
+            beam.trackable_number = series.iterations[index].get_attribute("trackable_number")
+            beam.stage_number = series.iterations[index].get_attribute("stage_number")
+            beam.location = series.iterations[index].get_attribute("location")
+            beam.num_bunches_in_train = series.iterations[index].get_attribute("num_bunches_in_train")
+            beam.bunch_separation = series.iterations[index].get_attribute("bunch_separation")
+        except:
+            beam.trackable_number = None
+            beam.stage_number = None
+            beam.location = None
+            beam.num_bunches_in_train = None
+            beam.bunch_separation = None
+        
+        return beam
+
+
     # ==================================================
     def imshow_plot(self, data, axes=None, extent=None, vmin=None, vmax=None, colmap='seismic', xlab=None, ylab=None, clab='', gridOn=False, origin='lower', interpolation=None, aspect='auto', log_cax=False, reduce_cax_pad=False):
         
@@ -1810,7 +1971,7 @@ def density_map_diags(self):
         self.distribution_plot_2D(arr1=zs, arr2=Es, weights=weights, hist_bins=hist_bins, hist_range=hist_range_energ, axes=axs[2][2], extent=extent_energ, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=energ_lab, clab=r'$\partial^2 N/\partial \xi \partial\mathcal{E}$ [$\mathrm{m}^{-1}$ $\mathrm{eV}^{-1}$]', origin='lower', interpolation='nearest')
 
 
-    # ==================================================
+        # ==================================================
     def print_summary(self):
 
         print('---------------------------------------------------')
@@ -1842,163 +2003,3 @@ def print_summary(self):
         print(f"rms beam length [um]:\t\t\t {self.bunch_length()*1e6 :.3f} \t\t")
         print(f"Peak current [kA]:\t\t\t {self.peak_current()/1e3 :.3f} \t\t")
         print('---------------------------------------------------')
-
-
-         ## SAVE AND LOAD BEAM
-    
-    def filename(self, runnable, beam_name):
-        return runnable.shot_path() + "/" + beam_name + "_" + str(self.trackable_number).zfill(3) + "_{:012.6F}".format(self.location) + ".h5"
-    
-    
-    # save beam (to OpenPMD format)
-    def save(self, runnable=None, filename=None, beam_name="beam", series=None):
-        
-        if len(self) == 0:
-            return
-        
-        # make new file if not provided
-        if series is None:
-
-            # open a new file
-            if runnable is not None:
-                filename = self.filename(runnable, beam_name)
-            
-            # open a new file
-            series = io.Series(filename, io.Access.create)
-            
-        
-            # add metadata
-            series.author = "ABEL (the Adaptable Beginning-to-End Linac simulation framework)"
-            series.date = datetime.now(timezone('CET')).strftime('%Y-%m-%d %H:%M:%S %z')
-
-        # make step (only one)
-        index = 0
-        iteration = series.iterations[index]
-        
-        # add attributes
-        iteration.set_attribute("time", self.location/SI.c)
-        for key, value in self.__dict__.items():
-            if not "__phasespace" in key:
-                iteration.set_attribute(key, value)
-                
-        # make beam record
-        particles = iteration.particles[beam_name]
-       
-        # generate datasets
-        dset_z = io.Dataset(self.zs().dtype, extent=self.zs().shape)
-        dset_x = io.Dataset(self.xs().dtype, extent=self.xs().shape)
-        dset_y = io.Dataset(self.ys().dtype, extent=self.ys().shape)
-        dset_zoff = io.Dataset(np.dtype('float64'), extent=[1])
-        dset_xoff = io.Dataset(np.dtype('float64'), extent=[1])
-        dset_yoff = io.Dataset(np.dtype('float64'), extent=[1])
-        dset_uz = io.Dataset(self.uzs().dtype, extent=self.uzs().shape)
-        dset_ux = io.Dataset(self.uxs().dtype, extent=self.uxs().shape)
-        dset_uy = io.Dataset(self.uys().dtype, extent=self.uys().shape)
-        dset_w = io.Dataset(self.weightings().dtype, extent=self.weightings().shape)
-        dset_id = io.Dataset(self.ids().dtype, extent=self.ids().shape)
-        dset_q = io.Dataset(np.dtype('float64'), extent=[1])
-        dset_m = io.Dataset(np.dtype('float64'), extent=[1])
-        
-        dset_n = io.Dataset(self.ids().dtype, extent=[1])
-        dset_f = io.Dataset(np.dtype('float64'), extent=[1])
-        
-        # prepare for writing
-        particles['position']['z'].reset_dataset(dset_z)
-        particles['position']['x'].reset_dataset(dset_x)
-        particles['position']['y'].reset_dataset(dset_y)
-        particles['positionOffset']['z'].reset_dataset(dset_zoff)
-        particles['positionOffset']['x'].reset_dataset(dset_xoff)
-        particles['positionOffset']['y'].reset_dataset(dset_yoff)
-        particles['momentum']['z'].reset_dataset(dset_uz)
-        particles['momentum']['x'].reset_dataset(dset_ux)
-        particles['momentum']['y'].reset_dataset(dset_uy)
-        particles['weighting'][io.Record_Component.SCALAR].reset_dataset(dset_w)
-        particles['id'][io.Record_Component.SCALAR].reset_dataset(dset_id)        
-        particles['charge'][io.Record_Component.SCALAR].reset_dataset(dset_q)
-        particles['mass'][io.Record_Component.SCALAR].reset_dataset(dset_m)
-        
-        # store data
-        particles['position']['z'].store_chunk(self.zs())
-        particles['position']['x'].store_chunk(self.xs())
-        particles['position']['y'].store_chunk(self.ys())
-        particles['positionOffset']['x'].make_constant(0.)
-        particles['positionOffset']['y'].make_constant(0.)
-        particles['positionOffset']['z'].make_constant(0.)
-        particles['momentum']['z'].store_chunk(self.uzs())
-        particles['momentum']['x'].store_chunk(self.uxs())
-        particles['momentum']['y'].store_chunk(self.uys())
-        particles['weighting'][io.Record_Component.SCALAR].store_chunk(self.weightings())
-        particles['id'][io.Record_Component.SCALAR].store_chunk(self.ids())
-        particles['charge'][io.Record_Component.SCALAR].make_constant(self.charge_sign()*SI.e)
-        particles['mass'][io.Record_Component.SCALAR].make_constant(SI.m_e)
-        
-        # set SI units (scaling factor)
-        particles['momentum']['z'].unit_SI = SI.m_e
-        particles['momentum']['x'].unit_SI = SI.m_e
-        particles['momentum']['y'].unit_SI = SI.m_e
-        
-        # set dimensional units
-        particles['position'].unit_dimension = {io.Unit_Dimension.L: 1}
-        particles['positionOffset'].unit_dimension = {io.Unit_Dimension.L: 1}
-        particles['momentum'].unit_dimension = {io.Unit_Dimension.L: 1, io.Unit_Dimension.M: 1, io.Unit_Dimension.T: -1}
-        particles['charge'].unit_dimension = {io.Unit_Dimension.T: 1, io.Unit_Dimension.I: 1}
-        particles['mass'].unit_dimension = {io.Unit_Dimension.M: 1}
-        
-        # save data to file
-        series.flush()
-        
-        return series
-        
-        
-    # load beam (from OpenPMD format)
-    @classmethod
-    def load(_, filename, beam_name='beam'):
-        
-        # load file
-        series = io.Series(filename, io.Access.read_only)
-        
-        # find index (use last one)
-        *_, index = series.iterations
-        
-        # get particle data
-        particles = series.iterations[index].particles[beam_name]
-        
-        # get attributes
-        charge = particles["charge"][io.Record_Component.SCALAR].get_attribute("value")
-        mass = particles["mass"][io.Record_Component.SCALAR].get_attribute("value")
-        
-        # extract phase space
-        ids = particles["id"][io.Record_Component.SCALAR].load_chunk()
-        weightings = particles["weighting"][io.Record_Component.SCALAR].load_chunk()
-        xs = particles['position']['x'].load_chunk()
-        ys = particles['position']['y'].load_chunk()
-        zs = particles['position']['z'].load_chunk()
-        pxs_unscaled = particles['momentum']['x'].load_chunk()
-        pys_unscaled = particles['momentum']['y'].load_chunk()
-        pzs_unscaled = particles['momentum']['z'].load_chunk()
-        series.flush()
-        
-        # apply SI scaling
-        pxs = pxs_unscaled * particles['momentum']['x'].unit_SI
-        pys = pys_unscaled * particles['momentum']['y'].unit_SI
-        pzs = pzs_unscaled * particles['momentum']['z'].unit_SI
-        
-        # make beam
-        beam = Beam()
-        beam.set_phase_space(Q=np.sum(weightings*charge), xs=xs, ys=ys, zs=zs, pxs=pxs, pys=pys, pzs=pzs, weightings=weightings)
-        
-        # add metadata to beam
-        try:
-            beam.trackable_number = series.iterations[index].get_attribute("trackable_number")
-            beam.stage_number = series.iterations[index].get_attribute("stage_number")
-            beam.location = series.iterations[index].get_attribute("location")
-            beam.num_bunches_in_train = series.iterations[index].get_attribute("num_bunches_in_train")
-            beam.bunch_separation = series.iterations[index].get_attribute("bunch_separation")
-        except:
-            beam.trackable_number = None
-            beam.stage_number = None
-            beam.location = None
-            beam.num_bunches_in_train = None
-            beam.bunch_separation = None
-        
-        return beam
\ No newline at end of file

From 5665394fd6c4facc0a5ba0f5e279aaed2dd44dd9 Mon Sep 17 00:00:00 2001
From: Kyrre Ness Sjobak <k.n.sjobak@fys.uio.no>
Date: Thu, 3 Apr 2025 15:54:41 +0200
Subject: [PATCH 49/51] Manually undo the moving of 'beam statistics' function
 from commit fcc2ce9ed879a5dcce4291138d9b8e5129080756

---
 abel/classes/beam.py | 464 +++++++++++++++++++++----------------------
 1 file changed, 229 insertions(+), 235 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index 2996ad47..0540596b 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -505,237 +505,6 @@ def norm_transverse_vector(self):
         vector[3,:] = self.uys()/SI.c
         return vector
 
-    
-
-    ## BEAM STATISTICS
-
-    def total_particles(self):
-        return int(np.nansum(self.weightings()))
-    
-    def charge(self):
-        return np.nansum(self.qs())
-    
-    def abs_charge(self):
-        return abs(self.charge())
-    
-    def charge_sign(self):
-        if self.charge() == 0:
-            return 1.0
-        else:
-            return self.charge()/abs(self.charge())
-    
-    def energy(self, clean=False):
-        return weighted_mean(self.Es(), self.weightings(), clean)
-    
-    def gamma(self, clean=False):
-        return weighted_mean(self.gammas(), self.weightings(), clean)
-    
-    def total_energy(self):
-        return SI.e * np.nansum(self.weightings()*self.Es())
-    
-    def energy_spread(self, clean=False):
-        return weighted_std(self.Es(), self.weightings(), clean)
-    
-    def rel_energy_spread(self, clean=False):
-        return self.energy_spread(clean)/self.energy(clean)
-    
-    def z_offset(self, clean=False):
-        return weighted_mean(self.zs(), self.weightings(), clean)
-    
-    def bunch_length(self, clean=False):
-        return weighted_std(self.zs(), self.weightings(), clean)
-    
-    def x_offset(self, clean=False):
-        return weighted_mean(self.xs(), self.weightings(), clean)
-    
-    def beam_size_x(self, clean=False):
-        return weighted_std(self.xs(), self.weightings(), clean)
-
-    def y_offset(self, clean=False):
-        return weighted_mean(self.ys(), self.weightings(), clean)
-
-    def beam_size_y(self, clean=False):
-        return weighted_std(self.ys(), self.weightings(), clean)
-    
-    def x_angle(self, clean=False):
-        return weighted_mean(self.xps(), self.weightings(), clean)
-    
-    def divergence_x(self, clean=False):
-        return weighted_std(self.xps(), self.weightings(), clean)
-
-    def y_angle(self, clean=False):
-        return weighted_mean(self.yps(), self.weightings(), clean)
-    
-    def divergence_y(self, clean=False):
-        return weighted_std(self.yps(), self.weightings(), clean)
-    
-    def ux_offset(self, clean=False):
-        return weighted_mean(self.uxs(), self.weightings(), clean)
-    
-    def uy_offset(self, clean=False):
-        return weighted_mean(self.uys(), self.weightings(), clean)
-    
-    def uz_offset(self, clean=False):
-        return weighted_mean(self.uzs(), self.weightings(), clean)
-
-    
-    def geom_emittance_x(self, clean=False):
-        return np.sqrt(np.linalg.det(weighted_cov(self.xs(), self.xps(), self.weightings(), clean)))
-    
-    def geom_emittance_y(self, clean=False):
-        return np.sqrt(np.linalg.det(weighted_cov(self.ys(), self.yps(), self.weightings(), clean)))
-    
-    def norm_emittance_x(self, clean=False):
-        return np.sqrt(np.linalg.det(weighted_cov(self.xs(), self.uxs()/SI.c, self.weightings(), clean)))
-    
-    def norm_emittance_y(self, clean=False):
-        return np.sqrt(np.linalg.det(weighted_cov(self.ys(), self.uys()/SI.c, self.weightings(), clean)))
-    
-    def beta_x(self, clean=False):
-        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
-        return covx[0,0]/np.sqrt(np.linalg.det(covx))
-    
-    def beta_y(self, clean=False):
-        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
-        return covy[0,0]/np.sqrt(np.linalg.det(covy))
-    
-    def alpha_x(self, clean=False):
-        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
-        return -covx[1,0]/np.sqrt(np.linalg.det(covx))
-    
-    def alpha_y(self, clean=False):
-        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
-        return -covy[1,0]/np.sqrt(np.linalg.det(covy))
-    
-    def gamma_x(self, clean=False):
-        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
-        return covx[1,1]/np.sqrt(np.linalg.det(covx))
-    
-    def gamma_y(self, clean=False):
-        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
-        return covy[1,1]/np.sqrt(np.linalg.det(covy))
-
-    def intrinsic_emittance(self):
-        covxy = np.cov(self.norm_transverse_vector(), aweights=self.weightings())
-        return np.sqrt(np.sqrt(np.linalg.det(covxy)))
-
-    def angular_momentum(self):
-        covxy = np.cov(self.norm_transverse_vector(), aweights=self.weightings())
-        det_covxy_cross = np.linalg.det(covxy[2:4,0:2])
-        return np.sign(covxy[3,0]-covxy[2,1])*np.sqrt(np.abs(det_covxy_cross))
-
-    def eigen_emittance_max(self):
-        return np.sqrt(self.norm_emittance_x()*self.norm_emittance_y()) + self.angular_momentum()
-
-    def eigen_emittance_min(self):
-        return np.sqrt(self.norm_emittance_x()*self.norm_emittance_y()) - self.angular_momentum()
-
-    def norm_amplitude_x(self, plasma_density=None, clean=False):
-        if plasma_density is not None:
-            beta_x = beta_matched(plasma_density, self.energy())
-            alpha_x = 0
-        else:
-            covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
-            emgx = np.sqrt(np.linalg.det(covx))
-            beta_x = covx[0,0]/emgx
-            alpha_x = -covx[1,0]/emgx
-        return np.sqrt(self.gamma()/beta_x)*np.sqrt(self.x_offset()**2 + (self.x_offset()*alpha_x + self.x_angle()*beta_x)**2)
-        
-    def norm_amplitude_y(self, plasma_density=None, clean=False):
-        if plasma_density is not None:
-            beta_y = beta_matched(plasma_density, self.energy())
-            alpha_y = 0
-        else:
-            covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
-            emgy = np.sqrt(np.linalg.det(covy))
-            beta_y = covy[0,0]/emgy
-            alpha_y = -covy[1,0]/emgy
-        return np.sqrt(self.gamma()/beta_y)*np.sqrt(self.y_offset()**2 + (self.y_offset()*alpha_y + self.y_angle()*beta_y)**2)
-        
-    def peak_density(self):  # TODO: this is only valid for Gaussian beams.
-        return (self.charge()/SI.e)/(np.sqrt(2*SI.pi)**3*self.beam_size_x()*self.beam_size_y()*self.bunch_length())
-    
-    def peak_current(self):
-        Is, _ = self.current_profile()
-        return max(abs(Is))
-    
-
-    ## BEAM HALO CLEANING (EXTREME OUTLIERS)
-    def remove_halo_particles(self, nsigma=20):
-        xfilter = np.abs(self.xs()-self.x_offset(clean=True)) > nsigma*self.beam_size_x(clean=True)
-        xpfilter = np.abs(self.xps()-self.x_angle(clean=True)) > nsigma*self.divergence_x(clean=True)
-        yfilter = np.abs(self.ys()-self.y_offset(clean=True)) > nsigma*self.beam_size_y(clean=True)
-        ypfilter = np.abs(self.yps()-self.y_angle(clean=True)) > nsigma*self.divergence_y(clean=True)
-        filter = np.logical_or(np.logical_or(xfilter, xpfilter), np.logical_or(yfilter, ypfilter))
-        del self[filter]
-
-    
-    ## BEAM PROJECTIONS
-    
-    def projected_density(self, fcn, bins=None):
-        if bins is None:
-            Nbins = int(np.sqrt(len(self)/2))
-            bins = np.linspace(min(fcn()), max(fcn()), Nbins)
-        counts, edges = np.histogram(fcn(), weights=self.qs(), bins=bins)
-        ctrs = (edges[0:-1] + edges[1:])/2
-        proj = counts/np.diff(edges)
-        return proj, ctrs
-        
-    def current_profile(self, bins=None):
-        return self.projected_density(self.ts, bins=bins)
-    
-    def longitudinal_num_density(self, bins=None):
-        dQdz, zs = self.projected_density(self.zs, bins=bins)
-        dNdz = dQdz / SI.e / self.charge_sign()
-        return dNdz, zs
-    
-    def energy_spectrum(self, bins=None):
-        return self.projected_density(self.Es, bins=bins)
-    
-    def rel_energy_spectrum(self, nom_energy=None, bins=None):
-        if nom_energy is None:
-            nom_energy = self.energy()
-        return self.projected_density(lambda: self.Es()/nom_energy-1, bins=bins)
-    
-    def transverse_profile_x(self, bins=None):
-        return self.projected_density(self.xs, bins=bins)
-    
-    def transverse_profile_y(self, bins=None):
-        return self.projected_density(self.ys, bins=bins)
-
-    def transverse_profile_xp(self, bins=None):
-        return self.projected_density(self.xps, bins=bins)
-    
-    def transverse_profile_yp(self, bins=None):
-        return self.projected_density(self.yps, bins=bins)
-    
-    ## phase spaces
-    
-    def phase_space_density(self, hfcn, vfcn, hbins=None, vbins=None):
-        self.remove_nans()
-        if hbins is None:
-            hbins = round(np.sqrt(len(self))/2)
-        if vbins is None:
-            vbins = round(np.sqrt(len(self))/2)
-        counts, hedges, vedges = np.histogram2d(hfcn(), vfcn(), weights=self.qs(), bins=(hbins, vbins))
-        hctrs = (hedges[0:-1] + hedges[1:])/2
-        vctrs = (vedges[0:-1] + vedges[1:])/2
-        density = (counts/np.diff(vedges)).T/np.diff(hedges)
-
-        #dx = np.diff(hedges)
-        #dy = np.diff(vedges)
-        #bin_areas = dx[:, None] * dy[None, :]
-        #density = counts/bin_areas
-        #print(np.sum(density*np.diff(vedges)*np.diff(hedges))/self.charge())
-        return density, hctrs, vctrs
-    
-    def density_lps(self, hbins=None, vbins=None):
-        return self.phase_space_density(self.zs, self.Es, hbins=hbins, vbins=vbins)
-    
-    def density_transverse(self, hbins=None, vbins=None):
-        return self.phase_space_density(self.xs, self.ys, hbins=hbins, vbins=vbins)
-
-
 ## Rotate the coordinate system of the beam
     # ==================================================
     def rotate_coord_sys_3D(self, axis1, angle1, axis2=np.array([0, 1, 0]), angle2=0.0, axis3=np.array([1, 0, 0]), angle3=0.0, invert=False):
@@ -1022,10 +791,7 @@ def x_tilt_angle(self, z_cutoff=None):
         return np.arctan(slope)
 
     
-    # # ==================================================
-    # def y_tilt_angle(self, clean=False):
-    #     return np.arcsin(self.uy_offset(clean=clean)/self.uz_offset(clean=clean))
-    
+    # ==================================================
     def y_tilt_angle(self, z_cutoff=None):
         "Retrieve the tilt angle of the beam in the zy-plane. WARNING: becomes unreliable around > 1e-4 rad."
         if z_cutoff is None:
@@ -1036,7 +802,235 @@ def y_tilt_angle(self, z_cutoff=None):
         slope, _ = np.polyfit(z_centroids, y_centroids, 1)
         
         return np.arctan(slope)
+
+    ## BEAM STATISTICS
+
+    def total_particles(self):
+        return int(np.nansum(self.weightings()))
+    
+    def charge(self):
+        return np.nansum(self.qs())
+    
+    def abs_charge(self):
+        return abs(self.charge())
+    
+    def charge_sign(self):
+        if self.charge() == 0:
+            return 1.0
+        else:
+            return self.charge()/abs(self.charge())
+    
+    def energy(self, clean=False):
+        return weighted_mean(self.Es(), self.weightings(), clean)
+    
+    def gamma(self, clean=False):
+        return weighted_mean(self.gammas(), self.weightings(), clean)
+    
+    def total_energy(self):
+        return SI.e * np.nansum(self.weightings()*self.Es())
+    
+    def energy_spread(self, clean=False):
+        return weighted_std(self.Es(), self.weightings(), clean)
+    
+    def rel_energy_spread(self, clean=False):
+        return self.energy_spread(clean)/self.energy(clean)
+    
+    def z_offset(self, clean=False):
+        return weighted_mean(self.zs(), self.weightings(), clean)
+    
+    def bunch_length(self, clean=False):
+        return weighted_std(self.zs(), self.weightings(), clean)
+    
+    def x_offset(self, clean=False):
+        return weighted_mean(self.xs(), self.weightings(), clean)
+    
+    def beam_size_x(self, clean=False):
+        return weighted_std(self.xs(), self.weightings(), clean)
+
+    def y_offset(self, clean=False):
+        return weighted_mean(self.ys(), self.weightings(), clean)
+
+    def beam_size_y(self, clean=False):
+        return weighted_std(self.ys(), self.weightings(), clean)
+    
+    def x_angle(self, clean=False):
+        return weighted_mean(self.xps(), self.weightings(), clean)
+    
+    def divergence_x(self, clean=False):
+        return weighted_std(self.xps(), self.weightings(), clean)
+
+    def y_angle(self, clean=False):
+        return weighted_mean(self.yps(), self.weightings(), clean)
+    
+    def divergence_y(self, clean=False):
+        return weighted_std(self.yps(), self.weightings(), clean)
+    
+    def ux_offset(self, clean=False):
+        return weighted_mean(self.uxs(), self.weightings(), clean)
+    
+    def uy_offset(self, clean=False):
+        return weighted_mean(self.uys(), self.weightings(), clean)
+    
+    def uz_offset(self, clean=False):
+        return weighted_mean(self.uzs(), self.weightings(), clean)
+
+    
+    def geom_emittance_x(self, clean=False):
+        return np.sqrt(np.linalg.det(weighted_cov(self.xs(), self.xps(), self.weightings(), clean)))
+    
+    def geom_emittance_y(self, clean=False):
+        return np.sqrt(np.linalg.det(weighted_cov(self.ys(), self.yps(), self.weightings(), clean)))
+    
+    def norm_emittance_x(self, clean=False):
+        return np.sqrt(np.linalg.det(weighted_cov(self.xs(), self.uxs()/SI.c, self.weightings(), clean)))
+    
+    def norm_emittance_y(self, clean=False):
+        return np.sqrt(np.linalg.det(weighted_cov(self.ys(), self.uys()/SI.c, self.weightings(), clean)))
+    
+    def beta_x(self, clean=False):
+        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
+        return covx[0,0]/np.sqrt(np.linalg.det(covx))
+    
+    def beta_y(self, clean=False):
+        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
+        return covy[0,0]/np.sqrt(np.linalg.det(covy))
+    
+    def alpha_x(self, clean=False):
+        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
+        return -covx[1,0]/np.sqrt(np.linalg.det(covx))
+    
+    def alpha_y(self, clean=False):
+        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
+        return -covy[1,0]/np.sqrt(np.linalg.det(covy))
+    
+    def gamma_x(self, clean=False):
+        covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
+        return covx[1,1]/np.sqrt(np.linalg.det(covx))
+    
+    def gamma_y(self, clean=False):
+        covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
+        return covy[1,1]/np.sqrt(np.linalg.det(covy))
+
+    def intrinsic_emittance(self):
+        covxy = np.cov(self.norm_transverse_vector(), aweights=self.weightings())
+        return np.sqrt(np.sqrt(np.linalg.det(covxy)))
+
+    def angular_momentum(self):
+        covxy = np.cov(self.norm_transverse_vector(), aweights=self.weightings())
+        det_covxy_cross = np.linalg.det(covxy[2:4,0:2])
+        return np.sign(covxy[3,0]-covxy[2,1])*np.sqrt(np.abs(det_covxy_cross))
+
+    def eigen_emittance_max(self):
+        return np.sqrt(self.norm_emittance_x()*self.norm_emittance_y()) + self.angular_momentum()
+
+    def eigen_emittance_min(self):
+        return np.sqrt(self.norm_emittance_x()*self.norm_emittance_y()) - self.angular_momentum()
+
+    def norm_amplitude_x(self, plasma_density=None, clean=False):
+        if plasma_density is not None:
+            beta_x = beta_matched(plasma_density, self.energy())
+            alpha_x = 0
+        else:
+            covx = weighted_cov(self.xs(), self.xps(), self.weightings(), clean)
+            emgx = np.sqrt(np.linalg.det(covx))
+            beta_x = covx[0,0]/emgx
+            alpha_x = -covx[1,0]/emgx
+        return np.sqrt(self.gamma()/beta_x)*np.sqrt(self.x_offset()**2 + (self.x_offset()*alpha_x + self.x_angle()*beta_x)**2)
+        
+    def norm_amplitude_y(self, plasma_density=None, clean=False):
+        if plasma_density is not None:
+            beta_y = beta_matched(plasma_density, self.energy())
+            alpha_y = 0
+        else:
+            covy = weighted_cov(self.ys(), self.yps(), self.weightings(), clean)
+            emgy = np.sqrt(np.linalg.det(covy))
+            beta_y = covy[0,0]/emgy
+            alpha_y = -covy[1,0]/emgy
+        return np.sqrt(self.gamma()/beta_y)*np.sqrt(self.y_offset()**2 + (self.y_offset()*alpha_y + self.y_angle()*beta_y)**2)
+        
+    def peak_density(self):  # TODO: this is only valid for Gaussian beams.
+        return (self.charge()/SI.e)/(np.sqrt(2*SI.pi)**3*self.beam_size_x()*self.beam_size_y()*self.bunch_length())
     
+    def peak_current(self):
+        Is, _ = self.current_profile()
+        return max(abs(Is))
+    
+
+    ## BEAM HALO CLEANING (EXTREME OUTLIERS)
+    def remove_halo_particles(self, nsigma=20):
+        xfilter = np.abs(self.xs()-self.x_offset(clean=True)) > nsigma*self.beam_size_x(clean=True)
+        xpfilter = np.abs(self.xps()-self.x_angle(clean=True)) > nsigma*self.divergence_x(clean=True)
+        yfilter = np.abs(self.ys()-self.y_offset(clean=True)) > nsigma*self.beam_size_y(clean=True)
+        ypfilter = np.abs(self.yps()-self.y_angle(clean=True)) > nsigma*self.divergence_y(clean=True)
+        filter = np.logical_or(np.logical_or(xfilter, xpfilter), np.logical_or(yfilter, ypfilter))
+        del self[filter]
+
+    
+    ## BEAM PROJECTIONS
+    
+    def projected_density(self, fcn, bins=None):
+        if bins is None:
+            Nbins = int(np.sqrt(len(self)/2))
+            bins = np.linspace(min(fcn()), max(fcn()), Nbins)
+        counts, edges = np.histogram(fcn(), weights=self.qs(), bins=bins)
+        ctrs = (edges[0:-1] + edges[1:])/2
+        proj = counts/np.diff(edges)
+        return proj, ctrs
+        
+    def current_profile(self, bins=None):
+        return self.projected_density(self.ts, bins=bins)
+    
+    def longitudinal_num_density(self, bins=None):
+        dQdz, zs = self.projected_density(self.zs, bins=bins)
+        dNdz = dQdz / SI.e / self.charge_sign()
+        return dNdz, zs
+    
+    def energy_spectrum(self, bins=None):
+        return self.projected_density(self.Es, bins=bins)
+    
+    def rel_energy_spectrum(self, nom_energy=None, bins=None):
+        if nom_energy is None:
+            nom_energy = self.energy()
+        return self.projected_density(lambda: self.Es()/nom_energy-1, bins=bins)
+    
+    def transverse_profile_x(self, bins=None):
+        return self.projected_density(self.xs, bins=bins)
+    
+    def transverse_profile_y(self, bins=None):
+        return self.projected_density(self.ys, bins=bins)
+
+    def transverse_profile_xp(self, bins=None):
+        return self.projected_density(self.xps, bins=bins)
+    
+    def transverse_profile_yp(self, bins=None):
+        return self.projected_density(self.yps, bins=bins)
+    
+    ## phase spaces
+    
+    def phase_space_density(self, hfcn, vfcn, hbins=None, vbins=None):
+        self.remove_nans()
+        if hbins is None:
+            hbins = round(np.sqrt(len(self))/2)
+        if vbins is None:
+            vbins = round(np.sqrt(len(self))/2)
+        counts, hedges, vedges = np.histogram2d(hfcn(), vfcn(), weights=self.qs(), bins=(hbins, vbins))
+        hctrs = (hedges[0:-1] + hedges[1:])/2
+        vctrs = (vedges[0:-1] + vedges[1:])/2
+        density = (counts/np.diff(vedges)).T/np.diff(hedges)
+
+        #dx = np.diff(hedges)
+        #dy = np.diff(vedges)
+        #bin_areas = dx[:, None] * dy[None, :]
+        #density = counts/bin_areas
+        #print(np.sum(density*np.diff(vedges)*np.diff(hedges))/self.charge())
+        return density, hctrs, vctrs
+    
+    def density_lps(self, hbins=None, vbins=None):
+        return self.phase_space_density(self.zs, self.Es, hbins=hbins, vbins=vbins)
+    
+    def density_transverse(self, hbins=None, vbins=None):
+        return self.phase_space_density(self.xs, self.ys, hbins=hbins, vbins=vbins)
+
     
     # ==================================================
     # TODO: Currently does not reproduce the correct peak density for Gaussian beams unless the bin numbers are adjusted manually.

From 35425313222583c12ba8205ce1815924d1a54a22 Mon Sep 17 00:00:00 2001
From: Kyrre Ness Sjobak <k.n.sjobak@fys.uio.no>
Date: Thu, 3 Apr 2025 16:01:11 +0200
Subject: [PATCH 50/51] Manually undo the moving of plotting functions from
 commit 74c833294e7492389d81b9888fbeaa044aa628d6

---
 abel/classes/beam.py | 159 +++++++++++++++++++++----------------------
 1 file changed, 79 insertions(+), 80 deletions(-)

diff --git a/abel/classes/beam.py b/abel/classes/beam.py
index 0540596b..a51a273e 100644
--- a/abel/classes/beam.py
+++ b/abel/classes/beam.py
@@ -1313,6 +1313,85 @@ def Ex_Ey(self, x_box_min, x_box_max, y_box_min, y_box_max, dx, dy, num_z_cells=
           
         
     
+    ## PLOTTING
+    def plot_current_profile(self):
+        dQdt, ts = self.current_profile()
+
+        fig, ax = plt.subplots()
+        fig.set_figwidth(6)
+        fig.set_figheight(4)        
+        ax.plot(ts*SI.c*1e6, np.abs(dQdt)/1e3)
+        ax.set_xlabel('z (um)')
+        ax.set_ylabel('Beam current (kA)')
+    
+    def plot_lps(self):
+        dQdzdE, zs, Es = self.density_lps()
+
+        fig, ax = plt.subplots()
+        fig.set_figwidth(8)
+        fig.set_figheight(5)  
+            
+        p = ax.pcolor(zs*1e6, Es/1e9, -dQdzdE*1e15, cmap=CONFIG.default_cmap, shading='auto')
+        ax.set_xlabel('z (um)')
+        ax.set_ylabel('E (GeV)')
+        ax.set_title('Longitudinal phase space')
+        cb = fig.colorbar(p)
+        cb.ax.set_ylabel('Charge density (pC/um/GeV)')
+        
+    def plot_trace_space_x(self):
+        dQdxdxp, xs, xps = self.phase_space_density(self.xs, self.xps)
+
+        fig, ax = plt.subplots()
+        fig.set_figwidth(8)
+        fig.set_figheight(5)  
+        p = ax.pcolor(xs*1e6, xps*1e3, -dQdxdxp*1e3, cmap=CONFIG.default_cmap, shading='auto')
+        ax.set_xlabel('x (um)')
+        ax.set_ylabel('x'' (mrad)')
+        ax.set_title('Horizontal trace space')
+        cb = fig.colorbar(p)
+        cb.ax.set_ylabel('Charge density (pC/um/mrad)')
+        
+    def plot_trace_space_y(self):
+        dQdydyp, ys, yps = self.phase_space_density(self.ys, self.yps)
+
+        fig, ax = plt.subplots()
+        fig.set_figwidth(8)
+        fig.set_figheight(5)  
+        p = ax.pcolor(ys*1e6, yps*1e3, -dQdydyp*1e3, cmap=CONFIG.default_cmap, shading='auto')
+        ax.set_xlabel('y (um)')
+        ax.set_ylabel('y'' (mrad)')
+        ax.set_title('Vertical trace space')
+        cb = fig.colorbar(p)
+        cb.ax.set_ylabel('Charge density (pC/um/mrad)')
+
+    def plot_transverse_profile(self):
+        dQdxdy, xs, ys = self.phase_space_density(self.xs, self.ys)
+
+        fig, ax = plt.subplots()
+        fig.set_figwidth(8)
+        fig.set_figheight(5)
+        p = ax.pcolor(xs*1e6, ys*1e6, -dQdxdy, cmap=CONFIG.default_cmap, shading='auto')
+        #p = ax.imshow(-dQdxdy, extent=[xs.min()*1e6, xs.max()*1e6, ys.min()*1e6, ys.max()*1e6], 
+        #   origin='lower', cmap=CONFIG.default_cmap, aspect='auto')
+        ax.set_xlabel('x (um)')
+        ax.set_ylabel('y (um)')
+        ax.set_title('Transverse profile')
+        cb = fig.colorbar(p)
+        cb.ax.set_ylabel('Charge density (pC/um^2)')
+
+    
+    # TODO: unfinished!
+    # def plot_bunch_pattern(self):
+        
+    #     fig, ax = plt.subplots()
+    #     fig.set_figwidth(6)
+    #     fig.set_figheight(4)
+    #     ax.plot(ts*SI.c*1e6, np.abs(dQdt)/1e3)
+    #     ax.set_xlabel('z (um)')
+    #     ax.set_ylabel('Beam current (kA)')
+        
+
+   
     ## CHANGE BEAM
     
     def accelerate(self, energy_gain=0, chirp=0, z_offset=0):
@@ -1505,86 +1584,6 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat
         else:
             return Es_final
         
-
-    
-     ## PLOTTING
-    def plot_current_profile(self):
-        dQdt, ts = self.current_profile()
-
-        fig, ax = plt.subplots()
-        fig.set_figwidth(6)
-        fig.set_figheight(4)        
-        ax.plot(ts*SI.c*1e6, np.abs(dQdt)/1e3)
-        ax.set_xlabel('z (um)')
-        ax.set_ylabel('Beam current (kA)')
-    
-    def plot_lps(self):
-        dQdzdE, zs, Es = self.density_lps()
-
-        fig, ax = plt.subplots()
-        fig.set_figwidth(8)
-        fig.set_figheight(5)  
-            
-        p = ax.pcolor(zs*1e6, Es/1e9, -dQdzdE*1e15, cmap=CONFIG.default_cmap, shading='auto')
-        ax.set_xlabel('z (um)')
-        ax.set_ylabel('E (GeV)')
-        ax.set_title('Longitudinal phase space')
-        cb = fig.colorbar(p)
-        cb.ax.set_ylabel('Charge density (pC/um/GeV)')
-        
-    def plot_trace_space_x(self):
-        dQdxdxp, xs, xps = self.phase_space_density(self.xs, self.xps)
-
-        fig, ax = plt.subplots()
-        fig.set_figwidth(8)
-        fig.set_figheight(5)  
-        p = ax.pcolor(xs*1e6, xps*1e3, -dQdxdxp*1e3, cmap=CONFIG.default_cmap, shading='auto')
-        ax.set_xlabel('x (um)')
-        ax.set_ylabel('x'' (mrad)')
-        ax.set_title('Horizontal trace space')
-        cb = fig.colorbar(p)
-        cb.ax.set_ylabel('Charge density (pC/um/mrad)')
-        
-    def plot_trace_space_y(self):
-        dQdydyp, ys, yps = self.phase_space_density(self.ys, self.yps)
-
-        fig, ax = plt.subplots()
-        fig.set_figwidth(8)
-        fig.set_figheight(5)  
-        p = ax.pcolor(ys*1e6, yps*1e3, -dQdydyp*1e3, cmap=CONFIG.default_cmap, shading='auto')
-        ax.set_xlabel('y (um)')
-        ax.set_ylabel('y'' (mrad)')
-        ax.set_title('Vertical trace space')
-        cb = fig.colorbar(p)
-        cb.ax.set_ylabel('Charge density (pC/um/mrad)')
-
-    def plot_transverse_profile(self):
-        dQdxdy, xs, ys = self.phase_space_density(self.xs, self.ys)
-
-        fig, ax = plt.subplots()
-        fig.set_figwidth(8)
-        fig.set_figheight(5)
-        p = ax.pcolor(xs*1e6, ys*1e6, -dQdxdy, cmap=CONFIG.default_cmap, shading='auto')
-        #p = ax.imshow(-dQdxdy, extent=[xs.min()*1e6, xs.max()*1e6, ys.min()*1e6, ys.max()*1e6], 
-        #   origin='lower', cmap=CONFIG.default_cmap, aspect='auto')
-        ax.set_xlabel('x (um)')
-        ax.set_ylabel('y (um)')
-        ax.set_title('Transverse profile')
-        cb = fig.colorbar(p)
-        cb.ax.set_ylabel('Charge density (pC/um^2)')
-
-    
-    # TODO: unfinished!
-    # def plot_bunch_pattern(self):
-        
-    #     fig, ax = plt.subplots()
-    #     fig.set_figwidth(6)
-    #     fig.set_figheight(4)
-    #     ax.plot(ts*SI.c*1e6, np.abs(dQdt)/1e3)
-    #     ax.set_xlabel('z (um)')
-    #     ax.set_ylabel('Beam current (kA)')
-        
-
   
     ## SAVE AND LOAD BEAM
     

From 95de216b259bb5e62c79d1b3e1feab7c8b921dfb Mon Sep 17 00:00:00 2001
From: Ben Chen <ben_chen93@hotmail.com>
Date: Thu, 3 Apr 2025 16:46:09 +0200
Subject: [PATCH 51/51] Changed beam reference file path in
 test_SourceFromFile2Beam() using os.sep.

---
 tests/test_sources.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tests/test_sources.py b/tests/test_sources.py
index 77de8df7..c7fe64ae 100644
--- a/tests/test_sources.py
+++ b/tests/test_sources.py
@@ -22,6 +22,7 @@
 
 import pytest
 from abel import *
+import os
 
 def check_beam_source_parameters(beam, source):
     assert np.allclose(beam.particle_mass, SI.m_e, rtol=1e-05, atol=1e-08)
@@ -183,7 +184,7 @@ def test_SourceCapsule2Beam():
 def test_SourceFromFile2Beam():
     "Check that ``SourceFromFile`` retuns returns the same ``Beam`` as the input ``Beam``."
 
-    beam_file = 'tests/data/test_StagePrtclTransWakeInstability_beamline/test_baseline_linac/shot_000/beam_003_00048.558626.h5'
+    beam_file = 'tests' + os.sep + 'data' + os.sep + 'test_StagePrtclTransWakeInstability_beamline' + os.sep + 'test_baseline_linac' + os.sep + 'shot_000' + os.sep + 'beam_003_00048.558626.h5'
 
     source = SourceFromFile()
     source.file = beam_file
@@ -193,10 +194,10 @@ def test_SourceFromFile2Beam():
 
     # Trigger exception for when a file does not exist
     with pytest.raises(FileNotFoundError):
-        file = 'tests/data/test_StagePrtclTransWakeInstability_beamline/test_baseline_linac/blabla.h5'
+        file = 'tests' + os.sep + 'data' + os.sep + 'test_StagePrtclTransWakeInstability_beamline' + os.sep + 'test_baseline_linac' + os.sep + 'blabla.h5'
         source = SourceFromFile(file=file)
     with pytest.raises(FileNotFoundError):
-        file = 'tests/data/test_StagePrtclTransWakeInstability_beamline/test_baseline_linac/shot_000/blabla.h5'
+        file = 'tests' + os.sep + 'data' + os.sep + 'test_StagePrtclTransWakeInstability_beamline' + os.sep + 'test_baseline_linac' + os.sep + 'shot_000' + os.sep + 'blabla.h5'
         source = SourceFromFile()
         source.file = file
 
