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 1/8] 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)

From c02a56643b63db6d09ad1410f440d1d2d00a48ad Mon Sep 17 00:00:00 2001
From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com>
Date: Wed, 6 Nov 2024 17:56:17 -0800
Subject: [PATCH 2/8] rearrange docs

---
 docs/examples/particle_examples.ipynb | 298 ++++++--------------------
 1 file changed, 71 insertions(+), 227 deletions(-)

diff --git a/docs/examples/particle_examples.ipynb b/docs/examples/particle_examples.ipynb
index 60afff8..fef7f5d 100644
--- a/docs/examples/particle_examples.ipynb
+++ b/docs/examples/particle_examples.ipynb
@@ -7,16 +7,6 @@
     "# openPMD beamphysics examples"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -387,57 +377,6 @@
     "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": {},
@@ -662,14 +601,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:18.461634Z",
-     "iopub.status.busy": "2024-10-17T23:20:18.461561Z",
-     "iopub.status.idle": "2024-10-17T23:20:18.462964Z",
-     "shell.execute_reply": "2024-10-17T23:20:18.462746Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "import h5py\n",
@@ -679,14 +611,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:18.464211Z",
-     "iopub.status.busy": "2024-10-17T23:20:18.464142Z",
-     "iopub.status.idle": "2024-10-17T23:20:18.473958Z",
-     "shell.execute_reply": "2024-10-17T23:20:18.473683Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "newh5file = \"particles.h5\"\n",
@@ -708,14 +633,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:18.475337Z",
-     "iopub.status.busy": "2024-10-17T23:20:18.475258Z",
-     "iopub.status.idle": "2024-10-17T23:20:18.478010Z",
-     "shell.execute_reply": "2024-10-17T23:20:18.477801Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "for key in [\"x\", \"px\", \"y\", \"py\", \"z\", \"pz\", \"t\", \"status\", \"weight\", \"id\"]:\n",
@@ -733,14 +651,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:18.479257Z",
-     "iopub.status.busy": "2024-10-17T23:20:18.479167Z",
-     "iopub.status.idle": "2024-10-17T23:20:18.483688Z",
-     "shell.execute_reply": "2024-10-17T23:20:18.483461Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "P2 == P"
@@ -756,14 +667,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:18.484951Z",
-     "iopub.status.busy": "2024-10-17T23:20:18.484870Z",
-     "iopub.status.idle": "2024-10-17T23:20:19.077238Z",
-     "shell.execute_reply": "2024-10-17T23:20:19.076925Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "P.write_astra(\"astra.dat\")"
@@ -772,14 +676,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:19.078673Z",
-     "iopub.status.busy": "2024-10-17T23:20:19.078564Z",
-     "iopub.status.idle": "2024-10-17T23:20:19.205207Z",
-     "shell.execute_reply": "2024-10-17T23:20:19.204803Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "!head astra.dat"
@@ -795,14 +692,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:19.207284Z",
-     "iopub.status.busy": "2024-10-17T23:20:19.206964Z",
-     "iopub.status.idle": "2024-10-17T23:20:19.217583Z",
-     "shell.execute_reply": "2024-10-17T23:20:19.217277Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "P.write(\"particles.h5\")"
@@ -827,14 +717,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:19.219281Z",
-     "iopub.status.busy": "2024-10-17T23:20:19.219174Z",
-     "iopub.status.idle": "2024-10-17T23:20:19.628355Z",
-     "shell.execute_reply": "2024-10-17T23:20:19.628099Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "P.plot(\"x\")"
@@ -850,14 +733,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:19.629649Z",
-     "iopub.status.busy": "2024-10-17T23:20:19.629563Z",
-     "iopub.status.idle": "2024-10-17T23:20:19.849704Z",
-     "shell.execute_reply": "2024-10-17T23:20:19.849459Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "P.slice_plot(\"norm_emit_x\", \"norm_emit_y\", ylim=(0, 1e-6))"
@@ -866,14 +742,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:19.850981Z",
-     "iopub.status.busy": "2024-10-17T23:20:19.850871Z",
-     "iopub.status.idle": "2024-10-17T23:20:20.158771Z",
-     "shell.execute_reply": "2024-10-17T23:20:20.158531Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "P.plot(\"z\", \"x\")"
@@ -889,14 +758,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:20.160286Z",
-     "iopub.status.busy": "2024-10-17T23:20:20.160188Z",
-     "iopub.status.idle": "2024-10-17T23:20:20.334573Z",
-     "shell.execute_reply": "2024-10-17T23:20:20.334321Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "P.slice_plot(\"sigma_x\", slice_key=\"Jx\")"
@@ -912,14 +774,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:20.335992Z",
-     "iopub.status.busy": "2024-10-17T23:20:20.335891Z",
-     "iopub.status.idle": "2024-10-17T23:20:20.643612Z",
-     "shell.execute_reply": "2024-10-17T23:20:20.643380Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "P.plot(\"x\", \"px\", ellipse=True)"
@@ -935,14 +790,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:20.645203Z",
-     "iopub.status.busy": "2024-10-17T23:20:20.645094Z",
-     "iopub.status.idle": "2024-10-17T23:20:21.062886Z",
-     "shell.execute_reply": "2024-10-17T23:20:21.062635Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "fig = P.plot(\"x\", return_figure=True)\n",
@@ -961,14 +809,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:21.064216Z",
-     "iopub.status.busy": "2024-10-17T23:20:21.064132Z",
-     "iopub.status.idle": "2024-10-17T23:20:21.139837Z",
-     "shell.execute_reply": "2024-10-17T23:20:21.139604Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "import copy\n",
@@ -990,14 +831,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:21.141053Z",
-     "iopub.status.busy": "2024-10-17T23:20:21.140946Z",
-     "iopub.status.idle": "2024-10-17T23:20:21.494659Z",
-     "shell.execute_reply": "2024-10-17T23:20:21.494422Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "P.plot(\"delta_z\", \"delta_p\", figsize=(8, 6))"
@@ -1013,14 +847,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:21.496329Z",
-     "iopub.status.busy": "2024-10-17T23:20:21.496222Z",
-     "iopub.status.idle": "2024-10-17T23:20:21.599826Z",
-     "shell.execute_reply": "2024-10-17T23:20:21.599600Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "H, edges = P.histogramdd(\"delta_z\", \"delta_p\", bins=(150, 150))\n",
@@ -1029,6 +856,57 @@
     "plt.imshow(H.T, origin=\"lower\", extent=extent, aspect=\"auto\", vmin=1e-15, cmap=cmap)"
    ]
   },
+  {
+   "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": {},
@@ -1041,14 +919,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:21.601117Z",
-     "iopub.status.busy": "2024-10-17T23:20:21.601028Z",
-     "iopub.status.idle": "2024-10-17T23:20:21.603036Z",
-     "shell.execute_reply": "2024-10-17T23:20:21.602808Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "from pmd_beamphysics import particle_paths\n",
@@ -1068,14 +939,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:21.604195Z",
-     "iopub.status.busy": "2024-10-17T23:20:21.604119Z",
-     "iopub.status.idle": "2024-10-17T23:20:21.606826Z",
-     "shell.execute_reply": "2024-10-17T23:20:21.606618Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "ppaths = particle_paths(h5)\n",
@@ -1092,14 +956,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:21.608025Z",
-     "iopub.status.busy": "2024-10-17T23:20:21.607926Z",
-     "iopub.status.idle": "2024-10-17T23:20:21.610352Z",
-     "shell.execute_reply": "2024-10-17T23:20:21.610162Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "ph5 = h5[ppaths[0]]\n",
@@ -1116,14 +973,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:21.611431Z",
-     "iopub.status.busy": "2024-10-17T23:20:21.611344Z",
-     "iopub.status.idle": "2024-10-17T23:20:21.615151Z",
-     "shell.execute_reply": "2024-10-17T23:20:21.614931Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "for component in all_components(ph5):\n",
@@ -1142,12 +992,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-10-17T23:20:21.616286Z",
-     "iopub.status.busy": "2024-10-17T23:20:21.616194Z",
-     "iopub.status.idle": "2024-10-17T23:20:21.618548Z",
-     "shell.execute_reply": "2024-10-17T23:20:21.618357Z"
-    },
     "tags": []
    },
    "outputs": [],

From 2e50472fc780adc996b41b28fbf695064afce808 Mon Sep 17 00:00:00 2001
From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com>
Date: Tue, 19 Nov 2024 19:41:05 -0800
Subject: [PATCH 3/8] Update pmd_beamphysics/particles.py

Co-authored-by: Ken Lauer <152229072+ken-lauer@users.noreply.github.com>
---
 pmd_beamphysics/particles.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py
index 5f0361f..2d150fe 100644
--- a/pmd_beamphysics/particles.py
+++ b/pmd_beamphysics/particles.py
@@ -1143,8 +1143,8 @@ 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:
-        -----------
+        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).
 

From e9bf158ed866eba5a2247e92c0c2575971ed8a32 Mon Sep 17 00:00:00 2001
From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com>
Date: Tue, 19 Nov 2024 19:41:13 -0800
Subject: [PATCH 4/8] Update pmd_beamphysics/particles.py

Co-authored-by: Ken Lauer <152229072+ken-lauer@users.noreply.github.com>
---
 pmd_beamphysics/particles.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py
index 2d150fe..52745f4 100644
--- a/pmd_beamphysics/particles.py
+++ b/pmd_beamphysics/particles.py
@@ -1151,8 +1151,8 @@ def fractional_split(self, fractions: Union[float, int, list], key: str):
         key : str
             The attribute of particles to be used for sorting and splitting (e.g., 'z' for longitudinal position).
 
-        Returns:
-        --------
+        Returns
+        -------
         particle_groups : list of ParticleGroup
             A list of ParticleGroup objects, each representing a subset of particles based on the specified fractions.
 

From a8928eee9be52c3485ef9a178bf9fa6cdb8a4236 Mon Sep 17 00:00:00 2001
From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com>
Date: Tue, 19 Nov 2024 19:41:20 -0800
Subject: [PATCH 5/8] Update pmd_beamphysics/particles.py

Co-authored-by: Ken Lauer <152229072+ken-lauer@users.noreply.github.com>
---
 pmd_beamphysics/particles.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py
index 52745f4..837a959 100644
--- a/pmd_beamphysics/particles.py
+++ b/pmd_beamphysics/particles.py
@@ -1156,8 +1156,8 @@ def fractional_split(self, fractions: Union[float, int, list], key: str):
         particle_groups : list of ParticleGroup
             A list of ParticleGroup objects, each representing a subset of particles based on the specified fractions.
 
-        Description:
-        ------------
+        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,

From 4834750193b7ad21d3602ec78a0be487f2f8505d Mon Sep 17 00:00:00 2001
From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com>
Date: Tue, 19 Nov 2024 19:41:34 -0800
Subject: [PATCH 6/8] Update pmd_beamphysics/particles.py

Co-authored-by: Ken Lauer <152229072+ken-lauer@users.noreply.github.com>
---
 pmd_beamphysics/particles.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py
index 837a959..4f7db76 100644
--- a/pmd_beamphysics/particles.py
+++ b/pmd_beamphysics/particles.py
@@ -1164,7 +1164,6 @@ def fractional_split(self, fractions: Union[float, int, list], key: str):
         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)):

From cbbf531417a2ebf3fa67510de11ac1c3657851fe Mon Sep 17 00:00:00 2001
From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com>
Date: Tue, 19 Nov 2024 19:43:49 -0800
Subject: [PATCH 7/8] fix docstring for fractions

---
 pmd_beamphysics/particles.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py
index 4f7db76..6c8d8ac 100644
--- a/pmd_beamphysics/particles.py
+++ b/pmd_beamphysics/particles.py
@@ -1145,7 +1145,7 @@ def fractional_split(self, fractions: Union[float, int, list], key: str):
 
         Parameters
         ----------
-        fractions : float, int, or list of float/int
+        fractions : float or list of float
             A fraction or a list of fractions used for splitting the particles. All values must be between 0 and 1 (exclusive).
 
         key : str
@@ -1174,8 +1174,8 @@ def fractional_split(self, fractions: Union[float, int, list], key: str):
             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]
+        ixs = np.argsort(self[key])
+        sorted_particles = self[ixs]
 
         # Sorted weights
         ws = sorted_particles.weight

From bd48d4392ff73323a51f6178463a18b80ae5c75b Mon Sep 17 00:00:00 2001
From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com>
Date: Wed, 20 Nov 2024 13:19:22 -0800
Subject: [PATCH 8/8] add test_fractional_split

---
 tests/test_particlegroup.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/test_particlegroup.py b/tests/test_particlegroup.py
index 5fccc5a..01c0061 100644
--- a/tests/test_particlegroup.py
+++ b/tests/test_particlegroup.py
@@ -85,3 +85,8 @@ def test_write_reload(tmp_path):
 
     P2.x += 1
     assert P != P2
+
+
+def test_fractional_split():
+    head, tail = P.fractional_split(0.5, "t")
+    head, core, tail = P.fractional_split((0.1, 0.9), "t")