desy-ml / cheetah

Fast and differentiable particle accelerator optics simulation for reinforcement learning and optimisation applications.
https://cheetah-accelerator.readthedocs.io
GNU General Public License v3.0
33 stars 13 forks source link

Implement opt-in merging of elements #88

Closed jank324 closed 1 year ago

jank324 commented 1 year ago

I had an issue with the LCLS lattice, where the speed really wasn't great. This was likely due to that lattice having a huge number of elements. However, in my RL environment, I only ever tough a few of them. Most elements remain constant all the time.

I've implemented a new element type that just has a custom transfer map. This can be useful for all kinds of things actually. For my problem in particular, it is possible to take consecutive elements that won't be changed and combine their transfer maps into such an element.

In addition with some other minor optimisations (removing inactive Markers and so on) this gave me around a 20x speedup. I'm currently thinking about whether to also include those optimisations in this PR.

All these optimisations are (currently) opt-in, i.e. you create a new optimised segment by calling a method of the old segment.

jank324 commented 1 year ago

Also there is an issue right now that over the LCLS lattice I am not getting exactly the correct Twiss parameters. It's small enough to possibly be a numerical error, but large enough that it might also be a proper error. I am investigating.

jank324 commented 1 year ago

Remaining todos for this PR: