diff --git a/docs/examples/19_particle_tracking_prt.ipynb b/docs/examples/19_particle_tracking_prt.ipynb index 165ac22e..87fcda91 100644 --- a/docs/examples/19_particle_tracking_prt.ipynb +++ b/docs/examples/19_particle_tracking_prt.ipynb @@ -4,7 +4,11 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Particle Tracking with PRT" + "# Particle Tracking with PRT\n", + "\n", + "## Setup\n", + "\n", + "### Packages" ] }, { @@ -17,12 +21,21 @@ "\n", "import flopy as fp\n", "import numpy as np\n", - "import pandas as pd\n", + "import geopandas as gpd\n", "import xarray as xr\n", - "\n", + "import pandas as pd\n", + "import matplotlib.pyplot as plt\n", + "import matplotlib as mpl\n", "import nlmod" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Logging" + ] + }, { "cell_type": "code", "execution_count": null, @@ -33,6 +46,13 @@ "nlmod.show_versions()" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Load base model" + ] + }, { "cell_type": "code", "execution_count": null, @@ -42,17 +62,30 @@ "model_ws = Path(\"./scratch_model\")\n", "model_name = \"from_scratch\"\n", "\n", - "# create new subdir for PRT model\n", - "prt_ws = Path(model_ws) / \"prt\"" + "ds = xr.open_dataset(model_ws / f\"{model_name}.nc\")\n", + "\n", + "# constants\n", + "wells = pd.DataFrame(\n", + " [[100, -50, -5, -10, -100.0], [200, 150, -20, -30, -300.0]],\n", + " columns=[\"x\", \"y\", \"top\", \"botm\", \"Q\"],\n", + " index=pd.Index([0, 1], name=\"well no.\"),\n", + ")\n", + "xyriv = [\n", + " (250.0, -500.0),\n", + " (300.0, -300.0),\n", + " (275.0, 0.0),\n", + " (200.0, 250.0),\n", + " (175.0, 500.0),\n", + "]" ] }, { - "cell_type": "code", - "execution_count": null, + "cell_type": "markdown", "metadata": {}, - "outputs": [], "source": [ - "ds = xr.open_dataset(model_ws / f\"{model_name}.nc\")" + "## PRT\n", + "\n", + "### Forward" ] }, { @@ -61,20 +94,24 @@ "metadata": {}, "outputs": [], "source": [ + "# create new subdir for PRT model\n", + "prt_ws = Path(model_ws) / \"prt_fw\"\n", + "\n", "# Create a new Simulation object\n", "simprt = nlmod.sim.sim(ds, sim_ws=prt_ws)\n", - "tdis = nlmod.sim.tdis(ds, simprt)\n", + "_ = nlmod.sim.tdis(ds, simprt)\n", "\n", "# Add PRT model\n", "prt = nlmod.prt.prt(ds, simprt)\n", "\n", "# DIS: discretization package\n", - "dis = nlmod.prt.dis(ds, prt)\n", + "_ = nlmod.prt.dis(ds, prt)\n", "\n", - "# MIP: matrix input package(?)\n", - "mip = nlmod.prt.mip(ds, prt, porosity=0.3)\n", + "# MIP: Model Input Package\n", + "_ = nlmod.prt.mip(ds, prt, porosity=0.3)\n", "\n", - "# PRP: particle release package\n", + "# PRP: particle release point package\n", + "# define particle release point every 11th cell in the first layer\n", "pdata = fp.modpath.ParticleData(\n", " [\n", " (k, i, j)\n", @@ -84,8 +121,8 @@ " ],\n", " structured=True,\n", ")\n", - "releasepts = list(pdata.to_prp(prt.modelgrid, global_xy=False))\n", - "prp = nlmod.prt.prp(ds, prt, packagedata=releasepts, perioddata={0: [\"FIRST\"]})\n", + "release_pts = list(pdata.to_prp(prt.modelgrid, global_xy=False))\n", + "prp = nlmod.prt.prp(ds, prt, packagedata=release_pts, perioddata={0: [\"FIRST\"]})\n", "\n", "# FMI: flow model interface\n", "fmi = nlmod.prt.fmi(ds, prt)\n", @@ -108,15 +145,17 @@ "metadata": {}, "outputs": [], "source": [ - "simprt.write_simulation()\n", - "simprt.run_simulation()" + "simprt.write_simulation(silent=True)\n", + "simprt.run_simulation(silent=True)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "Load pathline data" + "#### Load pathline data\n", + "\n", + "The read_pathlines function is a wrapper around pandas.read_csv but it adds the particle id for you. Additionallly it is usefull as documentation on the different istatus and ireason types." ] }, { @@ -125,14 +164,24 @@ "metadata": {}, "outputs": [], "source": [ - "pathlines = pd.read_csv(model_ws / \"prt\" / trackcsvfile_prt)" + "nlmod.prt.read_pathlines?" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pathlines_fw = nlmod.prt.read_pathlines(prt_ws / trackcsvfile_prt)\n", + "pathlines_fw.head()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "Plot results" + "#### Plot results" ] }, { @@ -141,39 +190,42 @@ "metadata": {}, "outputs": [], "source": [ - "pmv = fp.plot.PlotMapView(prt)\n", - "# pmv.plot_grid(lw=0.5, color=\"k\")\n", + "f, ax = plt.subplots(figsize=(8, 8))\n", + "pmv = fp.plot.PlotMapView(prt, ax=ax)\n", + "pmv.plot_grid(lw=0.1, color=\"k\")\n", "pmv.plot_pathline(\n", - " pathlines,\n", + " pathlines_fw,\n", " layer=\"all\",\n", " colors=\"0.7\",\n", " lw=1.0,\n", ")\n", "for ilay in [0, 1, 2]:\n", - " mask = pathlines[\"ilay\"] == (ilay + 1)\n", + " mask = pathlines_fw[\"ilay\"] == ilay\n", " pmv.plot_pathline(\n", - " pathlines.loc[mask],\n", - " layer=ilay + 1,\n", - " # colors=f\"C{ilay}\",\n", + " pathlines_fw.loc[mask],\n", + " layer=ilay - 3,\n", " lw=0.0,\n", " marker=\".\",\n", " ms=4,\n", " markercolor=f\"C{ilay}\",\n", " markerevery=5,\n", " )\n", - "pmv.plot_endpoint(pathlines, color=\"k\", marker=\".\", direction=\"starting\", zorder=10)\n", - "\n", - "# plot river\n", - "xyriv = [(250, -500), (300, -300), (275, 0), (200, 250), (175, 500)]\n", - "# plot well\n", - "pmv.ax.plot(\n", - " [xy[0] for xy in xyriv],\n", - " [xy[1] for xy in xyriv],\n", - " \"b-\",\n", - " lw=4,\n", - " label=\"river\",\n", - ")\n", - "pmv.ax.plot([100, 200], [-50, 150], \"rs\", ms=5)\n", + "pmv.plot_endpoint(pathlines_fw, color=\"k\", marker=\".\", direction=\"starting\", zorder=10)\n", + "\n", + "\n", + "# plot river and wells\n", + "def plot_river_and_wells(ax: plt.Axes):\n", + " ax.plot(\n", + " [xy[0] for xy in xyriv],\n", + " [xy[1] for xy in xyriv],\n", + " \"b-\",\n", + " lw=2,\n", + " label=\"river\",\n", + " )\n", + " ax.plot(wells[\"x\"].values, wells[\"y\"].values, \"rs\", ms=5)\n", + "\n", + "\n", + "plot_river_and_wells(ax=pmv.ax)\n", "handles = [\n", " pmv.ax.plot([], [], \"C0.\", ms=5, label=\"Layer 0\")[0],\n", " pmv.ax.plot([], [], \"C1.\", ms=5, label=\"Layer 1\")[0],\n", @@ -184,11 +236,338 @@ "pmv.ax.set_xlabel(\"X [m]\")\n", "pmv.ax.set_ylabel(\"Y [m]\");" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Backward and refined\n", + "\n", + "#### Refine and run gwf model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "refinement_extent = [0.0, 300.0, -100.0, 200.0]\n", + "refinement_features = (\n", + " [\n", + " nlmod.util.extent_to_polygon(refinement_extent),\n", + " ],\n", + " \"polygon\", # type of feature\n", + " 1, # refinement level\n", + ")\n", + "dsr = nlmod.grid.refine(ds, refinement_features=[refinement_features])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsr.attrs[\"model_name\"] = f\"{ds.attrs['model_name']}_ref\"\n", + "dsr.attrs[\"model_ws\"] = f\"./{ds.attrs['model_ws']}_ref\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sim = nlmod.sim.sim(dsr)\n", + "tdis = nlmod.sim.tdis(dsr, sim)\n", + "ims = nlmod.sim.ims(sim, complexity=\"SIMPLE\")\n", + "gwf = nlmod.gwf.gwf(dsr, sim)\n", + "disv = nlmod.gwf.disv(dsr, gwf)\n", + "npf = nlmod.gwf.npf(\n", + " dsr, gwf, save_flows=True, save_specific_discharge=True, save_saturation=True\n", + ")\n", + "ic = nlmod.gwf.ic(dsr, gwf, starting_head=1.0)\n", + "oc = nlmod.gwf.oc(dsr, gwf, save_head=True)\n", + "wel = nlmod.gwf.wells.wel_from_df(wells, gwf)\n", + "\n", + "riv_layer = 0 # add to first layer\n", + "bed_resistance = 0.1 # days\n", + "riv_stage = 1.0 # m NAP\n", + "riv_botm = -3.0 # m NAP\n", + "riv_data = nlmod.gwf.surface_water.rivdata_from_xylist(\n", + " gwf, xyriv, riv_layer, riv_stage, bed_resistance, riv_botm\n", + ")\n", + "\n", + "\n", + "def get_riv_cond_from_cid(icell2d: int, ds: xr.Dataset, bed_resistance: float) -> float:\n", + " \"\"\"Get river conductance from cell id.\"\"\"\n", + " area = ds.area.sel(icell2d=icell2d).values\n", + " # calculate conductance\n", + " return area / bed_resistance\n", + "\n", + "\n", + "riv_datac = [\n", + " [x[0], x[1], get_riv_cond_from_cid(x[0][1], dsr, x[2]), x[3]] for x in riv_data\n", + "]\n", + "riv = fp.mf6.ModflowGwfriv(gwf, stress_period_data={0: riv_datac})\n", + "\n", + "nlmod.sim.write_and_run(sim, dsr, silent=False)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Reverse heads and cellbudget files" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "hds = nlmod.gwf.get_headfile(ds=dsr, gwf=gwf, tdis=tdis)\n", + "hds_bw_path = hds.filename.parent / f\"{hds.filename.stem}_bkwd{hds.filename.suffix}\"\n", + "hds.reverse(hds_bw_path)\n", + "\n", + "cbc = nlmod.gwf.get_cellbudgetfile(ds=dsr, gwf=gwf, tdis=tdis)\n", + "cbc_bw_path = cbc.filename.parent / f\"{cbc.filename.stem}_bkwd{cbc.filename.suffix}\"\n", + "cbc.reverse(cbc_bw_path)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Create particle starts list\n", + "\n", + "Start a particle in each (active) cell in refined extent" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsr[\"z\"] = (dsr[\"top\"] + dsr[\"botm\"]) / 2\n", + "dsr[\"idomain\"] = (\n", + " (\"layer\", \"icell2d\"),\n", + " disv.idomain.array,\n", + ") # make sure to get idomain otherwhise you might start particles in inactive cells and everything crashes #beentheredonethat\n", + "\n", + "# map xy to cellids\n", + "nodes = dsr[[\"z\", \"idomain\"]].where(\n", + " (dsr.x >= refinement_extent[0])\n", + " & (dsr.x <= refinement_extent[1])\n", + " & (dsr.y >= refinement_extent[2])\n", + " & (dsr.y <= refinement_extent[3]),\n", + " drop=True,\n", + ")\n", + "ix = fp.utils.GridIntersect(gwf.modelgrid)\n", + "spatial_join = gpd.GeoDataFrame(\n", + " geometry=gpd.points_from_xy(nodes.x, nodes.y),\n", + ").sjoin(\n", + " gpd.GeoDataFrame(\n", + " {\"cellid\": ix.cellids},\n", + " geometry=ix.geoms,\n", + " ),\n", + " how=\"left\",\n", + ")\n", + "\n", + "# create release points\n", + "# define particle release point for every cell within refined extent\n", + "pid = 0\n", + "release_pts = []\n", + "for ilay in range(nodes.sizes[\"layer\"]):\n", + " nodes_lay = nodes.isel(layer=ilay)\n", + " active = (nodes_lay[\"idomain\"] > 0).values\n", + " j = spatial_join[\"cellid\"].values.astype(int)[active]\n", + " k = np.full_like(j, ilay, dtype=int)\n", + " x = nodes_lay[\"x\"].values[active]\n", + " y = nodes_lay[\"y\"].values[active]\n", + " z = nodes_lay[\"z\"].values[active]\n", + " pids = np.arange(pid, pid + len(k), dtype=int)\n", + " pid += len(k)\n", + " cellid = tuple(zip(k, j))\n", + " release_ptsi = list(zip(pids, cellid, x, y, z))\n", + " release_pts += release_ptsi\n", + "\n", + "print(f\"Number of release points: {len(release_pts)}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# create new subdir for PRT backwards model\n", + "prt_ws_bw = Path(dsr.attrs[\"model_ws\"]) / \"prt_bw\"\n", + "\n", + "# Create a new Simulation object\n", + "simprt_bw = nlmod.sim.sim(dsr, sim_ws=prt_ws_bw)\n", + "_ = nlmod.sim.tdis(dsr, simprt_bw)\n", + "\n", + "# Add PRT model\n", + "prt_bw = nlmod.prt.prt(dsr, simprt_bw, modelname=dsr.model_name)\n", + "\n", + "# Add DISV model\n", + "_ = nlmod.prt.disv(dsr, prt_bw)\n", + "\n", + "# MIP: Model Input Package\n", + "_ = nlmod.prt.mip(dsr, prt_bw, porosity=0.3)\n", + "\n", + "# PRP: particle release package\n", + "_ = nlmod.prt.prp(\n", + " dsr,\n", + " prt_bw,\n", + " packagedata=release_pts,\n", + " perioddata={0: [\"FIRST\"]},\n", + ")\n", + "\n", + "# OC: output control\n", + "trackcsvfile_prt_bw = f\"{prt_bw.name}bw.trk.csv\"\n", + "_ = nlmod.prt.oc(dsr, prt_bw, trackcsv_filerecord=trackcsvfile_prt_bw)\n", + "\n", + "# FMI: flow model interface\n", + "fmi_bw_packagedata = [\n", + " (\"GWFHEAD\", hds_bw_path),\n", + " (\"GWFBUDGET\", cbc_bw_path),\n", + "]\n", + "_ = nlmod.prt.fmi(dsr, prt_bw, packagedata=fmi_bw_packagedata)\n", + "\n", + "# EMS: explicit model solution\n", + "_ = nlmod.sim.ems(simprt_bw, model=prt_bw)\n", + "\n", + "# Write and run the simulation\n", + "simprt_bw.write_simulation()\n", + "success_bw_prt, _ = simprt_bw.run_simulation()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Load pathline data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pathlines_bw = nlmod.prt.read_pathlines(\n", + " prt_ws_bw / trackcsvfile_prt_bw, icell2d=dsr.icell2d.size\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Get pathlines of particles that started in a well (backward)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pathlines_bw0 = pathlines_bw.query(\"ireason==0\")\n", + "well_pids = []\n", + "for _, row in wel.stress_period_data.dataframe[0].iterrows():\n", + " cellid_layer = row[\"cellid_layer\"]\n", + " icell2d = row[\"cellid_cell\"]\n", + " well_pids.append(\n", + " pathlines_bw0.query(f\"ilay=={cellid_layer} and icell2d=={icell2d}\")\n", + " .loc[:, \"pid\"]\n", + " .to_numpy()[0]\n", + " )\n", + "pathlines_bww = pathlines_bw[pathlines_bw[\"pid\"].isin(well_pids)]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Plot results" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "f, ax = plt.subplots(figsize=(8, 8))\n", + "for pid, gr in pathlines_bww.groupby(\"pid\"):\n", + " for ilay, gr_lay in gr.groupby(\"ilay\"):\n", + " ax.plot(gr_lay[\"x\"], gr_lay[\"y\"], lw=1.0, color=f\"C{ilay}\", zorder=1)\n", + "\n", + "nlmod.plot.modelgrid(dsr, ax=ax, linewidth=0.1)\n", + "plot_river_and_wells(ax=ax)\n", + "ax.set_xlim(ds.extent[0], ds.extent[1])\n", + "ax.set_ylim(ds.extent[2], ds.extent[3])\n", + "ax.set_aspect(\"equal\")\n", + "ax.set_xlabel(\"X [m]\")\n", + "_ = ax.set_ylabel(\"Y [m]\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "f, axes = plt.subplots(3, 1, figsize=(8, 6), sharex=True, sharey=True)\n", + "\n", + "for pid, gr in pathlines_bw.groupby(\"pid\"):\n", + " t0 = gr[\"t\"].idxmin()\n", + " x0, y0, z0, ilay0 = gr.loc[t0, [\"x\", \"y\", \"z\", \"ilay\"]]\n", + " distance = np.sqrt((gr[\"x\"] - x0) ** 2 + (gr[\"y\"] - y0) ** 2 + (gr[\"z\"] - z0) ** 2)\n", + " if pid in well_pids:\n", + " plot_kwargs = {\"color\": \"k\", \"alpha\": 1.0, \"zorder\": 10, \"label\": f\"well {pid}\"}\n", + " else:\n", + " plot_kwargs = {\"color\": f\"C{ilay0}\", \"alpha\": 0.01}\n", + " axes[ilay0].plot(gr[\"t\"] / 365.25, distance, **plot_kwargs)\n", + "\n", + "for i, ax in enumerate(axes):\n", + " ax.xaxis.set_major_locator(mpl.ticker.MultipleLocator(250))\n", + " ax.yaxis.set_major_locator(mpl.ticker.MultipleLocator(100))\n", + " ax.set_ylabel(\"distance [m]\")\n", + " ax.set_title(f\"Layer {i}\")\n", + " ax.grid(True)\n", + " handles, labels = ax.get_legend_handles_labels()\n", + " if labels:\n", + " ax.legend(handles, labels, loc=\"lower right\", ncol=1, fontsize=8, frameon=True)\n", + " ax.set_xlim(0.0, 3000.0)\n", + " ax.set_ylim(0.0, 700.0)\n", + "_ = axes[-1].set_xlabel(\"time [year]\")" + ] } ], "metadata": { + "kernelspec": { + "display_name": "nlmod (3.13.5)", + "language": "python", + "name": "python3" + }, "language_info": { - "name": "python" + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.5" } }, "nbformat": 4, diff --git a/nlmod/prt/prt.py b/nlmod/prt/prt.py index df1ab987..d95fa878 100644 --- a/nlmod/prt/prt.py +++ b/nlmod/prt/prt.py @@ -1,7 +1,8 @@ import logging from pathlib import Path - +from pandas import read_csv, DataFrame import flopy as fp +from flopy.plot.plotutil import PRT_PATHLINE_DTYPE from ..gwf.gwf import _dis, _disv, _set_record from ..util import _get_value_from_ds_datavar @@ -90,7 +91,7 @@ def disv(ds, prt, length_units="METERS", pname="disv", **kwargs): def mip(ds, prt, porosity=None, **kwargs): - """Create matrix input package for groundwater transport model. + """Create model input package for particle tracking model. Parameters ---------- @@ -104,7 +105,7 @@ def mip(ds, prt, porosity=None, **kwargs): Returns ------- mip : flopy ModflowGwtmip - matrix input package + model input package """ logger.info("creating mf6 MIP") @@ -125,7 +126,7 @@ def mip(ds, prt, porosity=None, **kwargs): def prp(ds, prt, packagedata, perioddata, pname="prp", **kwargs): - """Create particle release package for groundwater transport model. + """Create particle release point package for particle tracking model. Parameters ---------- @@ -139,7 +140,7 @@ def prp(ds, prt, packagedata, perioddata, pname="prp", **kwargs): Returns ------- prp : flopy ModflowGwtprp - particle release package + particle release point package """ logger.info("creating mf6 PRP") prp_track_file = kwargs.pop("prp_track_file", f"{ds.model_name}.prp.trk") @@ -161,7 +162,7 @@ def prp(ds, prt, packagedata, perioddata, pname="prp", **kwargs): def fmi(ds, prt, packagedata=None, **kwargs): - """Create flow model interface package for groundwater transport model. + """Create flow model interface package for particle tracking model. Parameters ---------- @@ -215,3 +216,66 @@ def oc(ds, prt, save_budget=True, print_budget=False, **kwargs): **kwargs, ) return oc + + +def read_pathlines(path: str | Path, icell2d: int | None = None) -> DataFrame: + """Read PRT pathlines from (csv-)file and add particle ID. + + The columns in the pathlines file are: + - kper: stress period number + - kstp: time step number + - imdl: number of the model the particle originated in + - iprp: number of the particle release point (PRP) package the particle originated in + - irpt: release point number + - ilay: layer number + - icell: cell number + - izone: zone number + - istatus: particle status code + 0: particle was released + 1: particle is being actively tracked + 2: particle terminated at a boundary face + 3: particle terminated in a weak sink cell + 4: unused + 5: particle terminated in a cell with no exit face + 6: particle terminated in a stop zone + 7: particle terminated in an inactive cell + 8: particle terminated immediately upon release into a dry cell + 9: particle terminated in a subcell with no exit face + - ireason: reporting reason code (why the particle track record was saved) + 0: particle was released + 1: particle exited a cell + 2: time step ended + 3: particle terminated + 4: particle entered a weak sink cell + 5: user-specified tracking time + - trelease: particle release time + - t: particle tracking time + - x: particle x coordinate + - y: particle y coordinate + - z: particle z coordinate + - name: name of the particle release point + + Parameters + ---------- + path : Path or str + Path to the PRT pathlines file. + icell2d : int or None + Number of cells in the vertex grid. If provided, it will be used to + calculate the vertex cell index. E.g. `ds.icell2d.size` + + Returns + ------- + DataFrame + DataFrame containing the pathlines data. + """ + df = read_csv(path, dtype=PRT_PATHLINE_DTYPE) + df["ilay"] -= 1 + df["icell"] -= 1 + if icell2d is not None: + df["icell2d"] = df["icell"] - df["ilay"] * icell2d + + # identify particle ID and add as first column to the DataFrame + pid_cols = ["imdl", "iprp", "irpt", "trelease"] + pid = df.sort_values(pid_cols).groupby(pid_cols).ngroup() + df.insert(0, "pid", pid) + return df