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
27 stars 12 forks source link

Add distgen interface #149

Open jank324 opened 2 months ago

jank324 commented 2 months ago

It came up that adding a Cheetah-native interface directly into distgen would make it very easy to generate all kinds of beam distributions.

This issue is to discuss the design of what that interface should look like.

Until this is implemented, it should be possible to do distgen -> OpenPMD -> Astra -> Cheetah.

jank324 commented 2 months ago

So I would suggest to hit to birds with one stone here and add both an OpenPMD interface and distgen to ParticleBeam like in the following pseudocode.

@classmethod
def from_openpmd(cls, some_arg) -> ParticleBeam:
    # Some code to convert from OpenPMD to Cheetah
    ...

@classmethod
def from_distgen(cls, distgen_input: str | Path | dict) -> ParticleBeam:
    gen = Generator(generator_input, verbose=0)
    openpmd = gen.run()
    return cls.from_openpmd(openpmd)

I'm not familiar with distgen and this is only based on skimming over the examples. So maybe @cr-xu and @jp-ga do you have an opinion on this?

jp-ga commented 2 months ago

t to hit to birds with one stone here and add both an OpenPMD interface and distgen to

Sounds good to me. I did something similar with Bmad-X. See https://github.com/bmad-sim/Bmad-X/blob/main/bmadx/distgen_utils.py and https://github.com/bmad-sim/Bmad-X/blob/main/bmadx/pmd_utils.py