{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# SPH in Beltrami force field\n",
    "\n",
    "In this tutorial, you will set up and run a pressure-less SPH example in Struphy, then inspect particle evolution with simple diagnostics.\n",
    "\n",
    "This notebook focuses on pressure-less flow in a Beltrami force field (model `PressureLessSPH`). If you want to start from template parameter files, generate them from the command line:\n",
    "\n",
    "```\n",
    "struphy params PressureLessSPH\n",
    "```\n",
    "\n",
    "The notebook keeps everything in Python so you can adjust parameters interactively and rerun sections quickly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Pressure-Less Flow in a Beltrami Force Field\n",
    "\n",
    "We begin with a particle system in a 3D box domain $\\Omega \\subset \\mathbb{R}^3$.\n",
    "\n",
    "Given an initial velocity field $\\mathbf{u}_0(\\mathbf{x})$, we seek trajectories $(\\mathbf{x}_p, \\mathbf{v}_p)$ for particles $p=0,\\dots,N-1$ satisfying\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",
    "Here, $p \\in H^1(\\Omega)$ is a prescribed scalar field. In this example, we choose $p$ so that the force field has Beltrami structure.\n",
    "\n",
    "In practice, the workflow is: import APIs, define geometry and time stepping, set marker loading, configure propagators, run, and visualize.\n",
    "\n",
    "First, import the Struphy components used throughout this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "    BoundaryParameters,\n",
    "    LoadingParameters,\n",
    "    WeightsParameters,\n",
    "    SortingParameters,\n",
    "    SavingParameters,\n",
    ")\n",
    "from struphy import Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "### Step 1: Create Environment, Time Integrator, and Domain\n",
    "\n",
    "Set up a Cartesian (cuboid) domain and use Strang splitting for time integration. These choices define the simulation skeleton before any particle data are attached."
   ]
  },
  {
   "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": [
    "### Step 2: Define the Beltrami Background\n",
    "\n",
    "Next, define velocity, pressure, and density functions for the background equilibrium. We pass these to `GenericCartesianFluidEquilibrium`, which will later be used to initialize the particle state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct Beltrami flow\n",
    "import numpy as np\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",
    "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": [
    "### Step 3: Choose Grid and de Rham Options\n",
    "\n",
    "Even for particle-based dynamics, Struphy uses spline-based projections for some fields. Here we define a tensor-product grid and matching de Rham settings."
   ]
  },
  {
   "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": [
    "### Step 4: Instantiate the `PressureLessSPH` Simulation\n",
    "\n",
    "Create a lightweight model instance and combine it with the environment, domain, and discretization settings from the previous steps."
   ]
  },
  {
   "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": [
    "### Step 5: Configure Marker Loading and Saving\n",
    "\n",
    "Now set particle-marker parameters. In this example we:\n",
    "\n",
    "- load `Np=1000` particles\n",
    "- keep default weight construction\n",
    "- apply reflective boundaries in $x,y$ and periodic in $z$\n",
    "- save all markers via `n_markers=1.0`\n",
    "\n",
    "We also keep sorting simple with one box per direction for this first run."
   ]
  },
  {
   "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",
    "sorting_params = SortingParameters(boxes_per_dim=(1, 1, 1))\n",
    "saving_params = SavingParameters(n_markers=1.0)\n",
    "\n",
    "model.cold_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": "13",
   "metadata": {},
   "source": [
    "### Step 6: Set Propagator Options\n",
    "\n",
    "Configure the particle pushers. The key detail is to pass the Beltrami pressure potential to `push_v`, so the velocity update uses the intended force field."
   ]
  },
  {
   "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": [
    "### Step 7: Apply Initial Conditions\n",
    "\n",
    "Attach the Beltrami equilibrium as the background state for the cold fluid species. No additional perturbation is added in this tutorial run."
   ]
  },
  {
   "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": [
    "### Step 8: Run, Post-Process, and Visualize\n",
    "\n",
    "Execute the simulation pipeline in order:\n",
    "\n",
    "1. `run()` to advance particles in time\n",
    "2. `pproc()` to build post-processed outputs\n",
    "3. `load_plotting_data()` to bring diagnostics into memory\n",
    "\n",
    "Then plot marker trajectories to inspect how the Beltrami-driven flow evolves."
   ]
  },
  {
   "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": [
    "### Step 9: Repeat with Structured Marker Loading\n",
    "\n",
    "To compare loading strategies, run a second simulation in a separate folder (`sim_2`) so outputs stay isolated from the first run."
   ]
  },
  {
   "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, use tessellation-based loading so markers start on a regular grid-like pattern. This gives a useful contrast to random loading when you inspect the trajectory plots."
   ]
  },
  {
   "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",
    "sorting_params = SortingParameters(boxes_per_dim=(16, 16, 1))\n",
    "saving_params = SavingParameters(n_markers=1.0)\n",
    "\n",
    "model.cold_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=0.5\n",
    ")"
   ]
  },
  {
   "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"
   ]
  }
 ],
 "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
}