Register element readings as torch buffers for proper dtype conversion

CHANGELOG.md

@@ -19,6 +19,7 @@
 ### 🐛 Bug fixes
 
 - Fix issue where a space before a comma could cause the Elegant and Bmad converters to fail (see #327) (@jank324)
+- Fix issue of `BPM` and `Screen` not properly converting the `dtype` of their readings (#335) (@Hespe)
 
 ### 🐆 Other
 

cheetah/accelerator/bpm.py

@@ -20,11 +20,19 @@ class BPM(Element):
     :param name: Unique identifier of the element.
     """
 
-    def __init__(self, is_active: bool = False, name: Optional[str] = None) -> None:
-        super().__init__(name=name)
+    def __init__(
+        self,
+        is_active: bool = False,
+        name: Optional[str] = None,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> None:
+        super().__init__(name=name, device=device, dtype=dtype)
 
         self.is_active = is_active
-        self.reading = None
+        self.register_buffer(
+            "reading", torch.tensor(torch.nan, device=device, dtype=dtype)
+        )
 
     @property
     def is_skippable(self) -> bool:

cheetah/accelerator/screen.py

@@ -74,6 +74,14 @@ def __init__(
         self.register_buffer("misalignment", torch.tensor((0.0, 0.0), **factory_kwargs))
         self.register_buffer("kde_bandwidth", torch.clone(self.pixel_size[0]))
 
+        # NOTE: According to its type hint, the operation on resolution below is a
+        # no-op. However, this form is robust against accidentally passing a
+        # torch.Tensor, preventing crashes in some instances.
+        self.register_buffer(
+            "cached_reading",
+            torch.full((resolution[0], resolution[1]), torch.nan, **factory_kwargs),
+        )
+
         if pixel_size is not None:
             self.pixel_size = torch.as_tensor(pixel_size, **factory_kwargs)
         if misalignment is not None:
@@ -82,7 +90,6 @@ def __init__(
             self.kde_bandwidth = torch.as_tensor(kde_bandwidth, **factory_kwargs)
 
         self.set_read_beam(None)
-        self.cached_reading = None
 
     @property
     def is_skippable(self) -> bool:
@@ -190,8 +197,7 @@ def track(self, incoming: Beam) -> Beam:
 
     @property
     def reading(self) -> torch.Tensor:
-        image = None
-        if self.cached_reading is not None:
+        if not torch.all(torch.isnan(self.cached_reading)):
             return self.cached_reading
 
         read_beam = self.get_read_beam()
@@ -296,7 +302,12 @@ def set_read_beam(self, value: Beam) -> None:
         # prevent `nn.Module` from intercepting the read beam, which is itself an
         # `nn.Module`, and registering it as a submodule of the screen.
         self._read_beam = value
-        self.cached_reading = None
+        self.cached_reading = torch.full(
+            (self.resolution[0], self.resolution[1]),
+            torch.nan,
+            device=self.cached_reading.device,
+            dtype=self.cached_reading.dtype,
+        )
 
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]

tests/test_bpm.py

@@ -20,3 +20,21 @@ def test_no_tracking_error(is_bpm_active, beam_class):
     segment.my_bpm.is_active = is_bpm_active
 
     _ = segment.track(beam)
+
+
+def test_reading_dtype_conversion():
+    """Test that a dtype conversion is correctly reflected in the BPM reading."""
+    segment = cheetah.Segment(
+        elements=[
+            cheetah.Drift(length=torch.tensor(1.0), dtype=torch.float32),
+            cheetah.BPM(name="bpm", is_active=True, dtype=torch.float32),
+        ],
+    )
+    beam = cheetah.ParameterBeam.from_parameters(dtype=torch.float32)
+    assert segment.bpm.reading.dtype == torch.float32
+
+    segment.track(beam)
+    assert segment.bpm.reading.dtype == torch.float32
+
+    segment = segment.double()
+    assert segment.bpm.reading.dtype == torch.float64

tests/test_screen.py

@@ -132,3 +132,30 @@ def test_reading_shows_beam_ares(screen_method):
     assert segment.AREABSCR1.reading.shape == (2040, 2448)
     assert torch.all(segment.AREABSCR1.reading >= 0.0)
     assert torch.any(segment.AREABSCR1.reading > 0.0)
+
+
+def test_reading_dtype_conversion():
+    """Test that a dtype conversion is correctly reflected in the screen reading."""
+    segment = cheetah.Segment(
+        elements=[
+            cheetah.Drift(length=torch.tensor(1.0), dtype=torch.float32),
+            cheetah.Screen(name="screen", is_active=True, dtype=torch.float32),
+        ],
+    )
+    beam = cheetah.ParameterBeam.from_parameters(dtype=torch.float32)
+    assert segment.screen.reading.dtype == torch.float32
+
+    # Test generating new image
+    cloned = segment.clone()
+    cloned.track(beam)
+    cloned = cloned.double()
+    assert torch.all(torch.isnan(cloned.screen.cached_reading))
+    assert cloned.screen.reading.dtype == torch.float64
+
+    # Test reading from cache
+    segment.track(beam)
+    assert segment.screen.reading.dtype == torch.float32
+    assert segment.screen.cached_reading.dtype == torch.float32
+    segment = segment.double()
+    assert segment.screen.cached_reading.dtype == torch.float64
+    assert segment.screen.reading.dtype == torch.float64