{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Post-processing and standard plots\n",
    "\n",
    "This tutorial introduces the standardized post-processing interface. We run a small Vlasov–Ampère example, get its output as an autocomplete-friendly `Output`, and make the plots most commonly used to inspect a simulation.\n",
    "\n",
    "For a production run you can skip the simulation setup and open its output folder with `struphy.Output(\"path/to/sim\")` instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from IPython.display import HTML\n",
    "\n",
    "\n",
    "from struphy import (\n",
    "    BinningPlot,\n",
    "    BoundaryParameters,\n",
    "    ButcherTableau,\n",
    "    DerhamOptions,\n",
    "    EnvironmentOptions,\n",
    "    KernelDensityPlot,\n",
    "    LoadingParameters,\n",
    "    SavingParameters,\n",
    "    Simulation,\n",
    "    SortingParameters,\n",
    "    Time,\n",
    "    WeightsParameters,\n",
    "    domains,\n",
    "    equils,\n",
    "    grids,\n",
    "    maxwellians,\n",
    "    perturbations,\n",
    ")\n",
    "from struphy.models import Maxwell, ViscousEulerSPH, VlasovAmpereOneSpecies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Create a compact demonstration run\n",
    "\n",
    "Post-processing operates on a completed run. The small setup below saves an electric field, a few marker trajectories, scalar diagnostics, and a binned $(\\eta_1,v_1)$ distribution. These are the main output types handled by the plotting interface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model():\n",
    "    model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False)\n",
    "    model.em_fields.e_field.save_data = True\n",
    "    model.em_fields.phi.save_data = True\n",
    "    model.kinetic_ions.var.save_data = True\n",
    "\n",
    "    model.propagators.push_eta.options = model.propagators.push_eta.Options()\n",
    "    model.propagators.coupling_va.options = model.propagators.coupling_va.Options()\n",
    "    model.initial_poisson.options = model.initial_poisson.Options(stab_mat=\"M0\")\n",
    "\n",
    "    binplot = BinningPlot(\n",
    "        slice=\"e1_v1\",\n",
    "        n_bins=(32, 32),\n",
    "        ranges=((0.0, 1.0), (-5.0, 5.0)),\n",
    "    )\n",
    "    model.kinetic_ions.set_markers(\n",
    "        loading_params=LoadingParameters(ppc=32, seed=1234),\n",
    "        weights_params=WeightsParameters(control_variate=True),\n",
    "        boundary_params=BoundaryParameters(),\n",
    "        sorting_params=SortingParameters(boxes_per_dim=(4, 1, 1), do_sort=True),\n",
    "        saving_params=SavingParameters(n_markers=12, binning_plots=(binplot,)),\n",
    "    )\n",
    "\n",
    "    background = maxwellians.Maxwellian3D(n=(1.0, None))\n",
    "    model.kinetic_ions.var.add_background(background)\n",
    "    density_mode = perturbations.ModesCos(ls=(1,), amps=(1e-3,))\n",
    "    model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, density_mode)))\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "model = build_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "demo_tmp = tempfile.TemporaryDirectory(prefix=\"struphy_postprocessing_\")\n",
    "demo_root = demo_tmp.name\n",
    "\n",
    "env = EnvironmentOptions(\n",
    "    out_folders=demo_root,\n",
    "    sim_folder=\"vlasov_ampere_demo\",\n",
    "    save_restart=False,\n",
    ")\n",
    "sim = Simulation(\n",
    "    model=model,\n",
    "    env=env,\n",
    "    time_opts=Time(dt=0.05, Tend=5.0),\n",
    "    domain=domains.Cuboid(r1=2 * 3.141592653589793),\n",
    "    equil=equils.HomogenSlab(),\n",
    "    grid=grids.TensorProductGrid(num_elements=(16, 1, 1)),\n",
    "    derham_opts=DerhamOptions(degree=(2, 1, 1)),\n",
    ")\n",
    "out = sim.run(profiling_activated=True)\n",
    "print(f\"Raw output: {sim.env.path_out}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Process and load the output\n",
    "\n",
    "`sim.run()` returns an `Output`, which is also available later as `sim.output`. Scalars are read directly from the raw output; fields and particle products need post-processing, which runs with default options the first time they are accessed.\n",
    "\n",
    "To choose options, call `out.pproc()` first. `Output` evaluates saved FEEC fields and organizes particle diagnostics; `physical=True` additionally creates physical field components. Existing products made with the same options are reused, so re-running a cell is cheap.\n",
    "\n",
    "Individual products are standard `xarray.DataArray` objects with named dimensions, coordinates, units, and labels. Time is in Struphy units, in which the models' analytic results are written; seconds come along as the coordinate `t_seconds`, and `out.with_time_units(\"physical\")` gives an independent view with `t` in seconds. Arrays are loaded only when accessed. The saved configuration is available through `out.domain`, `out.model`, and `out.time_opts`; no simulation object is created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "out.pproc(physical=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "Products sit on the run under the species that produced them, so VS Code and interactive shells complete them as you type: `out.kinetic_ions.e1_v1_density.f`, `out.kinetic_ions.orbits`, `out.em_fields.phi_log`. To discover the products a particular namespace contains, inspect its flat, lazy `.catalog`: `list(out.kinetic_ions.catalog)`. The grouped views `out.fields`, `out.distributions`, `out.densities` and `out.orbits` show the same products by kind. In scripts, `out.evaluate(\"kinetic_ions/f\", dataset=\"e1_v1_density/f\")` looks up one explicitly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "out.info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(out.kinetic_ions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "### Reconstructed setup and initial conditions\n",
    "\n",
    "`out.info()` prints every evaluable key with a short description and the exact `out.evaluate(...)` call that loads it, followed by a few hints on selecting times and coordinates. The run configuration is not part of `info()`: in `out.metadata`, serialized initial conditions live on each variable under `model → species → variables → initial_conditions`. `out.initial_conditions` gives reconstructed objects, including saved Python functions and classes; `out.model` restores them onto the model variables. Unsupported definitions remain available in `out.metadata`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Model parameters:\", out.model.params)\n",
    "print(\"Kinetic variables:\", out.model.kinetic_ions.variables)\n",
    "print(\"Propagator options:\")\n",
    "for name, options in out.metadata[\"model\"][\"propagator_options\"].items():\n",
    "    print(f\"  {name}: {options}\")\n",
    "\n",
    "saved = out.metadata[\"model\"][\"species\"][\"kinetic_ions\"][\"variables\"][\"var\"][\"initial_conditions\"]\n",
    "print(\"Initial-condition entries:\", tuple(saved))\n",
    "initial = out.initial_conditions[\"kinetic_ions\"][\"var\"]\n",
    "print(\"Background distribution:\", initial[\"backgrounds\"])\n",
    "print(\"Initial distribution:\", initial[\"initial_condition\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "The reconstructed kinetic distributions are ordinary Struphy background objects. Here we evaluate the saved equilibrium and initial distribution along $\\eta_1$; their difference is the density perturbation that seeded the run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "background = initial[\"backgrounds\"]\n",
    "initial_distribution = initial[\"initial_condition\"]\n",
    "eta1 = np.linspace(0.0, 1.0, 256)\n",
    "zeros = np.zeros_like(eta1)\n",
    "n_background = np.asarray(background.n(eta1, zeros, zeros))\n",
    "n_initial = np.asarray(initial_distribution.n(eta1, zeros, zeros))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(eta1, n_background, label=\"background density\")\n",
    "ax.plot(eta1, n_initial, label=\"initial density\")\n",
    "ax.plot(eta1, n_initial - n_background, label=\"density perturbation\")\n",
    "ax.set(xlabel=r\"$\\eta_1$\", ylabel=\"density\", title=\"Saved kinetic initial condition\")\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "phase_space = out.kinetic_ions.e1_v1_density.f\n",
    "print(phase_space)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "### Evaluate saved splines directly: 1-D, 2-D, and 3-D\n",
    "\n",
    "`evaluate()` reads saved FEEC coefficients and evaluates the spline directly, without creating a post-processing field. With no `eta` arguments, it evaluates the full simulation grid at its cell centres. For cuts, pass a scalar for each direction to hold fixed; omitted directions stay on the cell centres of the simulation grid. Scalars, lists, NumPy arrays, and `range` objects can be mixed; every non-scalar input becomes a dimension of the tensor-product result. Thus one varying coordinate makes a 1-D line, two make a 2-D plane, and no coordinates makes a 3-D volume. The result is always an xarray array with logical `eta1`/`eta2`/`eta3` coordinates and mapped physical `X`/`Y`/`Z` coordinates.\n",
    "\n",
    "A representation conversion is applied after spline evaluation. Its input is inferred from the saved FEEC space, so `representation=` specifies only the target: `\"0\"`, `\"1\"`, `\"2\"`, `\"3\"`, `\"v\"`, or `\"norm\"`. Scalars default to `\"0\"`; vectors default to `\"norm\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1-D: a field line at eta2 = eta3 = 0.5.\n",
    "eta1_line = np.linspace(0.0, 1.0, 128)\n",
    "phi_line = out.evaluate(\n",
    "    \"em_fields/phi\",\n",
    "    eta1=eta1_line,\n",
    "    eta2=0.5,\n",
    "    eta3=0.5,\n",
    "    t=-1,\n",
    ")\n",
    "phi_line.plot()\n",
    "print(phi_line.dims, phi_line.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "For a 2-D plane, vary two coordinates and hold the third fixed. This is useful for a cross-section of a 3-D field even when the simulation was run on a coarser grid: the spline is evaluated at the requested points. For a 3-D volume, vary all three coordinates. Keep volume grids modest, then take a plane or line from the labeled result for plotting or further analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2-D: a logical eta1--eta2 plane at eta3 = 0.5.\n",
    "phi_plane = out.evaluate(\n",
    "    \"em_fields/phi\",\n",
    "    eta1=np.linspace(0.0, 1.0, 64),\n",
    "    eta2=np.linspace(0.0, 1.0, 48),\n",
    "    eta3=0.5,\n",
    "    t=-1,\n",
    ")\n",
    "phi_plane.plot(x=\"eta1\", y=\"eta2\")\n",
    "\n",
    "# 3-D: vary all three coordinates; keep volume grids modest.\n",
    "phi_volume = out.evaluate(\n",
    "    \"em_fields/phi\",\n",
    "    eta1=np.linspace(0.0, 1.0, 32),\n",
    "    eta2=np.linspace(0.0, 1.0, 16),\n",
    "    eta3=np.linspace(0.0, 1.0, 8),\n",
    "    t=-1,\n",
    ")\n",
    "print(phi_volume.dims, phi_volume.shape)\n",
    "\n",
    "# xarray plots a 2-D slice of the volume; choose the mid-plane by coordinate index.\n",
    "phi_volume.isel(eta3=phi_volume.sizes[\"eta3\"] // 2).plot(x=\"eta1\", y=\"eta2\")\n",
    "\n",
    "# Without eta arguments the spline is evaluated at the cell centres of the simulation grid,\n",
    "# here 16 x 1 x 1 cells, so this run's grid gives a line along eta1.\n",
    "phi_grid = out.evaluate(\"em_fields/phi\", t=-1)\n",
    "print(phi_grid.dims, phi_grid.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "### Inspecting `out` itself\n",
    "\n",
    "Most of what `out` exposes is generated on demand (`__getattr__`, cached properties), so `vars(out)` only shows a handful of private cache slots, not the products or methods. Use `dir(out)` for the flat list IPython's own tab-completion relies on, and `out.info()` for the full product tree. The same applies one level down: `out.model` is not a generic stub but the concrete model class of the run (`VlasovAmpereOneSpecies` here), so `dir(out.model)` and `print(out.model)` already show that model's own parameters directly — no separate `OutputVlasovAmpereOneSpecies`-style class is needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"vars(out):\", vars(out))  # only private cache slots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"dir(out):\", [name for name in dir(out) if not name.startswith(\"_\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(out.model)  # the concrete model class, with its own parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"dir(out.model):\", [name for name in dir(out.model) if not name.startswith(\"_\")][:15])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "### Products are xarray arrays\n",
    "\n",
    "A product is an `xarray.DataArray`, so xarray's own plotting already draws it, with the labels and units Struphy stored:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "phase_space.isel(t=-1).plot(x=\"eta1\", y=\"v1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26",
   "metadata": {},
   "source": [
    "Use xarray's plotting methods for ordinary one- and two-dimensional output. Select a saved time or spatial slice with `.isel()` or `.sel()` first, then call `.plot()` or `.plot.line()`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27",
   "metadata": {},
   "source": [
    "## Scalar overview and time series\n",
    "\n",
    "Field and particle products are xarray `DataArray` objects; `out.evaluate(\"scalars\")` returns an xarray `Dataset`. `out.evaluate(\"kinetic_ions/f\", dataset=\"e1_v1_density/f\")` looks a product up by name, which suits scripts and loops; attribute access is convenient interactively.\n",
    "\n",
    "Call `.plot()` for a scalar time series. Use a Matplotlib axes when combining several series or setting plot options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "out.scalars.electric_energy.plot.line(x=\"t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_fit = 2.0  # Struphy time units, like every time coordinate of this run\n",
    "energy = out.scalars.electric_energy.sel(t=slice(0.0, t_fit))\n",
    "fig, ax = plt.subplots()\n",
    "energy.plot.line(ax=ax, label=\"electric energy\")\n",
    "ax.set_yscale(\"log\")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30",
   "metadata": {},
   "source": [
    "## Two-dimensional data\n",
    "\n",
    "Choose the displayed dimensions with `x` and `y`, and pick one value for every other dimension by naming it: `t=\"last\"` (or `\"first\"`), `t=-1` for a position, and `t=0.35` for the nearest coordinate value. Arrays can also be sliced beforehand with xarray's `.isel()` and `.sel()`. `coords=\"physical\"` draws on the mapped coordinates instead of the logical ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "phase_space.isel(t=-1).plot(x=\"eta1\", y=\"v1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32",
   "metadata": {},
   "source": [
    "For a compact view of the evolution, select saved times and use xarray faceting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "phase_space.isel(t=np.linspace(0, phase_space.sizes[\"t\"] - 1, 5, dtype=int)).plot(\n",
    "    x=\"eta1\", y=\"v1\", col=\"t\", col_wrap=5\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34",
   "metadata": {},
   "source": [
    "## Selecting saved snapshots\n",
    "\n",
    "Use `.isel()` for index-based selection and `.sel()` for coordinate-based selection. This keeps selection explicit and works with every xarray operation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_phase_space = phase_space.isel(t=-1)\n",
    "final_phase_space.plot(x=\"eta1\", y=\"v1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36",
   "metadata": {},
   "source": [
    "Saved marker orbits sit under their species as an `xarray.Dataset` with one `(t, marker)` variable per quantity (`x`, `y`, `z`, velocities, `weight`); each variable's `description` attribute says what it is. Select one marker and plot its positions over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "metadata": {},
   "outputs": [],
   "source": [
    "orbit = out.kinetic_ions.orbits.isel(marker=0)\n",
    "orbit[[\"x\", \"y\", \"z\"]].to_dataarray(\"quantity\").plot.line(x=\"t\", hue=\"quantity\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38",
   "metadata": {},
   "source": [
    "A small collection of explicit snapshots is often more useful in a reproducible notebook than an interactive widget or animation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39",
   "metadata": {},
   "outputs": [],
   "source": [
    "phase_space.isel(t=[0, -1]).plot(x=\"eta1\", y=\"v1\", col=\"t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "phase_space.isel(t=-1).plot(ax=ax, x=\"eta1\", y=\"v1\")\n",
    "fig.savefig(os.path.join(demo_root, \"phase_space_final.png\"), bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41",
   "metadata": {},
   "source": [
    "The reconstructed equilibrium is available directly on the output handle for inspection and for model-specific analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(out.equil)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43",
   "metadata": {},
   "source": [
    "## Derived quantities\n",
    "\n",
    "xarray arithmetic computes derived quantities without special APIs. For example, subtract the first sample to obtain a drift and divide by it to obtain a relative error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_energy = out.scalars.total_energy\n",
    "energy_drift = total_energy - total_energy.isel(t=0)\n",
    "energy_error = abs(energy_drift) / abs(total_energy.isel(t=0))\n",
    "print(f\"largest drift of the total energy: {abs(energy_drift).max().item():.3e}\")\n",
    "\n",
    "energy_error.plot.line(x=\"t\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45",
   "metadata": {},
   "source": [
    "For custom diagnostics, first select the labeled subset needed for the calculation. Here a space-time field line is retained as an xarray object, ready for NumPy, SciPy, or another analysis package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46",
   "metadata": {},
   "outputs": [],
   "source": [
    "space_time_line = out.evaluate(\n",
    "    \"em_fields/phi\", eta1=np.linspace(0.0, 1.0, 64), eta2=0.5, eta3=0.5\n",
    ")\n",
    "print(space_time_line.dims, space_time_line.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47",
   "metadata": {},
   "source": [
    "## Reducing distribution functions\n",
    "\n",
    "A binned distribution usually has more dimensions than the question needs. Xarray reductions retain the remaining named dimensions, so averaging over `e1`, `e2` and `e3` turns the $(\\eta_1, v_1)$ product into $f(v_1, t)$. The mean is uniform in logical coordinates, which is a volume average on a Cartesian domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# average over whichever logical space directions the binning kept\n",
    "spatial = [dim for dim in (\"eta1\", \"eta2\", \"eta3\") if dim in phase_space.dims]\n",
    "f_of_v = phase_space.mean(spatial)\n",
    "print(f_of_v.dims)\n",
    "f_of_v.plot(x=\"t\", y=\"v1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49",
   "metadata": {},
   "source": [
    "Velocity moments are weighted xarray reductions. The bin widths and velocity coordinate remain labeled, making the density, mean velocity, and variance explicit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50",
   "metadata": {},
   "outputs": [],
   "source": [
    "dv1 = phase_space.v1.differentiate(\"v1\")\n",
    "density = (phase_space * dv1).sum(\"v1\")\n",
    "mean_v1 = (phase_space * phase_space.v1 * dv1).sum(\"v1\") / density\n",
    "variance_v1 = (phase_space * (phase_space.v1 - mean_v1) ** 2 * dv1).sum(\"v1\") / density\n",
    "mean_density = density.mean(spatial)\n",
    "mean_density.plot.line(x=\"t\")\n",
    "variance_v1.mean(spatial).plot.line(x=\"t\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51",
   "metadata": {},
   "source": [
    "## Physical units\n",
    "\n",
    "Products are in the normalization of the model. `out.units` holds the units of that normalization and `out.to_si(product)` converts a product: the time coordinate to seconds, the mapped coordinates `X`, `Y`, `Z` to meters and the velocities `v1`, `v2`, `v3` to m/s. The values are converted only when `unit=` names the unit the variable was normalized with (`\"x\"`, `\"B\"`, `\"n\"`, `\"v\"`, `\"t\"`, `\"p\"`, `\"rho\"`, `\"j\"` or `\"kBT\"`), or a number for a composite unit together with its `label`, because a product does not record which unit its variable uses. The original product is not modified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"1 length unit = {out.units.x} m, 1 velocity unit = {out.units.v:.4g} m/s, 1 time unit = {out.units.t:.4g} s\")\n",
    "\n",
    "phase_space_si = out.to_si(phase_space)\n",
    "print(phase_space_si.v1.attrs[\"units\"], phase_space_si.t.attrs[\"units\"])\n",
    "phase_space_si.isel(t=-1).plot(x=\"eta1\", y=\"v1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53",
   "metadata": {},
   "source": [
    "## Save standard output\n",
    "\n",
    "Products are xarray objects, so figures are saved with Matplotlib's `plt.savefig(path)` and arrays with xarray's own writers, for example `phi.to_netcdf(path)`. For a compact record of the run, `out.report()` writes the full scalar history as `scalars.csv` and a report listing the run metadata, every available product, and the dimensions and units of any products passed with `products=`. By default the files go to `post_processing/report/`; `format` is `\"markdown\"` or `\"html\"`, and the return value is the path of the report file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54",
   "metadata": {},
   "outputs": [],
   "source": [
    "report = out.report(products=[\"em_fields/phi\"], format=\"html\")\n",
    "print(\"Wrote:\")\n",
    "for path in (report, os.path.join(os.path.dirname(report), \"scalars.csv\")):\n",
    "    print(\" \", os.path.relpath(path, out.path_out))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55",
   "metadata": {},
   "source": [
    "## Comparing runs\n",
    "\n",
    "Time series accept arrays of other simulations, so comparing runs needs nothing special. Series are labelled by the run they come from, and the runs may have different time grids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_coarse = Simulation(\n",
    "    model=build_model(),\n",
    "    env=EnvironmentOptions(out_folders=demo_root, sim_folder=\"vlasov_ampere_coarse\", save_restart=False),\n",
    "    time_opts=Time(dt=0.1, Tend=5.0),\n",
    "    domain=domains.Cuboid(r1=2 * 3.141592653589793),\n",
    "    equil=equils.HomogenSlab(),\n",
    "    grid=grids.TensorProductGrid(num_elements=(16, 1, 1)),\n",
    "    derham_opts=DerhamOptions(degree=(2, 1, 1)),\n",
    ")\n",
    "out_coarse = sim_coarse.run(profiling_activated=True)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "out.scalars.electric_energy.plot.line(ax=ax, label=\"dt = 0.05\")\n",
    "out_coarse.scalars.electric_energy.plot.line(ax=ax, label=\"dt = 0.1\")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57",
   "metadata": {},
   "source": [
    "## Profiling\n",
    "\n",
    "A run started with `sim.run(profiling_activated=True)` records how long each region of the code takes, and `out.profile` reads that record back. Regions are the setup steps, every propagator (`prop: ...`), pusher, accumulation, compiled kernel (`kernel: ...`) and linear solve. Regions nest, so a region's time includes the regions it calls and the times of different regions must not be added up. `summary()` returns a dataset along the dimension `region`, sorted by total time; `table()` prints it. Filter with `prefix` and limit with `top`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(out.profile.table(top=8))\n",
    "\n",
    "kernels = out.profile.summary(prefix=\"kernel:\")\n",
    "print(kernels.total_time.to_series())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59",
   "metadata": {},
   "source": [
    "`compare()` puts the same statistic of several runs side by side, with runs whose region is missing as NaN. Here the two runs of the previous section differ only in the time step, so the number of calls per propagator halves for `dt = 0.1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60",
   "metadata": {},
   "outputs": [],
   "source": [
    "calls = out.profile.compare(out_coarse, metric=\"calls\", prefix=\"prop:\")\n",
    "print(calls)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61",
   "metadata": {},
   "source": [
    "## Other models\n",
    "\n",
    "The interface is the same for every model; only the products differ. Two more short runs show the two product types the Vlasov–Ampère demo does not have: SPH densities, and vector fields on a mapped domain."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62",
   "metadata": {},
   "source": [
    "### SPH densities\n",
    "\n",
    "A standing sound wave discretized with SPH markers. `KernelDensityPlot` reconstructs the density on a grid, which appears under `out.densities`, while `BinningPlot` produces the binned quantities under `out.distributions`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63",
   "metadata": {},
   "outputs": [],
   "source": [
    "sph_model = ViscousEulerSPH(with_B0=False, with_viscosity=False)\n",
    "sph_model.propagators.push_eta.options = sph_model.propagators.push_eta.Options(\n",
    "    butcher=ButcherTableau(algo=\"forward_euler\"),\n",
    ")\n",
    "sph_model.propagators.push_sph_p.options = sph_model.propagators.push_sph_p.Options(kernel_type=\"gaussian_1d\")\n",
    "sph_model.euler_fluid.set_markers(\n",
    "    loading_params=LoadingParameters(ppb=8, loading=\"tesselation\"),\n",
    "    weights_params=WeightsParameters(),\n",
    "    boundary_params=BoundaryParameters(),\n",
    "    sorting_params=SortingParameters(boxes_per_dim=(12, 1, 1), dims_mask=(True, False, False)),\n",
    "    saving_params=SavingParameters(\n",
    "        binning_plots=(BinningPlot(slice=\"e1\", n_bins=(32,), ranges=(0.0, 1.0)),),\n",
    "        kernel_density_plots=(KernelDensityPlot(pts_e1=41, pts_e2=1),),\n",
    "    ),\n",
    ")\n",
    "sph_model.euler_fluid.var.add_background(equils.ConstantVelocity())\n",
    "sph_model.euler_fluid.var.add_perturbation(del_n=perturbations.ModesSin(ls=(1,), amps=(1.0e-2,)))\n",
    "\n",
    "sph = Simulation(\n",
    "    model=sph_model,\n",
    "    env=EnvironmentOptions(out_folders=demo_root, sim_folder=\"sph_soundwave\", save_restart=False),\n",
    "    time_opts=Time(dt=0.03125, Tend=2.5, split_algo=\"Strang\"),\n",
    "    domain=domains.Cuboid(r1=2.5),\n",
    "    grid=None,\n",
    "    derham_opts=None,\n",
    ")\n",
    "out_sph = sph.run()\n",
    "\n",
    "print(\"densities:\", tuple(out_sph.density_catalog))\n",
    "print(\"binned:\", tuple(out_sph.distribution_catalog))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64",
   "metadata": {},
   "source": [
    "For a one-dimensional run, the clearest picture is a space-time map: the sweep dimension `t` may be used as a display axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65",
   "metadata": {},
   "outputs": [],
   "source": [
    "density = out_sph.euler_fluid.view_0.n.isel(eta2=0, eta3=0)\n",
    "density.plot(x=\"t\", y=\"eta1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66",
   "metadata": {},
   "source": [
    "Products are plain `xarray.DataArray` objects, so anything xarray can do works directly, for example profiles at selected times:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67",
   "metadata": {},
   "outputs": [],
   "source": [
    "density.isel(t=[0, len(density.t) // 4, len(density.t) // 2]).plot.line(x=\"eta1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68",
   "metadata": {},
   "source": [
    "### Vector fields on a mapped domain\n",
    "\n",
    "A coaxial waveguide mode of the Maxwell model, on an annulus. With `physical=True` the post-processing also computes the Cartesian field components (`*_xyz`), and `coords=\"physical\"` draws them on the mapped grid, with the plane chosen by `plane`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69",
   "metadata": {},
   "outputs": [],
   "source": [
    "a1, a2 = 2.326744, 3.686839\n",
    "\n",
    "maxwell_model = Maxwell()\n",
    "maxwell_model.propagators.maxwell.options = maxwell_model.propagators.maxwell.Options(algo=\"implicit\")\n",
    "maxwell_model.em_fields.e_field.add_perturbation(perturbations.CoaxialWaveguideElectric_r(m=3, a1=a1, a2=a2))\n",
    "maxwell_model.em_fields.e_field.add_perturbation(perturbations.CoaxialWaveguideElectric_theta(m=3, a1=a1, a2=a2))\n",
    "maxwell_model.em_fields.b_field.add_perturbation(perturbations.CoaxialWaveguideMagnetic(m=3, a1=a1, a2=a2))\n",
    "\n",
    "coaxial = Simulation(\n",
    "    model=maxwell_model,\n",
    "    env=EnvironmentOptions(out_folders=demo_root, sim_folder=\"coaxial\", save_restart=False),\n",
    "    time_opts=Time(dt=0.05, Tend=2.0),\n",
    "    domain=domains.HollowCylinder(a1=a1, a2=a2, Lz=2.0),\n",
    "    equil=equils.HomogenSlab(),\n",
    "    grid=grids.TensorProductGrid(num_elements=(24, 48, 1)),\n",
    "    derham_opts=DerhamOptions(degree=(2, 2, 1), bcs=((\"dirichlet\", \"dirichlet\"), None, None)),\n",
    ")\n",
    "out_coaxial = coaxial.run()\n",
    "out_coaxial.pproc(physical=True)\n",
    "\n",
    "print(\"fields:\", tuple(out_coaxial.field_catalog))\n",
    "print(\"dimensions:\", out_coaxial.em_fields.b_field_xyz.dims)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_coaxial.em_fields.b_field_xyz.isel(t=-1, component=2, eta3=0).plot(x=\"eta1\", y=\"eta2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_coaxial.em_fields.b_field_xyz.isel(component=2, eta3=0).plot(x=\"eta1\", y=\"eta2\", col=\"t\", col_wrap=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72",
   "metadata": {},
   "source": [
    "### Representation conversion on a torus\n",
    "\n",
    "A toroidal map makes the distinction between FEEC representations visible. The saved electric field is an H(curl) 1-form, so its source representation is inferred as `1`. We evaluate one poloidal line directly from saved spline coefficients, then request its native 1-form, normalized-vector, and Cartesian-vector representations. These differ away from a Cartesian map because the metric factors vary around the torus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73",
   "metadata": {},
   "outputs": [],
   "source": [
    "torus_model = Maxwell()\n",
    "torus_model.em_fields.e_field.save_data = True\n",
    "torus_model.em_fields.e_field.add_perturbation(\n",
    "    perturbations.ModesCos(ms=(1,), amps=(0.1,), given_in_basis=\"1\", comp=1)\n",
    ")\n",
    "\n",
    "torus = Simulation(\n",
    "    model=torus_model,\n",
    "    env=EnvironmentOptions(out_folders=demo_root, sim_folder=\"representation_torus\", save_restart=False),\n",
    "    time_opts=Time(dt=0.05, Tend=0.05),\n",
    "    domain=domains.HollowTorus(a1=0.2, a2=0.4, R0=1.0, tor_period=1),\n",
    "    equil=equils.HomogenSlab(),\n",
    "    grid=grids.TensorProductGrid(num_elements=(6, 12, 2)),\n",
    "    derham_opts=DerhamOptions(degree=(2, 2, 2), bcs=((\"dirichlet\", \"dirichlet\"), None, None)),\n",
    ")\n",
    "out_torus = torus.run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74",
   "metadata": {},
   "outputs": [],
   "source": [
    "eta2_line = np.linspace(0.0, 1.0, 256)\n",
    "component=0\n",
    "common = dict(eta1=0.7, eta2=eta2_line, eta3=0.0, t=-1, component=component)\n",
    "e_1 = out_torus.evaluate(\"em_fields/e_field\", representation=\"1\", **common)\n",
    "e_norm = out_torus.evaluate(\"em_fields/e_field\", representation=\"norm\", **common)\n",
    "e_v = out_torus.evaluate(\"em_fields/e_field\", representation=\"v\", **common)\n",
    "\n",
    "fig, ax = plt.subplots(ncols=3, figsize=(15, 5))\n",
    "i = 0\n",
    "for field, label in ((e_1, \"1-form\"), (e_norm, \"normalized vector\"), (e_v, \"Vector field\")):\n",
    "    print(label, field.isel(t=0))\n",
    "    ax[i].plot(eta2_line, field.isel(t=0), label=label)\n",
    "    \n",
    "    ax[i].set(xlabel=r\"$\\eta_2$\", ylabel=label, title=\"H(curl) field representations on a torus\")\n",
    "    i += 1\n",
    "# ax.legend()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75",
   "metadata": {},
   "source": [
    "## Apply the workflow to another run\n",
    "\n",
    "For an already completed simulation, possibly in a separate process without MPI, open its output folder:\n",
    "\n",
    "```python\n",
    "import struphy\n",
    "\n",
    "out = struphy.Output(\"/path/to/sim_1\").pproc(physical=True)\n",
    "out.domain, out.model.units  # reconstructed directly from saved metadata\n",
    "```\n",
    "\n",
    "Use `out.scalars`, `out.fields`, `out.distributions`, `out.orbits`, and `out.densities`. Attribute access is the normal interactive API; the corresponding `*_catalog` mappings are intended for generic loops and tooling."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.12.3)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}