1 files changed, 116 insertions, 0 deletions
diff --git a/Examples/Tests/AcceleratorLattice/analysis.py b/Examples/Tests/AcceleratorLattice/analysis.py
new file mode 100755
index 000000000..d2fd7f6ff
--- /dev/null
+++ b/Examples/Tests/AcceleratorLattice/analysis.py
@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+
+# Copyright 2022 David Grote
+#
+# This file is part of WarpX.
+#
+# License: BSD-3-Clause-LBNL
+
+
+"""
+This script tests the quad lattice
+The input file sets up a series of quadrupoles and propagates two particles through them.
+One particle is in the X plane, the other the Y plane.
+The final positions are compared to the analytic solutions.
+The motion is slow enough that relativistic effects are ignored.
+"""
+
+import os
+import sys
+
+import numpy as np
+from scipy.constants import c, e, m_e
+import yt
+
+yt.funcs.mylog.setLevel(0)
+sys.path.insert(1, '../../../../warpx/Regression/Checksum/')
+import checksumAPI
+
+filename = sys.argv[1]
+ds = yt.load( filename )
+ad = ds.all_data()
+
+gamma_boost = float(ds.parameters.get('warpx.gamma_boost', 1.))
+uz_boost = np.sqrt(gamma_boost*gamma_boost - 1.)*c
+
+# Fetch the final particle position
+xx_sim = ad['electron', 'particle_position_x'].v[0]
+zz_sim = ad['electron', 'particle_position_z'].v[0]
+ux_sim = ad['electron', 'particle_momentum_x'].v[0]/m_e
+
+if gamma_boost > 1.:
+    # The simulation data is in the boosted frame.
+    # Transform the z position to the lab frame.
+    time = ds.current_time.value
+    zz_sim = gamma_boost*zz_sim + uz_boost*time;
+
+# Fetch the quadrupole lattice data
+quad_starts = []
+quad_lengths = []
+quad_strengths_E = []
+z_location = 0.
+def read_lattice(rootname, z_location):
+    lattice_elements = ds.parameters.get(f'{rootname}.elements').split()
+    for element in lattice_elements:
+        element_type = ds.parameters.get(f'{element}.type')
+        if element_type == 'drift':
+            length = float(ds.parameters.get(f'{element}.ds'))
+            z_location += length
+        elif element_type == 'quad':
+            length = float(ds.parameters.get(f'{element}.ds'))
+            quad_starts.append(z_location)
+            quad_lengths.append(length)
+            quad_strengths_E.append(float(ds.parameters.get(f'{element}.dEdx')))
+            z_location += length
+        elif element_type == 'lattice':
+            z_location = read_lattice(element, z_location)
+    return z_location
+
+read_lattice('lattice', z_location)
+
+# Fetch the initial position of the particle
+x0 = [float(x) for x in ds.parameters.get('electron.single_particle_pos').split()]
+ux0 = [float(x)*c for x in ds.parameters.get('electron.single_particle_vel').split()]
+
+xx = x0[0]
+zz = x0[2]
+ux = ux0[0]
+uz = ux0[2]
+
+gamma = np.sqrt(uz**2/c**2 + 1.)
+vz = uz/gamma
+
+def applylens(x0, vx0, vz0, gamma, lens_length, lens_strength):
+    """Use analytic solution of a particle with a transverse dependent field"""
+    kb0 = np.sqrt(e/(m_e*gamma*vz0**2)*abs(lens_strength))
+    if lens_strength >= 0.:
+        x1 = x0*np.cos(kb0*lens_length) + (vx0/vz0)/kb0*np.sin(kb0*lens_length)
+        vx1 = vz0*(-kb0*x0*np.sin(kb0*lens_length) + (vx0/vz0)*np.cos(kb0*lens_length))
+    else:
+        x1 = x0*np.cosh(kb0*lens_length) + (vx0/vz0)/kb0*np.sinh(kb0*lens_length)
+        vx1 = vz0*(+kb0*x0*np.sinh(kb0*lens_length) + (vx0/vz0)*np.cosh(kb0*lens_length))
+    return x1, vx1
+
+# Integrate the particle using the analytic solution
+for i in range(len(quad_starts)):
+    z_lens = quad_starts[i]
+    vx = ux/gamma
+    dt = (z_lens - zz)/vz
+    xx = xx + dt*vx
+    xx, vx = applylens(xx, vx, vz, gamma, quad_lengths[i], quad_strengths_E[i])
+    ux = gamma*vx
+    zz = z_lens + quad_lengths[i]
+
+dt = (zz_sim - zz)/vz
+vx = ux/gamma
+xx = xx + dt*vx
+
+# Compare the analytic to the simulated final values
+print(f'Error in x position is {abs(np.abs((xx - xx_sim)/xx))}, which should be < 0.01')
+print(f'Error in x velocity is {abs(np.abs((ux - ux_sim)/ux))}, which should be < 0.002')
+
+assert abs(np.abs((xx - xx_sim)/xx)) < 0.01, Exception('error in x particle position')
+assert abs(np.abs((ux - ux_sim)/ux)) < 0.002, Exception('error in x particle velocity')
+
+test_name = os.path.split(os.getcwd())[1]
+checksumAPI.evaluate_checksum(test_name, filename)