diff --git a/src/abel/classes/beam.py b/src/abel/classes/beam.py index d01e28e4..067bc910 100644 --- a/src/abel/classes/beam.py +++ b/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)) diff --git a/src/abel/classes/beamline/beamline.py b/src/abel/classes/beamline/beamline.py index 6918a9ef..e820cb3d 100644 --- a/src/abel/classes/beamline/beamline.py +++ b/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$') diff --git a/src/abel/classes/source/source.py b/src/abel/classes/source/source.py index 1be6ae59..678d66ad 100644 --- a/src/abel/classes/source/source.py +++ b/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 diff --git a/src/abel/classes/stage/impl/stage_quasistatic_2d.py b/src/abel/classes/stage/impl/stage_quasistatic_2d.py index 0880a986..529d3cf8 100644 --- a/src/abel/classes/stage/impl/stage_quasistatic_2d.py +++ b/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) diff --git a/src/abel/physics_models/hills_equation.py b/src/abel/physics_models/hills_equation.py index b2153355..17ea4d49 100644 --- a/src/abel/physics_models/hills_equation.py +++ b/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 \ No newline at end of file + 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 diff --git a/src/abel/physics_models/spin_tracking.py b/src/abel/physics_models/spin_tracking.py index 5e72027a..5905c5e5 100644 --- a/src/abel/physics_models/spin_tracking.py +++ b/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 diff --git a/tests/test_Beam.py b/tests/test_Beam.py index a33f5425..91200408 100644 --- a/tests/test_Beam.py +++ b/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)