Merge pull request #88 from ChristopherMayes/divide_by_t

Divide by c option

docs/examples/labels.ipynb

@@ -205,7 +205,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.7"
+   "version": "3.13.0"
   }
  },
  "nbformat": 4,

docs/examples/particle_examples.ipynb

@@ -7,6 +7,15 @@
     "# openPMD beamphysics examples"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pmd_beamphysics import ParticleGroup"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -27,15 +36,6 @@
     "# Basic Usage"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from pmd_beamphysics import ParticleGroup"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -79,7 +79,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "a = P.plot(\"x\", \"px\", figsize=(8, 8))"
+    "P.plot(\"x\", \"px\", figsize=(8, 8))"
    ]
   },
   {
@@ -487,6 +487,31 @@
     "P.std(\"z\"), P.avg(\"t\"), set(P.t)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "There is a special key `z/c` to divide `z` by the speed of light to get units of `s`. This key also works with statistics:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "P[\"sigma_z/c\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "P.units(\"sigma_z/c\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -745,7 +770,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "P.plot(\"z\", \"x\")"
+    "P.plot(\"z/c\", \"x\")"
    ]
   },
   {

docs/examples/units.ipynb

@@ -296,7 +296,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.7"
+   "version": "3.13.0"
   }
  },
  "nbformat": 4,

pmd_beamphysics/labels.py

@@ -4,10 +4,10 @@
     # 'status'
     "t": "t",
     "energy": "E",
-    "kinetic_energy": r"E_{kinetic}",
+    "kinetic_energy": r"E_\text{kinetic}",
     # 'mass',
-    # 'higher_order_energy_spread',
-    # 'higher_order_energy',
+    "higher_order_energy_spread": r"\sigma_{E_{(2)}}",
+    "higher_order_energy": r"E_{(2)}",
     "Ex": "E_x",
     "Ey": "E_y",
     "Ez": "E_z",
@@ -35,20 +35,20 @@
     "gamma": r"\gamma",
     "theta": r"\theta",
     "charge": "Q",
-    "twiss_alpha_x": r"Twiss\ \alpha_x",
-    "twiss_beta_x": r"Twiss\ \beta_x",
-    "twiss_gamma_x": r"Twiss\ \gamma_x",
-    "twiss_eta_x": r"Twiss\ \eta_x",
-    "twiss_etap_x": r"Twiss\ \eta'_x",
-    "twiss_emit_x": r"Twiss\ \epsilon_{x}",
-    "twiss_norm_emit_x": r"Twiss\ \epsilon_{n, x}",
-    "twiss_alpha_y": r"Twiss\ \alpha_y",
-    "twiss_beta_y": r"Twiss\ \beta_y",
-    "twiss_gamma_y": r"Twiss\ \gamma_y",
-    "twiss_eta_y": r"Twiss\ \eta_y",
-    "twiss_etap_y": r"Twiss\ \eta'_y",
-    "twiss_emit_y": r"Twiss\ \epsilon_{y}",
-    "twiss_norm_emit_y": r"Twiss\ \epsilon_{n, y}",
+    "twiss_alpha_x": r"\text{Twiss}\ \alpha_x",
+    "twiss_beta_x": r"\text{Twiss}\ \beta_x",
+    "twiss_gamma_x": r"\text{Twiss}\ \gamma_x",
+    "twiss_eta_x": r"\text{Twiss}\ \eta_x",
+    "twiss_etap_x": r"\text{Twiss}\ \eta'_x",
+    "twiss_emit_x": r"\text{Twiss}\ \epsilon_{x}",
+    "twiss_norm_emit_x": r"\text{Twiss}\ \epsilon_{n, x}",
+    "twiss_alpha_y": r"\text{Twiss}\ \alpha_y",
+    "twiss_beta_y": r"\text{Twiss}\ \beta_y",
+    "twiss_gamma_y": r"\text{Twiss}\ \gamma_y",
+    "twiss_eta_y": r"\text{Twiss}\ \eta_y",
+    "twiss_etap_y": r"\text{Twiss}\ \eta'_y",
+    "twiss_emit_y": r"\text{Twiss}\ \epsilon_{y}",
+    "twiss_norm_emit_y": r"\text{Twiss}\ \epsilon_{n, y}",
     # 'species_charge',
     # 'weight',
     "average_current": r"I_{av}",
@@ -98,6 +98,7 @@ def texlabel(key: str):
     if key in TEXLABEL:
         return TEXLABEL[key]
 
+    # Operators
     for prefix in ["sigma_", "mean_", "min_", "max_", "ptp_", "delta_"]:
         if key.startswith(prefix):
             pre = prefix[:-1]
@@ -114,7 +115,7 @@ def texlabel(key: str):
                 return rf"\left<{tex0}\right>"
 
     if key.startswith("cov_"):
