{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# 4 - Smoothed-particle hydrodynamics\n",
    "\n",
    "Struphy provides several models for smoothed-particle hydrodynamics (SPH). In this tutorial, we shall explore\n",
    "\n",
    "1. Pressure-less fluid flow in a Beltrami force field (model `PressureLessSPH`)\n",
    "2. Gas expansion with SPH (model `ViscousEulerSPH`)\n",
    "\n",
    "The parameter files discussed below can be obtained with the console commands\n",
    "\n",
    "```\n",
    "struphy params PressureLessSPH\n",
    "struphy params ViscousEulerSPH \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Pressure-less fluid flow in  Beltrami force field\n",
    "\n",
    "Let $\\mathbf u_0(\\mathbf x)$ stand for an initial flow velocity field. Moreover,\n",
    "let $\\Omega \\subset \\mathbb R^3$ be a box (cuboid). We search for trajectories $(\\mathbf x_p, \\mathbf v_p): [0,T] \\to \\Omega \\times \\mathbb R^3$, $p = 0, \\ldots, N-1$ that 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 &= -\\nabla p(\\mathbf x_p) \\qquad && \\mathbf v_p(0) = \\mathbf u_0(\\mathbf x_p(0))\\,,\n",
    " \\end{aligned}\n",
    "$$\n",
    "\n",
    "where $p \\in H^1(\\Omega)$ is some given function. In what follows we shall set $p$ to give a Beltrami force field to guide the fluid particles.\n",
    "\n",
    "We start with the API imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from struphy import EnvironmentOptions, Time\n",
    "from struphy import domains\n",
    "from struphy import equils\n",
    "from struphy import grids\n",
    "from struphy import DerhamOptions\n",
    "from struphy import (\n",
    "    BinningPlot,\n",
    "    BoundaryParameters,\n",
    "    KernelDensityPlot,\n",
    "    LoadingParameters,\n",
    "    WeightsParameters,\n",
    ")\n",
    "from struphy import Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "We shall use the Strang splitting algorithm for the propagators, and simulate in Cartesian geometry (Cuboid):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# environment options\n",
    "env = EnvironmentOptions()\n",
    "\n",
    "# time stepping\n",
    "time_opts = Time(dt=0.02, Tend=4, split_algo=\"Strang\")\n",
    "\n",
    "# geometry\n",
    "l1 = -0.5\n",
    "r1 = 0.5\n",
    "l2 = -0.5\n",
    "r2 = 0.5\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": "5",
   "metadata": {},
   "source": [
    "The Beltrami flow can be specified as a lambda function and then passed to `GenericCartesianFluidEquilibrium`. This fluid equilibrium will be set as initial condition below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct Beltrami flow\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def u_fun(x, y, z):\n",
    "    ux = -np.cos(np.pi * x) * np.sin(np.pi * y)\n",
    "    uy = np.sin(np.pi * x) * np.cos(np.pi * y)\n",
    "    uz = 0 * x\n",
    "    return ux, uy, uz\n",
    "\n",
    "\n",
    "p_fun = lambda x, y, z: 0.5 * (np.sin(np.pi * x) ** 2 + np.sin(np.pi * y) ** 2)\n",
    "n_fun = lambda x, y, z: 1.0 + 0 * x\n",
    "\n",
    "# put the functions in a generic equilibrium container\n",
    "bel_flow = equils.GenericCartesianFluidEquilibrium(u_xyz=u_fun, p_xyz=p_fun, n_xyz=n_fun)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "We will also need a grid and Derham complex for projecting the fluid equilibrium onto a spline basis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "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": "9",
   "metadata": {},
   "source": [
    "Let us now create the instance of the simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import model\n",
    "from struphy.models import PressureLessSPH\n",
    "\n",
    "# light-weight model instance\n",
    "model = PressureLessSPH(epsilon=1.0)\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": "11",
   "metadata": {},
   "source": [
    "In the next step, the light-weight instance of the model is created. We set the parameter `epsilon` to 1.0, appearing in the propagator `PushVinEfield`. Moreover, we launch with 1000 particles and save 100 % percent of them through `n_markers=1.0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "loading_params = LoadingParameters(Np=1000)\n",
    "weights_params = WeightsParameters()\n",
    "boundary_params = BoundaryParameters(bc=(\"reflect\", \"reflect\", \"periodic\"))\n",
    "model.cold_fluid.set_markers(\n",
    "    loading_params=loading_params,\n",
    "    weights_params=weights_params,\n",
    "    boundary_params=boundary_params,\n",
    ")\n",
    "model.cold_fluid.set_sorting_boxes(boxes_per_dim=(1, 1, 1))\n",
    "model.cold_fluid.set_save_data(n_markers=1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "We now set the propagator options. Here, it is important to pass the pressure from the Beltrami flow as auxiliary field to `push_v`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "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",
    "phi = bel_flow.p0\n",
    "model.propagators.push_v.options = model.propagators.push_v.Options(phi=phi)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "The initial condition of the species is defined by a background function, plus possible perturbations, which we ignore here. The background of the species is set to the Beltrami flow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# background, perturbations and initial conditions\n",
    "model.cold_fluid.var.add_background(bel_flow)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "Let us start the run, post-process the raw data, load the post-processed data and plot the particle trajectories:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.pproc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.load_plotting_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(12, 28))\n",
    "\n",
    "orbits = sim.orbits.cold_fluid\n",
    "\n",
    "coloring = np.select(\n",
    "    [orbits[0, :, 0] <= -0.2, np.abs(orbits[0, :, 0]) < +0.2, orbits[0, :, 0] >= 0.2], [-1.0, 0.0, +1.0]\n",
    ")\n",
    "\n",
    "dt = time_opts.dt\n",
    "Nt = sim.t_grid.size - 1\n",
    "interval = Nt / 20\n",
    "plot_ct = 0\n",
    "for i in range(Nt):\n",
    "    if i % interval == 0:\n",
    "        print(f\"{i=}\")\n",
    "        plot_ct += 1\n",
    "        plt.subplot(5, 2, plot_ct)\n",
    "        ax = plt.gca()\n",
    "        plt.scatter(orbits[i, :, 0], orbits[i, :, 1], c=coloring)\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 == 10:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "Let us perform another simulation, similar to the previous one. We will save the results in the folder `sim_2`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# light-weight model instance\n",
    "model = PressureLessSPH(epsilon=1.0)\n",
    "\n",
    "# environment options\n",
    "env = EnvironmentOptions(sim_folder=\"sim_2\")\n",
    "\n",
    "# simulation instance\n",
    "sim_tess = sim.spawn_sister(model=model, env=env)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "This time, we shall draw the markers on a regular grid obtained from a tesselation of the domain: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "loading_params = LoadingParameters(ppb=4, loading=\"tesselation\")\n",
    "weights_params = WeightsParameters()\n",
    "boundary_params = BoundaryParameters(bc=(\"reflect\", \"reflect\", \"periodic\"))\n",
    "model.cold_fluid.set_markers(\n",
    "    loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params, bufsize=0.5\n",
    ")\n",
    "model.cold_fluid.set_sorting_boxes(boxes_per_dim=(16, 16, 1))\n",
    "model.cold_fluid.set_save_data(n_markers=1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "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",
    "phi = bel_flow.p0\n",
    "model.propagators.push_v.options = model.propagators.push_v.Options(phi=phi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# background, perturbations and initial conditions\n",
    "model.cold_fluid.var.add_background(bel_flow)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_tess.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_tess.pproc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_tess.load_plotting_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(12, 28))\n",
    "\n",
    "orbits = sim_tess.orbits.cold_fluid\n",
    "\n",
    "coloring = np.select(\n",
    "    [orbits[0, :, 0] <= -0.2, np.abs(orbits[0, :, 0]) < +0.2, orbits[0, :, 0] >= 0.2], [-1.0, 0.0, +1.0]\n",
    ")\n",
    "\n",
    "dt = time_opts.dt\n",
    "Nt = sim_tess.t_grid.size - 1\n",
    "interval = Nt / 20\n",
    "plot_ct = 0\n",
    "for i in range(Nt):\n",
    "    if i % interval == 0:\n",
    "        print(f\"{i=}\")\n",
    "        plot_ct += 1\n",
    "        plt.subplot(5, 2, plot_ct)\n",
    "        ax = plt.gca()\n",
    "        plt.scatter(orbits[i, :, 0], orbits[i, :, 1], c=coloring)\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 == 10:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32",
   "metadata": {},
   "source": [
    "## Gas expansion\n",
    "\n",
    "We use SPH to solve Euler's equations (model `ViscousEulerSPH`):\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",
    "where $S$ denotes the entropy per unit mass and the internal energy per unit mass is \n",
    "\n",
    "$$\n",
    "\\mathcal U(\\rho, S) = \\kappa(S) \\log \\rho\\,.\n",
    "$$\n",
    "\n",
    "The SPH discretization leads to ODEs for $N$ particles indexed by $p$,\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) \\qquad && \\mathbf v_p(0) = \\mathbf u(\\mathbf x_p(0))\\,,\n",
    " \\end{aligned}\n",
    "$$\n",
    "\n",
    "where the smoothed density reads\n",
    "\n",
    "$$\n",
    " \\rho^{N,h}(\\mathbf x) = \\sum_{j=1}^N w_j W_h(\\mathbf x - \\mathbf x_j)\\,,\n",
    "$$\n",
    "\n",
    "with weights $w_p = const.$ and where $W_h(\\mathbf x)$ is a suitable smoothing kernel.\n",
    "The velocity update is performed with the Propagator `PushVinSPHpressure`.\n",
    "\n",
    "We shall now compute a gas expansion in 2d (nonlinear example). First, check out some of the smoothing kernels available for SPH evaluations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "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": "34",
   "metadata": {},
   "source": [
    "We start with the generic imports and also import the model `ViscousEulerSPH`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "from struphy import DerhamOptions, EnvironmentOptions, Time\n",
    "from struphy import domains\n",
    "from struphy import grids\n",
    "from struphy import (\n",
    "    BoundaryParameters,\n",
    "    LoadingParameters,\n",
    "    WeightsParameters,\n",
    ")\n",
    "from struphy import Simulation\n",
    "\n",
    "# import model\n",
    "from struphy.models import ViscousEulerSPH"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36",
   "metadata": {},
   "source": [
    "Here, it is important to set the base unit `kBT` in order to derive the velocity unit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "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": "38",
   "metadata": {},
   "source": [
    "As background, which goes into the initial condition below, we define a Gaussian blob in the xy-plane:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39",
   "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": "40",
   "metadata": {},
   "source": [
    "We also need a grid and Derham complex for projecting the fluid background on a spline basis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "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": "42",
   "metadata": {},
   "source": [
    "Now we create the light-weight instance of `ViscousEulerSPH`, without the optional propagator for particles in a magnetic background field, and instantiate the simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43",
   "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": "44",
   "metadata": {},
   "source": [
    "Note as well that we shall reject particles whose weight is below a certain threshold in order to save computing time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45",
   "metadata": {},
   "outputs": [],
   "source": [
    "loading_params = LoadingParameters(ppb=400)\n",
    "weights_params = WeightsParameters(reject_weights=True, threshold=3e-3)\n",
    "boundary_params = BoundaryParameters()\n",
    "model.euler_fluid.set_markers(\n",
    "    loading_params=loading_params,\n",
    "    weights_params=weights_params,\n",
    "    boundary_params=boundary_params,\n",
    ")\n",
    "nx = 16\n",
    "ny = 16\n",
    "model.euler_fluid.set_sorting_boxes(boxes_per_dim=(nx, ny, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46",
   "metadata": {},
   "source": [
    "For visualization of the result, we want to save a binning plot and a kernel density plot (sph evaluation of the fluid density):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47",
   "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",
    "model.euler_fluid.set_save_data(\n",
    "    n_markers=1.0,\n",
    "    binning_plots=(bin_plot,),\n",
    "    kernel_density_plots=(kd_plot,),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48",
   "metadata": {},
   "source": [
    "We choose `gaussian_2d` as the smoothing kernel for sph evaluations during the pressure step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49",
   "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": "50",
   "metadata": {},
   "source": [
    "Now we set the initial condition and run the simulation. The time steps take long at the beginning, but get faster towards the end, when particles are spread out over the domain. The simulation launched in the console will be a lot faster than in the notebook, especially when using MPI, or compiling with GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51",
   "metadata": {},
   "outputs": [],
   "source": [
    "# background, perturbations and initial conditions\n",
    "model.euler_fluid.var.add_background(blob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.pproc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.load_plotting_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55",
   "metadata": {},
   "source": [
    "The above output tells us where to find the post-processed simulation data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56",
   "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": "57",
   "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": "58",
   "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
}