Improve recursive evolution plot

src/abel/classes/beam.py

@@ -1903,47 +1903,47 @@ 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 probe_evolution:  # TODO: This seems to be very clumsy. Consider re-writing.
+        if probe_evolution:
 
             # calculate evolution
-            num_evol_steps = max(20, min(400, round(2*L/(beta_matched(n0, self.energy()+min(0,np.mean(deltaEs)))))))
+            num_evol_steps = max(20, min(500, round(3*L/(beta_matched(n0, self.energy()+min(0,np.mean(deltaEs)))))))
             evol = SimpleNamespace()
-            evol.location = np.linspace(0, L, num_evol_steps)
-            evol.x, evol.ux, evol.ux, evol.y, evol.uy, evol.energy, evol.energy_spread, evol.rel_energy_spread, evol.beam_size_x, evol.beam_size_y, evol.emit_nx, evol.emit_ny, evol.beta_x, evol.beta_y = (np.empty(evol.location.shape) for _ in range(14))
 
             # select a given subset of the particles (for faster calculation)
-            sample_fraction = round(2*np.sqrt(len(self))) #
-            inds_sample = np.arange(0, len(self), sample_fraction, dtype=int) # not random, to be consistent across ramps and stages
-            xs_sample, uxs_sample = self.xs()[inds_sample], self.uxs()[inds_sample]
-            ys_sample, uys_sample = self.ys()[inds_sample], self.uys()[inds_sample]
-            gamma0s_sample, dgamma_ds_sample = gamma0s[inds_sample], dgamma_ds[inds_sample]
-            Es_sample, deltaEs_sample = self.Es()[inds_sample], deltaEs[inds_sample]
-
-            # go through steps
-            for i in range(num_evol_steps):
-
-                # evolve the beam (no radiation reaction)
-                xs_i, uxs_i = evolve_hills_equation_analytic(xs_sample-x0_driver, uxs_sample, evol.location[i], gamma0s_sample, dgamma_ds_sample, k_p(n0))
-                ys_i, uys_i = evolve_hills_equation_analytic(ys_sample-y0_driver, uys_sample, evol.location[i], gamma0s_sample, dgamma_ds_sample, k_p(n0))
-                Es_i = Es_sample+deltaEs_sample*i/(num_evol_steps-1)
-                uzs_i = energy2proper_velocity(Es_i)
-
-                # save the parameters
-                evol.x[i], evol.ux[i] = np.mean(xs_i), np.mean(uxs_i)
-                evol.y[i], evol.uy[i] = np.mean(ys_i), np.mean(uys_i)
-                evol.energy[i], evol.energy_spread[i] = np.mean(Es_i), np.std(Es_i)
-                evol.beam_size_x[i], evol.beam_size_y[i] = np.std(xs_i), np.std(ys_i)
-                evol.emit_nx[i] = np.sqrt(np.linalg.det(np.cov(xs_i, uxs_i/SI.c))) # only works for equal weight particles
-                evol.emit_ny[i] = np.sqrt(np.linalg.det(np.cov(ys_i, uys_i/SI.c))) # only works for equal weight particles
-                covx, covy = np.cov(xs_i, uxs_i/uzs_i), np.cov(ys_i, uys_i/uzs_i)
+            sample_fraction = round(4*np.sqrt(len(self))) #
+            inds_samp = np.arange(0, len(self), sample_fraction, dtype=int) # not random, to be consistent across ramps and stages
+            xs_samp, uxs_samp = self.xs()[inds_samp], self.uxs()[inds_samp]
+            ys_samp, uys_samp = self.ys()[inds_samp], self.uys()[inds_samp]
+            gamma0s_samp, dgamma_ds_samp = gamma0s[inds_samp], dgamma_ds[inds_samp]
+            Es_samp, deltaEs_samp = self.Es()[inds_samp], deltaEs[inds_samp]
+
+            # simulate evolution
+            from abel.physics_models.hills_equation import evolve_hills_equation_analytic_evolution
+            xs_all, uxs_all, gammas_all, ss = evolve_hills_equation_analytic_evolution(xs_samp-x0_driver, uxs_samp, L, gamma0s_samp, dgamma_ds_samp, k_p(n0), num_steps=num_evol_steps)
+            ys_all, uys_all, _, _ = evolve_hills_equation_analytic_evolution(ys_samp-y0_driver, uys_samp, L, gamma0s_samp, dgamma_ds_samp, k_p(n0), num_steps=num_evol_steps)
+            Es_all = gammas_all/(SI.e/(SI.m_e*SI.c**2))
+            uzs_all = energy2proper_velocity(Es_all)
+
+            # extract evolution
+            evol.location = ss
+            evol.x, evol.ux = np.mean(xs_all, axis=0), np.mean(uxs_all, axis=0)
+            evol.y, evol.uy = np.mean(ys_all, axis=0), np.mean(uys_all, axis=0)
+            evol.energy, evol.energy_spread = np.mean(Es_all, axis=0), np.std(Es_all, axis=0)
+            evol.rel_energy_spread = evol.energy_spread/evol.energy
+            evol.beam_size_x, evol.beam_size_y = np.std(xs_all, axis=0), np.std(ys_all, axis=0)
+
+            evol.emit_nx, evol.emit_ny, evol.beta_x, evol.beta_y = np.zeros_like(ss), np.zeros_like(ss), np.zeros_like(ss), np.zeros_like(ss)
+            for i in range(len(ss)):
+                evol.emit_nx[i] = np.sqrt(np.linalg.det(np.cov(xs_all[:,i], uxs_all[:,i]/SI.c))) # only works for equal weight particles
+                evol.emit_ny[i] = np.sqrt(np.linalg.det(np.cov(ys_all[:,i], uys_all[:,i]/SI.c))) # only works for equal weight particles
+                covx, covy = np.cov(xs_all[:,i], uxs_all[:,i]/uzs_all[:,i]), np.cov(ys_all[:,i], uys_all[:,i]/uzs_all[:,i])
                 evol.beta_x[i] = covx[0,0]/np.sqrt(np.linalg.det(covx))
                 evol.beta_y[i] = covy[0,0]/np.sqrt(np.linalg.det(covy))
