Add space-charge to Cheetah

cheetah/accelerator.py

@@ -30,7 +30,7 @@
 electron_mass_eV = torch.tensor(
  physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
 )
-
+epsilon_0 = torch.tensor(constants.epsilon_0)
 
 class Element(ABC, nn.Module):
  """
@@ -45,7 +45,7 @@ def __init__(self, name: Optional[str] = None) -> None:
  self.name = name if name is not None else generate_unique_name()
 
  def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
- r"""
+ """
  Generates the element's transfer map that describes how the beam and its
  particles are transformed when traveling through the element.
  The state vector consists of 6 values with a physical meaning:
@@ -299,6 +299,226 @@ def defining_features(self) -> list[str]:
  def __repr__(self) -> str:
  return f"{self.__class__.__name__}(length={repr(self.length)})"
 
+class SpaceChargeKick(Element):
+ """
+ Simulates space charge effects on a beam.
+ :param grid_points: Number of grid points in each dimension.
+ :param grid_dimensions: Dimensions of the grid in meters.
+ :param name: Unique identifier of the element.
+ """
+
+ def __init__(
+ self,
+ nx: Union[torch.Tensor, nn.Parameter,int],
+ ny: Union[torch.Tensor, nn.Parameter,int],
+ ns: Union[torch.Tensor, nn.Parameter,int],
+ dx: Union[torch.Tensor, nn.Parameter],
+ dy: Union[torch.Tensor, nn.Parameter],
+ ds: Union[torch.Tensor, nn.Parameter],
+ name: Optional[str] = None,
+ device=None,
+ dtype=torch.float32,
+ ) -> None:
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super().__init__(name=name)
+
+ self.nx = int(torch.as_tensor(nx, **factory_kwargs))
+ self.ny = int(torch.as_tensor(ny, **factory_kwargs))
+ self.ns = int(torch.as_tensor(ns, **factory_kwargs))
+ self.dx = torch.as_tensor(dx, **factory_kwargs) #in meters
+ self.dy = torch.as_tensor(dy, **factory_kwargs)
+ self.ds = torch.as_tensor(ds, **factory_kwargs)
+
+
+ def grid_shape(self) -> tuple[int,int,int]:
+ return (int(self.nx), int(self.ny), int(self.ns)) 
+
+
+ def grid_dimensions(self) -> torch.Tensor:
+ return torch.tensor([self.dx, self.dy, self.ds], device=self.dx.device)
+
+
+ def cell_size(self) -> torch.Tensor:
+ grid_shape = self.grid_shape()
+ grid_dimensions = self.grid_dimensions()
+ return 2*grid_dimensions / torch.tensor(grid_shape)
+
+
+ def space_charge_deposition(self, beam: ParticleBeam) -> torch.Tensor:
+ """
+ Deposition of the beam on the grid using fully vectorized computation.
+ """
+ grid_shape = self.grid_shape()
+ grid_dimensions = self.grid_dimensions()
+ cell_size = self.cell_size()
+
+ # Initialize the charge density grid
+ charge = torch.zeros(grid_shape, dtype=torch.float32)
+
+ # Get particle positions and charges
+ particle_pos = beam.particles[:, [0, 2, 4]]
+ particle_charge = beam.particle_charges
+
+ # Compute the normalized positions of the particles within the grid
+ normalized_pos = (particle_pos + grid_dimensions) / cell_size
+
+ # Find the indices of the lower corners of the cells containing the particles
+ cell_indices = torch.floor(normalized_pos).type(torch.long)
+
+ # Calculate the weights for all surrounding cells
+ offsets = torch.tensor([[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]])
+ surrounding_indices = cell_indices.unsqueeze(1) + offsets # Shape: (n_particles, 8, 3)
+ weights = 1 - torch.abs(normalized_pos.unsqueeze(1) - surrounding_indices) # Shape: (n_particles, 8, 3)
+ cell_weights = weights.prod(dim=2) # Shape: (n_particles, 8)
+
+ # Add the charge contributions to the cells
+ idx_x, idx_y, idx_s = surrounding_indices.view(-1, 3).T
+ valid_mask = (idx_x >= 0) & (idx_x < grid_shape[0]) & \
+ (idx_y >= 0) & (idx_y < grid_shape[1]) & \
+ (idx_s >= 0) & (idx_s < grid_shape[2])
+
+ # Accumulate the charge contributions
+ indices = torch.stack([idx_x[valid_mask], idx_y[valid_mask], idx_s[valid_mask]], dim=0)
+ repeated_charges = particle_charge.repeat_interleave(8)
+ values = (cell_weights.view(-1) * repeated_charges)[valid_mask]
+ charge.index_put_(tuple(indices), values, accumulate=True)
+ cell_volume = cell_size[0] * cell_size[1] * cell_size[2]
+
+ return charge/cell_volume # Normalize by the cell volume, so that the charge density is in C/m^3
+
+
+ def integrated_potential(self, x, y, s) -> torch.Tensor:
+ r = torch.sqrt(x**2 + y**2 + s**2)
+ G = (-0.5 * s**2 * torch.atan(x * y / (s * r))
+ -0.5 * y**2 * torch.atan(x * s / (y * r))
+ -0.5 * x**2 * torch.atan(y * s / (x * r))
+ + y * s * torch.asinh(x / torch.sqrt(y**2 + s**2))
+ + x * s * torch.asinh(y / torch.sqrt(x**2 + s**2))
+ + x * y * torch.asinh(s / torch.sqrt(x**2 + y**2)))
+ return G
+
+
+ def cyclic_rho(self,beam: ParticleBeam) -> torch.Tensor:
+ """
+ Compute the charge density on the grid using the cyclic deposition method.
+ """
+ grid_shape = self.grid_shape()
+ charge_density = self.space_charge_deposition(beam)
+
+ # Double the dimensions
+ new_dims = tuple(dim * 2 for dim in grid_shape)
+
+ # Create a new tensor with the doubled dimensions, filled with zeros
+ cyclic_charge_density = torch.zeros(new_dims)
+
+ # Copy the original charge_density values to the beginning of the new tensor
+ cyclic_charge_density[:charge_density.shape[0], :charge_density.shape[1], :charge_density.shape[2]] = charge_density
+ return cyclic_charge_density 
+
+ def IGF(self, beam: ParticleBeam) -> torch.Tensor:
+ gamma = beam.energy / rest_energy
+ dx, dy, ds = self.cell_size()[0], self.cell_size()[1], self.cell_size()[2] * gamma # ds is scaled by gamma
+ nx, ny, ns = self.grid_shape()
+
+ # Create coordinate grids
+ x = torch.arange(nx) * dx
+ y = torch.arange(ny) * dy
+ s = torch.arange(ns) * ds
+ x_grid, y_grid, s_grid = torch.meshgrid(x, y, s, indexing='ij')
+
+ # Compute the Green's function values
+ G_values = (
+ self.integrated_potential(x_grid + 0.5 * dx, y_grid + 0.5 * dy, s_grid + 0.5 * ds)
+ - self.integrated_potential(x_grid - 0.5 * dx, y_grid + 0.5 * dy, s_grid + 0.5 * ds)
+ - self.integrated_potential(x_grid + 0.5 * dx, y_grid - 0.5 * dy, s_grid + 0.5 * ds)
+ - self.integrated_potential(x_grid + 0.5 * dx, y_grid + 0.5 * dy, s_grid - 0.5 * ds)
+ + self.integrated_potential(x_grid + 0.5 * dx, y_grid - 0.5 * dy, s_grid - 0.5 * ds)
+ + self.integrated_potential(x_grid - 0.5 * dx, y_grid + 0.5 * dy, s_grid - 0.5 * ds)
+ + self.integrated_potential(x_grid - 0.5 * dx, y_grid - 0.5 * dy, s_grid + 0.5 * ds)
+ - self.integrated_potential(x_grid - 0.5 * dx, y_grid - 0.5 * dy, s_grid - 0.5 * ds)
+ )
+
+ # Initialize the grid with double dimensions
+ grid = torch.zeros(2 * nx, 2 * ny, 2 * ns)
+
+ # Fill the grid with G_values and its periodic copies
+ grid[:nx, :ny, :ns] = G_values
+ grid[nx+1:, :ny, :ns] = G_values[1:,:,:].flip(dims=[0]) # Reverse the x dimension, excluding the first element
+ grid[:nx, ny+1:, :ns] = G_values[:, 1:,:].flip(dims=[1]) # Reverse the y dimension, excluding the first element
+ grid[:nx, :ny, ns+1:] = G_values[:, :, 1:].flip(dims=[2]) # Reverse the s dimension, excluding the first element
+ grid[nx+1:, ny+1:, :ns] = G_values[1:, 1:,:].flip(dims=[0, 1]) # Reverse the x and y dimensions
+ grid[:nx, ny+1:, ns+1:] = G_values[:, 1:, 1:].flip(dims=[1, 2]) # Reverse the y and s dimensions
+ grid[nx+1:, :ny, ns+1:] = G_values[1:, :, 1:].flip(dims=[0, 2]) # Reverse the x and s dimensions
+ grid[nx+1:, ny+1:, ns+1:] = G_values[1:, 1:, 1:].flip(dims=[0, 1, 2]) # Reverse all dimensions
+
+ return grid
+
+
+ def solve_poisson_equation(self, beam: ParticleBeam) -> torch.Tensor: #works only for ParticleBeam at this stage
+ """
+ Solves the Poisson equation for the given charge density.
+ """
+ # Compute the charge density
+ charge_density = self.cyclic_rho(beam)
+
+ # Compute the Fourier transform of the charge density
+ charge_density_ft = torch.fft.fftn(charge_density)
+
+ # Compute the integrated Green's function
+ integrated_green_function = self.IGF(beam)
+
+ # Compute the integrated Green's function's Fourier transform
+ integrated_green_function_ft = torch.fft.fftn(integrated_green_function)
+
+ # Compute the Fourier transform of the potential
+ potential_ft = charge_density_ft * integrated_green_function_ft
+
+ # Compute the potential
+ potential = (1/4*torch.pi*epsilon_0)*torch.fft.ifftn(potential_ft).real
+
+ # Return the physical potential
+ return potential[:charge_density.shape[0]//2, :charge_density.shape[1]//2, :charge_density.shape[2]//2]
+
+
+ def split(self, resolution: torch.Tensor) -> list[Element]:
+ # TODO: Implement splitting for cavity properly, for now just returns the
+ # element itself
+ return [self]
+
+
+ def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
+ device = self.length.device
+ dtype = self.length.dtype
+
+ gamma = energy / rest_energy.to(device=device, dtype=dtype)
+ igamma2 = (
+ 1 / gamma**2
+ if gamma != 0
+ else torch.tensor(0.0, device=device, dtype=dtype)
+ )
+ beta = torch.sqrt(1 - igamma2)
+
+ tm = torch.eye(7, device=device, dtype=dtype)
+ tm[0, 1] = self.length
+ tm[2, 3] = self.length
+ tm[4, 5] = -self.length / beta**2 * igamma2
+
+ return tm
+
+ @property
+ def is_skippable(self) -> bool:
+ return True
+
+ def plot(self, ax: matplotlib.axes.Axes, s: float) -> None:
+ pass
+
+ @property
+ def defining_features(self) -> list[str]:
+ return super().defining_features + ["length"]
+
+ def __repr__(self) -> str:
+ return f"{self.__class__.__name__}(length={repr(self.length)})"
+
 
 class Quadrupole(Element):
  """

