Remove private indicator _ from ParameterBeam.mu and ParameterBeam.cov

Author: jank324

cheetah/accelerator/cavity.py

@@ -106,9 +106,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_cov = torch.matmul(
-                tm, torch.matmul(incoming._cov, tm.transpose(-2, -1))
+            outgoingmu = torch.matmul(tm, incoming.mu.unsqueeze(-1)).squeeze(-1)
+            outgoingcov = torch.matmul(
+                tm, torch.matmul(incoming.cov, tm.transpose(-2, -1))
             )
         else:  # ParticleBeam
             outgoing_particles = torch.matmul(incoming.particles, tm.transpose(-2, -1))
@@ -128,12 +128,12 @@ def _track_beam(self, incoming: Beam) -> Beam:
             )
 
             if isinstance(incoming, ParameterBeam):
-                outgoing_mu[..., 5] = incoming._mu[..., 5] * incoming.energy * beta0 / (
+                outgoingmu[..., 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]
+                outgoingcov[..., 5, 5] = incoming.cov[..., 5, 5]
             else:  # ParticleBeam
                 outgoing_particles[..., 5] = incoming.particles[
                     ..., 5
@@ -196,22 +196,22 @@ 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
+                outgoingmu[..., 4] = outgoingmu[..., 4] + (
+                    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
+                outgoingcov[..., 4, 4] = (
+                    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
+                outgoingcov[..., 4, 5] = (
+                    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]
+                outgoingcov[..., 5, 4] = outgoingcov[..., 4, 5]
             else:  # ParticleBeam
                 outgoing_particles[..., 4] = outgoing_particles[..., 4] + (
                     T566.unsqueeze(-1) * incoming.particles[..., 5] ** 2
@@ -223,12 +223,12 @@ def _track_beam(self, incoming: Beam) -> Beam:
 
         if isinstance(incoming, ParameterBeam):
             outgoing = ParameterBeam(
-                mu=outgoing_mu,
-                cov=outgoing_cov,
+                mu=outgoingmu,
+                cov=outgoingcov,
                 energy=outgoing_energy,
                 total_charge=incoming.total_charge,
-                device=outgoing_mu.device,
-                dtype=outgoing_mu.dtype,
+                device=outgoingmu.device,
+                dtype=outgoingmu.dtype,
             )
             return outgoing
         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

@@ -151,13 +151,13 @@ def track(self, incoming: Beam) -> Beam:
             copy_of_incoming = incoming.clone()
 
             if isinstance(incoming, ParameterBeam):
-                copy_of_incoming._mu, _ = torch.broadcast_tensors(
-                    copy_of_incoming._mu, self.misalignment[..., 0]
+                copy_of_incoming.mu, _ = torch.broadcast_tensors(
+                    copy_of_incoming.mu, self.misalignment[..., 0]
                 )
-                copy_of_incoming._mu = copy_of_incoming._mu.clone()
+                copy_of_incoming.mu = copy_of_incoming.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):
                 copy_of_incoming.particles, _ = torch.broadcast_tensors(
                     copy_of_incoming.particles,
@@ -178,8 +178,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),
                 )
@@ -208,30 +208,28 @@ 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
+            transversemu = torch.stack(
+                [read_beam.mu[..., 0], read_beam.mu[..., 2]], dim=-1
             )
-            transverse_cov = torch.stack(
+            transversecov = 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,
             )
-            dist = MultivariateNormal(
-                loc=transverse_mu, covariance_matrix=transverse_cov
-            )
+            dist = MultivariateNormal(loc=transversemu, covariance_matrix=transversecov)
 
             left = self.extent[0]
             right = self.extent[1]

cheetah/particles/parameter_beam.py

@@ -37,13 +37,13 @@ def __init__(
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
 
-        self.register_buffer("_mu", None)
-        self.register_buffer("_cov", None)
+        self.register_buffer("mu", None)
+        self.register_buffer("cov", None)
         self.register_buffer("energy", None)
         self.register_buffer("total_charge", torch.tensor(0.0, **factory_kwargs))
 
-        self._mu = torch.as_tensor(mu, **factory_kwargs)
-        self._cov = torch.as_tensor(cov, **factory_kwargs)
+        self.mu = torch.as_tensor(mu, **factory_kwargs)
+        self.cov = torch.as_tensor(cov, **factory_kwargs)
         self.energy = torch.as_tensor(energy, **factory_kwargs)
         if total_charge is not None:
             self.total_charge = torch.as_tensor(total_charge, **factory_kwargs)
@@ -467,8 +467,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
@@ -498,82 +498,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 ParameterBeam(
-            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

@@ -13,8 +13,8 @@
 from cheetah.utils import (
     elementwise_linspace,
     format_axis_with_prefixed_unit,
-    unbiased_weighted_covariance,
     unbiased_weighted_std,
+    unbiased_weightedcovariance,
     verify_device_and_dtype,
 )
 
@@ -263,8 +263,8 @@ def from_parameters(
         cov = cov.expand(*vector_shape, 6, 6)
         particles = torch.ones((*vector_shape, num_particles, 7), **factory_kwargs)
         distributions = [
-            MultivariateNormal(sample_mean, covariance_matrix=sample_cov)
-            for sample_mean, sample_cov in zip(mean.view(-1, 6), cov.view(-1, 6, 6))
+            MultivariateNormal(sample_mean, covariance_matrix=samplecov)
+            for sample_mean, samplecov in zip(mean.view(-1, 6), cov.view(-1, 6, 6))
         ]
         particles[..., :6] = torch.stack(
             [distribution.sample((num_particles,)) for distribution in distributions],
@@ -887,7 +887,7 @@ def transformed_to(
         mu_x, mu_px, mu_y, mu_py, mu_tau, mu_p = torch.broadcast_tensors(
             mu_x, mu_px, mu_y, mu_py, mu_tau, mu_p
         )
-        new_mu = torch.stack([mu_x, mu_px, mu_y, mu_py, mu_tau, mu_p], dim=-1)
+        newmu = torch.stack([mu_x, mu_px, mu_y, mu_py, mu_tau, mu_p], dim=-1)
         sigma_x, sigma_px, sigma_y, sigma_py, sigma_tau, sigma_p = (
             torch.broadcast_tensors(
                 sigma_x, sigma_px, sigma_y, sigma_py, sigma_tau, sigma_p
@@ -897,7 +897,7 @@ def transformed_to(
             [sigma_x, sigma_px, sigma_y, sigma_py, sigma_tau, sigma_p], dim=-1
         )
 
-        old_mu = torch.stack(
+        oldmu = torch.stack(
             [self.mu_x, self.mu_px, self.mu_y, self.mu_py, self.mu_tau, self.mu_p],
             dim=-1,
         )
@@ -915,10 +915,10 @@ def transformed_to(
 
         phase_space = self.particles[..., :6]
         phase_space = (
-            (phase_space.transpose(-2, -1) - old_mu.unsqueeze(-1))
+            (phase_space.transpose(-2, -1) - oldmu.unsqueeze(-1))
             / old_sigma.unsqueeze(-1)
             * new_sigma.unsqueeze(-1)
-            + new_mu.unsqueeze(-1)
+            + newmu.unsqueeze(-1)
         ).transpose(-2, -1)
 
         particles = torch.ones(*phase_space.shape[:-1], 7)
@@ -1546,7 +1546,7 @@ def cov_xpx(self) -> torch.Tensor:
         Returns the covariance between x and px. :math:`\sigma_{x, px}^2`.
         It is weighted by the survival probability of the particles.
         """
-        return unbiased_weighted_covariance(
+        return unbiased_weightedcovariance(
             self.x, self.px, weights=self.survival_probabilities, dim=-1
         )
 
@@ -1556,7 +1556,7 @@ def cov_ypy(self) -> torch.Tensor:
         Returns the covariance between y and py. :math:`\sigma_{y, py}^2`.
         It is weighted by the survival probability of the particles.
         """
-        return unbiased_weighted_covariance(
+        return unbiased_weightedcovariance(
             self.y, self.py, weights=self.survival_probabilities, dim=-1
         )
 
@@ -1566,7 +1566,7 @@ def cov_taup(self) -> torch.Tensor:
         Returns the covariance between tau and p. :math:`\sigma_{\tau, p}^2`.
         It is weighted by the survival probability of the particles.
         """
-        return unbiased_weighted_covariance(
+        return unbiased_weightedcovariance(
             self.tau, self.p, weights=self.survival_probabilities, dim=-1
         )
 

cheetah/utils/__init__.py

@@ -9,8 +9,8 @@
     format_axis_with_prefixed_unit,
 )
 from .statistics import (  # noqa: F401
-    unbiased_weighted_covariance,
     unbiased_weighted_std,
     unbiased_weighted_variance,
+    unbiased_weightedcovariance,
 )
 from .unique_name_generator import UniqueNameGenerator  # noqa: F401

cheetah/utils/statistics.py

@@ -1,7 +1,7 @@
 import torch
 
 
-def unbiased_weighted_covariance(
+def unbiased_weightedcovariance(
     input1: torch.Tensor, input2: torch.Tensor, weights: torch.Tensor, dim: int = None
 ) -> torch.Tensor:
     """