-
-            evol.rel_energy_spread = evol.energy_spread/evol.energy
-            evol.charge = self.charge()*np.ones(evol.location.shape)
-            evol.z = self.z_offset()*np.ones(evol.location.shape)
-            evol.bunch_length = self.bunch_length()*np.ones(evol.location.shape)
-            evol.plasma_density = n0*np.ones(evol.location.shape)
+
+            evol.z = self.z_offset()*np.ones_like(ss)
+            evol.charge = self.charge()*np.ones_like(ss)
+            evol.bunch_length = self.bunch_length()*np.ones_like(ss)
+            evol.plasma_density = n0*np.ones_like(ss)
 
 
         # calculate final positions and angles after betatron motion
@@ -1952,7 +1952,7 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat
         elif spin_tracking:
 
             # transverse motion evolution using Hill's equation (to get particle orbits and energy evolution)
-            from abel.physics_models.spin_tracking import evolve_hills_equation_analytic_evolution
+            from abel.physics_models.hills_equation import evolve_hills_equation_analytic_evolution
             xs_all, uxs_all, gammas_all, ss = evolve_hills_equation_analytic_evolution(self.xs()-x0_driver, self.uxs(), L, gamma0s, dgamma_ds, k_p(n0))
             ys_all, uys_all, _, _ = evolve_hills_equation_analytic_evolution(self.ys()-y0_driver, self.uys(), L, gamma0s, dgamma_ds, k_p(n0))
 

src/abel/classes/beamline/beamline.py

@@ -211,30 +211,48 @@ def plot_evolution_recursive(self, shot=None):
         # set up evolution data structure
         evol = SimpleNamespace()
         evol.location = np.array([])
-
+        location0 = 0.0
+        loc_lengths = np.array([])
+
         # first calculate the locations
         for trackable in self[shots[0]].trackables:
-            if hasattr(trackable, 'evolution') and hasattr(trackable.evolution, 'location'):
-                evol.location = np.append(evol.location, trackable.evolution.location)
-
+            if hasattr(trackable, 'evolution'):
+                if hasattr(trackable.evolution, 'beam'):
+                    track_evol = trackable.evolution.beam
+                else:
+                    track_evol = trackable.evolution
+                if hasattr(track_evol, 'location'):
+                    loc_lengths = np.append(loc_lengths, [len(track_evol.location)])
+                    evol.location = np.append(evol.location, track_evol.location - min(track_evol.location) + location0)
+                location0 = location0 + trackable.get_length()
+
         # declare the data structure to contain all the data
-        vals = np.zeros((8,len(shots),len(evol.location)))
-
+        vals = np.empty((8,len(shots),len(evol.location)))
+        vals.fill(np.nan)
+
         # extract all the data
         for i, shot in enumerate(shots):
             ind0 = 0
