Repository: desy-ml/cheetah (language: python)

Issue:
length cannot be a NN parameter
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


Discussion / hints:
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](https://github.com/desy-ml/cheetah/tree/164-length-cannot-be-a-nn-parameter)). 
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?


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 `Marker`s 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](https://github.com/desy-ml/cheetah/pull/116/commits/9b9abe480dad191d51d317f773e85d3377eda467).
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!

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 34279f5b..b3129d14 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -44,6 +44,7 @@ This is a major release with significant upgrades under the hood of Cheetah. Des
 - Fix NumPy 2 compatibility issues with PyTorch on Windows (see #220, #242) (@hespe)
 - Fix issue with Dipole hgap conversion in Bmad import (see #261) (@cr-xu)
 - Fix plotting for segments that contain tensors with `require_grad=True` (see #288) (@hespe)
+- Fix bug where `Element.length` could not be set as a `torch.nn.Parameter` (see #301) (@jank324, @hespe)
 
 ### 🐆 Other
 


Test patch:
diff --git a/tests/test_drift.py b/tests/test_drift.py
index bb5ba4b8..8e69d5f8 100644
--- a/tests/test_drift.py
+++ b/tests/test_drift.py
@@ -88,3 +88,23 @@ def test_drift_bmadx_tracking(dtype):
         atol=1e-14 if dtype == torch.float64 else 0.00001,
         rtol=1e-14 if dtype == torch.float64 else 1e-6,
     )
+
+
+def test_length_as_parameter():
+    """Test that the drift length can be set as a `torch.nn.Parameter`."""
+    length = torch.tensor(1.0)
+    parameter = torch.nn.Parameter(length)
+
+    # Create to equal drifts, one with Tensor, one with Parameter
+    drift = cheetah.Drift(length=length)
+    drift_parameter = cheetah.Drift(length=parameter)
+
+    incoming = cheetah.ParameterBeam.from_parameters()
+    outgoing = drift.track(incoming)
+    outgoing_parameter = drift_parameter.track(incoming)
+
+    # Check that all properties of the two outgoing beams are same
+    for buffer, buffer_parameter in zip(
+        outgoing.buffers(), outgoing_parameter.buffers()
+    ):
+        assert torch.allclose(buffer, buffer_parameter)