Assume more reasonable default values for total_charge (#377)

* Assume more reasonable default values for `total_charge`

* Update changelog

* Fix failing test

---------

Co-authored-by: Chenran Xu <xuchenran@gmail.com>

Author: jank324

CHANGELOG.md

@@ -9,6 +9,7 @@
 - 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/particles/parameter_beam.py

@@ -37,19 +37,22 @@ def __init__(
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
 
+        self.species = species if species is not None else Species("electron")
+
         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.register_buffer(
+            "total_charge",
+            torch.as_tensor(self.species.charge_coulomb, **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)
 
-        self.species = species if species is not None else Species("electron")
-
     @classmethod
     def from_parameters(
         cls,

cheetah/particles/particle_beam.py

@@ -564,6 +564,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)
@@ -593,7 +595,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)