Add 2D space charge element

[austin-hoover]

Description

This PR adds an element called SpaceChargeKick2D which applies 2D space charge kicks to the beam. It assumes the beam has an infinite length and uniform density in the longitudinal plane (line charges). The code is pretty much the same as the current SpaceChargeKick element. The integrated 2D Green function is taken from https://www.sciencedirect.com/science/article/abs/pii/S0021999104000282?via%3Dihub.

I have also renamed SpaceChargeKick to SpaceChargeKick3D.

Motivation and Context

• 2D space charge calculations are valid for long bunches.
• This is the simplest approach. For more realistic cases, one can weight the transverse kicks by the particle density in each longitudinal slice, or do a separate solution in each slice (https://github.com/PyORBIT-Collaboration/PyORBIT3/tree/main/src/spacecharge).
• This class also does not consider boundary conditions.
• Like the current SpaceChargeCalc element, the 2D grid is centered at zero and expanded based on the rms beam size, so it may not contain all particles.

• [x] I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

• [ ] Bug fix (non-breaking change which fixes an issue)
• [x] New feature (non-breaking change which adds functionality)
• [ ] Breaking change (fix or feature that would cause existing functionality to change)
• [ ] Documentation (update in the documentation)

Checklist

• [ ] I have updated the changelog accordingly (required).
• [x] My change requires a change to the documentation.
• [ ] I have updated the tests accordingly (required for a bug fix or a new feature).
• [ ] I have updated the documentation accordingly.
• [x] I have reformatted the code and checked that formatting passes (required).
• [x] I have have fixed all issues found by flake8 (required).
• [x] I have ensured that all pytest tests pass (required).
• [ ] I have run pytest on a machine with a CUDA GPU and made sure all tests pass (required).
• [x] I have checked that the documentation builds (required).

Note: We are using a maximum length of 88 characters per line.

[austin-hoover]

I've added a benchmark against the KV envelope equations in a drift and FODO lattice: https://github.com/austin-hoover/cheetah/tree/sc2d/tests/test_space_charge_kick_2d. Temporarily putting this in /tests directory.

[RemiLehe]

@austin-hoover This looks great, thanks for adding the new feature. Could check why the CI tests are failing and could you try to fix them? (See e.g. https://github.com/desy-ml/cheetah/actions/runs/18078378681/job/51438134719?pr=576) It seems that a number of failures are simply related to the renaming of SpaceChargeKick to SpaceChargeKick3D.

[RemiLehe]

In addition, I think it would be good to have tests for the SpaceChargeKick2D that run as part of the CI. This could be either the new tests that you added (if it is fast enough), or another test inspired e.g. by the existing CI tests for SpaceChargeKick3D.

A key aspect of the CI tests is that they need automated checks (typically with assert statements that result in a Python exception whenever the test criterion is not met: see e.g. https://github.com/desy-ml/cheetah/blob/master/tests/test_space_charge_kick.py#L126). My understanding is that the tests that you added currently produces plots, and require a human to look at them and check agreement, but do not yet have the automated assert statements. (Please correct me if I overlooked something here.)

[austin-hoover]

I'll check the CI failure. For the tests, ya these are only making plots right now. I'll add some proper tests mirroring SC3D. There is an expression for the radius of an expanding "cold" 2D KV distribution:

But integrating the KV envelope equations should be fast enough for the tests.

[austin-hoover]

Okay, I put tests under /tests/test_space_charge_kick_2d.py. The function test_kv_drift checks the that envelope equations and PIC solver agree on the rms beam sizes to within 0.1%. The rest of the tests are copied from /tests/test_space_charge_kick.py.

I didn't include the gradient tests. For a zero-emittance KV beam in free space, the envelope radius $r$ obeys $r'' = Q / r$ and gives the expression above. I'm not sure about the derivative of that. Maybe there's another case that could be used to check the gradient?

[austin-hoover]

We could compute derivative of single particle $x'$ with respect to $s$?

$$ x'' = \frac{dx'}{ds} = \frac{2Q}{(r_x + r_y)} \frac{x}{r_x}. $$

[aswaroop99]

Let me just share what I have via comment, as I'm a bit unclear if we're trying to generalize first or just trying to make cheetah functions for 1d and 2d. These functions copy the steps of the charge deposition function in space charge, but in 1d and 2d.

Here's what I have for 1d:

def deposit_counts_1d(z): Δ = (z_max - z_min) / Ngrid norm = (z - z_min) / Δ i0 = torch.floor(norm).long().clamp(0, Ngrid-1) w1 = (norm - i0.float()).clamp(0,1) w0 = 1 - w1 i1 = (i0+1).clamp(0, Ngrid-1) rho = torch.zeros(Ngrid, device=z.device) rho.scatter_add_(0, i0, w0) rho.scatter_add_(0, i1, w1) return rho

and this is what I have for 2d: def deposit_counts_2d(z,x,Nz,Nx,z0,z1,x0,x1): device=z.device Dz=(z1-z0)/Nz;Dx=(x1-x0)/Nx nz=(z-z0)/Dz;nx=(x-x0)/Dx i0=nz.floor().long();j0=nx.floor().long() wz1=(nz-i0.float()).clamp(0,1);wx1=(nx-j0.float()).clamp(0,1) wz0=1-wz1;wx0=1-wx1 i0=i0.clamp(0,Nz-1);j0=j0.clamp(0,Nx-1) i1=(i0+1).clamp(0,Nz-1);j1=(j0+1).clamp(0,Nx-1) rho2d=torch.zeros(Nz,Nx,device=device) rho2d.index_put_((i0,j0),wz0wx0,accumulate=True) rho2d.index_put_((i0,j1),wz0wx1,accumulate=True) rho2d.index_put_((i1,j0),wz1wx0,accumulate=True) rho2d.index_put_((i1,j1),wz1wx1,accumulate=True) return rho2d