{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c6e1066c",
   "metadata": {},
   "source": [
    "# Variable-density low-Mach closure - warm-bubble sanity check\n",
    "\n",
    "**This is a sanity check, not a benchmark.** There is no closed-form reference\n",
    "solution. The case validates two things about the variable-density low-Mach\n",
    "thermal closure (Taha et al. 2024), run with the in-collision bulk-viscosity\n",
    "stabilization (`models.LBM.bulk_viscosity`) enabled:\n",
    "\n",
    "1. **Stability** - the closure runs the full 6000 steps with no divergence; the\n",
    "   peak speed stays bounded and subsonic, i.e. there is no exponentially growing\n",
    "   reduced-pressure checkerboard / acoustic mode (with the bulk viscosity *off*\n",
    "   the same setup diverges within a few hundred steps).\n",
    "2. **Physical buoyancy** - a warm Gaussian bubble seeded low in the box, made\n",
    "   lighter by the equation of state `rho = P/(r T)`, rises (mean `uz > 0`) while\n",
    "   diffusing toward ambient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05173ed9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T18:52:42.158350Z",
     "iopub.status.busy": "2026-06-19T18:52:42.158283Z",
     "iopub.status.idle": "2026-06-19T18:52:43.073032Z",
     "shell.execute_reply": "2026-06-19T18:52:43.072699Z"
    }
   },
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import pathlib\n",
    "\n",
    "import h5py\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import nassu.viz as common\n",
    "from nassu.cfg.model import ConfigScheme\n",
    "\n",
    "common.use_style()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02b174b1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T18:52:43.074458Z",
     "iopub.status.busy": "2026-06-19T18:52:43.074342Z",
     "iopub.status.idle": "2026-06-19T18:52:43.097604Z",
     "shell.execute_reply": "2026-06-19T18:52:43.097281Z"
    }
   },
   "outputs": [],
   "source": [
    "def _find_project_root() -> pathlib.Path:\n",
    "    here = pathlib.Path.cwd().resolve()\n",
    "    for cand in [here, *here.parents]:\n",
    "        if (cand / \"pyproject.toml\").exists() and (cand / \"nassu\").is_dir():\n",
    "            return cand\n",
    "    raise RuntimeError(f\"Could not locate Nassu project root upward from {here}\")\n",
    "\n",
    "\n",
    "# Run from the project root: the simulation `save_path` (and every output path\n",
    "# read off the parsed config below) is repo-root-relative.\n",
    "project_root = _find_project_root()\n",
    "os.chdir(project_root)\n",
    "case_dir = project_root / \"validation/buoyant_flows/01_lowmach_bubble\"\n",
    "config_path = case_dir / \"01_lowmach_bubble.nassu.yaml\"\n",
    "\n",
    "sim_cfg = ConfigScheme.sim_cfgs_from_file(str(config_path))[0]\n",
    "n_steps = sim_cfg.n_steps\n",
    "print(f\"loaded {sim_cfg.name}: n_steps={n_steps}, domain={sim_cfg.domain.domain_size}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b85f3277",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T18:52:43.098609Z",
     "iopub.status.busy": "2026-06-19T18:52:43.098534Z",
     "iopub.status.idle": "2026-06-19T18:52:43.146726Z",
     "shell.execute_reply": "2026-06-19T18:52:43.146376Z"
    }
   },
   "outputs": [],
   "source": [
    "# Locate the instantaneous volume export off the parsed config and read each\n",
    "# snapshot. The volume export stores the full domain as a single block per time\n",
    "# group: `t<step>/block0/<macr>`, each a 3-D array in (z, y, x) layout.\n",
    "export = sim_cfg.output.exports[\"default\"].volumes[\"default\"].inst\n",
    "h5_path = sorted(glob.glob(str(export.base_path) + \"*.h5\"))[0]\n",
    "print(\"reading\", h5_path)\n",
    "\n",
    "\n",
    "def read_snapshots(path):\n",
    "    with h5py.File(path, \"r\") as h:\n",
    "        times = sorted({k.split(\"/\")[0] for k in h.keys()}, key=lambda s: float(s[1:]))\n",
    "        steps = np.array([float(t[1:]) for t in times])\n",
    "        data = {}\n",
    "        for t in times:\n",
    "            g = h[t][\"block0\"]\n",
    "            data[t] = {m: g[m][:] for m in (\"rho\", \"ux\", \"uy\", \"uz\", \"temp\")}\n",
    "    return steps, times, data\n",
    "\n",
    "\n",
    "steps, times, snaps = read_snapshots(h5_path)\n",
    "print(f\"{len(steps)} snapshots, steps {steps[0]:.0f} .. {steps[-1]:.0f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b3e5f09",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T18:52:43.147678Z",
     "iopub.status.busy": "2026-06-19T18:52:43.147600Z",
     "iopub.status.idle": "2026-06-19T18:52:43.161722Z",
     "shell.execute_reply": "2026-06-19T18:52:43.161389Z"
    }
   },
   "outputs": [],
   "source": [
    "# Per-snapshot diagnostics. All are reductions over the whole field, so the\n",
    "# (z, y, x) axis layout is immaterial here.\n",
    "umax, uz_mean, t_max, rho_min, rho_max = [], [], [], [], []\n",
    "any_nan = False\n",
    "for t in times:\n",
    "    s = snaps[t]\n",
    "    speed = np.sqrt(s[\"ux\"] ** 2 + s[\"uy\"] ** 2 + s[\"uz\"] ** 2)\n",
    "    umax.append(float(np.nanmax(speed)))\n",
    "    uz_mean.append(float(np.nanmean(s[\"uz\"])))\n",
    "    t_max.append(float(np.nanmax(s[\"temp\"])))\n",
    "    rho_min.append(float(np.nanmin(s[\"rho\"])))\n",
    "    rho_max.append(float(np.nanmax(s[\"rho\"])))\n",
    "    any_nan = any_nan or bool(np.isnan(s[\"rho\"]).any() or np.isnan(s[\"uz\"]).any())\n",
    "\n",
    "umax = np.array(umax)\n",
    "uz_mean = np.array(uz_mean)\n",
    "t_max = np.array(t_max)\n",
    "rho_min = np.array(rho_min)\n",
    "rho_max = np.array(rho_max)\n",
    "cs = 1.0 / np.sqrt(3.0)\n",
    "print(f\"peak |u| over run = {umax.max():.3e}  (Ma_max = {umax.max() / cs:.3e})\")\n",
    "print(f\"final mean uz     = {uz_mean[-1]:.3e}  (buoyant rise, > 0)\")\n",
    "print(f\"NaN encountered   = {any_nan}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22c8dc58",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71bc1d9f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T18:52:43.162798Z",
     "iopub.status.busy": "2026-06-19T18:52:43.162687Z",
     "iopub.status.idle": "2026-06-19T18:52:43.548212Z",
     "shell.execute_reply": "2026-06-19T18:52:43.546703Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(11, 4))\n",
    "\n",
    "ax[0].plot(steps, umax, \"-o\", ms=3, label=r\"$\\max|u|$\")\n",
    "ax[0].plot(steps, uz_mean, \"-s\", ms=3, label=r\"mean $u_z$ (buoyant rise)\")\n",
    "ax[0].axhline(0.1 * cs, ls=\"--\", c=\"grey\", lw=1, label=r\"$\\mathrm{Ma}=0.1$\")\n",
    "ax[0].set_xlabel(\"step\")\n",
    "ax[0].set_ylabel(\"lattice velocity\")\n",
    "ax[0].set_title(\"Velocity stays bounded, subsonic; mean $u_z>0$\")\n",
    "ax[0].legend()\n",
    "ax[0].grid(alpha=0.3)\n",
    "\n",
    "ax[1].plot(steps, t_max, \"-o\", ms=3, label=r\"$T_{\\max}$\")\n",
    "ax[1].plot(steps, rho_min, \"-s\", ms=3, label=r\"$\\rho_{\\min}$ (warm core)\")\n",
    "ax[1].plot(steps, rho_max, \"-^\", ms=3, label=r\"$\\rho_{\\max}$\")\n",
    "ax[1].axhline(1.0, ls=\"--\", c=\"grey\", lw=1, label=\"ambient\")\n",
    "ax[1].set_xlabel(\"step\")\n",
    "ax[1].set_ylabel(\"value\")\n",
    "ax[1].set_title(\"Bubble diffuses toward ambient; warm core stays lighter\")\n",
    "ax[1].legend()\n",
    "ax[1].grid(alpha=0.3)\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c970b11",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T18:52:43.549151Z",
     "iopub.status.busy": "2026-06-19T18:52:43.549052Z",
     "iopub.status.idle": "2026-06-19T18:52:43.783609Z",
     "shell.execute_reply": "2026-06-19T18:52:43.783318Z"
    }
   },
   "outputs": [],
   "source": [
    "# Vertical mid-plane (y = Ny/2) temperature slices: the warm bubble diffuses\n",
    "# and gently rises along -gravity (+z). Array layout is (z, y, x).\n",
    "sel = [times[0], times[len(times) // 3], times[2 * len(times) // 3], times[-1]]\n",
    "fig, axes = plt.subplots(1, len(sel), figsize=(13, 4.2), sharey=True)\n",
    "ny = snaps[times[0]][\"temp\"].shape[1]\n",
    "vmin, vmax = 1.0, t_max[0]\n",
    "for ax, t in zip(axes, sel):\n",
    "    sl = snaps[t][\"temp\"][:, ny // 2, :]  # (z, x)\n",
    "    im = ax.imshow(sl, origin=\"lower\", aspect=\"auto\", cmap=\"inferno\", vmin=vmin, vmax=vmax)\n",
    "    ax.set_title(f\"step {int(float(t[1:]))}\")\n",
    "    ax.set_xlabel(\"x\")\n",
    "axes[0].set_ylabel(\"z (gravity is -z)\")\n",
    "fig.colorbar(im, ax=axes, label=\"T\", fraction=0.025)\n",
    "fig.suptitle(\"Mid-plane temperature: warm bubble diffuses and rises\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8416f903",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Success criteria for this sanity check:\n",
    "\n",
    "- **No divergence** - all snapshots are finite (no NaN in `rho`/`u`).\n",
    "- **Bounded, subsonic velocity** - `max|u|` stays well below `Ma = 0.1`, with no\n",
    "  exponential growth (the checkerboard / acoustic mode is suppressed by the\n",
    "  in-collision bulk viscosity).\n",
    "- **Physical buoyant rise** - the mean vertical velocity is positive.\n",
    "- **Diffusive relaxation** - the peak temperature decays toward ambient and the\n",
    "  warm core stays EOS-consistently lighter than ambient (`rho_min < 1`).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4d32053",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T18:52:43.784573Z",
     "iopub.status.busy": "2026-06-19T18:52:43.784494Z",
     "iopub.status.idle": "2026-06-19T18:52:43.786784Z",
     "shell.execute_reply": "2026-06-19T18:52:43.786558Z"
    }
   },
   "outputs": [],
   "source": [
    "assert not any_nan, \"divergence: NaN found in the fields\"\n",
    "assert umax.max() < 0.1 * cs, f\"velocity not subsonic: max|u|={umax.max():.3e}\"\n",
    "assert umax.max() < 5e-3, f\"velocity grew unexpectedly large: {umax.max():.3e}\"\n",
    "assert uz_mean[-1] > 0, f\"no buoyant rise: final mean uz={uz_mean[-1]:.3e}\"\n",
    "assert t_max[-1] < t_max[0], \"peak temperature did not relax toward ambient\"\n",
    "assert rho_min.min() < 1.0, \"warm core is not lighter than ambient\"\n",
    "print(\"SANITY CHECK PASSED: stable, subsonic, buoyant, diffusing.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99fe3af2",
   "metadata": {},
   "source": [
    "## Version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f3b2c79",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T18:52:43.787509Z",
     "iopub.status.busy": "2026-06-19T18:52:43.787440Z",
     "iopub.status.idle": "2026-06-19T18:52:43.802900Z",
     "shell.execute_reply": "2026-06-19T18:52:43.802638Z"
    }
   },
   "outputs": [],
   "source": [
    "sim_info = sim_cfg.output.read_info()\n",
    "print(sim_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bb751ec",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e190e64e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T18:52:43.803641Z",
     "iopub.status.busy": "2026-06-19T18:52:43.803569Z",
     "iopub.status.idle": "2026-06-19T18:52:43.812096Z",
     "shell.execute_reply": "2026-06-19T18:52:43.811729Z"
    }
   },
   "outputs": [],
   "source": [
    "from IPython.display import Code\n",
    "\n",
    "Code(filename=str(config_path), language=\"yaml\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}