Add space-charge to Cheetah

[greglenerd]

Description

Adds a new SpaceChargeKick element that enables adding space charge to any other elements.

Motivation and Context

Implements #137.

See #137

• 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

[greglenerd]

these params might be changed: I was thinking of a way to integrate spacechargekicks as a setting (eg spacecharge = True) which would automatically build spacechargekick objects and incorporate them appropriately, as they don't represent any physical element of the accelerator. @cr-xu @jank324

[jank324]

I think at some point in the future, we should think about creating a Effects class and make SpaceChargeKick a subclass of that. But right now this doesn't make sense yet.

Nevertheless, here is how I think we should integrate adding space charge to elements right now in a way that can easily be extended in the future:

Basically what I would do is to add a method to the Element class that looks something like this:

class Element:

    ...

    def with_space_charge(self, resolution: float = 0.01, *args, **kwargs) -> Segment:
        splits = self.split(resolution)
        splits_with_space_charge = # List of [split, sc, split, sc, ..., split, sc] ... probably itertools has a nice way to create this ... *args and **kwargs go into SpaceChargeKick
        return Segment(elements=splits_with_space_charge, name=f"{self.name}_with_space_charge")

This should end up looking similar to what @RemiLehe proposed in the very beginning. So if you do

Drift(length=0.5, name="my_drift").with_space_charge(resolution=0.25, nx=64)

you get

Segment(
    name="my_drift_with_space_charge",
    elements=[
        Drift(length=0.25), SpaceChargeKick(nx=64), Drift(length=0.25), SpaceChargeKick(nx=64)
    ],
)

This example is probably not quite correct in many ways, but I think it illustrates my idea.

The nice thing about doing it this way would be that it automatically works for all elements that implement split correctly, and that it would be relatively straightforward to extend this for other collective effects in the future.

Does this make sense?

[ax3l]

An alternative thought for the low level implementation: One could also create sub- or mixing classes for thin (L=0) and thick (L>0) elements. For elements in ImpactX, we currently use these mixin classes, e.g. Drift vs. thin Multipole.

Space charge kicks could be thin elements :)

For the high level interface, I agree on the above syntax. You want to automatically slice this up, a property or method on how many slices on the sliced element is a nice user interface.

[RemiLehe]

Thanks a lot for this PR!

I added a few suggestions ; we can discuss them further offline.

[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.

[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.

[RemiLehe]

I would suggest to replace the term moments by something like xp_coordinates or beam_coordinates, or something similar, throughout the code.

The problem with moments is that it suggests either the idea of moments of a distribution (mean, rms, etc.), or the idea of momentum (px, py, pz). None of these are correct descriptions of what this array contains.

[RemiLehe]

Done.

[jank324]

@RemiLehe this will probably also have to be done in particle_beam.py, including maybe a rephrasing of the docstrings (?).

[RemiLehe]

Sounds good. Done in b2e45c0

[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]

🎉