-            for trackable in self[shot].trackables:
-                if hasattr(trackable, 'evolution') and hasattr(trackable.evolution, 'location'):
-                    inds = range(ind0, ind0+len(np.array(trackable.evolution.location)))
+            for j, trackable in enumerate(self[shot].trackables):
+                if hasattr(trackable, 'evolution'):
+                    if hasattr(trackable.evolution, 'beam'):
+                        track_evol = trackable.evolution.beam
+                    else:
+                        track_evol = trackable.evolution
+                    if not hasattr(track_evol, 'location'):
+                        continue
+                    inds = range(ind0, ind0 + round(loc_lengths[j]))
                     ind0 = max(inds)+1
-                    vals[0,i,inds] = np.array(trackable.evolution.energy)
-                    vals[1,i,inds] = np.array(trackable.evolution.charge)
-                    vals[2,i,inds] = np.array(trackable.evolution.emit_nx)
-                    vals[3,i,inds] = np.array(trackable.evolution.emit_ny)
-                    vals[4,i,inds] = np.array(trackable.evolution.rel_energy_spread)
-                    vals[5,i,inds] = np.array(trackable.evolution.bunch_length)
-                    vals[6,i,inds] = np.array(trackable.evolution.beam_size_x)
-                    vals[7,i,inds] = np.array(trackable.evolution.beam_size_y)
+                    ss = evol.location[inds]-min(evol.location[inds])
+                    ss_track_evol = track_evol.location-min(track_evol.location)
+                    vals[0,i,inds] = np.interp(ss, ss_track_evol, track_evol.energy)
+                    vals[1,i,inds] = np.interp(ss, ss_track_evol, track_evol.charge)
+                    vals[2,i,inds] = np.interp(ss, ss_track_evol, track_evol.emit_nx)
+                    vals[3,i,inds] = np.interp(ss, ss_track_evol, track_evol.emit_ny)
+                    vals[4,i,inds] = np.interp(ss, ss_track_evol, track_evol.rel_energy_spread)
+                    vals[5,i,inds] = np.interp(ss, ss_track_evol, track_evol.bunch_length)
+                    vals[6,i,inds] = np.interp(ss, ss_track_evol, track_evol.beam_size_x)
+                    vals[7,i,inds] = np.interp(ss, ss_track_evol, track_evol.beam_size_y)
 
         # get the mean
         val_mean = np.mean(vals, axis=1)
@@ -306,7 +324,7 @@ def plot_evolution_recursive(self, shot=None):
         ax = fig.add_subplot(gs[2, :])
         ax.fill_between(evol.location, (evol.beam_size_x + evol_std.beam_size_x)*1e6, (evol.beam_size_x-evol_std.beam_size_x)*1e6, color=col1, alpha=af, ec=None)
         ax.plot(evol.location, evol.beam_size_x*1e6, color=col1, label=r'$\sigma_x$')
-        ax.fill_between(evol.location, (evol.beam_size_y + evol_std.beam_size_y)*1e6, (evol.beam_size_y-evol_std.beam_size_y)*1e6, color=col1, alpha=af, ec=None)
+        ax.fill_between(evol.location, (evol.beam_size_y + evol_std.beam_size_y)*1e6, (evol.beam_size_y-evol_std.beam_size_y)*1e6, color=col2, alpha=af, ec=None)
         ax.plot(evol.location, evol.beam_size_y*1e6, color=col2, label=r'$\sigma_y$')
         ax.set_xlim(long_limits)
         ax.set_ylabel(r'Beam size, rms [$\mathrm{\mu}$m]')
@@ -317,7 +335,7 @@ def plot_evolution_recursive(self, shot=None):
         ax = fig.add_subplot(gs[3, :])
         ax.plot(evol.location, abs(evol.emit_ny[0]) * np.ones(evol.location.shape) * 1e6, ':', color=col2)
         ax.plot(evol.location, abs(evol.emit_nx[0]) * np.ones(evol.location.shape) * 1e6, ':', color=col1)
-        ax.fill_between(evol.location, (evol.emit_ny + evol_std.emit_ny)*1e6, (evol.emit_ny-evol_std.emit_ny)*1e6, color=col1, alpha=af, ec=None)
+        ax.fill_between(evol.location, (evol.emit_ny + evol_std.emit_ny)*1e6, (evol.emit_ny-evol_std.emit_ny)*1e6, color=col2, alpha=af, ec=None)
         ax.plot(evol.location, evol.emit_ny*1e6, color=col2, label=r'$\epsilon_y$')
         ax.fill_between(evol.location, (evol.emit_nx + evol_std.emit_nx)*1e6, (evol.emit_nx-evol_std.emit_nx)*1e6, color=col1, alpha=af, ec=None)
         ax.plot(evol.location, evol.emit_nx*1e6, color=col1, label=r'$\epsilon_x$')

src/abel/classes/source/source.py

