{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Gas Expansion with `ViscousEulerSPH`\n",
    "\n",
    "This tutorial solves Euler-type fluid dynamics with SPH using the `ViscousEulerSPH` model. The continuum equations are\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    " \\partial_t \\rho + \\nabla \\cdot (\\rho \\mathbf u) &= 0\\,,\n",
    " \\\\[2mm]\n",
    " \\rho(\\partial_t \\mathbf u + \\mathbf u \\cdot \\nabla \\mathbf u) &= - \\nabla \\left(\\rho^2 \\frac{\\partial \\mathcal U(\\rho, S)}{\\partial \\rho} \\right)\\,,\n",
    " \\\\[2mm]\n",
    " \\partial_t S + \\mathbf u \\cdot \\nabla S &= 0\\,.\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Here, $S$ is entropy per unit mass and the internal energy model is\n",
    "\n",
    "$$\n",
    "\\mathcal U(\\rho, S) = \\kappa(S)\\log\\rho\\,.\n",
    "$$\n",
    "\n",
    "After SPH discretization, particle trajectories satisfy\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    " \\dot{\\mathbf x}_p &= \\mathbf v_p\\,,\\qquad && \\mathbf x_p(0) = \\mathbf x_{p0}\\,,\n",
    " \\\\[2mm]\n",
    " \\dot{\\mathbf v}_p &= -\\kappa_p(0) \\sum_{i=1}^N w_i \\left(\\frac{1}{\\rho^{N,h}(\\mathbf x_p)} + \\frac{1}{\\rho^{N,h}(\\mathbf x_i)} \\right) \\nabla W_h(\\mathbf x_p-\\mathbf x_i)\\,,\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "with smoothed density\n",
    "\n",
    "$$\n",
    "\\rho^{N,h}(\\mathbf x)=\\sum_{j=1}^N w_j W_h(\\mathbf x-\\mathbf x_j)\\,.\n",
    "$$\n",
    "\n",
    "In Struphy, the pressure-driven velocity update is handled by `PushVinSPHpressure`.\n",
    "\n",
    "Before launching the full gas-expansion case, let us quickly inspect a few available smoothing kernels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from struphy.pic.sph_smoothing_kernels import gaussian_uni, linear_uni, trigonometric_uni\n",
    "\n",
    "x = np.linspace(-1, 1, 200)\n",
    "out1 = np.zeros_like(x)\n",
    "out2 = np.zeros_like(x)\n",
    "out3 = np.zeros_like(x)\n",
    "\n",
    "for i, xi in enumerate(x):\n",
    "    out1[i] = trigonometric_uni(xi, 1.0)\n",
    "    out2[i] = gaussian_uni(xi, 1.0)\n",
    "    out3[i] = linear_uni(xi, 1.0)\n",
    "plt.plot(x, out1, label=\"trigonometric\")\n",
    "plt.plot(x, out2, label=\"gaussian\")\n",
    "plt.plot(x, out3, label=\"linear\")\n",
    "plt.title(\"Some smoothing kernels\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "### Step 1: Import APIs and the `ViscousEulerSPH` Model\n",
    "\n",
    "Load all Struphy components and plotting utilities needed for this standalone gas-expansion tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from struphy import DerhamOptions, EnvironmentOptions, Time\n",
    "from struphy import domains\n",
    "from struphy import equils\n",
    "from struphy import grids\n",
    "from struphy import (\n",
    "    BinningPlot,\n",
    "    BoundaryParameters,\n",
    "    KernelDensityPlot,\n",
    "    LoadingParameters,\n",
    "    SavingParameters,\n",
    "    SortingParameters,\n",
    "    WeightsParameters,\n",
    ")\n",
    "from struphy import Simulation\n",
    "\n",
    "# import model\n",
    "from struphy.models import ViscousEulerSPH"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "Let us verfiy the model equations by calling the `pde` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ViscousEulerSPH.pde()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "### Step 2: Define Environment and Domain\n",
    "\n",
    "Set time step, final time, and simulation domain. These settings are tuned for a 2D expansion embedded in a cuboid geometry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# environment options\n",
    "env = EnvironmentOptions()\n",
    "\n",
    "# time stepping\n",
    "time_opts = Time(dt=0.04, Tend=1.6, split_algo=\"Strang\")\n",
    "\n",
    "# geometry\n",
    "l1 = -3.0\n",
    "r1 = 3.0\n",
    "l2 = -3.0\n",
    "r2 = 3.0\n",
    "l3 = 0.0\n",
    "r3 = 1.0\n",
    "domain = domains.Cuboid(l1=l1, r1=r1, l2=l2, r2=r2, l3=l3, r3=r3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "### Step 3: Build the Initial Density Blob\n",
    "\n",
    "Define a Gaussian density profile in the $(x,y)$ plane and wrap it as a fluid equilibrium object. This serves as the initial condition for the expansion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# gaussian initial blob\n",
    "import numpy as np\n",
    "\n",
    "T_h = 0.2\n",
    "gamma = 5 / 3\n",
    "n_fun = lambda x, y, z: np.exp(-(x**2 + y**2) / T_h) / 35\n",
    "\n",
    "blob = equils.GenericCartesianFluidEquilibrium(n_xyz=n_fun)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "### Step 4: Set Grid and de Rham Discretization\n",
    "\n",
    "Choose the spline grid and de Rham options used by the field projections in this SPH run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fluid equilibrium (can be used as part of initial conditions)\n",
    "equil = None\n",
    "\n",
    "# grid\n",
    "grid = grids.TensorProductGrid(num_elements=(64, 64, 1))\n",
    "\n",
    "# derham options\n",
    "bcs = ((\"free\", \"free\"), (\"free\", \"free\"), None)\n",
    "derham_opts = DerhamOptions(degree=(3, 3, 1), bcs=bcs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "### Step 5: Instantiate `ViscousEulerSPH` and Simulation\n",
    "\n",
    "Create a lightweight model with magnetic-background and viscosity terms disabled (`with_B0=False`, `with_viscosity=False`), then build the simulation object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# light-weight model instance\n",
    "model = ViscousEulerSPH(with_B0=False, with_viscosity=False)\n",
    "\n",
    "sim = Simulation(model,\n",
    "                 env=env,\n",
    "                 time_opts=time_opts,\n",
    "                 domain=domain,\n",
    "                 equil=equil,\n",
    "                 grid=grid,\n",
    "                 derham_opts=derham_opts,)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Step 6: Configure Marker Sampling\n",
    "\n",
    "Set dense particle-per-box loading and enable `reject_weights` to discard very small particle weights. This reduces cost while preserving the dominant density structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "loading_params = LoadingParameters(ppb=400)\n",
    "weights_params = WeightsParameters(reject_weights=True, threshold=3e-8)\n",
    "boundary_params = BoundaryParameters()\n",
    "nx = 16\n",
    "ny = 16\n",
    "sorting_params = SortingParameters(boxes_per_dim=(nx, ny, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "### Step 7: Add Diagnostics for Visualization\n",
    "\n",
    "Save both a binning plot and an SPH kernel-density plot so you can compare coarse gridded density against kernel-based reconstruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "bin_plot = BinningPlot(\n",
    "    slice=\"e1_e2\",\n",
    "    n_bins=(64, 64),\n",
    "    ranges=((0.0, 1.0), (0.0, 1.0)),\n",
    "    divide_by_jac=False,\n",
    ")\n",
    "pts_e1 = 100\n",
    "pts_e2 = 90\n",
    "kd_plot = KernelDensityPlot(pts_e1=pts_e1, pts_e2=pts_e2, pts_e3=1)\n",
    "\n",
    "saving_params = SavingParameters(\n",
    "    n_markers=1.0,\n",
    "    binning_plots=(bin_plot,),\n",
    "    kernel_density_plots=(kd_plot,),\n",
    ")\n",
    "\n",
    "model.euler_fluid.set_markers(\n",
    "    loading_params=loading_params,\n",
    "    weights_params=weights_params,\n",
    "    boundary_params=boundary_params,\n",
    "    sorting_params=sorting_params,\n",
    "    saving_params=saving_params,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "### Step 8: Select the SPH Kernel for Pressure Updates\n",
    "\n",
    "For the pressure propagator, choose `gaussian_2d` as smoothing kernel. You can swap this setting later to study kernel sensitivity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# propagator options\n",
    "from struphy import ButcherTableau\n",
    "\n",
    "butcher = ButcherTableau(algo=\"forward_euler\")\n",
    "model.propagators.push_eta.options = model.propagators.push_eta.Options(butcher=butcher)\n",
    "\n",
    "model.propagators.push_sph_p.options = model.propagators.push_sph_p.Options(kernel_type=\"gaussian_2d\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "### Step 9: Initialize, Run, and Load Results\n",
    "\n",
    "Attach the Gaussian background and execute the standard workflow: `run()`, `pproc()`, `load_plotting_data()`.\n",
    "\n",
    "Performance note: early time steps are typically slower because particles are highly concentrated; runtime usually improves as the cloud expands. Running from a console script (especially with MPI/GPU support) is faster than notebook execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# background, perturbations and initial conditions\n",
    "model.euler_fluid.var.add_background(blob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.pproc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.load_plotting_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25",
   "metadata": {},
   "source": [
    "### Step 10: Inspect Stored Outputs\n",
    "\n",
    "After post-processing, use the printed paths to locate simulation artifacts. The following cells compare analytical initialization, marker sampling, SPH reconstruction, and binned density diagnostics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# analytical functions\n",
    "n_xyz = blob.n_xyz\n",
    "n3 = blob.n3\n",
    "\n",
    "# grids\n",
    "x = np.linspace(l1, r1, pts_e1)\n",
    "y = np.linspace(l2, r2, pts_e2)\n",
    "xx, yy = np.meshgrid(x, y, indexing=\"ij\")\n",
    "ee1, ee2, ee3 = sim.n_sph.euler_fluid.view_0.grid_n_sph\n",
    "eta1 = ee1[:, 0, 0]\n",
    "eta2 = ee2[0, :, 0]\n",
    "bc_x = sim.f.euler_fluid.e1_e2_density.grid_e1\n",
    "bc_y = sim.f.euler_fluid.e1_e2_density.grid_e2\n",
    "\n",
    "# markers\n",
    "orbits = sim.orbits.euler_fluid\n",
    "positions = orbits[0, :, :3]\n",
    "weights = orbits[0, :, 6]\n",
    "\n",
    "# binning and sph eval\n",
    "n_sph = sim.n_sph.euler_fluid.view_0.n_sph[0]\n",
    "f_bin = sim.f.euler_fluid.e1_e2_density.f_binned[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(12, 15))\n",
    "\n",
    "# plots\n",
    "plt.subplot(3, 2, 1)\n",
    "plt.pcolor(xx, yy, n_fun(xx, yy, 0))\n",
    "plt.axis(\"square\")\n",
    "plt.title(\"n_xyz initial\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(3, 2, 2)\n",
    "plt.pcolor(eta1, eta2, n3(eta1, eta2, 0, squeeze_out=True).T)\n",
    "plt.axis(\"square\")\n",
    "plt.title(\"$\\hat{n}^{\\t{vol}}$ initial (volume form)\")\n",
    "plt.colorbar()\n",
    "\n",
    "make_scatter = True\n",
    "if make_scatter:\n",
    "    plt.subplot(3, 2, 3)\n",
    "    ax = plt.gca()\n",
    "    ax.set_xticks(np.linspace(l1, r1, nx + 1))\n",
    "    ax.set_yticks(np.linspace(l2, r2, ny + 1))\n",
    "    plt.tick_params(labelbottom=False)\n",
    "    coloring = weights\n",
    "    plt.scatter(positions[:, 0], positions[:, 1], c=coloring, s=0.25)\n",
    "    plt.grid(c=\"k\")\n",
    "    plt.axis(\"square\")\n",
    "    plt.title(\"$\\hat{n}^{\\t{vol}}$ initial scatter (random)\")\n",
    "    plt.xlim(l1, r1)\n",
    "    plt.ylim(l2, r2)\n",
    "    plt.colorbar()\n",
    "\n",
    "plt.subplot(3, 2, 4)\n",
    "ax = plt.gca()\n",
    "ax.set_xticks(np.linspace(0, 1, nx + 1))\n",
    "ax.set_yticks(np.linspace(0, 1.0, ny + 1))\n",
    "plt.tick_params(labelbottom=False)\n",
    "plt.pcolor(ee1[:, :, 0], ee2[:, :, 0], n_sph[:, :, 0])\n",
    "plt.grid()\n",
    "plt.axis(\"square\")\n",
    "plt.title(\"n_sph initial (random)\")\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(3, 2, 5)\n",
    "ax = plt.gca()\n",
    "plt.pcolor(bc_x, bc_y, f_bin)\n",
    "plt.axis(\"square\")\n",
    "plt.title(\"n_binned initial (random)\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "dt = time_opts.dt\n",
    "Nt = sim.t_grid.size - 1\n",
    "\n",
    "positions = orbits[:, :, :3]\n",
    "\n",
    "interval = Nt / 10\n",
    "plot_ct = 0\n",
    "\n",
    "plt.figure(figsize=(12, 24))\n",
    "for i in range(Nt):\n",
    "    if i % interval == 0:\n",
    "        print(f\"{i=}\")\n",
    "        plot_ct += 1\n",
    "        plt.subplot(4, 2, plot_ct)\n",
    "        ax = plt.gca()\n",
    "        coloring = weights\n",
    "        plt.scatter(positions[i, :, 0], positions[i, :, 1], c=coloring, s=0.25)\n",
    "        plt.axis(\"square\")\n",
    "        plt.title(\"n0_scatter\")\n",
    "        plt.xlim(l1, r1)\n",
    "        plt.ylim(l2, r2)\n",
    "        plt.colorbar()\n",
    "        plt.title(f\"Gas at t={i * dt}\")\n",
    "    if plot_ct == 8:\n",
    "        break"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env (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
}