Repository: roussel-ryan/gpsr
Pull request #32: Add class for generalized gpsr lattice
Author: login: lboult
Created: 2025-03-27T19:09:28Z | Merged: 2025-03-28T18:36:04Z

Added new class that allows the user to supply an arbitrary cheetah lattice with arbitrary variables and observables


Diff:
diff --git a/gpsr/modeling.py b/gpsr/modeling.py
index 37961e7..223b9cd 100644
--- a/gpsr/modeling.py
+++ b/gpsr/modeling.py
@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from copy import deepcopy
-from typing import Tuple
+from typing import Tuple, List
 
 import numpy as np
 import torch
@@ -15,6 +15,7 @@
     Screen,
 )
 from cheetah.particles import Beam
+from cheetah.accelerator import Element
 from gpsr.beams import BeamGenerator
 
 
@@ -195,3 +196,78 @@ def set_lattice_parameters(self, x: torch.Tensor) -> None:
         self.lattice.DIPOLE_TO_SCREEN.length.data = (
             self.l3 - self.l_bend / 2 / torch.cos(bend_angle)
         )
+
+
+class GenericGPSRLattice(GPSRLattice):
+    """
+    Attributes:
+        lattice: The base lattice structure used for beam tracking.
+        variable_elements: A list of tuples, where each tuple contains an element object and
+                          the name of the parameter to be varied as a string.
+        observable_elements: A list of elements that have the 'reading' property.
+    """
+
+    def __init__(
+        self,
+        lattice,
+        variable_elements: List[Tuple[Element, str]],
+        observable_elements: List[Element],
+    ):
+        """
+        Initializes the GPSRLattice instance.
+
+        Args:
+            lattice: The cheetah lattice used for beam tracking.
+            variable_elements: A list of tuples, where each tuple contains a cheetah element object and
+                              the name of the parameter to be varied as a string.
+            observable_elements: A list of elements that have the 'reading' property.
+        """
+        super().__init__()
+
+        for element in variable_elements:
+            if not hasattr(element[0], element[1]):
+                raise AttributeError(
+                    f"Variable element {element[0].name} does not have parameter '{element[1]}'."
+                )
+
+        for element in observable_elements:
+            if not hasattr(element, "reading"):
+                raise AttributeError(
+                    f"Observable element {element.name} does not have a 'reading' property."
+                )
+
+        self.lattice = lattice
+        self.variable_elements = variable_elements
+        self.observable_elements = observable_elements
+
+    def track_and_observe(self, beam) -> Tuple[Tensor, ...]:
+        """
+        Tracks a beam through the lattice and collects observations from designated elements.
+
+        Args:
+            beam: The beam object to be tracked.
+
+        Returns:
+            A tuple of tensors representing the observations from the observable elements.
+        """
+
+        # Compute the merged transfer maps for the lattice
+        merged_lattice = self.lattice.transfer_maps_merged(beam)
+
+        # Apply the merged lattice transformations to the beam
+        merged_lattice(beam)
+
+        # Collect observations from the observable elements
+        observations = tuple([element.reading for element in self.observable_elements])
+
+        return observations
+
+    def set_lattice_parameters(self, settings: torch.Tensor):
+        """
+        Sets the parameters of variable elements in the lattice.
+
+        Args:
+            settings: A tensor containing the new parameter values for the variable elements.
+        """
+        for i, element in enumerate(self.variable_elements, 0):
+            setattr(element[0], element[1], settings[..., i])
diff --git a/gpsr/tests/test_modeling.py b/gpsr/tests/test_modeling.py
index 75fc419..c0813d8 100644
--- a/gpsr/tests/test_modeling.py
+++ b/gpsr/tests/test_modeling.py
@@ -3,14 +3,16 @@
 import pytest
 import torch
 from cheetah import Segment, ParticleBeam
-from cheetah.accelerator import (
-    Quadrupole,
-    Drift,
-    Screen,
-)
+from cheetah.accelerator import Quadrupole, Drift, Screen, TransverseDeflectingCavity
 from cheetah.particles import Beam
 
-from gpsr.modeling import GPSR, GPSRLattice, GPSRQuadScanLattice, GPSR6DLattice
+from gpsr.modeling import (
+    GPSR,
+    GPSRLattice,
+    GPSRQuadScanLattice,
+    GPSR6DLattice,
+    GenericGPSRLattice,
+)
 
 
 class TestModeling:
@@ -201,3 +203,98 @@ def test_gpsr_forward(self):
         assert isinstance(results, tuple)
         assert len(results) == 1
         assert torch.equal(results[0], torch.tensor([1.0]))
