From 7ddd20471c21303bbdef4ba844d734afc82863ec Mon Sep 17 00:00:00 2001 From: Felipe Pena Date: Wed, 26 Mar 2025 17:02:24 +0200 Subject: [PATCH 01/11] Added an achromatic spectrometer class. ALso updated the plasma_lens_nonlinear_thin to accomodate for changing the axis (x or Y) of the nonlinearity --- abel/__init__.py | 2 + .../impl/plasma_lens_nonlinear_thin.py | 67 +++-- .../impl/spectrometer_achromatic_basic.py | 274 ++++++++++++++++++ .../achromatic/spectrometer_achromatic.py | 109 +++++++ 4 files changed, 433 insertions(+), 19 deletions(-) create mode 100644 abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py create mode 100644 abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py diff --git a/abel/__init__.py b/abel/__init__.py index f5aaad51..60719823 100644 --- a/abel/__init__.py +++ b/abel/__init__.py @@ -71,6 +71,8 @@ from .classes.bds.impl.bds_fbt import * from .classes.spectrometer.impl.spectrometer_facet_ocelot import * from .classes.spectrometer.impl.spectrometer_basic_clear import * +from .classes.spectrometer.impl.achromatic.spectrometer_achromatic import SpectrometerAchromatic +from .classes.spectrometer.impl.achromatic.impl.spectrometer_achromatic_basic import SpectrometerAchromaticBasic from .classes.ip.impl.ip_basic import * from .classes.ip.impl.ip_guineapig import * diff --git a/abel/classes/plasma_lens/impl/plasma_lens_nonlinear_thin.py b/abel/classes/plasma_lens/impl/plasma_lens_nonlinear_thin.py index a0a996dc..9aa5ccb9 100644 --- a/abel/classes/plasma_lens/impl/plasma_lens_nonlinear_thin.py +++ b/abel/classes/plasma_lens/impl/plasma_lens_nonlinear_thin.py @@ -2,28 +2,52 @@ from matplotlib import patches import numpy as np import scipy.constants as SI +import matplotlib.pyplot as plt class PlasmaLensNonlinearThin(PlasmaLens): - - def __init__(self, length=None, radius=None, current=None, rel_nonlinearity=0): + + def __init__(self, length=None, radius=None, current=None, rel_nonlinearity=0, nonlinearity_in_x=True): super().__init__(length, radius, current) - + # set nonlinearity (defined as R/Dx) self.rel_nonlinearity = rel_nonlinearity - - + self.nonlinearity_in_x = nonlinearity_in_x + + def remove_charge_outside(self, beam, plot=False): + + mask = np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius + + if plot: + + fig, ax = plt.subplots() + ax.add_patch(patches.Circle((self.offset_x, self.offset_y), self.radius, fill=False, linestyle='--')) + + ax.scatter(beam.xs(), beam.ys(), s=1, alpha=0.1, zorder=0) + ax.set_aspect('equal', 'box') + plt.show() + + rel_loss = mask.sum() / mask.size + print(f"Removing {rel_loss*100:.2f}% of the charge outside the lens") + + # remove charge outside the lens + del beam[mask] + + return beam + def track(self, beam, savedepth=0, runnable=None, verbose=False): # remove charge outside the lens (start) - del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius] - + # del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius] + self.remove_charge_outside(beam, plot=False) + # drift half the distance beam.transport(self.length/2) - + # remove charge outside the lens (middle) - del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius] - + # del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius] + self.remove_charge_outside(beam) + # get particles xs = beam.xs()-self.offset_x xps = beam.xps() @@ -34,16 +58,21 @@ def track(self, beam, savedepth=0, runnable=None, verbose=False): g0 = self.get_focusing_gradient() # calculate the nonlinearity - inv_Dx = self.rel_nonlinearity/self.radius - + if self.nonlinearity_in_x: + inv_Dx = self.rel_nonlinearity/self.radius + inv_Dy = 0.0 + else: + inv_Dx = 0.0 + inv_Dy = self.rel_nonlinearity/self.radius + # thin lens kick - Bx = g0*(ys + xs*ys*inv_Dx) - By = -g0*(xs + (xs**2 + ys**2)/2*inv_Dx) + Bx = g0*(ys + xs*ys*inv_Dx + (xs**2 + ys**2)/2*inv_Dy) + By = -g0*(xs + (xs**2 + ys**2)/2*inv_Dx + xs*ys*inv_Dy) # calculate the angular kicks delta_xp = self.length*(By*beam.charge_sign()*SI.c/beam.Es()) delta_yp = -self.length*(Bx*beam.charge_sign()*SI.c/beam.Es()) - + # set new beam positions and angles (shift back plasma-lens offsets) beam.set_xps(xps + delta_xp) beam.set_yps(yps + delta_yp) @@ -52,11 +81,11 @@ def track(self, beam, savedepth=0, runnable=None, verbose=False): beam.transport(self.length/2) # remove charge outside the lens (end) - del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius] - - return super().track(beam, savedepth, runnable, verbose) + # del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius] + self.remove_charge_outside(beam) + + return super().track(beam, savedepth, runnable, verbose) def get_focusing_gradient(self): return SI.mu_0 * self.current / (2*np.pi * self.radius**2) - \ No newline at end of file diff --git a/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py b/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py new file mode 100644 index 00000000..33e132a5 --- /dev/null +++ b/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py @@ -0,0 +1,274 @@ +""" +Module: spectrometer_achromatic_basic + +This module defines the SpectrometerAchromaticBasic class, which is a subclass of SpectrometerAchromatic. +It provides methods to initialize the spectrometer and track the beam. + +Classes +------- +SpectrometerAchromaticBasic + +Imports +------- +from abel import SpectrometerAchromatic, PlasmaLensNonlinearThin +import scipy.constants as SI +import numpy as np +""" + +from abel import SpectrometerAchromatic, PlasmaLensNonlinearThin +import scipy.constants as SI +import numpy as np + +class SpectrometerAchromaticBasic(SpectrometerAchromatic): + """ + A basic implementation of an achromatic spectrometer. + + Parameters + ---------- + field_dipole : float, optional + The magnetic field strength of the dipole (default is 1.0). + length_dipole : float, optional + The length of the dipole (default is 0.2). + length_drift_to_dipole : float, optional + The drift before the dipole (default is 0.3). + length_drift_dipole_to_lens : float, optional + The drift before the plasma lens (default is 0.2). + length_drift_lens_to_screen : float, optional + The drift before the screen (default is 0.3). + """ + + def __init__(self, + length_drift_to_dipole=0.2, # [m] + field_dipole=-0.85, # [T] + length_dipole=0.15, # [m] + length_drift_dipole_to_lens=0.15, # [m] + length_plasma_lens=20E-3, # [m] + radius_plasma_lens=1E-3, # [m] + length_drift_lens_to_screen=0.4, # [m] + imaging_energy=None, # [eV] + disable_lens_nonlinearity=True # [bool] + ): + """ + Initializes the SpectrometerAchromaticBasic with the given parameters. + + Parameters + ---------- + length_drift_to_dipole : float, optional + Length of the drift space to the dipole magnet in meters. Default is 0.2. + field_dipole : float, optional + Magnetic field strength of the dipole magnet in Tesla. + Negative means bending the beam downwards. + Default is -0.85. + length_dipole : float, optional + Length of the dipole magnet in meters. Default is 0.8. + length_drift_dipole_to_lens : float, optional + Length of the drift space from the dipole magnet to the lens in meters. Default is 0.2. + length_plasma_lens : float, optional + Length of the plasma lens in meters. Default is 20E-3. + radius_plasma_lens : float, optional + Radius of the plasma lens in meters. Default is 1E-3. + length_drift_lens_to_screen : float, optional + Length of the drift space from the lens to the screen in meters. Default is 0.4. + imaging_energy : float or None, optional + Imaging energy in electron volts. Default is None. + disable_lens_nonlinearity : bool, optional + Flag to disable lens nonlinearity. Default is True. + """ + + super().__init__(length_drift_to_dipole=length_drift_to_dipole, + field_dipole=field_dipole, + length_dipole=length_dipole, + length_drift_dipole_to_lens=length_drift_dipole_to_lens, + length_plasma_lens=length_plasma_lens, + radius_plasma_lens=radius_plasma_lens, + length_drift_lens_to_screen=length_drift_lens_to_screen, + imaging_energy=imaging_energy, + disable_lens_nonlinearity=disable_lens_nonlinearity) + + def track(self, beam0, savedepth=0, runnable=None, verbose=False): + """ + Tracks the beam through the spectrometer. + + Parameters + ---------- + beam0 : Beam + The initial beam to be tracked. + savedepth : int, optional + The depth at which to save the beam state (default is 0). + runnable : callable, optional + A callable to be run during tracking. + verbose : bool, optional + Flag to indicate if verbose output is desired (default is False). + + Returns + ------- + Beam + The tracked beam. + """ + + # transport the beam to the center of the dipole + beam0.transport(self.length_drift_to_dipole + self.length_dipole/2) + + # this function allows to plot the trace space for debugging + # beam0.plot_trace_space_y() + + # computing the bending angle of the dipole for each particle + # sin(bend_angle/2) = length_dipole*0.5 / bending_radius + # bending_radius = momentum / (B * e) + # for small angles sin(bend_angle/2) \approx bend_angle/2 + # bend_angle = length_dipole * B * e / momentum + # momentum [eV/c] + bend_angles = self.get_bending_angle(beam0.Es()) + # bend_angles = self.length_dipole * self.field_dipole * SI.c / beam0.Es() + + # if self.imaging_energy is None: + # energy_nom = beam0.energy() + # else: + energy_nom = beam0.energy() + + # the bending angle for the nominal beam energy + bend_angle_nom = self.get_bending_angle(energy_nom) + # bend_angle_nom = self.length_dipole * self.field_dipole * SI.c / energy_nom + + # the bending angle of each particle relative to the bending of + # the beam at nominal energy + bend_angles_rel = bend_angles - bend_angle_nom + + # updating the angles of the particles with the dipole kick + beam0.set_yps(beam0.yps() + bend_angles_rel) + + # transport the beam to the lens with the new angles of the particles + beam0.transport(self.length_dipole/2 + self.length_drift_dipole_to_lens) + + # calculate the focal length + length_object_to_lens = (self.length_drift_to_dipole + + self.length_dipole + + self.length_drift_dipole_to_lens + + self.length_plasma_lens/2) + length_lens_to_screen = (self.length_plasma_lens/2 + + self.length_drift_lens_to_screen) + focal_length = 1 / (1 / length_object_to_lens + + 1 / length_lens_to_screen) + + # calculate the lens strength based on the focal length + k = 1 / (focal_length * self.length_plasma_lens) + + if self.imaging_energy is None: + imaging_energy = beam0.energy() + else: + imaging_energy = self.imaging_energy + + # calculate the magnetic field gradient with the imaging energy + g = k / (SI.c/ imaging_energy) + + # calculate the current based on the magnetic field gradient + current = -g * 2 * np.pi * self.radius_plasma_lens**2 / (SI.mu_0) + + # calculate the dispersion introduced by the dipole at the position of + # the plasma lens + length_middle_dipole_to_lens = (self.length_dipole/2 + + self.length_drift_dipole_to_lens + + self.length_plasma_lens/2) + dispersion = self.get_dispersion(bend_angle=bend_angle_nom, + length=length_middle_dipole_to_lens) + + # calculate the nonlinearity of the plasma lens, which is + # defined by R/Dy (negative sign found by trial and error) + if not self.disable_lens_nonlinearity: + rel_nonlinearity = - self.radius_plasma_lens / dispersion + else: + rel_nonlinearity = 0.0 + + # use the already implemented PlasmaLensNonlinearThin class to calculate + # the kick produced by the plasma lens. The kicks are updated in the + # class beam + # nonlinearity_in_x=False as the dipole here disperses vertically + plasma_lens = PlasmaLensNonlinearThin(length=self.length_plasma_lens, + radius=self.radius_plasma_lens, + current=current, + rel_nonlinearity=rel_nonlinearity, + nonlinearity_in_x=False) + + beam = plasma_lens.track(beam=beam0) + + # transport the beam with the updated angles + length_middle_lens_to_screen = (self.length_plasma_lens/2 + + self.length_drift_lens_to_screen) + beam.transport(length_middle_lens_to_screen) + + # shift beam by the dispersion (special case for spectrometers) + # this will allow to plot the beam as if there was a screen + length_middle_dipole_to_screen = (self.length_dipole / 2 + + self.length_drift_dipole_to_lens + + self.length_plasma_lens + + self.length_drift_lens_to_screen) + # dispersion_on_screen = bend_angle_nom * length_middle_dipole_to_screen + dispersion_on_screen = self.get_dispersion(bend_angle=bend_angle_nom, + length= length_middle_dipole_to_screen) + beam.set_ys(beam.ys() + dispersion_on_screen) + + return super().track(beam, savedepth, runnable, verbose) + + def get_bending_angle(self, energy): + """ + Calculate the bending angle of the spectrometer for a given energy. + + Parameters: + ----------- + energy : float, optional + The energy of the particle beam in eV. + + Returns: + -------- + bend_angle : float + The calculated bending angle. + """ + + # computing the bending angle of the dipole for each particle + # sin(bend_angle/2) = length_dipole*0.5 / bending_radius + # bending_radius = momentum / (B * e) + # for small angles sin(bend_angle/2) \approx bend_angle/2 + # bend_angle = length_dipole * B * e / momentum + # momentum [eV/c] + bend_angle = self.length_dipole * self.field_dipole * SI.c / energy + + return bend_angle + + def get_dispersion(self, energy=None, bend_angle=None, length=None): + """ + Calculate the dispersion of the spectrometer for a given energy. + Either energy or bending angle (bend_angle) have to be given. + Parameters: + ----------- + energy : float, optional + The energy of the particle beam in eV. + bend_angle : float, optional + Bending angle of the dipole in rad. + length : float, optional + The distance to use for the calculation. If not provided, + it will be calculated as the sum of the lengths of the dipole/2, drift space + from dipole to lens, plasma lens, and drift space from lens to screen, + i.e., the distance from the center of the dipole to the screen. + Returns: + -------- + dispersion : float + The calculated dispersion value. + """ + + text = 'Please give either the bending angle *or* the energy.' + assert ((bend_angle is None and energy is not None) + or (bend_angle is not None and energy is None)), text + + if bend_angle is None: + bend_angle = self.get_bending_angle(energy) + + if length is None: + # assume lenth from center of dipole to the screen + length = (self.length_dipole / 2 + + self.length_drift_dipole_to_lens + + self.length_plasma_lens + + self.length_drift_lens_to_screen) + + dispersion = bend_angle * length + + return dispersion diff --git a/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py b/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py new file mode 100644 index 00000000..2595592f --- /dev/null +++ b/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py @@ -0,0 +1,109 @@ +from abc import abstractmethod +from abel import Spectrometer +import numpy as np +import matplotlib.pyplot as plt +from abel import CONFIG + + +class SpectrometerAchromatic(Spectrometer): + """ + Achromatic Spectrometer class. + + Parameters + ---------- + length_drift_to_dipole : float, optional + Length of the drift space to the dipole magnet in meters. Default is 0.2. + field_dipole : float, optional + Magnetic field strength of the dipole magnet in Tesla. + Negative means bending the beam downwards. + Default is -0.85. + length_dipole : float, optional + Length of the dipole magnet in meters. Default is 0.8. + length_drift_dipole_to_lens : float, optional + Length of the drift space from the dipole magnet to the lens in meters. Default is 0.2. + length_plasma_lens : float, optional + Length of the plasma lens in meters. Default is 20E-3. + radius_plasma_lens : float, optional + Radius of the plasma lens in meters. Default is 1E-3. + length_drift_lens_to_screen : float, optional + Length of the drift space from the lens to the screen in meters. Default is 0.4. + imaging_energy : float or None, optional + Imaging energy in electron volts. Default is None. + disable_lens_nonlinearity : bool, optional + Flag to disable lens nonlinearity. Default is True. + """ + + @abstractmethod + def __init__(self, + length_drift_to_dipole, + field_dipole, + length_dipole, + length_drift_dipole_to_lens, + length_plasma_lens, + radius_plasma_lens, + length_drift_lens_to_screen, + imaging_energy, + disable_lens_nonlinearity): + + self.length_drift_to_dipole = length_drift_to_dipole # [m] + self.field_dipole = field_dipole # [T] + self.length_dipole = length_dipole # [m] + + self.length_drift_dipole_to_lens = length_drift_dipole_to_lens # [m] + self.length_plasma_lens = length_plasma_lens # [m] + + self.imaging_energy = imaging_energy # [eV] + self.length_drift_lens_to_screen = length_drift_lens_to_screen # [m] + + self._radius_plasma_lens = radius_plasma_lens # [m] + self.disable_lens_nonlinearity = disable_lens_nonlinearity # [bool] + + super().__init__() + + def get_length(self): + """ + Calculate the total length of the spectrometer lattice. + + Returns + ------- + float + Total length of the spectrometer lattice in meters. + """ + total_length = (self.length_drift_to_dipole + + self.length_dipole + + self.length_drift_dipole_to_lens + + self.length_plasma_lens + + self.length_drift_lens_to_screen) + return total_length + + @property + def radius_plasma_lens(self): + return self._radius_plasma_lens + + @radius_plasma_lens.setter + def radius_plasma_lens(self, value): + assert value >= 0, 'Radius of the plasma lens must be greater than or equal to zero.' + self._radius_plasma_lens = value + + @abstractmethod + def track(self, beam, savedepth=0, runnable=None, verbose=False): + """ + Track the beam through the spectrometer. + + Parameters + ---------- + beam : object + The beam object to be tracked. + savedepth : int, optional + Depth of saved states during tracking. Default is 0. + runnable : object or None, optional + Runnable object for tracking. Default is None. + verbose : bool, optional + Flag to enable verbose output. Default is False. + + Returns + ------- + object + Result of the tracking process. + """ + return super().track(beam, savedepth, runnable, verbose) From aff853cec58736337b00b237d2aba2f4bcca5b75 Mon Sep 17 00:00:00 2001 From: Felipe Pena Date: Wed, 26 Mar 2025 17:04:13 +0200 Subject: [PATCH 02/11] Added a new source that produces bunches where the particles have discrete values in energy, and these values are defined by the user. No discretization in energy is also possible, where the energies are sampled from a uniform distribution. --- .../classes/source/impl/source_energy_comb.py | 177 ++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 abel/classes/source/impl/source_energy_comb.py diff --git a/abel/classes/source/impl/source_energy_comb.py b/abel/classes/source/impl/source_energy_comb.py new file mode 100644 index 00000000..8ada5f34 --- /dev/null +++ b/abel/classes/source/impl/source_energy_comb.py @@ -0,0 +1,177 @@ +import time +import numpy as np +import scipy.constants as SI +from abel import Source, Beam +from abel.utilities.beam_physics import generate_trace_space_xy, generate_symm_trace_space_xyz +from abel.utilities.relativity import energy2gamma + +class SourceEnergyComb(Source): + # this class behaves like the SourceBasic but has the added functionality + # of distributing the particle's energies within a given list of energies + + def __init__(self, length=0, num_particles=1000, energy=None, charge=0, rel_energy_spread=None, energy_spread=None, bunch_length=None, z_offset=0, x_offset=0, y_offset=0, x_angle=0, y_angle=0, emit_nx=0, emit_ny=0, beta_x=None, beta_y=None, alpha_x=0, alpha_y=0, angular_momentum=0, wallplug_efficiency=1, accel_gradient=None, symmetrize=False, symmetrize_6d=False, z_cutoff=None, energy_comb_delta=None): + + super().__init__(length, charge, energy, accel_gradient, wallplug_efficiency, x_offset, y_offset, x_angle, y_angle) + + self.rel_energy_spread = rel_energy_spread # [eV] + self.energy_spread = energy_spread + self.bunch_length = bunch_length # [m] + self.z_offset = z_offset # [m] + self.num_particles = num_particles + + self.emit_nx = emit_nx # [m rad] + self.emit_ny = emit_ny # [m rad] + self.beta_x = beta_x # [m] + self.beta_y = beta_y # [m] + self.alpha_x = alpha_x # [m] + self.alpha_y = alpha_y # [m] + + self.angular_momentum = angular_momentum + self.symmetrize = symmetrize + self.symmetrize_6d = symmetrize_6d + self.z_cutoff = z_cutoff + + # TODO: Make this an energy_diff or delta instead of a linspace + # the diff should be in percentage + # this will be easier for scans + self.energy_comb_delta = energy_comb_delta + + + def track(self, _=None, savedepth=0, runnable=None, verbose=False): + + # make empty beam + beam = Beam() + + # Lorentz gamma + gamma = energy2gamma(self.energy) + + # generate relative/absolute energy spreads + if self.rel_energy_spread is not None: + if self.energy_spread is None: + self.energy_spread = self.energy * self.rel_energy_spread + elif abs(self.energy_spread - self.energy * self.rel_energy_spread) > 0: + raise Exception("Both absolute and relative energy spread defined.") + + if self.symmetrize_6d is False: + + # horizontal and vertical phase spaces + xs, xps, ys, yps = generate_trace_space_xy(self.emit_nx/gamma, self.beta_x, self.alpha_x, self.emit_ny/gamma, self.beta_y, self.alpha_y, self.num_particles, self.angular_momentum/gamma, symmetrize=self.symmetrize) + + # longitudinal phase space + if self.symmetrize: + num_tiling = 4 + num_particles_actual = round(self.num_particles/num_tiling) + else: + num_particles_actual = self.num_particles + zs = np.random.normal(loc=self.z_offset, scale=self.bunch_length, size=num_particles_actual) + # Es = np.random.normal(loc=self.energy, scale=self.energy_spread, size=num_particles_actual) + # this should yield a uniform energy distribution with length 2*energy_spread around self.energy + Es = (np.random.random(size=num_particles_actual) - 0.5) * 2 * self.energy_spread + self.energy + if self.symmetrize: + zs = np.tile(zs, num_tiling) + Es = np.tile(Es, num_tiling) + + else: + xs, xps, ys, yps, zs, Es = generate_symm_trace_space_xyz(self.emit_nx/gamma, self.beta_x, self.alpha_x, self.emit_ny/gamma, self.beta_y, self.alpha_y, self.num_particles, self.bunch_length, self.energy_spread, self.angular_momentum/gamma) + + # add longitudinal offsets + zs += self.z_offset + Es += self.energy + + # create phase space + beam.set_phase_space(xs=xs, ys=ys, zs=zs, xps=xps, yps=yps, Es=Es, Q=self.charge) + + # Apply filter(s) if desired + if self.z_cutoff is not None: + beam = self.z_filter(beam) + + if self.energy_comb_delta is not None: + self.reorder_particle_energies(beam) + + # add jitters and offsets in super function + return super().track(beam, savedepth, runnable, verbose) + + def reorder_particle_energies(self, beam, plot=False): + + # takes in the beam and uses the attribute self.energy_comb_delta + # to create a list of energies that will be given to each beam particle + # to create a comb-like energy distribution for resolution estimates + + # the energies of all bunch particles + vals = beam.Es() + + energy_beam = beam.energy() + + # the energy delta is given as relative, thus multiply times the average + # particle energy + energy_delta = self.energy_comb_delta * energy_beam + + # half of the number of steps depending on the step size and the energy + # range of the beam particles + num_steps_half = int((vals.max() - vals.min()) * 0.5 + / energy_delta + 1) + + # go half the number of steps in both directions, with the bunch + # average energy in the middle + energy_steps = np.arange(-num_steps_half, num_steps_half+1) + energies_comb = energy_beam + energy_steps * energy_delta + + # find the index for each value that shows the closest value in energies + # this function requires "left" or "right" to be specified + # use "left" + idx = np.searchsorted(energies_comb, vals, side="left") + + # this is required as otherwise the energies_comb[idx] in the next line + # will throw an error + idx[idx == len(energies_comb)] = len(energies_comb) - 1 + + # find out if the value with the left-index or the value with the + # right-index value is closer and accordingly subtract the index value + idx = idx - (np.abs(vals - energies_comb[idx-1]) + < np.abs(vals - energies_comb[idx])) + + # energies_comb[idx] has all the energies that were closest for each + # particle. Set the particle energies to this new energies + beam.set_Es(energies_comb[idx]) + + if plot: + # print(np.diff(energies_comb)*1E-9) + beam.plot_lps() + + return + + + # ================================================== + # Filter out particles whose z < z_cutoff for testing instability etc. + def z_filter(self, beam): + xs = beam.xs() + ys = beam.ys() + zs = beam.zs() + pxs = beam.pxs() + pys = beam.pys() + pzs = beam.pzs() + weights = beam.weightings() + + # Apply the filter + bool_indices = (zs > self.z_cutoff) + zs_filtered = zs[bool_indices] + xs_filtered = xs[bool_indices] + ys_filtered = ys[bool_indices] + pxs_filtered = pxs[bool_indices] + pys_filtered = pys[bool_indices] + pzs_filtered = pzs[bool_indices] + weights_filtered = weights[bool_indices] + + # Initialise ABEL Beam object + beam_out = Beam() + + # Set the phase space of the ABEL beam + beam_out.set_phase_space(Q=np.sum(weights_filtered)*-SI.e, + xs=xs_filtered, + ys=ys_filtered, + zs=zs_filtered, + pxs=pxs_filtered, # Always use single particle momenta? + pys=pys_filtered, + pzs=pzs_filtered) + + return beam_out From 7a01c697d0c2ed5f853b4b8064852967bfcecccf Mon Sep 17 00:00:00 2001 From: Felipe Pena Date: Wed, 26 Mar 2025 17:12:19 +0200 Subject: [PATCH 03/11] Added a property 'l' for convenience that returns the length of the plasma lens --- abel/classes/plasma_lens/plasma_lens.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/abel/classes/plasma_lens/plasma_lens.py b/abel/classes/plasma_lens/plasma_lens.py index 26d55134..cc36c92e 100644 --- a/abel/classes/plasma_lens/plasma_lens.py +++ b/abel/classes/plasma_lens/plasma_lens.py @@ -4,12 +4,12 @@ import numpy as np class PlasmaLens(Trackable): - + @abstractmethod def __init__(self, length, radius, current, offset_x=0, offset_y=0): super().__init__() - + # common variables self.length = length self.radius = radius @@ -17,8 +17,8 @@ def __init__(self, length, radius, current, offset_x=0, offset_y=0): self.offset_x = offset_x self.offset_y = offset_y - - + + @abstractmethod def track(self, beam, savedepth=0, runnable=None, verbose=False): return super().track(beam, savedepth, runnable, verbose) @@ -26,10 +26,13 @@ def track(self, beam, savedepth=0, runnable=None, verbose=False): @abstractmethod def get_focusing_gradient(self): pass - + def get_length(self): return self.length def survey_object(self): return patches.Rectangle((0, -1), self.get_length(), 2) - \ No newline at end of file + + @property + def l(self): + return self.length From f8b2864d6627e7a3e2450440febe1b4be67defbe Mon Sep 17 00:00:00 2001 From: Felipe Pena Date: Wed, 26 Mar 2025 17:17:08 +0200 Subject: [PATCH 04/11] Added the possibility to edit the plot in plot_transverse_profile --- abel/classes/beam.py | 515 ++++++++++++++++++++++--------------------- 1 file changed, 259 insertions(+), 256 deletions(-) diff --git a/abel/classes/beam.py b/abel/classes/beam.py index d1d72215..5130c434 100644 --- a/abel/classes/beam.py +++ b/abel/classes/beam.py @@ -18,7 +18,7 @@ from matplotlib import pyplot as plt class Beam(): - + def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1, bunch_separation=0.0): # the phase space variable is private @@ -30,39 +30,39 @@ def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1, # bunch pattern information self.num_bunches_in_train = num_bunches_in_train self.bunch_separation = bunch_separation # [s] - + self.trackable_number = -1 # will increase to 0 after first tracking element self.stage_number = 0 - self.location = 0 - - + self.location = 0 + + # reset phase space def reset_phase_space(self, num_particles): self.__phasespace = np.zeros((8, num_particles)) - + # filter out macroparticles based on a mask (true means delete) def __delitem__(self, indices): if hasattr(indices, 'len'): if len(indices) == len(self): indices = np.where(indices) self.__phasespace = np.ascontiguousarray(np.delete(self.__phasespace, indices, 1)) - + # filter out nans def remove_nans(self): del self[np.isnan(self).any(axis=1)] - + # 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): - + # make empty phase space num_particles = len(xs) self.reset_phase_space(num_particles) - + # add positions self.set_xs(xs) self.set_ys(ys) self.set_zs(zs) - + # add momenta if uzs is None: if pzs is not None: @@ -70,57 +70,57 @@ def set_phase_space(self, Q, xs, ys, zs, uxs=None, uys=None, uzs=None, pxs=None, elif Es is not None: uzs = energy2proper_velocity(Es) self.__phasespace[5,:] = uzs - + if uxs is None: if pxs is not None: uxs = momentum2proper_velocity(pxs) elif xps is not None: uxs = xps * uzs self.__phasespace[3,:] = uxs - + if uys is None: if pys is not None: uys = momentum2proper_velocity(pys) elif yps is not None: uys = yps * uzs self.__phasespace[4,:] = uys - + # charge if weightings is None: self.__phasespace[6,:] = Q/num_particles else: self.__phasespace[6,:] = Q*weightings/np.sum(weightings) - + # ids self.__phasespace[7,:] = np.arange(num_particles) # single particle mass [kg] self.particle_mass = particle_mass - - + + # addition operator (add two beams using the + operator) def __add__(self, beam): return Beam(phasespace = np.append(self.__phasespace, beam.__phasespace, axis=1)) - + # in-place addition operator (add one beam to another using the += operator) def __iadd__(self, beam): - if beam is not None: + if beam is not None: self.__phasespace = np.append(self.__phasespace, beam.__phasespace, axis=1) return self - + # indexing operator (get single particle out) def __getitem__(self, index): return self.__phasespace[:,index] - + # "length" operator (number of macroparticles) def __len__(self): return self.__phasespace.shape[1] - + # string operator (called when printing) def __str__(self): return f"Beam: {len(self)} macroparticles, {self.charge()*1e9:.2f} nC, {self.energy()/1e9:.2f} GeV" - - + + ## BUNCH PATTERN def bunch_frequency(self) -> float: @@ -138,11 +138,11 @@ def train_duration(self) -> float: return None else: return self.bunch_separation * (self.num_bunches_in_train-1) - + def average_current_train(self) -> float: return self.charge()*self.bunch_frequency() - + ## BEAM ARRAYS # get phase space variables @@ -162,7 +162,7 @@ def qs(self): return self.__phasespace[6,:] def ids(self): return self.__phasespace[7,:] - + # set phase space variables def set_xs(self, xs): self.__phasespace[0,:] = xs @@ -176,22 +176,22 @@ def set_uys(self, uys): self.__phasespace[4,:] = uys def set_uzs(self, uzs): self.__phasespace[5,:] = uzs - + def set_xps(self, xps): self.set_uxs(xps*self.uzs()) def set_yps(self, yps): self.set_uys(yps*self.uzs()) def set_Es(self, Es): self.set_uzs(energy2proper_velocity(Es)) - + def set_qs(self, qs): self.__phasespace[6,:] = qs def set_ids(self, ids): self.__phasespace[7,:] = ids - + def weightings(self): return self.__phasespace[6,:]/(self.charge_sign()*SI.e) - + # copy another beam's macroparticle charge def copy_particle_charge(self, beam): self.set_qs(np.median(beam.qs())) @@ -201,18 +201,18 @@ def scale_charge(self, Q): def scale_energy(self, E): self.set_Es((E/self.energy())*self.Es()) - - + + def rs(self): return np.sqrt(self.xs()**2 + self.ys()**2) - + def pxs(self): return proper_velocity2momentum(self.uxs()) def pys(self): return proper_velocity2momentum(self.uys()) def pzs(self): return proper_velocity2momentum(self.uzs()) - + def xps(self): return self.uxs()/self.uzs() def yps(self): @@ -226,11 +226,11 @@ def deltas(self, pz0=None): if pz0 is None: pz0 = np.mean(self.pzs()) return self.pzs()/pz0 -1 - - + + def ts(self): return self.zs()/SI.c - + # vector of transverse positions and angles: (x, x', y, y') def transverse_vector(self): vector = np.zeros((4,len(self))) @@ -239,13 +239,13 @@ def transverse_vector(self): vector[2,:] = self.ys() vector[3,:] = self.yps() return vector - + # set phase space based on transverse vector: (x, x', y, y') def set_transverse_vector(self, vector): self.set_xs(vector[0,:]) self.set_xps(vector[1,:]) self.set_ys(vector[2,:]) - self.set_yps(vector[3,:]) + self.set_yps(vector[3,:]) def norm_transverse_vector(self): vector = np.zeros((4,len(self))) @@ -261,18 +261,18 @@ def norm_transverse_vector(self): 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): """ Rotates the coordinate system (passive transformation) of the beam first with angle1 around axis1, then with angle2 around axis2 and lastly with angle3 around axis3. - + Parameters ---------- axis1, axis2, axis3: 1x3 float nd arrays Unit vectors specifying the rotation axes. - + angle1, angle2, angle3: [rad] float Angles used for rotation of the beam's coordinate system around the respective axes. invert: bool Performs a standard passive transformation when False. If True, will perform an active transformation and can thus be used to invert the passive transformation. - + Returns ---------- Modified beam xs, ys, zs, uxs, uys and uzs. @@ -284,7 +284,7 @@ def rotate_coord_sys_3D(self, axis1, angle1, axis2=np.array([0, 1, 0]), angle2=0 if angle1 < -np.pi or angle1 > np.pi or angle2 < -np.pi or angle2 > np.pi or angle3 < -np.pi or angle3 > np.pi: raise ValueError('The rotation angles have to be in the interval [-pi, pi].') - + # Combine into (N, 3) arrays zs = self.zs() xs = self.xs() @@ -317,23 +317,23 @@ def rotate_coord_sys_3D(self, axis1, angle1, axis2=np.array([0, 1, 0]), angle2=0 self.set_uzs(rotated_uzs) self.set_uxs(rotated_uxs) self.set_uys(rotated_uys) - + # ================================================== def beam_alignment_angles(self): """ Calculates the angles for rotation around the y- and x-axis to align the z-axis to the beam proper velocity. - + Parameters ---------- ... - - + + Returns ---------- x_angle: [rad] float Used for rotating the beam's frame around the y-axis. - + y_angle: [rad] float Used for rotating the beam's frame around the x-axis. 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(). """ @@ -348,7 +348,7 @@ def beam_alignment_angles(self): # Calculate the angles to be used for beam rotation z_axis = np.array([1, 0, 0]) # Axis as an unit vector - + zx_projection = point_vec * np.array([1, 1, 0]) # The projection of the pointing vector onto the zx-plane. # Separate treatments for small angles to avoid numerical instability @@ -357,39 +357,39 @@ def beam_alignment_angles(self): else: x_angle = np.sign(point_vec[1]) * np.arccos( np.dot(zx_projection, z_axis) / np.linalg.norm(zx_projection) ) # The angle between zx_projection and z_axis. - + if np.abs(self.y_angle()) < 1e-4: y_angle = point_vec[2]/np.linalg.norm(zx_projection) else: rotated_zy_projection = np.array([ np.linalg.norm(zx_projection), 0, point_vec[2] ]) # The new pointing vector after its zx-projection has been aligned to the rotated z-axis. - + y_angle = np.sign(point_vec[2]) * np.arccos( np.dot(rotated_zy_projection, z_axis) / np.linalg.norm(rotated_zy_projection) ) # Note that due to the right hand rule, a positive y_angle corresponds to rotation from z-axis towards negative y. I.e. opposite sign convention of yps. return x_angle, y_angle - + # ================================================== 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. - + Parameters ---------- x_angle: [rad] float Angle to rotate the coordinate system with in the zx-plane. - + y_angle: [rad] float Angle to rotate the coordinate system with in the zy-plane. 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(). invert: bool Performs a standard passive transformation when False. If True, will perform an active transformation and can thus be used to invert the passive transformation. - - + + Returns ---------- Modified beam xs, ys, zs, uxs, uys and uzs. """ - + x_axis = np.array([0, 1, 0]) # Axis as an unit vector. Axis permutaton is zxy. y_axis = np.array([0, 0, 1]) @@ -398,24 +398,24 @@ def xy_rotate_coord_sys(self, x_angle=None, y_angle=None, invert=False): if y_angle is None: _, y_angle = self.beam_alignment_angles() y_angle = -y_angle - + self.rotate_coord_sys_3D(y_axis, x_angle, x_axis, y_angle, invert=invert) - + # ================================================== def add_pointing_tilts(self, align_x_angle=None, align_y_angle=None): """ Uses active transformation to tilt the beam in the zx- and zy-planes. - + Parameters ---------- align_x_angle: [rad] float Beam coordinates with in the zx-plane are rotated with this angle. - + 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(). - - + + Returns ---------- Modified beam xs, ys and zs. @@ -429,14 +429,14 @@ def add_pointing_tilts(self, align_x_angle=None, align_y_angle=None): y_axis = np.array([0, 0, 1]) x_axis = np.array([0, 1, 0]) - + zs = self.zs() xs = self.xs() ys = self.ys() # Combine into (N, 3) arrays coords = np.column_stack((zs, xs, ys)) - + # Create the rotation object rotation_y = Rot.from_rotvec(align_x_angle * y_axis) rotation_x = Rot.from_rotvec(align_y_angle * x_axis) @@ -472,7 +472,7 @@ def slice_centroids(self, beam_quant, bin_number=None, cut_off=None, make_plot=F make_plot: bool Flag for making plots. - + Returns ---------- beam_quant_slices: 1D float array @@ -481,7 +481,7 @@ def slice_centroids(self, beam_quant, bin_number=None, cut_off=None, make_plot=F z_centroids: [m] 1D float array z-coordinates of the beam slices. """ - + zs = self.zs() mean_z = self.z_offset() weights = self.weightings() @@ -507,7 +507,7 @@ def slice_centroids(self, beam_quant, bin_number=None, cut_off=None, make_plot=F # Beam slice zs _, edges = np.histogram(zs_roi, bins=bin_number) # Get the edges of the histogram of z with bin_number bins. z_centroids = (edges[0:-1] + edges[1:])/2 # Centres of the beam slices (z) - + # Compute the mean of beam_quant of all particles inside a z-bin beam_quant_centroids = np.empty(len(z_centroids)) for i in range(0,len(edges)-1): @@ -520,77 +520,77 @@ def slice_centroids(self, beam_quant, bin_number=None, cut_off=None, make_plot=F plt.scatter(zs*1e6, beam_quant) plt.plot(z_centroids*1e6, beam_quant_centroids, 'rx-') plt.xlabel(r'$\xi$ [$\mathrm{\mu}$m]') - + return beam_quant_centroids, z_centroids - - + + # ================================================== def x_tilt_angle(self, z_cutoff=None): if z_cutoff is None: z_cutoff = 1.5 * self.bunch_length() - + x_centroids, z_centroids = self.slice_centroids(self.xs(), cut_off=z_cutoff, make_plot=False) - + # Perform linear regression slope, _ = np.polyfit(z_centroids, x_centroids, 1) - + return np.arctan(slope) - + # ================================================== def y_tilt_angle(self, z_cutoff=None): 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) - + # Perform linear regression 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) @@ -599,58 +599,58 @@ def y_offset(self, clean=False): 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)) @@ -680,7 +680,7 @@ def norm_amplitude_x(self, plasma_density=None, clean=False): 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()) @@ -691,14 +691,14 @@ def norm_amplitude_y(self, plasma_density=None, clean=False): 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): @@ -709,9 +709,9 @@ def remove_halo_particles(self, nsigma=20): 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)) @@ -720,37 +720,37 @@ def projected_density(self, fcn, bins=None): 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: @@ -768,57 +768,57 @@ def phase_space_density(self, hfcn, vfcn, hbins=None, vbins=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. def charge_density_3D(self, zbins=None, xbins=None, ybins=None): """ Calculates the 3D charge density. - + Parameters ---------- zbins, xbins, ybins: [m] float or 1D float ndarray The bins along z(x,y). - + Returns ---------- - dQ_dxdydz: [C/m^3] 3D float ndarray + dQ_dxdydz: [C/m^3] 3D float ndarray Charge density of the beam. - - zctrs, xctrs, yctrs: [m] 1D float ndarray + + zctrs, xctrs, yctrs: [m] 1D float ndarray The centre positions of the bins of dQ_dxdydz. """ - + zs = self.zs() xs = self.xs() ys = self.ys() - + if zbins is None: zbins = round(np.sqrt(len(self))/2) if xbins is None: xbins = round(np.sqrt(len(self))/2) if ybins is None: ybins = round(np.sqrt(len(self))/2) - + # Create a 3D histogram counts, edges = np.histogramdd((zs, xs, ys), bins=(zbins, xbins, ybins), weights=self.qs()) edges_z = edges[0] edges_x = edges[1] edges_y = edges[2] - + # 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). - + # Calculate charge density per unit volume with np.errstate(divide='ignore', invalid='ignore'): # Handle division by zero dQ_dzdxdy = np.divide(counts, bin_volumes, out=np.zeros_like(counts), where=bin_volumes != 0) @@ -830,16 +830,16 @@ def charge_density_3D(self, zbins=None, xbins=None, ybins=None): xctrs = (edges_x[0:-1] + edges_x[1:])/2 yctrs = (edges_y[0:-1] + edges_y[1:])/2 - #print(np.sum(dQ_dzdxdy*bin_volumes)/self.charge()) - + #print(np.sum(dQ_dzdxdy*bin_volumes)/self.charge()) + return dQ_dzdxdy, zctrs, xctrs, yctrs, edges_z, edges_x, edges_y - + # ================================================== def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inner_off_diag, upper_outer_off_diag, lower_outer_off_diag, num_x_cells, num_unknowns, rhs, boundary_val): """ Applies Dirichlet boundary conditions and assemble the system matrix and the right hand side (source term) of the Poisson equation. - + Parameters ---------- main_diag: [m^-2] 1D float ndarray @@ -856,7 +856,7 @@ def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inne outer_inner_off_diag: [m^-2] 1D float ndarray The outer inne off-diagonal of the system matrix to be modified according to the boundary conditions. - + num_x_cells: float The number of cells in the x-direction. Determines the number of columns of the system matrix A. @@ -869,7 +869,7 @@ def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inne boundary_val: [V/m] float The value of the electric fields Ex and Ey at the simulation box boundary. - + Returns ---------- A: [m^-2] 2D float sparse matrix @@ -878,15 +878,15 @@ def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inne rhs: [V/m^3] 1D float ndarray The modified right hand side of the Poisson equation. """ - + # Set the right side boundary conditions rhs[num_x_cells-1::num_x_cells] = boundary_val # Set BC. Set every num_x_cells-th element starting from the num_x_cells-1 index to 1 - main_diag[num_x_cells-1::num_x_cells] = 1 # Set BC + main_diag[num_x_cells-1::num_x_cells] = 1 # Set BC upper_inner_off_diag[num_x_cells-1::num_x_cells] = 0 # Remove off-diagonal elements at boundaries lower_inner_off_diag[num_x_cells-2::num_x_cells] = 0 # Remove off-diagonal elements at boundaries upper_outer_off_diag[num_x_cells-1::num_x_cells] = 0 # Remove off-diagonal elements at boundaries lower_outer_off_diag[-1::-num_x_cells] = 0 # Remove off-diagonal elements at boundaries - + # Set the left side boundary conditions rhs[0::num_x_cells] = boundary_val main_diag[0::num_x_cells] = 1 @@ -894,14 +894,14 @@ def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inne lower_inner_off_diag[-num_x_cells::-num_x_cells] = 0 upper_outer_off_diag[0::num_x_cells] = 0 lower_outer_off_diag[-num_x_cells::-num_x_cells] = 0 - + # Set the top boundary conditions rhs[1:num_x_cells-1] = boundary_val main_diag[1:num_x_cells-1] = 1 upper_inner_off_diag[1:num_x_cells-1] = 0 lower_inner_off_diag[0:num_x_cells-2] = 0 upper_outer_off_diag[1:num_x_cells-1] = 0 - + # Set the bottom boundary conditions rhs[-num_x_cells+1:-1] = boundary_val main_diag[-num_x_cells+1:-1] = 1 @@ -915,7 +915,7 @@ def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inne A = sp.diags(diagonals, offsets, shape=(num_unknowns, num_unknowns), format="csr") return A, rhs - + # ================================================== def Ex_Ey_2D(self, num_x_cells, num_y_cells, charge_density_xy_slice, dx, dy, boundary_val=0.0): @@ -924,8 +924,8 @@ def Ex_Ey_2D(self, num_x_cells, num_y_cells, charge_density_xy_slice, dx, dy, bo dEx/dx + dEy/dy = 1/epsilon_0 * dQ/dzdxdy dEy/dx - dEx/dy = 0. - - + + Parameters ---------- num_x_cells, num_y_cells: float @@ -937,18 +937,18 @@ def Ex_Ey_2D(self, num_x_cells, num_y_cells, charge_density_xy_slice, dx, dy, bo dx, dy: [m] float Bin widths in x and y of the bins of dQ_dzdxdy. - boundary_val: [V/m] + boundary_val: [V/m] + - Returns ---------- - Ex: [V/m] 2D float array + Ex: [V/m] 2D float array x-conponent of electric field generated by the chosen beam slice. - Ey: [V/m] 2D float array + Ey: [V/m] 2D float array y-conponent of electric field generated by the chosen beam slice. """ - + num_rows, num_cols = charge_density_xy_slice.shape # Set up the system matrix @@ -995,7 +995,7 @@ def Ex_Ey(self, x_box_min, x_box_max, y_box_min, y_box_max, dx, dy, num_z_cells= x_box_max, y_box_max: [m] float The upper x(y) boundary of the simulation domain. Should be much larger than the plasma bubble radius. - + dx, dy: [m] float Bin widths in x and y of the bins of dQ_dzdxdy. @@ -1005,13 +1005,13 @@ def Ex_Ey(self, x_box_min, x_box_max, y_box_min, y_box_max, dx, dy, num_z_cells= boundary_val: [V/m] The values of the electric fields Ex and Ey at the simulation domain boundary. - + Returns ---------- - Ex: [V/m] 2D float array + Ex: [V/m] 2D float array x-conponent of electric field generated by the chosen beam slice. - Ey: [V/m] 2D float array + Ey: [V/m] 2D float array y-conponent of electric field generated by the chosen beam slice. zctrs, xctrs, yctrs: [m] 1D float ndarray @@ -1021,25 +1021,25 @@ def Ex_Ey(self, x_box_min, x_box_max, y_box_min, y_box_max, dx, dy, num_z_cells= # Check if the selected simulation boundaries are significantly larger than the beam extent xs = self.xs() ys = self.ys() - + if np.abs(x_box_min/xs.min()) < tolerance or np.abs(y_box_min/ys.min()) < tolerance or np.abs(x_box_max/xs.max()) < tolerance or np.abs(y_box_max/ys.max()) < tolerance: raise ValueError('Simulation box size is too small compared to beam size.') - + if num_z_cells is None: num_z_cells = round(np.sqrt(len(self))/2) z_box_max = np.max(self.zs()) z_box_min = np.min(self.zs()) zbins = np.linspace(z_box_min, z_box_max, num_z_cells+1) - + num_x_cells = int((x_box_max-x_box_min)/dx) num_y_cells = int((y_box_max-y_box_min)/dy) - + xbins = np.linspace(x_box_min, x_box_max, num_x_cells+1) ybins = np.linspace(y_box_min, y_box_max, num_y_cells+1) dQ_dzdxdy, zctrs, xctrs, yctrs, edges_z, edges_x, edges_y = self.charge_density_3D(zbins=zbins, xbins=xbins, ybins=ybins) - + Ex = np.zeros((num_z_cells, num_y_cells, num_x_cells)) Ey = np.zeros((num_z_cells, num_y_cells, num_x_cells)) @@ -1049,58 +1049,58 @@ def Ex_Ey(self, x_box_min, x_box_max, y_box_min, y_box_max, dx, dy, num_z_cells= # Calculate the fields for the charge density slice Ex_2d, Ey_2d = self.Ex_Ey_2D(num_x_cells, num_y_cells, charge_density_xy_slice, dx, dy, boundary_val=boundary_val) - + Ex[slice_idx,:,:] = Ex_2d Ey[slice_idx,:,:] = Ey_2d - + return Ex, Ey, zctrs, xctrs, yctrs - - - + + + ## PLOTTING def plot_current_profile(self): dQdt, ts = self.current_profile() fig, ax = plt.subplots() fig.set_figwidth(6) - fig.set_figheight(4) + 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) - + 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) + 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) + 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)') @@ -1108,49 +1108,52 @@ def plot_trace_space_y(self): cb = fig.colorbar(p) cb.ax.set_ylabel('Charge density (pC/um/mrad)') - def plot_transverse_profile(self): + def plot_transverse_profile(self, xlims=(None, None), ylims=(None, None), title=''): 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], + #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)') + ax.set_xlim(xlims) + ax.set_ylim(ylims) + ax.set_title(title) + - def plot_bunch_pattern(self): - + fig, ax = plt.subplots() fig.set_figwidth(6) - fig.set_figheight(4) + 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): - + # add energy and chirp - Es = self.Es() + energy_gain + Es = self.Es() + energy_gain Es = Es + np.sign(self.qs()) * (self.zs()-z_offset) * chirp self.set_uzs(energy2proper_velocity(Es)) - + # remove particles with subzero energy del self[Es < 0] - + def compress(self, R_56, nom_energy): zs = self.zs() + (1-self.Es()/nom_energy) * R_56 self.set_zs(zs) - + def scale_to_length(self, bunch_length): z_mean = self.z_offset() zs_scaled = z_mean + (self.zs()-z_mean)*bunch_length/self.bunch_length() @@ -1165,17 +1168,17 @@ def scale_norm_emittance_y(self, norm_emit_ny): 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)) - + # betatron damping (must be done before acceleration) def apply_betatron_damping(self, deltaE): gammasBoosted = energy2gamma(abs(self.Es()+deltaE)) betamag = np.sqrt(self.gammas()/gammasBoosted) self.magnify_beta_function(betamag) - - + + # magnify beta function (increase beam size, decrease divergence) def magnify_beta_function(self, beta_mag, axis_defining_beam=None): - + # calculate beam (not beta) magnification mag = np.sqrt(beta_mag) @@ -1189,19 +1192,19 @@ def magnify_beta_function(self, beta_mag, axis_defining_beam=None): y_offset = axis_defining_beam.y_offset() ux_offset = axis_defining_beam.ux_offset() uy_offset = axis_defining_beam.uy_offset() - + self.set_xs((self.xs()-x_offset)*mag + x_offset) self.set_ys((self.ys()-y_offset)*mag + y_offset) self.set_uxs((self.uxs()-ux_offset)/mag + ux_offset) self.set_uys((self.uys()-uy_offset)/mag + uy_offset) - + # transport in a drift def transport(self, L): self.set_xs(self.xs() + L*self.xps()) self.set_ys(self.ys() + L*self.yps()) - + def flip_transverse_phase_spaces(self, flip_momenta=True, flip_positions=False): if flip_momenta: self.set_uxs(-self.uxs()) @@ -1210,9 +1213,9 @@ def flip_transverse_phase_spaces(self, flip_momenta=True, flip_positions=False): self.set_xs(-self.xs()) 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): - + # remove particles with subzero and Nan energy neg_indices = self.Es() < 0 del self[neg_indices] @@ -1223,15 +1226,15 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat del self[nan_indices] if isinstance(deltaEs, np.ndarray) and len(deltaEs) > 1: deltaEs = deltaEs[~nan_indices] - + # determine initial and final Lorentz factor gamma0s = energy2gamma(self.Es()) Es_final = self.Es() + deltaEs gammas = energy2gamma(Es_final) dgamma_ds = (gammas-gamma0s)/L - + if calc_evolution: - + # calculate evolution num_evol_steps = max(20, min(400, round(2*L/(beta_matched(n0, self.energy()+min(0,np.mean(deltaEs))))))) evol = SimpleNamespace() @@ -1248,13 +1251,13 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat # 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) @@ -1271,15 +1274,15 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat 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) - - + + # calculate final positions and angles after betatron motion if radiation_reaction: xs, uxs, ys, uys, Es_final = evolve_betatron_motion(self.xs()-x0_driver, self.uxs(), self.ys()-y0_driver, self.uys(), L, gamma0s, dgamma_ds, k_p(n0)) else: xs, uxs = evolve_hills_equation_analytic(self.xs()-x0_driver, self.uxs(), L, gamma0s, dgamma_ds, k_p(n0)) ys, uys = evolve_hills_equation_analytic(self.ys()-y0_driver, self.uys(), L, gamma0s, dgamma_ds, k_p(n0)) - + # set new beam positions and angles (shift back driver offsets) self.set_xs(xs+x0_driver) self.set_uxs(uxs) @@ -1290,31 +1293,31 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat return Es_final, evol else: return Es_final - - + + ## 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') @@ -1322,16 +1325,16 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): # 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) @@ -1346,10 +1349,10 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): 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) @@ -1361,10 +1364,10 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): 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['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()) @@ -1379,42 +1382,42 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): 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() @@ -1425,16 +1428,16 @@ def load(_, filename, beam_name='beam'): 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") @@ -1448,13 +1451,13 @@ def load(_, filename, beam_name='beam'): 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): - + if axes is None: fig = plt.figure() # an empty figure with an axes ax = fig.add_axes([.15, .15, .75, .75]) @@ -1462,7 +1465,7 @@ def imshow_plot(self, data, axes=None, extent=None, vmin=None, vmax=None, colmap else: #ax = axes[0] # TODO: adjust colourbar axes #cbar_ax = axes[1] - + ax = axes cbar_ax = None @@ -1509,7 +1512,7 @@ def imshow_plot(self, data, axes=None, extent=None, vmin=None, vmax=None, colmap ax.set_ylabel(ylab) ax.set_xlabel(xlab) - + # ================================================== def distribution_plot_2D(self, arr1, arr2, weights=None, hist_bins=None, hist_range=None, axes=None, extent=None, vmin=None, vmax=None, colmap=CONFIG.default_cmap, xlab='', ylab='', clab='', origin='lower', interpolation='nearest', reduce_cax_pad=False): @@ -1524,20 +1527,20 @@ def distribution_plot_2D(self, arr1, arr2, weights=None, hist_bins=None, hist_ra hist_range[1] = [ arr2.min(), arr2.max() ] if extent is None: extent = hist_range[0] + hist_range[1] - + binned_data, zedges, xedges = np.histogram2d(arr1, arr2, hist_bins, hist_range, weights=weights) beam_hist2d = binned_data.T/np.diff(zedges)/np.diff(xedges) self.imshow_plot(beam_hist2d, axes=axes, extent=extent, vmin=vmin, vmax=vmax, colmap=colmap, xlab=xlab, ylab=ylab, clab=clab, gridOn=False, origin=origin, interpolation=interpolation, reduce_cax_pad=reduce_cax_pad) - + # ================================================== def density_map_diags(self): - + #colors = ['white', 'aquamarine', 'lightgreen', 'green'] #colors = ['white', 'forestgreen', 'limegreen', 'lawngreen', 'aquamarine', 'deepskyblue'] #bounds = [0, 0.2, 0.4, 0.8, 1] #cmap = LinearSegmentedColormap.from_list('my_cmap', colors, N=256) - + cmap = CONFIG.default_cmap # Macroparticles data @@ -1557,7 +1560,7 @@ def density_map_diags(self): xps_lab = r"$x'$ [mrad]" yps_lab = r"$y'$ [mrad]" energ_lab = r'$\mathcal{E}$ [GeV]' - + # Set up a figure with axes fig, axs = plt.subplots(nrows=3, ncols=3, layout='constrained', figsize=(5*3, 4*3)) fig.suptitle(r'$\Delta s=$' f'{format(self.location, ".2f")}' ' m') @@ -1573,7 +1576,7 @@ def density_map_diags(self): extent_zx = [i*1e6 for i in extent_zx] # [um] self.distribution_plot_2D(arr1=zs, arr2=xs, weights=weights, hist_bins=hist_bins, hist_range=hist_range, axes=axs[0][0], extent=extent_zx, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=xlab, clab=r'$\partial^2 N/\partial\xi \partial x$ [$\mathrm{m}^{-2}$]', origin='lower', interpolation='nearest') - + # 2D z-x' distribution hist_range_xps = [[None, None], [None, None]] @@ -1586,8 +1589,8 @@ def density_map_diags(self): extent_xps[3] = extent_xps[3]*1e3 # [mrad] self.distribution_plot_2D(arr1=zs, arr2=xps, weights=weights, hist_bins=hist_bins, hist_range=hist_range_xps, axes=axs[0][1], extent=extent_xps, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=xps_lab, clab=r"$\partial^2 N/\partial z \partial x'$ [$\mathrm{m}^{-1}$ $\mathrm{rad}^{-1}$]", origin='lower', interpolation='nearest') - - + + # 2D x-x' distribution hist_range_xxp = [[None, None], [None, None]] hist_range_xxp[0] = hist_range[1] @@ -1599,7 +1602,7 @@ def density_map_diags(self): extent_xxp[3] = extent_xxp[3]*1e3 # [mrad] self.distribution_plot_2D(arr1=xs, arr2=xps, weights=weights, hist_bins=hist_bins, hist_range=hist_range_xxp, axes=axs[0][2], extent=extent_xxp, vmin=None, vmax=None, colmap=cmap, xlab=xlab, ylab=xps_lab, clab=r"$\partial^2 N/\partial x\partial x'$ [$\mathrm{m}^{-1}$ $\mathrm{rad}^{-1}$]", origin='lower', interpolation='nearest') - + # 2D z-y distribution hist_range_zy = [[None, None], [None, None]] @@ -1609,7 +1612,7 @@ def density_map_diags(self): extent_zy = [i*1e6 for i in extent_zy] # [um] self.distribution_plot_2D(arr1=zs, arr2=ys, weights=weights, hist_bins=hist_bins, hist_range=hist_range_zy, axes=axs[1][0], extent=extent_zy, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=ylab, clab=r'$\partial^2 N/\partial\xi \partial y$ [$\mathrm{m}^{-2}$]', origin='lower', interpolation='nearest') - + # 2D z-y' distribution hist_range_yps = [[None, None], [None, None]] @@ -1620,9 +1623,9 @@ def density_map_diags(self): extent_yps[1] = extent_yps[1]*1e6 # [um] extent_yps[2] = extent_yps[2]*1e3 # [mrad] extent_yps[3] = extent_yps[3]*1e3 # [mrad] - + self.distribution_plot_2D(arr1=zs, arr2=yps, weights=weights, hist_bins=hist_bins, hist_range=hist_range_yps, axes=axs[1][1], extent=extent_yps, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=yps_lab, clab=r"$\partial^2 N/\partial z \partial y'$ [$\mathrm{m}^{-1}$ $\mathrm{rad}^{-1}$]", origin='lower', interpolation='nearest') - + # 2D y-y' distribution hist_range_yyp = [[None, None], [None, None]] @@ -1633,9 +1636,9 @@ def density_map_diags(self): extent_yyp[1] = extent_yyp[1]*1e6 # [um] extent_yyp[2] = extent_yyp[2]*1e3 # [mrad] extent_yyp[3] = extent_yyp[3]*1e3 # [mrad] - + self.distribution_plot_2D(arr1=ys, arr2=yps, weights=weights, hist_bins=hist_bins, hist_range=hist_range_yyp, axes=axs[1][2], extent=extent_yyp, vmin=None, vmax=None, colmap=cmap, xlab=ylab, ylab=yps_lab, clab=r"$\partial^2 N/\partial y\partial y'$ [$\mathrm{m}^{-1}$ $\mathrm{rad}^{-1}$]", origin='lower', interpolation='nearest') - + # 2D x-y distribution hist_range_xy = [[None, None], [None, None]] @@ -1645,7 +1648,7 @@ def density_map_diags(self): extent_xy = [i*1e6 for i in extent_xy] # [um] self.distribution_plot_2D(arr1=xs, arr2=ys, weights=weights, hist_bins=hist_bins, hist_range=hist_range_xy, axes=axs[2][0], extent=extent_xy, vmin=None, vmax=None, colmap=cmap, xlab=xlab, ylab=ylab, clab=r'$\partial^2 N/\partial x \partial y$ [$\mathrm{m}^{-2}$]', origin='lower', interpolation='nearest') - + # Energy distribution ax = axs[2][1] From 47282f66dfcad74249e087b2a0447dd9aa16acea Mon Sep 17 00:00:00 2001 From: Felipe Pena Date: Wed, 26 Mar 2025 19:43:49 +0200 Subject: [PATCH 05/11] Added the SourceEnergyComb in the __init__ file. --- abel/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/abel/__init__.py b/abel/__init__.py index 60719823..6245a820 100644 --- a/abel/__init__.py +++ b/abel/__init__.py @@ -42,6 +42,7 @@ from .classes.source.impl.source_basic import * from .classes.source.impl.source_trapezoid import SourceTrapezoid from .classes.source.impl.source_combiner import * +from .classes.source.impl.source_energy_comb import * from .classes.source.impl.source_from_file import * from .classes.source.impl.source_flattop import * from .classes.stage.impl.stage_basic import * From 9f869360756fe3b514f8a2ac4a3ef65ba4c026e7 Mon Sep 17 00:00:00 2001 From: Felipe Pena Date: Thu, 3 Apr 2025 10:51:49 +0300 Subject: [PATCH 06/11] Allows to use Experiment() class without a component, e.g., just a source and a spectrometer. --- .../beamline/impl/experiment/experiment.py | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/abel/classes/beamline/impl/experiment/experiment.py b/abel/classes/beamline/impl/experiment/experiment.py index b4dff5bc..08cdf4e1 100644 --- a/abel/classes/beamline/impl/experiment/experiment.py +++ b/abel/classes/beamline/impl/experiment/experiment.py @@ -9,33 +9,34 @@ import warnings class Experiment(Beamline): - + def __init__(self, linac=None, component=None, spectrometer=None, num_bunches_in_train=1, bunch_separation=0, rep_rate_trains=10): super().__init__(num_bunches_in_train, bunch_separation, rep_rate_trains) - + self.linac = linac self.component = component self.spectrometer = spectrometer - - - + + + # assemble the trackables def assemble_trackables(self): self.trackables = [] - + # add the linac/source assert(isinstance(self.linac, Linac) or isinstance(self.linac, Source)) self.trackables.append(self.linac) # add the experimental component - self.trackables.append(self.component) + if self.component is not None: + self.trackables.append(self.component) # add the spectrometer (if it exists if self.spectrometer is not None: assert(isinstance(self.spectrometer, Spectrometer)) self.trackables.append(self.spectrometer) - + # set the bunch train pattern etc. super().assemble_trackables() @@ -45,10 +46,10 @@ def energy_usage(self): def get_cost_breakdown(): return ('Experiment', 0) - + # density plots def plot_spectrometer_screen(self, xlims=None, ylims=None, E_calib = False, diverg = None, plot_m12 = False, savefig = None): - + # load phase space beam = self.final_beam @@ -60,15 +61,15 @@ def plot_spectrometer_screen(self, xlims=None, ylims=None, E_calib = False, dive if ylims is not None: ybins = np.linspace(min(ylims), max(ylims), num_bins) dQdxdy, xedges, yedges = beam.density_transverse(hbins=xbins, vbins=ybins) - + # prepare figure fig, ax = plt.subplots(2, 1, gridspec_kw={'height_ratios': [2, 1]}) fig.set_figwidth(8) fig.set_figheight(8) - + # current profile c0 = ax[0].pcolor(xedges*1e3, yedges * 1e3, abs(dQdxdy) * 1e3, cmap=CONFIG.default_cmap, shading='auto') - + # make plot ax[0].set_xlabel('x (mm)') ax[0].set_ylabel('y (mm)') @@ -77,25 +78,25 @@ def plot_spectrometer_screen(self, xlims=None, ylims=None, E_calib = False, dive cbar0 = fig.colorbar(c0, ax=ax[0], pad = 0.015*fig.get_figwidth()) cbar0.ax.set_ylabel('Charge density (nC ' r'$\mathrm{mm^{-2})}$') ax[0].set_ylim(min(yedges * 1e3), max(yedges * 1e3)) - + # calculate energy axis (E = E0*Dy/y) Dy_img = self.spectrometer.get_dispersion(energy=self.spectrometer.img_energy) E_times_Dy = self.spectrometer.img_energy*Dy_img - + # add imaging energies y_img = -E_times_Dy/self.spectrometer.img_energy y_img_y = -E_times_Dy/self.spectrometer.img_energy_y ax[0].axhline(y_img*1e3, color = 'black', linestyle = '--', linewidth = 1, label = 'Imaging energy x (GeV)') ax[0].axhline(y_img_y*1e3, color = 'black', linestyle = ':', linewidth = 1, label = 'Imaging energy y (GeV)') - + # add energy axis warnings.simplefilter('ignore', category=RuntimeWarning) ax2 = ax[0].secondary_yaxis('right', functions=(lambda y: -E_times_Dy/(y/1e3)/1e9, lambda y: -E_times_Dy/(y/1e3)/1e9)) ax2.set_ylabel('Energy (GeV)') - + if diverg is not None: energies = -E_times_Dy/yedges - m12s = self.spectrometer.get_m12(energies) + m12s = self.spectrometer.get_m12(energies) x_pos = m12s*diverg x_neg = -m12s*diverg ax[0].plot(x_pos*1e3, yedges*1e3, color = 'orange', label = str(diverg*1e3) + ' mrad butterfly', alpha = 0.7) @@ -104,12 +105,11 @@ def plot_spectrometer_screen(self, xlims=None, ylims=None, E_calib = False, dive ax[1].set_ylabel(r'$\mathrm{m_{12}}$') ax[1].set_xlabel('E (GeV)') ax[1].grid(False, which='major') - + ax[0].set_xlim(min(xedges*1e3), max(xedges*1e3)) ax[0].legend(loc = 'center right', fontsize = 6) - + if plot_m12 == False: fig.delaxes(ax[1]) if savefig is not None: plt.savefig(str(savefig), dpi = 700) - \ No newline at end of file From 61404d8f67bd15629f6b9d60790bfef7eb7e84e3 Mon Sep 17 00:00:00 2001 From: Felipe Pena Date: Thu, 3 Apr 2025 11:45:49 +0300 Subject: [PATCH 07/11] Added documentation to the new files --- .../classes/source/impl/source_energy_comb.py | 122 +++++++++++++- .../impl/spectrometer_achromatic_basic.py | 153 +++++++++++------- .../achromatic/spectrometer_achromatic.py | 111 ++++++++++--- 3 files changed, 307 insertions(+), 79 deletions(-) diff --git a/abel/classes/source/impl/source_energy_comb.py b/abel/classes/source/impl/source_energy_comb.py index 8ada5f34..1b4e2b2f 100644 --- a/abel/classes/source/impl/source_energy_comb.py +++ b/abel/classes/source/impl/source_energy_comb.py @@ -6,8 +6,73 @@ from abel.utilities.relativity import energy2gamma class SourceEnergyComb(Source): - # this class behaves like the SourceBasic but has the added functionality - # of distributing the particle's energies within a given list of energies + """ + A source class that extends the functionality of `SourceBasic` by distributing + particle energies within a given list of energies. This is useful for creating + comb-like energy distributions for resolution estimates. + + Parameters + ---------- + length : float, optional + Length of the source [m]. Default is 0. + + num_particles : int, optional + Number of particles in the beam. Default is 1000. + + energy : float, optional + Central energy of the beam [eV]. Default is None. + + charge : float, optional + Total charge of the beam [C]. Default is 0. + + rel_energy_spread : float, optional + Relative energy spread (fraction). Default is None. + + energy_spread : float, optional + Absolute energy spread [eV]. Default is None. + + bunch_length : float, optional + RMS bunch length [m]. Default is None. + + z_offset : float, optional + Longitudinal offset of the beam [m]. Default is 0. + + x_offset, y_offset : float, optional + Transverse offsets of the beam [m]. Default is 0. + + x_angle, y_angle : float, optional + Transverse angles of the beam [rad]. Default is 0. + + emit_nx, emit_ny : float, optional + Normalized emittances in x and y [m rad]. Default is 0. + + beta_x, beta_y : float, optional + Beta functions in x and y [m]. Default is None. + + alpha_x, alpha_y : float, optional + Alpha functions in x and y. Default is 0. + + angular_momentum : float, optional + Angular momentum of the beam. Default is 0. + + wallplug_efficiency : float, optional + Efficiency of the source. Default is 1. + + accel_gradient : float, optional + Acceleration gradient [V/m]. Default is None. + + symmetrize : bool, optional + Whether to symmetrize the transverse phase space. Default is False. + + symmetrize_6d : bool, optional + Whether to symmetrize the full 6D phase space. Default is False. + + z_cutoff : float, optional + Longitudinal cutoff for filtering particles [m]. Default is None. + + energy_comb_delta : float, optional + Relative energy step size for creating comb-like energy distributions. Default is None. + """ def __init__(self, length=0, num_particles=1000, energy=None, charge=0, rel_energy_spread=None, energy_spread=None, bunch_length=None, z_offset=0, x_offset=0, y_offset=0, x_angle=0, y_angle=0, emit_nx=0, emit_ny=0, beta_x=None, beta_y=None, alpha_x=0, alpha_y=0, angular_momentum=0, wallplug_efficiency=1, accel_gradient=None, symmetrize=False, symmetrize_6d=False, z_cutoff=None, energy_comb_delta=None): @@ -38,6 +103,29 @@ def __init__(self, length=0, num_particles=1000, energy=None, charge=0, rel_ener def track(self, _=None, savedepth=0, runnable=None, verbose=False): + """ + Tracks the beam through the source, generating its phase space and applying + energy comb distributions or filters if specified. + + Parameters + ---------- + _ : None, optional + Placeholder parameter. Default is None. + + savedepth : int, optional + Depth of saved tracking data. Default is 0. + + runnable : object, optional + Runnable object for tracking. Default is None. + + verbose : bool, optional + Whether to print verbose output. Default is False. + + Returns + ------- + Beam + The tracked beam object. + """ # make empty beam beam = Beam() @@ -92,6 +180,22 @@ def track(self, _=None, savedepth=0, runnable=None, verbose=False): return super().track(beam, savedepth, runnable, verbose) def reorder_particle_energies(self, beam, plot=False): + """ + Reorders the particle energies in the beam to create a comb-like energy + distribution based on `self.energy_comb_delta`. + + Parameters + ---------- + beam : Beam + The beam object whose particle energies will be reordered. + + plot : bool, optional + Whether to plot the longitudinal phase space after reordering. Default is False. + + Returns + ------- + None + """ # takes in the beam and uses the attribute self.energy_comb_delta # to create a list of energies that will be given to each beam particle @@ -144,6 +248,20 @@ def reorder_particle_energies(self, beam, plot=False): # ================================================== # Filter out particles whose z < z_cutoff for testing instability etc. def z_filter(self, beam): + """ + Filters out particles whose longitudinal positions (z) are below the + specified cutoff (`self.z_cutoff`). + + Parameters + ---------- + beam : Beam + The beam object to be filtered. + + Returns + ------- + Beam + A new beam object with particles filtered based on the z cutoff. + """ xs = beam.xs() ys = beam.ys() zs = beam.zs() diff --git a/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py b/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py index 33e132a5..83c2bbd1 100644 --- a/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py +++ b/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py @@ -1,12 +1,14 @@ """ Module: spectrometer_achromatic_basic -This module defines the SpectrometerAchromaticBasic class, which is a subclass of SpectrometerAchromatic. -It provides methods to initialize the spectrometer and track the beam. +This module defines the `SpectrometerAchromaticBasic` class, which is a +subclass of `SpectrometerAchromatic`. It provides methods to initialize +the spectrometer and track the beam. Classes ------- -SpectrometerAchromaticBasic +SpectrometerAchromaticBasic : SpectrometerAchromatic + A basic implementation of an achromatic spectrometer. Imports ------- @@ -25,52 +27,81 @@ class SpectrometerAchromaticBasic(SpectrometerAchromatic): Parameters ---------- + length_drift_to_dipole : float, optional + Length of the drift space to the dipole magnet [m]. Default is 0.2. + field_dipole : float, optional - The magnetic field strength of the dipole (default is 1.0). + Magnetic field strength of the dipole magnet [T]. Negative values + bend the beam downwards. Default is -0.85. + length_dipole : float, optional - The length of the dipole (default is 0.2). - length_drift_to_dipole : float, optional - The drift before the dipole (default is 0.3). + Length of the dipole magnet [m]. Default is 0.15. + length_drift_dipole_to_lens : float, optional - The drift before the plasma lens (default is 0.2). + Length of the drift space from the dipole magnet to the lens [m]. + Default is 0.15. + + length_plasma_lens : float, optional + Length of the plasma lens [m]. Default is 20E-3. + + radius_plasma_lens : float, optional + Radius of the plasma lens [m]. Default is 1E-3. + length_drift_lens_to_screen : float, optional - The drift before the screen (default is 0.3). + Length of the drift space from the lens to the screen [m]. Default + is 0.4. + + imaging_energy : float or None, optional + Imaging energy [eV]. Default is None. + + disable_lens_nonlinearity : bool, optional + Flag to disable lens nonlinearity. Default is True. """ def __init__(self, - length_drift_to_dipole=0.2, # [m] - field_dipole=-0.85, # [T] - length_dipole=0.15, # [m] - length_drift_dipole_to_lens=0.15, # [m] - length_plasma_lens=20E-3, # [m] - radius_plasma_lens=1E-3, # [m] - length_drift_lens_to_screen=0.4, # [m] - imaging_energy=None, # [eV] - disable_lens_nonlinearity=True # [bool] - ): + length_drift_to_dipole=0.2, + field_dipole=-0.85, + length_dipole=0.15, + length_drift_dipole_to_lens=0.15, + length_plasma_lens=20E-3, + radius_plasma_lens=1E-3, + length_drift_lens_to_screen=0.4, + imaging_energy=None, + disable_lens_nonlinearity=True): """ - Initializes the SpectrometerAchromaticBasic with the given parameters. + Initializes the `SpectrometerAchromaticBasic` with the given + parameters. Parameters ---------- length_drift_to_dipole : float, optional - Length of the drift space to the dipole magnet in meters. Default is 0.2. + Length of the drift space to the dipole magnet [m]. Default is + 0.2. + field_dipole : float, optional - Magnetic field strength of the dipole magnet in Tesla. - Negative means bending the beam downwards. - Default is -0.85. + Magnetic field strength of the dipole magnet [T]. Negative + values bend the beam downwards. Default is -0.85. + length_dipole : float, optional - Length of the dipole magnet in meters. Default is 0.8. + Length of the dipole magnet [m]. Default is 0.15. + length_drift_dipole_to_lens : float, optional - Length of the drift space from the dipole magnet to the lens in meters. Default is 0.2. + Length of the drift space from the dipole magnet to the lens + [m]. Default is 0.15. + length_plasma_lens : float, optional - Length of the plasma lens in meters. Default is 20E-3. + Length of the plasma lens [m]. Default is 20E-3. + radius_plasma_lens : float, optional - Radius of the plasma lens in meters. Default is 1E-3. + Radius of the plasma lens [m]. Default is 1E-3. + length_drift_lens_to_screen : float, optional - Length of the drift space from the lens to the screen in meters. Default is 0.4. + Length of the drift space from the lens to the screen [m]. + Default is 0.4. + imaging_energy : float or None, optional - Imaging energy in electron volts. Default is None. + Imaging energy [eV]. Default is None. + disable_lens_nonlinearity : bool, optional Flag to disable lens nonlinearity. Default is True. """ @@ -93,17 +124,20 @@ def track(self, beam0, savedepth=0, runnable=None, verbose=False): ---------- beam0 : Beam The initial beam to be tracked. + savedepth : int, optional - The depth at which to save the beam state (default is 0). + The depth at which to save the beam state. Default is 0. + runnable : callable, optional - A callable to be run during tracking. + A callable to be run during tracking. Default is None. + verbose : bool, optional - Flag to indicate if verbose output is desired (default is False). + Whether to print verbose output. Default is False. Returns ------- Beam - The tracked beam. + The tracked beam object. """ # transport the beam to the center of the dipole @@ -211,17 +245,17 @@ def track(self, beam0, savedepth=0, runnable=None, verbose=False): def get_bending_angle(self, energy): """ - Calculate the bending angle of the spectrometer for a given energy. + Calculates the bending angle of the spectrometer for a given energy. - Parameters: - ----------- - energy : float, optional - The energy of the particle beam in eV. + Parameters + ---------- + energy : float + The energy of the particle beam [eV]. - Returns: - -------- - bend_angle : float - The calculated bending angle. + Returns + ------- + float + The calculated bending angle [rad]. """ # computing the bending angle of the dipole for each particle @@ -236,23 +270,28 @@ def get_bending_angle(self, energy): def get_dispersion(self, energy=None, bend_angle=None, length=None): """ - Calculate the dispersion of the spectrometer for a given energy. - Either energy or bending angle (bend_angle) have to be given. - Parameters: - ----------- + Calculates the dispersion of the spectrometer for a given energy or + bending angle. + + Parameters + ---------- energy : float, optional - The energy of the particle beam in eV. + The energy of the particle beam [eV]. Either `energy` or + `bend_angle` must be provided. + bend_angle : float, optional - Bending angle of the dipole in rad. + Bending angle of the dipole [rad]. Either `bend_angle` or + `energy` must be provided. + length : float, optional - The distance to use for the calculation. If not provided, - it will be calculated as the sum of the lengths of the dipole/2, drift space - from dipole to lens, plasma lens, and drift space from lens to screen, - i.e., the distance from the center of the dipole to the screen. - Returns: - -------- - dispersion : float - The calculated dispersion value. + The distance to use for the calculation [m]. If not provided, + it defaults to the distance from the center of the dipole to + the screen. + + Returns + ------- + float + The calculated dispersion value [m]. """ text = 'Please give either the bending angle *or* the energy.' diff --git a/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py b/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py index 2595592f..1283a59e 100644 --- a/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py +++ b/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py @@ -9,27 +9,40 @@ class SpectrometerAchromatic(Spectrometer): """ Achromatic Spectrometer class. + This class models an achromatic spectrometer, which is used to analyze + particle beams by bending them through a magnetic field and focusing + them using a plasma lens. + Parameters ---------- - length_drift_to_dipole : float, optional - Length of the drift space to the dipole magnet in meters. Default is 0.2. - field_dipole : float, optional - Magnetic field strength of the dipole magnet in Tesla. - Negative means bending the beam downwards. - Default is -0.85. - length_dipole : float, optional - Length of the dipole magnet in meters. Default is 0.8. - length_drift_dipole_to_lens : float, optional - Length of the drift space from the dipole magnet to the lens in meters. Default is 0.2. - length_plasma_lens : float, optional - Length of the plasma lens in meters. Default is 20E-3. - radius_plasma_lens : float, optional - Radius of the plasma lens in meters. Default is 1E-3. - length_drift_lens_to_screen : float, optional - Length of the drift space from the lens to the screen in meters. Default is 0.4. - imaging_energy : float or None, optional - Imaging energy in electron volts. Default is None. - disable_lens_nonlinearity : bool, optional + length_drift_to_dipole : float + Length of the drift space to the dipole magnet [m]. Default is 0.2. + + field_dipole : float + Magnetic field strength of the dipole magnet [T]. Negative values + indicate bending the beam downwards. Default is -0.85. + + length_dipole : float + Length of the dipole magnet [m]. Default is 0.8. + + length_drift_dipole_to_lens : float + Length of the drift space from the dipole magnet to the lens [m]. + Default is 0.2. + + length_plasma_lens : float + Length of the plasma lens [m]. Default is 20E-3. + + radius_plasma_lens : float + Radius of the plasma lens [m]. Default is 1E-3. + + length_drift_lens_to_screen : float + Length of the drift space from the lens to the screen [m]. + Default is 0.4. + + imaging_energy : float or None + Imaging energy [eV]. Default is None. + + disable_lens_nonlinearity : bool Flag to disable lens nonlinearity. Default is True. """ @@ -44,7 +57,41 @@ def __init__(self, length_drift_lens_to_screen, imaging_energy, disable_lens_nonlinearity): + """ + Initializes the SpectrometerAchromatic object with the specified parameters. + + Parameters + ---------- + length_drift_to_dipole : float + Length of the drift space to the dipole magnet [m]. Default is 0.2. + + field_dipole : float + Magnetic field strength of the dipole magnet [T]. Negative values + indicate bending the beam downwards. Default is -0.85. + + length_dipole : float + Length of the dipole magnet [m]. Default is 0.8. + + length_drift_dipole_to_lens : float + Length of the drift space from the dipole magnet to the lens [m]. + Default is 0.2. + + length_plasma_lens : float + Length of the plasma lens [m]. Default is 20E-3. + + radius_plasma_lens : float + Radius of the plasma lens [m]. Default is 1E-3. + length_drift_lens_to_screen : float + Length of the drift space from the lens to the screen [m]. + Default is 0.4. + + imaging_energy : float or None + Imaging energy [eV]. Default is None. + + disable_lens_nonlinearity : bool + Flag to disable lens nonlinearity. Default is True. + """ self.length_drift_to_dipole = length_drift_to_dipole # [m] self.field_dipole = field_dipole # [T] self.length_dipole = length_dipole # [m] @@ -67,7 +114,7 @@ def get_length(self): Returns ------- float - Total length of the spectrometer lattice in meters. + Total length of the spectrometer lattice [m]. """ total_length = (self.length_drift_to_dipole + self.length_dipole @@ -78,10 +125,31 @@ def get_length(self): @property def radius_plasma_lens(self): + """ + Get the radius of the plasma lens. + + Returns + ------- + float + Radius of the plasma lens [m]. + """ return self._radius_plasma_lens @radius_plasma_lens.setter def radius_plasma_lens(self, value): + """ + Set the radius of the plasma lens. + + Parameters + ---------- + value : float + Radius of the plasma lens [m]. Must be greater than or equal to zero. + + Raises + ------ + AssertionError + If the radius is less than zero. + """ assert value >= 0, 'Radius of the plasma lens must be greater than or equal to zero.' self._radius_plasma_lens = value @@ -94,10 +162,13 @@ def track(self, beam, savedepth=0, runnable=None, verbose=False): ---------- beam : object The beam object to be tracked. + savedepth : int, optional Depth of saved states during tracking. Default is 0. + runnable : object or None, optional Runnable object for tracking. Default is None. + verbose : bool, optional Flag to enable verbose output. Default is False. From 29c55ab6f1bc83a59a1012a1eb50736d097d1fb0 Mon Sep 17 00:00:00 2001 From: Kyrre Ness Sjobak Date: Thu, 3 Apr 2025 16:09:29 +0200 Subject: [PATCH 08/11] Revert "Added the possibility to edit the plot in plot_transverse_profile" This reverts commit f8b2864d6627e7a3e2450440febe1b4be67defbe. --- abel/classes/beam.py | 515 +++++++++++++++++++++---------------------- 1 file changed, 256 insertions(+), 259 deletions(-) diff --git a/abel/classes/beam.py b/abel/classes/beam.py index 5130c434..d1d72215 100644 --- a/abel/classes/beam.py +++ b/abel/classes/beam.py @@ -18,7 +18,7 @@ from matplotlib import pyplot as plt class Beam(): - + def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1, bunch_separation=0.0): # the phase space variable is private @@ -30,39 +30,39 @@ def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1, # bunch pattern information self.num_bunches_in_train = num_bunches_in_train self.bunch_separation = bunch_separation # [s] - + self.trackable_number = -1 # will increase to 0 after first tracking element self.stage_number = 0 - self.location = 0 - - + self.location = 0 + + # reset phase space def reset_phase_space(self, num_particles): self.__phasespace = np.zeros((8, num_particles)) - + # filter out macroparticles based on a mask (true means delete) def __delitem__(self, indices): if hasattr(indices, 'len'): if len(indices) == len(self): indices = np.where(indices) self.__phasespace = np.ascontiguousarray(np.delete(self.__phasespace, indices, 1)) - + # filter out nans def remove_nans(self): del self[np.isnan(self).any(axis=1)] - + # 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): - + # make empty phase space num_particles = len(xs) self.reset_phase_space(num_particles) - + # add positions self.set_xs(xs) self.set_ys(ys) self.set_zs(zs) - + # add momenta if uzs is None: if pzs is not None: @@ -70,57 +70,57 @@ def set_phase_space(self, Q, xs, ys, zs, uxs=None, uys=None, uzs=None, pxs=None, elif Es is not None: uzs = energy2proper_velocity(Es) self.__phasespace[5,:] = uzs - + if uxs is None: if pxs is not None: uxs = momentum2proper_velocity(pxs) elif xps is not None: uxs = xps * uzs self.__phasespace[3,:] = uxs - + if uys is None: if pys is not None: uys = momentum2proper_velocity(pys) elif yps is not None: uys = yps * uzs self.__phasespace[4,:] = uys - + # charge if weightings is None: self.__phasespace[6,:] = Q/num_particles else: self.__phasespace[6,:] = Q*weightings/np.sum(weightings) - + # ids self.__phasespace[7,:] = np.arange(num_particles) # single particle mass [kg] self.particle_mass = particle_mass - - + + # addition operator (add two beams using the + operator) def __add__(self, beam): return Beam(phasespace = np.append(self.__phasespace, beam.__phasespace, axis=1)) - + # in-place addition operator (add one beam to another using the += operator) def __iadd__(self, beam): - if beam is not None: + if beam is not None: self.__phasespace = np.append(self.__phasespace, beam.__phasespace, axis=1) return self - + # indexing operator (get single particle out) def __getitem__(self, index): return self.__phasespace[:,index] - + # "length" operator (number of macroparticles) def __len__(self): return self.__phasespace.shape[1] - + # string operator (called when printing) def __str__(self): return f"Beam: {len(self)} macroparticles, {self.charge()*1e9:.2f} nC, {self.energy()/1e9:.2f} GeV" - - + + ## BUNCH PATTERN def bunch_frequency(self) -> float: @@ -138,11 +138,11 @@ def train_duration(self) -> float: return None else: return self.bunch_separation * (self.num_bunches_in_train-1) - + def average_current_train(self) -> float: return self.charge()*self.bunch_frequency() - + ## BEAM ARRAYS # get phase space variables @@ -162,7 +162,7 @@ def qs(self): return self.__phasespace[6,:] def ids(self): return self.__phasespace[7,:] - + # set phase space variables def set_xs(self, xs): self.__phasespace[0,:] = xs @@ -176,22 +176,22 @@ def set_uys(self, uys): self.__phasespace[4,:] = uys def set_uzs(self, uzs): self.__phasespace[5,:] = uzs - + def set_xps(self, xps): self.set_uxs(xps*self.uzs()) def set_yps(self, yps): self.set_uys(yps*self.uzs()) def set_Es(self, Es): self.set_uzs(energy2proper_velocity(Es)) - + def set_qs(self, qs): self.__phasespace[6,:] = qs def set_ids(self, ids): self.__phasespace[7,:] = ids - + def weightings(self): return self.__phasespace[6,:]/(self.charge_sign()*SI.e) - + # copy another beam's macroparticle charge def copy_particle_charge(self, beam): self.set_qs(np.median(beam.qs())) @@ -201,18 +201,18 @@ def scale_charge(self, Q): def scale_energy(self, E): self.set_Es((E/self.energy())*self.Es()) - - + + def rs(self): return np.sqrt(self.xs()**2 + self.ys()**2) - + def pxs(self): return proper_velocity2momentum(self.uxs()) def pys(self): return proper_velocity2momentum(self.uys()) def pzs(self): return proper_velocity2momentum(self.uzs()) - + def xps(self): return self.uxs()/self.uzs() def yps(self): @@ -226,11 +226,11 @@ def deltas(self, pz0=None): if pz0 is None: pz0 = np.mean(self.pzs()) return self.pzs()/pz0 -1 - - + + def ts(self): return self.zs()/SI.c - + # vector of transverse positions and angles: (x, x', y, y') def transverse_vector(self): vector = np.zeros((4,len(self))) @@ -239,13 +239,13 @@ def transverse_vector(self): vector[2,:] = self.ys() vector[3,:] = self.yps() return vector - + # set phase space based on transverse vector: (x, x', y, y') def set_transverse_vector(self, vector): self.set_xs(vector[0,:]) self.set_xps(vector[1,:]) self.set_ys(vector[2,:]) - self.set_yps(vector[3,:]) + self.set_yps(vector[3,:]) def norm_transverse_vector(self): vector = np.zeros((4,len(self))) @@ -261,18 +261,18 @@ def norm_transverse_vector(self): 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): """ Rotates the coordinate system (passive transformation) of the beam first with angle1 around axis1, then with angle2 around axis2 and lastly with angle3 around axis3. - + Parameters ---------- axis1, axis2, axis3: 1x3 float nd arrays Unit vectors specifying the rotation axes. - + angle1, angle2, angle3: [rad] float Angles used for rotation of the beam's coordinate system around the respective axes. invert: bool Performs a standard passive transformation when False. If True, will perform an active transformation and can thus be used to invert the passive transformation. - + Returns ---------- Modified beam xs, ys, zs, uxs, uys and uzs. @@ -284,7 +284,7 @@ def rotate_coord_sys_3D(self, axis1, angle1, axis2=np.array([0, 1, 0]), angle2=0 if angle1 < -np.pi or angle1 > np.pi or angle2 < -np.pi or angle2 > np.pi or angle3 < -np.pi or angle3 > np.pi: raise ValueError('The rotation angles have to be in the interval [-pi, pi].') - + # Combine into (N, 3) arrays zs = self.zs() xs = self.xs() @@ -317,23 +317,23 @@ def rotate_coord_sys_3D(self, axis1, angle1, axis2=np.array([0, 1, 0]), angle2=0 self.set_uzs(rotated_uzs) self.set_uxs(rotated_uxs) self.set_uys(rotated_uys) - + # ================================================== def beam_alignment_angles(self): """ Calculates the angles for rotation around the y- and x-axis to align the z-axis to the beam proper velocity. - + Parameters ---------- ... - - + + Returns ---------- x_angle: [rad] float Used for rotating the beam's frame around the y-axis. - + y_angle: [rad] float Used for rotating the beam's frame around the x-axis. 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(). """ @@ -348,7 +348,7 @@ def beam_alignment_angles(self): # Calculate the angles to be used for beam rotation z_axis = np.array([1, 0, 0]) # Axis as an unit vector - + zx_projection = point_vec * np.array([1, 1, 0]) # The projection of the pointing vector onto the zx-plane. # Separate treatments for small angles to avoid numerical instability @@ -357,39 +357,39 @@ def beam_alignment_angles(self): else: x_angle = np.sign(point_vec[1]) * np.arccos( np.dot(zx_projection, z_axis) / np.linalg.norm(zx_projection) ) # The angle between zx_projection and z_axis. - + if np.abs(self.y_angle()) < 1e-4: y_angle = point_vec[2]/np.linalg.norm(zx_projection) else: rotated_zy_projection = np.array([ np.linalg.norm(zx_projection), 0, point_vec[2] ]) # The new pointing vector after its zx-projection has been aligned to the rotated z-axis. - + y_angle = np.sign(point_vec[2]) * np.arccos( np.dot(rotated_zy_projection, z_axis) / np.linalg.norm(rotated_zy_projection) ) # Note that due to the right hand rule, a positive y_angle corresponds to rotation from z-axis towards negative y. I.e. opposite sign convention of yps. return x_angle, y_angle - + # ================================================== 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. - + Parameters ---------- x_angle: [rad] float Angle to rotate the coordinate system with in the zx-plane. - + y_angle: [rad] float Angle to rotate the coordinate system with in the zy-plane. 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(). invert: bool Performs a standard passive transformation when False. If True, will perform an active transformation and can thus be used to invert the passive transformation. - - + + Returns ---------- Modified beam xs, ys, zs, uxs, uys and uzs. """ - + x_axis = np.array([0, 1, 0]) # Axis as an unit vector. Axis permutaton is zxy. y_axis = np.array([0, 0, 1]) @@ -398,24 +398,24 @@ def xy_rotate_coord_sys(self, x_angle=None, y_angle=None, invert=False): if y_angle is None: _, y_angle = self.beam_alignment_angles() y_angle = -y_angle - + self.rotate_coord_sys_3D(y_axis, x_angle, x_axis, y_angle, invert=invert) - + # ================================================== def add_pointing_tilts(self, align_x_angle=None, align_y_angle=None): """ Uses active transformation to tilt the beam in the zx- and zy-planes. - + Parameters ---------- align_x_angle: [rad] float Beam coordinates with in the zx-plane are rotated with this angle. - + 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(). - - + + Returns ---------- Modified beam xs, ys and zs. @@ -429,14 +429,14 @@ def add_pointing_tilts(self, align_x_angle=None, align_y_angle=None): y_axis = np.array([0, 0, 1]) x_axis = np.array([0, 1, 0]) - + zs = self.zs() xs = self.xs() ys = self.ys() # Combine into (N, 3) arrays coords = np.column_stack((zs, xs, ys)) - + # Create the rotation object rotation_y = Rot.from_rotvec(align_x_angle * y_axis) rotation_x = Rot.from_rotvec(align_y_angle * x_axis) @@ -472,7 +472,7 @@ def slice_centroids(self, beam_quant, bin_number=None, cut_off=None, make_plot=F make_plot: bool Flag for making plots. - + Returns ---------- beam_quant_slices: 1D float array @@ -481,7 +481,7 @@ def slice_centroids(self, beam_quant, bin_number=None, cut_off=None, make_plot=F z_centroids: [m] 1D float array z-coordinates of the beam slices. """ - + zs = self.zs() mean_z = self.z_offset() weights = self.weightings() @@ -507,7 +507,7 @@ def slice_centroids(self, beam_quant, bin_number=None, cut_off=None, make_plot=F # Beam slice zs _, edges = np.histogram(zs_roi, bins=bin_number) # Get the edges of the histogram of z with bin_number bins. z_centroids = (edges[0:-1] + edges[1:])/2 # Centres of the beam slices (z) - + # Compute the mean of beam_quant of all particles inside a z-bin beam_quant_centroids = np.empty(len(z_centroids)) for i in range(0,len(edges)-1): @@ -520,77 +520,77 @@ def slice_centroids(self, beam_quant, bin_number=None, cut_off=None, make_plot=F plt.scatter(zs*1e6, beam_quant) plt.plot(z_centroids*1e6, beam_quant_centroids, 'rx-') plt.xlabel(r'$\xi$ [$\mathrm{\mu}$m]') - + return beam_quant_centroids, z_centroids - - + + # ================================================== def x_tilt_angle(self, z_cutoff=None): if z_cutoff is None: z_cutoff = 1.5 * self.bunch_length() - + x_centroids, z_centroids = self.slice_centroids(self.xs(), cut_off=z_cutoff, make_plot=False) - + # Perform linear regression slope, _ = np.polyfit(z_centroids, x_centroids, 1) - + return np.arctan(slope) - + # ================================================== def y_tilt_angle(self, z_cutoff=None): 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) - + # Perform linear regression 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) @@ -599,58 +599,58 @@ def y_offset(self, clean=False): 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)) @@ -680,7 +680,7 @@ def norm_amplitude_x(self, plasma_density=None, clean=False): 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()) @@ -691,14 +691,14 @@ def norm_amplitude_y(self, plasma_density=None, clean=False): 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): @@ -709,9 +709,9 @@ def remove_halo_particles(self, nsigma=20): 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)) @@ -720,37 +720,37 @@ def projected_density(self, fcn, bins=None): 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: @@ -768,57 +768,57 @@ def phase_space_density(self, hfcn, vfcn, hbins=None, vbins=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. def charge_density_3D(self, zbins=None, xbins=None, ybins=None): """ Calculates the 3D charge density. - + Parameters ---------- zbins, xbins, ybins: [m] float or 1D float ndarray The bins along z(x,y). - + Returns ---------- - dQ_dxdydz: [C/m^3] 3D float ndarray + dQ_dxdydz: [C/m^3] 3D float ndarray Charge density of the beam. - - zctrs, xctrs, yctrs: [m] 1D float ndarray + + zctrs, xctrs, yctrs: [m] 1D float ndarray The centre positions of the bins of dQ_dxdydz. """ - + zs = self.zs() xs = self.xs() ys = self.ys() - + if zbins is None: zbins = round(np.sqrt(len(self))/2) if xbins is None: xbins = round(np.sqrt(len(self))/2) if ybins is None: ybins = round(np.sqrt(len(self))/2) - + # Create a 3D histogram counts, edges = np.histogramdd((zs, xs, ys), bins=(zbins, xbins, ybins), weights=self.qs()) edges_z = edges[0] edges_x = edges[1] edges_y = edges[2] - + # 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). - + # Calculate charge density per unit volume with np.errstate(divide='ignore', invalid='ignore'): # Handle division by zero dQ_dzdxdy = np.divide(counts, bin_volumes, out=np.zeros_like(counts), where=bin_volumes != 0) @@ -830,16 +830,16 @@ def charge_density_3D(self, zbins=None, xbins=None, ybins=None): xctrs = (edges_x[0:-1] + edges_x[1:])/2 yctrs = (edges_y[0:-1] + edges_y[1:])/2 - #print(np.sum(dQ_dzdxdy*bin_volumes)/self.charge()) - + #print(np.sum(dQ_dzdxdy*bin_volumes)/self.charge()) + return dQ_dzdxdy, zctrs, xctrs, yctrs, edges_z, edges_x, edges_y - + # ================================================== def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inner_off_diag, upper_outer_off_diag, lower_outer_off_diag, num_x_cells, num_unknowns, rhs, boundary_val): """ Applies Dirichlet boundary conditions and assemble the system matrix and the right hand side (source term) of the Poisson equation. - + Parameters ---------- main_diag: [m^-2] 1D float ndarray @@ -856,7 +856,7 @@ def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inne outer_inner_off_diag: [m^-2] 1D float ndarray The outer inne off-diagonal of the system matrix to be modified according to the boundary conditions. - + num_x_cells: float The number of cells in the x-direction. Determines the number of columns of the system matrix A. @@ -869,7 +869,7 @@ def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inne boundary_val: [V/m] float The value of the electric fields Ex and Ey at the simulation box boundary. - + Returns ---------- A: [m^-2] 2D float sparse matrix @@ -878,15 +878,15 @@ def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inne rhs: [V/m^3] 1D float ndarray The modified right hand side of the Poisson equation. """ - + # Set the right side boundary conditions rhs[num_x_cells-1::num_x_cells] = boundary_val # Set BC. Set every num_x_cells-th element starting from the num_x_cells-1 index to 1 - main_diag[num_x_cells-1::num_x_cells] = 1 # Set BC + main_diag[num_x_cells-1::num_x_cells] = 1 # Set BC upper_inner_off_diag[num_x_cells-1::num_x_cells] = 0 # Remove off-diagonal elements at boundaries lower_inner_off_diag[num_x_cells-2::num_x_cells] = 0 # Remove off-diagonal elements at boundaries upper_outer_off_diag[num_x_cells-1::num_x_cells] = 0 # Remove off-diagonal elements at boundaries lower_outer_off_diag[-1::-num_x_cells] = 0 # Remove off-diagonal elements at boundaries - + # Set the left side boundary conditions rhs[0::num_x_cells] = boundary_val main_diag[0::num_x_cells] = 1 @@ -894,14 +894,14 @@ def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inne lower_inner_off_diag[-num_x_cells::-num_x_cells] = 0 upper_outer_off_diag[0::num_x_cells] = 0 lower_outer_off_diag[-num_x_cells::-num_x_cells] = 0 - + # Set the top boundary conditions rhs[1:num_x_cells-1] = boundary_val main_diag[1:num_x_cells-1] = 1 upper_inner_off_diag[1:num_x_cells-1] = 0 lower_inner_off_diag[0:num_x_cells-2] = 0 upper_outer_off_diag[1:num_x_cells-1] = 0 - + # Set the bottom boundary conditions rhs[-num_x_cells+1:-1] = boundary_val main_diag[-num_x_cells+1:-1] = 1 @@ -915,7 +915,7 @@ def Dirichlet_BC_system_matrix(self, main_diag, upper_inner_off_diag, lower_inne A = sp.diags(diagonals, offsets, shape=(num_unknowns, num_unknowns), format="csr") return A, rhs - + # ================================================== def Ex_Ey_2D(self, num_x_cells, num_y_cells, charge_density_xy_slice, dx, dy, boundary_val=0.0): @@ -924,8 +924,8 @@ def Ex_Ey_2D(self, num_x_cells, num_y_cells, charge_density_xy_slice, dx, dy, bo dEx/dx + dEy/dy = 1/epsilon_0 * dQ/dzdxdy dEy/dx - dEx/dy = 0. - - + + Parameters ---------- num_x_cells, num_y_cells: float @@ -937,18 +937,18 @@ def Ex_Ey_2D(self, num_x_cells, num_y_cells, charge_density_xy_slice, dx, dy, bo dx, dy: [m] float Bin widths in x and y of the bins of dQ_dzdxdy. - boundary_val: [V/m] - + boundary_val: [V/m] + Returns ---------- - Ex: [V/m] 2D float array + Ex: [V/m] 2D float array x-conponent of electric field generated by the chosen beam slice. - Ey: [V/m] 2D float array + Ey: [V/m] 2D float array y-conponent of electric field generated by the chosen beam slice. """ - + num_rows, num_cols = charge_density_xy_slice.shape # Set up the system matrix @@ -995,7 +995,7 @@ def Ex_Ey(self, x_box_min, x_box_max, y_box_min, y_box_max, dx, dy, num_z_cells= x_box_max, y_box_max: [m] float The upper x(y) boundary of the simulation domain. Should be much larger than the plasma bubble radius. - + dx, dy: [m] float Bin widths in x and y of the bins of dQ_dzdxdy. @@ -1005,13 +1005,13 @@ def Ex_Ey(self, x_box_min, x_box_max, y_box_min, y_box_max, dx, dy, num_z_cells= boundary_val: [V/m] The values of the electric fields Ex and Ey at the simulation domain boundary. - + Returns ---------- - Ex: [V/m] 2D float array + Ex: [V/m] 2D float array x-conponent of electric field generated by the chosen beam slice. - Ey: [V/m] 2D float array + Ey: [V/m] 2D float array y-conponent of electric field generated by the chosen beam slice. zctrs, xctrs, yctrs: [m] 1D float ndarray @@ -1021,25 +1021,25 @@ def Ex_Ey(self, x_box_min, x_box_max, y_box_min, y_box_max, dx, dy, num_z_cells= # Check if the selected simulation boundaries are significantly larger than the beam extent xs = self.xs() ys = self.ys() - + if np.abs(x_box_min/xs.min()) < tolerance or np.abs(y_box_min/ys.min()) < tolerance or np.abs(x_box_max/xs.max()) < tolerance or np.abs(y_box_max/ys.max()) < tolerance: raise ValueError('Simulation box size is too small compared to beam size.') - + if num_z_cells is None: num_z_cells = round(np.sqrt(len(self))/2) z_box_max = np.max(self.zs()) z_box_min = np.min(self.zs()) zbins = np.linspace(z_box_min, z_box_max, num_z_cells+1) - + num_x_cells = int((x_box_max-x_box_min)/dx) num_y_cells = int((y_box_max-y_box_min)/dy) - + xbins = np.linspace(x_box_min, x_box_max, num_x_cells+1) ybins = np.linspace(y_box_min, y_box_max, num_y_cells+1) dQ_dzdxdy, zctrs, xctrs, yctrs, edges_z, edges_x, edges_y = self.charge_density_3D(zbins=zbins, xbins=xbins, ybins=ybins) - + Ex = np.zeros((num_z_cells, num_y_cells, num_x_cells)) Ey = np.zeros((num_z_cells, num_y_cells, num_x_cells)) @@ -1049,58 +1049,58 @@ def Ex_Ey(self, x_box_min, x_box_max, y_box_min, y_box_max, dx, dy, num_z_cells= # Calculate the fields for the charge density slice Ex_2d, Ey_2d = self.Ex_Ey_2D(num_x_cells, num_y_cells, charge_density_xy_slice, dx, dy, boundary_val=boundary_val) - + Ex[slice_idx,:,:] = Ex_2d Ey[slice_idx,:,:] = Ey_2d - + return Ex, Ey, zctrs, xctrs, yctrs - - - + + + ## PLOTTING def plot_current_profile(self): dQdt, ts = self.current_profile() fig, ax = plt.subplots() fig.set_figwidth(6) - fig.set_figheight(4) + 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) - + 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) + 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) + 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)') @@ -1108,52 +1108,49 @@ def plot_trace_space_y(self): cb = fig.colorbar(p) cb.ax.set_ylabel('Charge density (pC/um/mrad)') - def plot_transverse_profile(self, xlims=(None, None), ylims=(None, None), title=''): + 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], + #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)') - ax.set_xlim(xlims) - ax.set_ylim(ylims) - ax.set_title(title) - + def plot_bunch_pattern(self): - + fig, ax = plt.subplots() fig.set_figwidth(6) - fig.set_figheight(4) + 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): - + # add energy and chirp - Es = self.Es() + energy_gain + Es = self.Es() + energy_gain Es = Es + np.sign(self.qs()) * (self.zs()-z_offset) * chirp self.set_uzs(energy2proper_velocity(Es)) - + # remove particles with subzero energy del self[Es < 0] - + def compress(self, R_56, nom_energy): zs = self.zs() + (1-self.Es()/nom_energy) * R_56 self.set_zs(zs) - + def scale_to_length(self, bunch_length): z_mean = self.z_offset() zs_scaled = z_mean + (self.zs()-z_mean)*bunch_length/self.bunch_length() @@ -1168,17 +1165,17 @@ def scale_norm_emittance_y(self, norm_emit_ny): 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)) - + # betatron damping (must be done before acceleration) def apply_betatron_damping(self, deltaE): gammasBoosted = energy2gamma(abs(self.Es()+deltaE)) betamag = np.sqrt(self.gammas()/gammasBoosted) self.magnify_beta_function(betamag) - - + + # magnify beta function (increase beam size, decrease divergence) def magnify_beta_function(self, beta_mag, axis_defining_beam=None): - + # calculate beam (not beta) magnification mag = np.sqrt(beta_mag) @@ -1192,19 +1189,19 @@ def magnify_beta_function(self, beta_mag, axis_defining_beam=None): y_offset = axis_defining_beam.y_offset() ux_offset = axis_defining_beam.ux_offset() uy_offset = axis_defining_beam.uy_offset() - + self.set_xs((self.xs()-x_offset)*mag + x_offset) self.set_ys((self.ys()-y_offset)*mag + y_offset) self.set_uxs((self.uxs()-ux_offset)/mag + ux_offset) self.set_uys((self.uys()-uy_offset)/mag + uy_offset) - + # transport in a drift def transport(self, L): self.set_xs(self.xs() + L*self.xps()) self.set_ys(self.ys() + L*self.yps()) - + def flip_transverse_phase_spaces(self, flip_momenta=True, flip_positions=False): if flip_momenta: self.set_uxs(-self.uxs()) @@ -1213,9 +1210,9 @@ def flip_transverse_phase_spaces(self, flip_momenta=True, flip_positions=False): self.set_xs(-self.xs()) 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): - + # remove particles with subzero and Nan energy neg_indices = self.Es() < 0 del self[neg_indices] @@ -1226,15 +1223,15 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat del self[nan_indices] if isinstance(deltaEs, np.ndarray) and len(deltaEs) > 1: deltaEs = deltaEs[~nan_indices] - + # determine initial and final Lorentz factor gamma0s = energy2gamma(self.Es()) Es_final = self.Es() + deltaEs gammas = energy2gamma(Es_final) dgamma_ds = (gammas-gamma0s)/L - + if calc_evolution: - + # calculate evolution num_evol_steps = max(20, min(400, round(2*L/(beta_matched(n0, self.energy()+min(0,np.mean(deltaEs))))))) evol = SimpleNamespace() @@ -1251,13 +1248,13 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat # 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) @@ -1274,15 +1271,15 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat 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) - - + + # calculate final positions and angles after betatron motion if radiation_reaction: xs, uxs, ys, uys, Es_final = evolve_betatron_motion(self.xs()-x0_driver, self.uxs(), self.ys()-y0_driver, self.uys(), L, gamma0s, dgamma_ds, k_p(n0)) else: xs, uxs = evolve_hills_equation_analytic(self.xs()-x0_driver, self.uxs(), L, gamma0s, dgamma_ds, k_p(n0)) ys, uys = evolve_hills_equation_analytic(self.ys()-y0_driver, self.uys(), L, gamma0s, dgamma_ds, k_p(n0)) - + # set new beam positions and angles (shift back driver offsets) self.set_xs(xs+x0_driver) self.set_uxs(uxs) @@ -1293,31 +1290,31 @@ def apply_betatron_motion(self, L, n0, deltaEs, x0_driver=0, y0_driver=0, radiat return Es_final, evol else: return Es_final - - + + ## 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') @@ -1325,16 +1322,16 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): # 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) @@ -1349,10 +1346,10 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): 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) @@ -1364,10 +1361,10 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): 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['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()) @@ -1382,42 +1379,42 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): 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() @@ -1428,16 +1425,16 @@ def load(_, filename, beam_name='beam'): 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") @@ -1451,13 +1448,13 @@ def load(_, filename, beam_name='beam'): 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): - + if axes is None: fig = plt.figure() # an empty figure with an axes ax = fig.add_axes([.15, .15, .75, .75]) @@ -1465,7 +1462,7 @@ def imshow_plot(self, data, axes=None, extent=None, vmin=None, vmax=None, colmap else: #ax = axes[0] # TODO: adjust colourbar axes #cbar_ax = axes[1] - + ax = axes cbar_ax = None @@ -1512,7 +1509,7 @@ def imshow_plot(self, data, axes=None, extent=None, vmin=None, vmax=None, colmap ax.set_ylabel(ylab) ax.set_xlabel(xlab) - + # ================================================== def distribution_plot_2D(self, arr1, arr2, weights=None, hist_bins=None, hist_range=None, axes=None, extent=None, vmin=None, vmax=None, colmap=CONFIG.default_cmap, xlab='', ylab='', clab='', origin='lower', interpolation='nearest', reduce_cax_pad=False): @@ -1527,20 +1524,20 @@ def distribution_plot_2D(self, arr1, arr2, weights=None, hist_bins=None, hist_ra hist_range[1] = [ arr2.min(), arr2.max() ] if extent is None: extent = hist_range[0] + hist_range[1] - + binned_data, zedges, xedges = np.histogram2d(arr1, arr2, hist_bins, hist_range, weights=weights) beam_hist2d = binned_data.T/np.diff(zedges)/np.diff(xedges) self.imshow_plot(beam_hist2d, axes=axes, extent=extent, vmin=vmin, vmax=vmax, colmap=colmap, xlab=xlab, ylab=ylab, clab=clab, gridOn=False, origin=origin, interpolation=interpolation, reduce_cax_pad=reduce_cax_pad) - + # ================================================== def density_map_diags(self): - + #colors = ['white', 'aquamarine', 'lightgreen', 'green'] #colors = ['white', 'forestgreen', 'limegreen', 'lawngreen', 'aquamarine', 'deepskyblue'] #bounds = [0, 0.2, 0.4, 0.8, 1] #cmap = LinearSegmentedColormap.from_list('my_cmap', colors, N=256) - + cmap = CONFIG.default_cmap # Macroparticles data @@ -1560,7 +1557,7 @@ def density_map_diags(self): xps_lab = r"$x'$ [mrad]" yps_lab = r"$y'$ [mrad]" energ_lab = r'$\mathcal{E}$ [GeV]' - + # Set up a figure with axes fig, axs = plt.subplots(nrows=3, ncols=3, layout='constrained', figsize=(5*3, 4*3)) fig.suptitle(r'$\Delta s=$' f'{format(self.location, ".2f")}' ' m') @@ -1576,7 +1573,7 @@ def density_map_diags(self): extent_zx = [i*1e6 for i in extent_zx] # [um] self.distribution_plot_2D(arr1=zs, arr2=xs, weights=weights, hist_bins=hist_bins, hist_range=hist_range, axes=axs[0][0], extent=extent_zx, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=xlab, clab=r'$\partial^2 N/\partial\xi \partial x$ [$\mathrm{m}^{-2}$]', origin='lower', interpolation='nearest') - + # 2D z-x' distribution hist_range_xps = [[None, None], [None, None]] @@ -1589,8 +1586,8 @@ def density_map_diags(self): extent_xps[3] = extent_xps[3]*1e3 # [mrad] self.distribution_plot_2D(arr1=zs, arr2=xps, weights=weights, hist_bins=hist_bins, hist_range=hist_range_xps, axes=axs[0][1], extent=extent_xps, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=xps_lab, clab=r"$\partial^2 N/\partial z \partial x'$ [$\mathrm{m}^{-1}$ $\mathrm{rad}^{-1}$]", origin='lower', interpolation='nearest') - - + + # 2D x-x' distribution hist_range_xxp = [[None, None], [None, None]] hist_range_xxp[0] = hist_range[1] @@ -1602,7 +1599,7 @@ def density_map_diags(self): extent_xxp[3] = extent_xxp[3]*1e3 # [mrad] self.distribution_plot_2D(arr1=xs, arr2=xps, weights=weights, hist_bins=hist_bins, hist_range=hist_range_xxp, axes=axs[0][2], extent=extent_xxp, vmin=None, vmax=None, colmap=cmap, xlab=xlab, ylab=xps_lab, clab=r"$\partial^2 N/\partial x\partial x'$ [$\mathrm{m}^{-1}$ $\mathrm{rad}^{-1}$]", origin='lower', interpolation='nearest') - + # 2D z-y distribution hist_range_zy = [[None, None], [None, None]] @@ -1612,7 +1609,7 @@ def density_map_diags(self): extent_zy = [i*1e6 for i in extent_zy] # [um] self.distribution_plot_2D(arr1=zs, arr2=ys, weights=weights, hist_bins=hist_bins, hist_range=hist_range_zy, axes=axs[1][0], extent=extent_zy, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=ylab, clab=r'$\partial^2 N/\partial\xi \partial y$ [$\mathrm{m}^{-2}$]', origin='lower', interpolation='nearest') - + # 2D z-y' distribution hist_range_yps = [[None, None], [None, None]] @@ -1623,9 +1620,9 @@ def density_map_diags(self): extent_yps[1] = extent_yps[1]*1e6 # [um] extent_yps[2] = extent_yps[2]*1e3 # [mrad] extent_yps[3] = extent_yps[3]*1e3 # [mrad] - + self.distribution_plot_2D(arr1=zs, arr2=yps, weights=weights, hist_bins=hist_bins, hist_range=hist_range_yps, axes=axs[1][1], extent=extent_yps, vmin=None, vmax=None, colmap=cmap, xlab=xilab, ylab=yps_lab, clab=r"$\partial^2 N/\partial z \partial y'$ [$\mathrm{m}^{-1}$ $\mathrm{rad}^{-1}$]", origin='lower', interpolation='nearest') - + # 2D y-y' distribution hist_range_yyp = [[None, None], [None, None]] @@ -1636,9 +1633,9 @@ def density_map_diags(self): extent_yyp[1] = extent_yyp[1]*1e6 # [um] extent_yyp[2] = extent_yyp[2]*1e3 # [mrad] extent_yyp[3] = extent_yyp[3]*1e3 # [mrad] - + self.distribution_plot_2D(arr1=ys, arr2=yps, weights=weights, hist_bins=hist_bins, hist_range=hist_range_yyp, axes=axs[1][2], extent=extent_yyp, vmin=None, vmax=None, colmap=cmap, xlab=ylab, ylab=yps_lab, clab=r"$\partial^2 N/\partial y\partial y'$ [$\mathrm{m}^{-1}$ $\mathrm{rad}^{-1}$]", origin='lower', interpolation='nearest') - + # 2D x-y distribution hist_range_xy = [[None, None], [None, None]] @@ -1648,7 +1645,7 @@ def density_map_diags(self): extent_xy = [i*1e6 for i in extent_xy] # [um] self.distribution_plot_2D(arr1=xs, arr2=ys, weights=weights, hist_bins=hist_bins, hist_range=hist_range_xy, axes=axs[2][0], extent=extent_xy, vmin=None, vmax=None, colmap=cmap, xlab=xlab, ylab=ylab, clab=r'$\partial^2 N/\partial x \partial y$ [$\mathrm{m}^{-2}$]', origin='lower', interpolation='nearest') - + # Energy distribution ax = axs[2][1] From 7cc40dc91c2eec47a5ab894570b534b111004ad4 Mon Sep 17 00:00:00 2001 From: Kyrre Ness Sjobak Date: Thu, 3 Apr 2025 16:12:38 +0200 Subject: [PATCH 09/11] Re-add the actual changes that was done in the whitespace-deletion-commit --- abel/classes/beam.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/abel/classes/beam.py b/abel/classes/beam.py index d1d72215..c41f23d2 100644 --- a/abel/classes/beam.py +++ b/abel/classes/beam.py @@ -1108,7 +1108,7 @@ def plot_trace_space_y(self): cb = fig.colorbar(p) cb.ax.set_ylabel('Charge density (pC/um/mrad)') - def plot_transverse_profile(self): + def plot_transverse_profile(self, xlims=(None, None), ylims=(None, None), title=''): dQdxdy, xs, ys = self.phase_space_density(self.xs, self.ys) fig, ax = plt.subplots() @@ -1123,6 +1123,9 @@ def plot_transverse_profile(self): cb = fig.colorbar(p) cb.ax.set_ylabel('Charge density (pC/um^2)') + ax.set_xlim(xlims) + ax.set_ylim(ylims) + ax.set_title(title) def plot_bunch_pattern(self): From 680d785349bd6bdcf7b67c93882ad92d3ba0929a Mon Sep 17 00:00:00 2001 From: Felipe Pena Date: Thu, 3 Apr 2025 19:26:24 +0300 Subject: [PATCH 10/11] Updated documentation --- .../classes/source/impl/source_energy_comb.py | 84 ++++++++----- .../impl/spectrometer_achromatic_basic.py | 96 +++++++------- .../achromatic/spectrometer_achromatic.py | 118 +++++++++++------- 3 files changed, 171 insertions(+), 127 deletions(-) diff --git a/abel/classes/source/impl/source_energy_comb.py b/abel/classes/source/impl/source_energy_comb.py index 1b4e2b2f..c59bb94c 100644 --- a/abel/classes/source/impl/source_energy_comb.py +++ b/abel/classes/source/impl/source_energy_comb.py @@ -1,3 +1,25 @@ +""" +Module: source_energy_comb + +This module defines the `SourceEnergyComb` class, which is a subclass +of `Source`. It provides methods to initialize the source and track +the beam with comb-like energy distributions. + +Classes +------- +`SourceEnergyComb` : `Source` + A source class for creating comb-like energy distributions. + +Imports +------- +import time +import numpy as np +import scipy.constants as SI +from abel import Source, Beam +from abel.utilities.beam_physics import generate_trace_space_xy, generate_symm_trace_space_xyz +from abel.utilities.relativity import energy2gamma +""" + import time import numpy as np import scipy.constants as SI @@ -9,45 +31,45 @@ class SourceEnergyComb(Source): """ A source class that extends the functionality of `SourceBasic` by distributing particle energies within a given list of energies. This is useful for creating - comb-like energy distributions for resolution estimates. +comb-like energy distributions for resolution estimates. Parameters ---------- - length : float, optional - Length of the source [m]. Default is 0. + length : [m] float, optional + Length of the source. Default is 0. num_particles : int, optional Number of particles in the beam. Default is 1000. - energy : float, optional - Central energy of the beam [eV]. Default is None. + energy : [eV] float, optional + Central energy of the beam. Default is ``None``. - charge : float, optional - Total charge of the beam [C]. Default is 0. + charge : [C] float, optional + Total charge of the beam. Default is 0. rel_energy_spread : float, optional - Relative energy spread (fraction). Default is None. + Relative energy spread (fraction). Default is ``None``. - energy_spread : float, optional - Absolute energy spread [eV]. Default is None. + energy_spread : [eV] float, optional + Absolute energy spread. Default is ``None``. - bunch_length : float, optional - RMS bunch length [m]. Default is None. + bunch_length : [m] float, optional + RMS bunch length. Default is ``None``. - z_offset : float, optional - Longitudinal offset of the beam [m]. Default is 0. + z_offset : [m] float, optional + Longitudinal offset of the beam. Default is 0. - x_offset, y_offset : float, optional - Transverse offsets of the beam [m]. Default is 0. + x_offset, y_offset : [m] float, optional + Transverse offsets of the beam. Default is 0. - x_angle, y_angle : float, optional - Transverse angles of the beam [rad]. Default is 0. + x_angle, y_angle : [rad] float, optional + Transverse angles of the beam. Default is 0. - emit_nx, emit_ny : float, optional - Normalized emittances in x and y [m rad]. Default is 0. + emit_nx, emit_ny : [m rad] float, optional + Normalized emittances in x and y. Default is 0. - beta_x, beta_y : float, optional - Beta functions in x and y [m]. Default is None. + beta_x, beta_y : [m] float, optional + Beta functions in x and y. Default is ``None``. alpha_x, alpha_y : float, optional Alpha functions in x and y. Default is 0. @@ -58,8 +80,8 @@ class SourceEnergyComb(Source): wallplug_efficiency : float, optional Efficiency of the source. Default is 1. - accel_gradient : float, optional - Acceleration gradient [V/m]. Default is None. + accel_gradient : [V/m] float, optional + Acceleration gradient. Default is ``None``. symmetrize : bool, optional Whether to symmetrize the transverse phase space. Default is False. @@ -67,11 +89,11 @@ class SourceEnergyComb(Source): symmetrize_6d : bool, optional Whether to symmetrize the full 6D phase space. Default is False. - z_cutoff : float, optional - Longitudinal cutoff for filtering particles [m]. Default is None. + z_cutoff : [m] float, optional + Longitudinal cutoff for filtering particles. Default is ``None``. energy_comb_delta : float, optional - Relative energy step size for creating comb-like energy distributions. Default is None. + Relative energy step size for creating comb-like energy distributions. Default is ``None``. """ def __init__(self, length=0, num_particles=1000, energy=None, charge=0, rel_energy_spread=None, energy_spread=None, bunch_length=None, z_offset=0, x_offset=0, y_offset=0, x_angle=0, y_angle=0, emit_nx=0, emit_ny=0, beta_x=None, beta_y=None, alpha_x=0, alpha_y=0, angular_momentum=0, wallplug_efficiency=1, accel_gradient=None, symmetrize=False, symmetrize_6d=False, z_cutoff=None, energy_comb_delta=None): @@ -109,14 +131,14 @@ def track(self, _=None, savedepth=0, runnable=None, verbose=False): Parameters ---------- - _ : None, optional - Placeholder parameter. Default is None. + _ : ``None``, optional + Placeholder parameter. Default is ``None``. savedepth : int, optional Depth of saved tracking data. Default is 0. runnable : object, optional - Runnable object for tracking. Default is None. + Runnable object for tracking. Default is ``None``. verbose : bool, optional Whether to print verbose output. Default is False. @@ -194,7 +216,7 @@ def reorder_particle_energies(self, beam, plot=False): Returns ------- - None + ``None`` """ # takes in the beam and uses the attribute self.energy_comb_delta diff --git a/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py b/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py index 83c2bbd1..74595b97 100644 --- a/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py +++ b/abel/classes/spectrometer/impl/achromatic/impl/spectrometer_achromatic_basic.py @@ -7,7 +7,7 @@ Classes ------- -SpectrometerAchromaticBasic : SpectrometerAchromatic +`SpectrometerAchromaticBasic` : `SpectrometerAchromatic` A basic implementation of an achromatic spectrometer. Imports @@ -27,35 +27,35 @@ class SpectrometerAchromaticBasic(SpectrometerAchromatic): Parameters ---------- - length_drift_to_dipole : float, optional - Length of the drift space to the dipole magnet [m]. Default is 0.2. + length_drift_to_dipole : [m] float, optional + Length of the drift space to the dipole magnet. Default is 0.2. - field_dipole : float, optional - Magnetic field strength of the dipole magnet [T]. Negative values + field_dipole : [T] float, optional + Magnetic field strength of the dipole magnet. Negative values bend the beam downwards. Default is -0.85. - length_dipole : float, optional - Length of the dipole magnet [m]. Default is 0.15. + length_dipole : [m] float, optional + Length of the dipole magnet. Default is 0.15. - length_drift_dipole_to_lens : float, optional - Length of the drift space from the dipole magnet to the lens [m]. + length_drift_dipole_to_lens : [m] float, optional + Length of the drift space from the dipole magnet to the lens. Default is 0.15. - length_plasma_lens : float, optional - Length of the plasma lens [m]. Default is 20E-3. + length_plasma_lens : [m] float, optional + Length of the plasma lens. Default is 20E-3. - radius_plasma_lens : float, optional - Radius of the plasma lens [m]. Default is 1E-3. + radius_plasma_lens : [m] float, optional + Radius of the plasma lens. Default is 1E-3. - length_drift_lens_to_screen : float, optional - Length of the drift space from the lens to the screen [m]. Default + length_drift_lens_to_screen : [m] float, optional + Length of the drift space from the lens to the screen. Default is 0.4. - imaging_energy : float or None, optional - Imaging energy [eV]. Default is None. + imaging_energy : [eV] float or None, optional + Imaging energy. Default is None. disable_lens_nonlinearity : bool, optional - Flag to disable lens nonlinearity. Default is True. + Flag to disable lens nonlinearity. Default is `True`. """ def __init__(self, @@ -74,36 +74,36 @@ def __init__(self, Parameters ---------- - length_drift_to_dipole : float, optional - Length of the drift space to the dipole magnet [m]. Default is + length_drift_to_dipole : [m] float, optional + Length of the drift space to the dipole magnet. Default is 0.2. - field_dipole : float, optional - Magnetic field strength of the dipole magnet [T]. Negative + field_dipole : [T] float, optional + Magnetic field strength of the dipole magnet. Negative values bend the beam downwards. Default is -0.85. - length_dipole : float, optional - Length of the dipole magnet [m]. Default is 0.15. + length_dipole : [m] float, optional + Length of the dipole magnet. Default is 0.15. - length_drift_dipole_to_lens : float, optional - Length of the drift space from the dipole magnet to the lens - [m]. Default is 0.15. + length_drift_dipole_to_lens : [m] float, optional + Length of the drift space from the dipole magnet to the lens. + Default is 0.15. - length_plasma_lens : float, optional - Length of the plasma lens [m]. Default is 20E-3. + length_plasma_lens : [m] float, optional + Length of the plasma lens. Default is 20E-3. - radius_plasma_lens : float, optional - Radius of the plasma lens [m]. Default is 1E-3. + radius_plasma_lens : [m] float, optional + Radius of the plasma lens. Default is 1E-3. - length_drift_lens_to_screen : float, optional - Length of the drift space from the lens to the screen [m]. + length_drift_lens_to_screen : [m] float, optional + Length of the drift space from the lens to the screen. Default is 0.4. - imaging_energy : float or None, optional - Imaging energy [eV]. Default is None. + imaging_energy : [eV] float or None, optional + Imaging energy. Default is None. disable_lens_nonlinearity : bool, optional - Flag to disable lens nonlinearity. Default is True. + Flag to disable lens nonlinearity. Default is `True`. """ super().__init__(length_drift_to_dipole=length_drift_to_dipole, @@ -122,21 +122,21 @@ def track(self, beam0, savedepth=0, runnable=None, verbose=False): Parameters ---------- - beam0 : Beam + beam0 : `Beam` object The initial beam to be tracked. savedepth : int, optional The depth at which to save the beam state. Default is 0. runnable : callable, optional - A callable to be run during tracking. Default is None. + A callable to be run during tracking. Default is ``None``. verbose : bool, optional - Whether to print verbose output. Default is False. + Whether to print verbose output. Default is `False`. Returns ------- - Beam + `Beam` object The tracked beam object. """ @@ -249,8 +249,8 @@ def get_bending_angle(self, energy): Parameters ---------- - energy : float - The energy of the particle beam [eV]. + energy : [eV] float + The energy of the particle beam. Returns ------- @@ -275,16 +275,16 @@ def get_dispersion(self, energy=None, bend_angle=None, length=None): Parameters ---------- - energy : float, optional - The energy of the particle beam [eV]. Either `energy` or + energy : [eV] float, optional + The energy of the particle beam. Either `energy` or `bend_angle` must be provided. - bend_angle : float, optional - Bending angle of the dipole [rad]. Either `bend_angle` or + bend_angle : [rad] float, optional + Bending angle of the dipole. Either `bend_angle` or `energy` must be provided. - length : float, optional - The distance to use for the calculation [m]. If not provided, + length : [m] float, optional + The distance to use for the calculation. If not provided, it defaults to the distance from the center of the dipole to the screen. diff --git a/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py b/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py index 1283a59e..48b8ab5a 100644 --- a/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py +++ b/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py @@ -1,3 +1,24 @@ +""" +Module: spectrometer_achromatic + +This module defines the `SpectrometerAchromatic` class, which is an +abstract base class for achromatic spectrometers. It provides methods +and properties to model the spectrometer's behavior. + +Classes +------- +`SpectrometerAchromatic` : `Spectrometer` + Abstract base class for achromatic spectrometers. + +Imports +------- +from abc import abstractmethod +from abel import Spectrometer +import numpy as np +import matplotlib.pyplot as plt +from abel import CONFIG +""" + from abc import abstractmethod from abel import Spectrometer import numpy as np @@ -7,7 +28,7 @@ class SpectrometerAchromatic(Spectrometer): """ - Achromatic Spectrometer class. + Abstract base class for achromatic spectrometers. This class models an achromatic spectrometer, which is used to analyze particle beams by bending them through a magnetic field and focusing @@ -15,35 +36,35 @@ class SpectrometerAchromatic(Spectrometer): Parameters ---------- - length_drift_to_dipole : float - Length of the drift space to the dipole magnet [m]. Default is 0.2. + length_drift_to_dipole : [m] float + Length of the drift space to the dipole magnet. Default is 0.2. - field_dipole : float - Magnetic field strength of the dipole magnet [T]. Negative values + field_dipole : [T] float + Magnetic field strength of the dipole magnet. Negative values indicate bending the beam downwards. Default is -0.85. - length_dipole : float - Length of the dipole magnet [m]. Default is 0.8. + length_dipole : [m] float + Length of the dipole magnet. Default is 0.8. - length_drift_dipole_to_lens : float - Length of the drift space from the dipole magnet to the lens [m]. + length_drift_dipole_to_lens : [m] float + Length of the drift space from the dipole magnet to the lens. Default is 0.2. - length_plasma_lens : float - Length of the plasma lens [m]. Default is 20E-3. + length_plasma_lens : [m] float + Length of the plasma lens. Default is 20E-3. - radius_plasma_lens : float - Radius of the plasma lens [m]. Default is 1E-3. + radius_plasma_lens : [m] float + Radius of the plasma lens. Default is 1E-3. - length_drift_lens_to_screen : float - Length of the drift space from the lens to the screen [m]. - Default is 0.4. + length_drift_lens_to_screen : [m] float + Length of the drift space from the lens to the screen. Default + is 0.4. - imaging_energy : float or None - Imaging energy [eV]. Default is None. + imaging_energy : [eV] float or None + Imaging energy. Default is None. disable_lens_nonlinearity : bool - Flag to disable lens nonlinearity. Default is True. + Flag to disable lens nonlinearity. Default is `True`. """ @abstractmethod @@ -58,39 +79,40 @@ def __init__(self, imaging_energy, disable_lens_nonlinearity): """ - Initializes the SpectrometerAchromatic object with the specified parameters. + Initializes the `SpectrometerAchromatic` object with the specified + parameters. Parameters ---------- - length_drift_to_dipole : float - Length of the drift space to the dipole magnet [m]. Default is 0.2. + length_drift_to_dipole : [m] float + Length of the drift space to the dipole magnet. Default is 0.2. - field_dipole : float - Magnetic field strength of the dipole magnet [T]. Negative values + field_dipole : [T] float + Magnetic field strength of the dipole magnet. Negative values indicate bending the beam downwards. Default is -0.85. - length_dipole : float - Length of the dipole magnet [m]. Default is 0.8. + length_dipole : [m] float + Length of the dipole magnet. Default is 0.8. - length_drift_dipole_to_lens : float - Length of the drift space from the dipole magnet to the lens [m]. + length_drift_dipole_to_lens : [m] float + Length of the drift space from the dipole magnet to the lens. Default is 0.2. - length_plasma_lens : float - Length of the plasma lens [m]. Default is 20E-3. + length_plasma_lens : [m] float + Length of the plasma lens. Default is 20E-3. - radius_plasma_lens : float - Radius of the plasma lens [m]. Default is 1E-3. + radius_plasma_lens : [m] float + Radius of the plasma lens. Default is 1E-3. - length_drift_lens_to_screen : float - Length of the drift space from the lens to the screen [m]. + length_drift_lens_to_screen : [m] float + Length of the drift space from the lens to the screen. Default is 0.4. - imaging_energy : float or None - Imaging energy [eV]. Default is None. + imaging_energy : [eV] float or None + Imaging energy. Default is None. disable_lens_nonlinearity : bool - Flag to disable lens nonlinearity. Default is True. + Flag to disable lens nonlinearity. Default is `True`. """ self.length_drift_to_dipole = length_drift_to_dipole # [m] self.field_dipole = field_dipole # [T] @@ -109,11 +131,11 @@ def __init__(self, def get_length(self): """ - Calculate the total length of the spectrometer lattice. + Calculates the total length of the spectrometer lattice. Returns ------- - float + total_length : [m] float Total length of the spectrometer lattice [m]. """ total_length = (self.length_drift_to_dipole @@ -126,11 +148,11 @@ def get_length(self): @property def radius_plasma_lens(self): """ - Get the radius of the plasma lens. + Gets the radius of the plasma lens. Returns ------- - float + self._radius_plasma_lens : [m] float Radius of the plasma lens [m]. """ return self._radius_plasma_lens @@ -138,12 +160,12 @@ def radius_plasma_lens(self): @radius_plasma_lens.setter def radius_plasma_lens(self, value): """ - Set the radius of the plasma lens. + Sets the radius of the plasma lens. Parameters ---------- - value : float - Radius of the plasma lens [m]. Must be greater than or equal to zero. + value : [m] float + Radius of the plasma lens. Must be greater than or equal to zero. Raises ------ @@ -156,7 +178,7 @@ def radius_plasma_lens(self, value): @abstractmethod def track(self, beam, savedepth=0, runnable=None, verbose=False): """ - Track the beam through the spectrometer. + Tracks the beam through the spectrometer. Parameters ---------- @@ -166,11 +188,11 @@ def track(self, beam, savedepth=0, runnable=None, verbose=False): savedepth : int, optional Depth of saved states during tracking. Default is 0. - runnable : object or None, optional - Runnable object for tracking. Default is None. + runnable : callable or None, optional + Runnable object for tracking. Default is ``None``. verbose : bool, optional - Flag to enable verbose output. Default is False. + Flag to enable verbose output. Default is ``False``. Returns ------- From 1ffafea62b54dd7c6d29856de483e0cfdece580c Mon Sep 17 00:00:00 2001 From: Felipe Pena Date: Thu, 3 Apr 2025 19:27:39 +0300 Subject: [PATCH 11/11] Found a typo --- .../spectrometer/impl/achromatic/spectrometer_achromatic.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py b/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py index 48b8ab5a..e0e50bdc 100644 --- a/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py +++ b/abel/classes/spectrometer/impl/achromatic/spectrometer_achromatic.py @@ -64,7 +64,7 @@ class SpectrometerAchromatic(Spectrometer): Imaging energy. Default is None. disable_lens_nonlinearity : bool - Flag to disable lens nonlinearity. Default is `True`. + Flag to disable lens nonlinearity. Default is ``True``. """ @abstractmethod @@ -112,7 +112,7 @@ def __init__(self, Imaging energy. Default is None. disable_lens_nonlinearity : bool - Flag to disable lens nonlinearity. Default is `True`. + Flag to disable lens nonlinearity. Default is ``True``. """ self.length_drift_to_dipole = length_drift_to_dipole # [m] self.field_dipole = field_dipole # [T]