Merge branch 'master' into fix-ares-rl-speed

Author: cr-xu

CHANGELOG.md

@@ -9,6 +9,8 @@
 - The covariance properties of the different beam classes were renamed from names like `cor_x` and `sigma_xpx` to consistent names like `cov_xpx` (see #331) (@jank324)
 - The signature of the `transfer_map` method of all element subclasses was extended by a non-optional `species` argument (see #276) (@cr-xu, @jank324, @Hespe)
 - `ParticleBeam.plot_distribution` allows for Seaborn-style passing of `axs` and returns the latter as well. In line with that change for the purpose of overlaying distributions, the `contour` argument of `ParticleBeam.plot_2d_distribution` was replaced by a `style` argument. (see #330) (@jank324)
+- The default values for `total_charge` in both beam classes are no longer `0.0` but more sensible values (see #377) (@jank324)
+- `ParameterBeam._mu` and `ParameterBeam._cov` were renamed to `ParameterBeam.mu` and `ParameterBeam.cov` (see #378) (@jank324)
 
 ### 🚀 Features
 

cheetah/accelerator/cavity.py

@@ -111,9 +111,9 @@ def _track_beam(self, incoming: Beam) -> Beam:
 
         tm = self.transfer_map(incoming.energy, incoming.species)
         if isinstance(incoming, ParameterBeam):
-            outgoing_mu = torch.matmul(tm, incoming._mu.unsqueeze(-1)).squeeze(-1)
+            outgoing_mu = torch.matmul(tm, incoming.mu.unsqueeze(-1)).squeeze(-1)
             outgoing_cov = torch.matmul(
-                tm, torch.matmul(incoming._cov, tm.transpose(-2, -1))
+                tm, torch.matmul(incoming.cov, tm.transpose(-2, -1))
             )
         else:  # ParticleBeam
             outgoing_particles = torch.matmul(incoming.particles, tm.transpose(-2, -1))
@@ -133,12 +133,12 @@ def _track_beam(self, incoming: Beam) -> Beam:
             )
 
             if isinstance(incoming, ParameterBeam):
-                outgoing_mu[..., 5] = incoming._mu[..., 5] * incoming.energy * beta0 / (
+                outgoing_mu[..., 5] = incoming.mu[..., 5] * incoming.energy * beta0 / (
                     outgoing_energy * beta1
                 ) + self.voltage * beta0 / (outgoing_energy * beta1) * (
-                    torch.cos(-incoming._mu[..., 4] * beta0 * k + phi) - torch.cos(phi)
+                    torch.cos(-incoming.mu[..., 4] * beta0 * k + phi) - torch.cos(phi)
                 )
-                outgoing_cov[..., 5, 5] = incoming._cov[..., 5, 5]
+                outgoing_cov[..., 5, 5] = incoming.cov[..., 5, 5]
             else:  # ParticleBeam
                 outgoing_particles[..., 5] = incoming.particles[
                     ..., 5
@@ -202,19 +202,19 @@ def _track_beam(self, incoming: Beam) -> Beam:
 
             if isinstance(incoming, ParameterBeam):
                 outgoing_mu[..., 4] = outgoing_mu[..., 4] + (
-                    T566 * incoming._mu[..., 5] ** 2
-                    + T556 * incoming._mu[..., 4] * incoming._mu[..., 5]
-                    + T555 * incoming._mu[..., 4] ** 2
+                    T566 * incoming.mu[..., 5] ** 2
+                    + T556 * incoming.mu[..., 4] * incoming.mu[..., 5]
+                    + T555 * incoming.mu[..., 4] ** 2
                 )
                 outgoing_cov[..., 4, 4] = (
-                    T566 * incoming._cov[..., 5, 5] ** 2
-                    + T556 * incoming._cov[..., 4, 5] * incoming._cov[..., 5, 5]
-                    + T555 * incoming._cov[..., 4, 4] ** 2
+                    T566 * incoming.cov[..., 5, 5] ** 2
+                    + T556 * incoming.cov[..., 4, 5] * incoming.cov[..., 5, 5]
+                    + T555 * incoming.cov[..., 4, 4] ** 2
                 )
                 outgoing_cov[..., 4, 5] = (
-                    T566 * incoming._cov[..., 5, 5] ** 2
-                    + T556 * incoming._cov[..., 4, 5] * incoming._cov[..., 5, 5]
-                    + T555 * incoming._cov[..., 4, 4] ** 2
+                    T566 * incoming.cov[..., 5, 5] ** 2
+                    + T556 * incoming.cov[..., 4, 5] * incoming.cov[..., 5, 5]
+                    + T555 * incoming.cov[..., 4, 4] ** 2
                 )
                 outgoing_cov[..., 5, 4] = outgoing_cov[..., 4, 5]
             else:  # ParticleBeam

cheetah/accelerator/element.py

@@ -68,8 +68,8 @@ def track(self, incoming: Beam) -> Beam:
         """
         if isinstance(incoming, ParameterBeam):
             tm = self.transfer_map(incoming.energy, incoming.species)
-            mu = torch.matmul(tm, incoming._mu.unsqueeze(-1)).squeeze(-1)
-            cov = torch.matmul(tm, torch.matmul(incoming._cov, tm.transpose(-2, -1)))
+            mu = torch.matmul(tm, incoming.mu.unsqueeze(-1)).squeeze(-1)
+            cov = torch.matmul(tm, torch.matmul(incoming.cov, tm.transpose(-2, -1)))
             return ParameterBeam(
                 mu,
                 cov,

cheetah/accelerator/screen.py

@@ -166,12 +166,12 @@ def track(self, incoming: Beam) -> Beam:
 
             if isinstance(incoming, ParameterBeam):
                 broadcasted_mu, _ = torch.broadcast_tensors(
-                    copy_of_incoming._mu, self.misalignment[..., 0]
+                    copy_of_incoming.mu, self.misalignment[..., 0]
                 )
-                copy_of_incoming._mu = broadcasted_mu.clone()
+                copy_of_incoming.mu = broadcasted_mu.clone()
 
-                copy_of_incoming._mu[..., 0] -= self.misalignment[..., 0]
-                copy_of_incoming._mu[..., 2] -= self.misalignment[..., 1]
+                copy_of_incoming.mu[..., 0] -= self.misalignment[..., 0]
+                copy_of_incoming.mu[..., 2] -= self.misalignment[..., 1]
             elif isinstance(incoming, ParticleBeam):
                 broadcasted_particles, _ = torch.broadcast_tensors(
                     copy_of_incoming.particles,
@@ -192,8 +192,8 @@ def track(self, incoming: Beam) -> Beam:
         if self.is_active and self.is_blocking:
             if isinstance(incoming, ParameterBeam):
                 return ParameterBeam(
-                    mu=incoming._mu,
-                    cov=incoming._cov,
+                    mu=incoming.mu,
+                    cov=incoming.cov,
                     energy=incoming.energy,
                     total_charge=torch.zeros_like(incoming.total_charge),
                 )
@@ -222,23 +222,23 @@ def reading(self) -> torch.Tensor:
                 dtype=self.misalignment.dtype,
             )
         elif isinstance(read_beam, ParameterBeam):
-            if torch.numel(read_beam._mu[..., 0]) > 1:
+            if torch.numel(read_beam.mu[..., 0]) > 1:
                 raise NotImplementedError(
                     "`Screen` does not support vectorization of `ParameterBeam`. "
                     "Please use `ParticleBeam` instead. If this is a feature you would "
                     "like to see, please open an issue on GitHub."
                 )
 
             transverse_mu = torch.stack(
-                [read_beam._mu[..., 0], read_beam._mu[..., 2]], dim=-1
+                [read_beam.mu[..., 0], read_beam.mu[..., 2]], dim=-1
             )
             transverse_cov = torch.stack(
                 [
                     torch.stack(
-                        [read_beam._cov[..., 0, 0], read_beam._cov[..., 0, 2]], dim=-1
+                        [read_beam.cov[..., 0, 0], read_beam.cov[..., 0, 2]], dim=-1
                     ),
                     torch.stack(
-                        [read_beam._cov[..., 2, 0], read_beam._cov[..., 2, 2]], dim=-1
+                        [read_beam.cov[..., 2, 0], read_beam.cov[..., 2, 2]], dim=-1
                     ),
                 ],
                 dim=-1,

cheetah/particles/parameter_beam.py

@@ -37,9 +37,11 @@ def __init__(
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
 
-        self.register_buffer_or_parameter("_mu", torch.as_tensor(mu, **factory_kwargs))
+        self.species = species if species is not None else Species("electron")
+  
+        self.register_buffer_or_parameter("mu", torch.as_tensor(mu, **factory_kwargs))
         self.register_buffer_or_parameter(
-            "_cov", torch.as_tensor(cov, **factory_kwargs)
+            "cov", torch.as_tensor(cov, **factory_kwargs)
         )
         self.register_buffer_or_parameter(
             "energy", torch.as_tensor(energy, **factory_kwargs)
@@ -51,8 +53,6 @@ def __init__(
             ),
         )
 
-        self.species = species if species is not None else Species("electron")
-
     @classmethod
     def from_parameters(
         cls,
@@ -471,8 +471,8 @@ def as_particle_beam(self, num_particles: int) -> "ParticleBeam":  # noqa: F821
             cov_taup=self.cov_taup,
             energy=self.energy,
             total_charge=self.total_charge,
-            device=self._mu.device,
-            dtype=self._mu.dtype,
+            device=self.mu.device,
+            dtype=self.mu.dtype,
         )
 
     def linspaced(self, num_particles: int) -> "ParticleBeam":  # noqa: F821
@@ -502,82 +502,82 @@ def linspaced(self, num_particles: int) -> "ParticleBeam":  # noqa: F821
             energy=self.energy,
             total_charge=self.total_charge,
             species=self.species,
-            device=self._mu.device,
-            dtype=self._mu.dtype,
+            device=self.mu.device,
+            dtype=self.mu.dtype,
         )
 
     @property
     def mu_x(self) -> torch.Tensor:
-        return self._mu[..., 0]
+        return self.mu[..., 0]
 
     @property
     def sigma_x(self) -> torch.Tensor:
-        return torch.sqrt(torch.clamp_min(self._cov[..., 0, 0], 1e-20))
+        return torch.sqrt(torch.clamp_min(self.cov[..., 0, 0], 1e-20))
 
     @property
     def mu_px(self) -> torch.Tensor:
-        return self._mu[..., 1]
+        return self.mu[..., 1]
 
     @property
     def sigma_px(self) -> torch.Tensor:
-        return torch.sqrt(torch.clamp_min(self._cov[..., 1, 1], 1e-20))
+        return torch.sqrt(torch.clamp_min(self.cov[..., 1, 1], 1e-20))
 
     @property
     def mu_y(self) -> torch.Tensor:
-        return self._mu[..., 2]
+        return self.mu[..., 2]
 
     @property
     def sigma_y(self) -> torch.Tensor:
-        return torch.sqrt(torch.clamp_min(self._cov[..., 2, 2], 1e-20))
+        return torch.sqrt(torch.clamp_min(self.cov[..., 2, 2], 1e-20))
 
     @property
     def mu_py(self) -> torch.Tensor:
-        return self._mu[..., 3]
+        return self.mu[..., 3]
 
     @property
     def sigma_py(self) -> torch.Tensor:
-        return torch.sqrt(torch.clamp_min(self._cov[..., 3, 3], 1e-20))
+        return torch.sqrt(torch.clamp_min(self.cov[..., 3, 3], 1e-20))
 
     @property
     def mu_tau(self) -> torch.Tensor:
-        return self._mu[..., 4]
+        return self.mu[..., 4]
 
     @property
     def sigma_tau(self) -> torch.Tensor:
-        return torch.sqrt(torch.clamp_min(self._cov[..., 4, 4], 1e-20))
+        return torch.sqrt(torch.clamp_min(self.cov[..., 4, 4], 1e-20))
 
     @property
     def mu_p(self) -> torch.Tensor:
-        return self._mu[..., 5]
+        return self.mu[..., 5]
 
     @property
     def sigma_p(self) -> torch.Tensor:
-        return torch.sqrt(torch.clamp_min(self._cov[..., 5, 5], 1e-20))
+        return torch.sqrt(torch.clamp_min(self.cov[..., 5, 5], 1e-20))
 
     @property
     def cov_xpx(self) -> torch.Tensor:
-        return self._cov[..., 0, 1]
+        return self.cov[..., 0, 1]
 
     @property
     def cov_ypy(self) -> torch.Tensor:
-        return self._cov[..., 2, 3]
+        return self.cov[..., 2, 3]
 
     @property
     def cov_taup(self) -> torch.Tensor:
-        return self._cov[..., 4, 5]
+        return self.cov[..., 4, 5]
 
     def clone(self) -> "ParameterBeam":
         return self.__class__(
-            mu=self._mu.clone(),
-            cov=self._cov.clone(),
+            mu=self.mu.clone(),
+            cov=self.cov.clone(),
             energy=self.energy.clone(),
             total_charge=self.total_charge.clone(),
         )
 
     def __repr__(self) -> str:
         return (
-            f"{self.__class__.__name__}(mu={repr(self._mu)}, "
-            + f"cov={repr(self._cov)}, "
+            f"{self.__class__.__name__}(mu={repr(self.mu)}, "
+            + f"cov={repr(self.cov)}, "
             + f"energy={repr(self.energy)}, "
             + f"total_charge={repr(self.total_charge)}, "
             + f"species={repr(self.species)})"

cheetah/particles/particle_beam.py

@@ -570,6 +570,8 @@ def make_linspaced(
         )
         factory_kwargs = {"device": device, "dtype": dtype}
 
+        species = species if species is not None else Species("electron")
+
         # Set default values without function call in function signature
         mu_x = mu_x if mu_x is not None else torch.tensor(0.0, **factory_kwargs)
         mu_px = mu_px if mu_px is not None else torch.tensor(0.0, **factory_kwargs)
@@ -599,7 +601,7 @@ def make_linspaced(
         total_charge = (
             total_charge
             if total_charge is not None
-            else torch.tensor(0.0, **factory_kwargs)
+            else species.charge_coulomb * num_particles
         )
         particle_charges = (
             torch.ones((*total_charge.shape, num_particles), **factory_kwargs)

tests/test_differentiable.py

@@ -67,13 +67,13 @@ def test_ea_incoming_parameter_beam():
         "tests/resources/ACHIP_EA1_2021.1351.001"
     )
 
-    incoming_beam._mu = nn.Parameter(incoming_beam._mu)
-    incoming_beam._cov = nn.Parameter(incoming_beam._cov)
+    incoming_beam.mu = nn.Parameter(incoming_beam.mu)
+    incoming_beam.cov = nn.Parameter(incoming_beam.cov)
 
     outgoing_beam = ea.track(incoming_beam)
 
-    assert outgoing_beam._mu.grad_fn is not None
-    assert outgoing_beam._cov.grad_fn is not None
+    assert outgoing_beam.mu.grad_fn is not None
+    assert outgoing_beam.cov.grad_fn is not None
 
 
 def test_ea_incoming_particle_beam():

tests/test_parameter_beam.py

@@ -121,8 +121,8 @@ def test_from_twiss_dtype():
     assert np.isclose(beam.emittance_y.cpu().numpy(), 3.497810737006068e-09)
     assert np.isclose(beam.energy.cpu().numpy(), 6e6)
 
-    assert beam._mu.dtype == torch.float64
-    assert beam._cov.dtype == torch.float64
+    assert beam.mu.dtype == torch.float64
+    assert beam.cov.dtype == torch.float64
 
 
 def test_conversion_to_and_from_particle_beam():