From 48b8909ab57b083e8b0f3c78c8c7bd7137411794 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Wed, 19 Jun 2024 10:48:26 +0200 Subject: [PATCH 1/3] Unify electron mass constants --- cheetah/accelerator/drift.py | 8 +------- cheetah/accelerator/horizontal_corrector.py | 8 +------- cheetah/accelerator/solenoid.py | 8 +------- cheetah/accelerator/undulator.py | 8 +------- cheetah/accelerator/vertical_corrector.py | 8 +------- cheetah/track_methods.py | 10 +++++----- 6 files changed, 10 insertions(+), 40 deletions(-) diff --git a/cheetah/accelerator/drift.py b/cheetah/accelerator/drift.py index 3218ec1f..1c29c9df 100644 --- 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/horizontal_corrector.py b/cheetah/accelerator/horizontal_corrector.py index ceaeab30..f1297101 100644 --- 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/solenoid.py b/cheetah/accelerator/solenoid.py index ae64e0ef..a92fc4be 100644 --- 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 index afeff4eb..6fa4b4ef 100644 --- 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 diff --git a/cheetah/accelerator/vertical_corrector.py b/cheetah/accelerator/vertical_corrector.py index 18795d10..47306e53 100644 --- a/cheetah/accelerator/vertical_corrector.py +++ b/cheetah/accelerator/vertical_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/track_methods.py b/cheetah/track_methods.py index 80a2216c..f7166b70 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -3,11 +3,11 @@ from typing import Optional import torch -from scipy import constants +from scipy.constants import physical_constants -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 +) def rotation_matrix(angle: torch.Tensor) -> torch.Tensor: @@ -56,7 +56,7 @@ def base_rmatrix( tilt = tilt if tilt is not None else torch.zeros_like(length) energy = energy if energy is not None else torch.zeros_like(length) - gamma = energy / REST_ENERGY.to(device=device, dtype=dtype) + gamma = energy / electron_mass_eV.to(device=device, dtype=dtype) igamma2 = torch.ones_like(length) igamma2[gamma != 0] = 1 / gamma[gamma != 0] ** 2 From 544ffd9a025be3edf353bd4c3c77dc70d3aa194a Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Wed, 19 Jun 2024 10:49:44 +0200 Subject: [PATCH 2/3] Remove unused constants and imports --- cheetah/accelerator/aperture.py | 11 ----------- cheetah/accelerator/bpm.py | 11 ----------- cheetah/accelerator/cavity.py | 5 ----- cheetah/accelerator/custom_transfer_map.py | 11 ----------- cheetah/accelerator/dipole.py | 11 ----------- cheetah/accelerator/element.py | 11 ----------- cheetah/accelerator/marker.py | 11 ----------- cheetah/accelerator/quadrupole.py | 11 ----------- cheetah/accelerator/rbend.py | 11 ----------- cheetah/accelerator/screen.py | 11 ----------- cheetah/accelerator/segment.py | 11 ----------- cheetah/particles/parameter_beam.py | 5 ----- cheetah/particles/particle_beam.py | 5 ----- 13 files changed, 125 deletions(-) diff --git a/cheetah/accelerator/aperture.py b/cheetah/accelerator/aperture.py index b179abdd..088f75f7 100644 --- 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 index acd446d7..5aede5ef 100644 --- 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 index 421aeb33..f5d74db9 100644 --- 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 index 64c01d9d..15fc6f7b 100644 --- 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 index ab3cd83a..2102dfa1 100644 --- 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/element.py b/cheetah/accelerator/element.py index b54f7f28..0f3a6297 100644 --- a/cheetah/accelerator/element.py +++ b/cheetah/accelerator/element.py @@ -3,8 +3,6 @@ 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 @@ -12,15 +10,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 Element(ABC, nn.Module): """ diff --git a/cheetah/accelerator/marker.py b/cheetah/accelerator/marker.py index f333058e..605c81df 100644 --- 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 index b87d0869..87022e87 100644 --- 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 index 765c87fe..39192ecd 100644 --- 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 index 6b76815a..7bd75609 100644 --- 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 index 68d6acf1..ee386720 100644 --- 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/particles/parameter_beam.py b/cheetah/particles/parameter_beam.py index 79378ce6..0e402ea3 100644 --- a/cheetah/particles/parameter_beam.py +++ b/cheetah/particles/parameter_beam.py @@ -2,14 +2,9 @@ import numpy as np import torch -from scipy.constants import physical_constants from .beam import Beam -electron_mass_eV = torch.tensor( - physical_constants["electron mass energy equivalent in MeV"][0] * 1e6 -) - class ParameterBeam(Beam): """ diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py index b6e014fa..a4f4d25e 100644 --- a/cheetah/particles/particle_beam.py +++ b/cheetah/particles/particle_beam.py @@ -1,15 +1,10 @@ from typing import Optional import torch -from scipy.constants import physical_constants from torch.distributions import MultivariateNormal from .beam import Beam -electron_mass_eV = torch.tensor( - physical_constants["electron mass energy equivalent in MeV"][0] * 1e6 -) - class ParticleBeam(Beam): """ From c9bcddd98367e21cddbcb1d887dca1a685eab3f3 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Wed, 19 Jun 2024 13:08:55 +0200 Subject: [PATCH 3/3] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d1169b44..85389670 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ - Split `accelerator` and `beam` into separate submodules (see #158) (@jank324) - 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) +- 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)