From df6057b3a2c9e539a847058a82828989fdabb420 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Tue, 22 Oct 2024 11:43:09 +0200 Subject: [PATCH 01/10] Clean up Beam class --- cheetah/particles/beam.py | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/cheetah/particles/beam.py b/cheetah/particles/beam.py index 427e0329..4bea4c7d 100644 --- a/cheetah/particles/beam.py +++ b/cheetah/particles/beam.py @@ -52,6 +52,8 @@ def from_parameters( cor_tau: Optional[torch.Tensor] = None, energy: Optional[torch.Tensor] = None, total_charge: Optional[torch.Tensor] = None, + device=None, + dtype=torch.float32, ) -> "Beam": """ Create beam that with given beam parameters. @@ -75,6 +77,9 @@ def from_parameters( :param cor_tau: Correlation between tau and p. :param energy: Reference energy of the beam in eV. :param total_charge: Total charge of the beam in C. + :param device: Device to create the beam on. If set to `"auto"` a CUDA GPU is + selected if available. The CPU is used otherwise. + :param dtype: Data type of the beam. """ raise NotImplementedError @@ -87,9 +92,9 @@ def from_twiss( beta_y: Optional[torch.Tensor] = None, alpha_y: Optional[torch.Tensor] = None, emittance_y: Optional[torch.Tensor] = None, - sigma_s: Optional[torch.Tensor] = None, + sigma_tau: Optional[torch.Tensor] = None, sigma_p: Optional[torch.Tensor] = None, - cor_s: Optional[torch.Tensor] = None, + cor_tau: Optional[torch.Tensor] = None, energy: Optional[torch.Tensor] = None, total_charge: Optional[torch.Tensor] = None, device=None, @@ -104,12 +109,15 @@ def from_twiss( :param beta_y: Beta function in y direction in meters. :param alpha_y: Alpha function in y direction in rad. :param emittance_y: Emittance in y direction in m*rad. - :param sigma_s: Sigma of the particle distribution in s direction in meters. - :param sigma_p: Sigma of the particle distribution in p direction in meters. - :param cor_s: Correlation of the particle distribution in s direction. + :param sigma_tau: Sigma of the particle distribution in longitudinal direction, + in meters. + :param sigma_p: Sigma of the particle distribution in p direction, + dimensionless. + :param cor_tau: Correlation between tau and p. :param energy: Energy of the beam in eV. :param total_charge: Total charge of the beam in C. - :param device: Device to create the beam on. + :param device: Device to create the beam on. If set to `"auto"` a CUDA GPU is + selected if available. The CPU is used otherwise. :param dtype: Data type of the beam. """ raise NotImplementedError @@ -140,6 +148,8 @@ def transformed_to( sigma_p: Optional[torch.Tensor] = None, energy: Optional[torch.Tensor] = None, total_charge: Optional[torch.Tensor] = None, + device=None, + dtype=torch.float32, ) -> "Beam": """ Create version of this beam that is transformed to new beam parameters. @@ -160,6 +170,9 @@ def transformed_to( dimensionless. :param energy: Reference energy of the beam in eV. :param total_charge: Total charge of the beam in C. + :param device: Device to create the transformed beam on. If set to `"auto"` a + CUDA GPU is selected if available. The CPU is used otherwise. + :param dtype: Data type of the transformed beam. """ # Figure out vector dimensions of the original beam and check that passed # arguments have the same vector dimensions. @@ -213,6 +226,8 @@ def transformed_to( sigma_p=sigma_p, energy=energy, total_charge=total_charge, + device=device, + dtype=dtype, ) @property @@ -264,7 +279,7 @@ def sigma_py(self) -> torch.Tensor: raise NotImplementedError @property - def mu_s(self) -> torch.Tensor: + def mu_tau(self) -> torch.Tensor: raise NotImplementedError @property From fbcf91795458b5a49d759c977ef619dfa1807f64 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Tue, 22 Oct 2024 12:06:15 +0200 Subject: [PATCH 02/10] Update beam docstrings --- cheetah/particles/parameter_beam.py | 3 +++ cheetah/particles/particle_beam.py | 7 ++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/cheetah/particles/parameter_beam.py b/cheetah/particles/parameter_beam.py index dcdbb8c0..ab509a11 100644 --- a/cheetah/particles/parameter_beam.py +++ b/cheetah/particles/parameter_beam.py @@ -307,6 +307,9 @@ def transformed_to( :param sigma_p: Sigma of the particle distribution in p, dimensionless. :param energy: Reference energy of the beam in eV. :param total_charge: Total charge of the beam in C. + :param device: Device to create the transformed beam on. If set to `"auto"` a + CUDA GPU is selected if available. The CPU is used otherwise. + :param dtype: Data type of the transformed beam. """ device = device if device is not None else self.mu_x.device dtype = dtype if dtype is not None else self.mu_x.dtype diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py index 5b83bab1..39e86aa3 100644 --- a/cheetah/particles/particle_beam.py +++ b/cheetah/particles/particle_beam.py @@ -24,6 +24,7 @@ class ParticleBeam(Beam): :param total_charge: Total charge of the beam in C. :param device: Device to move the beam's particle array to. If set to `"auto"` a CUDA GPU is selected if available. The CPU is used otherwise. + :param dtype: Data type of the generated particles. """ def __init__( @@ -99,6 +100,7 @@ def from_parameters( :total_charge: Total charge of the beam in C. :param device: Device to move the beam's particle array to. If set to `"auto"` a CUDA GPU is selected if available. The CPU is used otherwise. + :param dtype: Data type of the generated particles. """ # Set default values without function call in function signature @@ -310,7 +312,8 @@ def uniform_3d_ellipsoid( :param sigma_p: Sigma of the particle distribution in p, dimensionless. :param energy: Reference energy of the beam in eV. :param total_charge: Total charge of the beam in C. - :param device: Device to move the beam's particle array to. + :param device: Device to move the beam's particle array to. If set to `"auto"` a + CUDA GPU is selected if available. The CPU is used otherwise. :param dtype: Data type of the generated particles. :return: ParticleBeam with uniformly distributed particles inside an ellipsoid. @@ -450,6 +453,7 @@ def make_linspaced( :param energy: Energy of the beam in eV. :param device: Device to move the beam's particle array to. If set to `"auto"` a CUDA GPU is selected if available. The CPU is used otherwise. + :param dtype: Data type of the generated particles. """ # Set default values without function call in function signature @@ -582,6 +586,7 @@ def transformed_to( :param total_charge: Total charge of the beam in C. :param device: Device to move the beam's particle array to. If set to `"auto"` a CUDA GPU is selected if available. The CPU is used otherwise. + :param dtype: Data type of the transformed particles. """ device = device if device is not None else self.mu_x.device dtype = dtype if dtype is not None else self.mu_x.dtype From 260fc2ad710a42be154a00c4a2710c9d0c8a99e7 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Tue, 22 Oct 2024 12:17:03 +0200 Subject: [PATCH 03/10] Fix type hints for ParticleBeam --- cheetah/particles/particle_beam.py | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py index 39e86aa3..4942575f 100644 --- a/cheetah/particles/particle_beam.py +++ b/cheetah/particles/particle_beam.py @@ -56,7 +56,7 @@ def __init__( @classmethod def from_parameters( cls, - num_particles: Optional[torch.Tensor] = None, + num_particles: int = 1_000_000, mu_x: Optional[torch.Tensor] = None, mu_y: Optional[torch.Tensor] = None, mu_px: Optional[torch.Tensor] = None, @@ -104,9 +104,6 @@ def from_parameters( """ # Set default values without function call in function signature - num_particles = ( - num_particles if num_particles is not None else torch.tensor(100_000) - ) mu_x = mu_x if mu_x is not None else torch.tensor(0.0) mu_px = mu_px if mu_px is not None else torch.tensor(0.0) mu_y = mu_y if mu_y is not None else torch.tensor(0.0) @@ -190,7 +187,7 @@ def from_parameters( @classmethod def from_twiss( cls, - num_particles: Optional[torch.Tensor] = None, + num_particles: int = 1_000_000, beta_x: Optional[torch.Tensor] = None, alpha_x: Optional[torch.Tensor] = None, emittance_x: Optional[torch.Tensor] = None, @@ -230,9 +227,6 @@ def from_twiss( ), "Arguments must have the same shape." # Set default values without function call in function signature - num_particles = ( - num_particles if num_particles is not None else torch.tensor(1_000_000) - ) beta_x = beta_x if beta_x is not None else torch.full(shape, 0.0) alpha_x = alpha_x if alpha_x is not None else torch.full(shape, 0.0) emittance_x = emittance_x if emittance_x is not None else torch.full(shape, 0.0) @@ -278,7 +272,7 @@ def from_twiss( @classmethod def uniform_3d_ellipsoid( cls, - num_particles: Optional[torch.Tensor] = None, + num_particles: int = 1_000_000, radius_x: Optional[torch.Tensor] = None, radius_y: Optional[torch.Tensor] = None, radius_tau: Optional[torch.Tensor] = None, @@ -346,9 +340,6 @@ def uniform_3d_ellipsoid( # Set default values without function call in function signature # NOTE that this does not need to be done for values that are passed to the # Gaussian beam generation. - num_particles = ( - num_particles if num_particles is not None else torch.tensor(1_000_000) - ) radius_x = ( radius_x.expand(vector_shape) if radius_x is not None @@ -417,7 +408,7 @@ def uniform_3d_ellipsoid( @classmethod def make_linspaced( cls, - num_particles: Optional[torch.Tensor] = None, + num_particles: int = 10, mu_x: Optional[torch.Tensor] = None, mu_y: Optional[torch.Tensor] = None, mu_px: Optional[torch.Tensor] = None, @@ -457,7 +448,6 @@ def make_linspaced( """ # Set default values without function call in function signature - num_particles = num_particles if num_particles is not None else torch.tensor(10) mu_x = mu_x if mu_x is not None else torch.tensor(0.0) mu_px = mu_px if mu_px is not None else torch.tensor(0.0) mu_y = mu_y if mu_y is not None else torch.tensor(0.0) @@ -568,7 +558,6 @@ def transformed_to( """ Create version of this beam that is transformed to new beam parameters. - :param n: Number of particles to generate. :param mu_x: Center of the particle distribution on x in meters. :param mu_y: Center of the particle distribution on y in meters. :param mu_px: Center of the particle distribution on px, dimensionless. From 8beb510a87c46fc0146b8d4c2117090c0f4c0b7f Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Tue, 22 Oct 2024 12:44:53 +0200 Subject: [PATCH 04/10] Revert default particle count --- cheetah/particles/particle_beam.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py index 4942575f..035b2c44 100644 --- a/cheetah/particles/particle_beam.py +++ b/cheetah/particles/particle_beam.py @@ -56,7 +56,7 @@ def __init__( @classmethod def from_parameters( cls, - num_particles: int = 1_000_000, + num_particles: int = 100_000, mu_x: Optional[torch.Tensor] = None, mu_y: Optional[torch.Tensor] = None, mu_px: Optional[torch.Tensor] = None, From da706f68c69072ef242b1775abae8a1e657ee0e0 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Tue, 22 Oct 2024 13:00:58 +0200 Subject: [PATCH 05/10] Update changelog --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 65c9c24d..c5cb0888 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,8 +6,8 @@ This is a major release with significant upgrades under the hood of Cheetah. Des ### 🚨 Breaking Changes -- Cheetah is now vectorised. This means that you can run multiple simulations in parallel by passing a batch of beams and settings, resulting a number of interfaces being changed. For Cheetah developers this means that you now have to account for an arbitrary-dimensional tensor of most of the properties of you element, rather than a single value, vector or whatever else a property was before. (see #116, #157, #170, #172, #173, #198, #208, #213, #215, #218, #229, #233, #258, #265) (@jank324, @cr-xu, @hespe, @roussel-ryan) -- The fifth particle coordinate `s` is renamed to `tau`. Now Cheetah uses the canonical variables in phase space $(x,px=\frac{P_x}{p_0},y,py, \tau=c\Delta t, \delta=\Delta E/{p_0 c})$. In addition, the trailing "s" was removed from some beam property names (e.g. `beam.xs` becomes `beam.x`). (see #163) (@cr-xu) +- Cheetah is now vectorised. This means that you can run multiple simulations in parallel by passing a batch of beams and settings, resulting a number of interfaces being changed. For Cheetah developers this means that you now have to account for an arbitrary-dimensional tensor of most of the properties of you element, rather than a single value, vector or whatever else a property was before. (see #116, #157, #170, #172, #173, #198, #208, #213, #215, #218, #229, #233, #258, #265, #284) (@jank324, @cr-xu, @hespe, @roussel-ryan) +- The fifth particle coordinate `s` is renamed to `tau`. Now Cheetah uses the canonical variables in phase space $(x,px=\frac{P_x}{p_0},y,py, \tau=c\Delta t, \delta=\Delta E/{p_0 c})$. In addition, the trailing "s" was removed from some beam property names (e.g. `beam.xs` becomes `beam.x`). (see #163, #284) (@cr-xu, @hespe) - `Screen` no longer blocks the beam (by default). To return to old behaviour, set `Screen.is_blocking = True`. (see #208) (@jank324, @roussel-ryan) ### 🚀 Features From 3e36a1570e55d19b3502f7942a44ca78675a8a3c Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Mon, 28 Oct 2024 08:58:56 +0100 Subject: [PATCH 06/10] Test for plotting with gradients --- tests/test_plotting.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index c1b3bfe7..e59acf39 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -103,3 +103,17 @@ def test_reference_particle_plot_vectorized_2d(): # Run the plotting to see if it raises an exception segment.plot_overview(incoming=incoming, resolution=0.1, vector_idx=(0, 2)) + + +def test_plotting_with_gradients(): + """ + Test that plotting doesn't raise an exception for segments that contain tensors + that require gradients. + """ + segment = cheetah.Segment( + elements=[cheetah.Drift(length=torch.tensor(1.0, requires_grad=True))] + ) + beam = cheetah.ParameterBeam.from_parameters() + + segment.plot_overview(incoming=beam) + segment.plot_twiss(incoming=beam) From d14645e8f2e5c7f365f023963407e4f347fb957e Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Mon, 28 Oct 2024 09:06:33 +0100 Subject: [PATCH 07/10] Call tensor.detach before passing into matplotlib --- cheetah/accelerator/segment.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/cheetah/accelerator/segment.py b/cheetah/accelerator/segment.py index 290195e8..d9e70f44 100644 --- a/cheetah/accelerator/segment.py +++ b/cheetah/accelerator/segment.py @@ -400,7 +400,7 @@ def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> No dimension_reordered_ss[vector_idx] if stacked_ss.dim() > 1 else dimension_reordered_ss - ) + ).detach() ax.plot([0, plot_ss[-1]], [0, 0], "--", color="black") @@ -465,17 +465,17 @@ def plot_reference_particle_traces( dimensions_reordered_ss[vector_idx] if stacked_ss.dim() > 1 else dimensions_reordered_ss - ) + ).detach() plot_xs = ( dimension_reordered_xs[vector_idx] if stacked_xs.dim() > 2 else dimension_reordered_xs - ) + ).detach() plot_ys = ( dimension_reordered_ys[vector_idx] if stacked_ys.dim() > 2 else dimension_reordered_ys - ) + ).detach() for particle_idx in range(num_particles): axx.plot(plot_ss, plot_xs[particle_idx]) @@ -558,7 +558,7 @@ def plot_twiss( dimension_reordered_s_positions[vector_idx] if stacked_s_positions.dim() > 1 else dimension_reordered_s_positions - ) + ).detach() broadcast_beta_x = torch.broadcast_tensors(*beta_x) stacked_beta_x = torch.stack(broadcast_beta_x) @@ -567,7 +567,7 @@ def plot_twiss( dimension_reordered_beta_x[vector_idx] if stacked_beta_x.dim() > 2 else dimension_reordered_beta_x - ) + ).detach() broadcast_beta_y = torch.broadcast_tensors(*beta_y) stacked_beta_y = torch.stack(broadcast_beta_y) @@ -576,7 +576,7 @@ def plot_twiss( dimension_reordered_beta_y[vector_idx] if stacked_beta_y.dim() > 2 else dimension_reordered_beta_y - ) + ).detach() if ax is None: fig = plt.figure() From 6e76c9638642c4cd01c32fd93f30d307eac46a0a Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Mon, 28 Oct 2024 09:08:59 +0100 Subject: [PATCH 08/10] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 65c9c24d..e49e66e5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,7 @@ This is a major release with significant upgrades under the hood of Cheetah. Des - Fix issue in Bmad import where collimators had no length by interpreting them as `Drift` + `Aperture` (see #249) (@jank324) - 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) ### 🐆 Other From b57100102031960bdb6683755594e8c07ffb7ca3 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Mon, 28 Oct 2024 10:32:57 +0100 Subject: [PATCH 09/10] Make `Beam` an abstract class --- cheetah/particles/beam.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/cheetah/particles/beam.py b/cheetah/particles/beam.py index 4bea4c7d..a0ce6c91 100644 --- a/cheetah/particles/beam.py +++ b/cheetah/particles/beam.py @@ -1,3 +1,4 @@ +from abc import ABC, abstractmethod from typing import Optional import torch @@ -7,7 +8,7 @@ electron_mass_eV = physical_constants["electron mass energy equivalent in MeV"][0] * 1e6 -class Beam(nn.Module): +class Beam(ABC, nn.Module): r""" Parent class to represent a beam of particles. You should not instantiate this class directly, but use one of the subclasses. @@ -35,6 +36,7 @@ class directly, but use one of the subclasses. empty = "I'm an empty beam!" @classmethod + @abstractmethod def from_parameters( cls, mu_x: Optional[torch.Tensor] = None, @@ -84,6 +86,7 @@ def from_parameters( raise NotImplementedError @classmethod + @abstractmethod def from_twiss( cls, beta_x: Optional[torch.Tensor] = None, @@ -123,6 +126,7 @@ def from_twiss( raise NotImplementedError @classmethod + @abstractmethod def from_ocelot(cls, parray) -> "Beam": """ Convert an Ocelot ParticleArray `parray` to a Cheetah Beam. @@ -130,6 +134,7 @@ def from_ocelot(cls, parray) -> "Beam": raise NotImplementedError @classmethod + @abstractmethod def from_astra(cls, path: str, **kwargs) -> "Beam": """Load an Astra particle distribution as a Cheetah Beam.""" raise NotImplementedError @@ -247,50 +252,62 @@ def parameters(self) -> dict: } @property + @abstractmethod def mu_x(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def sigma_x(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def mu_px(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def sigma_px(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def mu_y(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def sigma_y(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def mu_py(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def sigma_py(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def mu_tau(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def sigma_tau(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def mu_p(self) -> torch.Tensor: raise NotImplementedError @property + @abstractmethod def sigma_p(self) -> torch.Tensor: raise NotImplementedError @@ -314,11 +331,13 @@ def p0c(self) -> torch.Tensor: return self.relativistic_beta * self.relativistic_gamma * electron_mass_eV @property + @abstractmethod def sigma_xpx(self) -> torch.Tensor: # The covariance of (x,px) ~ $\sigma_{xpx}$ raise NotImplementedError @property + @abstractmethod def sigma_ypy(self) -> torch.Tensor: raise NotImplementedError From a56ca5712448686e6f67c3abea79de16b6729155 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Mon, 28 Oct 2024 10:57:35 +0100 Subject: [PATCH 10/10] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index c5cb0888..48958146 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -51,6 +51,7 @@ This is a major release with significant upgrades under the hood of Cheetah. Des - Add CI runs for macOS (arm64) and Windows (see #226) (@cr-xu, @jank324, @hespe) - Clean up CI pipelines (see #243, #244) (@jank324) - Fix logo display in README (see #252) (@jank324) +- Made `Beam` an abstract class (see #284) (@hespe) ## [v0.6.3](https://github.com/desy-ml/cheetah/releases/tag/v0.6.3) (2024-03-28)