diff --git a/docs/examples/labels.ipynb b/docs/examples/labels.ipynb index e24e0ae..8041fad 100644 --- a/docs/examples/labels.ipynb +++ b/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, diff --git a/docs/examples/particle_examples.ipynb b/docs/examples/particle_examples.ipynb index fef7f5d..cb4a888 100644 --- a/docs/examples/particle_examples.ipynb +++ b/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\")" ] }, { diff --git a/docs/examples/units.ipynb b/docs/examples/units.ipynb index 1f9854c..a520ead 100644 --- a/docs/examples/units.ipynb +++ b/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, diff --git a/pmd_beamphysics/labels.py b/pmd_beamphysics/labels.py index 53d5b42..20c9022 100644 --- a/pmd_beamphysics/labels.py +++ b/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}$" diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py index 6c8d8ac..d990378 100644 --- a/pmd_beamphysics/particles.py +++ b/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}" diff --git a/pmd_beamphysics/plot.py b/pmd_beamphysics/plot.py index 4886ae0..2a043b9 100644 --- a/pmd_beamphysics/plot.py +++ b/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) diff --git a/pmd_beamphysics/units.py b/pmd_beamphysics/units.py index 8502e65..814d932 100644 --- a/pmd_beamphysics/units.py +++ b/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) :] diff --git a/tests/test_particlegroup.py b/tests/test_particlegroup.py index 01c0061..465efe6 100644 --- a/tests/test_particlegroup.py +++ b/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)