{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Velocity Diffusion with SPH Viscosity\n",
    "\n",
    "## Pure Viscous Momentum Diffusion in 1D\n",
    "\n",
    "This tutorial verifies the SPH viscous propagator by simulating pure velocity diffusion — the momentum equation with viscosity but **without pressure forces**. In this limit the equation reduces to\n",
    "\n",
    "$$\\partial_t u_1 = \\frac{4}{3}\\,\\mu\\,\\partial_{x}^2 u_1,$$\n",
    "\n",
    "whose exact solution for a sinusoidal initial condition is\n",
    "\n",
    "$$u_1(x, t) = A_0\\,\\sin\\!\\left(\\frac{2\\pi\\ell\\, x}{L}\\right)\\exp\\!\\left(-\\gamma\\, t\\right), \\qquad \\gamma = \\frac{4}{3}\\,\\mu\\,k^2,\\quad k = \\frac{2\\pi\\ell}{L}.$$\n",
    "\n",
    "The factor $4/3$ comes from the compressible viscous stress tensor:\n",
    "for a 1D plane wave $\\partial_i u_j$ only has components along $i=j=1$, so the deviatoric stress contributes $2\\mu\\,(1 - 1/3) = (4/3)\\mu$ times the velocity gradient.\n",
    "\n",
    "### Verification procedure\n",
    "\n",
    "1. Initialise a 1D particle distribution (tessellation loading) with a velocity perturbation $\\delta u_1 \\propto \\sin(2\\pi x/L)$ and **no pressure force** (`with_p=False`).\n",
    "2. Run the simulation for a short time $T = 0.1$ (decay e-folding $1/\\gamma \\approx 0.038$ at the chosen parameters).\n",
    "3. Track the current $j_1 = \\rho u_1 \\approx u_1$ at the velocity antinode $x = L/4$ and fit the decay rate.\n",
    "4. Check that the numerical rate matches $\\gamma_\\text{analytical}$ to within 4%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cunumpy as xp\n",
    "\n",
    "from struphy import (\n",
    "    BinningPlot,\n",
    "    BoundaryParameters,\n",
    "    EnvironmentOptions,\n",
    "    KernelDensityPlot,\n",
    "    LoadingParameters,\n",
    "    SavingParameters,\n",
    "    Simulation,\n",
    "    SortingParameters,\n",
    "    Time,\n",
    "    WeightsParameters,\n",
    "    domains,\n",
    "    equils,\n",
    "    perturbations,\n",
    ")\n",
    "from struphy.models import ViscousEulerSPH\n",
    "from struphy.ode.utils import ButcherTableau\n",
    "\n",
    "logger = logging.getLogger(\"struphy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "### Physical and Numerical Parameters\n",
    "\n",
    "We use a large viscosity $\\mu = 1$ so that the decay is fast enough to observe over a short simulation. Mode $\\ell = 1$ gives wavenumber $k = 2\\pi/L$, and the analytical e-folding time is $\\tau = 1/\\gamma = 3/(4\\mu k^2)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Physical parameters\n",
    "mu  = 1.0    # dynamic viscosity (large → fast decay, clearly observable)\n",
    "r1  = 1.0    # domain length (1D periodic)\n",
    "\n",
    "# Mode and analytical decay rate: gamma = (4/3)*mu*k^2\n",
    "ell = 1\n",
    "k   = 2.0 * np.pi * ell / r1\n",
    "gamma_analytical = mu * (4.0 / 3.0) * k**2\n",
    "\n",
    "# Numerical parameters\n",
    "nx       = 8    # boxes in x-direction\n",
    "ppb      = 100  # particles per box (high density for accurate diffusion)\n",
    "plot_pts = 11   # KDE evaluation points\n",
    "\n",
    "# Time stepping: Tend ~ 0.1 covers ~2.6 e-folding times\n",
    "dt   = 0.0025\n",
    "Tend = 0.1\n",
    "\n",
    "print(f\"Viscosity:             mu  = {mu}\")\n",
    "print(f\"Domain length:         L   = {r1}\")\n",
    "print(f\"Wave mode:             ell = {ell},  k = {k:.4f}\")\n",
    "print(f\"Analytical decay rate: gamma = (4/3)*mu*k^2 = {gamma_analytical:.4f}\")\n",
    "print(f\"E-folding time:        tau = 1/gamma = {1/gamma_analytical:.4f}\")\n",
    "print(f\"Total particles:       {ppb * nx}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "### Model Setup\n",
    "\n",
    "`with_p=False` disables the pressure propagator so only the viscous term acts. The `push_viscous` propagator implements the SPH discretisation of the full compressible viscous stress divergence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pressure-free, viscosity-only model\n",
    "model = ViscousEulerSPH(with_B0=False, with_p=False, with_viscosity=True)\n",
    "\n",
    "butcher = ButcherTableau(algo=\"forward_euler\")\n",
    "model.propagators.push_eta.options = model.propagators.push_eta.Options(butcher=butcher)\n",
    "model.propagators.push_viscous.options = model.propagators.push_viscous.Options(\n",
    "    kernel_type=\"gaussian_1d\", mu=mu\n",
    ")\n",
    "\n",
    "print(\"ViscousEulerSPH model configured (no pressure, with viscosity).\")\n",
    "print(f\"  push_viscous: gaussian_1d kernel, mu={mu}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "### Domain and Particle Markers\n",
    "\n",
    "A 1D periodic domain of length $L=1$. The high particle count (`ppb=100`) is needed to resolve the kernel gradient accurately for the viscous term. Two `BinningPlot` diagnostics are registered: one for density and one for the current $j_1 = \\rho u_1$, which tracks the velocity amplitude over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = domains.Cuboid(r1=r1)\n",
    "grid = None\n",
    "derham_opts = None\n",
    "\n",
    "loading_params  = LoadingParameters(ppb=ppb, loading=\"tesselation\")\n",
    "weights_params  = WeightsParameters()\n",
    "boundary_params = BoundaryParameters()\n",
    "sorting_params  = SortingParameters(\n",
    "    boxes_per_dim=(nx, 1, 1),\n",
    "    dims_mask=(True, False, False),\n",
    ")\n",
    "\n",
    "bin_plot    = BinningPlot(slice=\"e1\", n_bins=(16,), ranges=(0.0, 1.0))\n",
    "bin_plot_j1 = BinningPlot(slice=\"e1\", n_bins=(16,), ranges=(0.0, 1.0), output_quantity=\"current_1\")\n",
    "kd_plot     = KernelDensityPlot(pts_e1=plot_pts, pts_e2=1)\n",
    "saving_params = SavingParameters(\n",
    "    binning_plots=(bin_plot, bin_plot_j1),\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",
    ")\n",
    "\n",
    "print(f\"Domain: 1D periodic, r1={r1}\")\n",
    "print(f\"Particles: {ppb} ppb × {nx} boxes = {ppb * nx} total\")\n",
    "print(f\"Diagnostics: density + j1 (current) binning, {plot_pts} KDE evaluation points\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "### Initial Conditions\n",
    "\n",
    "A sinusoidal **velocity** perturbation $\\delta u_1 = 0.5 \\sin(2\\pi x / L)$ with no density perturbation. The large amplitude (0.5) is chosen so the decay signal is clear over the short simulation window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "background  = equils.ConstantVelocity(ux=0.0)\n",
    "model.euler_fluid.var.add_background(background)\n",
    "\n",
    "perturbation = perturbations.ModesSin(ls=(1,), amps=(0.5,))\n",
    "model.euler_fluid.var.add_perturbation(del_u1=perturbation)\n",
    "\n",
    "print(\"Background: uniform density n=1, zero mean velocity\")\n",
    "print(\"Perturbation: delta_u1 = 0.5 * sin(2*pi*x/L)  [mode l=1]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "### Simulation Setup and Execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_folder = os.path.join(os.getcwd(), \"struphy_verification_tests\")\n",
    "out_folders = os.path.join(test_folder, \"ViscousEulerSPH\")\n",
    "env = EnvironmentOptions(out_folders=out_folders, sim_folder=\"velocity_diffusion\")\n",
    "\n",
    "time_opts = Time(dt=dt, Tend=Tend, split_algo=\"Strang\")\n",
    "\n",
    "sim = Simulation(\n",
    "    model=model,\n",
    "    env=env,\n",
    "    time_opts=time_opts,\n",
    "    domain=domain,\n",
    "    grid=grid,\n",
    "    derham_opts=derham_opts,\n",
    ")\n",
    "\n",
    "print(f\"Running velocity diffusion: dt={dt}, Tend={Tend}, {ppb * nx} particles\")\n",
    "out = sim.run()\n",
    "print(\"Simulation complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "### Access Diagnostics\n",
    "\n",
    "Field data is post-processed lazily, the first time it is accessed through the `out` object returned by `run`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "density = out.evaluate(\"euler_fluid/n\", dataset=\"view_0/n\")\n",
    "ee1, ee2, ee3 = np.meshgrid(density.eta1, density.eta2, density.eta3, indexing=\"ij\")\n",
    "n_sph     = np.asarray(density)                                              # shape (Nt+1, plot_pts, 1, 1)\n",
    "j1_current = out.evaluate(\"euler_fluid/f\", dataset=\"e1_current_1/f\")\n",
    "j1_binned = np.asarray(j1_current)                                           # shape (Nt+1, n_bins)\n",
    "e1_binned = np.asarray(j1_current[\"eta1\"])                                     # logical x in [0, 1]\n",
    "n_binned  = np.asarray(out.evaluate(\"euler_fluid/f\", dataset=\"e1_density/f\"))           # shape (Nt+1, n_bins)\n",
    "\n",
    "Nt    = int(Tend / dt)\n",
    "times = np.linspace(0.0, Tend, Nt + 1)\n",
    "\n",
    "e1_np = np.asarray(e1_binned).flatten()\n",
    "\n",
    "print(f\"Loaded {Nt + 1} time snapshots\")\n",
    "print(f\"j1_binned shape: {np.asarray(j1_binned).shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Decay Rate Analysis\n",
    "\n",
    "The velocity antinode of mode $\\ell = 1$ lies at $x = L/4$. We extract the time series $j_1(t)$ at the nearest bin, fit $\\ln|j_1|$ vs. $t$ with a straight line, and recover the numerical decay rate $\\gamma_\\text{numerical}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Antinode at x = L/4 = 0.25\n",
    "idx_max = int(np.argmin(np.abs(e1_np - 0.25)))\n",
    "amplitude = np.asarray(j1_binned[:, idx_max]).flatten()\n",
    "\n",
    "# Analytical envelope\n",
    "A0 = amplitude[0]\n",
    "amplitude_analytical = A0 * np.exp(-gamma_analytical * times)\n",
    "\n",
    "# Fit log(|amplitude|) vs time — pure diffusion, no oscillations\n",
    "log_amp = np.log(np.abs(amplitude) + 1e-15)\n",
    "coeffs  = np.polyfit(times, log_amp, 1)\n",
    "gamma_numerical = -coeffs[0]\n",
    "\n",
    "rel_error = abs(gamma_numerical - gamma_analytical) / gamma_analytical\n",
    "\n",
    "print(f\"Analytical decay rate: gamma = (4/3)*mu*k^2 = {gamma_analytical:.4f}\")\n",
    "print(f\"Numerical  decay rate: gamma =                {gamma_numerical:.4f}\")\n",
    "print(f\"Relative error: {rel_error * 100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "### Visualisation\n",
    "\n",
    "**Left panel**: current $j_1 = \\rho u_1 \\approx u_1$ profiles (from the binned diagnostic) at equally spaced times. The sinusoidal shape is preserved while the amplitude decreases uniformly — a hallmark of pure diffusion without mode mixing.\n",
    "\n",
    "**Right panel**: semi-log plot of $|j_1|$ at the antinode $x = L/4$, comparing the numerical decay against the analytical exponential $e^{-\\gamma t}$ with $\\gamma = (4/3)\\mu k^2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_bin    = np.asarray(e1_binned).flatten() * r1   # physical bin positions\n",
    "j1_arr   = np.asarray(j1_binned)                  # shape (Nt+1, n_bins)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
    "\n",
    "# --- Left: j1 (velocity) profiles at equally spaced times ---\n",
    "ax = axes[0]\n",
    "n_snaps  = 8\n",
    "snap_ids = np.round(np.linspace(0, Nt, n_snaps)).astype(int)\n",
    "cmap_t   = plt.get_cmap(\"viridis\", n_snaps)\n",
    "for i, idx in enumerate(snap_ids):\n",
    "    ax.plot(x_bin, j1_arr[idx, :], color=cmap_t(i), linewidth=2,\n",
    "            label=f\"t={times[idx]:.3f}\")\n",
    "ax.set_xlabel(\"$x$\")\n",
    "ax.set_ylabel(r\"$j_1 = \\rho u_1$ (binned)\")\n",
    "ax.set_title(\"Velocity profiles $j_1(x)$ over time\")\n",
    "ax.set_ylim([-0.6, 0.6])\n",
    "ax.legend(fontsize=8)\n",
    "ax.grid(True, linestyle=\"--\", alpha=0.5)\n",
    "\n",
    "# --- Right: amplitude decay (semi-log) ---\n",
    "ax = axes[1]\n",
    "ax.semilogy(times, np.abs(amplitude), \"b-\", linewidth=1.2,\n",
    "            label=rf\"Numerical (fitted $\\gamma={gamma_numerical:.3f}$)\")\n",
    "ax.semilogy(times, np.abs(amplitude_analytical), \"k--\",\n",
    "            label=rf\"Analytical: $\\gamma={(4/3)*mu*k**2:.3f}$\")\n",
    "ax.set_xlabel(\"time\")\n",
    "ax.set_ylabel(rf\"$|j_1|$ at $x={e1_np[idx_max]*r1:.3f}$\")\n",
    "ax.set_title(\"Amplitude decay (log scale)\")\n",
    "ax.legend()\n",
    "ax.grid(True, which=\"both\", linestyle=\"--\", alpha=0.5)\n",
    "\n",
    "fig.suptitle(\n",
    "    rf\"Velocity diffusion: $\\mu={mu}$, $k={k:.3f}$, $\\gamma_\\mathrm{{anal}}={gamma_analytical:.3f}$\",\n",
    "    fontsize=12,\n",
    ")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "### Verification Check\n",
    "\n",
    "Two assertions are evaluated:\n",
    "1. The final velocity amplitude is effectively zero (diffusion has erased the mode).\n",
    "2. The fitted decay rate agrees with the analytical value to within 4%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_error  = float(xp.max(xp.abs(j1_binned[-1])))\n",
    "tol_final    = 0.0022\n",
    "tol_rate     = 0.04   # 4% relative error on the decay rate\n",
    "\n",
    "print(\"=== Velocity Diffusion Verification ===\")\n",
    "print(f\"  Final velocity amplitude:  {final_error:.4e}  (tolerance {tol_final:.0e})\")\n",
    "print(f\"  Decay rate relative error: {rel_error * 100:.2f}%  (tolerance {tol_rate*100:.0f}%)\")\n",
    "\n",
    "try:\n",
    "    assert final_error < tol_final, (\n",
    "        f\"Final amplitude {final_error:.4e} exceeds tolerance {tol_final:.0e}\"\n",
    "    )\n",
    "    print(\"\\n✓ Final amplitude check passed.\")\n",
    "except AssertionError as e:\n",
    "    print(f\"\\n✗ {e}\")\n",
    "\n",
    "try:\n",
    "    assert rel_error < tol_rate, (\n",
    "        f\"Decay rate {gamma_numerical:.4f} deviates {rel_error*100:.1f}% \"\n",
    "        f\"from analytical {gamma_analytical:.4f} (tolerance {tol_rate*100:.0f}%)\"\n",
    "    )\n",
    "    print(\"✓ Decay rate check passed.\")\n",
    "except AssertionError as e:\n",
    "    print(f\"✗ {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "### Conclusion\n",
    "\n",
    "This tutorial verified the SPH viscous propagator using pure velocity diffusion in 1D:\n",
    "\n",
    "- With **pressure disabled** (`with_p=False`), the system reduces to the heat equation for momentum, with exact exponential decay $e^{-\\gamma t}$ where $\\gamma = (4/3)\\mu k^2$.\n",
    "- The SPH kernel gradient approximation reproduces the diffusion coefficient to within 4%, confirming the correctness of the viscous stress tensor implementation.\n",
    "- The high particle count (`ppb=100`) is necessary for accurate kernel gradient estimates in the diffusion regime; fewer particles would overestimate dissipation.\n",
    "- This is a stringent test because it targets a single mechanism (viscosity) in isolation, without the competing dynamics of pressure waves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional cleanup\n",
    "if False:   # set to True to remove simulation output\n",
    "    shutil.rmtree(test_folder)\n",
    "    print(f\"Cleaned up {test_folder}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env (3.12.3.final.0)",
   "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
}