Add space-charge to Cheetah

[greglenerd]

This implements space-charge as outlined in https://github.com/desy-ml/cheetah/issues/137

Description

Draft: todo

Motivation and Context

See https://github.com/desy-ml/cheetah/issues/137

• [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).
• [ ] 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).
• [x] 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

[jank324]

Weirdly all GitHub Actions except for the docs build are currently not showing up on this PR. We need to keep an eye on this. It could just be a result of the merge conflicts with master, in which case we can deal with it once we get to that merge.

[cr-xu]

There appears to be also random failing of CI due to numerical issues? This probably should go away when running tests with more macroparticles.

[RemiLehe]

@cr-xu Yes, the tests where failing randomly, due to stochastic fluctuations in the initial particle distribution. In order to avoid this, I fixed the random seed in the test.

[jank324]

Okay ... the gradient compute in-place error is fixed, but in Replace obvious in-place operations with out-of-place alternatives I appear to have introduced a bug that just copies some lines of the particles tensor rather than returning the one updated with the space charge effects.

In the commit before that the tests pass just fine. 🤔

[jank324]

I've fixed the vectorisation now as well. However, there are two remaining issues, which I will highlight in comments below.

[jank324]

I think these two are the last remaining things before @RemiLehe can investigate the gradient-computations. Before merging, there are a few refactoring items we should probably take care of, but I would like to discuss them first.

[jank324]

Notes from our meeting about remaining things to do:

• [x] Implemented flatten hack to fix vectorisation and its gradients
• [x] Make betaref and gammaref properties of ParticleBeam
• [x] Make moments a property of ParticleBeam` and rename it (@RemiLehe you wanted to check what Ocelot calls it?)

[jank324]

So most issues are fixed. The only one I'm not really getting on top of is that the computation is still slightly wrong.

[cr-xu]

So most issues are fixed. The only one I'm not really getting on top of is that the computation is still slightly wrong.

I managed to fix that, but it seems that I can't push to the upstream branch

error: failed to push some refs to 'github.com:greglenerd/cheetah.git'

[jank324]

So most issues are fixed. The only one I'm not really getting on top of is that the computation is still slightly wrong.

I managed to fix that, but it seems that I can't push to the upstream branch

error: failed to push some refs to 'github.com:greglenerd/cheetah.git'

Huh ... odd ... I thought because you are Maintainer, you should have permissions for that. As a quick fix, maybe you can push to a temporary branch and then I merge it into the upstream one manually?

[cr-xu]

So most issues are fixed. The only one I'm not really getting on top of is that the computation is still slightly wrong.

I managed to fix that, but it seems that I can't push to the upstream branch

error: failed to push some refs to 'github.com:greglenerd/cheetah.git'

Huh ... odd ... I thought because you are Maintainer, you should have permissions for that. As a quick fix, maybe you can push to a temporary branch and then I merge it into the upstream one manually?

Okay I pushed the fix to fix-compute-forces-in-space-charge

[jank324]

So most issues are fixed. The only one I'm not really getting on top of is that the computation is still slightly wrong.

I managed to fix that, but it seems that I can't push to the upstream branch

error: failed to push some refs to 'github.com:greglenerd/cheetah.git'

Huh ... odd ... I thought because you are Maintainer, you should have permissions for that. As a quick fix, maybe you can push to a temporary branch and then I merge it into the upstream one manually?

Okay I pushed the fix to fix-compute-forces-in-space-charge

Is merged!

[jank324]

I've been implementing the plot function, but I'm not actually sure, where the space charge effect is applied relative to the element's position. I thought it's effect_length backwards from the s position where the element is located. So the following plotting code

import matplotlib.pyplot as plt
import torch

import cheetah


quad_length = torch.tensor([0.3])
segment = cheetah.Segment(
    elements=[
        cheetah.Drift(torch.tensor([0.5])),
        cheetah.Quadrupole(quad_length / 6),
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 3),
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 3),
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 6),
        cheetah.Drift(torch.tensor([0.5])),
    ]
)

ax = plt.subplot(111)
segment.plot(ax, s=0.0)

with the current implementation of SpaceChargeKick.plot

https://github.com/desy-ml/cheetah/blob/fea6c839100ff3cde7759c5e15dcbee8dc5dd929/cheetah/accelerator/space_charge_kick.py#L616-L624

generates this image:

Does that make sense?

[jank324]

I've gone through now and resolved all comments that are already dealt with or out-of-date. Here is what remains to do before we can merge:

• [x] Make sure the plotting is correct (https://github.com/desy-ml/cheetah/pull/142#issuecomment-2178491914)
• [x] Consider if there is a better name for moments (https://github.com/desy-ml/cheetah/pull/142#discussion_r1645149168)
• [x] Answer the question about the use of torch.long (https://github.com/desy-ml/cheetah/pull/142#discussion_r1644755047)
• [x] Implement defining_features (https://github.com/desy-ml/cheetah/pull/142#discussion_r1626159492)
• [x] Implement __repr__ (https://github.com/desy-ml/cheetah/pull/142#discussion_r1626160494)

If you think of anything else, please edit this post and add it here.

[RemiLehe]

@cr-xu @jank324 This is awesome! Thanks a lot for fixing the space-charge test, along with differentiation and vectorization! You guys are so fast!

Regarding a replacement for moments, here are a few suggestions:

• cheetah_coords_to_xp
• cheetah_coords_to_x_y_z_px_py_pz
• canonical_coords_to_xyz_pxpypz Any preferences? (For reference, here is the corresponding Ocelot function, but may want to use a different name.)

Regarding plotting: could we simply plot a vertical line (as if the space-charge kick was a thin element), instead of a rectangle?

Regarding torch.long, I think we could remove it and use regular ints instead. I will push a corresponding commit. We do need to cast to some integer type (e.g. int or long) after using floor, since floor returns a floating-point type, which is not appropriate for indexing arrays later on.

[jank324]

Regarding a replacement for moments, here are a few suggestions:

• cheetah_coords_to_xp
• cheetah_coords_to_x_y_z_px_py_pz
• canonical_coords_to_xyz_pxpypz Any preferences? (For reference, here is the corresponding Ocelot function, but may want to use a different name.)

I have no opinion regarding the xp, x_y_z_px_py_pz or xyz_pxpypz question. But I would have the function just to_xyz_pxpypz (or any of the three options) because in the end it would be used as

moments = cheetah_beam.to_xyz_pxpypz()
another_cheetah_beam = ParameterBeam.from_xyz_pxpypz(moments, ...)

Regarding plotting: could we simply plot a vertical line (as if the space-charge kick was a thin element), instead of a rectangle?

A vertical line might actually be more intuitive, yes. Still, can someone explain why we do

segment = cheetah.Segment(
    elements=[
        cheetah.Quadrupole(quad_length / 6),
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 3),
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 3),
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 6),
    ]
)

rather than

segment = cheetah.Segment(
    elements=[
        cheetah.Quadrupole(quad_length / 3),
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 3),
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 3),
        cheetah.SpaceChargeKick(quad_length / 3),
    ]
)

?

[cr-xu]

Regarding a replacement for moments, here are a few suggestions:

• cheetah_coords_to_xp
• cheetah_coords_to_x_y_z_px_py_pz
• canonical_coords_to_xyz_pxpypz Any preferences? (For reference, [here](https://github.com/ocelot-> collab/ocelot/blob/master/ocelot/cpbd/coord_transform.py#L57) is the corresponding Ocelot function, but may want to use a different name.)

I have no opinion regarding the xp, x_y_z_px_py_pz or xyz_pxpypz question. But I would have the function just to_xyz_pxpypz (or any of the three options) because in the end it would be used as

moments = cheetah_beam.to_xyz_pxpypz()
another_cheetah_beam = ParameterBeam.from_xyz_pxpypz(moments, ...)

to_xyz_pxpypz and from_xyz_pxpypz sounds good

Regarding plotting: could we simply plot a vertical line (as if the space-charge kick was a thin element), instead of a rectangle?

A vertical line might actually be more intuitive, yes. Still, can someone explain why we do

segment = cheetah.Segment(
   elements=[
       cheetah.Quadrupole(quad_length / 6),
       cheetah.SpaceChargeKick(quad_length / 3),
       cheetah.Quadrupole(quad_length / 3),
       cheetah.SpaceChargeKick(quad_length / 3),
       cheetah.Quadrupole(quad_length / 3),
       cheetah.SpaceChargeKick(quad_length / 3),
       cheetah.Quadrupole(quad_length / 6),
   ]
)

rather than

segment = cheetah.Segment(
   elements=[
       cheetah.Quadrupole(quad_length / 3),
       cheetah.SpaceChargeKick(quad_length / 3),
       cheetah.Quadrupole(quad_length / 3),
       cheetah.SpaceChargeKick(quad_length / 3),
       cheetah.Quadrupole(quad_length / 3),
       cheetah.SpaceChargeKick(quad_length / 3),
   ]
)

?

I believe this is just to make the SC apply in the middle of the split drifts. They way you suggested would effectly only apply SC kicks 2 times as the last one is at the end of the tracking. In the limit of large n_split they would be the same

[jank324]

I believe this is just to make the SC apply in the middle of the split drifts. They way you suggested would effectly only apply SC kicks 2 times as the last one is at the end of the tracking. In the limit of large n_split they would be the same

Ah ... so what about the following?

segment = cheetah.Segment(
    elements=[
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 3),
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 3),
        cheetah.SpaceChargeKick(quad_length / 3),
        cheetah.Quadrupole(quad_length / 3),
    ]
)

What I would just like to understand is that, if I place a SpaceChargeKick at position s with effect_length = l, it is as if space charge is considered on what interval? [s - l, s], [s, s + l] or [s - l / 2, s + l / 2]? Or does it not matter? Does the question even make sense?

[cr-xu]

What I would just like to understand is that, if I place a SpaceChargeKick at position s with effect_length = l, it is as if space charge is considered on what interval? [s - l, s], [s, s + l] or [s - l / 2, s + l / 2]? Or does it not matter? Does the question even make sense?

Hmm, I would say it doesn't make sense for small $l$, but makes a difference when the split is coarse. In the end $R_{sc}*R_{drift} \neq R_{sc}*R_{drift}$ because the transformations are noncommutative. But in the limit of $l \rightarrow 0$ then the differences become negligible. Let it happen in the middle ([s-2/l, s+2/l]) is probably just a better approximation (similar as for numerical integration)?

Example: a beam initially with $p_x=0$ would have $p_x=p$ after SC effect, so it would be more accurate to track half the drift with $x_p=0$ and half with $p_x=p$ ( [Drift(0.5l), SC_kick(l), Drift(0.5l)]) than the other two options:

• track the whole section with $x_p=0$ ( [Drift(l), SC_kick(l)]), and
• track the whole section with $x_p=p$ ( [SC_kick(l), Drift(l)])

[jank324]

What I would just like to understand is that, if I place a SpaceChargeKick at position s with effect_length = l, it is as if space charge is considered on what interval? [s - l, s], [s, s + l] or [s - l / 2, s + l / 2]? Or does it not matter? Does the question even make sense?

Hmm, I would say it doesn't make sense for small l, but makes a difference when the split is coarse. In the end Rsc∗Rdrift≠Rsc∗Rdrift because the transformations are noncommutative. But in the limit of l→0 then the differences become negligible. Let it happen in the middle ([s-2/l, s+2/l]) is probably just a better approximation (similar as for numerical integration)?

Example: a beam initially with px=0 would have px=p after SC effect, so it would be more accurate to track half the drift with xp=0 and half with px=p ( [Drift(0.5l), SC_kick(l), Drift(0.5l)]) than the other two options:

• track the whole section with xp=0 ( [Drift(l), SC_kick(l)]), and
• track the whole section with xp=p ( [SC_kick(l), Drift(l)])

I think I get it now. Thanks!

[jank324]

Regarding plotting: could we simply plot a vertical line (as if the space-charge kick was a thin element), instead of a rectangle?

Done! It looks like this now:

[jank324]

If I didn't miss anything, this PR should be done. 🎉 I've requested final reviews from everyone and then we can merge.

[RemiLehe]

Regarding a replacement for moments, here are a few suggestions:

• cheetah_coords_to_xp
• cheetah_coords_to_x_y_z_px_py_pz
• canonical_coords_to_xyz_pxpypz Any preferences? (For reference, [here](https://github.com/ocelot-> collab/ocelot/blob/master/ocelot/cpbd/coord_transform.py#L57) is the corresponding Ocelot function, but may want to use a different name.)

I have no opinion regarding the xp, x_y_z_px_py_pz or xyz_pxpypz question. But I would have the function just to_xyz_pxpypz (or any of the three options) because in the end it would be used as

moments = cheetah_beam.to_xyz_pxpypz()
another_cheetah_beam = ParameterBeam.from_xyz_pxpypz(moments, ...)

to_xyz_pxpypz and from_xyz_pxpypz sounds good

Regarding plotting: could we simply plot a vertical line (as if the space-charge kick was a thin element), instead of a rectangle?

A vertical line might actually be more intuitive, yes. Still, can someone explain why we do

segment = cheetah.Segment(
   elements=[
       cheetah.Quadrupole(quad_length / 6),
       cheetah.SpaceChargeKick(quad_length / 3),
       cheetah.Quadrupole(quad_length / 3),
       cheetah.SpaceChargeKick(quad_length / 3),
       cheetah.Quadrupole(quad_length / 3),
       cheetah.SpaceChargeKick(quad_length / 3),
       cheetah.Quadrupole(quad_length / 6),
   ]
)

rather than

segment = cheetah.Segment(
   elements=[
       cheetah.Quadrupole(quad_length / 3),
       cheetah.SpaceChargeKick(quad_length / 3),
       cheetah.Quadrupole(quad_length / 3),
       cheetah.SpaceChargeKick(quad_length / 3),
       cheetah.Quadrupole(quad_length / 3),
       cheetah.SpaceChargeKick(quad_length / 3),
   ]
)

?

I believe this is just to make the SC apply in the middle of the split drifts. They way you suggested would effectly only apply SC kicks 2 times as the last one is at the end of the tracking. In the limit of large n_split they would be the same

@cr-xu Yes, that is correct. The way in which we apply the space-charge kicks can be seen as an instance of Strang splitting, which under reasonable conditions, is second-order accurate which respect to the length of the splitting intervals.

[jank324]

So, it turns out scipy 1.14.0 requires Python>=3.10. That's why I wasn't getting the new version. So if GitHub Actions actually runs the 3.9 test process, it should actually pass.

I guess this one early sign that 3.9 is losing support. NumPy actually recommends only supporting 3.10+ starting 5 April 2024. I would keep 3.9 in for now though. Last I checked DOOCS was only compatible with 3.9.

[ax3l]

🎉