@@ -114,6 +114,8 @@ def __init__(self, length=0.0, charge=None, energy=None, accel_gradient=None, wa
         self.spin_polarization = spin_polarization
         self.spin_polarization_direction = spin_polarization_direction
 
+        self.evolution = SimpleNamespace()
+
 
     @abstractmethod
     def track(self, beam, savedepth=0, runnable=None, verbose=False):
@@ -174,11 +176,22 @@ def track(self, beam, savedepth=0, runnable=None, verbose=False):
 
         if do_beam_rotation:
             beam.add_pointing_tilts(rotation_angle_x, rotation_angle_y)
-
-            # if not np.isclose(beam.x_angle(), beam.x_tilt_angle(), rtol=1e-2, atol=0.0) or not np.isclose(beam.y_angle(), beam.y_tilt_angle(), rtol=1e-2, atol=0.0):
-            #     pointing_error_string = 'The beam may not have been accurately aligned to its propagation direction.\n' + 'beam x_angle: ' + str(beam.x_angle()) + '\nbeam x_tilt_angle: ' + str(beam.x_tilt_angle()) + '\nbeam y_angle: ' + str(beam.y_angle()) + '\nbeam y_tilt_angle: ' + str(beam.y_tilt_angle())
-            #     warnings.warn(pointing_error_string)
-
+
+        # save evolution (output)
+        self.evolution.location = np.array([self.get_length()])
+        self.evolution.x, self.evolution.ux = np.array([beam.x_offset()]), np.array([np.mean(beam.uxs())])
+        self.evolution.y, self.evolution.uy = np.array([beam.y_offset()]), np.array([np.mean(beam.uys())])
+        self.evolution.energy, self.evolution.energy_spread = np.array([beam.energy()]), np.array([beam.energy_spread()])
+        self.evolution.rel_energy_spread = self.evolution.energy_spread/self.evolution.energy
+        self.evolution.beam_size_x, self.evolution.beam_size_y = np.array([beam.beam_size_x()]), np.array([beam.beam_size_y()])
+        self.evolution.emit_nx = np.array([beam.norm_emittance_x()])
+        self.evolution.emit_ny = np.array([beam.norm_emittance_y()])
+        self.evolution.beta_x = np.array([beam.beta_x()])
+        self.evolution.beta_y = np.array([beam.beta_y()])
+        self.evolution.z = np.array([beam.z_offset()])
+        self.evolution.charge = np.array([beam.charge()])
+        self.evolution.bunch_length = np.array([beam.bunch_length()])
+
         # set metadata
         beam.location = 0
         beam.stage_number = 0

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

@@ -218,9 +218,9 @@ def main_tracking_procedure(self, beam0, driver0):
         driver_waket = wake_t_bunch2beam(bunches[0][-1])
 
         # save evolution of the beam and driver
-        if self.probe_evolution:
-            self._driver_evolution = wake_t.diagnostics.analyze_bunch_list(bunches[0])
-            self._beam_evolution = wake_t.diagnostics.analyze_bunch_list(bunches[1])
+        #if self.probe_evolution:
+        #    self._driver_evolution = wake_t.diagnostics.analyze_bunch_list(bunches[0])
+        #    self._beam_evolution = wake_t.diagnostics.analyze_bunch_list(bunches[1])
 
         # Store wakefield and beam quantities
         self.__save_initial_step(tmpfolder)

src/abel/physics_models/hills_equation.py

@@ -121,4 +121,42 @@ def evolve_hills_equation_analytic(x0, ux0, L, gamma0, dgamma_ds, kp):
     # convert angles back to proper velocities
     ux = xp * gamma2proper_velocity(gamma)
 
-    return x, ux
+    return x, ux
+
+
+def evolve_hills_equation_analytic_evolution(x0, ux0, L, gamma0, dgamma_ds, kp, num_steps=None):
+
+    import scipy.special as scispec
+    from abel.utilities.relativity import gamma2proper_velocity
+
+    # find number of steps
+    if num_steps is None:
+        beta_mat = np.sqrt(2*np.mean(gamma0))/kp
+        num_steps = round(L/beta_mat*10)
+
+    # define propagation axis
+    s = np.linspace(0, L, num_steps).T
+
+    # find initial angles
+    xp0 = ux0 / gamma2proper_velocity(gamma0)
+
+    # calculate energy evolution
+    gamma = np.outer(gamma0, np.ones_like(s)) + np.outer(dgamma_ds, s)
+
+    # calculate x evolution
+    C = np.sqrt(2) * kp / dgamma_ds
+    A0 = C * np.sqrt(gamma0)
+    Di = (kp**2 * x0 * scispec.iv(1, A0 * 1j) + A0 * 1j * dgamma_ds * xp0 * scispec.iv(0, A0 * 1j))
+    Dk = (kp**2 * x0 * scispec.kv(1, A0 * 1j) - A0 * 1j * dgamma_ds * xp0 * scispec.kv(0, A0 * 1j))
+    E = kp**2 * (scispec.iv(1, A0 * 1j) * scispec.kv(0, A0 * 1j) +
+                 scispec.iv(0, A0 * 1j) * scispec.kv(1, A0 * 1j)) * 1j
+    A = np.outer(C, np.ones_like(s)) * np.sqrt(gamma)
+    x = np.real(1j * (np.outer(Di,np.ones_like(s)) * scispec.kv(0, A * 1j) + np.outer(Dk,np.ones_like(s)) * scispec.iv(0, A * 1j)) / np.outer(E,np.ones_like(s)))
+
+    # calculate momentum evolution
+    K = dgamma_ds * C**2 / (2 * E)
+    xp = -np.real( np.outer(K,np.ones_like(s)) / A *
+                  (np.outer(Dk,np.ones_like(s)) * scispec.iv(1, A * 1j) - np.outer(Di,np.ones_like(s)) * scispec.kv(1, A * 1j)))
+    ux = xp * gamma2proper_velocity(gamma)
+
+    return x, ux, gamma, s

src/abel/physics_models/spin_tracking.py

@@ -12,46 +12,6 @@
 Simulates particle motion and tracks spin motion using the T-BMT equation
 """
 
-
-def evolve_hills_equation_analytic_evolution(x0, ux0, L, gamma0, dgamma_ds, kp=None, g=None, num_steps=None):
-
-    # set equivalent plasma density from magnetic field gradient
-    if kp is None:
-        kp = np.sqrt(2 * g * SI.e / (SI.m_e * SI.c))
-
-    # find number of steps
-    if num_steps is None:
-        beta_mat = np.sqrt(2*np.mean(gamma0))/kp
-        num_steps = round(L/beta_mat*10)
-
-    # define propagation axis
-    s = np.linspace(0, L, num_steps).T
-
-    # find initial angles
-    xp0 = ux0 / gamma2proper_velocity(gamma0)
-
-    # calculate energy evolution
-    gamma = np.outer(gamma0, np.ones_like(s)) + np.outer(dgamma_ds, s)
-
-    # calculate x evolution
-    C = np.sqrt(2) * kp / dgamma_ds
-    A0 = C * np.sqrt(gamma0)
-    Di = (kp**2 * x0 * scispec.iv(1, A0 * 1j) + A0 * 1j * dgamma_ds * xp0 * scispec.iv(0, A0 * 1j))
-    Dk = (kp**2 * x0 * scispec.kv(1, A0 * 1j) - A0 * 1j * dgamma_ds * xp0 * scispec.kv(0, A0 * 1j))
-    E = kp**2 * (scispec.iv(1, A0 * 1j) * scispec.kv(0, A0 * 1j) +
-                 scispec.iv(0, A0 * 1j) * scispec.kv(1, A0 * 1j)) * 1j
-    A = np.outer(C, np.ones_like(s)) * np.sqrt(gamma)
-    x = np.real(1j * (np.outer(Di,np.ones_like(s)) * scispec.kv(0, A * 1j) + np.outer(Dk,np.ones_like(s)) * scispec.iv(0, A * 1j)) / np.outer(E,np.ones_like(s)))
-
-    # calculate momentum evolution
-    K = dgamma_ds * C**2 / (2 * E)
-    xp = -np.real( np.outer(K,np.ones_like(s)) / A *
-                  (np.outer(Dk,np.ones_like(s)) * scispec.iv(1, A * 1j) - np.outer(Di,np.ones_like(s)) * scispec.kv(1, A * 1j)))
-    ux = xp * gamma2proper_velocity(gamma)
-
-    return x, ux, gamma, s
-
-
 def plasma_E_field(x, y, k_p, Ez=0.0):
     Er_coeff = SI.m_e * (k_p**2) * SI.c**2 / (2*SI.e)
     Ex = Er_coeff * x

tests/test_Beam.py

@@ -2389,13 +2389,13 @@ def test_apply_betatron_motion_emitt_pres():
 
     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(np.std(Es_final)/Es_final.mean(), 1e-20, rtol=0.0, atol=1e-15)
     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(Es_final.mean(), evol.energy[-1], rtol=1e-15, atol=0.0)
+    assert np.allclose(evol.energy_spread[-1], 1e-20, rtol=0.0, atol=1e-15)
     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)