+
+    def test_generic_gpsr_lattice_initialization(self):
+        TDC = TransverseDeflectingCavity(
+            length=torch.tensor(1.0),
+            voltage=torch.tensor(100.0),
+            phase=torch.tensor(0.0),
+            name="TDC",
+        )
+        d1 = Drift(length=torch.tensor(1.0))
+        d2 = Drift(length=torch.tensor(1.0))
+        d3 = Drift(length=torch.tensor(1.0))
+
+        q1 = Quadrupole(length=torch.tensor(1.0), k1=torch.tensor(0.5), name="q1")
+
+        screen_1 = Screen(name="screen_1")
+        screen_2 = Screen(name="screen_2")
+
+        seg = Segment([TDC, d1, screen_1, d2, q1, d3, screen_2])
+
+        var_elements = [(seg.q1, "k1"), (seg.TDC, "voltage")]
+        obs_elements = [seg.screen_1, seg.screen_2]
+
+        lattice = GenericGPSRLattice(
+            seg, variable_elements=var_elements, observable_elements=obs_elements
+        )
+
+        assert isinstance(lattice.lattice, Segment)
+        assert lattice.variable_elements == var_elements
+        assert lattice.observable_elements == obs_elements
+
+    def test_generic_gpsr_lattice_set_lattice_parameters(self):
+        TDC = TransverseDeflectingCavity(
+            length=torch.tensor(1.0),
+            voltage=torch.tensor(100.0),
+            phase=torch.tensor(0.0),
+            name="TDC",
+        )
+        d1 = Drift(length=torch.tensor(1.0))
+        d2 = Drift(length=torch.tensor(1.0))
+        d3 = Drift(length=torch.tensor(1.0))
+
+        q1 = Quadrupole(length=torch.tensor(1.0), k1=torch.tensor(0.5), name="q1")
+
+        screen_1 = Screen(name="screen_1")
+        screen_2 = Screen(name="screen_2")
+
+        seg = Segment([TDC, d1, screen_1, d2, q1, d3, screen_2])
+
+        var_elements = [(seg.q1, "k1"), (seg.TDC, "voltage")]
+        obs_elements = [seg.screen_1, seg.screen_2]
+
+        lattice = GenericGPSRLattice(
+            seg, variable_elements=var_elements, observable_elements=obs_elements
+        )
+        x = torch.tensor([[0.1, 0.5], [0.1, 0.5]])
+        lattice.set_lattice_parameters(x)
+
+        assert torch.equal(lattice.lattice.q1.k1, torch.tensor([0.1, 0.1]))
+        assert torch.equal(lattice.lattice.TDC.voltage, torch.tensor([0.5, 0.5]))
+
+    def test_generic_gpsr_lattice_track_and_observe(self):
+        TDC = TransverseDeflectingCavity(
+            length=torch.tensor(1.0),
+            voltage=torch.tensor(100.0),
+            phase=torch.tensor(0.0),
+            name="TDC",
+        )
+        d1 = Drift(length=torch.tensor(1.0))
+        d2 = Drift(length=torch.tensor(1.0))
+        d3 = Drift(length=torch.tensor(1.0))
+
+        q1 = Quadrupole(length=torch.tensor(1.0), k1=torch.tensor(0.5), name="q1")
+
+        screen_1 = Screen(name="screen_1")
+        screen_2 = Screen(name="screen_2")
+
+        seg = Segment([TDC, d1, screen_1, d2, q1, d3, screen_2])
+
+        var_elements = [(seg.q1, "k1"), (seg.TDC, "voltage")]
+        obs_elements = [seg.screen_1, seg.screen_2]
+
+        lattice = GenericGPSRLattice(
+            seg, variable_elements=var_elements, observable_elements=obs_elements
+        )
+
+        beam = ParticleBeam(energy=torch.tensor(1e6), particles=torch.rand((10, 7)))
+
+        x = torch.tensor([[0.1, 0.5], [0.1, 0.5]])
+        lattice.set_lattice_parameters(x)
+        observations = lattice.track_and_observe(beam)
+
+        assert isinstance(observations, tuple)
+        assert len(observations) == 2
+        assert observations[0].shape == seg.screen_1.reading.shape
+        assert observations[1].shape == seg.screen_2.reading.shape


Reviews:
user: login: roussel-ryan
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-27T19:32:00Z
user: login: lboult
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T20:46:06Z
user: login: roussel-ryan
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-27T20:50:49Z
user: login: lboult
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T23:03:23Z
user: login: roussel-ryan
state: APPROVED
author_association: OWNER
submitted_at: 2025-03-28T18:35:57Z

Review comments:
diff_hunk: @@ -195,3 +196,75 @@ def set_lattice_parameters(self, x: torch.Tensor) -> None:
         self.lattice.DIPOLE_TO_SCREEN.length.data = (
             self.l3 - self.l_bend / 2 / torch.cos(bend_angle)
         )
