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
32 stars 13 forks source link

length cannot be a NN parameter #164

Open RTSandberg opened 4 months ago

RTSandberg commented 4 months ago

I have been unable to let length be an optimizable Pytorch.nn.Parameter without corrupting the transfer map. I think it is because every element inherits a length from the Element class. https://github.com/desy-ml/cheetah/blob/3fe16a797ccb141abd2c13be4342a693fe0b2585/cheetah/accelerator/element.py#L32

If I comment this out, then I am able to assign length as a Parameter, but this then causes some tests to fail (I believe it is every CI test using transfer maps).

I haven't found a good fix yet but it would be nice to allow length to be a parameter

cr-xu commented 4 months ago

Hi, indeed the length is causing quite some trouble. nn.Parameter doesn't like to re-register some attribute which is already defined in the class.

I implemented a quick and dirty way to circumvent this (see the new branch). Now there should be no problem defining the length as a trainable nn.Parameter, and the tests are passing.

The fix was basically moving the default length definition into the __init__. The issue with that is that Segment actually has a different length definition since it's adding all the lengths from its elements. I now added _initialize_length in the base Element class and redefine it in Segment so that it doesn't clash with the length being a property.

@jank324 I'm not sure if it's the best solution though, any opinion?

jank324 commented 4 months ago

Hmmm ... yeah ... we already struggled with this in #143 (which was part of #116).

Back then, some elements like Marker didn't have a length. Conceptually, I think this makes a lot of sense. But we had issues with broadcasting to the correct dimensions in situations where we needed to assume a length. I think an example of this would be transfer map reduction.

We decided as a fix just to give the Element class a length such that every element has one, that is Markers for example simply have length = 0.0. Conceptually, I think this is not as clean, but it still makes sense. This does not play nicely with the registration of nn.Parameter, which I believe is the only reason why we didn't allow length to be optimisable.

I think that now there are new considerations for this with #142. Here it makes even less sense for the element to have a length property. In fact, I think it would be too easily confused with what will probably be called something like effect_length.

A possible solution for this would be to just let every element define its own length again, i.e. no length in the base class. In my opinion this would be the cleanest solution. We would just have to find a good solution for #143. Admittedly, the original solution we found for this may have been just that little bit too hacky. 🙄

... so probably reverse and find a cleaner alternative for 9b9abe4.

RTSandberg commented 4 months ago

It sounds like you have thought through this some already and it's a bit of a thorny issue; I look forward to the fix!