add BmadQuadrupole element

[jp-ga]

Description

Add BmadQuadrupole element by interfacing with Bmad-X tracking routines.

Motivation and Context

Closes #139

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

[jank324]

As a general comment. I'm not sure I like this quick-and-dirty drop-in way of integrating Bmad-X elements. However, I think it makes sense to test them this way. In the long-term we should think about how this can be clean up better.

Some questions we should ask ourselves:

• Do we really want a separate bmadx module? This means that code is suddenly in a very different place than the other tracking code.
• Should there be a single Quadrupole class that has something like a tracking_method: Literal["linear", "bmadx"] argument?

Maybe slightly unrelated but I saw it in this PR:

• I think it's not very nice to have something like C = scipy.constants.speed_of_light at the top of a file. Ideally, one would just from scripy.constants import speed_of_light and use that. I think this might cause issues sometimes, when a torch.Tensor is strictly needed, but I cannot produce an example right now. Maybe we should add a cheetah.constants module that takes scripy.constnts and always returns torch.Tensors. But again, I'm not sure we actually need that ... just something to think about.

[jp-ga]

One small thing to keep in mind, bmadx tracking agrees to single precision with bmadx. I was expecting it to agree to double precision, but it is good enough for now.

[cr-xu]

Would it make sense to just integrate these two methods also in ParticleBeam like the other transformations?
I think p0c can also be provided as a property.

[jp-ga]

I was thinking to include the coordinate transforms in here instead: https://github.com/desy-ml/cheetah/blob/master/cheetah/converters/bmad.py

This way it could also be used to transform beams for both Bmad and Bmad-X.

Any thoughts?

[jank324]

@jp-ga, the question came up with @roussel-ryan, if this or #163 should be merged first. I think a main question in that was if the changes in #163 might make conversions in this PR redundant. Do you have an opinion on this?

[jank324]

@jp-ga Can you have a brief look at the split function I wrote a comment about?

@roussel-ryan Can you check if the test I added seems to be about the error you ran into?

Otherwise, I think this is ready to merge!

[jp-ga]

I'm not sure if we should add now how Quadrupole.split handles the num_steps argument.

@jank324 I thought about this back when developing Bmad-X, and at the end I decided that it was better to keep the number of steps in the quad different and independent from the number of elements when splitting. For example, if you have a quad with num_steps=3 and you split the element in 5, the code should do 5 quadrupoles with L/5 length, each of which with num_steps=3 (doing 3 drift-kick-drift steps).

In the future I would like to implement fringe fields, so the splitting should have entrance fringe at the first element's entrance and exit fringe at the last element's exit, but we can leave this for a future PR.

[jank324]

I'm not sure if we should add now how Quadrupole.split handles the num_steps argument.

@jank324 I thought about this back when developing Bmad-X, and at the end I decided that it was better to keep the number of steps in the quad different and independent from the number of elements when splitting. For example, if you have a quad with num_steps=3 and you split the element in 5, the code should do 5 quadrupoles with L/5 length, each of which with num_steps=3 (doing 3 drift-kick-drift steps).

In the future I would like to implement fringe fields, so the splitting should have entrance fringe at the first element's entrance and exit fringe at the last element's exit, but we can leave this for a future PR.

@jp-ga So the behaviour you would like to have now is that each of the 5 child quadrupoles has 3 splits, just like their parent had 3 splits? In that case, don't we need to add the line highlighted below in the split method?

element = Quadrupole(
    torch.min(resolution, remaining),
    self.k1,
    misalignment=self.misalignment,
    tilt=self.tilt,
    num_steps=self.num_steps,   # <-- Add this line?
    dtype=self.length.dtype,
    device=self.length.device,
)

[roussel-ryan]

@jank324 I'm not sure if the test you added will raise the issue I ran into, I think my error had to do with the ellipsoid beam generation, but I'm not sure if that test will run into the same issue. I think it may be likely through since I remember that the different methods used the same code

[jp-ga]

I'm not sure if we should add now how Quadrupole.split handles the num_steps argument.

@jank324 I thought about this back when developing Bmad-X, and at the end I decided that it was better to keep the number of steps in the quad different and independent from the number of elements when splitting. For example, if you have a quad with num_steps=3 and you split the element in 5, the code should do 5 quadrupoles with L/5 length, each of which with num_steps=3 (doing 3 drift-kick-drift steps).
In the future I would like to implement fringe fields, so the splitting should have entrance fringe at the first element's entrance and exit fringe at the last element's exit, but we can leave this for a future PR.

@jp-ga So the behaviour you would like to have now is that each of the 5 child quadrupoles has 3 splits, just like their parent had 3 splits? In that case, don't we need to add the line highlighted below in the split method?

element = Quadrupole(
    torch.min(resolution, remaining),
    self.k1,
    misalignment=self.misalignment,
    tilt=self.tilt,
    num_steps=self.num_steps,   # <-- Add this line?
    dtype=self.length.dtype,
    device=self.length.device,
)

@jank324 Exactly, I will add it soon.

[cr-xu]

One final comment: Is it intended to keep the dev/ folder? If the notebook is useful, shall we move it to some other places, like a documentation page?
Otherwise it's good to go!

[jank324]

One final comment: Is it intended to keep the dev/ folder? If the notebook is useful, shall we move it to some other places, like a documentation page? Otherwise it's good to go!

I'm also not sure about the dev/ folder. I left it because I also sometimes write notebooks to develop features and I often worry that I will have to write them twice because I always delete them at some point once I'm done and I never commit them. dev/ would be a clean way to keep them around.

On the other hand, I've never actually felt like I needed to rewrite a notebook for some reason and we might not want to clutter up the repository and history like this.

[jank324]

So ... to document the decision: The dev/ folder remains for now, but the long-term goal should be to not have it. Interesting examples should instead, for example, appear in the documentation.