Implement tracker (#75)

* Add time as input for the field gathering

* Change field gathering to be additive

* Implement null components more efficiently

* Implement gathering of multiple fields

* Make `update` method exclusive of `NumericalField`

* Implement new `Tracker` class

* Implement new tracker in `PlasmaStage`

* Remove unused imports. Fix bug.

* Implement more efficient optimized `dt`

* Add docstrings

* Implement new progress bar

* Fix formatting

* Implement `name` in `PlasmaStage`

* Fix formatting

* Make sure laser is evolved until `t_final` + fix

* Fix bug in adaptive time step

* Add new test

This test used to fail with the old pipeline. With the new `Tracker`
it now works as expected.

* Implement name in other plasma elements

* Fix bug

* Fix bug

* Accept list of fields in diagnostics

* Implement external fields and no wakefield

* Implement linear b_theta field

* Update active plasma lens to new implementation

* Add `_pre_gather` method to `AnalyticField`

* Update docstring

* Remake `FocusingBlowoutField` as `AnalyticField`

* Implement more efficient auto dt

* Update `PlasmaRamp` implementation

* Implement simple blowout as `AnalyticField`

* Add simple blowout test

* Implement custom blowout as `AnalyticField`

* Add custom blowout test

* Update implementation of fluid model

* Update list of plasma wakefields

* Remove old wakefields

* Update tests

* Evolve laser using correct plasma density

* Rename `AnalyticField` to `AnalyticalField`

tests/test_active_plasma_lens.py

@@ -29,7 +29,7 @@ def test_active_plasma_lens():
     apl.track(bunch)
     bunch_params = analyze_bunch(bunch)
     gamma_x = bunch_params['gamma_x']
-    assert approx(gamma_x, rel=1e-10) == 92.14018794551832
+    assert approx(gamma_x, rel=1e-10) == 92.14017315271572
 
 
 def test_active_plasma_lens_with_wakefields():
@@ -55,7 +55,7 @@ def test_active_plasma_lens_with_wakefields():
     apl = ActivePlasmaLens(
         1e-2, 1000, wakefields=True, n_out=3, density=1e23,
         wakefield_model='quasistatic_2d', r_max=200e-6, xi_min=-30e-6,
-        xi_max=30e-6, n_r=200, n_xi=120, dz_fields=0.1e-3, ppc=10
+        xi_max=30e-6, n_r=200, n_xi=120, dz_fields=0.2e-3, ppc=4
     )
 
     # Do tracking.
@@ -64,7 +64,7 @@ def test_active_plasma_lens_with_wakefields():
     # Analyze and check results.
     bunch_params = analyze_bunch(bunch)
     gamma_x = bunch_params['gamma_x']
-    assert approx(gamma_x, rel=1e-10) == 77.48240050479696
+    assert approx(gamma_x, rel=1e-10) == 77.32005041793602
 
 
 if __name__ == '__main__':

tests/test_custom_blowout.py

@@ -0,0 +1,84 @@
+from pytest import approx
+import numpy as np
+import matplotlib.pyplot as plt
+from aptools.plotting.quick_diagnostics import slice_analysis
+
+from wake_t import PlasmaStage, GaussianPulse
+from wake_t.utilities.bunch_generation import get_matched_bunch
+from wake_t.diagnostics import analyze_bunch, analyze_bunch_list
+
+
+def test_custom_blowout_wakefield(make_plots=False):
+    """Check that the `'custom_blowout'` wakefield model works as expected.
+
+    A matched bunch with high emittance is propagated through a 3 cm plasma
+    stage. Due to the high emittance and Ez slope, the beam develops not only
+    a large chirp but also an uncorrelated energy spread. The test checks that
+    the final projected energy spread is always the same.
+    """
+    # Set numpy random seed to get reproducible results
+    np.random.seed(1)
+
+    # Create laser driver.
+    laser = GaussianPulse(100e-6, l_0=800e-9, w_0=50e-6, a_0=3,
+                        tau=30e-15, z_foc=0.)
+
+    # Create bunch (matched to a blowout at a density of 10^{23} m^{-3}).
+    en = 10e-6  # m
+    ene = 200  # units of beta*gamma
+    ene_sp = 0.3  # %
+    xi_c = laser.xi_c - 55e-6  # m
+    s_t = 10  # fs
+    q_tot = 100  # pC
+    n_part = 3e4
+    k_x = 1e6
+    bunch = get_matched_bunch(en, en, ene, ene_sp, s_t, xi_c, q_tot, n_part,
+                              k_x=k_x)
+
+    # Create plasma stage.
+    plasma = PlasmaStage(
+        3e-2, 1e23, laser=laser, wakefield_model='custom_blowout',
+        lon_field=-10e9, lon_field_slope=1e15, foc_strength=k_x,
+        xi_fields=xi_c, n_out=50, bunch_pusher='boris')
+
+    # Do tracking.
+    bunch_list = plasma.track(bunch)
+
+    bunch_params = analyze_bunch(bunch)
+    rel_ene_sp = bunch_params['rel_ene_spread']
+    assert approx(rel_ene_sp, rel=1e-10) == 0.21192531095086517
+
+    if make_plots:
+        # Analyze bunch evolution.
+        params_evolution = analyze_bunch_list(bunch_list)
+
+
+        # Quick plot of results.
+        z = params_evolution['prop_dist'] * 1e2
+        fig_1 = plt.figure()
+        plt.subplot(411)
+        plt.plot(z, params_evolution['beta_x']*1e3)
+        plt.tick_params(axis='x', which='both', labelbottom=False)
+        plt.ylabel("$\\beta_x$ [mm]")
+        plt.subplot(412)
+        plt.plot(z, params_evolution['emitt_x']*1e6)
+        plt.tick_params(axis='x', which='both', labelbottom=False)
+        plt.ylabel("$\\epsilon_{nx}$ [$\\mu$m]")
+        plt.subplot(413)
+        plt.plot(z, params_evolution['rel_ene_spread']*100)
+        plt.tick_params(axis='x', which='both', labelbottom=False)
+        plt.ylabel("$\\frac{\\Delta \\gamma}{\\gamma}$ [%]")
+        plt.subplot(414)
+        plt.plot(z, params_evolution['avg_ene'])
+        plt.xlabel("z [mm]")
+        plt.ylabel("$\\gamma$")
+        plt.tight_layout()
+        fig_2 = plt.figure()
+        slice_analysis(
+            bunch.x, bunch.y, bunch.xi, bunch.px, bunch.py, bunch.pz,
+            bunch.q, fig=fig_2)
+        plt.show()
+
+
+if __name__ == "__main__":
+    test_custom_blowout_wakefield(make_plots=True)

tests/test_field_gathering.py

@@ -1,4 +1,5 @@
 import numpy as np
+import scipy.constants as ct
 from wake_t.particles.interpolation import (
     gather_field_cyl_linear, gather_main_fields_cyl_linear)
 from wake_t.utilities.bunch_generation import get_matched_bunch
@@ -72,17 +73,31 @@ def test_gather_main_fields_cyl_linear():
     y_part = np.zeros_like(x_part)
     z_part = Z.flatten()
 
+    # Preallocate field arrays
+    n_part = x_part.shape[0]
+    ex = np.zeros(n_part)
+    ey = np.zeros(n_part)
+    ez = np.zeros(n_part)
+    bx = np.zeros(n_part)
+    by = np.zeros(n_part)
+    bz = np.zeros(n_part)
+
     # Gather field
-    wx_part, wy_part, ez_part = gather_main_fields_cyl_linear(
-        f, f, z_min, z_max, r_min+dr/2, r_max, dz, dr, x_part, y_part, z_part)
+    gather_main_fields_cyl_linear(
+        f, f, f, z_min, z_max, r_min+dr/2, r_max, dz, dr,
+        x_part, y_part, z_part, ex, ey, ez, bx, by, bz)
 
-    wx_part = np.reshape(wx_part, (n_z, n_r))
-    wy_part = np.reshape(wy_part, (n_z, n_r))
-    ez_part = np.reshape(ez_part, (n_z, n_r))
+    ex = np.reshape(ex, (n_z, n_r))
+    ey = np.reshape(ey, (n_z, n_r))
+    ez = np.reshape(ez, (n_z, n_r))
+    bx = np.reshape(bx, (n_z, n_r))
+    by = np.reshape(by, (n_z, n_r))
+    bz = np.reshape(bz, (n_z, n_r))
 
     # Check
-    np.testing.assert_array_almost_equal(wx_part, f[2:-2, 2:-2])
-    np.testing.assert_array_almost_equal(ez_part, f[2:-2, 2:-2])
+    np.testing.assert_array_almost_equal(ex, f[2:-2, 2:-2])
+    np.testing.assert_array_almost_equal(ez, f[2:-2, 2:-2])
+    np.testing.assert_array_almost_equal(by, f[2:-2, 2:-2])
 
 
 def test_gather_main_fields_cyl_linear_at_bunch():
@@ -123,15 +138,18 @@ def test_gather_main_fields_cyl_linear_at_bunch():
     f = np.zeros((n_z+4, n_r+4))
     f[2:-2, 2:-2] = np.sin(R/r_max-1) * np.sin(Z/z_max-1)
 
+    # Get preallocated field arrays
+    ex, ey, ez, bx, by, bz = bunch.get_field_arrays()
+
     # Gather field
-    wx_part, wy_part, ez_part = gather_main_fields_cyl_linear(
-        f, f, z_min, z_max, r_min+dr/2, r_max, dz, dr,
-        bunch.x, bunch.y, bunch.xi)
+    gather_main_fields_cyl_linear(
+        f, f, f, z_min, z_max, r_min+dr/2, r_max, dz, dr,
+        bunch.x, bunch.y, bunch.xi, ex, ey, ez, bx, by, bz)
 
     # Check
-    np.testing.assert_almost_equal(np.sum(wx_part), 15.709780320676144)
-    np.testing.assert_almost_equal(np.sum(wy_part), -75.57014220169886)
-    np.testing.assert_almost_equal(np.sum(ez_part), 6798.124201568496)
+    np.testing.assert_almost_equal(np.sum(ex / ct.c + by), 15.709780373078333)
+    np.testing.assert_almost_equal(np.sum(ey / ct.c - bx), -75.57014245377374)
+    np.testing.assert_almost_equal(np.sum(ez), 6798.124201568496)
 
 
 if __name__ == '__main__':

tests/test_parabolic_profile.py

@@ -0,0 +1,99 @@
+from copy import deepcopy
+
+import numpy as np
+import scipy.constants as ct
+
+from wake_t import GaussianPulse, PlasmaStage, Beamline
+from wake_t.utilities.bunch_generation import get_matched_bunch
+
+
+def test_variable_parabolic_coefficient():
+    """
+    Checks that a z-dependent parabolic coefficient works as expected.
+    A piecewise function for the coefficient is defined, which should be
+    equivalent to tracking multiple plasma stages each having a different
+    parabolic coefficient.
+    """
+    # Plasma density.
+    n_p = 1e23
+
+    # Grid and field parameters.
+    xi_max = 0
+    xi_min = -100e-6
+    r_max = 100e-6
+    Nxi = 100
+    Nr = 100
+    dz_fields = 1e-3
+
+    # Laser parameters.
+    a0 = 1
+    w0 = 30e-6
+    tau = 25e-15
+    l0 = 0.8e-6
+
+    # Guiding plasma channel
+    r_e = ct.e**2 / (4. * np.pi * ct.epsilon_0 * ct.m_e * ct.c**2)
+    rel_delta_n_over_w2 = 1. / (np.pi * r_e * w0**4 * n_p)
+
+    # Length and parabolic coefficient of each section (stretch).
+    L_stretches = [1e-2, 1e-2, 1e-2, 1e-2, 1e-2]
+    pc_stretches = rel_delta_n_over_w2 * np.array([1, 2, 1, 0.5, 2])
+
+    # Define z-dependent parabolic coefficient.
+    def parabolic_coefficient(z):
+        z_stretch_end = np.cumsum(L_stretches)
+        i = np.sum(np.float32(z_stretch_end) <= np.float32(z))
+        if i == len(L_stretches):
+            i -= 1
+        return pc_stretches[i]
+
+    # Create identical laser pulses for each case.
+    laser = GaussianPulse(-50e-6, a0, w0, tau, z_foc=0., l_0=l0)
+    laser_1 = deepcopy(laser)
+    laser_2 = deepcopy(laser)
+
+    # Create identical bunches for each case.
+    bunch = get_matched_bunch(
+        1e-6, 1e-6, 200, 1, 3, laser.xi_c - 30e-6, 1e-6, 1e4, n_p=n_p)
+    bunch_1 = deepcopy(bunch)
+    bunch_2 = deepcopy(bunch)
+
+    # Create single plasma stage (containing all sections).
+    plasma_single = PlasmaStage(
+        np.sum(L_stretches), n_p, wakefield_model='quasistatic_2d', n_out=10,
+        xi_min=xi_min, xi_max=xi_max, r_max=r_max, r_max_plasma=r_max,
+        laser=laser_1, laser_evolution=True, n_r=Nr, n_xi=Nxi, ppc=2,
+        dz_fields=dz_fields, parabolic_coefficient=parabolic_coefficient)
+
+    # Track single plasma.
+    plasma_single.track(bunch_1)
+
+    # Create set of plasma stages, one per section.
+    sub_stages = []
+    for i, (l, pc) in enumerate(zip(L_stretches, pc_stretches)):
+        stage = PlasmaStage(
+            l, n_p, wakefield_model='quasistatic_2d', n_out=3,
+            xi_min=xi_min, xi_max=xi_max, r_max=r_max, r_max_plasma=r_max,
+            laser=laser_2, laser_evolution=True, n_r=Nr, n_xi=Nxi, ppc=2,
+            dz_fields=dz_fields, parabolic_coefficient=pc)
+        sub_stages.append(stage)
+    plasma_multi = Beamline(sub_stages)
+
+    # Track through set of stages.
+    plasma_multi.track(bunch_2)
+
+    # Get final envelope of both lasers.
+    a_env_1 = laser_1.get_envelope()
+    a_mod_1 = np.abs(a_env_1)
+    a_phase_1 = np.angle(a_env_1)
+    a_env_2 = laser_2.get_envelope()
+    a_mod_2 = np.abs(a_env_2)
+    a_phase_2 = np.angle(a_env_2)
+
+    # Check that both envelopes are equal.
+    np.testing.assert_almost_equal(a_mod_2, a_mod_1)
+    np.testing.assert_almost_equal(a_phase_2, a_phase_1)
+
+
+if __name__ == "__main__":
+    test_variable_parabolic_coefficient()

tests/test_ramps.py

@@ -33,7 +33,7 @@ def test_downramp():
     downramp.track(bunch)
     bunch_params = analyze_bunch(bunch)
     beta_x = bunch_params['beta_x']
-    assert beta_x == 0.009757683520255687
+    assert beta_x == 0.009756658297049697
 
 
 def test_upramp():
@@ -64,7 +64,7 @@ def test_upramp():
     downramp.track(bunch)
     bunch_params = analyze_bunch(bunch)
     beta_x = bunch_params['beta_x']
-    assert beta_x == 0.0007641795179724504
+    assert beta_x == 0.0007642922463413662
 
 
 if __name__ == '__main__':

tests/test_simple_blowout.py

@@ -0,0 +1,83 @@
+from pytest import approx
+import numpy as np
+import matplotlib.pyplot as plt
+from aptools.plotting.quick_diagnostics import slice_analysis
+
+from wake_t import PlasmaStage, GaussianPulse
+from wake_t.utilities.bunch_generation import get_matched_bunch
+from wake_t.diagnostics import analyze_bunch, analyze_bunch_list
+
+
+def test_simple_blowout_wakefield(make_plots=False):
+    """Check that the `'simple_blowout'` wakefield model works as expected.
+
+    A matched bunch with high emittance is propagated through a 3 cm plasma
+    stage. Due to the high emittance and Ez slope, the beam develops not only
+    a large chirp but also an uncorrelated energy spread. The test checks that
+    the final projected energy spread is always the same.
+    """
+    # Set numpy random seed to get reproducible results
+    np.random.seed(1)
+
+    # Create laser driver.
+    laser = GaussianPulse(100e-6, l_0=800e-9, w_0=50e-6, a_0=3,
+                        tau=30e-15, z_foc=0.)
+
+
+    # Create bunch (matched to a blowout at a density of 10^{23} m^{-3}).
+    en = 10e-6  # m
+    ene = 200  # units of beta*gamma
+    ene_sp = 0.3  # %
+    xi_c = laser.xi_c - 55e-6  # m
+    s_t = 10  # fs
+    q_tot = 100  # pC
+    n_part = 3e4
+    bunch = get_matched_bunch(en, en, ene, ene_sp, s_t, xi_c, q_tot, n_part,
+                            n_p=1e23)
+
+    # Create plasma stage.
+    plasma = PlasmaStage(
+        3e-2, 1e23, laser=laser, wakefield_model='simple_blowout', n_out=50,
+        bunch_pusher='boris')
+
+    # Do tracking.
+    bunch_list = plasma.track(bunch)
+
+    bunch_params = analyze_bunch(bunch)
+    rel_ene_sp = bunch_params['rel_ene_spread']
+    assert approx(rel_ene_sp, rel=1e-10) == 0.36376504595288617
+
+    if make_plots:
+        # Analyze bunch evolution.
+        params_evolution = analyze_bunch_list(bunch_list)
+
+
+        # Quick plot of results.
+        z = params_evolution['prop_dist'] * 1e2
+        fig_1 = plt.figure()
+        plt.subplot(411)
+        plt.plot(z, params_evolution['beta_x']*1e3)
+        plt.tick_params(axis='x', which='both', labelbottom=False)
+        plt.ylabel("$\\beta_x$ [mm]")
+        plt.subplot(412)
+        plt.plot(z, params_evolution['emitt_x']*1e6)
+        plt.tick_params(axis='x', which='both', labelbottom=False)
+        plt.ylabel("$\\epsilon_{nx}$ [$\\mu$m]")
+        plt.subplot(413)
+        plt.plot(z, params_evolution['rel_ene_spread']*100)
+        plt.tick_params(axis='x', which='both', labelbottom=False)
+        plt.ylabel("$\\frac{\\Delta \\gamma}{\\gamma}$ [%]")
+        plt.subplot(414)
+        plt.plot(z, params_evolution['avg_ene'])
+        plt.xlabel("z [mm]")
+        plt.ylabel("$\\gamma$")
+        plt.tight_layout()
+        fig_2 = plt.figure()
+        slice_analysis(
+            bunch.x, bunch.y, bunch.xi, bunch.px, bunch.py, bunch.pz,
+            bunch.q, fig=fig_2)
+        plt.show()
+
+
+if __name__ == "__main__":
+    test_simple_blowout_wakefield(make_plots=True)