tests/test_space_charge_kick.py

@@ -0,0 +1,43 @@
+import pytest
+import torch
+
+import cheetah
+
+def test_charge_deposition():
+ """
+ Test that the charge deposition is correct for a particle beam. The first test checks that the total charge is preserved, and the second test checks that the charge is deposited in the correct grid cells.
+ """
+ space_charge_kick = cheetah.SpaceChargeKick(nx=32,ny=32,ns=32,dx=3e-9,dy=3e-9,ds=2e-6)
+ incoming_beam = cheetah.ParticleBeam.from_parameters(
+ num_particles=torch.tensor(1000),
+ sigma_xp=torch.tensor(2e-7),
+ sigma_yp=torch.tensor(2e-7),
+ )
+ total_charge = incoming_beam.total_charge
+ space_charge_grid = space_charge_kick.space_charge_deposition(incoming_beam) 
+
+ assert torch.isclose(space_charge_grid.sum() * space_charge_kick.grid_resolution ** 3, torch.tensor(total_charge), atol=1e-12) # grid_resolution is a parameter of the space charge kick #Total charge is preserved
+
+ # something similar to the read function in the CIC code should be implemented
+ assert outgoing_beam.sigma_y > incoming_beam.sigma_y
+
+
+@pytest.mark.skip(
+ reason="Requires rewriting Element and Beam member variables to be buffers."
+)
+def test_device_like_torch_module():
+ """
+ Test that when changing the device, Drift reacts like a `torch.nn.Module`.
+ """
+ # There is no point in running this test, if there aren't two different devices to
+ # move between
+ if not torch.cuda.is_available():
+ return
+
+ element = cheetah.Drift(length=torch.tensor(0.2), device="cuda")
+
+ assert element.length.device.type == "cuda"
+
+ element = element.cpu()
+
+ assert element.length.device.type == "cpu"