Add 2D space charge element by austin-hoover · Pull Request #576 · desy-ml/cheetah

austin-hoover · 2025-09-28T18:46:14Z

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.

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)
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).
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.
I have reformatted the code and checked that formatting passes (required).
I have have fixed all issues found by flake8 (required).
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).
I have checked that the documentation builds (required).

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

austin-hoover · 2025-09-28T18:48:38Z

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.

roussel-ryan · 2025-10-02T19:56:11Z

+            ),
+        )
+
+    def _deposit_charge_on_grid(


could you write general functions for depositing charge on 1/2/3 D grids and put it in a separate? This is something that we want to do for wakefields, images as well

Yes I could. @RemiLehe @cr-xu did you start this already? I saw the Issue #543.

Not yet. It would be great if you can add this feature.

Hey Austin, Ryan mentioned this to me in our last meeting. I have code for 1d and 2d grids from our wakefields work which I can send over. I think it may be a good idea to start with 3 separate charge deposition functions, as the 1d case is especially simplified.

Great! Want to start a new PR with your work?

Are we wanting to utilize a grid deposition function in this PR, or later on?

I think it would make sense to first finalize/merge this PR, and then later refactor the code to expose a general grid deposition function that could be used in other parts of the code (incl. in the code that generates beam images).

RemiLehe · 2025-10-07T12:53:15Z

@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 · 2025-10-07T12:59:41Z

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 · 2025-10-07T13:21:30Z

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 · 2025-10-07T20:07:05Z

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 · 2025-10-07T21:39:55Z

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 · 2025-10-08T20:26:22Z

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

cr-xu · 2026-03-05T18:32:30Z

@jank324 I was looking for the charge deposition feature and stumbled on this, are there more fixes to be done for this PR? It's been sitting for quite a while now. Otherwise I would say we should merge it first and work on refactoring the charge deposition function later, as @RemiLehe suggested.
@roussel-ryan also has a custom charge deposition branch in his fork (master...roussel-ryan:cheetah:screen-charge-dep). We can probably use that for the general refactor?

jank324 · 2026-03-06T08:48:19Z



-class SpaceChargeKick(Element):
+class SpaceChargeKick3D(Element):


This is a breaking change. I'm undecided if we should just accept the breaking change or handle it in some way? (e.g. stub class with informative error or something like that). There are pros and cons for both. What are your opinions?

jank324 · 2026-03-06T10:37:27Z

+from cheetah import Segment, SpaceChargeKick2D
+
+
+def get_lorentz_factors(rest_energy: float, kin_energy: float) -> tuple[float, float]:


There seems to be a huge amount of boilerplate code for the tests here. I would prefer if this mimicked the 3D space charge kick tests, i.e. if it was basically the same but with minor modifications to work for SpaceChargeKick2D.

At the moment these tests also still fail.

So I really think it would be nicer to match the tests for the 3D element.

Copilot

Pull request overview

This PR introduces a new 2D space-charge kick element (SpaceChargeKick2D) for long-bunch / line-charge approximations, and renames the existing space-charge kick implementation to SpaceChargeKick3D to distinguish the solvers.

Changes:

Added SpaceChargeKick2D element (2D integrated Green function solver) and accompanying tests.
Renamed SpaceChargeKick → SpaceChargeKick3D and updated exports/usages across the library and tests.
Updated changelog to document the new feature and the breaking rename.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
`cheetah/accelerator/space_charge_kick_2d.py`	New 2D space charge kick element implementation.
`cheetah/accelerator/space_charge_kick_3d.py`	Renames the class to `SpaceChargeKick3D` and minor numeric/style adjustments.
`cheetah/accelerator/__init__.py`	Exports `SpaceChargeKick2D` and `SpaceChargeKick3D` from the accelerator package.
`cheetah/__init__.py`	Re-exports `SpaceChargeKick2D` and `SpaceChargeKick3D` at top-level `cheetah.*`.
`tests/test_space_charge_kick_2d.py`	Adds a dedicated 2D space-charge test suite, including an envelope-comparison test.
`tests/test_space_charge_kick_3d.py`	Updates tests to use `SpaceChargeKick3D`.
`tests/test_vectorized.py`	Updates element coverage list to include both 2D and 3D space-charge elements.
`tests/test_elements.py`	Updates xfail set to include both new space-charge element classes.
`tests/conftest.py`	Updates default element parameterization for both `SpaceChargeKick2D` and `SpaceChargeKick3D`.
`CHANGELOG.md`	Documents the breaking rename and new 2D feature.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

Copilot · 2026-03-06T10:58:43Z

+        beam_length = (
+            (beam.particles[..., 4]).max() - (beam.particles[..., 4]).min()
+        ).abs()
+        charge_density = (
+            self._array_rho(beam, xp_coordinates, cell_size, grid_dimensions)
+            / beam_length
+        )


In _solve_poisson_equation, beam_length is computed via a global .max()/.min() over all batches and particles. For vectorized beams this collapses the batch dimension, so every beam instance shares the same beam_length, which will incorrectly scale the charge density (and can vary depending on other batch entries). Compute the longitudinal extent per beam (reduce only over the particle dimension) and broadcast accordingly before dividing.

Suggested change

beam_length = (

(beam.particles[..., 4]).max() - (beam.particles[..., 4]).min()

).abs()

charge_density = (

self._array_rho(beam, xp_coordinates, cell_size, grid_dimensions)

/ beam_length

)

# Compute longitudinal extent per beam (reduce only over particle dimension)

z_coords = beam.particles[..., 4]

z_max = z_coords.max(dim=-1).values

z_min = z_coords.min(dim=-1).values

beam_length = (z_max - z_min).abs()

charge_density = self._array_rho(

beam, xp_coordinates, cell_size, grid_dimensions

)

# Broadcast beam_length to match charge_density shape for per-beam scaling

while beam_length.dim() < charge_density.dim():

beam_length = beam_length.unsqueeze(-1)

charge_density = charge_density / beam_length

I don't quite understand why @copilot wants to modify so many more things. I just fixed the vectorisation issue as such ... I think.

Copilot · 2026-03-06T10:58:44Z

+    def sample(self, n: int) -> np.ndarray:
+        x = np.random.normal(size=(n, 4))
+        x /= np.linalg.norm(x, axis=1)[:, None]
+        x /= np.std(x, axis=0)
+
+        S = self.cov()
+        L = np.linalg.cholesky(S)
+        x = np.matmul(x, L.T)
+        return x


This test uses NumPy RNG (np.random.normal) but the test suite fixture only seeds Torch RNGs. That makes the test nondeterministic and can cause flaky failures. Please seed NumPy (e.g., via a fixture) or generate the sample using Torch so it is covered by the existing deterministic seeding.

Copilot

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

rho has shape (n, bx, by) and charge_density has shape (n,). So charge_density needs to be reshaped to (n, 1, 1) before dividing.

austin-hoover · 2026-03-06T17:00:58Z

The failure is now coming from /tests/test_elements.py.

E       FileNotFoundError: [Errno 2] No such file or directory: 'tests/resources/consistency_expected_outgoing/SpaceChargeKick3D_ParticleBeam_default.pkl'

austin-hoover · 2026-03-06T17:15:12Z

Regarding the benchmarks: the overhead is to integrate the KV envelope equations. (You could probably condense it if you wanted to.) The tests in test_spacecharge_kick_3d.py use the analytic solution for a "cold" uniform-density ellipsoid in a drift. The solution for the 2D case involves the error functions above, so I think you'd need to implement that function and then decide what to do for the gradient tests. I don't have time to look at this over the next couple months so I leave it to you all.

austin-hoover added 2 commits September 28, 2025 14:01

Add SpaceChargeKick2D class and rename SpaceChargeKick3D

9680e87

Add KV envelope benchmarks

5012875

austin-hoover mentioned this pull request Sep 29, 2025

Add 2D space charge element #577

Open

roussel-ryan reviewed Oct 2, 2025

View reviewed changes

austin-hoover added 2 commits October 7, 2025 11:06

Change SpaceChargeKick to SpaceChargeKick3D/2D in tests

2f46258

Add SpaceChargeKick2D to conftest.py

b76a7db

austin-hoover marked this pull request as draft October 7, 2025 15:27

austin-hoover added 13 commits October 7, 2025 15:03

Fix tests

0a1847d

Remove unused files

4530a14

Move sc2d test to tests/ directory

705a21d

Format

e8fcbd6

Add test_vectorized

8fc1cd9

Add test_incoming_beam_not_modified

f20d5d6

Add test_does_not_break_segment_length

c942804

Copy more tests from sc3d

0a6a777

Format

a0afa32

Edit docstrings

2ff793a

Format

ff48c96

Format

cfbf272

Format

212b372

austin-hoover marked this pull request as ready for review October 7, 2025 20:07

jank324 added the enhancement New feature or request label Oct 10, 2025

jank324 reviewed Mar 6, 2026

View reviewed changes

Comment thread tests/conftest.py Outdated

Apply suggestion from @jank324

5017a49

jank324 reviewed Mar 6, 2026

View reviewed changes

Comment thread tests/test_elements.py Outdated

Apply suggestion from @jank324

a364a17

jank324 reviewed Mar 6, 2026

View reviewed changes

Comment thread tests/test_vectorized.py Outdated

jank324 added 4 commits March 6, 2026 10:12

Apply suggestion from @jank324

870657a

Rename original space charge kick test file

8ead4d5

Match 2D space charge style to 3D implementation and other Cheetah code

b0ea112

Minor code matching between space charge kick test files

de676fd

jank324 requested changes Mar 6, 2026

View reviewed changes

jank324 added 3 commits March 6, 2026 11:39

Fix formatting error

7dcfd3a

Add changelog entries

fd249c0

Merge branch 'master' into sc2d

b349e4c

jank324 requested review from Copilot March 6, 2026 10:52

Copilot started reviewing on behalf of jank324 March 6, 2026 10:53 View session

Copilot AI reviewed Mar 6, 2026

View reviewed changes

jank324 assigned austin-hoover Mar 6, 2026

Copilot AI reviewed Mar 6, 2026

View reviewed changes

jank324 and others added 5 commits March 6, 2026 12:12

Apply suggestion from @Copilot

3ff5b12

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Fix vectorisation issue pointed out by @Copilot

d638c42

Fix error message wording

6c03c54

Merge branch 'master' into sc2d

86b9c50

Reshape beam_length array

1113565

rho has shape (n, bx, by) and charge_density has shape (n,). So charge_density needs to be reshaped to (n, 1, 1) before dividing.

jank324 added 3 commits March 10, 2026 17:35

Merge branch 'master' into sc2d

4c0bb6e

Fix how beam_length is reshaped

1172bcc

Replace 2D space charge kicks with those from 3D

321dbf7



		class SpaceChargeKick(Element):
		class SpaceChargeKick3D(Element):

		from cheetah import Segment, SpaceChargeKick2D


		def get_lorentz_factors(rest_energy: float, kin_energy: float) -> tuple[float, float]:

-        beam_length = (
-            (beam.particles[..., 4]).max() - (beam.particles[..., 4]).min()
-        ).abs()
-        charge_density = (
-            self._array_rho(beam, xp_coordinates, cell_size, grid_dimensions)
-            / beam_length
-        )
+        # Compute longitudinal extent per beam (reduce only over particle dimension)
+        z_coords = beam.particles[..., 4]
+        z_max = z_coords.max(dim=-1).values
+        z_min = z_coords.min(dim=-1).values
+        beam_length = (z_max - z_min).abs()
+        charge_density = self._array_rho(
+            beam, xp_coordinates, cell_size, grid_dimensions
+        )
+        # Broadcast beam_length to match charge_density shape for per-beam scaling
+        while beam_length.dim() < charge_density.dim():
+            beam_length = beam_length.unsqueeze(-1)
+        charge_density = charge_density / beam_length

Conversation

austin-hoover commented Sep 28, 2025 • edited by jank324 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Types of changes

Checklist

Uh oh!

austin-hoover commented Sep 28, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

RemiLehe commented Oct 7, 2025

Uh oh!

RemiLehe commented Oct 7, 2025

Uh oh!

austin-hoover commented Oct 7, 2025

Uh oh!

austin-hoover commented Oct 7, 2025

Uh oh!

austin-hoover commented Oct 7, 2025

Uh oh!

aswaroop99 commented Oct 8, 2025

Uh oh!

cr-xu commented Mar 5, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

austin-hoover commented Mar 6, 2026

Uh oh!

austin-hoover commented Mar 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

austin-hoover commented Sep 28, 2025 •

edited by jank324

Loading