Merge pull request #189 from Hespe/179-electron-mass-constants

Unify constants for electron mass
desy-ml · Jun 19, 2024 · a35fc03 · a35fc03
2 parents e2d0a7d + 19f8b41
commit a35fc03
Show file tree

Hide file tree

Showing 20 changed files with 11 additions and 165 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,6 +26,7 @@
 - Update reference from arXiv preprint to PRAB publication (see #166) (@jank324)
 - Rename converter modules to the respective name of the accelerator code (see #167) (@jank324)
 - Added imports to the code example in the README (see #188) (@jank324)
+- Refactor definitions of physical constants (see #189) (@hespe)
 
 ## [v0.6.3](https://github.com/desy-ml/cheetah/releases/tag/v0.6.3) (2024-03-28)
 

diff --git a/cheetah/accelerator/aperture.py b/cheetah/accelerator/aperture.py
@@ -3,8 +3,6 @@
 import matplotlib.pyplot as plt
 import torch
 from matplotlib.patches import Rectangle
-from scipy import constants
-from scipy.constants import physical_constants
 from torch import Size, nn
 
 from cheetah.particles import Beam, ParticleBeam
@@ -14,15 +12,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
-electron_mass_eV = torch.tensor(
-    physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-)
-
 
 class Aperture(Element):
     """

diff --git a/cheetah/accelerator/bpm.py b/cheetah/accelerator/bpm.py
@@ -4,8 +4,6 @@
 import matplotlib.pyplot as plt
 import torch
 from matplotlib.patches import Rectangle
-from scipy import constants
-from scipy.constants import physical_constants
 from torch import Size
 
 from cheetah.particles import Beam, ParameterBeam, ParticleBeam
@@ -15,15 +13,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
-electron_mass_eV = torch.tensor(
-    physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-)
-
 
 class BPM(Element):
     """

diff --git a/cheetah/accelerator/cavity.py b/cheetah/accelerator/cavity.py
@@ -14,11 +14,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
 electron_mass_eV = torch.tensor(
     physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
 )

diff --git a/cheetah/accelerator/custom_transfer_map.py b/cheetah/accelerator/custom_transfer_map.py
@@ -2,8 +2,6 @@
 
 import matplotlib.pyplot as plt
 import torch
-from scipy import constants
-from scipy.constants import physical_constants
 from torch import Size, nn
 
 from cheetah.particles import Beam
@@ -13,15 +11,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
-electron_mass_eV = torch.tensor(
-    physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-)
-
 
 class CustomTransferMap(Element):
     """

diff --git a/cheetah/accelerator/dipole.py b/cheetah/accelerator/dipole.py
@@ -4,8 +4,6 @@
 import numpy as np
 import torch
 from matplotlib.patches import Rectangle
-from scipy import constants
-from scipy.constants import physical_constants
 from torch import Size, nn
 
 from cheetah.track_methods import base_rmatrix, rotation_matrix
@@ -15,15 +13,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
-electron_mass_eV = torch.tensor(
-    physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-)
-
 
 class Dipole(Element):
     """

diff --git a/cheetah/accelerator/drift.py b/cheetah/accelerator/drift.py
@@ -2,7 +2,6 @@
 
 import matplotlib.pyplot as plt
 import torch
-from scipy import constants
 from scipy.constants import physical_constants
 from torch import Size, nn
 
@@ -12,11 +11,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
 electron_mass_eV = torch.tensor(
     physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
 )
@@ -53,7 +47,7 @@ def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
         device = self.length.device
         dtype = self.length.dtype
 
-        gamma = energy / rest_energy.to(device=device, dtype=dtype)
+        gamma = energy / electron_mass_eV.to(device=device, dtype=dtype)
         igamma2 = torch.zeros_like(gamma)  # TODO: Effect on gradients?
         igamma2[gamma != 0] = 1 / gamma[gamma != 0] ** 2
         beta = torch.sqrt(1 - igamma2)

diff --git a/cheetah/accelerator/element.py b/cheetah/accelerator/element.py
@@ -3,24 +3,13 @@
 
 import matplotlib.pyplot as plt
 import torch
-from scipy import constants
-from scipy.constants import physical_constants
 from torch import nn
 
 from cheetah.particles import Beam, ParameterBeam, ParticleBeam
 from cheetah.utils import UniqueNameGenerator
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
-electron_mass_eV = torch.tensor(
-    physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-)
-
 
 class Element(ABC, nn.Module):
     """

diff --git a/cheetah/accelerator/horizontal_corrector.py b/cheetah/accelerator/horizontal_corrector.py
@@ -4,7 +4,6 @@
 import numpy as np
 import torch
 from matplotlib.patches import Rectangle
-from scipy import constants
 from scipy.constants import physical_constants
 from torch import Size, nn
 
@@ -14,11 +13,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
 electron_mass_eV = torch.tensor(
     physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
 )
@@ -57,7 +51,7 @@ def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
         device = self.length.device
         dtype = self.length.dtype
 
-        gamma = energy / rest_energy.to(device=device, dtype=dtype)
+        gamma = energy / electron_mass_eV.to(device=device, dtype=dtype)
         igamma2 = torch.zeros_like(gamma)  # TODO: Effect on gradients?
         igamma2[gamma != 0] = 1 / gamma[gamma != 0] ** 2
         beta = torch.sqrt(1 - igamma2)

diff --git a/cheetah/accelerator/marker.py b/cheetah/accelerator/marker.py
@@ -2,8 +2,6 @@
 
 import matplotlib.pyplot as plt
 import torch
-from scipy import constants
-from scipy.constants import physical_constants
 from torch import Size
 
 from cheetah.particles import Beam
@@ -13,15 +11,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
-electron_mass_eV = torch.tensor(
-    physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-)
-
 
 class Marker(Element):
     """

diff --git a/cheetah/accelerator/quadrupole.py b/cheetah/accelerator/quadrupole.py
@@ -4,8 +4,6 @@
 import numpy as np
 import torch
 from matplotlib.patches import Rectangle
-from scipy import constants
-from scipy.constants import physical_constants
 from torch import Size, nn
 
 from cheetah.track_methods import base_rmatrix, misalignment_matrix
@@ -15,15 +13,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
-electron_mass_eV = torch.tensor(
-    physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-)
-
 
 class Quadrupole(Element):
     """

diff --git a/cheetah/accelerator/rbend.py b/cheetah/accelerator/rbend.py
@@ -1,8 +1,6 @@
 from typing import Optional, Union
 
 import torch
-from scipy import constants
-from scipy.constants import physical_constants
 from torch import nn
 
 from cheetah.utils import UniqueNameGenerator
@@ -11,15 +9,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
-electron_mass_eV = torch.tensor(
-    physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-)
-
 
 class RBend(Dipole):
     """

diff --git a/cheetah/accelerator/screen.py b/cheetah/accelerator/screen.py
@@ -4,8 +4,6 @@
 import matplotlib.pyplot as plt
 import torch
 from matplotlib.patches import Rectangle
-from scipy import constants
-from scipy.constants import physical_constants
 from torch import Size, nn
 from torch.distributions import MultivariateNormal
 
@@ -16,15 +14,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
-electron_mass_eV = torch.tensor(
-    physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-)
-
 
 class Screen(Element):
     """

diff --git a/cheetah/accelerator/segment.py b/cheetah/accelerator/segment.py
@@ -5,8 +5,6 @@
 import matplotlib
 import matplotlib.pyplot as plt
 import torch
-from scipy import constants
-from scipy.constants import physical_constants
 from torch import Size, nn
 
 from cheetah.converters.bmad import convert_bmad_lattice
@@ -22,15 +20,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
-electron_mass_eV = torch.tensor(
-    physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-)
-
 
 class Segment(Element):
     """

diff --git a/cheetah/accelerator/solenoid.py b/cheetah/accelerator/solenoid.py
@@ -3,7 +3,6 @@
 import matplotlib.pyplot as plt
 import torch
 from matplotlib.patches import Rectangle
-from scipy import constants
 from scipy.constants import physical_constants
 from torch import Size, nn
 
@@ -14,11 +13,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
 electron_mass_eV = torch.tensor(
     physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
 )
@@ -66,7 +60,7 @@ def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
         device = self.length.device
         dtype = self.length.dtype
 
-        gamma = energy / rest_energy.to(device=device, dtype=dtype)
+        gamma = energy / electron_mass_eV.to(device=device, dtype=dtype)
         c = torch.cos(self.length * self.k)
         s = torch.sin(self.length * self.k)
 

diff --git a/cheetah/accelerator/undulator.py b/cheetah/accelerator/undulator.py
@@ -3,7 +3,6 @@
 import matplotlib.pyplot as plt
 import torch
 from matplotlib.patches import Rectangle
-from scipy import constants
 from scipy.constants import physical_constants
 from torch import Size, nn
 
@@ -13,11 +12,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-rest_energy = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge  # electron mass
-)
 electron_mass_eV = torch.tensor(
     physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
 )
@@ -53,7 +47,7 @@ def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
         device = self.length.device
         dtype = self.length.dtype
 
-        gamma = energy / rest_energy.to(device=device, dtype=dtype)
+        gamma = energy / electron_mass_eV.to(device=device, dtype=dtype)
         igamma2 = (
             1 / gamma**2
             if gamma != 0