diff --git a/CHANGELOG.md b/CHANGELOG.md index 32697a7d4..45d3982be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ ### 🚨 Breaking Changes +- The `SpaceChargeKick` element has been renamed to `SpaceChargeKick3D` in order to reflect the difference to the new `SpaceChargeKick2D` element (see #576) (@austin-hoover, @RemiLehe, @jank324) + ### 🚀 Features - Improve the speed of `SpaceChargeKick` by up to 2x by replacing its custom Cloud-in-Cell implementation with the new general implementation (see #653) (@jank324) @@ -45,6 +47,7 @@ ### 🚀 Features - All plotting functions in `Segment` now accept an optional axes or figure object with an interface loosely mimicking that of [_Seaborn_](https://seaborn.pydata.org). If they are passed one, they use it, otherwise they create one themselves. Either way, they return the axes or figure they used. (see #604) (@jank324) +- Add a `SpaceChargeKick2D` element that applies a space charge kick using the 2D integrated Green function method (see #576) (@austin-hoover, @RemiLehe, @jank324) ### 🐛 Bug fixes diff --git a/cheetah/__init__.py b/cheetah/__init__.py index 8b48f8628..d5741c368 100644 --- a/cheetah/__init__.py +++ b/cheetah/__init__.py @@ -18,6 +18,7 @@ Sextupole, Solenoid, SpaceChargeKick, + SpaceChargeKick2D, Superimposed, TransverseDeflectingCavity, Undulator, diff --git a/cheetah/accelerator/__init__.py b/cheetah/accelerator/__init__.py index 62becab10..c4c45e4e9 100644 --- a/cheetah/accelerator/__init__.py +++ b/cheetah/accelerator/__init__.py @@ -15,6 +15,7 @@ from .sextupole import Sextupole # noqa: F401 from .solenoid import Solenoid # noqa: F401 from .space_charge_kick import SpaceChargeKick # noqa: F401 +from .space_charge_kick_2d import SpaceChargeKick2D # noqa: F401 from .superimposed import Superimposed # noqa: F401 from .transverse_deflecting_cavity import TransverseDeflectingCavity # noqa: F401 from .undulator import Undulator # noqa: F401 diff --git a/cheetah/accelerator/space_charge_kick.py b/cheetah/accelerator/space_charge_kick.py index 45abdc5d5..0e8422fc8 100644 --- a/cheetah/accelerator/space_charge_kick.py +++ b/cheetah/accelerator/space_charge_kick.py @@ -306,7 +306,7 @@ def _solve_poisson_equation( integrated_green_function, dim=[1, 2, 3] ) potential_ft = charge_density_ft * integrated_green_function_ft - potential = (1.0 / (4 * torch.pi * epsilon_0)) * torch.fft.irfftn( + potential = (1.0 / (4.0 * torch.pi * epsilon_0)) * torch.fft.irfftn( potential_ft, dim=[1, 2, 3] ).real diff --git a/cheetah/accelerator/space_charge_kick_2d.py b/cheetah/accelerator/space_charge_kick_2d.py new file mode 100644 index 000000000..fabfa9c9b --- /dev/null +++ b/cheetah/accelerator/space_charge_kick_2d.py @@ -0,0 +1,486 @@ +import matplotlib.pyplot as plt +import torch +from scipy.constants import elementary_charge, epsilon_0, speed_of_light + +from cheetah.accelerator.element import Element +from cheetah.particles import ParticleBeam +from cheetah.utils.cloud_in_cell import cloud_in_cell_charge_deposition + + +class SpaceChargeKick2D(Element): + """ + Applies the effect of space charge over a length `effect_length`, on the + **momentum** (i.e. divergence and energy spread) of the beam. The positions are + unmodified; this is meant to be combined with another lattice element (e.g. `Drift`) + that does modify the positions, but does not take into account space charge. The 2D + integrated Green function method (https://doi.org/10.1016/j.jcp.2004.01.008) is used + to compute the effect of space charge. + + Overview of the method: + - Compute the beam charge density on a grid. + - Convolve the charge density with a Green function (the integrated green function) + to find the potential `phi` on the grid. The convolution uses the Hockney method + for open boundaries (allocate 2x larger arrays and perform convolution using + FFTs). + - Compute the corresponding electromagnetic fields and Lorentz force on the grid. + - Interpolate the Lorentz force to the particles and update their momentum. + + This is a true 2D solver; we assume a uniform density in the longitudinal plane + (line charges). + + :param effect_length: Length over which the effect is applied in meters. + :param grid_shape: Number of grid points in (x, y) directions. + :param grid_extent_x: Dimensions of the grid on which to compute space-charge, as + multiples of sigma of the beam in the x direction (dimensionless). + :param grid_extent_y: Dimensions of the grid on which to compute space-charge, as + multiples of sigma of the beam in the y direction (dimensionless). + :param name: Unique identifier of the element. + :param sanitize_name: Whether to sanitise the name to be a valid Python variable + name. This is needed if you want to use the `segment.element_name` syntax to + access the element in a segment. + :param metadata: Dictionary of arbitrary, serialisable annotations attached to the + element (e.g. control-system addresses or PVs). This information is *not* used + in simulation and may contain any extra data the user wants to store along with + the lattice. See :doc:`/examples/including_metadata` for more information. + """ + + def __init__( + self, + effect_length: torch.Tensor, + grid_shape: tuple[int, int] = (32, 32), + # TODO: Simplify these to a single tensor? + grid_extent_x: torch.Tensor | None = None, + grid_extent_y: torch.Tensor | None = None, + name: str | None = None, + sanitize_name: bool = False, + metadata: dict | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> None: + factory_kwargs = {"device": device, "dtype": dtype} + + super().__init__( + name=name, sanitize_name=sanitize_name, metadata=metadata, **factory_kwargs + ) + + self.grid_shape = grid_shape + + self.register_buffer_or_parameter("effect_length", effect_length) + # In multiples of sigma + self.register_buffer_or_parameter( + "grid_extent_x", + ( + grid_extent_x + if grid_extent_x is not None + else torch.tensor(3.0, **factory_kwargs) + ), + ) + self.register_buffer_or_parameter( + "grid_extent_y", + ( + grid_extent_y + if grid_extent_y is not None + else torch.tensor(3.0, **factory_kwargs) + ), + ) + + def _integrated_potential(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """Computes the Green function at point (x, y).""" + integrated_potential = -0.5 * ( + -3.0 * x * y + + x.square() * (y / x).atan() + + y.square() * (x / y).atan() + + x * y * (x.square() + y.square()).log() + ) + return integrated_potential + + def _array_rho( + self, + beam: ParticleBeam, + xp_coordinates: torch.Tensor, + cell_size: torch.Tensor, + grid_dimensions: torch.Tensor, + ) -> torch.Tensor: + """ + Allocates a 2x larger array in all dimensions (to perform Hockney's method), and + copies the charge density in one of the "quadrants". + """ + charge_grid = cloud_in_cell_charge_deposition( + positions=xp_coordinates[..., [0, 2]], + bins=self.grid_shape, + extent=torch.stack([-grid_dimensions, grid_dimensions], dim=-1), + charges=beam.particle_charges * beam.survival_probabilities, + ) + + # Normalise by the cell volume to get density + inv_cell_volume = cell_size.prod(dim=-1).reciprocal() + charge_density = charge_grid * inv_cell_volume[..., None, None] + + # Create a new tensor with the doubled dimensions, filled with zeros + new_dims = tuple(2 * dim for dim in self.grid_shape) + new_charge_density = beam.particles.new_zeros( + beam.particles.shape[:-2] + new_dims + ) + + # Copy the original charge_density values to the beginning of the new tensor + new_charge_density[ + ..., + : charge_density.shape[-2], + : charge_density.shape[-1], + ] = charge_density + + return new_charge_density + + def _integrated_green_function( + self, beam: ParticleBeam, cell_size: torch.Tensor + ) -> torch.Tensor: + """ + Computes the Integrated Green Function (IGF) in the 2x larger array, + as needed for the Hockney method. + """ + dx, dy = cell_size[..., 0], cell_size[..., 1] + num_grid_points_x, num_grid_points_y = self.grid_shape + + # Create coordinate grids + x = torch.arange(num_grid_points_x, device=beam.particles.device) + y = torch.arange(num_grid_points_y, device=beam.particles.device) + ix_grid, iy_grid = torch.meshgrid(x, y, indexing="ij") + x_grid = ix_grid[None, :, :] * dx[..., None, None] # Shape: [..., nx, ny] + y_grid = iy_grid[None, :, :] * dy[..., None, None] # Shape: [..., nx, ny] + + # Compute the Green's function values + G_values = ( + self._integrated_potential( + x_grid + 0.5 * dx[..., None, None], + y_grid + 0.5 * dy[..., None, None], + ) + - self._integrated_potential( + x_grid - 0.5 * dx[..., None, None], + y_grid + 0.5 * dy[..., None, None], + ) + - self._integrated_potential( + x_grid + 0.5 * dx[..., None, None], + y_grid - 0.5 * dy[..., None, None], + ) + + self._integrated_potential( + x_grid - 0.5 * dx[..., None, None], + y_grid - 0.5 * dy[..., None, None], + ) + ) + + # Initialize the grid with double dimensions + green_func_values = beam.particles.new_zeros( + (*beam.particles.shape[:-2], 2 * num_grid_points_x, 2 * num_grid_points_y) + ) + + # Fill the grid with G_values and its periodic copies + green_func_values[..., :num_grid_points_x, :num_grid_points_y] = G_values + + # Reverse x, excluding the first element + green_func_values[..., num_grid_points_x + 1 :, :num_grid_points_y] = G_values[ + ..., 1:, : + ].flip(dims=[-2]) + + # Reverse y, excluding the first element + green_func_values[..., :num_grid_points_x, num_grid_points_y + 1 :] = G_values[ + ..., :, 1: + ].flip(dims=[-1]) + + # Reverse x and y + green_func_values[..., num_grid_points_x + 1 :, num_grid_points_y + 1 :] = ( + G_values[..., 1:, 1:].flip(dims=[-2, -1]) + ) + + return green_func_values + + def _solve_poisson_equation( + self, + beam: ParticleBeam, + xp_coordinates: torch.Tensor, + cell_size, + grid_dimensions, + ) -> torch.Tensor: # Works only for ParticleBeam at this stage + """ + Solves the Poisson equation for the given charge density, using FFT convolution. + """ + # Line density + beam_length = ( + (beam.particles[..., 4]).max(dim=-1).values + - (beam.particles[..., 4]).min(dim=-1).values + ).abs() + charge_density = self._array_rho( + beam, xp_coordinates, cell_size, grid_dimensions + ) / beam_length.unsqueeze(-1).unsqueeze(-1) + charge_density_ft = torch.fft.rfftn(charge_density, dim=[1, 2]) + integrated_green_function = self._integrated_green_function(beam, cell_size) + integrated_green_function_ft = torch.fft.rfftn( + integrated_green_function, dim=[1, 2] + ) + potential_ft = charge_density_ft * integrated_green_function_ft + potential = (1.0 / (2.0 * torch.pi * epsilon_0)) * torch.fft.irfftn( + potential_ft, dim=[1, 2] + ).real + + # Return the physical potential + return potential[ + ..., : charge_density.shape[-2] // 2, : charge_density.shape[-1] // 2 + ] + + def _E_plus_vB_field( + self, + beam: ParticleBeam, + xp_coordinates: torch.Tensor, + cell_size: torch.Tensor, + grid_dimensions: torch.Tensor, + ) -> torch.Tensor: + """ + Computes the force field from the potential and the particle positions and + velocities, as in https://doi.org/10.1063/1.2837054. + """ + inv_cell_size = cell_size.reciprocal() + igamma2 = torch.zeros_like(beam.relativistic_gamma) + igamma2[beam.relativistic_gamma != 0] = ( + beam.relativistic_gamma[beam.relativistic_gamma != 0].square().reciprocal() + ) + potential = self._solve_poisson_equation( + beam, xp_coordinates, cell_size, grid_dimensions + ) + + grad_x = torch.zeros_like(potential) + grad_y = torch.zeros_like(potential) + + # Compute the gradients of the potential, using central differences, with 0 + # boundary conditions + grad_x[..., 1:-1, :] = (potential[..., 2:, :] - potential[..., :-2, :]) * ( + 0.5 * inv_cell_size[..., 0, None, None] + ) + grad_y[..., :, 1:-1] = (potential[..., :, 2:] - potential[..., :, :-2]) * ( + 0.5 * inv_cell_size[..., 1, None, None] + ) + + # Scale the gradients with lorentz factor + grad_x = -igamma2[..., None, None] * grad_x + grad_y = -igamma2[..., None, None] * grad_y + + return grad_x, grad_y + + def _compute_forces( + self, + beam: ParticleBeam, + xp_coordinates: torch.Tensor, + cell_size: torch.Tensor, + grid_dimensions: torch.Tensor, + ) -> torch.Tensor: + """ + Interpolates the space charge force from the grid onto the macroparticles. + Reciprocal function of _deposit_charge_on_grid. `beam` needs to have a flattened + vector shape. + """ + grad_x, grad_y = self._E_plus_vB_field( + beam, xp_coordinates, cell_size, grid_dimensions + ) + grid_shape = self.grid_shape + interpolated_forces = beam.particles.new_zeros( + (*beam.particles.shape[:-1], 2) + ) # (..., num_particles, 2) + + # Get particle positions + particle_positions = xp_coordinates[..., [0, 2]] + normalized_positions = ( + particle_positions + grid_dimensions.unsqueeze(-2) + ) / cell_size.unsqueeze(-2) + + # Find indices of the lower corners of the cells containing the particles + cell_indices = normalized_positions.floor().to(torch.int) + + # Calculate the weights for all surrounding cells + offsets = torch.tensor( + [[0, 0], [0, 1], [1, 0], [1, 1]], device=cell_indices.device + ) + surrounding_indices = cell_indices.unsqueeze(-2) + offsets.unsqueeze( + -3 + ) # Shape:(..., num_particles, 4, 2) + weights = ( + 1.0 - (normalized_positions.unsqueeze(-2) - surrounding_indices).abs() + ) # Shape: (..., num_particles, 4, 2) + cell_weights = weights.prod(dim=-1) # Shape: (..., num_particles, 4) + + # Extract forces from the grids + surrounding_indices_flattened = surrounding_indices.flatten( + start_dim=-3, end_dim=-2 + ) # Shape: (..., num_particles * 4, 2) + idx_vector = ( + torch.arange(cell_indices.shape[0], device=cell_indices.device) + .repeat(4 * beam.particles.shape[-2], 1) + .T + ) # Shape: (..., num_particles * 4) + idx_x = surrounding_indices_flattened[..., 0] + idx_y = surrounding_indices_flattened[..., 1] + valid_mask = ( + (idx_x >= 0) + & (idx_x < grid_shape[0]) + & (idx_y >= 0) + & (idx_y < grid_shape[1]) + ) + + # Keep dimensions, and set F to zero if non-valid + force_indices = ( + idx_vector, + idx_x.clamp(min=0, max=grid_shape[0] - 1), + idx_y.clamp(min=0, max=grid_shape[1] - 1), + ) + + Fx_values = grad_x[force_indices].where(valid_mask, 0) + Fy_values = grad_y[force_indices].where(valid_mask, 0) + + # Compute interpolated forces + # Cell weights validation is taken care of by the F_x, F_y, F_z values + cell_weights_with_e = cell_weights.flatten(start_dim=-2) * elementary_charge + values_x = cell_weights_with_e * Fx_values + values_y = cell_weights_with_e * Fy_values + + forces_to_add = torch.stack([values_x, values_y], dim=-1) + + index_tensor = ( + torch.arange(beam.num_particles, device=beam.particles.device) + .repeat_interleave(4) + .unsqueeze(0) + .unsqueeze(-1) + .expand(beam.particles.shape[0], 4 * beam.particles.shape[-2], 2) + ) + + # Add the forces to the particles + accumulated_forces = torch.scatter_add( + interpolated_forces, dim=1, index=index_tensor, src=forces_to_add + ) + + return accumulated_forces + + def track(self, incoming: ParticleBeam) -> ParticleBeam: + """ + Tracks particles through the element. The input must be a `ParticleBeam`. + + :param incoming: Beam of particles entering the element. + :returns: Beam of particles exiting the element. + """ + assert isinstance( + incoming, ParticleBeam + ), "SpaceChargeKick2D tracking is currently only supported for `ParticleBeam`." + + # This flattening is a hack to only think about one vector dimension in the + # following code. It is reversed at the end of the function. + + # Make sure that the incoming beam has at least one vector dimension by + # broadcasting with a dummy dimension (1,). + vector_shape = torch.broadcast_shapes( + incoming.particles.shape[:-2], + incoming.energy.shape, + incoming.particle_charges.shape[:-1], + incoming.survival_probabilities.shape[:-1], + (1,), + ) + vectorized_incoming = ParticleBeam( + particles=torch.broadcast_to( + incoming.particles, (*vector_shape, incoming.num_particles, 7) + ), + energy=torch.broadcast_to(incoming.energy, vector_shape), + particle_charges=torch.broadcast_to( + incoming.particle_charges, (*vector_shape, incoming.num_particles) + ), + survival_probabilities=torch.broadcast_to( + incoming.survival_probabilities, + (*vector_shape, incoming.num_particles), + ), + species=incoming.species, + device=incoming.particles.device, + dtype=incoming.particles.dtype, + ) + + flattened_incoming = ParticleBeam( + particles=vectorized_incoming.particles.flatten(end_dim=-3), + energy=vectorized_incoming.energy.flatten(end_dim=-1), + particle_charges=vectorized_incoming.particle_charges.flatten(end_dim=-2), + survival_probabilities=( + vectorized_incoming.survival_probabilities.flatten(end_dim=-2) + ), + species=incoming.species, + device=vectorized_incoming.particles.device, + dtype=vectorized_incoming.particles.dtype, + ) + flattened_length_effect = self.effect_length.flatten(end_dim=-1) + + # Compute useful quantities + grid_dimensions = torch.stack( + [ + self.grid_extent_x * flattened_incoming.sigma_x, + self.grid_extent_y * flattened_incoming.sigma_y, + ], + dim=-1, + ) + cell_size = ( + 2 + * grid_dimensions + / torch.tensor( + self.grid_shape, + device=grid_dimensions.device, + dtype=grid_dimensions.dtype, + ) + ) + dt = flattened_length_effect / ( + speed_of_light * flattened_incoming.relativistic_beta + ) + + # Change coordinates to apply the space charge effect + xp_coordinates = flattened_incoming.to_xyz_pxpypz() + forces = self._compute_forces( + flattened_incoming, xp_coordinates, cell_size, grid_dimensions + ) + xp_coordinates[..., 1] = xp_coordinates[..., 1] + forces[..., 0] * dt.unsqueeze( + -1 + ) + xp_coordinates[..., 3] = xp_coordinates[..., 3] + forces[..., 1] * dt.unsqueeze( + -1 + ) + + # Reverse the flattening of the vector dimensions + outgoing_vector_shape = torch.broadcast_shapes( + incoming.particles.shape[:-2], + incoming.energy.shape, + incoming.particle_charges.shape[:-1], + incoming.survival_probabilities.shape[:-1], + self.effect_length.shape, + ) + outgoing = ParticleBeam.from_xyz_pxpypz( + xp_coordinates=xp_coordinates.reshape( + (*outgoing_vector_shape, incoming.num_particles, 7) + ), + energy=incoming.energy, + particle_charges=incoming.particle_charges, + survival_probabilities=incoming.survival_probabilities, + s=incoming.s, + species=incoming.species, + ) + + return outgoing + + @property + def is_skippable(self) -> bool: + return False + + def plot( + self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None + ) -> plt.Axes: + ax = ax or plt.subplot(111) + + plot_s = s[vector_idx] if s.dim() > 0 else s + + ax.axvline(plot_s, ymin=0.01, ymax=0.99, color="orange", linestyle="-") + + @property + def defining_features(self) -> list[str]: + return super().defining_features + [ + "effect_length", + "grid_shape", + "grid_extent_x", + "grid_extent_y", + ] diff --git a/tests/conftest.py b/tests/conftest.py index 2c6c49dac..ffe24c14f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -127,6 +127,7 @@ } }, cheetah.SpaceChargeKick: {"default": {"effect_length": torch.tensor(1.0)}}, + cheetah.SpaceChargeKick2D: {"default": {"effect_length": torch.tensor(1.0)}}, cheetah.Superimposed: { "default": { "base_element": cheetah.Quadrupole( diff --git a/tests/resources/consistency_expected_outgoing/SpaceChargeKick2D_ParticleBeam_default.pkl b/tests/resources/consistency_expected_outgoing/SpaceChargeKick2D_ParticleBeam_default.pkl new file mode 100644 index 000000000..a3300b00f Binary files /dev/null and b/tests/resources/consistency_expected_outgoing/SpaceChargeKick2D_ParticleBeam_default.pkl differ diff --git a/tests/test_elements.py b/tests/test_elements.py index 6e38adfcf..93e6757c7 100644 --- a/tests/test_elements.py +++ b/tests/test_elements.py @@ -102,7 +102,12 @@ def test_particle_beam_tracking_with_device_and_dtype(element, device, dtype): @pytest.mark.for_every_element( "element", xfail_if=lambda element: isinstance( - element, (cheetah.SpaceChargeKick, cheetah.TransverseDeflectingCavity) + element, + ( + cheetah.SpaceChargeKick, + cheetah.SpaceChargeKick2D, + cheetah.TransverseDeflectingCavity, + ), ) or ( isinstance( diff --git a/tests/test_space_charge_kick.py b/tests/test_space_charge_kick.py index 170bdb795..ecba35fae 100644 --- a/tests/test_space_charge_kick.py +++ b/tests/test_space_charge_kick.py @@ -42,14 +42,18 @@ def test_cold_uniform_beam_expansion(energy): energy=energy, radius_x=R0, radius_y=R0, - radius_tau=R0 / gamma / beta, # Duration of the beam in in the lab frame + radius_tau=R0 / gamma / beta, # Duration of the beam in the lab frame sigma_px=torch.tensor(1e-15), sigma_py=torch.tensor(1e-15), sigma_p=torch.tensor(1e-15), ) # Compute section length that results in a doubling of the beam size - kappa = 1 + (torch.tensor(2).sqrt() / 4) * (3 + 2 * torch.tensor(2).sqrt()).log() + kappa = ( + 1.0 + + (torch.tensor(2.0).sqrt() / 4.0) + * (3.0 + 2.0 * torch.tensor(2.0).sqrt()).log() + ) Nb = incoming.total_charge / elementary_charge section_length = beta * gamma * kappa * (R0.pow(3) / (Nb * electron_radius)).sqrt() @@ -66,9 +70,9 @@ def test_cold_uniform_beam_expansion(energy): ) outgoing = segment.track(incoming) - assert torch.isclose(outgoing.sigma_x, 2 * incoming.sigma_x, rtol=2e-2) - assert torch.isclose(outgoing.sigma_y, 2 * incoming.sigma_y, rtol=2e-2) - assert torch.isclose(outgoing.sigma_tau, 2 * incoming.sigma_tau, rtol=2e-2) + assert torch.isclose(outgoing.sigma_x, 2.0 * incoming.sigma_x, rtol=2e-2) + assert torch.isclose(outgoing.sigma_y, 2.0 * incoming.sigma_y, rtol=2e-2) + assert torch.isclose(outgoing.sigma_tau, 2.0 * incoming.sigma_tau, rtol=2e-2) def test_vectorized_cold_uniform_beam_expansion(): @@ -95,14 +99,18 @@ def test_vectorized_cold_uniform_beam_expansion(): energy=energy, radius_x=R0, radius_y=R0, - radius_tau=R0 / gamma / beta, # Duration of the beam in in the lab frame + radius_tau=R0 / gamma / beta, # Duration of the beam in the lab frame sigma_px=torch.tensor(1e-15), sigma_py=torch.tensor(1e-15), sigma_p=torch.tensor(1e-15), ) # Compute section length - kappa = 1 + (torch.tensor(2).sqrt() / 4) * (3 + 2 * torch.tensor(2).sqrt()).log() + kappa = ( + 1.0 + + (torch.tensor(2.0).sqrt() / 4.0) + * (3.0 + 2.0 * torch.tensor(2.0).sqrt()).log() + ) Nb = incoming.total_charge / elementary_charge section_length = beta * gamma * kappa * (R0.pow(3) / (Nb * electron_radius)).sqrt() @@ -119,9 +127,9 @@ def test_vectorized_cold_uniform_beam_expansion(): ) outgoing = segment.track(incoming) - assert torch.allclose(outgoing.sigma_x, 2 * incoming.sigma_x, rtol=2e-2) - assert torch.allclose(outgoing.sigma_y, 2 * incoming.sigma_y, rtol=2e-2) - assert torch.allclose(outgoing.sigma_tau, 2 * incoming.sigma_tau, rtol=2e-2) + assert torch.allclose(outgoing.sigma_x, 2.0 * incoming.sigma_x, rtol=2e-2) + assert torch.allclose(outgoing.sigma_y, 2.0 * incoming.sigma_y, rtol=2e-2) + assert torch.allclose(outgoing.sigma_tau, 2.0 * incoming.sigma_tau, rtol=2e-2) def test_vectorized(): @@ -226,7 +234,9 @@ def test_gradient_value_backward_ad(): # Compute section length that results in a doubling of the beam size electron_radius = torch.tensor(physical_constants["classical electron radius"][0]) kappa = ( - 1 + (torch.tensor(2.0).sqrt() / 4) * (3 + 2 * torch.tensor(2.0).sqrt()).log() + 1.0 + + (torch.tensor(2.0).sqrt() / 4.0) + * (3.0 + 2.0 * torch.tensor(2.0).sqrt()).log() ) Nb = incoming_beam.total_charge / constants.elementary_charge segment_length = beta * gamma * kappa * (R0.pow(3) / (Nb * electron_radius)).sqrt() @@ -290,7 +300,9 @@ def test_gradient_value_forward_ad(): # Compute section length that results in a doubling of the beam size electron_radius = torch.tensor(physical_constants["classical electron radius"][0]) kappa = ( - 1 + (torch.tensor(2.0).sqrt() / 4) * (3 + 2 * torch.tensor(2.0).sqrt()).log() + 1.0 + + (torch.tensor(2.0).sqrt() / 4.0) + * (3.0 + 2.0 * torch.tensor(2.0).sqrt()).log() ) Nb = incoming_beam.total_charge / constants.elementary_charge segment_length = beta * gamma * kappa * (R0.pow(3) / (Nb * electron_radius)).sqrt() diff --git a/tests/test_space_charge_kick_2d.py b/tests/test_space_charge_kick_2d.py new file mode 100644 index 000000000..475bdab17 --- /dev/null +++ b/tests/test_space_charge_kick_2d.py @@ -0,0 +1,409 @@ +import math + +import pytest +import torch +import torch.autograd.forward_ad as fwAD +from scipy import constants +from scipy.constants import physical_constants +from torch import nn + +import cheetah +from cheetah.utils import compute_relativistic_factors + + +# Run the test below for both the ultra-relativistic case +# (250 MeV) and the non-relativistic case (1 MeV). +@pytest.mark.parametrize( + "energy", + [torch.tensor(2.5e8), torch.tensor(1e6)], + ids=["ultra-relativistic", "non-relativistic"], +) +def test_cold_uniform_beam_expansion(energy): + """ + Tests that that a cold uniform beam doubles in size in both dimensions when + travelling through a drift section with space_charge. (cf ImpactX test: + https://impactx.readthedocs.io/en/latest/usage/examples/expanding_beam/README.html) + See Free Expansion of a Cold Uniform Bunch in + https://accelconf.web.cern.ch/hb2023/papers/thbp44.pdf + """ + # Simulation parameters + R0 = torch.tensor(0.001) + rest_energy = torch.tensor( + constants.electron_mass + * constants.speed_of_light**2 + / constants.elementary_charge + ) + elementary_charge = torch.tensor(constants.elementary_charge) + electron_radius = torch.tensor(physical_constants["classical electron radius"][0]) + gamma = energy / rest_energy + beta = (1 - gamma.square().reciprocal()).sqrt() + + incoming = cheetah.ParticleBeam.uniform_3d_ellipsoid( + num_particles=100_000, + total_charge=torch.tensor(1e-8), + energy=energy, + radius_x=R0, + radius_y=R0, + radius_tau=R0 / gamma / beta, # Duration of the beam in the lab frame + sigma_px=torch.tensor(1e-15), + sigma_py=torch.tensor(1e-15), + sigma_p=torch.tensor(1e-15), + ) + + # Compute section length that results in a doubling of the beam size + kappa = 1.352 * beta.reciprocal().sqrt() + Nb = incoming.total_charge / elementary_charge + section_length = beta * gamma * kappa * (R0.pow(3) / (Nb * electron_radius)).sqrt() + + segment = cheetah.Segment( + elements=[ + cheetah.Drift(section_length / 6), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 3), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 3), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 6), + ] + ) + outgoing = segment.track(incoming) + + assert torch.isclose(outgoing.sigma_x, 2.0 * incoming.sigma_x, rtol=2e-2) + assert torch.isclose(outgoing.sigma_y, 2.0 * incoming.sigma_y, rtol=2e-2) + assert torch.isclose(outgoing.sigma_tau, incoming.sigma_tau, rtol=2e-2) + + +def test_vectorized_cold_uniform_beam_expansion(): + """ + Same as `test_cold_uniform_beam_expansion` but testing that all results in a + vectorised setup are correct. + """ + # Simulation parameters + R0 = torch.tensor(0.001) + energy = torch.tensor(2.5e8) + rest_energy = torch.tensor( + constants.electron_mass + * constants.speed_of_light**2 + / constants.elementary_charge + ) + elementary_charge = torch.tensor(constants.elementary_charge) + electron_radius = torch.tensor(physical_constants["classical electron radius"][0]) + gamma = energy / rest_energy + beta = (1 - gamma.square().reciprocal()).sqrt() + + incoming = cheetah.ParticleBeam.uniform_3d_ellipsoid( + num_particles=100_000, + total_charge=torch.tensor(1e-8).repeat(3, 2), + energy=energy, + radius_x=R0, + radius_y=R0, + radius_tau=R0 / gamma / beta, # Duration of the beam in the lab frame + sigma_px=torch.tensor(1e-15), + sigma_py=torch.tensor(1e-15), + sigma_p=torch.tensor(1e-15), + ) + + # Compute section length + kappa = 1.352 * beta.reciprocal().sqrt() + Nb = incoming.total_charge / elementary_charge + section_length = beta * gamma * kappa * (R0.pow(3) / (Nb * electron_radius)).sqrt() + + segment = cheetah.Segment( + elements=[ + cheetah.Drift(section_length / 6), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 3), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 3), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 6), + ] + ) + outgoing = segment.track(incoming) + + assert torch.allclose(outgoing.sigma_x, 2.0 * incoming.sigma_x, rtol=2e-2) + assert torch.allclose(outgoing.sigma_y, 2.0 * incoming.sigma_y, rtol=2e-2) + assert torch.allclose(outgoing.sigma_tau, incoming.sigma_tau, rtol=2e-2) + + +def test_vectorized(): + """Tests that the space charge kick can be applied to a vectorized beam.""" + # Simulation parameters + section_length = torch.tensor(0.42) + R0 = torch.tensor(0.001) + energy = torch.tensor(2.5e8) + rest_energy = torch.tensor( + constants.electron_mass + * constants.speed_of_light**2 + / constants.elementary_charge + ) + gamma = energy / rest_energy + beta = (1 - gamma.square().reciprocal()).sqrt() + + incoming = cheetah.ParticleBeam.uniform_3d_ellipsoid( + num_particles=10_000, + total_charge=torch.tensor([[1e-9, 2e-9], [3e-9, 4e-9], [5e-9, 6e-9]]), + energy=energy.expand([3, 2]), + radius_x=R0.expand([3, 2]), + radius_y=R0.expand([3, 2]), + radius_tau=R0.expand([3, 2]) / gamma / beta, + # Duration of the beam in the lab frame + sigma_px=torch.tensor(1e-15).expand([3, 2]), + sigma_py=torch.tensor(1e-15).expand([3, 2]), + sigma_p=torch.tensor(1e-15).expand([3, 2]), + ) + + segment = cheetah.Segment( + elements=[ + cheetah.Drift(section_length / 6), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 3), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 3), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 6), + ] + ) + + outgoing = segment.track(incoming) + + assert outgoing.particles.shape == (3, 2, 10_000, 7) + + +def test_incoming_beam_not_modified(): + """ + Tests that the incoming beam is not modified when calling the track method. + """ + incoming_beam = cheetah.ParticleBeam.from_parameters( + num_particles=10_000, sigma_px=torch.tensor(2e-7), sigma_py=torch.tensor(2e-7) + ) + # Initial beam properties + incoming_beam_before = incoming_beam.particles + + section_length = torch.tensor(1.0) + segment_space_charge = cheetah.Segment( + elements=[ + cheetah.Drift(section_length / 6), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 3), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 3), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 6), + ] + ) + # Calling the track method + segment_space_charge.track(incoming_beam) + + # Final beam properties + incoming_beam_after = incoming_beam.particles + + assert torch.allclose(incoming_beam_before, incoming_beam_after) + + +def test_gradient_value_backward_ad(): + """ + Tests that the gradient of the track method is computed accurately. Using PyTorch's + default backward mode automatic differentiation. + """ + # Simulation parameters + R0 = torch.tensor(0.001) + energy = torch.tensor(2.5e8) + species = cheetah.Species("electron") + gamma, _, beta = compute_relativistic_factors(energy, species.mass_eV) + + incoming_beam = cheetah.ParticleBeam.uniform_3d_ellipsoid( + num_particles=100_000, + total_charge=torch.tensor(1e-8), + energy=energy, + radius_x=R0, + radius_y=R0, + radius_tau=R0 / gamma / beta, # Duration of the beam in the lab frame + sigma_px=torch.tensor(1e-15), + sigma_py=torch.tensor(1e-15), + sigma_p=torch.tensor(1e-15), + species=species, + ) + + # Compute section length that results in a doubling of the beam size + electron_radius = torch.tensor(physical_constants["classical electron radius"][0]) + kappa = 1.352 * beta.reciprocal().sqrt() + Nb = incoming_beam.total_charge / constants.elementary_charge + segment_length = beta * gamma * kappa * (R0.pow(3) / (Nb * electron_radius)).sqrt() + + segment_length = nn.Parameter(segment_length) + segment = cheetah.Segment( + elements=[ + cheetah.Drift(segment_length / 6), + cheetah.SpaceChargeKick2D(segment_length / 3), + cheetah.Drift(segment_length / 3), + cheetah.SpaceChargeKick2D(segment_length / 3), + cheetah.Drift(segment_length / 3), + cheetah.SpaceChargeKick2D(segment_length / 3), + cheetah.Drift(segment_length / 6), + ] + ) + + # Track the beam + outgoing_beam = segment.track(incoming_beam) + + # Compute the gradient ... would throw an error if in-place operations are used + outgoing_beam.sigma_x.backward() + + # Check that the gradient is correct by comparing the derivative of the beam radius + # as a function of the segment length + dsigma_dlength = segment_length.grad + # For a sphere, the radius is sqrt(5) bigger than sigma_x + dradius_dlength = 5**0.5 * dsigma_dlength + # Theoretical formula obtained by conservation of energy in the beam frame, + # scaled by 1.25 due to the parabolic longitudinal profile. + expected_dradius_dlength = 1.25 * (2.0 * math.log(2.0) / beta).sqrt() * (Nb * electron_radius / R0).sqrt() / gamma + + assert torch.allclose(dradius_dlength, expected_dradius_dlength, rtol=0.1) + + +def test_gradient_value_forward_ad(): + """ + Tests that the gradient of the track method is computed accurately. Using PyTorch's + forward mode automatic differentiation. + + See: https://pytorch.org/tutorials/intermediate/forward_ad_usage.html + """ + # Simulation parameters + R0 = torch.tensor(0.001) + energy = torch.tensor(2.5e8) + species = cheetah.Species("electron") + gamma, _, beta = compute_relativistic_factors(energy, species.mass_eV) + + incoming_beam = cheetah.ParticleBeam.uniform_3d_ellipsoid( + num_particles=100_000, + total_charge=torch.tensor(1e-8), + energy=energy, + radius_x=R0, + radius_y=R0, + radius_tau=R0 / gamma / beta, # Duration of the beam in the lab frame + sigma_px=torch.tensor(1e-15), + sigma_py=torch.tensor(1e-15), + sigma_p=torch.tensor(1e-15), + species=species, + ) + + # Compute section length that results in a doubling of the beam size + electron_radius = torch.tensor(physical_constants["classical electron radius"][0]) + kappa = 1.352 * beta.reciprocal().sqrt() + Nb = incoming_beam.total_charge / constants.elementary_charge + segment_length = beta * gamma * kappa * (R0.pow(3) / (Nb * electron_radius)).sqrt() + + tangent = torch.ones_like(segment_length) + + with fwAD.dual_level(): + segment_length = fwAD.make_dual(segment_length, tangent) + + segment = cheetah.Segment( + elements=[ + cheetah.Drift(segment_length / 6), + cheetah.SpaceChargeKick2D(segment_length / 3), + cheetah.Drift(segment_length / 3), + cheetah.SpaceChargeKick2D(segment_length / 3), + cheetah.Drift(segment_length / 3), + cheetah.SpaceChargeKick2D(segment_length / 3), + cheetah.Drift(segment_length / 6), + ] + ) + + # Track the beam + outgoing_beam = segment.track(incoming_beam) + beam_size = outgoing_beam.sigma_x + + # Check that the gradient is correct by comparing the derivative of the beam + # radius as a function of the segment length + dsigma_dlength = fwAD.unpack_dual(beam_size).tangent + # For a sphere, the radius is sqrt(5) bigger than sigma_x + dradius_dlength = 5**0.5 * dsigma_dlength + # Theoretical formula obtained by conservation of energy in the beam frame, + # scaled by 1.25 due to the parabolic longitudinal profile. + expected_dradius_dlength = 1.25 * (2.0 * math.log(2.0) / beta).sqrt() * (Nb * electron_radius / R0).sqrt() / gamma + + assert torch.allclose(dradius_dlength, expected_dradius_dlength, rtol=0.1) + + +def test_does_not_break_segment_length(): + """ + Test that the computation of a `Segment`'s length does not break when + `SpaceChargeKick2D` is used. + """ + section_length = torch.tensor(1.0) + segment = cheetah.Segment( + elements=[ + cheetah.Drift(section_length / 6), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 3), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 3), + cheetah.SpaceChargeKick2D(section_length / 3), + cheetah.Drift(section_length / 6), + ] + ) + + assert segment.length.shape == torch.Size([]) + assert torch.allclose(segment.length, torch.tensor(1.0)) + + +def test_space_charge_with_ares_astra_beam(): + """ + Tests running space charge through a 1m drift with an Astra beam from the ARES + linac. This test is added because running this code would throw an error: + `IndexError: index -38 is out of bounds for dimension 3 with size 32`. + """ + segment = cheetah.Segment( + [ + cheetah.Drift(length=torch.tensor(1.0)), + cheetah.SpaceChargeKick2D(effect_length=torch.tensor(1.0)), + ] + ) + beam = cheetah.ParticleBeam.from_astra("tests/resources/ACHIP_EA1_2021.1351.001") + + _ = segment.track(beam) + + +def test_space_charge_with_aperture_cutoff(): + """ + Tests that the space charge kick is correctly applied only to surviving particles, + by comparing the results with and without an aperture that results in beam losses. + """ + segment = cheetah.Segment( + elements=[ + cheetah.Drift(length=torch.tensor(0.2)), + cheetah.Aperture( + x_max=torch.tensor(1e-4), + y_max=torch.tensor(1e-4), + shape="rectangular", + is_active=False, + name="aperture", + ), + cheetah.Drift(length=torch.tensor(0.25)), + cheetah.SpaceChargeKick2D(effect_length=torch.tensor(0.5)), + cheetah.Drift(length=torch.tensor(0.25)), + ] + ) + incoming_beam = cheetah.ParticleBeam.from_parameters( + num_particles=10_000, + total_charge=torch.tensor(1e-9), + mu_x=torch.tensor(5e-5), + sigma_px=torch.tensor(1e-4), + sigma_py=torch.tensor(1e-4), + ) + + # Track with inactive aperture + outgoing_beam_without_aperture = segment.track(incoming_beam) + + # Activate the aperture and track the beam + segment.aperture.is_active = True + outgoing_beam_with_aperture = segment.track(incoming_beam) + + # Check that with particle loss the space charge kick is different + assert not torch.allclose( + outgoing_beam_with_aperture.particles, outgoing_beam_without_aperture.particles + ) + # Check that the number of surviving particles is less than the initial number + assert outgoing_beam_with_aperture.survival_probabilities.sum(dim=-1).max() < 10_000 diff --git a/tests/test_vectorized.py b/tests/test_vectorized.py index be72138e0..aeaa74aeb 100644 --- a/tests/test_vectorized.py +++ b/tests/test_vectorized.py @@ -348,6 +348,7 @@ def test_vectorized_screen_2d(BeamClass, method): cheetah.Screen, cheetah.Segment, cheetah.SpaceChargeKick, + cheetah.SpaceChargeKick2D, cheetah.Superimposed, ), ),