search

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

Automatic broadcasting #138

Open jank324 opened 3 months ago

jank324 commented 3 months ago

In #116, we added a vectorised way of using Cheetah. This means that instead of either one of these

elem1 = Drift(length=0.3142)
elem2 = Drift(length=torch.tensor(0.3142))

you have to remember to set properties in the correct dimension even if you are not using multiple dimensions or only a single value for that property

elemt3 = Drift(length=torch.tensor([0.3142])

By introducing automatic broadcasting (in a similar way to how PyTorch does it), we want to be able to use either one of the examples in the first code block while tracking, for example, a beam with 1000 different properties.