From 4ff5fca67b8e073bbe34f191798ceb72dbc0b4a7 Mon Sep 17 00:00:00 2001 From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com> Date: Wed, 6 Nov 2024 17:43:59 -0800 Subject: [PATCH] Add fractional_split method --- docs/examples/particle_examples.ipynb | 396 +++++++------------------- pmd_beamphysics/particles.py | 62 ++++ 2 files changed, 161 insertions(+), 297 deletions(-) diff --git a/docs/examples/particle_examples.ipynb b/docs/examples/particle_examples.ipynb index f96320c..60afff8 100644 --- a/docs/examples/particle_examples.ipynb +++ b/docs/examples/particle_examples.ipynb @@ -10,14 +10,17 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:16.493563Z", - "iopub.status.busy": "2024-10-17T23:20:16.493231Z", - "iopub.status.idle": "2024-10-17T23:20:16.756224Z", - "shell.execute_reply": "2024-10-17T23:20:16.755978Z" - } - }, + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, "outputs": [], "source": [ "# Nicer plotting\n", @@ -37,14 +40,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:16.757736Z", - "iopub.status.busy": "2024-10-17T23:20:16.757618Z", - "iopub.status.idle": "2024-10-17T23:20:17.066823Z", - "shell.execute_reply": "2024-10-17T23:20:17.066533Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "from pmd_beamphysics import ParticleGroup" @@ -53,14 +49,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.068335Z", - "iopub.status.busy": "2024-10-17T23:20:17.068183Z", - "iopub.status.idle": "2024-10-17T23:20:17.075538Z", - "shell.execute_reply": "2024-10-17T23:20:17.075293Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P = ParticleGroup(\"data/bmad_particles2.h5\")\n", @@ -70,14 +59,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.097947Z", - "iopub.status.busy": "2024-10-17T23:20:17.097804Z", - "iopub.status.idle": "2024-10-17T23:20:17.100584Z", - "shell.execute_reply": "2024-10-17T23:20:17.100339Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.energy" @@ -86,14 +68,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.101830Z", - "iopub.status.busy": "2024-10-17T23:20:17.101729Z", - "iopub.status.idle": "2024-10-17T23:20:17.104641Z", - "shell.execute_reply": "2024-10-17T23:20:17.104411Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P[\"mean_energy\"], P.units(\"mean_energy\")" @@ -102,14 +77,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.105799Z", - "iopub.status.busy": "2024-10-17T23:20:17.105706Z", - "iopub.status.idle": "2024-10-17T23:20:17.117036Z", - "shell.execute_reply": "2024-10-17T23:20:17.116807Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.where(P.x < P[\"mean_x\"])" @@ -118,14 +86,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.118234Z", - "iopub.status.busy": "2024-10-17T23:20:17.118138Z", - "iopub.status.idle": "2024-10-17T23:20:17.568107Z", - "shell.execute_reply": "2024-10-17T23:20:17.567825Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "a = P.plot(\"x\", \"px\", figsize=(8, 8))" @@ -134,14 +95,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.570337Z", - "iopub.status.busy": "2024-10-17T23:20:17.570233Z", - "iopub.status.idle": "2024-10-17T23:20:17.813384Z", - "shell.execute_reply": "2024-10-17T23:20:17.813006Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.write_elegant(\"elegant_particles.txt\", verbose=True)" @@ -164,14 +118,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.818500Z", - "iopub.status.busy": "2024-10-17T23:20:17.818325Z", - "iopub.status.idle": "2024-10-17T23:20:17.821375Z", - "shell.execute_reply": "2024-10-17T23:20:17.820860Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.x" @@ -187,14 +134,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.823608Z", - "iopub.status.busy": "2024-10-17T23:20:17.823471Z", - "iopub.status.idle": "2024-10-17T23:20:17.826057Z", - "shell.execute_reply": "2024-10-17T23:20:17.825834Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.gamma" @@ -210,14 +150,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.827703Z", - "iopub.status.busy": "2024-10-17T23:20:17.827508Z", - "iopub.status.idle": "2024-10-17T23:20:17.832477Z", - "shell.execute_reply": "2024-10-17T23:20:17.832144Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "len(P), P[\"n_particle\"]" @@ -240,14 +173,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.840376Z", - "iopub.status.busy": "2024-10-17T23:20:17.840165Z", - "iopub.status.idle": "2024-10-17T23:20:17.844354Z", - "shell.execute_reply": "2024-10-17T23:20:17.843905Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.avg(\"gamma\"), P.std(\"p\")" @@ -263,14 +189,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.846001Z", - "iopub.status.busy": "2024-10-17T23:20:17.845906Z", - "iopub.status.idle": "2024-10-17T23:20:17.851579Z", - "shell.execute_reply": "2024-10-17T23:20:17.851273Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.cov(\"x\", \"px\", \"y\", \"kinetic_energy\")" @@ -286,14 +205,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.853476Z", - "iopub.status.busy": "2024-10-17T23:20:17.853188Z", - "iopub.status.idle": "2024-10-17T23:20:17.860581Z", - "shell.execute_reply": "2024-10-17T23:20:17.860286Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P[\"sigma_x\"], P[\"sigma_energy\"], P[\"min_y\"], P[\"norm_emit_x\"], P[\"norm_emit_4d\"]" @@ -309,14 +221,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.862271Z", - "iopub.status.busy": "2024-10-17T23:20:17.862162Z", - "iopub.status.idle": "2024-10-17T23:20:17.866196Z", - "shell.execute_reply": "2024-10-17T23:20:17.865886Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P[\"cov_x__kinetic_energy\"]" @@ -332,14 +237,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.867645Z", - "iopub.status.busy": "2024-10-17T23:20:17.867546Z", - "iopub.status.idle": "2024-10-17T23:20:17.874795Z", - "shell.execute_reply": "2024-10-17T23:20:17.874415Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "H, edges = P.histogramdd(\"t\", \"delta_pz\", bins=(5, 10))\n", @@ -358,14 +256,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.876510Z", - "iopub.status.busy": "2024-10-17T23:20:17.876297Z", - "iopub.status.idle": "2024-10-17T23:20:17.916863Z", - "shell.execute_reply": "2024-10-17T23:20:17.916585Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "ss = P.slice_statistics(\"norm_emit_x\")\n", @@ -382,14 +273,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.918784Z", - "iopub.status.busy": "2024-10-17T23:20:17.918628Z", - "iopub.status.idle": "2024-10-17T23:20:17.986930Z", - "shell.execute_reply": "2024-10-17T23:20:17.986674Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "ss = P.slice_statistics(\"norm_emit_x\", \"norm_emit_y\", \"twiss\")\n", @@ -419,14 +303,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.988506Z", - "iopub.status.busy": "2024-10-17T23:20:17.988404Z", - "iopub.status.idle": "2024-10-17T23:20:17.995209Z", - "shell.execute_reply": "2024-10-17T23:20:17.994828Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.twiss(\"x\")" @@ -442,14 +319,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:17.996575Z", - "iopub.status.busy": "2024-10-17T23:20:17.996481Z", - "iopub.status.idle": "2024-10-17T23:20:18.090868Z", - "shell.execute_reply": "2024-10-17T23:20:18.090588Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.twiss(\"xy\", fraction=0.95)" @@ -465,14 +335,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.092851Z", - "iopub.status.busy": "2024-10-17T23:20:18.092732Z", - "iopub.status.idle": "2024-10-17T23:20:18.101285Z", - "shell.execute_reply": "2024-10-17T23:20:18.100955Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P2 = P.twiss_match(beta=30, alpha=-3, plane=\"x\")\n", @@ -493,14 +356,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.102934Z", - "iopub.status.busy": "2024-10-17T23:20:18.102784Z", - "iopub.status.idle": "2024-10-17T23:20:18.128067Z", - "shell.execute_reply": "2024-10-17T23:20:18.127790Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.resample()" @@ -516,14 +372,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.129460Z", - "iopub.status.busy": "2024-10-17T23:20:18.129369Z", - "iopub.status.idle": "2024-10-17T23:20:18.137052Z", - "shell.execute_reply": "2024-10-17T23:20:18.136774Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.resample(1000)" @@ -532,19 +381,63 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.138293Z", - "iopub.status.busy": "2024-10-17T23:20:18.138204Z", - "iopub.status.idle": "2024-10-17T23:20:18.396743Z", - "shell.execute_reply": "2024-10-17T23:20:18.396500Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.resample(1000).plot(\"x\", \"px\", bins=100)" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Splitting\n", + "\n", + "It is often useful to split particles along a key dimension for analysis. This method will split into chunks with approximately equal numbers of particles." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "P.split(n_chunks=3, key=\"z\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fractional split will use weights to partition the particles. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "P.fractional_split(fractions=[0.1, 0.9], key=\"z\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This is useful for splitting particles into head, core, and tail parts. Here, the 5% of the charge is in the tail, 90% is in the core, and 5 % is in the head:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for p in P.fractional_split(key=\"z\", fractions=[0.05, 0.95]):\n", + " p.plot(\"z\", \"energy\", bins=100, figsize=(3, 3))" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -558,14 +451,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.398128Z", - "iopub.status.busy": "2024-10-17T23:20:18.398013Z", - "iopub.status.idle": "2024-10-17T23:20:18.400353Z", - "shell.execute_reply": "2024-10-17T23:20:18.400088Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "(\n", @@ -580,14 +466,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.401556Z", - "iopub.status.busy": "2024-10-17T23:20:18.401460Z", - "iopub.status.idle": "2024-10-17T23:20:18.403371Z", - "shell.execute_reply": "2024-10-17T23:20:18.403136Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.units(\"mean_energy\")" @@ -596,14 +475,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.404591Z", - "iopub.status.busy": "2024-10-17T23:20:18.404499Z", - "iopub.status.idle": "2024-10-17T23:20:18.406349Z", - "shell.execute_reply": "2024-10-17T23:20:18.406148Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "str(P.units(\"cov_x__kinetic_energy\"))" @@ -621,14 +493,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.407490Z", - "iopub.status.busy": "2024-10-17T23:20:18.407417Z", - "iopub.status.idle": "2024-10-17T23:20:18.410062Z", - "shell.execute_reply": "2024-10-17T23:20:18.409837Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.std(\"z\"), P.std(\"t\")" @@ -644,14 +509,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.411207Z", - "iopub.status.busy": "2024-10-17T23:20:18.411135Z", - "iopub.status.idle": "2024-10-17T23:20:18.413231Z", - "shell.execute_reply": "2024-10-17T23:20:18.413014Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "t0 = P.avg(\"t\")\n", @@ -668,14 +526,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.414455Z", - "iopub.status.busy": "2024-10-17T23:20:18.414376Z", - "iopub.status.idle": "2024-10-17T23:20:18.417177Z", - "shell.execute_reply": "2024-10-17T23:20:18.416948Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.drift_to_t(t0)" @@ -691,14 +542,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.418394Z", - "iopub.status.busy": "2024-10-17T23:20:18.418317Z", - "iopub.status.idle": "2024-10-17T23:20:18.424567Z", - "shell.execute_reply": "2024-10-17T23:20:18.424343Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.std(\"z\"), P.avg(\"t\"), set(P.t)" @@ -718,14 +562,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.425709Z", - "iopub.status.busy": "2024-10-17T23:20:18.425638Z", - "iopub.status.idle": "2024-10-17T23:20:18.427775Z", - "shell.execute_reply": "2024-10-17T23:20:18.427538Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P.status[0:10] = 0\n", @@ -742,14 +579,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.428889Z", - "iopub.status.busy": "2024-10-17T23:20:18.428814Z", - "iopub.status.idle": "2024-10-17T23:20:18.448042Z", - "shell.execute_reply": "2024-10-17T23:20:18.447783Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P0 = P.where(P.status == 0)\n", @@ -767,14 +597,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.449394Z", - "iopub.status.busy": "2024-10-17T23:20:18.449317Z", - "iopub.status.idle": "2024-10-17T23:20:18.451377Z", - "shell.execute_reply": "2024-10-17T23:20:18.451179Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P2 = P1.copy()" @@ -790,14 +613,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.452540Z", - "iopub.status.busy": "2024-10-17T23:20:18.452467Z", - "iopub.status.idle": "2024-10-17T23:20:18.454606Z", - "shell.execute_reply": "2024-10-17T23:20:18.454384Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P2.charge = 9.8765e-12\n", @@ -814,14 +630,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.455747Z", - "iopub.status.busy": "2024-10-17T23:20:18.455676Z", - "iopub.status.idle": "2024-10-17T23:20:18.457416Z", - "shell.execute_reply": "2024-10-17T23:20:18.457207Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "\"id\" in P2" @@ -837,14 +646,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:20:18.458547Z", - "iopub.status.busy": "2024-10-17T23:20:18.458474Z", - "iopub.status.idle": "2024-10-17T23:20:18.460471Z", - "shell.execute_reply": "2024-10-17T23:20:18.460235Z" - } - }, + "metadata": {}, "outputs": [], "source": [ "P2.id, \"id\" in P2" @@ -1374,7 +1176,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.7" + "version": "3.13.0" }, "vscode": { "interpreter": { diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py index bca939b..5f0361f 100644 --- a/pmd_beamphysics/particles.py +++ b/pmd_beamphysics/particles.py @@ -2,6 +2,7 @@ import os import pathlib from copy import deepcopy +from typing import Union import numpy as np from h5py import File @@ -1138,6 +1139,67 @@ def slice_plot( def split(self, n_chunks=100, key="z"): return split_particles(self, n_chunks=n_chunks, key=key) + def fractional_split(self, fractions: Union[float, int, list], key: str): + """ + Split particles based on a given array key and a list of specified fractions or a single fraction. + + Parameters: + ----------- + fractions : float, int, or list of float/int + A fraction or a list of fractions used for splitting the particles. All values must be between 0 and 1 (exclusive). + + key : str + The attribute of particles to be used for sorting and splitting (e.g., 'z' for longitudinal position). + + Returns: + -------- + particle_groups : list of ParticleGroup + A list of ParticleGroup objects, each representing a subset of particles based on the specified fractions. + + Description: + ------------ + This function splits the given group of particles into multiple subsets based on the provided attribute (e.g., position). + The splits are determined such that each specified fraction of the total particle weights is separated. + The function first sorts the particles by the specified key, computes the cumulative sum of weights, + and determines the split values. It then returns a list of ParticleGroup objects representing the split subsets. + + """ + particles = self + + # Ensure fractions is a list + if isinstance(fractions, (float, int)): + fractions = [fractions] + + # Validate fraction values + if any(f <= 0 or f >= 1 for f in fractions): + raise ValueError("All fraction values must be between 0 and 1 (exclusive)") + + # Sort particles by the specified key + ixs = np.argsort(particles[key]) + sorted_particles = particles[ixs] + + # Sorted weights + ws = sorted_particles.weight + total_weight = np.sum(ws) + cw = np.cumsum(ws) / total_weight + + # Use vectorized searchsorted to determine split indices + fractions = np.array(fractions) + split_indices = np.searchsorted(cw, fractions, side="right") + + # Create ParticleGroup subsets for each split + particle_groups = [] + previous_index = 0 + for isplit in split_indices: + particle_groups.append(sorted_particles[previous_index:isplit]) + previous_index = isplit + + # Add the remaining particles to the last group + if previous_index < len(sorted_particles): + particle_groups.append(sorted_particles[previous_index:]) + + return particle_groups + def copy(self): """Returns a deep copy""" return deepcopy(self)