Charge property

xsuite · Oct 13, 2023 · 5d61b21 · 5d61b21
1 parent 49feab9
commit 5d61b21
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/xplt/particles.py b/xplt/particles.py
@@ -70,7 +70,8 @@ def _init_particle_mixin(
             Jx = Prop("$J_x$", unit=f"({Prop.get('x').unit})^2/({Prop.get('betx').unit})"),       # Action Jx
             Jy = Prop("$J_y$", unit=f"({Prop.get('y').unit})^2/({Prop.get('bety').unit})"),       # Action Jy
             Θx = Prop("$Θ_x$", unit=f"rad"),  # angle
-            Θy = Prop("$Θ_y$", unit=f"rad"),
+            Θy = Prop("$Θ_y$", unit=f"rad"),  # angle
+            q  = Prop("q",     unit=Prop.get('q0').unit),
             # fmt: on
         )
         kwargs["display_units"] = defaults(
@@ -216,6 +217,9 @@ def _get_masked(self, particles, key, mask=None):
                     time = time + turn / frev
                 return np.array(time)
 
+            if key == "q":
+                return get(particles, "q0") * self._get_masked(particles, "charge_ratio", mask)
+
             # otherwise fail
             raise
 

diff --git a/xplt/units.py b/xplt/units.py
@@ -100,7 +100,7 @@ def get(p, custom_properties=None):
     chi     = Prop("$\\chi$",      unit="1",                                ),  # q/ q0 * m0/m = qratio / mratio
 
     #mass0   = Prop("m_\\mathrm{ref}",         unit='eV',                    ),   # Reference rest mass
-    #q0      = Prop("q_\\mathrm{ref}",         unit='e',                     ),   # Reference charge
+    q0      = Prop("q_\\mathrm{ref}",         unit='e',                     ),   # Reference charge
     #p0c     = Prop("p_\\mathrm{ref}c",        unit='eV',                    ),   # Reference momentum
     #energy0 = Prop("E_\\mathrm{ref}",         unit='eV',                    ),   # Reference energy (total energy)
     #gamma0  = Prop("\\gamma_\\mathrm{ref}",   unit="1",                     ),   # Reference relativistic gamma