{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# 2D Dam Break Simulation with SPH\n",
    "\n",
    "## Collapse of a Fluid Column under Gravity\n",
    "\n",
    "This tutorial simulates a classic **2D dam break**: a dense fluid column confined to the left quarter of a closed box is released at $t=0$. Gravity pulls the fluid downward and the pressure gradient drives a horizontal spreading wave. The SPH method naturally handles the large interface deformation and free-surface dynamics.\n",
    "\n",
    "### Physical Setup\n",
    "\n",
    "The domain is a $1 \\times 1$ box with **reflective walls** on all sides. At $t = 0$ the fluid occupies the region $x \\in [0, L/4]$ at density $\\rho_\\text{high}$; the rest of the box is effectively empty (markers with negligible weight are removed). As the fluid evolves:\n",
    "\n",
    "1. The column collapses under downward gravity $g_y$.\n",
    "2. The pressure gradient ($\\propto c_s^2 \\nabla\\rho$ in the isothermal model) drives a shock-like spreading wave.\n",
    "3. Viscosity damps small-scale velocity fluctuations and stabilises the SPH.\n",
    "4. Reflections off the walls cause the fluid to slosh and eventually reach a hydrostatic equilibrium with density concentrated at the bottom.\n",
    "\n",
    "### SPH parameters for compressible free-surface flows\n",
    "\n",
    "The isothermal equation of state $p = \\kappa \\rho$ is used with a **weak compressibility** coefficient $\\kappa = c_s^2 \\ll 1$. This keeps the flow subsonic (Mach number $\\text{Ma} = U_\\text{max}/c_s \\ll 1$) while allowing large density variations — the *Weakly Compressible SPH* (WCSPH) regime.\n",
    "\n",
    "Markers with very small weights (corresponding to the near-vacuum region) are removed by the `reject_weights` option before the simulation starts.\n",
    "\n",
    "### What to expect\n",
    "\n",
    "- The dense column hits the right wall, runs up it, and splashes back.\n",
    "- Multiple reflections produce complex sloshing dynamics.\n",
    "- At late times ($t \\gtrsim 2$) the fluid settles near the bottom, forming a stable layer.\n",
    "- As a qualitative verification we check that **no markers escape the closed box** throughout the simulation."
   ]
  },
  {
   "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",
    "\n",
    "from struphy import (\n",
    "    BoundaryParameters,\n",
    "    EnvironmentOptions,\n",
    "    KernelDensityPlot,\n",
    "    LoadingParameters,\n",
    "    SavingParameters,\n",
    "    Simulation,\n",
    "    SortingParameters,\n",
    "    Time,\n",
    "    WeightsParameters,\n",
    "    domains,\n",
    "    equils,\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",
    "The free-fall time from height $H/2$ is $\\sqrt{2(H/2)/g_y} \\approx 0.32$, so $T_\\text{end} = 3$ covers roughly 9 free-fall times. The sound speed $c_s = \\sqrt{\\kappa} \\approx 0.45$, giving Mach number $\\text{Ma} = \\sqrt{2 g_y H/2}/c_s \\approx 0.71$ — weakly supersonic, which is acceptable for WCSPH with small $\\kappa$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Physical parameters\n",
    "kappa  = 0.2    # isothermal coefficient (= c_s^2); weak compressibility\n",
    "mu     = 0.05   # dynamic viscosity (small, for stability)\n",
    "g_y    = 10.0   # gravitational acceleration (downward = -y direction)\n",
    "r1     = 1.0    # domain width  (x)\n",
    "r2     = 1.0    # domain height (y)\n",
    "n_high = 0.1    # initial density of the fluid column\n",
    "\n",
    "# Derived\n",
    "c_s     = kappa**0.5\n",
    "U_max   = (2.0 * g_y * r2 / 2.0)**0.5   # rough free-fall velocity scale\n",
    "Ma      = U_max / c_s\n",
    "\n",
    "# Numerical parameters\n",
    "nx       = 8    # boxes per spatial dimension\n",
    "ppb      = 32   # particles per box\n",
    "plot_pts = 21   # KDE evaluation points per dimension\n",
    "\n",
    "# Time stepping: CFL limit ~ h/c_s = (r1/nx)/c_s\n",
    "dt   = 0.02\n",
    "Tend = 3.0\n",
    "\n",
    "print(f\"Sound speed:      c_s  = {c_s:.3f}\")\n",
    "print(f\"Max velocity est: U_max = {U_max:.3f}  (Mach = {Ma:.2f})\")\n",
    "print(f\"Fluid density:    n_high = {n_high}  (vacuum elsewhere)\")\n",
    "print(f\"Time stepping:    dt={dt}, Tend={Tend}, {int(Tend/dt)} steps\")\n",
    "print(f\"Total particles:  ~{ppb * nx * nx // 4} (left quarter of domain)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "### Model Setup\n",
    "\n",
    "All three physical effects are active: pressure (`with_p=True`), viscosity (`with_viscosity=True`). The 2D Gaussian kernel is used. Gravity enters as a downward vector `(0, -g_y, 0)` in the pressure propagator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ViscousEulerSPH(with_B0=False, with_p=True, 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_sph_p.options   = model.propagators.push_sph_p.Options(\n",
    "    kernel_type=\"gaussian_2d\",\n",
    "    gravity=(0.0, -g_y, 0.0),   # downward gravity\n",
    "    kappa=kappa,\n",
    ")\n",
    "model.propagators.push_viscous.options = model.propagators.push_viscous.Options(\n",
    "    kernel_type=\"gaussian_2d\",\n",
    "    mu=mu,\n",
    ")\n",
    "\n",
    "print(\"ViscousEulerSPH configured (pressure + viscosity + gravity).\")\n",
    "print(f\"  kappa={kappa}, mu={mu}, g_y={g_y} (downward)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "### Domain, Boundary Conditions and Diagnostics\n",
    "\n",
    "The box has **reflective walls** on all sides (`bc=\"reflect\"`) with SPH mirror ghost particles (`bc_sph=\"mirror\"`). The `reject_weights` option removes near-vacuum markers (those with weight below the threshold) so that only the dense left-column particles are simulated.\n",
    "\n",
    "The `n_markers=1.0` option in `SavingParameters` saves every marker orbit at each time step, enabling post-hoc visualisation of particle trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = domains.Cuboid(r1=r1, r2=r2)\n",
    "\n",
    "loading_params  = LoadingParameters(ppb=ppb, loading=\"tesselation\")\n",
    "weights_params  = WeightsParameters(reject_weights=True, threshold=1e-6)\n",
    "boundary_params = BoundaryParameters(\n",
    "    bc    =(\"reflect\", \"reflect\",  \"periodic\"),\n",
    "    bc_sph=(\"mirror\",  \"mirror\",   \"periodic\"),\n",
    ")\n",
    "sorting_params = SortingParameters(\n",
    "    boxes_per_dim=(nx, nx, 1),\n",
    "    dims_mask=(True, True, False),\n",
    ")\n",
    "\n",
    "kd_plot = KernelDensityPlot(pts_e1=plot_pts, pts_e2=plot_pts, pts_e3=1)\n",
    "saving_params = SavingParameters(\n",
    "    n_markers=1.0,              # save all marker positions every step\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",
    "    bufsize=2,\n",
    ")\n",
    "\n",
    "print(f\"2D closed box [{r1}×{r2}], reflective walls on all sides\")\n",
    "print(\"Vacuum markers (weight < 1e-6) removed before simulation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "### Initial Conditions\n",
    "\n",
    "The `step_function_xy` density profile places $\\rho = \\rho_\\text{high}$ where $x < L/4$ and $y < H$, and near-vacuum ($\\sim 10^{-8}$) elsewhere. The near-vacuum markers are then removed by the `reject_weights` filter, leaving only the dense left-column particles.\n",
    "\n",
    "We colour each marker by its initial $x$-position (normalised to $[0, 1]$ within the column) to track how the fluid mixes during the dam break."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "background = equils.ConstantVelocity(\n",
    "    density_profile=\"step_function_xy\",\n",
    "    n=n_high,\n",
    "    upper_x=r1 / 4,   # dense column: x < r1/4\n",
    "    upper_y=r2,\n",
    ")\n",
    "model.euler_fluid.var.add_background(background)\n",
    "\n",
    "print(f\"Initial condition: dense column (n={n_high}) for x < {r1/4}\")\n",
    "print(\"Near-vacuum markers (x > r1/4) will be removed by reject_weights\")"
   ]
  },
  {
   "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=\"dam_break\")\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=None,\n",
    "    derham_opts=None,\n",
    ")\n",
    "\n",
    "print(f\"Running 2D dam break: dt={dt}, Tend={Tend}\")\n",
    "sim.run()\n",
    "print(\"Simulation complete.\")\n",
    "\n",
    "sim.pproc()\n",
    "print(\"Post-processing complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "### Load Diagnostics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.load_plotting_data()\n",
    "\n",
    "# KDE density field: shape (Nt+1, pts_e1, pts_e2, 1)\n",
    "ee1, ee2, ee3 = sim.n_sph.euler_fluid.view_0.grid_n_sph\n",
    "n_sph = sim.n_sph.euler_fluid.view_0.n_sph\n",
    "\n",
    "# Marker orbits: shape (Nt_orb, n_markers, n_attrs)\n",
    "# attrs for vdim=2: [x, y, z, v1, v2, w, diag, id]\n",
    "orbits = np.asarray(sim.orbits.euler_fluid)\n",
    "\n",
    "Nt      = int(Tend / dt)\n",
    "times   = np.linspace(0.0, Tend, Nt + 1)\n",
    "Nt_orb  = orbits.shape[0]\n",
    "t_orbit = np.linspace(0.0, Tend, Nt_orb)\n",
    "\n",
    "X = np.asarray(ee1)[:, :, 0] * r1   # physical x, shape (pts_e1, pts_e2)\n",
    "Y = np.asarray(ee2)[:, :, 0] * r2   # physical y\n",
    "n_arr = np.asarray(n_sph)            # (Nt+1, pts_e1, pts_e2, 1)\n",
    "\n",
    "# Colour each marker by its initial x position within the column\n",
    "x_init = orbits[0, :, 0]\n",
    "c_val  = x_init / (r1 / 4.0)        # 0 = left wall, 1 = dam face\n",
    "\n",
    "print(f\"KDE field shape:   {n_arr.shape}\")\n",
    "print(f\"Marker orbits:     {orbits.shape}  [{Nt_orb} snapshots, {orbits.shape[1]} markers]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Visualisation: Density Field Snapshots\n",
    "\n",
    "Twelve equally spaced snapshots show the KDE density field (colour map) overlaid with marker positions (coloured by initial $x$). The colour gradient from blue (left wall) to red (dam face) reveals how the fluid column mixes as it spreads across the domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "snapshot_inds = np.round(np.linspace(0, Nt, 12)).astype(int)\n",
    "orb_inds      = np.round(np.linspace(0, Nt_orb - 1, 12)).astype(int)\n",
    "vmax_plot     = float(np.max(n_arr))\n",
    "\n",
    "fig, axes = plt.subplots(4, 3, figsize=(15, 12), sharex=True, sharey=True)\n",
    "im = None\n",
    "for ax, idx, oidx in zip(axes.flatten(), snapshot_inds, orb_inds):\n",
    "    n_2d = n_arr[idx, :, :, 0]\n",
    "    im = ax.pcolormesh(X, Y, n_2d, vmin=0.0, vmax=vmax_plot, cmap=\"Blues\", shading=\"auto\")\n",
    "    ax.scatter(\n",
    "        orbits[oidx, :, 0],\n",
    "        orbits[oidx, :, 1],\n",
    "        c=c_val, cmap=\"autumn\", s=3,\n",
    "        vmin=0.0, vmax=1.0, alpha=0.7,\n",
    "    )\n",
    "    ax.set_title(f\"$t = {times[idx]:.2f}$\", fontsize=10)\n",
    "    ax.set_aspect(\"equal\")\n",
    "\n",
    "for ax in axes[-1, :]:\n",
    "    ax.set_xlabel(\"$x$\")\n",
    "for ax in axes[:, 0]:\n",
    "    ax.set_ylabel(\"$y$\")\n",
    "\n",
    "if im is not None:\n",
    "    fig.colorbar(im, ax=axes.ravel().tolist(), label=r\"$\\rho$ (KDE)\", shrink=0.6)\n",
    "\n",
    "fig.suptitle(\n",
    "    rf\"2D dam break: $\\kappa={kappa}$, $g_y={g_y}$, $\\mu={mu}$, {nx}\\times{nx} boxes\",\n",
    "    fontsize=12,\n",
    ")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "### Visualisation: Final Marker State\n",
    "\n",
    "Show the final particle configuration to check that the fluid has settled into a stable layer at the bottom of the box."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 6))\n",
    "sc = ax.scatter(\n",
    "    orbits[-1, :, 0],\n",
    "    orbits[-1, :, 1],\n",
    "    c=c_val, cmap=\"autumn\", s=8, vmin=0.0, vmax=1.0,\n",
    ")\n",
    "ax.set_xlim(0.0, r1)\n",
    "ax.set_ylim(0.0, r2)\n",
    "ax.set_xlabel(\"$x$\")\n",
    "ax.set_ylabel(\"$y$\")\n",
    "ax.set_title(rf\"Final state at $t = {Tend:.1f}$\")\n",
    "ax.set_aspect(\"equal\")\n",
    "plt.colorbar(sc, ax=ax, label=\"Initial position (0=left wall, 1=dam face)\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "### Verification Check\n",
    "\n",
    "Since the dam break has no simple analytical steady state, the verification is a **domain-bounds assertion**: no marker should escape the closed reflective box. A 1% tolerance accounts for the finite displacement during a single time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_all = orbits[:, :, 0]\n",
    "y_all = orbits[:, :, 1]\n",
    "\n",
    "x_min, x_max = float(np.min(x_all)), float(np.max(x_all))\n",
    "y_min, y_max = float(np.min(y_all)), float(np.max(y_all))\n",
    "\n",
    "print(\"=== Dam Break Domain Bounds Verification ===\")\n",
    "print(f\"  x range: [{x_min:.4f}, {x_max:.4f}]  (domain [0, {r1}])\")\n",
    "print(f\"  y range: [{y_min:.4f}, {y_max:.4f}]  (domain [0, {r2}])\")\n",
    "\n",
    "tol = 0.01   # 1% of domain size\n",
    "\n",
    "try:\n",
    "    assert x_min >= -tol * r1 and x_max <= (1.0 + tol) * r1, (\n",
    "        f\"Markers escaped x-domain: x in [{x_min:.4f}, {x_max:.4f}]\"\n",
    "    )\n",
    "    print(\"\\n✓ x-domain bounds check passed.\")\n",
    "except AssertionError as e:\n",
    "    print(f\"\\n✗ {e}\")\n",
    "\n",
    "try:\n",
    "    assert y_min >= -tol * r2 and y_max <= (1.0 + tol) * r2, (\n",
    "        f\"Markers escaped y-domain: y in [{y_min:.4f}, {y_max:.4f}]\"\n",
    "    )\n",
    "    print(\"✓ y-domain bounds check passed.\")\n",
    "except AssertionError as e:\n",
    "    print(f\"✗ {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "### Conclusion\n",
    "\n",
    "This tutorial demonstrated the SPH method for a free-surface flow problem:\n",
    "\n",
    "- **Weakly Compressible SPH (WCSPH)**: the isothermal equation of state with small $\\kappa$ provides a pressure response that keeps the flow subsonic while allowing large density variations — the hallmark of free-surface SPH.\n",
    "- **Mirror ghost particles** enforce reflective boundary conditions on all four walls without special treatment for free surfaces.\n",
    "- **Reject-weights** cleanly removes near-vacuum markers before the simulation, avoiding spurious SPH interactions in the void region.\n",
    "- **Marker colouring** by initial position reveals the mixing and transport patterns during the collapse and subsequent sloshing.\n",
    "- The domain-bounds assertion verifies that the boundary conditions work correctly throughout the highly dynamic simulation.\n",
    "\n",
    "The dam break is a standard benchmark for validating SPH implementations. More quantitative comparisons (run-up height, wave arrival time) can be made against published experimental data or higher-resolution SPH simulations."
   ]
  },
  {
   "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
}