+
+
+class GeneralizedGPSRLattice(GPSRLattice):
+    """    
+    Attributes:
+        lattice: The base lattice structure used for beam tracking.
+        variable_elements: A list of tuples, where each tuple contains an element object and 
+                          the name of the parameter to be varied as a string.
+        observable_elements: A list of elements that have the 'reading' property.
+    """
+    
+    def __init__(self, 
+                 lattice,
+                 variable_elements: List[Tuple[Element, str]] = [],
path: gpsr/modeling.py
user: login: roussel-ryan
body: can we replace the tuples with strings, such as `Q003.k1`. Python getattr might be able to interpret this properly or you can split into requisite parts via python `.split(".")`
created_at: 2025-03-27T19:29:45Z
author_association: OWNER
diff_hunk: @@ -195,3 +196,75 @@ def set_lattice_parameters(self, x: torch.Tensor) -> None:
         self.lattice.DIPOLE_TO_SCREEN.length.data = (
             self.l3 - self.l_bend / 2 / torch.cos(bend_angle)
         )
+
+
+class GeneralizedGPSRLattice(GPSRLattice):
+    """    
+    Attributes:
+        lattice: The base lattice structure used for beam tracking.
+        variable_elements: A list of tuples, where each tuple contains an element object and 
+                          the name of the parameter to be varied as a string.
+        observable_elements: A list of elements that have the 'reading' property.
+    """
+    
+    def __init__(self, 
+                 lattice,
+                 variable_elements: List[Tuple[Element, str]] = [],
path: gpsr/modeling.py
user: login: roussel-ryan
body: also this should probably be a required argument until we run into a case where its not required, same for observable elements below
created_at: 2025-03-27T19:30:37Z
author_association: OWNER
diff_hunk: @@ -195,3 +196,75 @@ def set_lattice_parameters(self, x: torch.Tensor) -> None:
         self.lattice.DIPOLE_TO_SCREEN.length.data = (
             self.l3 - self.l_bend / 2 / torch.cos(bend_angle)
         )
+
+
+class GeneralizedGPSRLattice(GPSRLattice):
+    """    
+    Attributes:
+        lattice: The base lattice structure used for beam tracking.
+        variable_elements: A list of tuples, where each tuple contains an element object and 
+                          the name of the parameter to be varied as a string.
+        observable_elements: A list of elements that have the 'reading' property.
+    """
+    
+    def __init__(self, 
+                 lattice,
+                 variable_elements: List[Tuple[Element, str]] = [],
+                 observable_elements: List[Element] = []):
+        """
+        Initializes the GPSRLattice instance.
+        
+        Args:
+            lattice: The cheetah lattice used for beam tracking.
+            variable_elements: A list of tuples, where each tuple contains a cheetah element object and 
+                              the name of the parameter to be varied as a string.
+            observable_elements: A list of elements that have the 'reading' property.
+        """
+        super().__init__()
+
+        for element in variable_elements:
+            if not hasattr(element[0], element[1]):
+                raise AttributeError(f"Variable element {element[0].name} does not have parameter '{element[1]}'.")
+        
+        for element in observable_elements:
+            if not hasattr(element, 'reading'):
+                raise AttributeError(f"Observable element {element.name} does not have a 'reading' property.")
+        
+        self.lattice = lattice
+        self.variable_elements = variable_elements
+        self.observable_elements = observable_elements
+    
+    def track_and_observe(self, beam) -> Tuple[Tensor, ...]:
+        """
+        Tracks a beam through the lattice and collects observations from designated elements.
+        
+        Args:
+            beam: The beam object to be tracked.
+        
+        Returns:
+            A tuple of tensors representing the observations from the observable elements.
+        """
+        
+        # Assign unit charge to all beam particles
+        beam.particle_charges = torch.ones(beam.x.shape)
path: gpsr/modeling.py
user: login: roussel-ryan
body: this has been fixed in the most recent version of cheetah so we don't need
created_at: 2025-03-27T19:31:25Z
author_association: OWNER
diff_hunk: @@ -195,3 +196,75 @@ def set_lattice_parameters(self, x: torch.Tensor) -> None:
         self.lattice.DIPOLE_TO_SCREEN.length.data = (
             self.l3 - self.l_bend / 2 / torch.cos(bend_angle)
         )
+
+
+class GeneralizedGPSRLattice(GPSRLattice):
+    """    
+    Attributes:
+        lattice: The base lattice structure used for beam tracking.
+        variable_elements: A list of tuples, where each tuple contains an element object and 
+                          the name of the parameter to be varied as a string.
+        observable_elements: A list of elements that have the 'reading' property.
+    """
+    
+    def __init__(self, 
+                 lattice,
+                 variable_elements: List[Tuple[Element, str]] = [],
path: gpsr/modeling.py
user: login: lboult
body: Agreed with it being a required argument. I can make them into strings, I thought though that having the first elements of the tuples as actual cheetah elements might be nice for error checking/ type checking though... I'm willing to concede the point though :)
created_at: 2025-03-27T20:46:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -195,3 +196,75 @@ def set_lattice_parameters(self, x: torch.Tensor) -> None:
         self.lattice.DIPOLE_TO_SCREEN.length.data = (
             self.l3 - self.l_bend / 2 / torch.cos(bend_angle)
         )
+
+
+class GeneralizedGPSRLattice(GPSRLattice):
+    """    
+    Attributes:
+        lattice: The base lattice structure used for beam tracking.
+        variable_elements: A list of tuples, where each tuple contains an element object and 
+                          the name of the parameter to be varied as a string.
+        observable_elements: A list of elements that have the 'reading' property.
+    """
+    
+    def __init__(self, 
+                 lattice,
+                 variable_elements: List[Tuple[Element, str]] = [],
path: gpsr/modeling.py
user: login: roussel-ryan
body: ok thats fine, lets have them be tuples with cheetah elements
created_at: 2025-03-27T20:50:49Z
author_association: OWNER
diff_hunk: @@ -195,3 +196,75 @@ def set_lattice_parameters(self, x: torch.Tensor) -> None:
         self.lattice.DIPOLE_TO_SCREEN.length.data = (
             self.l3 - self.l_bend / 2 / torch.cos(bend_angle)
         )
+
+
+class GeneralizedGPSRLattice(GPSRLattice):
+    """    
+    Attributes:
+        lattice: The base lattice structure used for beam tracking.
+        variable_elements: A list of tuples, where each tuple contains an element object and 
+                          the name of the parameter to be varied as a string.
+        observable_elements: A list of elements that have the 'reading' property.
+    """
+    
+    def __init__(self, 
+                 lattice,
+                 variable_elements: List[Tuple[Element, str]] = [],
+                 observable_elements: List[Element] = []):
+        """
+        Initializes the GPSRLattice instance.
+        
+        Args:
+            lattice: The cheetah lattice used for beam tracking.
+            variable_elements: A list of tuples, where each tuple contains a cheetah element object and 
+                              the name of the parameter to be varied as a string.
+            observable_elements: A list of elements that have the 'reading' property.
+        """
+        super().__init__()
+
+        for element in variable_elements:
+            if not hasattr(element[0], element[1]):
+                raise AttributeError(f"Variable element {element[0].name} does not have parameter '{element[1]}'.")
+        
+        for element in observable_elements:
+            if not hasattr(element, 'reading'):
+                raise AttributeError(f"Observable element {element.name} does not have a 'reading' property.")
+        
+        self.lattice = lattice
+        self.variable_elements = variable_elements
+        self.observable_elements = observable_elements
+    
+    def track_and_observe(self, beam) -> Tuple[Tensor, ...]:
+        """
+        Tracks a beam through the lattice and collects observations from designated elements.
+        
+        Args:
+            beam: The beam object to be tracked.
+        
+        Returns:
+            A tuple of tensors representing the observations from the observable elements.
+        """
+        
+        # Assign unit charge to all beam particles
+        beam.particle_charges = torch.ones(beam.x.shape)
path: gpsr/modeling.py
user: login: lboult
body: Done
created_at: 2025-03-27T23:03:23Z
author_association: CONTRIBUTOR

PR comments:
user: login: roussel-ryan
created_at: 2025-03-27T19:35:13Z
body: I will fix the workflows, although you will have to fix the linting of the new file you have added. You can install ruff via pip and run `ruff format .` in the gpsr directory to run formatting. welcome so software dev!
author_association: OWNER
user: login: roussel-ryan
created_at: 2025-03-27T21:01:48Z
body: can you also add a test in the tests directory for the generalized class (same file as the other models) that tests some of the basic functionality?
author_association: OWNER
user: login: lboult
created_at: 2025-03-27T23:09:36Z
body: Added unit tests now, based mainly on the 6dgpsr examples in there...

There are some tests in the examples that I didn't understand and so didn't add equivalents for in these tests, e.g:

```
        with pytest.raises(RuntimeError):
            lattice.track_and_observe(beam)
```

and

```
        assert torch.equal(observations[0], torch.eye(2))
        assert torch.equal(observations[1], torch.eye(2))
```

in ```test_gpsr_6d_lattice_track_and_observe``` : perhaps you have idea if they are needed / how they can be included in these tests?

author_association: CONTRIBUTOR
user: login: roussel-ryan
created_at: 2025-03-28T14:20:24Z
body: I think the tests look good to me, just need to fix the formatting issues, make sure you pull in recent changes of the github actions from the main branch
author_association: OWNER