Add space-charge to Cheetah #142

greglenerd · 2024-04-03T22:41:44Z

This implements space-charge as outlined in #137

Description

Draft: todo

Motivation and Context

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

… of all the IGF solver functions.

…ome bugs.

cheetah/accelerator.py

Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>

cheetah/accelerator.py

…ed. The code is not yet tested. (last test E+vB field)

…test (expansion of a cold uniform beam).

greglenerd · 2024-05-03T18:49:16Z

cheetah/accelerator.py

@@ -299,6 +307,356 @@ 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.


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

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?

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.

greglenerd · 2024-05-03T18:57:09Z

cheetah/accelerator.py

+class SpaceChargeKick(Element):
+    """
+    Simulates space charge effects on a beam.
+    :param length: Length of the element in meters.


As for now, time discretization is manually implemented through the lenght. If spacechargekick objects are built automatically, it might be interesting to play on this parameter.

greglenerd · 2024-05-03T18:59:25Z

cheetah/accelerator.py

+    def is_skippable(self) -> bool:
+        return False
+
+    def plot(self, ax: matplotlib.axes.Axes, s: float) -> None:


I need to understand better the plot functions of cheetah to implement something here. The only thing I know is there will be an issue given the fact that spacechargekick elements add 'artificial' lenght to the beam. @cr-xu @jank324

@greglenerd So I think we need to be careful with the length property. In all existing elements, length represents the physical length of the beam pipe and Cheetah actually relies on this assumption in some places. I would suggest that SpaceChargeKick.length = 0.0 all the time (because it doesn't add length to the beam pipe), and that the property of SpaceChargeKick currently named length is renamed to something like effect_length. Does that make sense?

Based on this, I think SpaceChargeKick could be plotted with something like plt.axvspan as a background highlight from s - effect_length to s.

RemiLehe

Thanks a lot for this PR!

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

cheetah/accelerator.py

RemiLehe · 2024-05-05T14:54:17Z

cheetah/accelerator.py

+        nx: Union[torch.Tensor, nn.Parameter,int]=32,
+        ny: Union[torch.Tensor, nn.Parameter,int]=32,
+        ns: Union[torch.Tensor, nn.Parameter,int]=32,
+        dx: Union[torch.Tensor, nn.Parameter]=3,
+        dy: Union[torch.Tensor, nn.Parameter]=3,
+        ds: Union[torch.Tensor, nn.Parameter]=3,


@cr-xu @jank324
Could you comment on whether nx, ny, dx, dy etc. need to be tensors (instead regular ints and floats)?

Good question. I think their role is very similar to that of Screen.resolution. So I think no one will every want get the gradient with respect to any of them (?), meaning they don't have to be nn.Parameter. I think we could get away with making them int and float. But, in NumPy, if you interleave operations on NumPy types and Python types, computations slow down significantly. I don't know to what extend this is also an issue in PyTorch, but just to avoid it, I would make them all torch.Tensor.

cheetah/accelerator.py

cheetah/particles.py

tests/test_space_charge_kick.py

jank324 · 2024-05-06T08:23:16Z

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.

Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>

jank324 · 2024-06-04T14:35:15Z

tests/test_space_charge_kick.py

+    outgoing_beam = segment_space_charge.track(incoming)
+
+    # Final beam properties
+    sig_xo = outgoing_beam.sigma_x


This still needs to be done

jank324 · 2024-06-04T14:35:42Z

tests/test_space_charge_kick.py

+        sigma_yp=torch.tensor([2e-7]),
+    )
+    # Initial beam properties
+    incoming_particles0 = incoming_beam.particles


This also still need doing

jank324 · 2024-06-04T14:36:46Z

tests/test_space_charge_kick.py

+    # Final beam properties
+    incoming_particles1 = incoming_beam.particles
+
+    torch.set_printoptions(precision=16)


Why was this done? I think this line should probably be removed. I think it's not needed and tests shouldn't change global configs when they run.

jank324 · 2024-06-04T14:38:03Z

tests/test_space_charge_kick.py

+        3 + 2 * torch.sqrt(torch.tensor(2))
+    )
+    Nb = total_charge / elementary_charge
+    L = beta * gamma * kappa * torch.sqrt(R0**3 / (Nb * electron_radius))


I would avoid capital and single letter variable names ... I guess this should be section_length?

jank324 · 2024-06-04T14:40:46Z

cheetah/particles/particle_beam.py

@@ -269,6 +269,140 @@ def from_twiss(
            dtype=dtype,
        )

+    @classmethod
+    def uniform_3d_ellispoid(


This is a duplicate of the method below.

I think @greglenerd pasted this in here while Chenran and I were still working on it, but it was since merged into master and this PR. So this copy of the method is no longer needed.

Note that the two copies are not identical. The lower one should be kept.