-        subkeys = key.strip("cov_").split("__")
+        subkeys = key.removeprefix("cov_").split("__")
         tex0 = texlabel(subkeys[0])
         tex1 = texlabel(subkeys[1])
         return rf"\left<{tex0}, {tex1}\right>"
@@ -124,7 +125,7 @@ def texlabel(key: str):
         x, _, prefix = nice_array(wavelength)
         return rf"\mathrm{{bunching~at}}~{x:.1f}~\mathrm{{ {prefix}m }}"
 
-    return None
+    return rf"\mathrm{{ {key} }}"
 
 
 def mathlabel(*keys, units=None, tex=True):
@@ -167,6 +168,7 @@ def mathlabel(*keys, units=None, tex=True):
         label_list = [texlabel(key) or rf"\mathrm{{ {key} }}" for key in keys]
         label = ", ".join(label_list)
         if units:
+            units = units.replace("*", r"{\cdot}")
             label = rf"{label}~(\mathrm{{ {units} }} )"
 
         return rf"${label}$"

pmd_beamphysics/particles.py

@@ -783,8 +783,12 @@ def __getitem__(self, key):
         if not isinstance(key, str):
             return particle_parts(self, key)
 
+        # 'z/c' special case
+        if key == "z/c":
+            return self["z"] / (c_light)
+
         if key.startswith("cov_"):
-            subkeys = key[4:].split("__")
+            subkeys = key.removeprefix("cov_").split("__")
             assert (
                 len(subkeys) == 2
             ), f"Too many properties in covariance request: {key}"

pmd_beamphysics/plot.py

@@ -340,7 +340,7 @@ def marginal_plot(
         _, hist_prefix = nice_scale_prefix(hist_f / f1)
         ax_marg_x.set_ylabel(f"{hist_prefix}A")
     else:
-        ax_marg_x.set_ylabel(mathlabel(f"{hist_prefix}C/{ux}"))  # Always use tex
+        ax_marg_x.set_ylabel(f"{hist_prefix}" + mathlabel(f"C/{ux}"))  # Always use tex
 
     # Side histogram
     # Old method:
@@ -351,7 +351,7 @@ def marginal_plot(
     hist_width = np.diff(bin_edges)
     hist_y, hist_f, hist_prefix = nice_array(hist / hist_width)
     ax_marg_y.barh(hist_x, hist_y, hist_width, color="gray")
-    ax_marg_y.set_xlabel(mathlabel(f"{hist_prefix}C/{uy}"))  # Always use tex
+    ax_marg_y.set_xlabel(f"{hist_prefix}" + mathlabel(f"C/{uy}"))  # Always use tex
 
     # Turn off tick labels on marginals
     plt.setp(ax_marg_x.get_xticklabels(), visible=False)

pmd_beamphysics/units.py

@@ -438,7 +438,7 @@ def plottable_array(x, nice=True, lim=None):
 PARTICLEGROUP_UNITS = {}
 for k in ["n_particle", "status", "id", "n_alive", "n_dead"]:
     PARTICLEGROUP_UNITS[k] = unit("1")
-for k in ["t"]:
+for k in ["t", "z/c"]:
     PARTICLEGROUP_UNITS[k] = unit("s")
 for k in [
     "energy",
@@ -493,6 +493,7 @@ def pg_units(key):
     if key in PARTICLEGROUP_UNITS:
         return PARTICLEGROUP_UNITS[key]
 
+    # Operators
     for prefix in ["sigma_", "mean_", "min_", "max_", "ptp_", "delta_"]:
         if key.startswith(prefix):
             nkey = key[len(prefix) :]

tests/test_particlegroup.py

@@ -8,12 +8,13 @@
 
 ARRAY_KEYS = """
 x y z px py pz t status weight id
+z/c
 p energy kinetic_energy xp yp higher_order_energy
 r theta pr ptheta
 Lz
 gamma beta beta_x beta_y beta_z
 x_bar px_bar Jx Jy
-charge
+weight
 
 """.split()
 
@@ -34,6 +35,9 @@ def array_key(request):
     return request.param
 
 
+array_key2 = array_key
+
+
 @pytest.fixture(params=OPERATORS)
 def operator(request):
     return request.param
@@ -44,13 +48,9 @@ def test_operator(operator, array_key):
     P[key]
 
 
-# This is probably uneccessary:
-# @pytest.fixture(params=array_keys)
-# def array_key2(request):
-#     return request.param
-#
-# def test_cov(array_key, array_key2):
-#     P[f'cov_{array_key}__{array_key}']
+def test_cov_(array_key, array_key2):
+    key = f"cov_{array_key}__{array_key2}"
+    P[key]
 
 
 @pytest.fixture(params=SPECIAL_STATS)